epimodels package

Subpackages

Submodules

epimodels.ensembles module

epimodels.exceptions module

Custom exceptions for epimodels.

exception epimodels.exceptions.ValidationError[source]

Bases: Exception

Raised when parameter or initial condition validation fails.

epimodels.formulas module

epimodels.interventions module

epimodels.io module

Save/load epidemic models to JSON or YAML.

Serializes the model specification — family, class name, state variables, parameter symbols and the last-run parameter values, and optionally the simulation traces — so models can be stored, shared and reconstructed. Class lookup goes through the model registry (epimodels.registry.get_model()), falling back to the module path recorded at save time for custom models.

Example

>>> from epimodels.continuous import SIR
>>> from epimodels.io import save_model, load_model
>>>
>>> model = SIR()
>>> model([1000, 1, 0], [0, 50], 1001, {"beta": 2, "gamma": 0.1})
>>> save_model(model, "sir_run.json", include_traces=True)
>>> clone = load_model("sir_run.json")
>>> clone.param_values["beta"]
2
epimodels.io.load_model(path: str | Path) → BaseModel[source]

Load a model saved with save_model().

Returns:

A reconstructed model instance.

epimodels.io.model_from_dict(data: dict[str, Any]) → BaseModel[source]

Reconstruct a model from its serialized specification.

The registry is consulted first; if the model class is not registered (e.g. a custom model), the module recorded at save time is imported.

Returns:

A new model instance with parameter values (and traces, if present in the dict) restored.

epimodels.io.model_to_dict(model: BaseModel, include_traces: bool = False) → dict[str, Any][source]

Serialize a model specification to a plain dict.

Parameters:

  • model – The model instance.

  • include_traces – Whether to embed the simulation traces (converted to nested lists).

Returns:

JSON/YAML-compatible dict describing the model.

epimodels.io.save_model(model: BaseModel, path: str | Path, include_traces: bool = False) → Path[source]

Save a model specification to a .json or .yaml/.yml file.

YAML support requires pyyaml (pip install pyyaml).

Returns:

The path the model was written to.

epimodels.network module

epimodels.registry module

Model registry: look up models by name across model families.

The library ships three model families (continuous, discrete, stochastic) whose class names collide (e.g. SIR exists in all three). The registry provides a string-based, collision-aware API:

Example

>>> from epimodels.registry import get_model, list_models
>>> SIR = get_model("SIR", family="continuous")
>>> model = SIR()
>>> list_models(family="discrete")
{'discrete': ['Influenza', 'SIR', ...]}

Custom models can be registered with the register_model() decorator:

>>> from epimodels.registry import register_model
>>> from epimodels.continuous import ContinuousModel
>>> @register_model("MySIR", family="custom")
... class MySIR(ContinuousModel):
...     ...
epimodels.registry.get_model(name: str, family: str = 'any') → type[BaseModel][source]

Look up a model class by name.

Parameters:

  • name – Model class name, e.g. "SIR"

  • family – Family to search: "continuous", "discrete", "stochastic", "custom" or "any" (default). When "any" and the name exists in multiple families, the first match in precedence order (continuous, discrete, stochastic, custom) is returned; pass an explicit family to disambiguate.

Returns:

The model class (not an instance).

Raises:

  • KeyError – If the model is not found (message lists available models).

  • ValueError – If an unknown family is requested.

epimodels.registry.list_models(family: str | None = None) → dict[str, list[str]][source]

List registered model names.

Parameters:

family – Restrict to one family (default: all families).

Returns:

Dict mapping family name to a sorted list of model names.

epimodels.registry.register_model(name: str | None = None, family: str = 'custom') → Callable[[M], M][source]

Class decorator to register a model in the registry.

Parameters:

  • name – Registry name (defaults to the class name)

  • family – Family label (default: “custom”)

Example

>>> @register_model("MySIR", family="custom")
... class MySIR(ContinuousModel):
...     ...
epimodels.registry.unregister_model(name: str, family: str) → None[source]

Remove a model from the registry (mainly useful for tests).

epimodels.rt module

epimodels.sde module

epimodels.solvers module

Module contents

class epimodels.BaseModel[source]

Bases: object

Base class for all models both discrete and continuous

Supports two validation modes: 1. Simple mode (backward compatible): Uses parameters and state_variables dicts 2. Rich mode (new): Uses parameter_specs and variable_specs for enhanced validation

add_constraint(constraint: ModelConstraint) → None[source]

Add a cross-parameter constraint.

Parameters:

constraint – ModelConstraint object

copy(include_traces: bool = False) → BaseModel[source]

Create a copy of the model configuration.

A deep copy is performed so that mutable state (parameter dicts, state variables, specs, formulas, solver) is not shared between the original and the copy.

Parameters:

include_traces – If True, copy simulation results too

Returns:

New model instance with same configuration

define_parameter(spec: ParameterSpec) → None[source]

Register a parameter specification.

Parameters:

spec – ParameterSpec object defining the parameter

define_variable(spec: VariableSpec) → None[source]

Register a state variable specification.

Parameters:

spec – VariableSpec object defining the variable

fit(data: Any, params_to_fit: Any = None, *, times: Any = None, total_population: float | None = None, method: str = 'mle', variable_mapping: dict[str, str] | None = None, likelihood: str = 'normal', sigma: float | dict[str, float] = 1.0, **kwargs) → Any[source]

Fit this model to observed data.

Dispatches to epimodels.fitting.fit_model() (method="mle") or epimodels.fitting.bayes.fit_model_bayesian() (method="bayes").

Parameters:

  • data – For MLE: dict mapping series names to observed values (plus times), or a DataFrame (see fit_model). For Bayes: the same dict+times combination, or a prepared Dataset (in which case params_to_fit must be a list of ParameterSpec).

  • params_to_fit – For MLE: dict mapping parameter names to (lower, upper) bounds. For Bayes with raw dict data: same.

  • times – Observation times (required when data is a dict).

  • total_population – Total population size.

  • method – "mle" (default) or "bayes".

  • variable_mapping – Optional mapping from series names to state variable names (MLE path).

  • likelihood – Observation model for Bayes (“normal”, “poisson”, “negative_binomial”).

  • sigma – Noise level for the Bayes observation model.

  • kwargs – Forwarded to the underlying fitting function.

Returns:

FittingResult (MLE) or BayesianFitResult (Bayes).

Example

>>> result = model.fit(
...     {"I": observed_I}, times=t,
...     params_to_fit={"beta": (0.1, 5.0), "gamma": (0.01, 1.0)},
...     total_population=10000,
... )
>>> result.best_params
model_constraints: list[ModelConstraint]
model_type: str | None
name: str | None
param_values: dict[str, Any]
parameter_specs: dict[str, ParameterSpec]
parameter_table(latex: bool = False) → dict | str[source]
parameters: dict[str, str]
plot_traces(vars: list | None = None) → None[source]

Plots the simulations :param vars: variables to plot

reset() → None[source]

Clear simulation results and parameter values.

simulate(inits: list[float], trange: list[float], totpop: float, params: dict[str, Any], **kwargs) → BaseModel[source]

Run the model and return it, for chaining.

Convenience wrapper around __call__():

>>> model.simulate([1000, 1, 0], [0, 100], 1001,
...                {"beta": 2, "gamma": 0.1}).plot_traces()
Parameters:

kwargs – Forwarded to the model call (e.g. t_eval, solver).

Returns:

self, with traces populated.

state_variables: dict[str, str]
summary() → dict[str, float | int][source]

Return epidemic summary statistics.

Requires simulation results to be available.

Returns:

Dictionary with statistics: - peak_I: Maximum number of infectious individuals - peak_time: Time at which I is maximum - final_S: Final number of susceptible individuals - final_R: Final number of removed/recovered individuals - attack_rate: Proportion of population that was infected - duration: Time until I drops below 1 (if applicable)

Raises:

ValueError – If no simulation has been run

symbolic_model: Any | None
to_dataframe() → pd.DataFrame[source]

Return simulation results as a pandas DataFrame.

Returns:

DataFrame with time and state variable columns

Raises:
to_dict() → dict[str, Any][source]

Return a deep copy of simulation traces.

Returns:

Dictionary with time and state variable arrays

traces: dict[str, Any]
validate_initial_conditions(inits: list[float], totpop: float) → None[source]

Validate initial conditions using either rich specifications or simple validation.

If variable_specs are defined, uses rich validation with: - Bounds checking - Non-negativity - Constraint evaluation

Otherwise falls back to simple validation (backward compatible).

Parameters:

  • inits – List of initial condition values

  • totpop – Total population

Raises:

ValidationError – If validation fails

validate_parameters(params: dict[str, Any]) → None[source]

Validate parameters using either rich specifications or simple validation.

If parameter_specs are defined, uses rich validation with: - Type checking - Domain/bounds validation - Constraint evaluation

Otherwise falls back to simple validation (backward compatible).

Non-numeric parameters (like lists for imported cases) are allowed for advanced model features.

Parameters:

params – Dictionary of parameter values

Raises:

ValidationError – If validation fails

validate_time_range(trange: list[float]) → None[source]

Validate time range input.

Parameters:

trange – Time range [t0, tf]

Raises:

ValidationError – If validation fails

variable_specs: dict[str, VariableSpec]
class epimodels.BetaGammaR0Mixin[source]

Bases: object

Mixin providing a basic reproduction number R0 = beta/gamma.

For models whose R0 has this simple form, inheriting from this mixin avoids duplicating the same property in every model class.

property R0: float | None

Basic reproduction number R0 = beta/gamma.

Returns:

Basic reproduction number, or None if parameters not set

param_values: dict[str, Any]
class epimodels.BetaRR0Mixin[source]

Bases: object

Mixin providing a basic reproduction number R0 = beta/r.

For discrete models whose recovery rate is named r.

property R0: float | None

Basic reproduction number R0 = beta/r.

Returns:

Basic reproduction number, or None if parameters not set

param_values: dict[str, Any]
exception epimodels.FormulaExtractionError(model_name: str, reason: str, suggestion: str = '')[source]

Bases: Exception

Raised when automatic formula extraction fails for a ContinuousModel.

This typically occurs when the model’s _model method uses constructs that cannot be symbolically executed (e.g., loops, conditionals).

model_name

Name of the model that failed extraction

reason

Description of why extraction failed

suggestion

Suggested fix for the user

epimodels.get_model(name: str, family: str = 'any') → type[source]

Look up a model class by name across model families.

Thin wrapper around epimodels.registry.get_model().

Parameters:

  • name – Model class name, e.g. “SIR”

  • family – “continuous”, “discrete”, “stochastic”, “custom” or “any”

epimodels.list_models(family: str | None = None) → dict[source]

List registered model names, grouped by family.

Thin wrapper around epimodels.registry.list_models().