Add changes to initial conditions to draw functions

README.md

@@ -62,7 +62,7 @@ Does other simulation software exist in Python? Sure, but most of them hold your
     - Version 0.2.4 (2023-12-04, PR #62)
         > Validated the use of Python 3.11. Efficiency gains in simulation of jump processes. Ommitted dependency on Numba. All changes related to publishing our software manuscript in Journal of Computational Science. Improved nomenclature in model defenition.
     - IN PROGRESS: 0.2.5
-        > Validated the use of Python 3.12. Validated pySODM on macOS Sonoma 14.5. 'draw functions' only have 'parameters' as mandatory input, followed by an arbitrary number of additional parameters (PR #75). Tutorial environment can now be found in `tutorial_env.yml` and was renamed `PYSODM-TUTORIALS` (PR #76). Users can choose when the simulation starts when calibrating a model (PR #92).
+        > Validated the use of Python 3.12. Validated pySODM on macOS Sonoma 14.5. 'draw functions' only have 'parameters' as mandatory input, followed by an arbitrary number of additional parameters (PR #75). Tutorial environment can now be found in `tutorial_env.yml` and was renamed `PYSODM-TUTORIALS` (PR #76). Users can choose when the simulation starts when calibrating a model (PR #92). Draw functions extended to changing the initial condition, be sure to have `parameters` followed by `initial_states` as the first two arguments of a draw function (PR #95).
 - Version 0.1 (2022-12-23, PR #14)
     > Application pySODM to three use cases. Documentation website. Unit tests for ODE, JumpProcess and calibration. 
     - Version 0.1.1 (2023-01-09, PR #20)

docs/enzyme_kinetics.md

@@ -251,14 +251,14 @@ Finally, we can use the *draw functions* to propagate the parameter samples in o
 ```python
 if __name__ == '__main__':
 
-    def draw_fcn(parameters, samples):
+    def draw_fcn(parameters, initial_states, samples):
         # Always draw correlated samples at the SAME INDEX! 
         idx, parameters['Vf_Ks'] = random.choice(list(enumerate(samples['Vf_Ks'])))
         parameters['R_AS'] = samples['R_AS'][idx]
         parameters['R_AW'] = samples['R_AW'][idx]
         parameters['R_Es'] = samples['R_Es'][idx]
         parameters['K_eq'] = samples['K_eq'][idx]
-        return parameters
+        return parameters, initial_states
 
     # Loop over datasets
     for i,df in enumerate(data):
@@ -428,13 +428,13 @@ f = open(os.path.join(os.path.dirname(__file__),'data/username_SAMPLES_2023-06-0
 samples = json.load(f)
 
 # Define draw function
-def draw_fcn(parameters, samples):
+def draw_fcn(parameters, initial_states, samples):
     idx, parameters['Vf_Ks'] = random.choice(list(enumerate(samples['Vf_Ks'])))
     parameters['R_AS'] = samples['R_AS'][idx]
     parameters['R_AW'] = samples['R_AW'][idx]
     parameters['R_Es'] = samples['R_Es'][idx]
     parameters['K_eq'] = samples['K_eq'][idx]
-    return parameters
+    return parameters, initial_states
 ```
 
 A first experiment with the continuous flow reactor was performed using a reaction mixture containing 30 mM D-glucose, 60 mM lauric acid and 28 mM water. The reactants were pumped through the reactor at a flow rate of 0.2 mL/min, resulting in an average residence time of 13.5 minutes. After ten retention times, when the outlet concentration had stabilized, three samples were withdrawn at the outlet. Then, the reactor was cut short by 0.10 m, and again three samples were withdrawn at the outlet. In this way, the reactant profile acrosss the reactor length was obtained. I omit the code to replicate the following figures from this documentation as no new concepts are introduced beyond this point. Our packed-bed model does an adequate job at describing the data.

docs/quickstart.md

@@ -245,21 +245,21 @@ class SIR(JumpProcess):
 
 ### Draw functions
 
-The simulation functions of the `ODE` and `JumpProcess` classes can be used to perform {math}`N` repeated simulations by using the optional argument `N`. A *draw function* can be used to instruct the model on how to alter the value of a model parameters during consecutive model runs, thereby offering a powerful tool for sensitivity analysis and uncertainty propagation.
+The simulation functions of the `ODE` and `JumpProcess` classes can be used to perform {math}`N` repeated simulations by using the optional argument `N`. A *draw function* can be used to instruct the model on how to alter the model parameters and the initial states during consecutive model runs, thereby offering a powerful tool for sensitivity analysis and uncertainty propagation.
 
-Draw functions **always** take the dictionary of model parameters, `parameters` as their first argument, input checks are used to ensure you provide the correct name and type. This can be followed an arbitrary number of user-defined parameters, which are supplied to the `sim()` function by using its `draw_function_kwargs` argument. The output of a draw function is **always** the dictionary of model parameters, without alterations to the dictionaries keys (no new parameters introduced or parameters deleted). In the example below, we attempt to draw a model parameter `gamma` randomly from 0 to 5,
+Draw functions **always** take the dictionary of model `parameters` as their first argument and the dictionary of `initial_states` as their second argument. This can be followed an arbitrary number of user-defined parameters, which must be supplied to the `sim()` function through the `draw_function_kwargs` argument. The output of a draw function is **always** the dictionary of model `parameters` and the dictionary of `initial states`, without alterations to the dictionaries keys (no new introductions or deletions). In the example below, we attempt to draw a model parameter `gamma` randomly from 0 to 5 and don't make alterations to the initial states,
 
 ```python
 # make an example of a dictionary containing samples of a parameter `gamma`
 samples = np.random.uniform(low=0, high=5, n=100)
 
 # define a 'draw function'
-def draw_function(parameters, samples):
-    """ randomly selects a sample of 'gamma' from the provided dictionary of samples and
-        assigns it to the dictionary of model parameters
+def draw_function(parameters, initial_states, samples):
+    """ randomly selects a sample of 'gamma' from the provided dictionary of `samples` and
+        assigns it to the dictionary of model `parameters`
     """
     parameters['gamma'] = np.random.choice(samples)
-    return parameters
+    return parameters, initial_states
 
 # simulate 10 trajectories, exploit 10 cores to speed up the computation
 out = model.sim(121, N=10, draw_function=draw_function, draw_function_kwargs={'samples': samples}, processes=10)
@@ -281,18 +281,18 @@ Data variables:
     R           (draws, age_groups, time) float64 0.0 0.3439 ... 684.0 684.0
 ```
 
-This example can also be coded up by drawing the random values within the *draw function*,
+This example can also be coded up by drawing the random samples of `gamma` within the *draw function*,
 
 ```python
 # define a 'draw function'
-def draw_function(parameters):
+def draw_function(parameters, initial_states):
     parameters['gamma'] = np.random.uniform(low=0, high=5)
-    return parameters
+    return parameters, initial_states
 
 # simulate the model
 out = model.sim(121, N=10, draw_function=draw_function)
 ```
-**_NOTE_** Internally, a call to `draw_function` is made within the `sim()` function, where it is given the dictionary of model parameters `parameters`, followed by the `draw_function_kwargs`.
+**_NOTE_** Internally, a call to `draw_function` is made within the `sim()` function, where it is given the dictionary of model `parameters`, `initial_states`, followed by the `draw_function_kwargs`.
 
 ### Time-dependent parameter functions
 

docs/workflow.md

@@ -347,21 +347,21 @@ plt.close()
 
 Next, let's visualize how well our simple SIR model fits the data. To this end, we'll simulate the model a number of times, and we'll update the value of `beta` with a sample from its posterior probability distribution obtained using the MCMC. Then, we'll add the noise introduced by observing the epidemiological process to each simulated model trajectory using `add_negative_binomial()`. Finally, we'll visualize the individual model trajectories and the data to asses the goodness-of-fit.
 
-To repeatedly simulate a model and update a parameter value in each consecutive run, you can use pySODM's *draw function*. These functions always takes `parameters` as its first input argument, representing the dictionary of model parameters. This can then be followed by an arbitrary number of user-defined parameters to aid in the draw function. A draw function must always return the model parameters dictionary `parameters`, without alterations to the dictionaries keys. The draw function defines how values of parameters can change during consecutive simulations of the model, i.e. it updates parameter values in the model parameters dictionary. This feature is useful to perform sensitivty analysis.
+To repeatedly simulate a model and update a parameter value in each consecutive run, you can use pySODM's *draw function*. These functions **always** take the model `parameters` as its first input argument, and the `initial_states` as the second argument, followed by an arbitrary number of user-defined parameters to aid in the draw function. A draw function must always return the dictionary of model `parameters` and the dictionary of `initial_states`, without alterations to the dictionaries keys. The draw function defines how parameters and initial conditions can change during consecutive simulations of the model, i.e. it updates parameter values in the model parameters dictionary. This feature is useful to perform sensitivty analysis.
 
 In this example, we'll use a draw function to replace `beta` with a random value obtained from its posterior probability distribution obtained during calibration. We accomplish this by defining a draw function with one additional argument, `samples`, which is a list containing the samples of the posterior probability distribution of `beta`. `np.random.choice()` is used to sample a random value of `beta` and assign it to the model parameteres dictionary,
 
 ```python
 # Define draw function
-def draw_fcn(parameters, samples):
+def draw_fcn(parameters, initial_states, samples):
     parameters['beta'] = np.random.choice(np.array(samples))
-    return parameters
+    return parameters, initial_states
 ```
 
 To use this draw function, you provide four additional arguments to the `sim()` function,
 1. `N`: the number of repeated simulations,
 2. `draw_function`: the draw function,
-2. `draw_function_kwargs`: a dictionary containing all parameters of the draw function not equal to `parameters` (passed internally by the `sim()` function).
+2. `draw_function_kwargs`: a dictionary containing all parameters of the draw function not equal to `parameters` or `initial_states`.
 4. `processes`: the number of cores to divide the `N` simulations over.
 
 As demonstrated in the quickstart example, the `xarray` containing the model output will now contain an additional dimension to accomodate the repeated simulations: `draws`.
@@ -448,10 +448,10 @@ To simulate ramp-like adoptation of measures, we can add the number of additiona
 
 ```python
 # Define draw function
-def draw_fcn(parameters, samples, ramp_length):
+def draw_fcn(parameters, initial_states, samples, ramp_length):
     parameters['beta'] = np.random.choice(samples)
     parameters['start_measures'] += pd.Timedelta(days=np.random.triangular(left=0,mode=0,right=ramp_length))
-    return parameters
+    return parameters, initial_states
 ```
 
 Don't forget to add the new parameter `ramp_length` to the dictionary of `draw_function_kwargs` when simulating the model! Gradually adopting the preventive measures results in a more realistic simulation,