EPI-NODE

EpiNode is a physics-informed neural ODE model for epidemiological forecasting. It decomposes the observed infection series into Trend, Season, and Residual (TSR) components, evolves each via controlled ODEs, fuses the latent states, and uses a parameter network to produce time-varying SIRS (or SIR) rate parameters (β, γ, δ). The model then integrates the SIRS/SIR equations to forecast susceptible (S), infected (I), and recovered (R) compartments.

Overview

• TSR decomposition: Uses Variational Mode Decomposition (VMD) to split the incidence series into trend, seasonal, and residual modes.
• Controlled ODEs: Three ODEFuncTSR modules (one per TSR component) are driven by the decomposed signals; optional time-delay embedding can be applied to the controls.
• ParamNet: MLP that maps the fused latent vector to β, γ, δ within dataset-specific bounds (presets for flu, flu-multi, cls, covid, sirs).
• SIRS physics: Fourth-order Runge–Kutta integration of the compartmental ODEs with learned β, γ, δ (and optional δ=0 for SIR).

Training minimizes MSE (with optional ramp weighting near the train/test split) on the infected fraction I over the training segment; metrics and plots are computed on the forecast segment.

Requirements

• Python 3.x
• PyTorch
• torchdiffeq (ODE integration)
• vmdpy (VMD for TSR)
• NumPy, Pandas, Matplotlib, scipy, scikit-learn

Install dependencies:

pip install torch torchdiffeq vmdpy numpy pandas matplotlib scipy scikit-learn

Optional: torchcde is imported in helper.py (for potential future use); install with pip install torchcde if you see import errors.

Project Structure

EPI-NODE/
├── epi_node.py      # Main model: TSR decomposition, ODEs, ParamNet, SIRS, training & runner
├── data_loader.py   # Synthetic (SIR/SIRS/SEIRS) and real (ILI, CLS) data loaders
├── helper.py        # Metrics, plotting, windowed RMSE, forecast horizon, peak metrics, seeding
├── model_io.py      # Model registry, save/load checkpoints
├── README.md
└── Datasets/
    ├── ILI/         # CDC ILINet data (e.g. CDC_ILINet_HHS_Region.csv) for flu-multi
    └── CLS/         # Delphi COVID-like symptoms CSV (or use Datasets/COVID-19/ as in default path)

Datasets

Name	Description	Loader
sir_fixed	Synthetic SIR, fixed β, γ	load_sir_fixed_data
sirs_fixed	Synthetic SIRS, fixed β, γ, δ	load_sirs_fixed_data
sirs_varying	Synthetic SIRS, time-varying β, γ, δ	load_sirs_varying_data
seirs_fixed	Synthetic SEIRS, fixed parameters	load_seirs_fixed_data
flu-multi	CDC ILINet HHS Region ILI (%)	load_flu_multi_data
cls	Delphi COVID-like symptoms (US)	load_cls_data

• ILI: Place CDC_ILINet_HHS_Region.csv in Datasets/ILI/.
• CLS: Default path in code is Datasets/COVID-19/delphi_COVIDcast_COVID-Like Symptoms in Community_US.csv. If you use Datasets/CLS/, pass the correct file_path when calling load_cls_data.

Usage

Run all experiments (main)

From the project root:

python epi_node.py

This runs EpiNode for multiple seeds and datasets with dataset-specific train/test split fractions. Results are written under:

Results/EpiNode/<data_name>/seed_<seed>/

Single-dataset runner

Use run_epinode() for one dataset and custom splits:

from epi_node import run_tsr_epinode

run_epinode(
    data_name="sirs_fixed",
    t_split_fracs=(0.3),
    seed=21,
    epochs=250,
    lr=2e-3,
    dim_trend=32,
    dim_season=32,
    dim_residual=32,
    results_root="Results/EpiNode",
    use_delay=True,
    tau=1,
    m=3,
)

Main Scripts

• epi_node.py: Time-delay embedding, TSR (VMD), ParamNet, SIRS/SIR dynamics, ODEFuncTSR, train_epinode_model(), run_tsr_epinode(), and if __name__ == "__main__" batch runs.
• data_loader.py: generate_sir_data, generate_sirs_data, generate_seirs_data; loaders for synthetic and real (ILI, CLS) data; returns S, I, R, β, γ, δ, time tensor, split index, and I0.
• helper.py: set_seed, compute_metrics (RMSE, MAE, MAPE, KL), save_plots (loss, I forecast, S/I/R, parameters), rmse_windows, forecast_horizon, peak_metrics, save_training_artifacts.
• model_io.py: MODEL_REGISTRY, register_model, save_model, load_model for EpiNode checkpoints.

Outputs

For each run, the following are saved under the run directory (e.g. Results/EpiNode/cls/seed_21/cls_split_0.50/):

• Plots: loss_curve.png, forecast_I.png, true_vs_forecast.png, parameters_over_time.png
• Metrics: metrics.json (RMSE, MAE, MAPE, KL divergence)
• Advanced: advanced_metrics.txt (windowed RMSE, forecast horizon, peak metrics)
• Checkpoint: best_model.pt (EpiNode components + meta)
• Summary: epinode_summary_<data_name>.csv aggregates metrics across split fractions (when using run_tsr_epinode).

Key Parameters

Parameter	Description
frac	Train fraction of the time series; forecast is the remainder.
dim_trend / dim_season / dim_residual	Latent dimensions for the three TSR ODEs.
tsr_method	Decomposition method; only "vmd" is implemented.
tsr_kwargs	VMD options, e.g. {"K": 5, "alpha": 2000}.
use_delay	If True, apply time-delay embedding to TSR controls.
tau, m	Delay embedding: lag tau, embedding dimension m.
use_sir	If True, use SIR (δ=0); otherwise SIRS.
ode_method	Integration method for latent ODEs (e.g. "rk4").

License

See repository license if applicable.