diff --git a/README.md b/README.md index 5fd0a81e..bc5c0a3e 100644 --- a/README.md +++ b/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) diff --git a/docs/enzyme_kinetics.md b/docs/enzyme_kinetics.md index bf4621cd..d5a4ba1e 100644 --- a/docs/enzyme_kinetics.md +++ b/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. diff --git a/docs/quickstart.md b/docs/quickstart.md index 6040ae69..49010e63 100644 --- a/docs/quickstart.md +++ b/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 diff --git a/docs/workflow.md b/docs/workflow.md index 17034479..73d7e507 100644 --- a/docs/workflow.md +++ b/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, diff --git a/src/pySODM/models/base.py b/src/pySODM/models/base.py index 4b34507e..544f0c15 100644 --- a/src/pySODM/models/base.py +++ b/src/pySODM/models/base.py @@ -367,12 +367,16 @@ def _sim_single(self, time, actual_start_date=None, method='tau_leap', tau=0.5, return _output_to_xarray_dataset(output, self.state_shapes, self.dimensions_per_state, self.state_coordinates, actual_start_date) - def _mp_sim_single(self, drawn_parameters, seed, time, actual_start_date, method, tau, output_timestep): + def _mp_sim_single(self, drawn_parameters, drawn_initial_states, seed, time, actual_start_date, method, tau, output_timestep): """ A Multiprocessing-compatible wrapper for _sim_single, assigns the drawn dictionary and runs _sim_single """ + # set sampled parameters/initial states self.parameters.update(drawn_parameters) + self.initial_states.update(drawn_initial_states) + # set seed np.random.seed(seed) + # simulate model return self._sim_single(time, actual_start_date, method, tau, output_timestep) def sim(self, time, warmup=0, N=1, draw_function=None, draw_function_kwargs={}, processes=None, method='tau_leap', tau=1, output_timestep=1): @@ -430,31 +434,31 @@ def sim(self, time, warmup=0, N=1, draw_function=None, draw_function_kwargs={}, # Input checks related to draw functions if draw_function: # validate function - validate_draw_function(draw_function, draw_function_kwargs, self.parameters) + validate_draw_function(draw_function, draw_function_kwargs, copy.deepcopy(self.parameters), copy.deepcopy(self.initial_states), self.state_shapes) - # Copy parameter dictionary --> dict is global - cp = copy.deepcopy(self.parameters) # Construct list of drawn dictionaries drawn_dictionaries=[] - for n in range(N): - cp_draws=copy.deepcopy(self.parameters) + for _ in range(N): if draw_function: - out={} # Need because of global dictionaries and voodoo magic - out.update(draw_function(self.parameters,**draw_function_kwargs)) - drawn_dictionaries.append(out) + drawn_dictionaries.append(draw_function(copy.deepcopy(self.parameters), copy.deepcopy(self.initial_states), **draw_function_kwargs)) else: - drawn_dictionaries.append({}) - self.parameters=cp_draws + drawn_dictionaries.append([{},{}]) + drawn_parameters = [tpl[0] for tpl in drawn_dictionaries] + drawn_initial_states = [tpl[1] for tpl in drawn_dictionaries] + + # save a copy before altering to reset after simulation + cp_pars = copy.deepcopy(self.parameters) + cp_init_states = copy.deepcopy(self.initial_states) # Run simulations if processes: # Needed with get_context("fork").Pool(processes) as p: # 'fork' instead of 'spawn' to run on Apple Silicon seeds = np.random.randint(0, 2**32, size=N) # requires manual reseeding of the random number generators used in the stochastic algorithms in every child process - output = p.starmap(partial(self._mp_sim_single, time=time, actual_start_date=actual_start_date, method=method, tau=tau, output_timestep=output_timestep), zip(drawn_dictionaries, seeds)) + output = p.starmap(partial(self._mp_sim_single, time=time, actual_start_date=actual_start_date, method=method, tau=tau, output_timestep=output_timestep), zip(drawn_parameters, drawn_initial_states, seeds)) else: output=[] - for dictionary in drawn_dictionaries: - output.append(self._mp_sim_single(dictionary, np.random.randint(0, 2**32, size=1), time, actual_start_date, method=method, tau=tau, output_timestep=output_timestep)) + for pars, init_states in zip(drawn_parameters, drawn_initial_states): + output.append(self._mp_sim_single(pars, init_states, np.random.randint(0, 2**32, size=1), time, actual_start_date, method=method, tau=tau, output_timestep=output_timestep)) # Append results out = output[0] @@ -462,7 +466,8 @@ def sim(self, time, warmup=0, N=1, draw_function=None, draw_function_kwargs={}, out = xarray.concat([out, xarr], "draws") # Reset parameter dictionary - self.parameters = cp + self.parameters = cp_pars + self.initial_states = cp_init_states return out @@ -662,11 +667,14 @@ def _sim_single(self, time, actual_start_date=None, method='RK23', rtol=5e-3, ou # Map to variable names return _output_to_xarray_dataset(output, self.state_shapes, self.dimensions_per_state, self.state_coordinates, actual_start_date) - def _mp_sim_single(self, drawn_parameters, time, actual_start_date, method, rtol, output_timestep, tau): + def _mp_sim_single(self, drawn_parameters, drawn_initial_states, time, actual_start_date, method, rtol, output_timestep, tau): """ - A `multiprocessing`-compatible wrapper for `_sim_single`, assigns the drawn dictionary and runs `_sim_single` + A `multiprocessing`-compatible wrapper for `_sim_single`, assigns the drawn dictionaries and runs `_sim_single` """ + # set sampled parameters/initial states self.parameters.update(drawn_parameters) + self.initial_states.update(drawn_initial_states) + # simulate model out = self._sim_single(time, actual_start_date, method, rtol, output_timestep, tau) return out @@ -720,7 +728,7 @@ def sim(self, time, warmup=0, N=1, draw_function=None, draw_function_kwargs={}, output: xarray.Dataset Simulation output """ - + # Input checks on solution settings validate_solution_methods_ODE(rtol, method, tau) # Input checks on supplied simulation time @@ -728,33 +736,33 @@ def sim(self, time, warmup=0, N=1, draw_function=None, draw_function_kwargs={}, # Input checks related to draw functions if draw_function: # validate function - validate_draw_function(draw_function, draw_function_kwargs, self.parameters) + validate_draw_function(draw_function, draw_function_kwargs, copy.deepcopy(self.parameters), copy.deepcopy(self.initial_states), self.state_shapes) # provinding 'N' but no draw function: wasteful of resources if ((N != 1) & (draw_function==None)): raise ValueError('attempting to perform N={0} repeated simulations without using a draw function'.format(N)) - # Construct list of drawn dictionaries - # Copy parameter dictionary --> dict is global - cp = copy.deepcopy(self.parameters) + # Construct list of drawn parameters and initial states drawn_dictionaries=[] - for n in range(N): - cp_draws=copy.deepcopy(self.parameters) + for _ in range(N): if draw_function: - out={} # Need because of global dictionaries and voodoo magic - out.update(draw_function(self.parameters,**draw_function_kwargs)) - drawn_dictionaries.append(out) + drawn_dictionaries.append(draw_function(copy.deepcopy(self.parameters), copy.deepcopy(self.initial_states), **draw_function_kwargs)) else: - drawn_dictionaries.append({}) - self.parameters=cp_draws + drawn_dictionaries.append([{},{}]) + drawn_parameters = [tpl[0] for tpl in drawn_dictionaries] + drawn_initial_states = [tpl[1] for tpl in drawn_dictionaries] + + # save a copy before altering to reset after simulation + cp_pars = copy.deepcopy(self.parameters) + cp_init_states = copy.deepcopy(self.initial_states) # Run simulations if processes: # Needed with get_context("fork").Pool(processes) as p: - output = p.map(partial(self._mp_sim_single, time=time, actual_start_date=actual_start_date, method=method, rtol=rtol, output_timestep=output_timestep, tau=tau), drawn_dictionaries) + output = p.starmap(partial(self._mp_sim_single, time=time, actual_start_date=actual_start_date, method=method, rtol=rtol, output_timestep=output_timestep, tau=tau), zip(drawn_parameters, drawn_initial_states)) else: output=[] - for dictionary in drawn_dictionaries: - output.append(self._mp_sim_single(dictionary, time, actual_start_date, method=method, rtol=rtol, output_timestep=output_timestep, tau=tau)) + for pars, init_states in zip(drawn_parameters, drawn_initial_states): + output.append(self._mp_sim_single(pars, init_states, time, actual_start_date, method=method, rtol=rtol, output_timestep=output_timestep, tau=tau)) # Append results out = output[0] @@ -762,7 +770,8 @@ def sim(self, time, warmup=0, N=1, draw_function=None, draw_function_kwargs={}, out = xarray.concat([out, xarr], "draws") # Reset parameter dictionary - self.parameters = cp + self.parameters = cp_pars + self.initial_states = cp_init_states return out diff --git a/src/pySODM/models/validation.py b/src/pySODM/models/validation.py index 2bc93108..c1e8e5f3 100644 --- a/src/pySODM/models/validation.py +++ b/src/pySODM/models/validation.py @@ -118,9 +118,12 @@ def validate_solution_methods_JumpProcess(method, tau): "discrete timestep 'tau' must be of type int or float" ) -def validate_draw_function(draw_function, draw_function_kwargs, parameters): - """Validates the draw function's input and output. For use in the sim() functions of the ODE and JumpProcess classes (base.py). +def validate_draw_function(draw_function, draw_function_kwargs, parameters, initial_states, state_shapes): + """ + Validates the draw function's input and output. Used in the sim() functions of the ODE and JumpProcess classes (base.py). + Makes a call to `validate_initial_states` + input ----- @@ -133,12 +136,12 @@ def validate_draw_function(draw_function, draw_function_kwargs, parameters): parameters: dict the dictionary of model parameters - - output - ------ - - parameters: dict - an updated dictionary of model parameters + + initial_states: dict + the dictionary of initial model states + + state_shapes: dict + contains the shape of every model state. """ # check that the draw_function is a function @@ -150,38 +153,62 @@ def validate_draw_function(draw_function, draw_function_kwargs, parameters): args = list(inspect.signature(draw_function).parameters.keys()) if args[0] != "parameters": raise ValueError( - f"your draw function '{draw_function.__name__}' must have 'parameters' as its first input. Its current inputs are: {args}" + f"your draw function '{draw_function.__name__}' must have 'parameters' as its first input. Its current inputs are: '{args}'" + ) + # check that it's second argument is named 'initial_states' + args = list(inspect.signature(draw_function).parameters.keys()) + if args[1] != "initial_states": + raise ValueError( + f"your draw function '{draw_function.__name__}' must have 'initial_states' as its second input. Its current inputs are: '{args}'" ) - # check that draw_function_kwargs is a 'dict' + # check that `draw_function_kwargs` is a 'dict' if not isinstance(draw_function_kwargs, dict): raise TypeError( - f"your draw function kwargs must be of type 'dict' but are of type {type(draw_function_kwargs)}" + f"your `draw_function_kwargs` must be of type 'dict' but are of type '{type(draw_function_kwargs)}'" ) # if draw_functions_kwargs is an empty dict and draw_function has additional kwargs the user has most likely forgotten to pass draw_function_kwargs to the sim() function - if ((len(args[1:]) > 0) & (len(list(draw_function_kwargs.keys())) == 0)): + if ((len(args[2:]) > 0) & (len(list(draw_function_kwargs.keys())) == 0)): raise ValueError( - f"the draw function '{draw_function.__name__}' has {len(args[1:])} arguments in addition to the mandatory 'parameters' argument\n" - f"have you forgotten to pass 'draw_function_kwargs' to the sim() function?" + f"the draw function '{draw_function.__name__}' has {len(args[2:])} arguments in addition to the mandatory 'parameters' and 'initial_states' arguments\n" + f"have you forgotten to pass `draw_function_kwargs` to the sim() function?" ) - # check that it's keys have the same name as the inputs of draw_function that follow 'parameters' - if set(args[1:]) != set(list(draw_function_kwargs.keys())): + # check that it's keys have the same name as the inputs of draw_function that follow `parameters` and `initial_states` + if set(args[2:]) != set(list(draw_function_kwargs.keys())): raise ValueError( f"incorrect arguments passed to draw function '{draw_function.__name__}'\n" - "keys missing in 'draw_function_kwargs': {0}. redundant keys: {1}".format(set(args[1:]).difference(list(draw_function_kwargs.keys())), set(list(draw_function_kwargs.keys())).difference(set(args[1:]))) + "keys missing in 'draw_function_kwargs': {0}. redundant keys: {1}".format(set(args[2:]).difference(list(draw_function_kwargs.keys())), set(list(draw_function_kwargs.keys())).difference(set(args[2:]))) ) # call draw function and check its outputs - cp_draws=copy.deepcopy(parameters) - d = draw_function(parameters, **draw_function_kwargs) - parameters = cp_draws - if not isinstance(d, dict): + output = draw_function(copy.deepcopy(parameters), copy.deepcopy(initial_states), **draw_function_kwargs) + # check if it returns two outputs + if not ((isinstance(output, tuple)) & (len(output) == 2)): + raise TypeError(f"a draw function must return two dictionaries: 1) parameters, 2) initial_states") + # check they're both dicts + if not ((isinstance(output[0], dict)) & (isinstance(output[1], dict))): raise TypeError( - f"a draw function must return a dictionary. found type {type(d)}" + "a draw function must return two dictionaries: 1) parameters, 2) initial_states\n" + f"found the following types: ('{type(output[0])}', '{type(output[1])}')." + ) + # verify keys are the same on input/output 'parameters' + if set(output[0].keys()) != set(parameters.keys()): + raise ValueError( + f"a draw function must return two dictionaries: 1) parameters, 2) initial_states" + f"keys in first output dictionary of draw function '{draw_function.__name__}' must match the keys in input dictionary 'parameters'.\n" + "keys missing in draw function output: {0}. redundant keys: {1}".format(set(parameters.keys()).difference(set(output[0].keys())), set(output[0].keys()).difference(set(parameters.keys()))) ) - if set(d.keys()) != set(parameters.keys()): + # verify keys are the same on input/output 'initial_states' + if set(output[1].keys()) != set(initial_states.keys()): raise ValueError( - f"keys in output dictionary of draw function '{draw_function.__name__}' must match the keys in input dictionary 'parameters'.\n" - "keys missing in draw function output: {0}. redundant keys: {1}".format(set(parameters.keys()).difference(set(d.keys())), set(d.keys()).difference(set(parameters.keys()))) + f"a draw function must return two dictionaries." + f"keys in second output dictionary of draw function '{draw_function.__name__}' must match the keys in input dictionary 'initial_states'.\n" + "keys missing in draw function output: {0}. redundant keys: {1}".format(set(initial_states.keys()).difference(set(output[1].keys())), set(output[1].keys()).difference(set(initial_states.keys()))) ) + # verify the initial states sizes + try: + _ = validate_initial_states(state_shapes, output[1]) + except Exception as e: + error_message = f"your draw function did not return a valid initial state.\nfound error: {str(e)}" + raise RuntimeError(error_message) from e def fill_initial_state_with_zero(state_names, initial_states): for state in state_names: diff --git a/src/tests/test_JumpProcess.py b/src/tests/test_JumpProcess.py index 5a52695e..648d51b3 100644 --- a/src/tests/test_JumpProcess.py +++ b/src/tests/test_JumpProcess.py @@ -526,8 +526,8 @@ def test_draw_function(): coordinates = {'age_groups': ['0-20','20-120']} # correct draw function without additional arguments - def draw_function(parameters): - return parameters + def draw_function(parameters, initial_states): + return parameters, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) @@ -536,8 +536,8 @@ def draw_function(parameters): assert 'draws' in list(output.sizes.keys()) # correct draw function with additional arguments - def draw_function(parameters, par_1, par_2): - return parameters + def draw_function(parameters, initial_states, par_1, par_2): + return parameters, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) @@ -545,15 +545,9 @@ def draw_function(parameters, par_1, par_2): # assert dimension 'draws' is present in output assert 'draws' in list(output.sizes.keys()) - # wrong draw function: not a function - time = [0, 10] - model = SIRstratified(initial_states, parameters, coordinates=coordinates) - with pytest.raises(TypeError, match="a 'draw function' must be callable"): - model.sim(time, draw_function='bliblablu', N=5) - # correct draw function but too few arguments provided in draw_function_kwargs - def draw_function(parameters, par_1, par_2): - return parameters + def draw_function(parameters, initial_states, par_1, par_2): + return parameters, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) @@ -561,54 +555,107 @@ def draw_function(parameters, par_1, par_2): model.sim(time, draw_function=draw_function, draw_function_kwargs={'par_1': 0}, N=5) # correct draw function but too much arguments in draw_function_kwargs - def draw_function(parameters): - return parameters + def draw_function(parameters, initial_states): + return parameters, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) with pytest.raises(ValueError, match="incorrect arguments passed to draw function"): model.sim(time, draw_function=draw_function, draw_function_kwargs={'par_1': 0}, N=5) + # correct draw function but draw_function_kwargs is not a dict + def draw_function(parameters, initial_states): + return parameters, initial_states + # simulate model + time = [0, 10] + model = SIRstratified(initial_states, parameters, coordinates=coordinates) + with pytest.raises(TypeError, match="your `draw_function_kwargs` must be of type 'dict' but are of type"): + model.sim(time, draw_function=draw_function, draw_function_kwargs='bliblablu', N=5) + # correct draw function with extra args but user forgets to provide draw_function_kwargs to sim() - def draw_function(parameters, par_1, par_2): + def draw_function(parameters, initial_states, par_1, par_2): + return parameters, initial_states + # simulate model + time = [0, 10] + model = SIRstratified(initial_states, parameters, coordinates=coordinates) + with pytest.raises(ValueError, match="the draw function 'draw_function' has 2 arguments in addition to the mandatory 'parameters' and 'initial_states' arguments"): + model.sim(time, draw_function=draw_function, N=5) + + # wrong draw function: not a function + time = [0, 10] + model = SIRstratified(initial_states, parameters, coordinates=coordinates) + with pytest.raises(TypeError, match="a 'draw function' must be callable"): + model.sim(time, draw_function='bliblablu', N=5) + + # wrong draw function: only one output + def draw_function(parameters, initial_states): return parameters # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) - with pytest.raises(ValueError, match="the draw function 'draw_function' has 2 arguments in addition to the mandatory 'parameters' argument"): + with pytest.raises(TypeError, match="a draw function must return two dictionaries"): model.sim(time, draw_function=draw_function, N=5) # wrong draw function: first input argument is not 'parameters' - def draw_function(par_1, parameters): - return parameters + def draw_function(par_1, parameters, initial_states): + return parameters, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) with pytest.raises(ValueError, match="must have 'parameters' as its first input. Its current inputs are"): model.sim(time, draw_function=draw_function, draw_function_kwargs={'par_1': 0}, N=5) - - # wrong draw function: return a scalar - def draw_function(parameters): - return 5 + + # wrong draw function: second input argument is not 'initial_states' + def draw_function(parameters, par_1, initial_states): + return parameters, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) - with pytest.raises(TypeError, match="a draw function must return a dictionary. found type"): + with pytest.raises(ValueError, match="must have 'initial_states' as its second input."): + model.sim(time, draw_function=draw_function, draw_function_kwargs={'par_1': 0}, N=5) + + # wrong draw function: 2 outputs but wrong type on one of them + def draw_function(parameters, initial_states): + return 5, initial_states + # simulate model + time = [0, 10] + model = SIRstratified(initial_states, parameters, coordinates=coordinates) + with pytest.raises(TypeError, match="a draw function must return two dictionaries"): model.sim(time, draw_function=draw_function, N=5) - # wrong draw function: put a new keys in model parameters dictionary that doesn't represent a model parameter - def draw_function(parameters): + # wrong draw function: put a new key in model parameters dictionary that doesn't represent a model parameter + def draw_function(parameters, initial_states): parameters['bliblublo'] = 5 - return parameters + return parameters, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) - with pytest.raises(ValueError, match="keys in output dictionary of draw function 'draw_function' must match the keys in input dictionary 'parameters'."): + with pytest.raises(ValueError, match="must match the keys in input dictionary 'parameters'"): model.sim(time, draw_function=draw_function, N=5) - # user provides N but no draw function (differs from ODE) - def draw_function(parameters): - return parameters + # wrong draw function: put a new key in initial states dictionary that doesn't represent a model state + def draw_function(parameters, initial_states): + initial_states['bliblublo'] = 5 + return parameters, initial_states + # simulate model + time = [0, 10] + model = SIRstratified(initial_states, parameters, coordinates=coordinates) + with pytest.raises(ValueError, match="keys in second output dictionary of draw function 'draw_function' must match the keys in input dictionary 'initial_states'"): + model.sim(time, draw_function=draw_function, N=5) + + # wrong draw function: wrong initial state size + def draw_function(parameters, initial_states): + initial_states['S'] = np.array([100, 100, 100]) + return parameters, initial_states + # simulate model + time = [0, 10] + model = SIRstratified(initial_states, parameters, coordinates=coordinates) + with pytest.raises(RuntimeError, match="your draw function did not return a valid initial state."): + model.sim(time, draw_function=draw_function, N=5) + + # no error: user provides N but no draw function (differs from ODE) + def draw_function(parameters, initial_states): + return parameters, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) @@ -616,7 +663,7 @@ def draw_function(parameters): # assert dimension 'draws' is present in output assert 'draws' in list(output.sizes.keys()) - # or + # no error: user provides N, no draw function but draw_function_kwargs output = model.sim(time, draw_function_kwargs={'arg_1': 0}, N=100) # assert dimension 'draws' is present in output assert 'draws' in list(output.sizes.keys()) \ No newline at end of file diff --git a/src/tests/test_ODE.py b/src/tests/test_ODE.py index a6c55942..9492bf7e 100644 --- a/src/tests/test_ODE.py +++ b/src/tests/test_ODE.py @@ -518,8 +518,8 @@ def test_draw_function(): coordinates = {'age_groups': ['0-20','20-120']} # correct draw function without additional arguments - def draw_function(parameters): - return parameters + def draw_function(parameters, initial_states): + return parameters, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) @@ -528,8 +528,8 @@ def draw_function(parameters): assert 'draws' in list(output.sizes.keys()) # correct draw function with additional arguments - def draw_function(parameters, par_1, par_2): - return parameters + def draw_function(parameters, initial_states, par_1, par_2): + return parameters, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) @@ -537,15 +537,9 @@ def draw_function(parameters, par_1, par_2): # assert dimension 'draws' is present in output assert 'draws' in list(output.sizes.keys()) - # wrong draw function: not a function - time = [0, 10] - model = SIRstratified(initial_states, parameters, coordinates=coordinates) - with pytest.raises(TypeError, match="a 'draw function' must be callable"): - model.sim(time, draw_function='bliblablu', N=5) - # correct draw function but too few arguments provided in draw_function_kwargs - def draw_function(parameters, par_1, par_2): - return parameters + def draw_function(parameters, initial_states, par_1, par_2): + return parameters, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) @@ -553,49 +547,102 @@ def draw_function(parameters, par_1, par_2): model.sim(time, draw_function=draw_function, draw_function_kwargs={'par_1': 0}, N=5) # correct draw function but too much arguments in draw_function_kwargs - def draw_function(parameters): - return parameters + def draw_function(parameters, initial_states): + return parameters, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) with pytest.raises(ValueError, match="incorrect arguments passed to draw function"): model.sim(time, draw_function=draw_function, draw_function_kwargs={'par_1': 0}, N=5) + # correct draw function but draw_function_kwargs is not a dict + def draw_function(parameters, initial_states): + return parameters, initial_states + # simulate model + time = [0, 10] + model = SIRstratified(initial_states, parameters, coordinates=coordinates) + with pytest.raises(TypeError, match="your `draw_function_kwargs` must be of type 'dict' but are of type"): + model.sim(time, draw_function=draw_function, draw_function_kwargs='bliblablu', N=5) + # correct draw function with extra args but user forgets to provide draw_function_kwargs to sim() - def draw_function(parameters, par_1, par_2): + def draw_function(parameters, initial_states, par_1, par_2): + return parameters, initial_states + # simulate model + time = [0, 10] + model = SIRstratified(initial_states, parameters, coordinates=coordinates) + with pytest.raises(ValueError, match="the draw function 'draw_function' has 2 arguments in addition to the mandatory 'parameters' and 'initial_states' arguments"): + model.sim(time, draw_function=draw_function, N=5) + + # wrong draw function: not a function + time = [0, 10] + model = SIRstratified(initial_states, parameters, coordinates=coordinates) + with pytest.raises(TypeError, match="a 'draw function' must be callable"): + model.sim(time, draw_function='bliblablu', N=5) + + # wrong draw function: only one output + def draw_function(parameters, initial_states): return parameters # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) - with pytest.raises(ValueError, match="the draw function 'draw_function' has 2 arguments in addition to the mandatory 'parameters' argument"): + with pytest.raises(TypeError, match="a draw function must return two dictionaries"): model.sim(time, draw_function=draw_function, N=5) # wrong draw function: first input argument is not 'parameters' - def draw_function(par_1, parameters): - return parameters + def draw_function(par_1, parameters, initial_states): + return parameters, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) with pytest.raises(ValueError, match="must have 'parameters' as its first input. Its current inputs are"): model.sim(time, draw_function=draw_function, draw_function_kwargs={'par_1': 0}, N=5) - - # wrong draw function: return a scalar - def draw_function(parameters): - return 5 + + # wrong draw function: second input argument is not 'initial_states' + def draw_function(parameters, par_1, initial_states): + return parameters, initial_states + # simulate model + time = [0, 10] + model = SIRstratified(initial_states, parameters, coordinates=coordinates) + with pytest.raises(ValueError, match="must have 'initial_states' as its second input."): + model.sim(time, draw_function=draw_function, draw_function_kwargs={'par_1': 0}, N=5) + + # wrong draw function: 2 outputs but wrong type on one of them + def draw_function(parameters, initial_states): + return 5, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) - with pytest.raises(TypeError, match="a draw function must return a dictionary. found type"): + with pytest.raises(TypeError, match="a draw function must return two dictionaries"): model.sim(time, draw_function=draw_function, N=5) - # wrong draw function: put a new keys in model parameters dictionary that doesn't represent a model parameter - def draw_function(parameters): + # wrong draw function: put a new key in model parameters dictionary that doesn't represent a model parameter + def draw_function(parameters, initial_states): parameters['bliblublo'] = 5 - return parameters + return parameters, initial_states + # simulate model + time = [0, 10] + model = SIRstratified(initial_states, parameters, coordinates=coordinates) + with pytest.raises(ValueError, match="must match the keys in input dictionary 'parameters'"): + model.sim(time, draw_function=draw_function, N=5) + + # wrong draw function: put a new key in initial states dictionary that doesn't represent a model state + def draw_function(parameters, initial_states): + initial_states['bliblublo'] = 5 + return parameters, initial_states + # simulate model + time = [0, 10] + model = SIRstratified(initial_states, parameters, coordinates=coordinates) + with pytest.raises(ValueError, match="keys in second output dictionary of draw function 'draw_function' must match the keys in input dictionary 'initial_states'"): + model.sim(time, draw_function=draw_function, N=5) + + # wrong draw function: wrong initial state size + def draw_function(parameters, initial_states): + initial_states['S'] = np.array([100, 100, 100]) + return parameters, initial_states # simulate model time = [0, 10] model = SIRstratified(initial_states, parameters, coordinates=coordinates) - with pytest.raises(ValueError, match="keys in output dictionary of draw function 'draw_function' must match the keys in input dictionary 'parameters'."): + with pytest.raises(RuntimeError, match="your draw function did not return a valid initial state."): model.sim(time, draw_function=draw_function, N=5) # user provides N but no draw function diff --git a/tutorials/SIR/workflow_tutorial.py b/tutorials/SIR/workflow_tutorial.py index 047f0f4a..5ee7d384 100644 --- a/tutorials/SIR/workflow_tutorial.py +++ b/tutorials/SIR/workflow_tutorial.py @@ -171,9 +171,9 @@ def integrate(t, S, I, R, beta, gamma): ############################## # Define draw function - def draw_fcn(parameters, samples): + def draw_fcn(parameters, initial_states, samples): parameters['beta'] = np.random.choice(samples) - return parameters + return parameters, initial_states # Simulate model out = model.sim([start_date, end_date+pd.Timedelta(days=2*28)], N=100, draw_function=draw_fcn, @@ -207,10 +207,10 @@ def lower_infectivity(t, states, param, start_measures): return param/2 # 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 # Attach its arguments to the parameter dictionary model.parameters.update({'start_measures': end_date}) diff --git a/tutorials/SIR_SI/calibration.py b/tutorials/SIR_SI/calibration.py index 8df63b5d..5eab0b94 100644 --- a/tutorials/SIR_SI/calibration.py +++ b/tutorials/SIR_SI/calibration.py @@ -142,10 +142,10 @@ # Define draw function import random - def draw_fcn(parameters, samples): + def draw_fcn(parameters, initial_states, samples): idx, parameters['beta'] = random.choice(list(enumerate(samples['beta']))) parameters['alpha'] = np.array([slice[idx] for slice in samples['alpha']]) - return parameters + return parameters, initial_states # Simulate model N = 100 out = model.sim([start_date, end_date+60], N=N, draw_function=draw_fcn, draw_function_kwargs={'samples': samples_dict}, processes=processes) diff --git a/tutorials/enzyme_kinetics/calibrate_intrinsic_kinetics.py b/tutorials/enzyme_kinetics/calibrate_intrinsic_kinetics.py index f072c73a..aaf806eb 100644 --- a/tutorials/enzyme_kinetics/calibrate_intrinsic_kinetics.py +++ b/tutorials/enzyme_kinetics/calibrate_intrinsic_kinetics.py @@ -10,10 +10,9 @@ ## Load required packages ## ############################ -import sys,os +import os import random import datetime -import emcee import pandas as pd import numpy as np import multiprocessing as mp @@ -140,13 +139,13 @@ ## Visualize result ## ###################### - 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 # Loop over datasets for i,df in enumerate(data): diff --git a/tutorials/enzyme_kinetics/simulate_1D_PFR.py b/tutorials/enzyme_kinetics/simulate_1D_PFR.py index c53ddb03..5d269099 100644 --- a/tutorials/enzyme_kinetics/simulate_1D_PFR.py +++ b/tutorials/enzyme_kinetics/simulate_1D_PFR.py @@ -38,13 +38,13 @@ samples_dict = 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 ################### ## Load the data ## diff --git a/tutorials/influenza_1718/calibration.py b/tutorials/influenza_1718/calibration.py index dff40346..587c6561 100644 --- a/tutorials/influenza_1718/calibration.py +++ b/tutorials/influenza_1718/calibration.py @@ -216,11 +216,11 @@ ###################### # Define draw function - def draw_fcn(parameters, samples): + def draw_fcn(parameters, initial_states, samples): # Sample model parameters idx, parameters['beta'] = random.choice(list(enumerate(samples['beta']))) parameters['f_ud'] = np.array([slice[idx] for slice in samples['f_ud']]) - return parameters + return parameters, initial_states # Simulate model out = model.sim([start_visualisation, end_visualisation], N=n, tau=tau, output_timestep=1, draw_function=draw_fcn, draw_function_kwargs={'samples': samples_dict}, processes=processes)