epydemix.model package

Submodules

epydemix.model.epimodel module

class epydemix.model.epimodel.EpiModel(name: str = 'EpiModel', compartments: List | None = None, parameters: Dict | None = None, population_name: str = 'epydemix_population', population_data_path: str | None = None, contact_layers: List | None = None, contacts_source: str | None = None, age_group_mapping: Dict | None = None, supported_contacts_sources: Dict | None = None, use_default_population: bool = True, default_population_size: int = 100000, data_version: str = 'v1.2.0', attribute: str = 'age')

Bases: object

EpiModel: A compartmental epidemic model simulator

Example

>>> model = EpiModel(compartments=["S", "I", "R"], parameters={"beta": 0.3, "gamma": 0.1})
>>> model.add_transition(source="S", target="I", params=("beta", "I"))
>>> model.add_transition(source="I", target="R", params="gamma")
>>> results = model.run_simulations(
...     start_date="2020-01-01",
...     end_date="2020-12-31",
...     Nsim=100
... )

add_compartments(compartments: List | object) → None

Adds compartments to the epidemic model.

Parameters:

compartments (list or object) – A list of compartments or a single compartment to be added.

Raises:

ValueError – If a compartment is already in the model or is repeated in the input.

Returns:

None

add_intervention(layer_name: str, start_date: str | Timestamp, end_date: str | Timestamp, reduction_factor: float | None = None, new_matrix: ndarray | None = None, name: str = '') → None

Adds an intervention to the epidemic model.

Parameters:

• layer_name (str) – The name of the layer to which the intervention applies.

• start_date (str or datetime) – The start date of the intervention.

• end_date (str or datetime) – The end date of the intervention.

• reduction_factor (float, optional) – The factor by which to reduce the contact matrix. Default is None.

• new_matrix (np.ndarray, optional) – A new contact matrix to use during the intervention. Default is None.

• name (str, optional) – The name of the intervention. Default is an empty string.

Raises:

ValueError – If neither reduction_factor nor new_matrix is provided.

Returns:

None

add_parameter(parameter_name: str | None = None, value: Any = None, parameters_dict: Dict | None = None) → None

Adds new parameters to the model.

Parameters:

• parameter_name (str, optional) – The name of the parameter to add. Ignored if parameters_dict is provided.

• value (any, optional) – The value of the parameter to add. Ignored if parameters_dict is provided.

• parameters_dict (dict, optional) – A dictionary of parameter names and values. If provided, parameter_name and value are ignored and the parameters are updated using this dictionary.

Raises:

ValueError – If neither parameter_name/value nor parameters_dict is provided.

Returns:

None

add_transition(source: str, target: str, kind: str, params: Any) → None

Adds a transition to the epidemic model.

Parameters:

• source (str) – The source compartment of the transition.

• target (str) – The target compartment of the transition.

• kind (str) – The kind of transition (e.g., spontaneous or mediated). Custom kinds must be registered with register_transition_kind before being used here.

• params (Any) – The parameters involved in the transition.

Raises:

ValueError – If the source or target is not in the compartments list, or if the kind is not a registered transition kind.

Returns:

None

apply_intervention(intervention: Dict, simulation_dates: List[Timestamp]) → None

Applies an intervention to the contact matrices for specified simulation dates.

Parameters:

• intervention (dict) – A dictionary containing intervention details with the following keys: - “layer” (str): The name of the layer to which the intervention applies. - “start_date” (pd.Timestamp): The start date of the intervention. - “end_date” (pd.Timestamp): The end date of the intervention. - “reduction_factor” (float, optional): The factor by which to reduce the contact matrix. - “new_matrix” (np.ndarray, optional): A new contact matrix to use during the intervention.

• simulation_dates (list of pd.Timestamp) – A list of dates for which the simulation is run.

Raises:

ValueError – If neither reduction_factor nor new_matrix is provided in the intervention.

Returns:

None

clear_compartments() → None

This method resets the compartments list and clears the compartments_idx dictionary. Transitions are cleared as well, since they refer to compartments that no longer exist.

Returns:

None

clear_interventions() → None

Clears all interventions from the model.

This method resets the interventions list, removing all existing interventions.

Returns:

None

clear_overrides() → None

Clears all parameter overrides from the model.

Returns:

None

clear_parameters() → None

Clears all parameters from the model’s parameters dictionary.

Returns:

None

clear_transitions() → None

Clears all transitions from the model.

This method resets the transitions_list and transitions_idx and reinitializes the transitions dictionary, clearing all existing transitions while keeping the structure intact for each compartment.

Returns:

None

compute_contact_reductions(simulation_dates: List[Timestamp]) → None

Computes the contact reductions for a population over the given simulation dates.

This function applies interventions to the contact matrices and computes the overall contact matrix for each date in the simulation period.

Parameters:

simulation_dates (list of pd.Timestamp) – A list of dates over which the simulation is run.

Returns:

The function updates the instance variable self.Cs with the contact matrices for each date,

including the overall contact matrix after applying interventions.

Return type:

None

create_default_initial_conditions(percentage_in_agents: float = 0.0005) → Dict[str, ndarray]

Creates default initial conditions for the epidemic model. If initial conditions are not provided, assigns a percentage of the population to agent compartments and the rest to source compartments, considering different age groups.

Parameters:

• percentage_in_agents (float) – Percentage of population to place in agent compartments. Defaults to 0.05%.

• initial_conditions (dict, optional) – A dictionary specifying initial conditions. If None, this function creates default conditions.

Returns:

A dictionary with initial conditions for each compartment, with values as arrays representing different age groups.

Return type:

dict

Raises:

ValueError – If the model has no mediated transitions, since the defaults are inferred from them.

delete_override(parameter_name: str) → None

Deletes overrides for a specific parameter.

Parameters:

parameter_name (str) – The name of the parameter whose overrides are to be deleted.

Returns:

None

delete_parameter(parameter_name: str) → Any

Deletes a parameter from the dictionary of the model’s parameters.

Parameters:

parameter_name (str) – The name of the parameter to delete.

Returns:

The value of the deleted parameter.

Return type:

any

Raises:

KeyError – If the parameter with the given name is not found.

get_parameter(parameter_name: str) → Any

Retrieves the value of a specified parameter.

Parameters:

parameter_name (str) – The name of the parameter to retrieve.

Returns:

The value of the specified parameter.

Return type:

any

Raises:

KeyError – If the parameter with the given name is not found.

import_epydemix_population(population_name: str, population_data_path: str | None = None, contact_layers: List | None = None, contacts_source: str | None = None, age_group_mapping: Dict | None = None, supported_contacts_sources: Dict | None = None, data_version: str = 'v1.2.0', attribute: str = 'age') → None

Sets or updates the population for the model.

Parameters:

• population_name (str) – The name of the population to set.

• population_data_path (str or None, optional) – The path to the population data file. If None, data may be fetched online.

• contact_layers (list or None, optional) – A list of contact layers to load for the population. Defaults to [“school”, “work”, “home”, “community”] if None.

• contacts_source (str or None, optional) – The source of contact data. If None, the function will attempt to determine the source automatically.

• age_group_mapping (dict or None, optional) – A mapping for age groups. If None, a default mapping is used.

• supported_contacts_sources (list or None, optional) – A list of supported contact data sources (e.g., “prem_2017”, “prem_2021”). Defaults to [“prem_2017”, “prem_2021”, “mistry_2021”, “litvinova_2025”] for age, [“litvinova_2025”] for sex and race_ethnicity if None.

• data_version (str, optional) – The git tag/version of the epydemix-data repository. Defaults to “v1.2.0”.

• attribute (str, optional) – The demographic attribute layer. Defaults to “age”.

Returns:

None

property n_compartments: int

Get total number of compartments.

property n_transitions: int

Get total number of transitions.

override_parameter(start_date: str, end_date: str, parameter_name: str, value: Any) → None

Adds an override for a parameter with the specified name.

Parameters:

• start_date (str) – The start date for the override period.

• end_date (str) – The end date for the override period.

• parameter_name (str) – The name of the parameter to override.

• value (any) – The value to override the parameter with.

Returns:

None

register_transition_kind(kind: str, function: Callable)

Registers a new transition kind providing a function to compute the rate of the transition.

Parameters:

• kind (str) – The kind of transition (e.g., spontaneous or mediated).

• function (Callable) – The function to register.

Returns:

None

remove_compartment(compartment: str) → List[Transition]

Removes a compartment from the model, together with every transition that involves it: transitions leaving it, entering it, and mediated transitions in which it is the agent.

Transitions of custom kinds that refer to the compartment only through their params cannot be detected and must be removed with remove_transition.

Parameters:

compartment (str) – The compartment to remove.

Raises:

ValueError – If the compartment is not in the model.

Returns:

The transitions removed along with the compartment.

Return type:

List[Transition]

remove_transition(source: str, target: str, kind: str | None = None, agent: str | None = None) → List[Transition]

Removes the transitions going from source to target.

A model can have several transitions between the same two compartments (e.g. infection mediated by two different infectious compartments). By default all of them are removed; use kind and/or agent to remove only some.

Parameters:

• source (str) – The source compartment of the transitions.

• target (str) – The target compartment of the transitions.

• kind (str, optional) – Only remove transitions of this kind (e.g. “mediated”).

• agent (str, optional) – Only remove mediated transitions whose agent (the compartment mediating the transition, e.g. “I” in params=(“beta”, “I”)) is this compartment.

Raises:

ValueError – If no transition matches.

Returns:

The removed transitions.

Return type:

List[Transition]

run_simulations(start_date: str | Timestamp = '2020-01-01', end_date: str | Timestamp = '2020-12-31', initial_conditions_dict: Dict[str, ndarray] | None = None, Nsim: int = 100, percentage_in_agents: float = 0.0005, dt: float | None = 1.0, resample_frequency: str | None = 'D', resample_aggregation_compartments: str | dict | None = 'last', resample_aggregation_transitions: str | dict | None = 'sum', fill_method: str | None = 'ffill', apply_linear_approximation: bool = False, rng: int | Generator | None = None) → SimulationResults

Simulates the epidemic model multiple times over the given time period.

Parameters:

• start_date (str or pd.Timestamp) – The start date of the simulation. Default is “2020-01-01”.

• end_date (str or pd.Timestamp) – The end date of the simulation. Default is “2020-12-31”.

• initial_conditions_dict (dict, optional) – A dictionary of initial conditions for the simulation.

• Nsim (int, optional) – The number of simulation runs to perform (default is 100).

• percentage_in_agents (float, optional) – The percentage of the population to initialize in the agents compartment.

• dt (float, optional) – The time step for the simulation, expressed in days. Default is 1 (day).

• resample_frequency (str, optional) – The frequency at which to resample the simulation results. Default is “D” (daily).

• resample_aggregation_compartments (str, optional) – The aggregation method to use when resampling the compartments. Default is “last”.

• resample_aggregation_transitions (str, optional) – The aggregation method to use when resampling the transitions. Default is “sum”.

• fill_method (str, optional) – Method to fill NaN values after resampling. Default is “ffill”.

• apply_linear_approximation (bool, optional) – Whether to use linear approximation to the probabilities. Default is False.

• rng (int or np.random.Generator, optional) – Seed or random number generator. Default is None.

Returns:

An object containing all simulation trajectories.

Return type:

SimulationResults

Raises:

RuntimeError – If the simulation fails.

set_population(population: Population) → None

Sets the population for the epidemic model.

Parameters:

population (Population) – The population object to set.

Returns:

None

epydemix.model.epimodel.compute_mediated_transition_rate(params, data)

Compute the rate of a mediated transition.

Parameters:

• params – The parameters of the transition provided by the user.

• data – A dictionary containing the data needed for the transition. - parameters: The model parameters - t: The current time step - comp_indices: The indices of the compartments - contact_matrix: The contact matrix - pop: The population in different compartments - pop_sizes: The population sizes

epydemix.model.epimodel.compute_spontaneous_transition_rate(params, data)

Compute the rate of a spontaneous transition.

Parameters:

• params – The parameters of the transition provided by the user.

• data – A dictionary containing the data needed for the transition. - parameters: The model parameters - t: The current time step

Returns:

The rate of the transition

epydemix.model.epimodel.evaluate_rate_expression(expr: str, data: Dict) → ndarray

Evaluate a rate written as an expression of the model parameters (e.g. “beta * (1 - eff)”).

The parameters are fixed for the whole simulation, so the result (one row per time step) does not depend on the current step. When data has an “expression_cache” dictionary, as set up by stochastic_simulation, each expression is evaluated once per simulation and reused at every step. Without it, the expression is evaluated at every call.

Parameters:

• expr – The expression to evaluate.

• data – A dictionary containing the data needed for the transition. - parameters: The model parameters - expression_cache (optional): Expressions already evaluated in this simulation

Returns:

The evaluated expression for all time steps

epydemix.model.epimodel.simulate(epimodel, start_date: str | Timestamp = '2020-01-01', end_date: str | Timestamp = '2020-12-31', initial_conditions_dict: Dict[str, ndarray] | None = None, percentage_in_agents: float = 0.0005, dt: float | None = 1.0, resample_frequency: str | None = 'D', resample_aggregation_compartments: str | dict | None = 'last', resample_aggregation_transitions: str | dict | None = 'sum', fill_method: str | None = 'ffill', apply_linear_approximation: bool = False, rng: int | Generator | None = None, contact_matrices: List[Dict[str, ndarray]] | None = None, simulation_dates: List[Timestamp] | None = None, **kwargs) → Trajectory

Runs a simulation of the epidemic model over the specified simulation dates.

Parameters:

• epimodel (EpiModel) – The epidemic model instance to simulate.

• start_date (str or pd.Timestamp) – The start date of the simulation. Default is “2020-01-01”.

• end_date (str or pd.Timestamp) – The end date of the simulation. Default is “2020-12-31”.

• initial_conditions_dict (dict, optional) – A dictionary of initial conditions for the simulation.

• percentage_in_agents (float, optional) – The percentage of the population to initialize in the agents compartment.

• dt (float, optional) – The time step for the simulation, expressed in days. Default is 1 (day).

• resample_frequency (str, optional) – The frequency at which to resample the simulation results. Default is “D” (daily).

• resample_aggregation_compartments (str, optional) – The aggregation method to use when resampling the compartments. Default is “last”.

• resample_aggregation_transitions (str, optional) – The aggregation method to use when resampling the transitions. Default is “sum”.

• fill_method (str, optional) – The method to use when filling NaN values after resampling. Default is “ffill”.

• apply_linear_approximation (bool, optional) – Whether to use linear approximation to the probabilities. Default is False.

• rng (int or np.random.Generator, optional) – Seed or random number generator. Default is None.

• contact_matrices (list, optional) – A list of contact matrices for the simulation. Default is None.

• simulation_dates (list, optional) – A list of simulation dates. Default is None.

• **kwargs – Additional parameters to overwrite model parameters during the simulation.

Returns:

The trajectory of the simulation

Return type:

Trajectory

Raises:

ValueError – If the model has no transitions defined.

epydemix.model.epimodel.stochastic_simulation(T: int, contact_matrices: List[Dict[str, ndarray]], epimodel, parameters: Dict, initial_conditions: ndarray, dt: float, apply_linear_approximation: bool = False, rng: int | Generator | None = None) → ndarray

Run a stochastic simulation of the epidemic model.

Parameters:

• T – Number of time steps

• contact_matrices – Pre-computed list of contact matrices dictionaries (key is the layer, value is the contact matrix)

• epimodel – The epidemic model

• parameters – Model parameters

• initial_conditions – Initial population distribution

• dt – Time step size

• apply_linear_approximation (bool, optional) – Whether to use linear approximation to the probabilities. Default is False.

• rng (int or np.random.Generator, optional) – Seed or random number generator. Default is None.

epydemix.model.epimodel.validate_transition_function(func: Callable) → None

Validates that a transition rate function has the correct signature and parameters.

Parameters:

func – The transition rate function to validate

Raises:

ValueError – If the function signature doesn’t match requirements

epydemix.model.predefined_models module

epydemix.model.predefined_models.add_outcome(model: EpiModel, outcome: str, mortality_rate: float, hospitalization_rate: float, hospitalization_recovery_rate: float) → EpiModel

Add an outcome compartment (Dead or Hospitalized) branching from the Infected compartment.

mortality_rate / hospitalization_rate are the fraction of the Infected outflow that goes to Dead / Hospitalized; the existing Infected outflow (e.g. Infected -> Recovered) is rescaled by (1 - rate) accordingly.

epydemix.model.predefined_models.add_vaccination(model: EpiModel, vaccination_rate: float, vaccine_efficacy: float) → EpiModel

Add a vaccination compartment (Susceptible → Vaccinated → Infected/Exposed at reduced rate).

epydemix.model.predefined_models.add_waning_immunity(model: EpiModel, waning_rate: float) → EpiModel

Add waning immunity (Recovered → Susceptible) to an existing model.

epydemix.model.predefined_models.create_seiar(transmission_rate: float, incubation_rate: float, recovery_rate: float, asymptomatic_fraction: float, asymptomatic_recovery_rate: float, asymptomatic_relative_infectivity: float) → EpiModel

Create a SEIAR model with symptomatic and asymptomatic infectious compartments.

epydemix.model.predefined_models.create_seir(transmission_rate: float, incubation_rate: float, recovery_rate: float) → EpiModel

Create a SEIR model with the given transmission rate, incubation rate, and recovery rate.

epydemix.model.predefined_models.create_sir(transmission_rate: float, recovery_rate: float) → EpiModel

Create a SIR model with the given transmission rate and recovery rate.

epydemix.model.predefined_models.create_sis(transmission_rate: float, recovery_rate: float) → EpiModel

Create a SIS model with the given transmission rate and recovery rate.

epydemix.model.predefined_models.load_predefined_model(model_name: str, transmission_rate: float = 0.3, recovery_rate: float = 0.1, incubation_rate: float = 0.2, asymptomatic_fraction: float = 0.4, asymptomatic_recovery_rate: float = 0.14, asymptomatic_relative_infectivity: float = 0.5, waning_immunity: bool = False, waning_rate: float = 0.0027397260273972603, vaccination: bool = False, vaccination_rate: float = 0.01, vaccine_efficacy: float = 0.9, outcome: str = None, mortality_rate: float = 0.01, hospitalization_rate: float = 0.01, hospitalization_recovery_rate: float = 0.1) → EpiModel

Load a predefined epidemic model with optional modular extensions.

Parameters:

• model_name (str) – Backbone model. One of “SIR”, “SEIR”, “SIS”, “SEIAR”.

• transmission_rate (float) – Rate of transmission. Default 0.3.

• recovery_rate (float) – Rate of recovery from Infected. Default 0.1.

• incubation_rate (float) – Rate of progression from Exposed to Infected (SEIR, SEIAR). Default 0.2.

• asymptomatic_fraction (float) – Fraction of exposed who become asymptomatic (SEIAR only). Default 0.4.

• asymptomatic_recovery_rate (float) – Recovery rate for asymptomatic individuals (SEIAR only). Default 0.14.

• asymptomatic_relative_infectivity (float) – Infectivity of asymptomatics relative to symptomatics (SEIAR only). Default 0.5.

• waning_immunity (bool) – Add waning immunity (R → S). Not compatible with SIS. Default False.

• waning_rate (float) – Rate of immunity waning. Default 1/365 (~1 year). Interpreted as the inverse of the average immunity duration in days.

• vaccination (bool) – Add vaccination compartment (S → Vaccinated → Infected/Exposed at reduced rate; breakthrough infections go to Exposed if the backbone has an Exposed compartment, otherwise Infected). Default False.

• vaccination_rate (float) – Rate of vaccination from Susceptible. Default 0.01.

• vaccine_efficacy (float) – Fraction by which vaccine reduces transmission. Default 0.9.

• outcome (str or None) – Track a disease outcome. One of None, “deaths”, “hospitalization”. Default None.

• mortality_rate (float) – Fraction of the Infected outflow that goes to Dead, rescaling the existing Infected outflow (e.g. recovery_rate) by (1 - mortality_rate). Used when outcome=”deaths”. Default 0.01.

• hospitalization_rate (float) – Fraction of the Infected outflow that goes to Hospitalized, rescaling recovery_rate by (1 - hospitalization_rate). Used when outcome=”hospitalization”. Default 0.01.

• hospitalization_recovery_rate (float) – Rate of recovery from Hospitalized. Used when outcome=”hospitalization”. Default 0.1.

All rate parameters accept scalars, 1D arrays of shape (T,) for time-varying values, or 2D arrays of shape (T, G) for age-stratified values, consistent with the rest of the epydemix parameter system.

Returns:

Configured epidemic model.

Return type:

EpiModel

Raises:

ValueError – If model_name is not recognised, or a module is incompatible with the chosen backbone.

Examples

SIRS: load_predefined_model(“SIR”, waning_immunity=True) SEIRS: load_predefined_model(“SEIR”, waning_immunity=True) SEIR-V: load_predefined_model(“SEIR”, vaccination=True) SIRD: load_predefined_model(“SIR”, outcome=”deaths”) SEIRD: load_predefined_model(“SEIR”, outcome=”deaths”) SEIRH: load_predefined_model(“SEIR”, outcome=”hospitalization”)

epydemix.model.simulation_output module

class epydemix.model.simulation_output.Trajectory(compartments: Dict[str, ndarray], transitions: Dict[str, ndarray], dates: List[Timestamp], compartment_idx: Dict[str, int], transitions_idx: Dict[str, int], parameters: Dict[str, Any])

Bases: object

Class to store a single trajectory data.

compartments

Dictionary mapping compartment names to arrays of shape (timesteps,)

Type:

Dict[str, np.ndarray]

transitions

Dictionary mapping transition names to arrays of shape (timesteps,)

Type:

Dict[str, np.ndarray]

dates

List of simulation dates

Type:

List[pd.Timestamp]

compartment_idx

Dictionary mapping compartment names to indices

Type:

Dict[str, int]

transitions_idx

Dictionary mapping transition names to indices

Type:

Dict[str, int]

parameters

Dictionary of parameters used in the simulation

Type:

Dict[str, Any]

compartment_idx: Dict[str, int]

compartments: Dict[str, ndarray]

dates: List[Timestamp]

parameters: Dict[str, Any]

resample(freq: str, method_compartments: str = 'last', method_transitions: str = 'sum', fill_method: str = 'ffill') → None

Resample trajectory to new frequency.

Parameters:

• freq (str) – Frequency for resampling (e.g., ‘D’ for daily, ‘W’ for weekly)

• method_compartments (str) – Aggregation method for compartments. Default is ‘last’

• method_transitions (str) – Aggregation method for transitions. Default is ‘sum’

• fill_method (str) – Method to fill NaN values after resampling. Options are: - ‘ffill’: Forward fill (use last valid observation) - ‘bfill’: Backward fill (use next valid observation) - ‘interpolate’: Linear interpolation between points Default is ‘ffill’.

Raises:

ValueError – If fill_method is not one of [‘ffill’, ‘bfill’, ‘interpolate’]

transitions: Dict[str, ndarray]

transitions_idx: Dict[str, int]

epydemix.model.simulation_results module

class epydemix.model.simulation_results.SimulationResults(trajectories: List[Trajectory], parameters: Dict[str, Any])

Bases: object

Class to store and manage multiple simulation results.

trajectories

List of simulation trajectories

Type:

List[Trajectory]

parameters

Dictionary of parameters used in the simulations

Type:

Dict[str, Any]

property Nsim: int

Number of simulations.

property compartment_idx: Dict[str, int]

Compartment indices.

property dates: List[Timestamp]

Simulation dates.

get_quantiles(stacked: Dict[str, ndarray], quantiles: List[float] | None = None, ignore_nan: bool = False) → DataFrame

Compute quantiles across all trajectories.

Parameters:

• stacked – Dictionary of stacked trajectory arrays

• quantiles – List of quantile values to compute. If None, defaults to [0.025, 0.05, 0.25, 0.5, 0.75, 0.95, 0.975]

• ignore_nan – If True, use np.nanquantile to ignore NaN values. Defaults to False. When enabled, a warning is issued if any time point has >50% NaN values, as quantiles may be unreliable with small sample sizes.

get_quantiles_compartments(quantiles: List[float] | None = None, ignore_nan: bool = False, variables: List[str] | None = None) → DataFrame

Compute quantiles across all trajectories for compartments. The name of the columns are the compartments names and the demographic groups, in the following format: {compartment_name}_{demographic_group}. For example, the column I_total contains the quantiles of the number of infected (“I”) individuals across all demographic groups (“total”).

Parameters:

• quantiles – List of quantile values to compute. If None, defaults to [0.025, 0.05, 0.25, 0.5, 0.75, 0.95, 0.975]

• ignore_nan – If True, use np.nanquantile to ignore NaN values. Defaults to False.

• variables – List of compartment names to include. If None, all compartments are included.

get_quantiles_transitions(quantiles: List[float] | None = None, ignore_nan: bool = False, variables: List[str] | None = None) → DataFrame

Compute quantiles across all trajectories for transitions. The name of the columns are the transitions names and the demographic groups, in the following format: {source_compartment_name}_to_{target_compartment_name}_{demographic_group}. For example, the column S_to_I_total contains the quantiles of the number of individuals transitioning from susceptible (“S”) to infected (“I”) individuals across all demographic groups (“total”).

Parameters:

• quantiles – List of quantile values to compute. If None, defaults to [0.025, 0.05, 0.25, 0.5, 0.75, 0.95, 0.975]

• ignore_nan – If True, use np.nanquantile to ignore NaN values. Defaults to False.

• variables – List of transition names to include. If None, all transitions are included.

get_stacked_compartments(variables: List[str] | None = None) → Dict[str, ndarray]

Get trajectories stacked into arrays of shape (Nsim, timesteps).

Parameters:

variables – List of compartment names to include. If None, all compartments are included.

get_stacked_transitions(variables: List[str] | None = None) → Dict[str, ndarray]

Get trajectories stacked into arrays of shape (Nsim, timesteps).

Parameters:

variables – List of transition names to include. If None, all transitions are included.

parameters: Dict[str, Any]

resample(freq: str, method: str = 'last', fill_method: str = 'ffill') → SimulationResults

Resample all trajectories to new frequency.

trajectories: List[Trajectory]

epydemix.model.transition module

class epydemix.model.transition.Transition(source: str, target: str, kind: str, params: Any)

Bases: object

Represents a transition in an epidemic model.

A transition defines a change from one compartment to another with a specified rate. Optionally, it can involve a specific agent.

source

The source compartment of the transition.

Type:

str

target

The target compartment of the transition.

Type:

str

kind

The kind of transition (e.g., spontaneous or mediated).

Type:

str

params

The parameters involved in the transition.

Type:

Any

Module contents

class epydemix.model.EpiModel(name: str = 'EpiModel', compartments: List | None = None, parameters: Dict | None = None, population_name: str = 'epydemix_population', population_data_path: str | None = None, contact_layers: List | None = None, contacts_source: str | None = None, age_group_mapping: Dict | None = None, supported_contacts_sources: Dict | None = None, use_default_population: bool = True, default_population_size: int = 100000, data_version: str = 'v1.2.0', attribute: str = 'age')

Bases: object

EpiModel: A compartmental epidemic model simulator

Example

>>> model = EpiModel(compartments=["S", "I", "R"], parameters={"beta": 0.3, "gamma": 0.1})
>>> model.add_transition(source="S", target="I", params=("beta", "I"))
>>> model.add_transition(source="I", target="R", params="gamma")
>>> results = model.run_simulations(
...     start_date="2020-01-01",
...     end_date="2020-12-31",
...     Nsim=100
... )

add_compartments(compartments: List | object) → None

Adds compartments to the epidemic model.

Parameters:

compartments (list or object) – A list of compartments or a single compartment to be added.

Raises:

ValueError – If a compartment is already in the model or is repeated in the input.

Returns:

None

add_intervention(layer_name: str, start_date: str | Timestamp, end_date: str | Timestamp, reduction_factor: float | None = None, new_matrix: ndarray | None = None, name: str = '') → None

Adds an intervention to the epidemic model.

Parameters:

• layer_name (str) – The name of the layer to which the intervention applies.

• start_date (str or datetime) – The start date of the intervention.

• end_date (str or datetime) – The end date of the intervention.

• reduction_factor (float, optional) – The factor by which to reduce the contact matrix. Default is None.

• new_matrix (np.ndarray, optional) – A new contact matrix to use during the intervention. Default is None.

• name (str, optional) – The name of the intervention. Default is an empty string.

Raises:

ValueError – If neither reduction_factor nor new_matrix is provided.

Returns:

None

add_parameter(parameter_name: str | None = None, value: Any = None, parameters_dict: Dict | None = None) → None

Adds new parameters to the model.

Parameters:

• parameter_name (str, optional) – The name of the parameter to add. Ignored if parameters_dict is provided.

• value (any, optional) – The value of the parameter to add. Ignored if parameters_dict is provided.

• parameters_dict (dict, optional) – A dictionary of parameter names and values. If provided, parameter_name and value are ignored and the parameters are updated using this dictionary.

Raises:

ValueError – If neither parameter_name/value nor parameters_dict is provided.

Returns:

None

add_transition(source: str, target: str, kind: str, params: Any) → None

Adds a transition to the epidemic model.

Parameters:

• source (str) – The source compartment of the transition.

• target (str) – The target compartment of the transition.

• kind (str) – The kind of transition (e.g., spontaneous or mediated). Custom kinds must be registered with register_transition_kind before being used here.

• params (Any) – The parameters involved in the transition.

Raises:

ValueError – If the source or target is not in the compartments list, or if the kind is not a registered transition kind.

Returns:

None

apply_intervention(intervention: Dict, simulation_dates: List[Timestamp]) → None

Applies an intervention to the contact matrices for specified simulation dates.

Parameters:

• intervention (dict) – A dictionary containing intervention details with the following keys: - “layer” (str): The name of the layer to which the intervention applies. - “start_date” (pd.Timestamp): The start date of the intervention. - “end_date” (pd.Timestamp): The end date of the intervention. - “reduction_factor” (float, optional): The factor by which to reduce the contact matrix. - “new_matrix” (np.ndarray, optional): A new contact matrix to use during the intervention.

• simulation_dates (list of pd.Timestamp) – A list of dates for which the simulation is run.

Raises:

ValueError – If neither reduction_factor nor new_matrix is provided in the intervention.

Returns:

None

clear_compartments() → None

This method resets the compartments list and clears the compartments_idx dictionary. Transitions are cleared as well, since they refer to compartments that no longer exist.

Returns:

None

clear_interventions() → None

Clears all interventions from the model.

This method resets the interventions list, removing all existing interventions.

Returns:

None

clear_overrides() → None

Clears all parameter overrides from the model.

Returns:

None

clear_parameters() → None

Clears all parameters from the model’s parameters dictionary.

Returns:

None

clear_transitions() → None

Clears all transitions from the model.

This method resets the transitions_list and transitions_idx and reinitializes the transitions dictionary, clearing all existing transitions while keeping the structure intact for each compartment.

Returns:

None

compute_contact_reductions(simulation_dates: List[Timestamp]) → None

Computes the contact reductions for a population over the given simulation dates.

This function applies interventions to the contact matrices and computes the overall contact matrix for each date in the simulation period.

Parameters:

simulation_dates (list of pd.Timestamp) – A list of dates over which the simulation is run.

Returns:

The function updates the instance variable self.Cs with the contact matrices for each date,

including the overall contact matrix after applying interventions.

Return type:

None

create_default_initial_conditions(percentage_in_agents: float = 0.0005) → Dict[str, ndarray]

Creates default initial conditions for the epidemic model. If initial conditions are not provided, assigns a percentage of the population to agent compartments and the rest to source compartments, considering different age groups.

Parameters:

• percentage_in_agents (float) – Percentage of population to place in agent compartments. Defaults to 0.05%.

• initial_conditions (dict, optional) – A dictionary specifying initial conditions. If None, this function creates default conditions.

Returns:

A dictionary with initial conditions for each compartment, with values as arrays representing different age groups.

Return type:

dict

Raises:

ValueError – If the model has no mediated transitions, since the defaults are inferred from them.

delete_override(parameter_name: str) → None

Deletes overrides for a specific parameter.

Parameters:

parameter_name (str) – The name of the parameter whose overrides are to be deleted.

Returns:

None

delete_parameter(parameter_name: str) → Any

Deletes a parameter from the dictionary of the model’s parameters.

Parameters:

parameter_name (str) – The name of the parameter to delete.

Returns:

The value of the deleted parameter.

Return type:

any

Raises:

KeyError – If the parameter with the given name is not found.

get_parameter(parameter_name: str) → Any

Retrieves the value of a specified parameter.

Parameters:

parameter_name (str) – The name of the parameter to retrieve.

Returns:

The value of the specified parameter.

Return type:

any

Raises:

KeyError – If the parameter with the given name is not found.

import_epydemix_population(population_name: str, population_data_path: str | None = None, contact_layers: List | None = None, contacts_source: str | None = None, age_group_mapping: Dict | None = None, supported_contacts_sources: Dict | None = None, data_version: str = 'v1.2.0', attribute: str = 'age') → None

Sets or updates the population for the model.

Parameters:

• population_name (str) – The name of the population to set.

• population_data_path (str or None, optional) – The path to the population data file. If None, data may be fetched online.

• contact_layers (list or None, optional) – A list of contact layers to load for the population. Defaults to [“school”, “work”, “home”, “community”] if None.

• contacts_source (str or None, optional) – The source of contact data. If None, the function will attempt to determine the source automatically.

• age_group_mapping (dict or None, optional) – A mapping for age groups. If None, a default mapping is used.

• supported_contacts_sources (list or None, optional) – A list of supported contact data sources (e.g., “prem_2017”, “prem_2021”). Defaults to [“prem_2017”, “prem_2021”, “mistry_2021”, “litvinova_2025”] for age, [“litvinova_2025”] for sex and race_ethnicity if None.

• data_version (str, optional) – The git tag/version of the epydemix-data repository. Defaults to “v1.2.0”.

• attribute (str, optional) – The demographic attribute layer. Defaults to “age”.

Returns:

None

property n_compartments: int

Get total number of compartments.

property n_transitions: int

Get total number of transitions.

override_parameter(start_date: str, end_date: str, parameter_name: str, value: Any) → None

Adds an override for a parameter with the specified name.

Parameters:

• start_date (str) – The start date for the override period.

• end_date (str) – The end date for the override period.

• parameter_name (str) – The name of the parameter to override.

• value (any) – The value to override the parameter with.

Returns:

None

register_transition_kind(kind: str, function: Callable)

Registers a new transition kind providing a function to compute the rate of the transition.

Parameters:

• kind (str) – The kind of transition (e.g., spontaneous or mediated).

• function (Callable) – The function to register.

Returns:

None

remove_compartment(compartment: str) → List[Transition]

Removes a compartment from the model, together with every transition that involves it: transitions leaving it, entering it, and mediated transitions in which it is the agent.

Transitions of custom kinds that refer to the compartment only through their params cannot be detected and must be removed with remove_transition.

Parameters:

compartment (str) – The compartment to remove.

Raises:

ValueError – If the compartment is not in the model.

Returns:

The transitions removed along with the compartment.

Return type:

List[Transition]

remove_transition(source: str, target: str, kind: str | None = None, agent: str | None = None) → List[Transition]

Removes the transitions going from source to target.

A model can have several transitions between the same two compartments (e.g. infection mediated by two different infectious compartments). By default all of them are removed; use kind and/or agent to remove only some.

Parameters:

• source (str) – The source compartment of the transitions.

• target (str) – The target compartment of the transitions.

• kind (str, optional) – Only remove transitions of this kind (e.g. “mediated”).

• agent (str, optional) – Only remove mediated transitions whose agent (the compartment mediating the transition, e.g. “I” in params=(“beta”, “I”)) is this compartment.

Raises:

ValueError – If no transition matches.

Returns:

The removed transitions.

Return type:

List[Transition]

run_simulations(start_date: str | Timestamp = '2020-01-01', end_date: str | Timestamp = '2020-12-31', initial_conditions_dict: Dict[str, ndarray] | None = None, Nsim: int = 100, percentage_in_agents: float = 0.0005, dt: float | None = 1.0, resample_frequency: str | None = 'D', resample_aggregation_compartments: str | dict | None = 'last', resample_aggregation_transitions: str | dict | None = 'sum', fill_method: str | None = 'ffill', apply_linear_approximation: bool = False, rng: int | Generator | None = None) → SimulationResults

Simulates the epidemic model multiple times over the given time period.

Parameters:

• start_date (str or pd.Timestamp) – The start date of the simulation. Default is “2020-01-01”.

• end_date (str or pd.Timestamp) – The end date of the simulation. Default is “2020-12-31”.

• initial_conditions_dict (dict, optional) – A dictionary of initial conditions for the simulation.

• Nsim (int, optional) – The number of simulation runs to perform (default is 100).

• percentage_in_agents (float, optional) – The percentage of the population to initialize in the agents compartment.

• dt (float, optional) – The time step for the simulation, expressed in days. Default is 1 (day).

• resample_frequency (str, optional) – The frequency at which to resample the simulation results. Default is “D” (daily).

• resample_aggregation_compartments (str, optional) – The aggregation method to use when resampling the compartments. Default is “last”.

• resample_aggregation_transitions (str, optional) – The aggregation method to use when resampling the transitions. Default is “sum”.

• fill_method (str, optional) – Method to fill NaN values after resampling. Default is “ffill”.

• apply_linear_approximation (bool, optional) – Whether to use linear approximation to the probabilities. Default is False.

• rng (int or np.random.Generator, optional) – Seed or random number generator. Default is None.

Returns:

An object containing all simulation trajectories.

Return type:

SimulationResults

Raises:

RuntimeError – If the simulation fails.

set_population(population: Population) → None

Sets the population for the epidemic model.

Parameters:

population (Population) – The population object to set.

Returns:

None

class epydemix.model.SimulationResults(trajectories: List[Trajectory], parameters: Dict[str, Any])

Bases: object

Class to store and manage multiple simulation results.

trajectories

List of simulation trajectories

Type:

List[Trajectory]

parameters

Dictionary of parameters used in the simulations

Type:

Dict[str, Any]

property Nsim: int

Number of simulations.

property compartment_idx: Dict[str, int]

Compartment indices.

property dates: List[Timestamp]

Simulation dates.

get_quantiles(stacked: Dict[str, ndarray], quantiles: List[float] | None = None, ignore_nan: bool = False) → DataFrame

Compute quantiles across all trajectories.

Parameters:

• stacked – Dictionary of stacked trajectory arrays

• quantiles – List of quantile values to compute. If None, defaults to [0.025, 0.05, 0.25, 0.5, 0.75, 0.95, 0.975]

• ignore_nan – If True, use np.nanquantile to ignore NaN values. Defaults to False. When enabled, a warning is issued if any time point has >50% NaN values, as quantiles may be unreliable with small sample sizes.

get_quantiles_compartments(quantiles: List[float] | None = None, ignore_nan: bool = False, variables: List[str] | None = None) → DataFrame

Compute quantiles across all trajectories for compartments. The name of the columns are the compartments names and the demographic groups, in the following format: {compartment_name}_{demographic_group}. For example, the column I_total contains the quantiles of the number of infected (“I”) individuals across all demographic groups (“total”).

Parameters:

• quantiles – List of quantile values to compute. If None, defaults to [0.025, 0.05, 0.25, 0.5, 0.75, 0.95, 0.975]

• ignore_nan – If True, use np.nanquantile to ignore NaN values. Defaults to False.

• variables – List of compartment names to include. If None, all compartments are included.

get_quantiles_transitions(quantiles: List[float] | None = None, ignore_nan: bool = False, variables: List[str] | None = None) → DataFrame

Compute quantiles across all trajectories for transitions. The name of the columns are the transitions names and the demographic groups, in the following format: {source_compartment_name}_to_{target_compartment_name}_{demographic_group}. For example, the column S_to_I_total contains the quantiles of the number of individuals transitioning from susceptible (“S”) to infected (“I”) individuals across all demographic groups (“total”).

Parameters:

• quantiles – List of quantile values to compute. If None, defaults to [0.025, 0.05, 0.25, 0.5, 0.75, 0.95, 0.975]

• ignore_nan – If True, use np.nanquantile to ignore NaN values. Defaults to False.

• variables – List of transition names to include. If None, all transitions are included.

get_stacked_compartments(variables: List[str] | None = None) → Dict[str, ndarray]

Get trajectories stacked into arrays of shape (Nsim, timesteps).

Parameters:

variables – List of compartment names to include. If None, all compartments are included.

get_stacked_transitions(variables: List[str] | None = None) → Dict[str, ndarray]

Get trajectories stacked into arrays of shape (Nsim, timesteps).

Parameters:

variables – List of transition names to include. If None, all transitions are included.

parameters: Dict[str, Any]

resample(freq: str, method: str = 'last', fill_method: str = 'ffill') → SimulationResults

Resample all trajectories to new frequency.

trajectories: List[Trajectory]

class epydemix.model.Transition(source: str, target: str, kind: str, params: Any)

Bases: object

Represents a transition in an epidemic model.

A transition defines a change from one compartment to another with a specified rate. Optionally, it can involve a specific agent.

source

The source compartment of the transition.

Type:

str

target

The target compartment of the transition.

Type:

str

kind

The kind of transition (e.g., spontaneous or mediated).

Type:

str

params

The parameters involved in the transition.

Type:

Any

epydemix.model.load_predefined_model(model_name: str, transmission_rate: float = 0.3, recovery_rate: float = 0.1, incubation_rate: float = 0.2, asymptomatic_fraction: float = 0.4, asymptomatic_recovery_rate: float = 0.14, asymptomatic_relative_infectivity: float = 0.5, waning_immunity: bool = False, waning_rate: float = 0.0027397260273972603, vaccination: bool = False, vaccination_rate: float = 0.01, vaccine_efficacy: float = 0.9, outcome: str = None, mortality_rate: float = 0.01, hospitalization_rate: float = 0.01, hospitalization_recovery_rate: float = 0.1) → EpiModel

Load a predefined epidemic model with optional modular extensions.

Parameters:

• model_name (str) – Backbone model. One of “SIR”, “SEIR”, “SIS”, “SEIAR”.

• transmission_rate (float) – Rate of transmission. Default 0.3.

• recovery_rate (float) – Rate of recovery from Infected. Default 0.1.

• incubation_rate (float) – Rate of progression from Exposed to Infected (SEIR, SEIAR). Default 0.2.

• asymptomatic_fraction (float) – Fraction of exposed who become asymptomatic (SEIAR only). Default 0.4.

• asymptomatic_recovery_rate (float) – Recovery rate for asymptomatic individuals (SEIAR only). Default 0.14.

• asymptomatic_relative_infectivity (float) – Infectivity of asymptomatics relative to symptomatics (SEIAR only). Default 0.5.

• waning_immunity (bool) – Add waning immunity (R → S). Not compatible with SIS. Default False.

• waning_rate (float) – Rate of immunity waning. Default 1/365 (~1 year). Interpreted as the inverse of the average immunity duration in days.

• vaccination (bool) – Add vaccination compartment (S → Vaccinated → Infected/Exposed at reduced rate; breakthrough infections go to Exposed if the backbone has an Exposed compartment, otherwise Infected). Default False.

• vaccination_rate (float) – Rate of vaccination from Susceptible. Default 0.01.

• vaccine_efficacy (float) – Fraction by which vaccine reduces transmission. Default 0.9.

• outcome (str or None) – Track a disease outcome. One of None, “deaths”, “hospitalization”. Default None.

• mortality_rate (float) – Fraction of the Infected outflow that goes to Dead, rescaling the existing Infected outflow (e.g. recovery_rate) by (1 - mortality_rate). Used when outcome=”deaths”. Default 0.01.

• hospitalization_rate (float) – Fraction of the Infected outflow that goes to Hospitalized, rescaling recovery_rate by (1 - hospitalization_rate). Used when outcome=”hospitalization”. Default 0.01.

• hospitalization_recovery_rate (float) – Rate of recovery from Hospitalized. Used when outcome=”hospitalization”. Default 0.1.

All rate parameters accept scalars, 1D arrays of shape (T,) for time-varying values, or 2D arrays of shape (T, G) for age-stratified values, consistent with the rest of the epydemix parameter system.

Returns:

Configured epidemic model.

Return type:

EpiModel

Raises:

ValueError – If model_name is not recognised, or a module is incompatible with the chosen backbone.

Examples

SIRS: load_predefined_model(“SIR”, waning_immunity=True) SEIRS: load_predefined_model(“SEIR”, waning_immunity=True) SEIR-V: load_predefined_model(“SEIR”, vaccination=True) SIRD: load_predefined_model(“SIR”, outcome=”deaths”) SEIRD: load_predefined_model(“SEIR”, outcome=”deaths”) SEIRH: load_predefined_model(“SEIR”, outcome=”hospitalization”)

epydemix.model.simulate(epimodel, start_date: str | Timestamp = '2020-01-01', end_date: str | Timestamp = '2020-12-31', initial_conditions_dict: Dict[str, ndarray] | None = None, percentage_in_agents: float = 0.0005, dt: float | None = 1.0, resample_frequency: str | None = 'D', resample_aggregation_compartments: str | dict | None = 'last', resample_aggregation_transitions: str | dict | None = 'sum', fill_method: str | None = 'ffill', apply_linear_approximation: bool = False, rng: int | Generator | None = None, contact_matrices: List[Dict[str, ndarray]] | None = None, simulation_dates: List[Timestamp] | None = None, **kwargs) → Trajectory

Runs a simulation of the epidemic model over the specified simulation dates.

Parameters:

• epimodel (EpiModel) – The epidemic model instance to simulate.

• start_date (str or pd.Timestamp) – The start date of the simulation. Default is “2020-01-01”.

• end_date (str or pd.Timestamp) – The end date of the simulation. Default is “2020-12-31”.

• initial_conditions_dict (dict, optional) – A dictionary of initial conditions for the simulation.

• percentage_in_agents (float, optional) – The percentage of the population to initialize in the agents compartment.

• dt (float, optional) – The time step for the simulation, expressed in days. Default is 1 (day).

• resample_frequency (str, optional) – The frequency at which to resample the simulation results. Default is “D” (daily).

• resample_aggregation_compartments (str, optional) – The aggregation method to use when resampling the compartments. Default is “last”.

• resample_aggregation_transitions (str, optional) – The aggregation method to use when resampling the transitions. Default is “sum”.

• fill_method (str, optional) – The method to use when filling NaN values after resampling. Default is “ffill”.

• apply_linear_approximation (bool, optional) – Whether to use linear approximation to the probabilities. Default is False.

• rng (int or np.random.Generator, optional) – Seed or random number generator. Default is None.

• contact_matrices (list, optional) – A list of contact matrices for the simulation. Default is None.

• simulation_dates (list, optional) – A list of simulation dates. Default is None.

• **kwargs – Additional parameters to overwrite model parameters during the simulation.

Returns:

The trajectory of the simulation

Return type:

Trajectory

Raises:

ValueError – If the model has no transitions defined.