Software Design

Initializer, Params, Runner, Interpreter

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.

Initializer

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. Implementations of the AbstractInitializer will be specific to the virus and required data sources.

Parameters classes

A parameters class is defined by its ability to return a dictionary of parameter names and values. These values can be user-defined in the configuration file, dependent on user-defined parameters, or sampled from a prior distribution provided by the user. Additionally, Mechanisms like vaccination uptake or external population introductions are python functions which return values when given time values. Regardless of their form, parameters classes ensure that each parameter is prepared for use by the Runner.

Inferer

The MechanisticInferer is a subclass of the AbstractParameters class that handles parameter passing specifically for inference purposes. In this case, users can define prior distributions for parameter values and use an inference method like MCMC to estimate the parameters based on observed data.

For most inferable parameters (excluding things like age structure), their static values can be replaced with a JSON object that references a numpyro.distributions class. As an example, let's consider sampling the R0 (basic reproduction number) of one strain in a three-strain model. The STRAIN_R0s parameter would be defined as follows:

 "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.

StaticValueParameters

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

Runner

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 ODEs to retrieve a timeline of compartment sizes at each timestep and package them all into a diffrax.Solution object for analysis.

Interpreter

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.