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
                            }
                        }
                    }
                }
            }
        }
    ],

Users have the flexibility to define any numpyro distribution by using 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

The StaticValueParameters class is a simpler class that does not include inference methods. It takes static values and lists of parameters as inputs, creates any downstream parameters that directly depend on the given parameter values, and organizes them in a structured manner for the MechanisticRunner.

Runner

The Runner class is responsible for executing a single instance of the model's ordinary differential equations (ODEs). It requires a dictionary of parameters from the AbstractParameters class and an initial state from the AbstractInitializer class (except in the case of the inferer, which may modify the initial state). Additionally, the runner takes an ODE function that calculates the rates of change at each timestep. You can find example ODE functions in the /model_odes folder. The runner iteratively calls these ODEs to obtain a timeline of compartment sizes at each timestep and organizes them into a diffrax.Solution object for further analysis.

Interpreter

The SolutionInterpreter class is responsible for plotting and saving the diffrax.Solution object generated by the MechanisticRunner. It provides functionality to visualize the solution object. For detailed instructions on how to plot solution objects, please refer to the documentation in the plotting section.