Model Fitting

Epimodels provides a comprehensive parameter inference framework through the epimodels.fitting module. It allows you to fit model parameters to observed epidemiological data using a variety of loss functions and optimization algorithms.

Quick Start

from epimodels.continuous import SIR
from epimodels.fitting import fit_model, Dataset

model = SIR()
dataset = Dataset()
dataset.add_series("I", times=[0, 1, 2, 3, 5, 7, 10], values=[1, 3, 8, 20, 50, 80, 60])

result = fit_model(
    model, dataset,
    params={"beta": (0.1, 5.0), "gamma": (0.01, 1.0)},
    initial_conditions=[1000, 1, 0],
    time_range=[0, 10],
)
print(result.best_params)  # {'beta': ..., 'gamma': ...}
print(result.loss)         # Final loss value
result.fitted_model.plot_traces()

Dataset Management

The Dataset class manages observed data:

from epimodels.fitting import Dataset

dataset = Dataset()
dataset.add_series("I", times=[0, 1, 2, 3], values=[1, 5, 20, 50])

# From pandas DataFrame
dataset.from_dataframe(df, time_col="date", value_col="cases", state_var="I")

# With uncertainty
dataset.add_series("I", times=[0, 1, 2], values=[1, 5, 20], uncertainty=[0.5, 2, 5])

Loss Functions

Choose an appropriate loss function for your data type:

Loss Function	Best For
SumOfSquaredErrors	General purpose, default choice
WeightedSSE	When some state variables are more important
PoissonLikelihood	Count data (case reports)
NegativeBinomialLikelihood	Overdispersed count data
NormalLikelihood	Continuous data with Gaussian noise
HuberLoss	Robust to outliers
CustomLoss	User-defined objectives
LogLikelihood	Generic log-likelihood (user-supplied function)

Optimizers

Optimizer	Methods	Notes
ScipyOptimizer	L-BFGS-B, BFGS, Nelder-Mead, Powell, differential_evolution	CPU, most methods available
JAXOptimizer	Adam, SGD, RMSprop	GPU-accelerated via JAX
NevergradOptimizer	Derivative-free algorithms	No gradients needed
MultiStartOptimizer	Multi-start with Latin Hypercube / Sobol sampling	Avoids local minima

Advanced Usage

Multi-start optimization:

from epimodels.fitting import ModelFitter, MultiStartOptimizer, ScipyOptimizer

fitter = ModelFitter(
    model=SIR(),
    dataset=dataset,
    params={"beta": (0.1, 5.0), "gamma": (0.01, 1.0)},
)
fitter.set_optimizer(MultiStartOptimizer(
    base_optimizer=ScipyOptimizer(method="L-BFGS-B"),
    n_starts=10,
))
result = fitter.fit()

Fitting Initial Conditions

You can also fit initial conditions alongside parameters using InitialConditionSpec:

from epimodels.fitting import ModelFitter, InitialConditionSpec

fitter = ModelFitter(
    model=SIR(),
    dataset=dataset,
    params={"beta": (0.1, 5.0), "gamma": (0.01, 1.0)},
    initial_conditions=[
        InitialConditionSpec("S", bounds=(500, 1000)),
        InitialConditionSpec("I", bounds=(1, 100)),
    ],
    total_population=1000,
)
result = fitter.fit()
print(result.best_initial_conditions)

Profile likelihood for confidence intervals:

result = fitter.fit()
ci = fitter.profile_likelihood("beta", result.best_params)
print(f"beta 95% CI: {ci.lower_bound:.3f} - {ci.upper_bound:.3f}")

Examples

See the Jupyter notebooks for detailed examples:

• Model_Fitting.ipynb – Basic fitting workflow

• SISLogistic_fitting.ipynb – Fitting SISLogistic model

• SISLogistic_fitting_real_data.ipynb – Fitting with real epidemiological data

• SIRS_parameter_inference.ipynb – SIRS inference with bounded optimization

API Reference