API reference

Inference of Compartmental Models Toolbox.

class icomo.CompModel(Comp_dict)

Bases: object

Class to help building a compartmental model.

The model is built by adding flows between compartments. The model is then compiled into a function that can be used in an ODE.

__init__(Comp_dict)

Initialize the CompModel class.

Parameters:

Comp_dict (dict) – Dictionary of compartments. Keys are the names of the compartments and values are floats or ndarrays that represent their value.

flow(start_comp, end_comp, rate, label=None, end_comp_is_list=False)

Add a flow from start_comp to end_comp with rate flow.

Parameters:

• start_comp (str or list) – Key of the start compartment. Can also be a list of keys in which case an identical flow is added from each compartment of the list to end_com

• end_comp (str) – Key of the end compartment

• rate (float or ndarray) – rate of the flow to add between compartments, is multiplied by start_comp, so it should be broadcastable whith it.

• label (str, optional) – label of the edge between the compartments that will be used when displaying a graph of the compartmental model.

• end_comp_is_list (bool, default: False) – If True, end_comp points to a list of compartments, and the outflow is added to the first of them. This is typically the case when the end compartment is used with a Erlang kernel.

Returns:

None

erlang_flow(start_comp, end_comp, rate, label=None, end_comp_is_list=False)

Add a flow with erlang kernel from start_comp to end_comp with rate flow.

Parameters:

• start_comp (str) – start compartments of the flow. The length the list is the shape of the Erlang kernel.

• end_comp (str) – end compartment of the flow

• rate (float or ndarray) – rate of the flow, equal to the inverse of the mean time spent in the Erlang kernel.

• label (str, optional) – label of the edge between the compartments that will be used when displaying a graph of the compartmental model.

• end_comp_is_list (bool, default: False) – If True, end_comp points to a list of compartments, and the outflow is added to the first of them. This is typically the case when the end compartment is used with a Erlang kernel.

property dy

Returns the derivative of the compartments.

This should be returned by the function that defines the system of ODEs.

Returns:

dComp (dict)

view_graph(on_display=True)

Display a graph of the compartmental model.

Parameters:

on_display (bool) – If True, the graph is displayed in the notebook, otherwise it is saved as a pdf in the current folder and opened with the default pdf viewer.

Returns:

None

class icomo.ODEIntegrator(ts_out, ts_solver=None, ts_arg=None, interp='cubic', solver=None, t_0=None, t_1=None, **kwargs)

Bases: object

Creates an integrator for compartmental models.

The integrator is a function that takes as input a function that returns the derivatives of the variables of the system of differential equations, and returns a function that solves the system of differential equations. For the initilization of the integrator object, the timesteps of the solver and the output have to be specified.

__init__(ts_out, ts_solver=None, ts_arg=None, interp='cubic', solver=None, t_0=None, t_1=None, **kwargs)

Initialize the ODEIntegrator class.

Parameters:

• ts_out (array-like) – The timesteps at which the output is returned.

• ts_solver (array-like or None) – The timesteps at which the solver will be called. If None, it is set to ts_out.

• ts_arg (array-like or None) – The timesteps at which the time-dependent argument of the system of differential equations are given. If None, it is set to ts_solver.

• interp (str) – The interpolation method used to interpolate_pytensor the time-dependent argument of the system of differential equations. Can be “cubic” or “linear”.

• solver (diffrax.AbstractStepSizeController) – The solver used to integrate the system of differential equations. Default is diffrax.Tsit5(), a 5th order Runge-Kutta method.

• t_0 (float or None) – The initial time of the integration. If None, it is set to ts_solve[0].

• t_1 (float or None) – The final time of the integration. If None, it is set to ts_solve[-1].

• **kwargs – Arguments passed to the solver, see diffrax.diffeqsolve() for more details.

get_func(ODE, list_keys_to_return=None)

Return a function that solves the system of differential equations.

Parameters:

• ODE (function(t, y, args)) – A function that returns the derivatives of the variables of the system of differential equations. The function has to take as input the time t, the variables y and the arguments args=(arg_t, constant_args) of the system of differential equations. t is a float, y is a list or dict of floats or ndarrays, or in general, a pytree, see jax.tree_util for more details. The return value of the function has to be a pytree/list/dict with the same structure as y.

• list_keys_to_return (list of str or None, default is None) – The keys of the variables of the system of differential equations that are returned by the integrator. If set, the integrator returns a list of the variables of the system of differential equations in the order of the keys. If None, the output is returned as is.

Returns:

integrator (function(y0, arg_t=None, constant_args=None)) – A function that solves the system of differential equations and returns the output at the specified timesteps. The function takes as input y0 the initial values of the variables of the system of differential equations, the time-dependent argument of the system of differential equations arg_t, and the constant arguments constant_args of the system of differential equations. t, y0 and (arg_t, constant_args) are passed to the ODE function as its three arguments. If arg_t is None, only constant_args are passed to the ODE function and vice versa, without being in a tuple.

get_op(ODE, return_shapes=((),), list_keys_to_return=None, name=None)

Return a pytensor operator that solves the system of differential equations.

Same as get_func, but returns a pytensor operator that can be used in a pymc model. Beware that for this operator the output of the integration of the ODE can only be a single or a list variables. If the output is a dict, set list_keys_to_return to specify the keys of the variables that are returned by the integrator. These return values aren’t allowed to be further nested.

Parameters:

• ODE (function(t, y, args)) – A function that returns the derivatives of the variables of the system of differential equations. The function has to take as input the time t, the variables y and the arguments args=(arg_t, constant_args) of the system of differential equations. t is a float, y is a list or dict of floats or ndarrays, or in general, a pytree, see jax.tree_util for more details. The return value of the function has to be a pytree/list/dict with the same structure as y.

• return_shapes (tuple of tuples, default is ((),)) – Depreceated, the return shape had to be specified before, now it is inferred automatically. This argument isn’t used anymore.

• list_keys_to_return (list of str or None, default is None) – The keys of the variables of the system of differential equations that will be chosen to be returned by the integrator. If None, the output is returned as is.

• name – The name under which the operator is registered in pymc.

Returns:

pytensor_op (pytensor.graph.op.Op) – A pytensor operator that can be used in a pymc.Model.

icomo.interpolate_pytensor(ts_in, ts_out, y, method='cubic', ret_gradients=False, name=None)

Interpolate the time-dependent variable y at the timesteps ts_out.

Parameters:

• ts_in (array-like) – The timesteps at which the time-dependent variable is given.

• ts_out (array-like) – The timesteps at which the time-dependent variable should be interpolated.

• y (array-like) – The time-dependent variable.

• method (str) – The interpolation method used. Can be “cubic” or “linear”.

• ret_gradients (bool) – If True, the gradients of the interpolated variable are returned. Default is False.

Returns:

y_interp (array-like) – The interpolated variable at the timesteps ts_out.

icomo.interpolation_func(ts, x, method='cubic')

Return a diffrax-interpolation function that can be used to interpolate pytensors.

Parameters:

• ts (array-like) – The timesteps at which the time-dependent variable is given.

• x (array-like) – The time-dependent variable.

• method – The interpolation method used. Can be “cubic” or “linear”.

Returns:

interp (diffrax.CubicInterpolation or diffrax.LinearInterpolation) – The interpolation function. Call interp.evaluate(t) to evaluate the interpolated variable at time t. t can be a float or an array-like.

icomo.erlang_kernel(inflow, Vars, rate)

Model the compartments delayed by an Erlang kernel.

Utility function to model an Erlang kernel for a compartmental model. The shape is determined by the length of Vars.

Parameters:

• inflow (float or ndarray) – The inflow into the first compartment of the kernel

• Vars (list of floats or list of ndarrays) – The compartments of the kernel

• rate (float or ndarray) – The rate of the kernel, 1/rate is the mean time spent in total in all compartments

Returns:

• dVars (list of floats or list of ndarrays) – The derivatives of the compartments

• out_flow (float or ndarray) – The outflow from the last compartment

icomo.delayed_copy(initial_var, delayed_vars, tau_delay)

Return the derivative to model the delayed copy of a compartment.

The delay has the form of an Erlang kernel with shape parameter len(delayed_vars).

Parameters:

• initial_var (float or ndarray) – The compartment that is copied

• delayed_vars (list of floats or list of ndarrays) – List of compartments that are delayed copies of initial_var, the last element is the compartment which has the same total content as initial_var over time, but is delayed by tau_delay.

• tau_delay (float or ndarray) – The mean delay of the copy.

Returns:

d_delayed_vars (list of floats or list of ndarrays) – The derivatives of the delayed compartments

icomo.SIR(t, y, args)

Differential equations of an SIR model.

Parameters:

• t (float) – time variable

• y (dict) – dictionary of compartments. Has to include keys “S”, “I”, “R”. The value of “S” is the number of susceptible individuals, “I” is the number of infected individuals and “R” is the number of recovered individuals.

• args (tuple) – tuple of arguments. The first argument is the function β(t) that describes the infection rate. The second argument is a dictionary of constant arguments. It has to include the key “gamma” which is the recovery rate and the key “N” which is the total population size.

Returns:

dy (dict) – dictionary of derivatives with keys “S”, “I”, “R”.

icomo.Erlang_SEIRS(t, y, args)

Return the derivatives of the Erlang SEIRS model.

This function defines the differential equations of an SEIRS model with Erlang distributed latent, infectious and recovering (recovered-to-susceptible) periods.

Parameters:

• t (float) – time variable

• y (dict) – dictionary of compartments. Has to include keys “S”, “Es”, “Is”, “Rs”. The value of “Es”, “Is”, “Rs” have to be a list of length n where n is the number of compartments/shape of the Erlang distribution/kernel.

• args (tuple: (callable, dict)) – Callable is a function with argument t that returns the value of the transmission rate beta(t) dict is the constant arguments of the model. Has to include keys “N”, “rate_latent”, “rate_infectious” and “rate_recovery”.

Returns:

dComp (dict) – Same as y, but with the derivatives of the compartments.

icomo.priors_for_cps(cp_dim, time_dim, name_positions, name_magnitudes, name_durations, beta_magnitude=1, sigma_magnitude_fix=None, dist_magnitudes=pymc.Normal, absolute_magnitude_parametrization=False, empirical_bayes_hyper_sigma=False, centered_parametrization=False, model=None)

Create priors for changepoints.

Their positions are uniformly distributed between the first and last timepoint. The magnitudes are sampled from a hierarchical prior with a beta distribution. The durations are sampled from a normal distribution with mean equal to the mean distance between changepoints and standard deviation equal to the standard deviation of the distance between changepoints.

Parameters:

• cp_dim (str) – Dimension of the pymc.Model for the changepoints. Define it by passing coords={cp_dim: np.arange(num_cps)} to pymc.Model at creation. The length of this dimension determines the number of changepoints.

• time_dim (str) – Dimension of the pymc.Model for the time.

• name_positions (str) – Name under which the positions of the changepoints are stored in pymc.Model

• name_magnitudes (str) – Name under which the magnitudes of the changepoints are stored in pymc.Model

• name_durations (str) – Name under which the durations of the changepoints are stored in pymc.Model

• beta_magnitude (float, default=1) – Beta parameter of the hierarchical prior for the magnitudes

• sigma_magnitude_fix (float, default=None) – If not None, the standard deviation from which the magnitudes are sampled is fixed

• dist_magnitudes (pymc.Distribution, default=pm.Normal) – Distribution from which the magnitudes are sampled. Can for example be functools.partial(pm.StudentT, nu=4) to sample from a StudentT distribution for a more robust model.

• absolute_magnitude_parametrization (bool, default=False) – Whether to use an parametrization that is absolute or relative to previous values for the magnitudes.

• empirical_bayes_hyper_sigma (bool, default=False) – Whether to set the standard deviation of the hierarchical magnitudes to the maximum likelihood estimate instead of sampling. Corresponds to an empirical Bayes approach.

• centered_parametrization (bool, default=False) – Whether to use a centered parametrization for the hierarchical priors of the magnitudes.

• model (pymc.Model, default=None) – pm.Model in which the priors are created. If None, the pm.Model is taken from the context.

Returns:

positions, magnitudes, durations (pytensor.Variable) – Variables of dim cp_dim that define the positions, magnitudes and durations of the changepoints

icomo.sigmoidal_changepoints(ts_out, positions_cp, magnitudes_cp, durations_cp, reorder_cps=False)

Modulation of a time series by sigmoidal changepoints.

The changepoints are defined by their position, magnitude and duration. The resulting equation is:

\[f(t) = \sum_{i=1}^{\mathrm{num_cps}} \frac{\mathrm{magnitudes\_cp}[i]}{1 + exp(-4 \cdot \mathrm{slope}[i] \cdot (t - \mathrm{positions_cp}[i]))}\]

where slope[i] = magnitudes_cp[i] / durations_cp[i].

Parameters:

• t_out (1d-array) – timepoints where modulation is evaluated, shape: (time, )

• t_cp (nd-array) – timepoints of the changepoints, shape: (num_cps, further dims…)

• magnitudes (nd-array) – magnitude of the changepoints, shape: (num_cps, further dims…)

• durations (nd-array) – magnitude of the changepoints, shape: (num_cps, further dims…)

• reorder_cps (bool, default=False) – reorder changepoints such that their timepoints are linearly increasing

Returns:

modulation_t ((n+1)d-array) – shape: (time, further dims…)

icomo.hierarchical_priors(name, dims, beta=1, fix_hyper_sigma=None, dist_values=pymc.Normal, centered_parametrization=False, empirical_sigma=False)

Create hierarchical priors for a variable.

Create an n-dimensional variable with hierarchical prior with name name and dimensions dims. The prior is a normal distribution with standard deviation sigma which is sample from a half-cauchy distribution with beta parameter beta.

Parameters:

• name (str) – Name under which the variable is stored in pm.Model

• dims (tuple of str) – Dimensions over which the variable is defined. Define a dimension by passing coords={dim_name: np.arange(size)} to pm.Model at creation.

• beta (float, default=1) – Beta parameter of the half-cauchy distribution

• fix_hyper_sigma (float, default=None) – If not None, the standard deviation from which the variable is sampled is fixed

• dist_values (pymc.Distribution, default=pm.Normal) – Distribution from which the values are sampled. Can for example be functools.partial(pm.StudentT, nu=4) to sample from a StudentT distribution for a more robust model.

• centered_parametrization (bool, default=False) – Whether to use a centered or non-centered parametrization for the hierarchical prior

• empirical_sigma (bool, default=False) – Whether to set the standard deviation of the normal distribution to the standard deviation of the sampled value. Corresponds to an empirical Bayes approach.

Returns:

values (pm.Variable) – pm.Variable for the variable with name name

icomo.jax2pytensor(jaxfunc, output_shape_def=None, args_for_graph='all', name=None, static_argnames=(), dtype=None)

Return a pytensor from a jax jittable function.

It requires to define the output types of the returned values as pytensor types. A unique name should also be passed in case the name of the jaxfunc is identical to some other node. The design of this function is based on https://www.pymc-labs.io/blog-posts/jax-functions-in-pymc-3-quick-examples/

Parameters:

• jaxfunc (jax jittable function) – function for which the node is created, can return multiple tensors as a tuple. It is required that all return values are able to transformed to pytensor.TensorVariable.

• output_shape_def (function) – Function that returns the shape of the output. If None, the shape is expected to be the same as the shape of the args_for_graph arguments. If not None, the function should return a tuple of shapes, it will receive as input the shapes of the args_for_graph as tuples. Shapes are defined as tuples of integers or None.

• args_for_graph (list of str or "all") – If “all”, all arguments except arguments passed via **kwargs are used for the graph. Otherwise specify a list of argument names to use for the graph.

• name (str) – Name of the created pytensor Op, defaults to the name of the passed function. Only used internally in the pytensor graph.

• dtype (pytensor type) – dtype of the in- and outputs, if None, it is inferred from the input types, by upcasting all inputs.

Returns:

• A Pytensor Op which can be used in a pm.Model as function, is differentiable

• and compilable with both JAX and C backend.