-
Notifications
You must be signed in to change notification settings - Fork 0
Software Design
Behind the hood, the mechanistic model is actually a collection of four software classes, the initializer, Parameters, Runner, and Interpreter. Each interacting with one another in structured ways to avoid interdependence as much as possible.
The role of an AbstractInitializer and its child classes is to set up the initial state of the model. Each compartment, in the order specified by the COMPARTMENT_IDX enum, must appear in a tuple when the get_initial_state() method is called. Returning the initial state in the correct form, as specified by the model structure is the only requirement of the AbstractInitializer class, as all implementations of the initializer will be specific to the virus and data sources necessary.
A parameters class is defined by its ability to return a dictionary of parameter names and values. These values may be defined by the user in the configuration file, directly depend on a parameter defined by the user, or sampled from a prior distribution provided by the user. Furthermore, Mechanisms like vaccination uptake or external population introductions are python functions which return values when given time values. Regardless of the form, a parameters class must prepare each parameter within so it is useable by the Runner.
A subset of the AbstractParameters class, the MechanisticInferer handles parameter passing in the special case of inference. In this case, users may define prior distributions they expect parameter values to fall inside of, and leave it to an inference method such as MCMC to infer the parameters using some observed data.
Most inferable parameters (you cannot infer things like age structure) are able to have their static values swapped out with a JSON object that points to a numpyro.distributions class. For example, here we want to sample the R0 of one strain in a 3 strain model, so our STRAIN_R0s parameter looks like
"STRAIN_R0s": [
1.2,
1.8,
{
"distribution": "TransformedDistribution",
"params": {
"base_distribution": {
"distribution": "Beta",
"params": {
"concentration1": 8,
"concentration0": 2
}
},
"transforms": {
"transform": "AffineTransform",
"params": {
"loc": 2.5,
"scale": 1,
"domain": {
"constraint": "interval",
"params": {
"lower_bound": 0,
"upper_bound": 1
}
}
}
}
}
}
],
Currently users are able to define any numpyro distribution with the distribution key, any transformation with the transform key, and any constraint with the constraint key. All parameters used to initialize these priors are passed via the params keyword. Any JSON object that does not follow this pattern of keywords will be loaded as a normal python dictionary.
A much simpler class, the StaticValueParameters class does not implement inference methods, and takes static values and lists of parameters, creating any necessary downstream parameters that depend directly on the value of given parameters, and packaging them neatly for the MechanisticRunner
A relatively simple class, the runner is responsible for actually running a single instance of model odes. It takes as input a dictionary of parameters from AbstractParameters and an initial state from the AbstractInitializer (Except for in the inferer which may also modify initial state). Along with the parameters and initial state, the runner importantly takes an ODE function which returns the rates of change at a single timestep. Example ODE functions are visible in the /model_odes folder. The runner will string together may calls to the ode model to retrieve a timeline of compartment sizes at each timestep and package them all into a diffrax.Solution object for analysis.
the SolutionInterpreter class is responsible for plotting and saving the diffrax.Solution object returned by the MechanisticRunner. For how to plot solution objects, see the Plotting section.