model

This module contains the main disease modelling class.

class summer2.model.BackendType

Bases: object

JAX = 'jax'
PYTHON = 'python'
class summer2.model.CompartmentalModel(times: Tuple[int, int], compartments: List[str], infectious_compartments: List[str], timestep: float = 1.0, ref_date: Optional[datetime] = None)

Bases: object

A compartmental disease model

This model defines a set of compartments which each contain a population. Disease dynamics are defined by a set of flows which link the compartments together. The model is run over a period of time, starting from some initial conditions to predict the future state of a disease.

Parameters
  • times – The start and end times.

  • compartments – The compartments to simulate.

  • infectious_compartments – The compartments which are counted as infectious.

  • time_step (optional) – The timesteps to return results for. This request does not

  • solver (affect the ODE) –

  • 1. (but is used for the stochastic solver. Defaults to) –

times

The times that the model will simulate.

Type

np.ndarray

compartments

The model’s compartments.

Type

List[Compartment]

initial_population

The model’s starting population. The indices of this array match to the compartments attribute. This is zero by default, but should be set with the set_initial_population method.

Type

np.ndarray

outputs

The values of each compartment for each requested timestep. For C compartments and T timesteps this will be a TxC matrix. The column indices of this array will match up with compartments and the row indices will match up with times.

Type

np.ndarray

derived_outputs

Additional results that are calculated from outputs for each timestep.

Type

Dict[str, np.ndarray]

add_computed_value_func(name: str, func: Function)
add_crude_birth_flow(name: str, birth_rate: Union[float, Callable[[float], float]], dest: str, dest_strata: Optional[Dict[str, str]] = None, expected_flow_count: Optional[int] = None)

Adds a crude birth rate flow to the model. The number of births will be determined by the product of the birth rate and total population.

Parameters
  • name – The name of the new flow.

  • birth_rate – The fractional crude birth rate per timestep.

  • dest – The name of the destination compartment.

  • dest_strata (optional) – A whitelist of strata to filter the destination compartments.

  • expected_flow_count (optional) – Used to assert that a particular number of flows are created.

add_death_flow(name: str, death_rate: Union[float, Callable[[float], float]], source: str, source_strata: Optional[Dict[str, str]] = None, expected_flow_count: Optional[int] = None)

Adds a flow where people die and leave the compartment, reducing the total population.

Parameters
  • name – The name of the new flow.

  • death_rate – The fractional death rate per timestep.

  • source – The name of the source compartment.

  • source_strata (optional) – A whitelist of strata to filter the source compartments.

  • expected_flow_count (optional) – Used to assert that a particular number of flows are created.

add_importation_flow(name: str, num_imported: Union[float, Callable[[float], float]], dest: str, split_imports: bool, dest_strata: Optional[Dict[str, str]] = None, expected_flow_count: Optional[int] = None)

Adds an importation flow to the model, where people enter the destination compartment from outside the system. The number of people imported per timestep is completely determined by num_imported.

Parameters
  • name – The name of the new flow.

  • num_imported – The number of people arriving per timestep.

  • dest – The name of the destination compartment.

  • split_imports – Whether to split num_imported amongst existing destination compartments (True), or add the full value to each (False)

  • dest_strata (optional) – A whitelist of strata to filter the destination compartments.

  • expected_flow_count (optional) – Used to assert that a particular number of flows are created.

add_infection_density_flow(name: str, contact_rate: Union[float, Callable[[float], float]], source: str, dest: str, source_strata: Optional[Dict[str, str]] = None, dest_strata: Optional[Dict[str, str]] = None, expected_flow_count: Optional[int] = None)

Adds a flow that infects people using an “infection density” contact rate, which is when the infectious multiplier is determined by the number of infectious people.

Parameters
  • name – The name of the new flow.

  • contact_rate – The contact rate.

  • source – The name of the source compartment.

  • dest – The name of the destination compartment.

  • source_strata (optional) – A whitelist of strata to filter the source compartments.

  • dest_strata (optional) – A whitelist of strata to filter the destination compartments.

  • expected_flow_count (optional) – Used to assert that a particular number of flows are created.

add_infection_frequency_flow(name: str, contact_rate: Union[float, Callable[[float], float]], source: str, dest: str, source_strata: Optional[Dict[str, str]] = None, dest_strata: Optional[Dict[str, str]] = None, expected_flow_count: Optional[int] = None)

Adds a flow that infects people using an “infection frequency” contact rate, which is when the infectious multiplier is determined by the proportion of infectious people to the total population.

Parameters
  • name – The name of the new flow.

  • contact_rate – The effective contact rate.

  • source – The name of the source compartment.

  • dest – The name of the destination compartment.

  • source_strata (optional) – A whitelist of strata to filter the source compartments.

  • dest_strata (optional) – A whitelist of strata to filter the destination compartments.

  • expected_flow_count (optional) – Used to assert that a particular number of flows are created.

add_replacement_birth_flow(name: str, dest: str, dest_strata: Optional[Dict[str, str]] = None, expected_flow_count: Optional[int] = None)

Adds a death-replacing birth flow to the model. The number of births will replace the total number of deaths each year

Parameters
  • name – The name of the new flow.

  • dest – The name of the destination compartment.

  • dest_strata (optional) – A whitelist of strata to filter the destination compartments.

  • expected_flow_count (optional) – Used to assert that a particular number of flows are created.

add_transition_flow(name: str, fractional_rate: Union[float, Callable[[float], float]], source: str, dest: str, source_strata: Optional[Dict[str, str]] = None, dest_strata: Optional[Dict[str, str]] = None, expected_flow_count: Optional[int] = None, absolute=False)

Adds a flow transferring people from a source to a destination based on the population of the source compartment and the fractional flow rate.

Parameters
  • name – The name of the new flow.

  • fractional_rate – The fraction of people that transfer per timestep.

  • source – The name of the source compartment.

  • dest – The name of the destination compartment.

  • source_strata (optional) – A whitelist of strata to filter the source compartments.

  • dest_strata (optional) – A whitelist of strata to filter the destination compartments.

  • expected_flow_count (optional) – Used to assert that a particular number of flows are created.

add_universal_death_flows(name: str, death_rate: Union[float, Callable[[float], float]])

Adds a universal death rate flow to every compartment in the model. The number of deaths per compartment will be determined by the product of the death rate and the compartment population.

Parameters
  • name – The name for the new flow(s).

  • death_rate – The fractional death rate per timestep.

adjust_population_split(strat: str, dest_filter: dict, proportions: dict)

Adjust the initial population to redistribute the population for a particular stratification, over a subset of some other strata

Parameters
  • strat (str) – The stratification to redistribute over

  • dest_filter (dict) – Subset of (other) strata to filter the split by

  • proportions (dict) – Proportions of new split (must have all strata specified)

finalize()
get_computed_value_keys()
get_default_parameters() dict
get_derived_outputs_df()
get_epoch()
get_initial_population(parameters: Optional[dict] = None)
get_input_parameters() set
get_matching_compartments(name: str, strata: dict)
get_object_graph(obj) ComputeGraph

Return the ComputeGraph for this object

Parameters

obj – Any summer object with a graph_key

get_outputs_df()
get_runner(parameters: dict, dyn_params: Optional[List] = None, jit=True, include_full_outputs=True, **backend_args)
get_stratification(name: str) Stratification

Return Stratification matching name, or None if not found

Parameters

name – Name of stratificaton to get

Returns

Stratification

init_population_with_graphobject(init_pop)

Set the initial population with a graphobject (Function, Data or array) The output of this object must be an array matching the model’s final compartments, and therefore no further stratification is possible after this function is called

Parameters

init_pop – GraphObject or array type

query_compartments(query: Optional[dict] = None, tags: Optional[List] = None, as_idx=False)
query_flows(flow_name: Optional[str] = None, source: Optional[dict] = None, dest: Optional[dict] = None, tags: Optional[List] = None)
request_aggregate_output(name: str, sources: List[Union[str, DerivedOutput]], save_results: bool = True) DerivedOutput

Adds a derived output to the model’s results. The output will be the aggregate of other derived outputs.

Parameters
  • name – The name of the derived output.

  • sources – The names of the derived outputs to aggregate.

  • save_results (optional) – Whether to save or discard the results.

request_computed_value_output(name: str, save_results: bool = True) DerivedOutput

Save a computed value process output to derived outputs

Parameters
  • name (str) – Name (key) of computed value process

  • save_results (bool, optional) – Save outputs (or discard if False)

Returns

The resulting DerivedOutput

request_cumulative_output(name: str, source: Union[str, DerivedOutput], start_time: Optional[int] = None, save_results: bool = True) DerivedOutput

Adds a derived output to the model’s results. The output will be the accumulated values of another derived output over the model’s time period.

Parameters
  • name – The name of the derived output.

  • source – The name of the derived outputs to accumulate.

  • start_time (optional) – The time to start accumulating from, defaults to model start time

  • save_results (optional) – Whether to save or discard the results.

request_function_output(name: str, func: Function, save_results: bool = True) DerivedOutput

Request a generic Function output

Parameters
  • name – The name of the derived output

  • func – The Function, whose args are GraphObjects

  • save_results – Whether these results should be available in the final output

request_output_for_compartments(name: str, compartments: List[str], strata: Optional[Dict[str, str]] = None, save_results: bool = True) DerivedOutput

Adds a derived output to the model’s results. The output will be the aggregate population of the requested compartments at the at each timestep.

Parameters
  • name – The name of the derived output.

  • compartments – The name of the compartments to track.

  • strata (optional) – A whitelist of strata to filter the compartments.

  • save_results (optional) – Whether to save or discard the results.

request_output_for_flow(name: str, flow_name: str, source_strata: Optional[Dict[str, str]] = None, dest_strata: Optional[Dict[str, str]] = None, save_results: bool = True, raw_results: bool = False) DerivedOutput

Adds a derived output to the model’s results. The output will be the value of the requested flow at each timestep.

Parameters
  • name – The name of the derived output.

  • flow_name – The name of the flow to track.

  • source_strata (optional) – A whitelist of strata to filter the source compartments.

  • dest_strata (optional) – A whitelist of strata to filter the destination compartments.

  • save_results (optional) – Whether to save or discard the results. Defaults to True.

  • raw_results (optional) – Whether to use raw interpolated flow rates, or post-process them so that they’re more representative of the changes in compartment sizes. Defaults to False.

request_track_modelled_value(name: str, f: GraphObject) DerivedOutput

Save the output of a computegraph Function as a derived output

Parameters
  • name – Name (key) of derived output

  • f – The GraphObject to save

Returns

The resulting DerivedOutput

run(parameters: Optional[dict] = None, solver: str = 'solve_ivp', backend_args: Optional[dict] = None, rebuild: bool = False, **kwargs)

Runs the model over the provided time span, calculating the outputs and the derived outputs. The model calculates the outputs by solving an ODE which is defined by the initial population and the inter-compartmental flows.

Note: The default ODE solver used to produce the model’s outputs does not necessarily evaluate every requested timestep. This adaptive solver can skip over times, or double back when trying to characterize the ODE. The final results are produced by interpolating the solution produced by the ODE solver. This means that model dynamics that only occur in short time periods may not be reflected in the outputs.

Parameters
  • solver (optional) – The ODE solver to use, defaults to SciPy’s IVP solver.

  • **kwargs (optional) – Extra arguments to supplied to the solver, see summer.solver

  • details. (for) –

set_baseline(baseline)

Set a baseline model to be used for this (scenario) run Sets initial population values to the baseline values for this model’s start time Cumulative and relative outputs will refer to the baseline

Parameters

baseline (CompartmentalModel) – The baseline model to be used as reference

set_default_parameters(parameters: dict)
set_derived_outputs_whitelist(whitelist: List[str])

Request that we should only calculate a subset of the model’s derived outputs. This can be useful when you only care about some results and you want to cut down on runtime. For example, we may only need some derived outputs for calibration, but may need more later when we want to know all the dyanmics that the model actually showed.

Parameters

whitelist – A list of the derived output names to calculate, ignoring all others.

set_initial_population(distribution: Dict[str, float], force=False)

Sets the initial population of the model, which is zero by default. Use this method before strafication to determine population by stratification split, or alternatively use init_population_with_graphobject to set the entire population at once

Parameters

distribution – A map of populations to be assigned to compartments.

set_validation_enabled(validate: bool)

Set this to False in order to turn of some (potentially expensive) runtime validation E.g. In calibration, leave it enabled for the first iteration, to catch any structural model issues, but then disable for subsequent iterations

stratify_with(strat: Stratification)

Apply the stratification to the model’s flows and compartments.

Parameters

strat – The stratification to apply.

class summer2.model.ModelResults(model: CompartmentalModel, run_func, runner_dict=None)

Bases: object

get_derived_outputs_df()
get_outputs_df()
run(parameters: dict, filter=True, expand=True, ret_raw=True)
Run the model for a set of input parameters

This is mostly a wrapper for the AuTuMN toolkit, and is not threadsafe! New clients should call run_model instead

Parameters
  • parameters – Input parameters for this run

  • filter – Filter the input parameters dictionary so that only valid input params are passed to the model

  • expand – Expand a nested dictionary into flat parameters

  • ret_raw – Return the outputs of the (jax) _run_func directly

Returns

(optional) returned jax structure from _run_func

run_model(parameters)

_summary_

Parameters
  • parameters (_type_) – _description_

  • out_format (str, optional) – _description_. Defaults to “pandas”.

Returns

_description_

Return type

_type_