dynode.simulation.odes.simulate

dynode.simulation.odes.simulate(ode: Callable[[Array | ndarray | bool | number | bool | int | float | complex, Tuple[Array, ...], PyTree], Tuple[Array, ...]], duration_days: int, initial_state: Tuple[Array, ...], ode_parameters: AbstractODEParams, solver_parameters: SolverParams, sub_save_indices: Tuple[int, ...] | None = None, save_step: int = 1) → Solution

Solve model ODEs for tf days using initial_state and args parameters.

Parameters

ode: ODEs

a callable that takes in a numeric time, a compartment state, and the passed ode_parameters in that order and returns the gradients of the compartment state at that numeric time.

initial_stateCompartmentState

tuple of jax arrays representing the compartments modeled by ODEs in their initial states at t=0.

solver_parametersSolverParams

solver specific parameters that dictate how the ODE solver works.

ode_parametersAbstractODEParams

ode specific parameters that dictate transmission, protection, strain introduction etc. Specific ODE classes will likely require subclasses of AbstractODEParams for their usecase.

duration_daysint, Optional

number of days to solve ODEs for, by default 100 days

sub_save_indicesTuple[int, …]

tuple of initial_state indices specifying which compartments to save states for in sol.ys. sub_save_indices is optional and by default set to None.

save_step: int

value that lets you increment your time step at which a state is saved. If for example you would like to run your solution for duration_days = 100 but only save a state weekly you would pass save_step = 7. save_step is optional by default it is set to 1 which will have no effect.

Returns

diffrax.Solution

Solution object, sol.ys containing compartment states for each day including t=0 and t=duration_days. For more information on whats included within diffrax.Solution see: https://docs.kidger.site/diffrax/api/solution/

Raises

TypeError

initial_state must only contain jax.Array types.

Examples

sub_save_indices and save_step example:

solution: Solution = simulate(

…, sub_save_indices=(config.idx.s, config.idx.r), save_step=7,

)