Repository: twallema/pySODM (language: python)

Issue:
Order of non-stratified vs. stratified parameters in the `integrate` function shouldn't matter when initializing a model
<!-- If you have a question or issue about the `covid-19 model`, please use following template. -->

## Suggestion

Consider an SIR model stratified into two `age groups` (0-20 yo, 20-120 yo) with two parameters: `beta`, the per-contact chance of infection, and `gamma`, the recovery rate of the infected individuals. In the example `beta` is stratified and `gamma` is not.

```
class SIRstratified(ODEModel):

    # state variables and parameters
    state_names = ['S', 'I', 'R']
    parameter_names = ['gamma']
    parameters_stratified_names = ['beta']
    stratification_names = ['age_groups']

    @staticmethod
    def integrate(t, S, I, R, beta, gamma):
        """Basic SIR model"""
        # Model equations
        N = S + I + R
        dS = -beta*S*I/N
        dI = beta*S*I/N - gamma*I
        dR = gamma*I
        return dS, dI, dR

parameters = {"gamma": 0.2, "beta": np.array([0.8, 0.9])}
initial_states = {"S": [600_000 - 20, 400_000 - 10], "I": [20, 10], "R": [0, 0]}
coordinates = {'age_groups': ['0-20','20-120']}
model = SIRstratified(initial_states, parameters, coordinates=coordinates)
```
Providing the parameter `beta` before `gamma` yields the following error:
```
ValueError: The parameters in the 'integrate' function definition do not match the parameter_names + parameters_stratified_names + stratification: ['beta', 'gamma'] vs ['gamma', 'beta']
```
Ideally, the order in which the parameters are given should not matter when initializing the model.



Discussion / hints:

Fix patch:
diff --git a/docs/introduction.md b/docs/introduction.md
index 1dc9075..0d2c6d4 100644
--- a/docs/introduction.md
+++ b/docs/introduction.md
@@ -51,6 +51,8 @@ The following features will be implemented in future versions of pySODM,
 
 - Currently, a *time-dependent parameter function* takes only the parameter to be varied `param` as an input. Perhaps the entire parameter dictionary should be given instead for more flexibility?
 
+- Currently, a parameter can only have one *time-depedent parameter function*. Perhaps a list containing multiple function can be supplied? A workaround is to split the parameter in two parameters and multiply them in the integrate function. Very low priority.
+
 ### Versions
 
 - version 0.1 (2022-12-23) 
diff --git a/docs/models.md b/docs/models.md
index 9793c7b..1fcb18f 100644
--- a/docs/models.md
+++ b/docs/models.md
@@ -1,6 +1,6 @@
 # Models
 
-Standard usage of pySODM involves building an ODE or SDE model. For practical examples of initializing models, we refer to the [tutorials](workflow.md).
+Standard usage of pySODM involves building an ODE or SDE model. For practical examples of initializing models, we refer to the [tutorials](quickstart.md).
 
 ## base.py
 
@@ -12,7 +12,7 @@ The ODEModel class inherits several attributes from the model class defined by t
 
 * **state_names** (list) - Contains the names (type: str) of the states.
 * **parameter_names** (list) - Contains the names (type: str) of the non-stratified parameters. Non-stratified parameters are not subject to input checks and can thus be of any datatype/size.
-* **integrate** (function) - Function to compute the differentials of every model state. The integrate function must have the timestep `t` as its first input, followed by the `state_names`, then the `parameter_names` and then the `parameter_stratified_names`, all in the correct order. The integrate function must return a differential for every model state of the correct size and in the same order as `state_names`. The integrate function must be a static method (include `@staticmethod`).
+* **integrate** (function) - Function to compute the differentials of every model state. The integrate function must have the timestep `t` as its first input, followed by the `state_names`, `parameter_names` and `parameter_stratified_names` (their order is not important). The integrate function must return a differential for every model state of the correct size and **in the same order as `state_names`**. The integrate function must be a static method (include `@staticmethod`).
 * **stratification_names** (list) - optional - Contains the names of the stratification axes. The names given here become the dimensions in the simulation output `xarray` Dataset.
 * **parameter_stratified_names** (list) - optional - Contains the names of the stratified parameters. Stratified parameters are subject to input checks. They must be a list or a 1D np.ndarray with a length equal to the number of coordinates of the stratification axis. 
     * For one stratification axes: list contains strings - `['stratpar_1', 'stratpar_2']`
@@ -58,7 +58,7 @@ The SDEModel class inherits several attributes from the model class defined by t
 
 * **state_names** (list) - Contains the names (type: str) of the states.
 * **parameter_names** (list) - Contains the names (type: str) of the non-stratified parameters. Non-stratified parameters are not subject to input checks and can thus be of any datatype/size.
-* **compute_rates** (function) - Function returning the rates of transitioning between the model states. `compute_rates()` must have the timestep `t` as its first input, followed by the `state_names`, then the `parameter_names` and then the `parameter_stratified_names`, all in the correct order. `compute_rates()` must be a static method (include `@staticmethod`). The output of `compute_rates()` must be a dictionary. Its keys must be valid model states, a rate is only needed for the states undergoing a transitioning. The corresponding values must be a list containing the rates of the possible transitionings of the state. In this way, a model state can have multiple transitionings.
+* **compute_rates** (function) - Function returning the rates of transitioning between the model states. `compute_rates()` must have the timestep `t` as its first input, followed by the `state_names`, `parameter_names` and the `parameter_stratified_names` (their order is not important). `compute_rates()` must be a static method (include `@staticmethod`). The output of `compute_rates()` must be a dictionary. Its keys must be valid model states, a rate is only needed for the states undergoing a transitioning. The corresponding values must be a list containing the rates of the possible transitionings of the state. In this way, a model state can have multiple transitionings.
 * **apply_transitionings** (function) - Function to update the states with the number of drawn transitionings. `apply_transitionings()` must have the timestep `t` as its first input, followed by the solver timestep `tau`, follwed by the dictionary containing the transitionings `transitionings`, then followed by the model states and parameters similarily to `compute_rates()`.
 * **stratification_names** (list) - optional - Contains the names of the stratification axes. The names given here become the dimensions in the simulation output `xarray` Dataset.
 * **parameter_stratified_names** (list) - optional - Contains the names of the stratified parameters. Stratified parameters are subject to input checks. They must be a list or a 1D np.ndarray with a length equal to the number of coordinates of the stratification axis. 
diff --git a/docs/quickstart.md b/docs/quickstart.md
index c3103ad..0bd5543 100644
--- a/docs/quickstart.md
+++ b/docs/quickstart.md
@@ -40,10 +40,14 @@ model = SIR(states={'S': 1000, 'I': 1}, parameters={'beta': 0.35, 'gamma': 5})
 Simulate the model using the `sim()` method. The solver method and tolerance, as well as the timesteps in the output can be tailored. pySODM supports the use of `datetime` types as timesteps.
 
 ```python
+# Timesteps
 out = model.sim(121)
+
+# Dates
+out = model.sim(['2022-12-01', '2023-05-01'])
 ```
 
-Results are sent to an `xarray.Dataset`, for more information on indexing and selecting data, [see](https://docs.xarray.dev/en/stable/user-guide/indexing.html).
+Results are sent to an `xarray.Dataset`, for more information on indexing and selecting data using `xarray`, [see](https://docs.xarray.dev/en/stable/user-guide/indexing.html).
 ```bash
 <xarray.Dataset>
 Dimensions:  (time: 122)
@@ -53,10 +57,10 @@ Data variables:
     S        (time) float64 1e+03 999.6 999.2 998.7 ... 287.2 287.2 287.2 287.2
     I        (time) float64 1.0 1.161 1.348 1.565 ... 0.1455 0.1316 0.1192
     R        (time) float64 0.0 0.2157 0.4662 0.7569 ... 713.6 713.7 713.7 713.7
-
-
 ```
 
+![SIR_time](/_static/figs/workflow/SIR_time.png)
+
 ## Stratifications: adding age groups to the SIR model
 
 To transform our SIR model's states in 1D vectors, referred to as a *stratification* throughout the documentation, add the `stratification_names` keyword to the model declaration.
@@ -144,7 +148,7 @@ Data variables:
 
 ### Time-dependent parameter function
 
-Parameters can also be varied through the course of one simulation using a *time-dependent parameter function*, which can be an arbitrarily complex function. A generic time-dependent parameter function has the timestep, the dictionary of model states and the value of the varied parameter as its inputs. Additionally, the function can take any number of arguments.
+Parameters can also be varied through the course of one simulation using a *time-dependent parameter function*, which can be an arbitrarily complex function. A generic time-dependent parameter function has the timestep, the dictionary of model states and the value of the varied parameter as its inputs. Additionally, the function can take any number of arguments (including other model parameters).
 
 ```python
 def vary_my_parameter(t, states, param, an_additional_parameter):
diff --git a/docs/workflow.md b/docs/workflow.md
index 6c83ec8..0db0479 100644
--- a/docs/workflow.md
+++ b/docs/workflow.md
@@ -83,15 +83,15 @@ or schematically,
 
 <img src="./_static/figs/workflow/flowchart_SIR.png" width="500" />
 
-The model has three states: 1) The number of individuals susceptible to the disease (S), 2) the number of infectious individuals (I), and 3) the number of removed individuals (R). The model has two parameters: 1) `beta`, the rate of transmission and, 2) `gamma`, the duration of infectiousness. Building a model is based on the class inheritance, the user must first load the `ODEModel` or `SDEModel` class from  `~/src/models/base.by`. Then, the user must define his/her own class which must contain (minimally),
+The model has three states: 1) The number of individuals susceptible to the disease (S), 2) the number of infectious individuals (I), and 3) the number of removed individuals (R). The model has two parameters: 1) `beta`, the rate of transmission and, 2) `gamma`, the duration of infectiousness. Building a model is based on the class inheritance, the user must first load the `ODEModel` class from  `~/src/models/base.by`. Then, the user must define his/her own class which must contain (minimally),
 - A list containing the state names `state_names`,
 - A list containing the parameter names `parameter_names`,
 - An `integrate()` function where the differentials of the model are computed,
 
 and take the `ODEModel` class as its input. Checkout the documentation of the ODEModel class [here](models.md). There are some important formatting requirements to the integrate function, which are verified when the model is initialized,
 1. The integrate function must have the timestep `t` as its first input
-2. The timestep `t` is first followed by the states, then the parameters in the correct order
-3. The integrate function must return a differential for every model state
+2. The model states and parameters must also be given as input, their order does not make a difference
+3. The integrate function must return a differential for every model state, arranged in the same order as the state names defined in `state_names`
 4. The integrate function must be a static method (include `@staticmethod`)
 
 ```
@@ -151,29 +151,6 @@ plt.close()
 
 ![SIR_date](/_static/figs/workflow/SIR_date.png)
 
-The `sim()` method can also be run with timesteps as coordinates of the time axis (int/float). By convention, the name of the time axis in the `xarray` output is equal to `date` when using dates and `time` when using timesteps.
-```
-# Simulate from t=0 until t=121
-out = model.sim([0, 121])
-
-# Is equivalent to:
-out = model.sim(121)
-
-# But now the time axis is named 'time'
-fig,ax=plt.subplots(figsize=(12,4))
-ax.plot(out['time'], out['S'], color='green', label='Susceptible')
-ax.plot(out['time'], out['I'], color='red', label='Infectious')
-ax.plot(out['time'], out['R'], color='black', label='Removed')
-ax.xaxis.set_major_locator(plt.MaxNLocator(3))
-ax.legend()
-plt.show()
-plt.close()
-```
-
-![SIR_time](/_static/figs/workflow/SIR_time.png)
-
-The user can acces the integration methods and relative tolerance of `scipy.solve_ivp()` by supplying the `method` and `rtol` arguments to the [`sim()` function](models.md). The user can determine the output timestep frequency by changing the optional `output_timestep` argument.
-
 ### Calibrating the model
 
 #### The posterior probability 
diff --git a/src/pySODM/models/base.py b/src/pySODM/models/base.py
index 9966c0e..90bf875 100644
--- a/src/pySODM/models/base.py
+++ b/src/pySODM/models/base.py
@@ -13,7 +13,7 @@
 from functools import partial
 from scipy.integrate import solve_ivp
 from pySODM.models.utils import int_to_date
-from pySODM.models.validation import validate_draw_function, validate_simulation_time, validate_stratifications, validate_time_dependent_parameters, validate_ODEModel, validate_SDEModel
+from pySODM.models.validation import merge_parameter_names_parameter_stratified_names, validate_draw_function, validate_simulation_time, validate_stratifications, validate_time_dependent_parameters, validate_ODEModel, validate_SDEModel, check_duplicates
 
 class SDEModel:
     """
@@ -24,11 +24,11 @@ class SDEModel:
     To initialise the model, provide following inputs:
 
     states : dictionary
-        contains the initial values of all non-zero model states
-        e.g. {'S': N, 'E': np.ones(n_stratification)} with N being the total population and n_stratifications the number of stratified layers
-        initialising zeros is thus not required
+        contains the initial values of all non-zero model states, f.i. for an SEIR model,
+        e.g. {'S': [1000, 1000, 1000], 'E': [1, 1, 1]
+        initialising zeros is not required
     parameters : dictionary
-        containing the values of all parameters (both stratified and not)
+        values of all parameters (both stratified and not)
     time_dependent_parameters : dictionary, optional
         Optionally specify a function for time-dependent parameters. The
         signature of the function should be ``fun(t, states, param, ...)`` taking
@@ -52,6 +52,18 @@ def __init__(self, states, parameters, coordinates=None, time_dependent_paramete
         self.coordinates = coordinates
         self.time_dependent_parameters = time_dependent_parameters
 
+        # Merge parameter_names and parameter_stratified_names
+        self.parameter_names_merged = merge_parameter_names_parameter_stratified_names(self.parameter_names, self.parameter_stratified_names)
+
+        # Duplicates in lists containing names of states/parameters/stratified parameters/stratifications?
+        check_duplicates(self.state_names, 'state_names')
+        try:
+            check_duplicates(self.parameter_names_merged, 'parameter_names + parameter_stratified_names')
+        except:
+            check_duplicates(self.parameter_names, 'parameter_names')
+        if self.stratification_names:
+            check_duplicates(self.stratification_names, 'stratification_names')
+
         # Validate and compute the stratification sizes
         self.stratification_size = validate_stratifications(self.stratification_names, self.coordinates)
 
@@ -62,9 +74,9 @@ def __init__(self, states, parameters, coordinates=None, time_dependent_paramete
             self._function_parameters = []
         
         # Validate the model
-        self.initial_states, self.parameters, self._n_function_params = validate_SDEModel(states, parameters, coordinates, self.stratification_size, self.state_names,
-                                                                                            self.parameter_names, self.parameter_stratified_names, self._function_parameters,
-                                                                                            self.compute_rates, self.apply_transitionings)
+        self.initial_states, self.parameters, self._n_function_params, self._extra_params = validate_SDEModel(states, parameters, coordinates, self.stratification_size, self.state_names,
+                                                                                                              self.parameter_names, self.parameter_stratified_names, self._function_parameters,
+                                                                                                              self.compute_rates, self.apply_transitionings)
 
     # Overwrite integrate class
     @staticmethod
@@ -270,20 +282,17 @@ def func(t, y, pars={}):
                         func_params = {key: params[key] for key in self._function_parameters[i]}
                         params[param] = param_func(date, states, pars[param], **func_params)
 
-                # ----------------------------------
-                # construct list of model parameters
-                # ----------------------------------
+                # -------------------------------------------------------------
+                #  throw parameters of TDPFs out of model parameters dictionary
+                # -------------------------------------------------------------
 
-                if self._n_function_params > 0:
-                    model_pars = list(params.values())[:-self._n_function_params]
-                else:
-                    model_pars = list(params.values())
+                params = {k:v for k,v in params.items() if ((k not in self._extra_params) or (k in self.parameter_names_merged))}
 
                 # -------------
                 # compute rates
                 # -------------
 
-                rates = self.compute_rates(t, *y_reshaped, *model_pars)
+                rates = self.compute_rates(t, **states, **params)
 
                 # --------------
                 # Gillespie step
@@ -298,7 +307,7 @@ def func(t, y, pars={}):
                 # update states 
                 # -------------
 
-                dstates = self.apply_transitionings(t, tau, transitionings, *y_reshaped, *model_pars)
+                dstates = self.apply_transitionings(t, tau, transitionings, **states, **params)
 
                 return np.array(dstates).flatten(), tau
 
@@ -489,11 +498,11 @@ class ODEModel:
     To initialise the model, provide following inputs:
 
     states : dictionary
-        contains the initial values of all non-zero model states
-        e.g. {'S': N, 'E': np.ones(n_stratification)} with N being the total population and n_stratifications the number of stratified layers
-        initialising zeros is thus not required
+        contains the initial values of all non-zero model states, f.i. for an SEIR model,
+        e.g. {'S': [1000, 1000, 1000], 'E': [1, 1, 1]
+        initialising zeros is not required
     parameters : dictionary
-        containing the values of all parameters (both stratified and not)
+        values of all parameters (both stratified and not)
     time_dependent_parameters : dictionary, optional
         Optionally specify a function for time-dependent parameters. The
         signature of the function should be ``fun(t, states, param, ...)`` taking
@@ -518,6 +527,20 @@ def __init__(self, states, parameters, coordinates=None, time_dependent_paramete
         self.coordinates = coordinates
         self.time_dependent_parameters = time_dependent_parameters
 
+        # Merge parameter_names and parameter_stratified_names
+        self.parameter_names_merged = merge_parameter_names_parameter_stratified_names(self.parameter_names, self.parameter_stratified_names)
+
+        # Duplicates in lists containing names of states/parameters/stratified parameters/stratifications?
+        check_duplicates(self.state_names, 'state_names')
+        try:
+            check_duplicates(self.parameter_names_merged, 'parameter_names + parameter_stratified_names')
+        except:
+            check_duplicates(self.parameter_names, 'parameter_names')
+        if self.stratification_names:
+            check_duplicates(self.stratification_names, 'stratification_names')
+        if self.state_2d:
+            check_duplicates(self.state_2d, 'state_2d')
+
         # Validate and compute the stratification sizes
         self.stratification_size = validate_stratifications(self.stratification_names, self.coordinates)
 
@@ -528,9 +551,9 @@ def __init__(self, states, parameters, coordinates=None, time_dependent_paramete
             self._function_parameters = []
 
         # Validate the model
-        self.initial_states, self.parameters, self._n_function_params = validate_ODEModel(states, parameters, coordinates, self.stratification_size, self.state_names,
-                                                                                            self.parameter_names, self.parameter_stratified_names, self._function_parameters,
-                                                                                            self._create_fun, self.integrate, self.state_2d)
+        self.initial_states, self.parameters, self._n_function_params, self._extra_params = validate_ODEModel(states, parameters, coordinates, self.stratification_size, self.state_names,
+                                                                                                            self.parameter_names, self.parameter_stratified_names, self._function_parameters,
+                                                                                                            self._create_fun, self.integrate, self.state_2d)
 
         # Experimental: added to use 2D states for the Economic IO model
         if self.state_2d:
@@ -558,13 +581,13 @@ def func(t, y, pars={}):
                 for size in self.stratification_size:
                     size_lst.append(size)
                 y_reshaped = y.reshape(tuple(size_lst))
-                state_params = dict(zip(self.state_names, y_reshaped))
+                states = dict(zip(self.state_names, y_reshaped))
             else:
                 # incoming y -> different reshape for 1D vs 2D variables  (2)
                 y_1d, y_2d = np.split(y, [self.split_point])
                 y_1d = y_1d.reshape(((len(self.state_names) - 1), self.stratification_size[0]))
                 y_2d = y_2d.reshape((self.stratification_size[0], self.stratification_size[0]))
-                state_params = state_params = dict(zip(self.state_names, [y_1d,y_2d]))
+                states  = dict(zip(self.state_names, [y_1d,y_2d]))
 
             # --------------------------------------
             # update time-dependent parameter values
@@ -579,26 +602,23 @@ def func(t, y, pars={}):
                     date = t
                 for i, (param, param_func) in enumerate(self.time_dependent_parameters.items()):
                     func_params = {key: params[key] for key in self._function_parameters[i]}
-                    params[param] = param_func(date, state_params, pars[param], **func_params)
+                    params[param] = param_func(date, states, pars[param], **func_params)
 
-            # ----------------------------------
-            # construct list of model parameters
-            # ----------------------------------
+            # -------------------------------------------------------------
+            #  throw parameters of TDPFs out of model parameters dictionary
+            # -------------------------------------------------------------
 
-            if self._n_function_params > 0:
-                model_pars = list(params.values())[:-self._n_function_params]
-            else:
-                model_pars = list(params.values())
+            params = {k:v for k,v in params.items() if ((k not in self._extra_params) or (k in self.parameter_names_merged))}
 
             # -------------------
             # perform integration
             # -------------------
 
             if not self.state_2d:
-                dstates = self.integrate(t, *y_reshaped, *model_pars)
+                dstates = self.integrate(t, **states, **params)
                 return np.array(dstates).flatten()
             else:
-                dstates = self.integrate(t, *y_1d, y_2d, *model_pars)
+                dstates = self.integrate(t, *y_1d, y_2d, **params)
                 return np.concatenate([np.array(state).flatten() for state in dstates])
 
         return func
@@ -614,7 +634,8 @@ def _sim_single(self, time, actual_start_date=None, method='RK23', rtol=5e-3, ou
             for state in self.state_2d:
                 self.initial_states[state] = self.initial_states[state].flatten()
 
-        # Initial conditions must be one long list of values
+        # Scipy can only take flattened list of states
+        # TODO: rearrange y0 in order of self.state_names --> ALREADY DONE IN INITIALIZATION!
         y0 = list(itertools.chain(*self.initial_states.values()))
         while np.array(y0).ndim > 1:
             y0 = list(itertools.chain(*y0))
diff --git a/src/pySODM/models/validation.py b/src/pySODM/models/validation.py
index 6e2503a..bfeb606 100644
--- a/src/pySODM/models/validation.py
+++ b/src/pySODM/models/validation.py
@@ -136,16 +136,16 @@ def validate_parameter_function(func):
     else:
         return keywords[3:]
 
-def validate_time_dependent_parameters(parameter_names, parameters_stratified_names,time_dependent_parameters):
+def validate_time_dependent_parameters(parameter_names, parameter_stratified_names,time_dependent_parameters):
 
     extra_params = []
     all_param_names = parameter_names.copy()
 
-    if parameters_stratified_names:
-        if not isinstance(parameters_stratified_names[0], list):
-            all_param_names += parameters_stratified_names
+    if parameter_stratified_names:
+        if not isinstance(parameter_stratified_names[0], list):
+            all_param_names += parameter_stratified_names
         else:
-            for lst in parameters_stratified_names:
+            for lst in parameter_stratified_names:
                 all_param_names.extend(lst)
             
     for param, func in time_dependent_parameters.items():
@@ -233,8 +233,32 @@ def validate_stratified_parameters(values, name, object_name,i,stratification_si
             )
         )
 
+def check_duplicates(lst, name):
+    """A function raising an error if lst contains duplicates"""
+    seen = set()
+    dupes = [x for x in lst if x in seen or seen.add(x)]
+    if dupes:
+        raise ValueError(
+            f"List '{name}' contains duplicates: {dupes}"
+        )
+
+def merge_parameter_names_parameter_stratified_names(parameter_names, parameter_stratified_names):
+    """ A function to merge the 'parameter_names' and 'parameter_stratified_names' lists"""
+    merged_params = parameter_names.copy()
+    if parameter_stratified_names:
+        if not isinstance(parameter_stratified_names[0], list):
+            if len(parameter_stratified_names) == 1:
+                merged_params += parameter_stratified_names
+            else:
+                for stratified_names in parameter_stratified_names:
+                    merged_params += [stratified_names,]
+        else:
+            for stratified_names in parameter_stratified_names:
+                merged_params += stratified_names
+    return merged_params
+
 def validate_ODEModel(initial_states, parameters, coordinates, stratification_size, state_names, parameter_names,
-                        parameters_stratified_names, _function_parameters, _create_fun, integrate_func, state_2d=None):
+                        parameter_stratified_names, _function_parameters, _create_fun, integrate_func, state_2d=None):
     """
     This does some basic validation of the model + initialization:
 
@@ -245,70 +269,51 @@ def validate_ODEModel(initial_states, parameters, coordinates, stratification_si
 
     2) Validation of the actual initialization with initial values for the
     states and parameter values.
-    TODO: For now, we require that those are passed in the exact same
-    order, but this requirement could in principle be relaxed, if we ensure
-    to pass the states and parameters as keyword arguments and not as
-    positional arguments to the `integrate` function.
+
     """
 
     # Validate Model class definition (the integrate function)
+    # First argument should always be 't'
     sig = inspect.signature(integrate_func)
     keywords = list(sig.parameters.keys())
     if keywords[0] != "t":
         raise ValueError(
             "The first argument of the integrate function should always be 't'"
         )
-    else:
-        start_index = 1
-
-    # Get names of states and parameters that follow after 't' or 't' and 'l'
-    N_states = len(state_names)
-    integrate_states = keywords[start_index : start_index + N_states]
-    if integrate_states != state_names:
+    # Add parameters and stratified parameters to one list: specified_params
+    specified_params = merge_parameter_names_parameter_stratified_names(parameter_names, parameter_stratified_names)
+    # Check if all the states and parameters are present
+    if set(specified_params + state_names) != set(keywords[1:]):
+        # Extract redundant and missing parameters
+        redundant_all = set(keywords[1:]).difference(set(specified_params + state_names))
+        missing_all = set(specified_params + state_names).difference(set(keywords[1:]))
+        # Let's split the missing variables in parameters/states for extra clarity
+        for state_name in state_names:
+            missing_states = [name for name in missing_all if name in state_names]
+        for parameter_name in parameter_names:
+            missing_parameters = [name for name in missing_all if name in parameter_names]
         raise ValueError(
-            "The states in the 'integrate' function definition do not match "
-            "the state_names: {0} vs {1}".format(integrate_states, state_names)
+            "The provided state names and parameters don't match the parameters and states of the integrate function. "
+            "Missing parameters: {0}. Missing states: {1}. "
+            "Redundant: {2}. ".format(missing_parameters, missing_states, list(redundant_all))
         )
-    integrate_params = keywords[start_index + N_states :]
-    specified_params = parameter_names.copy()
-
-    if parameters_stratified_names:
-        if not isinstance(parameters_stratified_names[0], list):
-            if len(parameters_stratified_names) == 1:
-                specified_params += parameters_stratified_names
-            else:
-                for stratified_names in parameters_stratified_names:
-                    specified_params += [stratified_names,]
-        else:
-            for stratified_names in parameters_stratified_names:
-                specified_params += stratified_names
-
-    if integrate_params != specified_params:
-        raise ValueError(
-            "The parameters in the 'integrate' function definition do not match "
-            "the parameter_names + parameters_stratified_names + stratification: "
-            "{0} vs {1}".format(integrate_params, specified_params)
-        )
-
     # additional parameters from time-dependent parameter functions
     # are added to specified_params after the above check
-
     if _function_parameters:
-        extra_params = [item for sublist in _function_parameters for item in sublist]
-
-        # TODO check that it doesn't duplicate any existing parameter (completed?)
-        # Line below removes duplicate arguments in time dependent parameter functions
-        extra_params = OrderedDict((x, True) for x in extra_params).keys()
-        specified_params += extra_params
+        _extra_params = [item for sublist in _function_parameters for item in sublist]
+        # Remove duplicate arguments in time dependent parameter functions
+        _extra_params = OrderedDict((x, True) for x in _extra_params).keys()
+        specified_params += _extra_params
         len_before = len(specified_params)
         # Line below removes duplicate arguments with integrate defenition
         specified_params = OrderedDict((x, True) for x in specified_params).keys()
         len_after = len(specified_params)
         # Line below computes number of integrate arguments used in time dependent parameter functions
         n_duplicates = len_before - len_after
-        _n_function_params = len(extra_params) - n_duplicates
+        _n_function_params = len(_extra_params) - n_duplicates
     else:
         _n_function_params = 0
+        _extra_params = []
 
     # Validate the params
     if set(parameters.keys()) != set(specified_params):
@@ -328,23 +333,23 @@ def validate_ODEModel(initial_states, parameters, coordinates, stratification_si
             )
 
     # the size of the stratified parameters
-    if parameters_stratified_names:
+    if parameter_stratified_names:
         i = 0
-        if not isinstance(parameters_stratified_names[0], list):
-            if len(parameters_stratified_names) == 1:
-                for param in parameters_stratified_names:
+        if not isinstance(parameter_stratified_names[0], list):
+            if len(parameter_stratified_names) == 1:
+                for param in parameter_stratified_names:
                     validate_stratified_parameters(
                             parameters[param], param, "stratified parameter",i,stratification_size,coordinates
                         )
                 i = i + 1
             else:
-                for param in parameters_stratified_names:
+                for param in parameter_stratified_names:
                     validate_stratified_parameters(
                             parameters[param], param, "stratified parameter",i,stratification_size,coordinates
                         )
                 i = i + 1
         else:
-            for stratified_names in parameters_stratified_names:
+            for stratified_names in parameter_stratified_names:
                 for param in stratified_names:
                     validate_stratified_parameters(
                         parameters[param], param, "stratified parameter",i,stratification_size,coordinates
@@ -388,11 +393,11 @@ def validate_ODEModel(initial_states, parameters, coordinates, stratification_si
             raise ValueError(
                 "The return value of the integrate function does not have the correct length."
             )
-    
-    return initial_states, parameters, _n_function_params
+
+    return initial_states, parameters, _n_function_params, list(_extra_params)
 
 def validate_SDEModel(initial_states, parameters, coordinates, stratification_size, state_names, parameter_names,
-                        parameters_stratified_names, _function_parameters, compute_rates_func, apply_transitionings_func):
+                        parameter_stratified_names, _function_parameters, compute_rates_func, apply_transitionings_func):
     """
     This does some basic validation of the model + initialization:
 
@@ -403,10 +408,7 @@ def validate_SDEModel(initial_states, parameters, coordinates, stratification_si
 
     2) Validation of the actual initialization with initial values for the
     states and parameter values.
-    TODO: For now, we require that those are passed in the exact same
-    order, but this requirement could in principle be relaxed, if we ensure
-    to pass the states and parameters as keyword arguments and not as
-    positional arguments to the `integrate` function.
+
     """
 
     #############################
@@ -422,57 +424,42 @@ def validate_SDEModel(initial_states, parameters, coordinates, stratification_si
         raise ValueError(
             "The first argument of the 'compute_rates' function should always be 't'"
         )
-    else:
-        start_index = 1
-
-    # Get names of states and parameters that follow after 't'
-    N_states = len(state_names)
-    compute_rates_states = keywords[start_index : start_index + N_states]
-    if compute_rates_states != state_names:
-        raise ValueError(
-            "The states in the 'compute_rates' function definition do not match "
-            "the provided 'state_names': {0} vs {1}".format(compute_rates_states, state_names)
-        )
-    compute_rates_params = keywords[start_index + N_states :]
-    specified_params = parameter_names.copy()
-
-    if parameters_stratified_names:
-        if not isinstance(parameters_stratified_names[0], list):
-            if len(parameters_stratified_names) == 1:
-                specified_params += parameters_stratified_names
-            else:
-                for stratified_names in parameters_stratified_names:
-                    specified_params += [stratified_names,]
-        else:
-            for stratified_names in parameters_stratified_names:
-                specified_params += stratified_names
-
-    if compute_rates_params != specified_params:
+    # Add parameters and stratified parameters to one list: specified_params
+    specified_params = merge_parameter_names_parameter_stratified_names(parameter_names, parameter_stratified_names)
+    # Check if all the states and parameters are present
+    if set(specified_params + state_names) != set(keywords[1:]):
+        # Extract redundant and missing parameters
+        redundant_all = set(keywords[1:]).difference(set(specified_params + state_names))
+        missing_all = set(specified_params + state_names).difference(set(keywords[1:]))
+        # Let's split the missing variables in parameters/states for extra clarity
+        for state_name in state_names:
+            missing_states = [name for name in missing_all if name in state_names]
+        for parameter_name in parameter_names:
+            missing_parameters = [name for name in missing_all if name in parameter_names]
         raise ValueError(
-            "The parameters in the 'compute_rates' function definition do not match "
-            "the provided 'parameter_names' + 'parameters_stratified_names': "
-            "{0} vs {1}".format(compute_rates_params, specified_params)
+            "The provided state names and parameters don't match the parameters and states of the compute_rates function. "
+            "Missing parameters: {0}. Missing states: {1}. "
+            "Redundant: {2}. ".format(missing_parameters, missing_states, list(redundant_all))
         )
-
+    # Save a copy to validate the model later on
+    specified_params_wo_TDPF_pars = specified_params.copy()
     # additional parameters from time-dependent parameter functions
     # are added to specified_params after the above check
-
     if _function_parameters:
-        extra_params = [item for sublist in _function_parameters for item in sublist]
-
-        # TODO check that it doesn't duplicate any existing parameter (completed?)
-        # Line below removes duplicate arguments in time dependent parameter functions
-        extra_params = OrderedDict((x, True) for x in extra_params).keys()
-        specified_params += extra_params
+        _extra_params = [item for sublist in _function_parameters for item in sublist]
+        # Remove duplicate arguments in time dependent parameter functions
+        _extra_params = OrderedDict((x, True) for x in _extra_params).keys()
+        specified_params += _extra_params
         len_before = len(specified_params)
         # Line below removes duplicate arguments with integrate defenition
         specified_params = OrderedDict((x, True) for x in specified_params).keys()
         len_after = len(specified_params)
         # Line below computes number of integrate arguments used in time dependent parameter functions
         n_duplicates = len_before - len_after
-        _n_function_params = len(extra_params) - n_duplicates
+        _n_function_params = len(_extra_params) - n_duplicates
     else:
         _n_function_params = 0
+        _extra_params = []
 
     # Validate the params
     if set(parameters.keys()) != set(specified_params):
@@ -502,58 +489,41 @@ def validate_SDEModel(initial_states, parameters, coordinates, stratification_si
         raise ValueError(
             "The third argument of the 'apply_transitionings' function should always be 'transitionings'"
         )
-    else:
-        start_index = 3
 
-    # Get names of states and parameters that follow after 't'
-    N_states = len(state_names)
-    apply_transitionings_states = keywords[start_index : start_index + N_states]
-    if apply_transitionings_states != state_names:
+    # Add parameters and stratified parameters to one list: specified_params
+    specified_params = merge_parameter_names_parameter_stratified_names(parameter_names, parameter_stratified_names)
+    # Check if all the states and parameters are present
+    if set(specified_params + state_names) != set(keywords[3:]):
+        # Extract redundant and missing parameters
+        redundant_all = set(keywords[3:]).difference(set(specified_params + state_names))
+        missing_all = set(specified_params + state_names).difference(set(keywords[3:]))
+        # Let's split the missing variables in parameters/states for extra clarity
+        for state_name in state_names:
+            missing_states = [name for name in missing_all if name in state_names]
+        for parameter_name in parameter_names:
+            missing_parameters = [name for name in missing_all if name in parameter_names]
         raise ValueError(
-            "The states in the 'apply_transitionings' function definition do not match "
-            "the provided 'state_names': {0} vs {1}".format(apply_transitionings_states, state_names)
+            "The provided state names and parameters don't match the parameters and states of the apply_transitionings function. "
+            "Missing parameters: {0}. Missing states: {1}. "
+            "Redundant: {2}. ".format(missing_parameters, missing_states, list(redundant_all))
         )
-    apply_transitionings_params = keywords[start_index + N_states :]
-    specified_params = parameter_names.copy()
-
-    if parameters_stratified_names:
-        if not isinstance(parameters_stratified_names[0], list):
-            if len(parameters_stratified_names) == 1:
-                specified_params += parameters_stratified_names
-            else:
-                for stratified_names in parameters_stratified_names:
-                    specified_params += [stratified_names,]
-        else:
-            for stratified_names in parameters_stratified_names:
-                specified_params += stratified_names
-
-    if apply_transitionings_params != specified_params:
-        raise ValueError(
-            "The parameters in the 'apply_transitionings' function definition do not match "
-            "the provided 'parameter_names' + 'parameters_stratified_names': "
-            "{0} vs {1}".format(apply_transitionings_params, specified_params)
-        )
-
     # additional parameters from time-dependent parameter functions
     # are added to specified_params after the above check
-
     if _function_parameters:
-        extra_params = [item for sublist in _function_parameters for item in sublist]
-
-        # TODO check that it doesn't duplicate any existing parameter (completed?)
-        # Line below removes duplicate arguments in time dependent parameter functions
-        extra_params = OrderedDict((x, True) for x in extra_params).keys()
-        specified_params += extra_params
+        _extra_params = [item for sublist in _function_parameters for item in sublist]
+        # Remove duplicate arguments in time dependent parameter functions
+        _extra_params = OrderedDict((x, True) for x in _extra_params).keys()
+        specified_params += _extra_params
         len_before = len(specified_params)
         # Line below removes duplicate arguments with integrate defenition
         specified_params = OrderedDict((x, True) for x in specified_params).keys()
         len_after = len(specified_params)
         # Line below computes number of integrate arguments used in time dependent parameter functions
         n_duplicates = len_before - len_after
-        _n_function_params = len(extra_params) - n_duplicates
+        _n_function_params = len(_extra_params) - n_duplicates
     else:
         _n_function_params = 0
-
+        _extra_params = []
     # Validate the params
     if set(parameters.keys()) != set(specified_params):
         raise ValueError(
@@ -573,34 +543,41 @@ def validate_SDEModel(initial_states, parameters, coordinates, stratification_si
             set(parameters.keys()).difference(set(parameters_compute_rates.keys())))
         )
 
-    # After building the list of all model parameters, verify no parameters 't' was used
+    # After building the list of all model parameters, verify no parameters 't'/'tau' or 'transitionings' were used
     if 't' in parameters:
         raise ValueError(
             "Parameter name 't' is reserved for the simulation time.\nPlease verify no model parameters named 't' are present in the model parameters dictionary or in the time-dependent parameter functions."
             )
-
+    if 'tau' in parameters:
+        raise ValueError(
+            "Parameter name 'tau' is reserved for the simulation time.\nPlease verify no model parameters named 't' are present in the model parameters dictionary or in the time-dependent parameter functions."
+            )
+    if 'transitionings' in parameters:
+        raise ValueError(
+            "Parameter name 'transitionings' is reserved for the simulation time.\nPlease verify no model parameters named 't' are present in the model parameters dictionary or in the time-dependent parameter functions."
+            )
     ###############################################################################
     ## Validate the initial_states / stratified params having the correct length ##
     ###############################################################################
 
     # the size of the stratified parameters
-    if parameters_stratified_names:
+    if parameter_stratified_names:
         i = 0
-        if not isinstance(parameters_stratified_names[0], list):
-            if len(parameters_stratified_names) == 1:
-                for param in parameters_stratified_names:
+        if not isinstance(parameter_stratified_names[0], list):
+            if len(parameter_stratified_names) == 1:
+                for param in parameter_stratified_names:
                     validate_stratified_parameters(
                             parameters[param], param, "stratified parameter",i,stratification_size,coordinates
                         )
                 i = i + 1
             else:
-                for param in parameters_stratified_names:
+                for param in parameter_stratified_names:
                     validate_stratified_parameters(
                             parameters[param], param, "stratified parameter",i,stratification_size,coordinates
                         )
                 i = i + 1
         else:
-            for stratified_names in parameters_stratified_names:
+            for stratified_names in parameter_stratified_names:
                 for param in stratified_names:
                     validate_stratified_parameters(
                         parameters[param], param, "stratified parameter",i,stratification_size,coordinates
@@ -641,13 +618,10 @@ def validate_SDEModel(initial_states, parameters, coordinates, stratification_si
 
     # compute_rates
     # ~~~~~~~~~~~~~
-
+    # TODO: Throw out _extra_params here
+    parameters_wo_TDPF_pars = {param: parameters[param] for param in specified_params_wo_TDPF_pars}
     # Call the function with initial values to check if the function returns the right format of dictionary
-    if _n_function_params > 0:
-        parameters_list = list(parameters.values())[:-_n_function_params]
-    else:
-        parameters_list = list(parameters.values())
-    rates = compute_rates_func(10, *initial_states.values(), *parameters_list)
+    rates = compute_rates_func(10, **initial_states, **parameters_wo_TDPF_pars)
     # Check if a dictionary is returned
     if not isinstance(rates, dict):
         raise TypeError("Output of function 'compute_rates' should be of type dictionary")
@@ -674,7 +648,7 @@ def validate_SDEModel(initial_states, parameters, coordinates, stratification_si
     # apply_transitionings
     # ~~~~~~~~~~~~~~~~~~~~
 
-    new_states = apply_transitionings_func(10, 1, rates, *initial_states.values(), *parameters_list)
+    new_states = apply_transitionings_func(10, 1, rates, **initial_states, **parameters_wo_TDPF_pars)
     # Check
     if len(list(new_states)) != len(state_names):
         raise ValueError(f"The number of outputs of function 'apply_transitionings_func' ({len(list(new_states))}) is not equal to the number of states ({len(state_names)})")
@@ -687,4 +661,4 @@ def validate_SDEModel(initial_states, parameters, coordinates, stratification_si
     # Reset initial states
     initial_states={k: v[:] for k, v in initial_states_copy.items()}
 
-    return initial_states, parameters, _n_function_params
+    return initial_states, parameters, _n_function_params, list(_extra_params)
diff --git a/tutorials/influenza_1718/calibration.py b/tutorials/influenza_1718/calibration.py
index 6db9ab3..fe227d8 100644
--- a/tutorials/influenza_1718/calibration.py
+++ b/tutorials/influenza_1718/calibration.py
@@ -54,7 +54,7 @@
 ##############
 
 alpha = 0.03 # Overdispersion of data
-N = 20 # Repeated simulations
+N = 10 # Repeated simulations
 start_calibration = start_date 
 end_calibration = pd.Timestamp('2018-03-01')
 identifier = 'twallema_2018-03-01'
@@ -119,7 +119,7 @@
 
     # Variables
     processes = int(os.getenv('SLURM_CPUS_ON_NODE', mp.cpu_count()/2))
-    n_pso = 50
+    n_pso = 20
     multiplier_pso = 20
     # Define dataset
     data=[df_influenza[start_calibration:end_calibration], ]


Test patch:
diff --git a/src/tests/test_ODEModel.py b/src/tests/test_ODEModel.py
index 844e97f..c73b413 100644
--- a/src/tests/test_ODEModel.py
+++ b/src/tests/test_ODEModel.py
@@ -244,7 +244,7 @@ def test_model_stratified_init_validation():
         SIRstratified(initial_states2, parameters, coordinates=coordinates)
 
     # validate model class itself
-    msg = "The parameters in the 'integrate' function definition do not match"
+    msg = "The provided state names and parameters don't match the parameters and states of the integrate function"
     SIRstratified.parameter_names = ["gamma", "alpha"]
     with pytest.raises(ValueError, match=msg):
         SIRstratified(initial_states, parameters, coordinates=coordinates)
@@ -323,10 +323,11 @@ def compliance_func(t, states, param, gamma):
             return param * gamma
 
     parameters = {"gamma": 0.2, "beta": np.array([0.8, 0.9])}
+    #msg="are used both in a time-dependent parameter function and in the integrate function"
+    #with pytest.raises(ValueError, match=msg):
     model2 = SIRstratified(initial_states, parameters, coordinates=coordinates,
                            time_dependent_parameters={'beta': compliance_func})
-    output2 = model2.sim(time)
-
+    output2 = model2.sim(time)      
 # TDPF on a regular parameter
 def test_TDPF_nonstratified():
     # states, parameters, coordinates
diff --git a/src/tests/test_SDEModel.py b/src/tests/test_SDEModel.py
index 035123c..3a07b62 100644
--- a/src/tests/test_SDEModel.py
+++ b/src/tests/test_SDEModel.py
@@ -146,16 +146,16 @@ def test_model_init_validation():
 
     # validate model class itself
     SIR.state_names = ["S", "R"]
-    with pytest.raises(ValueError, match="The states in the 'compute_rates' function definition do not match the provided 'state_names'"):
+    with pytest.raises(ValueError, match="The provided state names and parameters don't match the parameters and states of the compute_rates function."):
         SIR(initial_states, parameters)
 
     SIR.state_names = ["S", "II", "R"]
-    with pytest.raises(ValueError, match="The states in the 'compute_rates' function definition"):
+    with pytest.raises(ValueError, match="The provided state names and parameters don't match the parameters and states of the compute_rates function."):
         SIR(initial_states, parameters)
 
     SIR.state_names = ["S", "I", "R"]
     SIR.parameter_names = ['beta', 'alpha']
-    with pytest.raises(ValueError, match="The parameters in the 'compute_rates' function"):
+    with pytest.raises(ValueError, match="The provided state names and parameters don't match the parameters and states of the compute_rates function."):
         SIR(initial_states, parameters)
 
     # ensure to set back to correct ones
@@ -271,7 +271,7 @@ def test_model_stratified_init_validation():
         SIRstratified(initial_states2, parameters, coordinates=coordinates)
 
     # validate model class itself
-    msg = "The parameters in the 'compute_rates' function definition do not match"
+    msg = "The provided state names and parameters don't match the parameters and states of the"
     SIRstratified.parameter_names = ["gamma", "alpha"]
     with pytest.raises(ValueError, match=msg):
         SIRstratified(initial_states, parameters, coordinates=coordinates)