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.

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

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

Raises:

ValueError – If the source or target is not in the compartments list.

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.

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

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

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: 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 (np.random.Generator, optional) – 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.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: 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 (np.random.Generator, optional) – 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: 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 (np.random.Generator, optional) – 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.

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

Add a vaccination compartment (Susceptible → Vaccinated → Infected 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 at reduced rate). 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) – Rate of death from Infected. Used when outcome=”deaths”. Default 0.01.

• hospitalization_rate (float) – Rate of hospitalization from Infected. 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.

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

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

Raises:

ValueError – If the source or target is not in the compartments list.

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.

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

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

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: 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 (np.random.Generator, optional) – 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 at reduced rate). 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) – Rate of death from Infected. Used when outcome=”deaths”. Default 0.01.

• hospitalization_rate (float) – Rate of hospitalization from Infected. 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: 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 (np.random.Generator, optional) – 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.