Lower level

These functions are used internally. They are documented for people extending epydemix and are not needed for normal use.

Simulation engine

epydemix.model.epimodel.stochastic_simulation(T, contact_matrices, epimodel, parameters, initial_conditions, dt, apply_linear_approximation=False, rng=None)

Run a stochastic simulation of the epidemic model.

Parameters:

• T (int) – Number of time steps

• contact_matrices (List[Dict[str, ndarray]]) – Pre-computed list of contact matrices dictionaries (key is the layer, value is the contact matrix)

• epimodel – The epidemic model

• parameters (Dict) – Model parameters

• initial_conditions (ndarray) – Initial population distribution

• dt (float) – 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.

Return type:

ndarray

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.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.evaluate_rate_expression(expr, data)

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 (str) – The expression to evaluate.

• data (Dict) – 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

Return type:

ndarray

epydemix.model.epimodel.validate_transition_function(func)

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

Parameters:

func (Callable) – The transition rate function to validate

Raises:

ValueError – If the function signature doesn’t match requirements

Return type:

None

Building predefined models

epydemix.model.predefined_models.create_sir(transmission_rate, recovery_rate)

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

Parameters:

• transmission_rate (float)

• recovery_rate (float)

Return type:

EpiModel

epydemix.model.predefined_models.create_seir(transmission_rate, incubation_rate, recovery_rate)

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

Parameters:

• transmission_rate (float)

• incubation_rate (float)

• recovery_rate (float)

Return type:

EpiModel

epydemix.model.predefined_models.create_sis(transmission_rate, recovery_rate)

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

Parameters:

• transmission_rate (float)

• recovery_rate (float)

Return type:

EpiModel

epydemix.model.predefined_models.create_seiar(transmission_rate, incubation_rate, recovery_rate, asymptomatic_fraction, asymptomatic_recovery_rate, asymptomatic_relative_infectivity)

Create a SEIAR model with symptomatic and asymptomatic infectious compartments.

Parameters:

• transmission_rate (float)

• incubation_rate (float)

• recovery_rate (float)

• asymptomatic_fraction (float)

• asymptomatic_recovery_rate (float)

• asymptomatic_relative_infectivity (float)

Return type:

EpiModel

epydemix.model.predefined_models.add_waning_immunity(model, waning_rate)

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

Parameters:

• model (EpiModel)

• waning_rate (float)

Return type:

EpiModel

epydemix.model.predefined_models.add_vaccination(model, vaccination_rate, vaccine_efficacy)

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

Parameters:

• model (EpiModel)

• vaccination_rate (float)

• vaccine_efficacy (float)

Return type:

EpiModel

epydemix.model.predefined_models.add_outcome(model, outcome, mortality_rate, hospitalization_rate, hospitalization_recovery_rate)

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.

Parameters:

• model (EpiModel)

• outcome (str)

• mortality_rate (float)

• hospitalization_rate (float)

• hospitalization_recovery_rate (float)

Return type:

EpiModel

Population helpers

epydemix.population.population.aggregate_matrix(initial_matrix, old_population, new_population, age_group_mapping, old_age_groups_idx, new_age_group_idx)

Aggregates a contact matrix based on new demographic groupings.

Parameters:

• initial_matrix (np.ndarray) – The initial contact matrix (rates) between old demographic groups.

• old_population (np.ndarray) – The population sizes of the old demographic groups.

• new_population (np.ndarray) – The population sizes of the new aggregated demographic groups.

• age_group_mapping (Dict[str, list]) – A dictionary mapping new demographic group names to lists of old group names.

• old_age_groups_idx (Dict[str, int]) – A dictionary mapping old age group names to their indices in the contact matrix.

• new_age_group_idx (Dict[str, int]) – A dictionary mapping new age group names to their indices in the aggregated matrix.

Returns:

The aggregated contact matrix (rates) for the new demographic groups.

Return type:

np.ndarray

epydemix.population.population.aggregate_demographic(data, grouping)

Aggregates demographic data based on a grouping dictionary.

Parameters:

• data (pd.DataFrame) – A DataFrame containing demographic data with columns ‘group_name’ and ‘value’.

• grouping (Dict[str, List[str]]) – A dictionary where keys are new group names and values are lists of original group names to aggregate.

Returns:

A DataFrame with two columns: ‘group_name’ and ‘value’, where ‘value’ is the sum of the ‘value’ column from the original DataFrame for each new group.

Return type:

pd.DataFrame

epydemix.population.population.map_age_groups_to_idx(age_group_mapping, old_age_groups_idx, new_age_group_idx)

Maps old age groups to new age groups using index mappings.

Parameters:

• age_group_mapping (Dict[str, List[str]]) – A dictionary where keys are new age groups, and values are lists of old age groups.

• old_age_groups_idx (Dict[str, int]) – A dictionary mapping old age group names to their respective indices.

• new_age_group_idx (Dict[str, int]) – A dictionary mapping new age group names to their respective indices.

Returns:

A dictionary mapping old age group indices to new age group indices.

Return type:

Dict[int, int]

ABC-SMC building blocks

class epydemix.utils.abc_smc_utils.Perturbation(param_name)

abstractmethod propose(x, rng=None)

Propose a new value from the current value x.

Draw randomness from rng (a seed or np.random.Generator) so that proposals are reproducible

abstractmethod pdf(x, center)

Evaluate the PDF of the kernel.

abstractmethod update(particles, weights, param_names)

Update the kernel parameters based on particles and weights.

class epydemix.utils.abc_smc_utils.DefaultPerturbationContinuous(param_name)

Componenent-wise normal perturbation kernel with adaptive standard deviation (Beaumont et al. (2009)).

propose(x, rng=None)

Propose a new value based on the current value.

pdf(x, center)

Evaluate the PDF of the kernel.

update(particles, weights, param_names)

Update the standard deviation based on previous generation variance.

class epydemix.utils.abc_smc_utils.DefaultPerturbationDiscrete(param_name, prior, jump_probability=0.3)

propose(x, rng=None)

Propose a new value for the discrete parameter.

pdf(x, center)

Transition probability for the discrete parameter.

update(particles, weights, param_names)

Update jump_probability or other characteristics if needed.

epydemix.utils.abc_smc_utils.sample_prior(priors, param_names, rng=None)

Samples a parameter set from the given prior distributions. priors: dictionary mapping parameter names to scipy.stats distributions param_names: list of parameter names to maintain consistent order rng: optional np.random.Generator (or seed) used for sampling Returns: list of sampled parameter values in the order of param_names

epydemix.utils.abc_smc_utils.compute_effective_sample_size(weights)

Computes the effective sample size (ESS) of a set of weights.

Parameters:

weights (np.ndarray) – A NumPy array of weights for the particles.

Returns:

The effective sample size, calculated as the reciprocal of the sum of squared weights.

Return type:

float

epydemix.utils.abc_smc_utils.weighted_quantile(values, weights, quantile)

Compute the weighted quantile of a dataset.

Parameters:

• values (Union[np.ndarray, list]) – Data values for which the quantile is to be computed.

• weights (Union[np.ndarray, list]) – Weights associated with the data values.

• quantile (float) – The quantile to compute, must be between 0 and 1.

Returns:

The weighted quantile value.

Return type:

float

Raises:

ValueError – If quantile is not between 0 and 1.

General utilities

epydemix.utils.utils.validate_parameter_shape(key, value, T, n_age)

Validates the shape of the input value based on its type and expected dimensions.

Parameters:

• key (str) – The key associated with the value in the parameters dictionary.

• value (Union[np.ndarray, Iterable]) – The value to be validated, which can be a NumPy array or an iterable.

• T (int) – The expected length of the first dimension.

• n_age (int) – The expected length of the second dimension.

Raises:

ValueError – If the value does not meet the required shape criteria or if the type is unsupported.

Return type:

None

epydemix.utils.utils.resize_parameter(value, T, n_age)

Resizes the input value to have the shape (T, n_age).

Parameters:

• value (Union[np.ndarray, int, float]) – The value to be resized, which can be a NumPy array or a scalar.

• T (int) – The length of the first dimension.

• n_age (int) – The length of the second dimension.

Returns:

A 2D array with shape (T, n_age).

Return type:

np.ndarray

epydemix.utils.utils.create_definitions(parameters, T, n_age)

Generates a dictionary where each value is a 2D array based on the input dictionary and the parameters provided.

Parameters:

• parameters (Dict[str, Union[np.ndarray, int, float, Iterable]]) – A dictionary where values can be either scalars or iterables.

• T (int) – The length of the first dimension of the arrays to be created.

• n_age (int) – The length of the second dimension of the arrays to be created.

Returns:

A dictionary where keys are the same as in parameters and values are 2D arrays of shape (T, n_age).

Return type:

Dict[str, np.ndarray]

Raises:

ValueError – If any parameter value does not meet the required shape criteria.

epydemix.utils.utils.format_simulation_output(compartments_evolution, transitions_evolution, compartments_idx, transitions_idx, demographics)

Format simulation output into dictionary with proper naming.

Parameters:

• compartments_evolution (ndarray) – Array of shape (timesteps, n_compartments, n_demographics)

• transitions_evolution (ndarray) – Array of shape (timesteps, n_transitions, n_demographics)

• compartments_idx (Dict[str, int]) – Dictionary mapping compartment names to their indices

• transitions_idx (Dict[str, int]) – Dictionary mapping transition names to their indices

• demographics (List[str]) – List of demographic group names

Return type:

Dict[str, ndarray]

epydemix.utils.utils.str_to_date(date_str)

Converts a date string in the format ‘YYYY-MM-DD’ to a datetime.date object.

Parameters:

date_str (str) – A string representing a date in the format ‘YYYY-MM-DD’.

Returns:

A datetime.date object corresponding to the input string.

Return type:

datetime.date

Raises:

ValueError – If the input string is not in the correct format.

epydemix.utils.utils.apply_overrides(definitions, overrides, dates)

Applies parameter overrides to the definitions based on the specified date ranges.

Parameters:

• definitions (dict) – A dictionary where keys are parameter names and values are 2D arrays representing the parameter values over time and demographics.

• overrides (dict) – A dictionary where keys are parameter names and values are lists of override specifications. Each override specification is a dictionary containing: - ‘start_date’ (str): The start date of the override period in ‘YYYY-MM-DD’ format. - ‘end_date’ (str): The end date of the override period in ‘YYYY-MM-DD’ format. - ‘value’ (np.ndarray or scalar): The value to override within the specified date range.

• dates (list) – A list of datetime.date objects corresponding to the time steps in the definitions arrays.

Returns:

A dictionary with the same keys as definitions, but with values updated according to the overrides.

Return type:

dict

Raises:

ValueError – If the override values do not match the expected shape for the specified date ranges.

epydemix.utils.utils.generate_unique_string(length=12)

Generates a random unique string containing only letters.

Parameters:

length (int, optional) – The length of the generated string. Defaults to 12.

Returns:

A random unique string containing both uppercase and lowercase letters.

Return type:

str

epydemix.utils.utils.compute_days(start_date, end_date)

Computes the number of days between two dates.

Parameters:

• start_date (Union[str, pd.Timestamp]) – The start date in string format (e.g., “YYYY-MM-DD”) or as a pandas Timestamp.

• end_date (Union[str, pd.Timestamp]) – The end date in string format (e.g., “YYYY-MM-DD”) or as a pandas Timestamp.

Returns:

The number of days between the start date and end date, inclusive.

Return type:

int

epydemix.utils.utils.evaluate(expr, env)

Evaluates the expression with the given environment, allowing only whitelisted operations.

This function uses a private copy of the base evaluation model that also whitelists the ‘Mult’ (multiplication) and ‘Pow’ (power) operations, ensuring that only these operations are permitted during the evaluation.

Parameters:

• expr (str) – The expression to evaluate. It is expected to be a string containing the mathematical expression to be evaluated.

• env (dict) – The environment containing variable values. Keys should be variable names and values should be their corresponding numeric values.

Returns:

The result of evaluating the expression. The result type depends on the expression

and its evaluation.

Return type:

any

Raises:

EvalException – If there is an error in evaluating the expression, such as an invalid operation or an undefined variable.

epydemix.utils.utils.compute_simulation_dates(start_date, end_date, steps=None, dt=1.0)

Compute simulation dates including sub-day intervals.

Parameters:

• start_date (str | date | datetime64) – Start date

• end_date (str | date | datetime64) – End date

• steps (int | None) – Number of steps (if None, computed from dt)

• dt (float) – Time step in days (e.g., 1/3 for 8-hour intervals)

Returns:

numpy array of datetime64 objects for each simulation step

Return type:

ndarray

epydemix.utils.utils.apply_initial_conditions(epimodel, initial_conditions_dict)

Applies initial conditions to the compartments of an epidemiological model.

Parameters:

• epimodel (EpiModel) – An instance of an epidemiological model containing compartments and population information.

• **kwargs – Keyword arguments where each key is a compartment name and each value is an array or list specifying the initial population for that compartment.

Returns:

A 2D array where rows correspond to compartments and columns correspond to demographic groups,

representing the initial conditions of the model. The shape of the array is (number_of_compartments, number_of_demographic_groups).

Return type:

np.ndarray

Raises:

ValueError – If initial_conditions_dict contains keys that are not compartments of the model.

epydemix.utils.utils.convert_to_2Darray(lst)

Converts a list into a 2D NumPy array.

Parameters:

lst (List[Any]) – A list of elements to be converted into a 2D array. Elements can be of any type.

Returns:

A 2D NumPy array with shape (1, len(lst)), where len(lst) is the length of the input list.

Return type:

np.ndarray

epydemix.utils.utils.combine_simulation_outputs(simulation_outputs_list)

Combines multiple simulation outputs into a single dictionary by appending new outputs to existing keys.

Parameters:

simulation_outputs_list (List[Dict[str, np.ndarray]]) – A list of dictionaries containing the latest simulation output to be combined. Each dictionary contains keys corresponding to compartment-demographic names and values are arrays of simulation results.

Returns:

A dictionary where keys are compartment-demographic names and values are lists of arrays.

Each list contains simulation results accumulated from multiple runs.

Return type:

Dict[str, List[np.ndarray]]

epydemix.utils.utils.multinomial(n, rates, stay_idx, mask, dt, apply_linear_approximation=False, rng=None, probs_out=None)

Multinomial sample with a ‘stay’ compartment. Probability computation is JIT-compiled via Numba; the draw uses the numpy Generator passed in to preserve reproducibility.

Parameters:

• n (int) – number of trials

• rates (np.ndarray) – array of rates

• stay_idx (int) – index of the stay compartment

• mask (np.ndarray) – boolean array selecting the ‘leave’ destinations

• dt (float) – time step size

• apply_linear_approximation (bool) – whether to apply a linear approximation to the probabilities

• rng (int, np.random.Generator, or None) – seed or random number generator

• probs_out (np.ndarray) – unused, kept for API compatibility

Returns:

array of multinomial samples

Return type:

np.ndarray

epydemix.utils.utils.get_initial_conditions_dict(Nk, perc_dict)

Helper function to get initial conditions dictionary from percentage dictionary.

Parameters:

perc_dict (dict) – Dictionary with percentage of population for each compartment

Returns:

Initial conditions dictionary

Return type:

dict