Models

EpiLearn 0.1.0 ships 65 models, from linear regression to pretrained time-series foundation models. They are grouped into three families by the kind of input they consume:

epilearn.models.Temporal – one series at a time (no graph).
epilearn.models.Spatial – one graph snapshot at a time (used by Detection).
epilearn.models.SpatialTemporal – a series and a graph.

One interface for every model

Whatever the family, a model is constructed from the same keywords:

num_features, num_timesteps_input, num_timesteps_output, device

plus num_nodes for the SpatialTemporal family (the Spatial family takes num_features, num_classes, device). Everything model-specific — hidden size, dropout, number of layers, a Hugging Face checkpoint name — is an extra keyword with a default.

Because of that, you rarely construct a model yourself. You hand the class to a task as prototype=, and the task reads the data shape and builds the model for you. Swapping architectures is then a one-line change, even across families:

import torch
from epilearn.data import Dataset
from epilearn.models.Temporal import GRUModel, RidgeModel     # deep net, scikit-learn
from epilearn.models.SpatialTemporal import STGCN             # graph model
from epilearn.tasks.forecast import Forecast
from epilearn.utils import transforms

lookback, horizon = 12, 3
toy = Dataset()
toy.load_toy_dataset()

for prototype, needs_graph in [(GRUModel, False), (RidgeModel, False), (STGCN, True)]:
    torch.manual_seed(42)
    # graph models see the adjacency; purely temporal models do not
    data = Dataset(x=toy.x.clone(), y=toy.y.clone(),
                   graph=toy.graph if needs_graph else None,
                   timestamps=toy.timestamps)
    data.set_transforms(transforms.Compose(
        {"features": [transforms.normalize_feat()],
         "target": [transforms.normalize_target()]}))
    task = Forecast(prototype=prototype, lookback=lookback, horizon=horizon, device='cpu')
    result = task.rolling_train(data, train_size=300, val_size=60, test_size=60,
                                max_folds=1, train_loss='mse', epochs=10, batch_size=64)
    print(f"{prototype.__name__:12s} RMSE={result['aggregate_metrics']['rmse_mean']:.4f} "
          f"coverage={result['aggregate_metrics']['coverage_mean']:.2f}")

Three families, one loop body:

GRUModel     RMSE=0.6171 coverage=0.73
RidgeModel   RMSE=0.5807 coverage=0.72
STGCN        RMSE=1.2556 coverage=0.73

(One fold and 10 epochs: a smoke test of the interface, not a benchmark. For real comparisons use the benchmark, which runs the same protocol over many models with more folds and use_optuna=True.)

The other tasks work the same way: NowcastTask and ScenarioTask take a prototype= too (NowcastTask is single-region, so give it a Temporal model), and Detection takes a Spatial model.

Models that are not drop-in prototypes

A handful of classes predate the shared interface and do not accept the keyword set a task builds. Passing them as prototype= raises TypeError:

Class	Use instead / how to use
`Temporal.Compartmental.SIR` / `SIS` / `SEIR`, `NetworkSIR` / `NetworkSIS` / `NetworkSEIR`	Forward-simulation modules with no `fit()`. For fitting an SIR/SIS/SEIR curve inside a task, use `SIRModel` / `SISModel` / `SEIRModel` from `Temporal.CompartmentalModel`.
`SpatialTemporal.DMP`, `SpatialTemporal.NetSIR`	Mechanistic simulators (`num_nodes`, `horizon`, rate parameters). Call them directly; they have no `fit()`.
`Spatial.GAT`, `Spatial.SAGE`, `Spatial.GIN`	They reject the `num_nodes` keyword that `Detection` passes. Use `Spatial.GCN` as a `prototype=`, or construct GAT/SAGE/GIN yourself and call `.fit()` (see the `Spatial` code example below).

Two further classes are exported but return {"mean": ..., "std": ...} instead of a single tensor, so they do not satisfy the task contract: Temporal.GRU_u.GRUModel (exported as GRU_u_Model) and SpatialTemporal.DSTGCN_u.DSTGCN (exported as DSTGCN_u). They are heteroscedastic variants of GRUModel and DSTGCN that also predict a per-step standard deviation; use them directly, or use rolling_train’s conformal intervals for calibrated uncertainty instead.

Temporal Models

Deep time-series models

Modern sequence architectures. All of them are pure PyTorch and train from scratch.

class epilearn.models.Temporal.GRU.GRUModel(num_features, num_timesteps_input, num_timesteps_output, nhids=256, dropout=0.5, use_norm=False, device='cpu', **kwargs)

Single-layer Gated Recurrent Unit (GRU) Network

Parameters:

num_features (int) – Number of features in the input data.
num_timesteps_input (int) – Number of input timesteps.
num_timesteps_output (int) – Number of output timesteps to predict.
nhid (int, optional) – Number of hidden units in the GRU layer. Default: 256.
dropout (float, optional) – Dropout rate for the GRU layer. Default: 0.5.
use_norm (bool, optional) – Whether to use Layer Normalization after the GRU layer. Default: False.

Returns:

A tensor of shape (batch_size, num_timesteps_output) representing the predicted values for the future timesteps. Each element corresponds to a predicted value for a future timestep.

Return type:

torch.Tensor

forward(x, **kwargs)

Parameters:: x (torch.Tensor) – The input tensor for the model. Expected shape is (batch_size, num_timesteps_input, num_features), where batch_size is the number of samples in the batch, num_timesteps_input is the number of input timesteps, and num_features is the number of features for each timestep.
Returns:: The output of the model, a tensor of shape (batch_size, num_timesteps_output), representing the predicted values for the future timesteps. Each element corresponds to a predicted value for a future timestep.
Return type:: torch.Tensor

class epilearn.models.Temporal.LSTM.LSTMModel(num_features, num_timesteps_input, num_timesteps_output, nhid=256, dropout=0.5, use_norm=False, **kwargs)

Long Short-Term Memory (LSTM) Model

Parameters:

num_features (int) – Number of features in each timestep of the input data.
num_timesteps_input (int) – Number of timesteps considered for each input sample.
num_timesteps_output (int) – Number of output timesteps to predict.
nhid (int, optional) – Number of hidden units in the LSTM layers. Default: 256.
dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.5.
use_norm (bool, optional) – Whether to apply layer normalization after the LSTM layers. Default: False.

Returns:

A tensor of shape (batch_size, num_timesteps_output) representing the predicted values for the future timesteps. This tensor is the output from the last timestep processed through a linear layer to predict the desired number of future timesteps.

Return type:

torch.Tensor

forward(x, **kwargs)

Parameters:: x (torch.Tensor) – The input tensor for the model. Expected shape is (batch_size, num_timesteps_input, num_features), where batch_size is the number of samples in the batch, num_timesteps_input is the number of input timesteps, and num_features is the number of features for each timestep.
Returns:: The output of the model, a tensor of shape (batch_size, num_timesteps_output), representing the predicted values for the future timesteps. Each element corresponds to a predicted value for a future timestep.
Return type:: torch.Tensor

class epilearn.models.Temporal.CNN.CNNModel(num_features, num_timesteps_input, num_timesteps_output, conv1_hid=16, conv2_hid=32, kernel_size=3, linear_hid=128, dropout=0.5, device='cpu', **kwargs)

Convolutional Neural Network for Time Series Forecasting

Parameters:

num_features (int) – Number of features in the input data.
num_timesteps_input (int) – Number of input timesteps.
num_timesteps_output (int) – Number of output timesteps to predict.
conv1_hid (int, optional) – Number of filters in first convolutional layer. Default: 16.
conv2_hid (int, optional) – Number of filters in second convolutional layer. Default: 32.
kernel_size (int, optional) – Kernel size for convolutional layers. Default: 3.
linear_hid (int, optional) – Number of hidden units in the linear layer. Default: 128.
dropout (float, optional) – Dropout rate. Default: 0.5.

Returns:

A tensor of shape (batch_size, num_timesteps_output) representing the predicted values for the future timesteps. Each element corresponds to a predicted value for a future timestep.

Return type:

torch.Tensor

forward(x, **kwargs)

Parameters:: x (torch.Tensor) – The input tensor for the model. Expected shape is (batch_size, num_timesteps_input, num_features), where batch_size is the number of samples in the batch, num_timesteps_input is the number of input timesteps, and num_features is the number of features for each timestep.
Returns:: The output of the model, a tensor of shape (batch_size, num_timesteps_output), representing the predicted values for the future timesteps. Each element corresponds to a predicted value for a future timestep.
Return type:: torch.Tensor

class epilearn.models.Temporal.MLP.MLPModel(num_features, num_timesteps_input, num_timesteps_output, hidden_dims=[256, 128], dropout=0.2, activation='relu', use_batch_norm=False, device='cpu', **kwargs)

Multi-Layer Perceptron (MLP) Model for Time Series Forecasting

A simple feedforward neural network that flattens the input time series and processes it through multiple fully connected layers.

Parameters:

num_features (int) – Number of features in each timestep of the input data.
num_timesteps_input (int) – Number of timesteps considered for each input sample.
num_timesteps_output (int) – Number of output timesteps to predict.
hidden_dims (list of int, optional) – List of hidden layer dimensions. Default: [256, 128].
dropout (float, optional) – Dropout rate for regularization. Default: 0.2.
activation (str, optional) – Activation function (‘relu’, ‘gelu’, ‘tanh’). Default: ‘relu’.
use_batch_norm (bool, optional) – Whether to use batch normalization after each hidden layer. Default: False.

Returns:

A tensor of shape (batch_size, num_timesteps_output) representing the predicted values for the future timesteps.

Return type:

torch.Tensor

Examples

>>> model = MLPModel(num_features=3, num_timesteps_input=24, num_timesteps_output=12)
>>> x = torch.randn(32, 24, 3)  # batch_size=32
>>> output = model(x)
>>> output.shape
torch.Size([32, 12])

forward(x, **kwargs)

Parameters:: x (torch.Tensor) – The input tensor for the model. Expected shape is (batch_size, num_timesteps_input, num_features), where batch_size is the number of samples in the batch, num_timesteps_input is the number of input timesteps, and num_features is the number of features for each timestep.
Returns:: The output of the model, a tensor of shape (batch_size, num_timesteps_output), representing the predicted values for the future timesteps.
Return type:: torch.Tensor

initialize(): Initialize model parameters using Xavier/Glorot initialization for weights and zeros for biases.

class epilearn.models.Temporal.Dlinear.DlinearModel(num_features, num_timesteps_input, num_timesteps_output, moving_avg_window=25, nhid=None, device='cpu', **kwargs)

Dynamic Linear Model

Parameters:

num_features (int) – Number of features in each timestep of the input data.
num_timesteps_input (int) – Number of timesteps considered for each input sample.
num_timesteps_output (int) – Number of output timesteps to predict.
moving_avg_window (int, optional) – Kernel size for the moving average decomposition. Default: 25.
nhid (int, optional) – Hidden dimension for intermediate transformation. If None, uses direct transformation. Default: None.

decomposition

Method to decompose the time series data into seasonal and trend components.

Type:: function

Linear_Transform

Linear transformation layer to project the decomposed input data to the output space.

Type:: torch.nn.Linear

Returns:: A tensor of shape (batch_size, num_timesteps_output) representing the predicted values for the future timesteps. Each element corresponds to a predicted value for a specific future timestep. The output is averaged across the feature dimension to reduce it to a single predictive value per timestep.
Return type:: torch.Tensor

forward(x, **kwargs)

Parameters:: x (torch.Tensor) – The input tensor representing time series data for each batch. Expected shape is (batch_size, num_timesteps_input, num_features), where batch_size is the number of samples in the batch, num_timesteps_input is the number of input time steps, and num_features represents features at each time step.
Returns:: The output of the model, a tensor of shape (batch_size, num_timesteps_output) that represents the predicted values for the future time steps, reduced to a single predictive value per time step by averaging across the feature dimension.
Return type:: torch.Tensor

class epilearn.models.Temporal.PatchTST.PatchTSTModel(num_features, num_timesteps_input, num_timesteps_output, patch_len=4, stride=1, d_model=128, n_heads=8, num_layers=1, dim_feedforward=16, dropout=0.1, activation='gelu', norm_first=True, device='cpu', **kwargs)

PatchTST: A Time Series Forecasting Model using Patching and Transformers

This implementation follows the PatchTST architecture with channel independence, where each feature/channel is processed independently through the transformer.

Parameters:

num_features (int) – Number of features in each timestep of the input data.
num_timesteps_input (int) – Number of timesteps considered for each input sample.
num_timesteps_output (int) – Number of output timesteps to predict.
patch_len (int, optional) – Length of each patch. Default: 4.
stride (int, optional) – Stride for patch extraction. Default: 1.
d_model (int, optional) – Dimension of the model (embedding dimension). Default: 128.
n_heads (int, optional) – Number of attention heads in transformer. Default: 8.
num_layers (int, optional) – Number of transformer encoder layers. Default: 1.
dim_feedforward (int, optional) – Dimension of feedforward network in transformer. Default: 256.
dropout (float, optional) – Dropout rate. Default: 0.1.
activation (str, optional) – Activation function (‘relu’ or ‘gelu’). Default: ‘gelu’.
norm_first (bool, optional) – If True, layer norm is applied before attention/feedforward. Default: True.

Returns:

A tensor of shape (batch_size, num_timesteps_output) representing the predicted values for the future timesteps.

Return type:

torch.Tensor

forward(x, **kwargs)

Parameters:: x (torch.Tensor) – The input tensor for the model. Expected shape is (batch_size, num_timesteps_input, num_features), where batch_size is the number of samples in the batch, num_timesteps_input is the number of input timesteps, and num_features is the number of features for each timestep.
Returns:: The output of the model, a tensor of shape (batch_size, num_timesteps_output), representing the predicted values for the future timesteps.
Return type:: torch.Tensor

initialize(): Initialize model parameters

class epilearn.models.Temporal.iTransformer.iTransformerModel(num_features, num_timesteps_input, num_timesteps_output, num_nodes=1, d_model=64, n_heads=4, e_layers=2, d_ff=None, dropout=0.1, activation='relu', device='cpu', **kwargs)

iTransformer: Inverted Transformer for Time Series Forecasting

This model inverts the traditional transformer approach by treating channels (features/variables) as tokens instead of time steps. This allows the model to capture correlations between different channels effectively.

For spatiotemporal data with shape (batch, time, nodes, features), the model automatically flattens nodes and features into a single channel dimension.

Parameters:

num_features (int) – Number of features per node (for temporal) or total channels (for spatiotemporal).
num_timesteps_input (int) – Number of input timesteps (lookback window).
num_timesteps_output (int) – Number of output timesteps to predict (horizon).
num_nodes (int, optional) – Number of nodes for spatiotemporal data. Default: 1 (temporal only).
d_model (int, optional) – Dimension of the model embeddings. Default: 64.
n_heads (int, optional) – Number of attention heads. Default: 4.
e_layers (int, optional) – Number of encoder layers. Default: 2.
d_ff (int, optional) – Dimension of feedforward network. Default: 4 * d_model (256 for default d_model=64).
dropout (float, optional) – Dropout rate. Default: 0.1.
activation (str, optional) – Activation function (‘relu’ or ‘gelu’). Default: ‘relu’.
device (str, optional) – Device to run model on. Default: ‘cpu’.
Shapes (Input)
------------
Temporal ((batch, time, features))
Spatiotemporal ((batch, time, nodes, features) -> automatically flattened to (batch, time, nodes*features))
Shape (Output)
------------
(batch
column) (horizon) - returns prediction for target variable (last)

forward(x, **kwargs)

Forward pass of iTransformer (aligned with tslib implementation).

Parameters:: x (torch.Tensor) – Input tensor. Can be: - Temporal: (batch, time, features) - Spatiotemporal: (batch, time, nodes, features)
Returns:: Predictions: - 3D input: (batch, horizon) - 4D input: (batch, nodes, horizon)
Return type:: torch.Tensor

initialize(): Re-initialize model weights using Xavier uniform.

predict(feature, graph=None, states=None, dynamic_graph=None)

Make predictions.

Parameters:: feature (torch.Tensor) – Input features.
Returns:: Predictions of shape (batch, horizon)
Return type:: torch.Tensor

class epilearn.models.Temporal.TSMixer.TSMixerModel(num_features: int, num_timesteps_input: int, num_timesteps_output: int, num_nodes: int | None = None, d_model: int = 64, e_layers: int = 2, dropout: float = 0.1, device: str = 'cpu', **kwargs)

TSMixer: Channel-dependent time series forecasting model.

This model treats each channel as a separate entity and applies both temporal mixing (across time) and channel mixing (across features) using simple MLP blocks.

For spatiotemporal data: - Input (batch, time, nodes, features) is flattened to (batch, time, nodes*features) - The model treats nodes*features as the total number of channels - Output is (batch, horizon) averaged across all channels

Parameters:

num_features – Number of input features per node
num_timesteps_input – Length of input sequence
num_timesteps_output – Length of output sequence (horizon)
num_nodes – Number of spatial nodes (optional, for spatiotemporal)
d_model – Hidden dimension for mixing layers
e_layers – Number of residual blocks
dropout – Dropout rate
device – Device to run the model on

forward(x, A_q=None, A_h=None, **kwargs)

Forward pass.

Parameters:

x – Input tensor - 3D: (batch, time, features) - 4D: (batch, time, nodes, features)
A_q – Adjacency matrices (not used, for compatibility)
A_h – Adjacency matrices (not used, for compatibility)

Returns:

Output tensor

3D input: (batch, horizon)
4D input: (batch, nodes, horizon)

Return type:

out

initialize(): Re-initialize model weights.

predict(feature, graph=None, states=None, dynamic_graph=None): Make predictions.

class epilearn.models.Temporal.FreTS.FreTSModel(num_features: int, num_timesteps_input: int, num_timesteps_output: int, num_nodes: int | None = None, embed_size: int = 128, hidden_size: int = 256, channel_independence: bool = False, sparsity_threshold: float = 0.01, device: str = 'cpu', **kwargs)

FreTS: Frequency-domain MLP model for time series forecasting.

This model operates primarily in the frequency domain: 1. Embeds each time point with learnable embeddings 2. Applies frequency-domain MLPs on both temporal and channel dimensions 3. Uses FFT for temporal/channel transformations

For spatiotemporal data: - Input (batch, time, nodes, features) is flattened to (batch, time, nodes*features) - The model treats nodes*features as the total number of channels - Output is (batch, horizon) averaged across all channels

Parameters:

num_features – Number of input features per node
num_timesteps_input – Length of input sequence
num_timesteps_output – Length of output sequence (horizon)
num_nodes – Number of spatial nodes (optional, for spatiotemporal)
embed_size – Embedding dimension
hidden_size – Hidden dimension for output projection
channel_independence – If True, skip channel mixing (faster but less expressive)
sparsity_threshold – Threshold for soft shrinkage
device – Device to run the model on

FreMLP(B, nd, dimension, x, r, i, rb, ib)

Frequency-domain MLP.

Applies complex-valued linear transformation in frequency domain.

Parameters:

B – Batch size
nd – Number of elements in non-FFT dimension
dimension – Dimension along which FFT was taken
x – Complex tensor after FFT
r – Real and imaginary weight matrices
i – Real and imaginary weight matrices
rb – Real and imaginary biases
ib – Real and imaginary biases

Returns:

Complex tensor after MLP

Return type:

y

MLP_channel(x, B, N, L)

Frequency channel learner.

Applies FFT along channel dimension, MLP in frequency domain, then IFFT back.

MLP_temporal(x, B, N, L)

Frequency temporal learner.

Applies FFT along time dimension, MLP in frequency domain, then IFFT back.

forward(x, A_q=None, A_h=None, **kwargs)

Forward pass.

Parameters:

x – Input tensor - 3D: (batch, time, features) - 4D: (batch, time, nodes, features)
A_q – Adjacency matrices (not used, for compatibility)
A_h – Adjacency matrices (not used, for compatibility)

Returns:

Output tensor

3D input: (batch, horizon)
4D input: (batch, nodes, horizon)

Return type:

out

initialize(): Re-initialize model weights.

predict(feature, graph=None, states=None, dynamic_graph=None): Make predictions.

tokenEmb(x)

Embed each value with learnable embeddings.

Parameters:: x – [B, L, N] input tensor
Returns:: [B, N, L, D] where D is embed_size
Return type:: embedded

Epidemic-specific deep models

Deep models whose architecture encodes epidemiological structure. EINNModel (Rodriguez et al., AAAI 2023) is a GRU encoder/decoder over SEIR compartments with a learnable SEIR ODE as a soft regularizer; EpiDeepModel (Adhikari et al., KDD 2019) clusters historical seasons with a dual autoencoder and attends over them; CALINetModel (Kamarthi et al., AAAI 2022) distils a pretrained EpiDeepModel into a lightweight target module.

class epilearn.models.Temporal.EINN.EINNModel(device='cpu', lookback=16, horizon=4, n_features=1, num_features=None, num_timesteps_input=None, num_timesteps_output=None, hidden_dim=40, n_states=5, n_layers=1, ode_lambda=0.1, lr=0.001, **kwargs)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

predict(feature, **kwargs)

Returns:: If the model returns a tensor, returns the tensor. If the model returns a dictionary (e.g., with uncertainty estimates), returns the dictionary with all tensors moved to CPU.
Return type:: torch.FloatTensor or dict

class epilearn.models.Temporal.EpiDeep.EpiDeepModel(device='cpu', lookback=16, horizon=4, n_features=1, num_features=None, num_timesteps_input=None, num_timesteps_output=None, hidden_dim=20, n_clusters=4, seq_len=5, encode_layers=None, mapping_layers=None, pretrain_epochs=200, alpha_cluster=0.1, alpha_pred=10.0, lr=0.001, **kwargs)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

predict(feature, **kwargs)

Returns:: If the model returns a tensor, returns the tensor. If the model returns a dictionary (e.g., with uncertainty estimates), returns the dictionary with all tensors moved to CPU.
Return type:: torch.FloatTensor or dict

class epilearn.models.Temporal.CALINet.CALINetModel(device='cpu', lookback=16, horizon=4, n_features=1, num_features=None, num_timesteps_input=None, num_timesteps_output=None, hidden_dim=20, n_clusters=4, calib_dim=16, seq_len=5, pretrain_epochs=200, alpha_kd=0.1, alpha_recon=0.1, kd_warmup=10, finetune_source=True, lr=0.001, **kwargs)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

predict(feature, **kwargs)

Returns:: If the model returns a tensor, returns the tensor. If the model returns a dictionary (e.g., with uncertainty estimates), returns the dictionary with all tensors moved to CPU.
Return type:: torch.FloatTensor or dict

Foundation models

Pretrained time-series foundation models, used zero-shot by default. Each class is a thin wrapper that tokenises the lookback window and calls the pretrained backend. The Base / Large / Small classes subclass the plain one and only pin a different checkpoint (e.g. MoiraiBaseModel is MoiraiModel with model_name='Salesforce/moirai-1.0-R-base'); ChronosBoltModel is the encoder-only, faster Chronos variant.

Note

The backends are optional dependencies:

pip install epilearn[chronos]     # also [moirai], [moment], [timesfm]

They are imported lazily, inside the wrapper’s _load_model(). Constructing the model therefore always works; the first forward() raises ImportError with the exact pip command if the backend is missing, e.g. chronos-forecasting is required for ChronosModel. Install it with: pip install chronos-forecasting. The benchmark checks this up front (epilearn.benchmark.check_foundation_deps) and reports such models as skipped rather than failing the run.

Backend per family: Chronos* needs chronos-forecasting, Moirai* needs uni2ts, Moment* needs momentfm, TimesFM needs timesfm.

class epilearn.models.Temporal.Chronos.ChronosModel(num_features, num_timesteps_input, num_timesteps_output, model_name='amazon/chronos-t5-small', num_samples=20, temperature=1.0, top_k=50, top_p=1.0, device='cpu', **kwargs)

Chronos Time Series Foundation Model wrapper for EpiLearn.

This model uses Amazon’s pretrained Chronos model for zero-shot or fine-tuned time series forecasting. By default, it uses the pretrained model without fine-tuning (zero-shot inference).

Parameters:

num_features (int) – Number of features in the input data.
num_timesteps_input (int) – Number of input timesteps (lookback window).
num_timesteps_output (int) – Number of output timesteps to predict (horizon).
model_name (str, optional) –
Chronos model variant to use. Default: ‘amazon/chronos-t5-small’. Options: ‘amazon/chronos-t5-tiny’, ‘amazon/chronos-t5-mini’,

’amazon/chronos-t5-small’, ‘amazon/chronos-t5-base’, ‘amazon/chronos-t5-large’
num_samples (int, optional) – Number of sample paths to generate for probabilistic forecasts. Default: 20.
temperature (float, optional) – Temperature for sampling. Lower values give more deterministic outputs. Default: 1.0.
top_k (int, optional) – Top-k sampling parameter. Default: 50.
top_p (float, optional) – Top-p (nucleus) sampling parameter. Default: 1.0.
device (str, optional) – Device to run the model on. Default: ‘cpu’.

Returns:

Predicted values of shape (batch_size, num_timesteps_output).

Return type:

torch.Tensor

fit(train_input, train_target, train_states=None, train_graph=None, train_dynamic_graph=None, val_input=None, val_target=None, val_states=None, val_graph=None, val_dynamic_graph=None, loss='mse', epochs=1, batch_size=10, lr=0.001, weight_decay=0, initialize=True, verbose=False, patience=10, **kwargs)

Fit method for Chronos (zero-shot, no training required).

Chronos is a pretrained model and performs zero-shot forecasting. This method simply ensures the model is loaded.

forward(x, **kwargs)

Forward pass using Chronos for prediction.

Parameters:: x (torch.Tensor) – Input tensor of shape (batch_size, num_timesteps_input, num_features). Only the first feature (target column) is used for prediction.
Returns:: Predicted values of shape (batch_size, num_timesteps_output).
Return type:: torch.Tensor

initialize(): Initialize the model (loads pretrained weights).

predict(feature, graph=None, states=None, dynamic_graph=None)

Make predictions using the Chronos model.

Parameters:: feature (torch.Tensor) – Input features of shape (batch_size, num_timesteps_input, num_features).
Returns:: Predictions of shape (batch_size, num_timesteps_output).
Return type:: torch.Tensor

class epilearn.models.Temporal.Chronos.ChronosBoltModel(num_features, num_timesteps_input, num_timesteps_output, model_name='amazon/chronos-bolt-small', device='cpu', **kwargs)

Chronos-Bolt: Faster variant of Chronos using encoder-only architecture.

Chronos-Bolt models are more efficient variants that provide significant speedup while maintaining competitive forecasting accuracy.

Parameters:

num_features (int) – Number of features in the input data.
num_timesteps_input (int) – Number of input timesteps (lookback window).
num_timesteps_output (int) – Number of output timesteps to predict (horizon).
model_name (str, optional) –
Chronos-Bolt model variant. Default: ‘amazon/chronos-bolt-small’. Options: ‘amazon/chronos-bolt-tiny’, ‘amazon/chronos-bolt-mini’,

’amazon/chronos-bolt-small’, ‘amazon/chronos-bolt-base’
device (str, optional) – Device to run the model on. Default: ‘cpu’.

fit(train_input, train_target, **kwargs): Zero-shot inference - just load the model.

forward(x, **kwargs): Forward pass using Chronos-Bolt.

initialize(): Initialize the model.

predict(feature, graph=None, states=None, dynamic_graph=None): Make predictions.

class epilearn.models.Temporal.Moirai.MoiraiModel(num_features, num_timesteps_input, num_timesteps_output, model_name='Salesforce/moirai-1.0-R-small', num_samples=100, patch_size='auto', device='cpu', **kwargs)

MOIRAI Time Series Foundation Model wrapper for EpiLearn.

This model uses Salesforce’s pretrained MOIRAI model for zero-shot time series forecasting. MOIRAI supports multivariate time series and variable prediction horizons.

Parameters:

num_features (int) – Number of features in the input data.
num_timesteps_input (int) – Number of input timesteps (lookback window).
num_timesteps_output (int) – Number of output timesteps to predict (horizon).
model_name (str, optional) –
MOIRAI model variant to use. Default: ‘Salesforce/moirai-1.0-R-small’. Options: ‘Salesforce/moirai-1.0-R-small’, ‘Salesforce/moirai-1.0-R-base’,

’Salesforce/moirai-1.0-R-large’
num_samples (int, optional) – Number of sample paths for probabilistic forecasts. Default: 100.
patch_size (str or int, optional) – Patch size for the model. Default: ‘auto’.
device (str, optional) – Device to run the model on. Default: ‘cpu’.

Returns:

Predicted values of shape (batch_size, num_timesteps_output).

Return type:

torch.Tensor

fit(train_input, train_target, train_states=None, train_graph=None, train_dynamic_graph=None, val_input=None, val_target=None, val_states=None, val_graph=None, val_dynamic_graph=None, loss='mse', epochs=1, batch_size=10, lr=0.001, weight_decay=0, initialize=True, verbose=False, patience=10, **kwargs)

Fit method for MOIRAI (zero-shot, no training required).

MOIRAI is a pretrained model and performs zero-shot forecasting.

forward(x, **kwargs)

Forward pass using MOIRAI for prediction.

Parameters:: x (torch.Tensor) – Input tensor of shape (batch_size, num_timesteps_input, num_features).
Returns:: Predicted values of shape (batch_size, num_timesteps_output).
Return type:: torch.Tensor

initialize(): Initialize the model (loads pretrained weights).

predict(feature, graph=None, states=None, dynamic_graph=None)

Make predictions using the MOIRAI model.

Parameters:: feature (torch.Tensor) – Input features of shape (batch_size, num_timesteps_input, num_features).
Returns:: Predictions of shape (batch_size, num_timesteps_output).
Return type:: torch.Tensor

class epilearn.models.Temporal.Moirai.MoiraiBaseModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **kwargs): MOIRAI Base model variant (91M parameters).

class epilearn.models.Temporal.Moirai.MoiraiLargeModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **kwargs): MOIRAI Large model variant (311M parameters).

class epilearn.models.Temporal.Moment.MomentModel(num_features, num_timesteps_input, num_timesteps_output, model_name='AutonLab/MOMENT-1-large', device='cpu', **kwargs)

MOMENT Time Series Foundation Model wrapper for EpiLearn.

Uses the pre-trained reconstruction head with short_forecast for zero-shot time series forecasting (no fine-tuning needed).

Parameters:

num_features (int) – Number of features in the input data.
num_timesteps_input (int) – Number of input timesteps (lookback window).
num_timesteps_output (int) – Number of output timesteps to predict (horizon).
model_name (str, optional) – HuggingFace model id. Default: 'AutonLab/MOMENT-1-large'.
device (str, optional) – Device to run the model on. Default: 'cpu'.

Returns:

Predicted values of shape (batch_size, num_timesteps_output).

Return type:

torch.Tensor

fit(train_input, train_target, train_states=None, train_graph=None, train_dynamic_graph=None, val_input=None, val_target=None, val_states=None, val_graph=None, val_dynamic_graph=None, loss='mse', epochs=1, batch_size=10, lr=0.001, weight_decay=0, initialize=True, verbose=False, patience=10, **kwargs): Zero-shot — just ensure the model is loaded.

forward(x, **kwargs)

Forward pass — zero-shot forecasting via MOMENT’s short_forecast.

Parameters:: x (torch.Tensor) – (batch, num_timesteps_input, num_features)
Returns:: (batch, num_timesteps_output)
Return type:: torch.Tensor

initialize(): Initialize the model (loads pretrained weights).

predict(feature, graph=None, states=None, dynamic_graph=None): Make predictions using the MOMENT model.

class epilearn.models.Temporal.Moment.MomentSmallModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **kwargs): MOMENT Small model variant.

class epilearn.models.Temporal.Moment.MomentBaseModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **kwargs): MOMENT Base model variant.

class epilearn.models.Temporal.TimesFM.TimesFMModel(num_features, num_timesteps_input, num_timesteps_output, model_name='google/timesfm-2.0-500m-pytorch', freq=1, device='cpu', **kwargs)

TimesFM (Google) wrapper for EpiLearn — zero-shot time series forecasting.

Parameters:

num_features (int) – Number of features in the input data.
num_timesteps_input (int) – Number of input timesteps (lookback window).
num_timesteps_output (int) – Number of output timesteps to predict (horizon).
model_name (str, optional) – HuggingFace checkpoint id. Default: 'google/timesfm-2.0-500m-pytorch'.
freq (int, optional) – Frequency indicator: 0 = high (≤ daily), 1 = medium (weekly / monthly), 2 = low (quarterly +). Default: 1 (weekly epidemiological data).
device (str, optional) – 'cpu' or 'cuda'. Default: 'cpu'.

Returns:

Point forecasts of shape (batch_size, num_timesteps_output).

Return type:

torch.Tensor

fit(train_input, train_target, train_states=None, train_graph=None, train_dynamic_graph=None, val_input=None, val_target=None, val_states=None, val_graph=None, val_dynamic_graph=None, loss='mse', epochs=1, batch_size=10, lr=0.001, weight_decay=0, initialize=True, verbose=False, patience=10, **kwargs): Zero-shot — just ensure the model is loaded.

forward(x, **kwargs)

Forward pass — zero-shot forecasting with TimesFM.

Parameters:: x (torch.Tensor) – (batch, num_timesteps_input, num_features)
Returns:: (batch, num_timesteps_output)
Return type:: torch.Tensor

initialize(): Initialize the model (loads pretrained weights).

predict(feature, graph=None, states=None, dynamic_graph=None): Make predictions using TimesFM.

Statistical models and baselines

Classical statistics and the reference baselines every serious comparison needs. SeasonalNaiveModel repeats the value from the same point in the previous cycle – hard to beat on strongly seasonal data.

RKINowcastModel and NobBSModel are nowcasting baselines and expect a reporting triangle, so use them with NowcastTask; inside Forecast they raise NotImplementedError: Subclasses must implement _forecast_single_series. RKINowcastModel (An der Heiden & Hamouda, 2020) learns a completion CDF \(F(d)\) during fit and divides each partial count by it; NobBSModel (McGough et al., 2020) estimates the delay distribution per sample and smooths the log-incidence curve, so it needs no training data at all.

class epilearn.models.Temporal.StatsModel.ARIMAModel(num_features, num_timesteps_input, num_timesteps_output, order=(1, 0), seasonal_order=(0, 0, 0, 0), trend='c', device='cpu', p=None, d=None, q=None, **kwargs)

Autoregressive Integrated Moving Average (ARIMA) Model. Wrapper around statsmodels ARIMA for univariate time series.

Parameters:

num_features (int) – Number of features (only first feature is used for univariate ARIMA)
num_timesteps_input (int) – Number of input timesteps
num_timesteps_output (int) – Number of output timesteps to predict
order (tuple) – (p, d, q) order of the ARIMA model. If only (p, d) provided, q defaults to 0
seasonal_order (tuple, optional) – (P, D, Q, s) seasonal order of the ARIMA model. Defaults to (0, 0, 0, 0) for no seasonality.
trend (str) – Trend parameter (‘c’ for constant, ‘n’ for none)

class epilearn.models.Temporal.StatsModel.VARMAXModel(num_features, num_timesteps_input, num_timesteps_output, order=(1, 0), trend='c', device='cpu', p=None, q=None, rki_correct=False, **kwargs)

Vector Autoregression Moving-Average with eXogenous variables (VARMAX) Model. Wrapper around statsmodels VARMAX for multivariate time series.

Parameters:

num_features (int) – Number of features in each timestep
num_timesteps_input (int) – Number of input timesteps
num_timesteps_output (int) – Number of output timesteps to predict
order (tuple) – (p, q) order of the VARMAX model
trend (str) – Trend parameter (‘c’ for constant, ‘n’ for none)

class epilearn.models.Temporal.StatsModel.SeasonalNaiveModel(num_features, num_timesteps_input, num_timesteps_output, season_length=None, device='cpu', **kwargs)

Seasonal Naive Forecasting Model.

A simple baseline that forecasts by repeating the values from the same season in the previous cycle. For example, with weekly data and annual seasonality (52 weeks), the forecast for week t is the actual value from week t-52.

This is a strong baseline for seasonal data and often outperforms more complex methods when seasonality is the dominant pattern.

Parameters:

num_features (int) – Number of features (only last feature/target is used)
num_timesteps_input (int) – Number of input timesteps (lookback window)
num_timesteps_output (int) – Number of output timesteps to predict (forecast horizon)
season_length (int) – Length of the seasonal cycle (e.g., 52 for weekly data with annual seasonality, 12 for monthly data with annual seasonality, 7 for daily data with weekly seasonality). Default: 52 (annual seasonality for weekly epidemic data).
device (str) – Device for compatibility (‘cpu’ or ‘cuda’)

Examples

For weekly COVID-19 data with annual patterns: >>> model = SeasonalNaiveModel(num_features=4, num_timesteps_input=53, … num_timesteps_output=4, season_length=52) >>> # Forecast for week t uses actual value from week t-52

Notes

If lookback < season_length, falls back to naive (last value) forecast
If season_length is not provided, uses lookback as season length (assumes 1 cycle in window)
Very fast to compute (no model fitting required)
Provides a strong baseline for seasonal data

class epilearn.models.Temporal.StatsModel.RKINowcastModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **kwargs)

RKI-style delay-adjusted nowcasting (An der Heiden & Hamouda, 2020).

During fit(): learns the completion CDF F(d) = E[count_at_delay_d / final_count] from training data where final counts (targets) are known.

During predict(): for each target day, divides the most-mature partial observation by F(d) to estimate the final count.

Parameters:

num_features (int) – Standard model dimensions (set automatically by the framework).
num_timesteps_input (int) – Standard model dimensions (set automatically by the framework).
num_timesteps_output (int) – Standard model dimensions (set automatically by the framework).

fit(train_input, train_target, **kwargs): Estimate completion CDF from training (features, targets) pairs.

class epilearn.models.Temporal.StatsModel.NobBSModel(num_features, num_timesteps_input, num_timesteps_output, smoothing=0.3, device='cpu', **kwargs)

Simplified NobBS nowcasting (McGough et al., PLOS Comp Bio, 2020).

Estimates the delay distribution per sample from the observed triangle, then corrects partial counts and applies random-walk smoothing on the log-incidence curve. No training data required — purely generative.

Parameters:

num_features (int) – Standard model dimensions.
num_timesteps_input (int) – Standard model dimensions.
num_timesteps_output (int) – Standard model dimensions.
smoothing (float) – Exponential-smoothing weight α ∈ (0, 1]. Lower → more smoothing (stronger random-walk prior). Default 0.3.

fit(**kwargs): No-op — NobBS estimates everything per-sample during predict().

Note

epilearn.models.Temporal.ARIMA still imports VARMAXModel and ARIMAModel for backwards compatibility, but the module moved to epilearn.models.Temporal.StatsModel in 0.1.0. Prefer the new path.

scikit-learn regressors

Each wrapper flattens the lookback window into a feature vector and fits one scikit-learn regressor per horizon step. They need no GPU and are the cheapest sanity check available.

class epilearn.models.Temporal.ScikitModel.LinearRegressionModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **model_params)

Linear Regression Model for time series forecasting. Wrapper around sklearn.linear_model.LinearRegression.

Parameters:

num_features (int) – Number of features in each timestep
num_timesteps_input (int) – Number of input timesteps (lookback window)
num_timesteps_output (int) – Number of output timesteps to predict (forecast horizon)
fit_intercept (bool, optional) – Whether to calculate the intercept (default: True)

class epilearn.models.Temporal.ScikitModel.RidgeModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **model_params)

Ridge Regression Model with L2 regularization. Wrapper around sklearn.linear_model.Ridge.

Parameters:: alpha (float, optional) – Regularization strength (default: 1.0)

class epilearn.models.Temporal.ScikitModel.LassoModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **model_params)

Lasso Regression Model with L1 regularization. Wrapper around sklearn.linear_model.Lasso.

Parameters:: alpha (float, optional) – Regularization strength (default: 1.0)

class epilearn.models.Temporal.ScikitModel.ElasticNetModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **model_params)

ElasticNet Regression Model with L1 and L2 regularization. Wrapper around sklearn.linear_model.ElasticNet.

Parameters:

alpha (float, optional) – Regularization strength (default: 1.0)
l1_ratio (float, optional) – Mix ratio between L1 and L2 (default: 0.5)

class epilearn.models.Temporal.ScikitModel.RandomForestModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **model_params)

Random Forest Regressor for time series forecasting. Wrapper around sklearn.ensemble.RandomForestRegressor.

Parameters:

n_estimators (int, optional) – Number of trees (default: 100)
max_depth (int, optional) – Maximum depth of trees (default: None)
min_samples_split (int, optional) – Minimum samples required to split (default: 2)

class epilearn.models.Temporal.ScikitModel.GradientBoostingModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **model_params)

Gradient Boosting Regressor for time series forecasting. Wrapper around sklearn.ensemble.GradientBoostingRegressor.

Parameters:

n_estimators (int, optional) – Number of boosting stages (default: 100)
learning_rate (float, optional) – Learning rate (default: 0.1)
max_depth (int, optional) – Maximum depth of trees (default: 3)

class epilearn.models.Temporal.ScikitModel.SVRModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **model_params)

Support Vector Regressor for time series forecasting. Wrapper around sklearn.svm.SVR.

Parameters:

kernel (str, optional) – Kernel type (‘linear’, ‘poly’, ‘rbf’, ‘sigmoid’) (default: ‘rbf’)
C (float, optional) – Regularization parameter (default: 1.0)
epsilon (float, optional) – Epsilon in epsilon-SVR (default: 0.1)

class epilearn.models.Temporal.ScikitModel.KNNModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **model_params)

K-Nearest Neighbors Regressor for time series forecasting. Wrapper around sklearn.neighbors.KNeighborsRegressor.

Parameters:

n_neighbors (int, optional) – Number of neighbors (default: 5)
weights (str, optional) – Weight function (‘uniform’, ‘distance’) (default: ‘uniform’)

class epilearn.models.Temporal.ScikitModel.DecisionTreeModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **model_params)

Decision Tree Regressor for time series forecasting. Wrapper around sklearn.tree.DecisionTreeRegressor.

Parameters:

max_depth (int, optional) – Maximum depth of tree (default: None)
min_samples_split (int, optional) – Minimum samples required to split (default: 2)

Compartmental models

These follow the same per-sample fitting interface as the statistical models: fit() is a no-op, and for every sample predict() integrates the ODE with scipy.integrate.odeint, fits the rate parameters to that lookback window with L-BFGS-B, and simulates horizon steps forward. They are ordinary prototype= models.

class epilearn.models.Temporal.CompartmentalModel.SIRModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **kwargs): Susceptible → Infected → Recovered.

class epilearn.models.Temporal.CompartmentalModel.SISModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **kwargs): Susceptible → Infected → Susceptible (no permanent immunity).

class epilearn.models.Temporal.CompartmentalModel.SEIRModel(num_features, num_timesteps_input, num_timesteps_output, device='cpu', **kwargs): Susceptible → Exposed → Infected → Recovered.

For richer mechanistic simulation (SEIRS, vaccination, waning immunity, time-varying parameter schedules), see epilearn.utils.compartmental_models.

Compartmental simulation modules

epilearn.models.Temporal.Compartmental holds the original nn.Module compartmental blocks. They hold their rates in nn.Linear weights and roll an initial compartment vector forward one step at a time. They have no fit() and do not accept num_features etc., so they cannot be passed to a task – call them directly:

import torch
from epilearn.models.Temporal.Compartmental import SIR

model = SIR(horizon=5, infection_rate=0.3, recovery_rate=0.1, population=1000)
trajectory = model(torch.tensor([990., 10., 0.]))   # initial S, I, R
print(trajectory.shape)                             # torch.Size([5, 3])

class epilearn.models.Temporal.Compartmental.SIR(horizon=None, infection_rate=0.01, recovery_rate=0.038, population=None)

Susceptible-Infected-Recovered (SIR) Model

Parameters:

horizon (int, optional) – Number of future time steps to simulate. If None, a single step is simulated unless overridden in the forward method.
infection_rate (float, optional) – Initial infection rate parameter, representing the rate at which susceptible individuals become infected. Default: 0.01.
recovery_rate (float, optional) – Initial recovery rate parameter, representing the rate at which infected individuals recover. Default: 0.038.
population (int, optional) – Total population considered in the model. If None, the sum of the initial conditions (susceptible, infected, recovered) is used as the total population.

beta

Linear layer with no bias to model the infection rate dynamically.

Type:: torch.nn.Linear

gamma

Linear layer with no bias to model the recovery rate dynamically.

Type:: torch.nn.Linear

Returns:: A tensor of shape (horizon, 3), representing the predicted number of susceptible, infected, and recovered individuals at each timestep. Each row corresponds to a timestep, with the columns representing susceptible, infected, and recovered counts respectively.
Return type:: torch.Tensor

forward(x, steps=1)

Parameters:

x (torch.Tensor) – The initial condition tensor for the model. Expected shape is (3,), where the elements represent the number of susceptible (S), infected (I), and recovered (R) individuals respectively.
steps (int, optional) – Number of future time steps to simulate. If horizon is specified during initialization and not None, it overrides this parameter. Default is 1 if horizon is None.

Returns:

A tensor of shape (steps, 3), representing the predicted number of susceptible, infected, and recovered individuals at each timestep. Each row corresponds to a timestep, with the columns representing susceptible, infected, and recovered counts respectively.

Return type:

torch.Tensor

class epilearn.models.Temporal.Compartmental.SIS(horizon=None, infection_rate=None, recovery_rate=None, population=None)

Susceptible-Infected-Susceptible (SIS) Model

Parameters:

horizon (int, optional) – Number of future time steps to simulate. If None, a single step is simulated unless overridden in the forward method.
infection_rate (float, optional) – Infection rate parameter, representing the rate at which susceptible individuals become infected. If None, must be initialized separately.
recovery_rate (float, optional) – Recovery rate parameter, representing the rate at which infected individuals recover and return to the susceptible state. If None, must be initialized separately.
population (int, optional) – Total population considered in the model. If None, the sum of the initial conditions (susceptible and infected) is used as the total population.

beta

Linear layer with no bias to model the infection rate dynamically.

Type:: torch.nn.Linear

gamma

Linear layer with no bias to model the recovery rate dynamically.

Type:: torch.nn.Linear

Returns:: A tensor of shape (horizon, 2), representing the predicted number of susceptible and infected individuals at each timestep. Each row corresponds to a timestep, with the columns representing the susceptible and infected counts respectively.
Return type:: torch.Tensor

forward(x, steps=1, **kwargs)

Parameters:

x (torch.Tensor) – The initial condition tensor for the model, representing the initial numbers of susceptible (S) and infected (I) individuals. Expected shape is (2,), where x[0] is the number of susceptible and x[1] is the number of infected individuals at the start.
steps (int, optional) – Number of future time steps to simulate. If horizon is specified during initialization and not None, it overrides this parameter. Default is 1 if horizon is None.

Returns:

A tensor of shape (steps, 2), representing the predicted number of susceptible and infected individuals at each timestep. Each row corresponds to a timestep, with the first column representing susceptible and the second column representing infected counts.

Return type:

torch.Tensor

class epilearn.models.Temporal.Compartmental.SEIR(horizon=None, infection_rate=None, recovery_rate=None, cure_rate=None, latency=None, population=None)

Susceptible-Exposed-Infected-Recovered (SEIR) Model

Parameters:

horizon (int, optional) – Number of future time steps to simulate. If None, a single step is simulated unless overridden in the forward method.
infection_rate (float, optional) – Infection rate parameter, representing the rate at which susceptible individuals become exposed. If None, must be initialized separately.
recovery_rate (float, optional) – Recovery rate parameter, representing the rate at which infected individuals recover. If None, must be initialized separately.
cure_rate (float, optional) – Natural immunity rate parameter, representing the rate at which individuals (across S, E, I, R compartments) return to susceptible due to loss of immunity. If None, must be initialized separately.
latency (float, optional) – Latency rate parameter, representing the rate at which exposed individuals become infected. If None, must be initialized separately.
population (int, optional) – Total population considered in the model. If None, the sum of the initial conditions (susceptible, exposed, infected, recovered) is used as the total population.

beta

Linear layer with no bias to dynamically model the infection rate.

Type:: torch.nn.Linear

gamma

Linear layer with no bias to dynamically model the recovery rate.

Type:: torch.nn.Linear

mu

Linear layer with no bias to model the natural immunity rate.

Type:: torch.nn.Linear

a

Linear layer with no bias to model the latency rate.

Type:: torch.nn.Linear

Returns:: A tensor of shape (horizon, 4), representing the predicted number of susceptible, exposed, infected, and recovered individuals at each timestep. Each row corresponds to a timestep, with the columns representing the counts of susceptible, exposed, infected, and recovered individuals respectively.
Return type:: torch.Tensor

forward(x, steps=1)

Parameters:

x (torch.Tensor) – The initial condition tensor for the model, representing the initial numbers of susceptible (S), exposed (E), infected (I), and recovered (R) individuals. Expected shape is (4,), where elements correspond to S, E, I, and R counts.
steps (int, optional) – Number of future time steps to simulate. If horizon is specified during initialization and not None, it overrides this parameter. Default is 1 if horizon is None.

Returns:

A tensor of shape (steps, 4), representing the predicted number of susceptible, exposed, infected, and recovered individuals at each timestep. Each row corresponds to a timestep, with columns representing the counts of susceptible, exposed, infected, and recovered individuals respectively.

Return type:

torch.Tensor

class epilearn.models.Temporal.Compartmental.NetworkSIR(num_nodes, horizon=None, infection_rate=0.01, recovery_rate=0.038, population=None)

Network-based SIR (Susceptible-Infected-Recovered) Model with spatial coupling.

Each node represents a location/region, and the adjacency matrix defines how infections spread between connected nodes.

Parameters:

num_nodes (int) – Number of nodes in the network.
horizon (int, optional) – Number of future time steps to simulate.
infection_rate (float, optional) – Initial infection rate parameter. Default: 0.01.
recovery_rate (float, optional) – Initial recovery rate parameter. Default: 0.038.
population (int, optional) – Total population. If None, computed from initial conditions.

forward(x, adj, steps=1)

Parameters:

x (torch.Tensor) – Initial state tensor with shape (num_nodes, 3) representing S, I, R for each node.
adj (torch.Tensor) – Adjacency matrix with shape (num_nodes, num_nodes).
steps (int, optional) – Number of future time steps to simulate.

Returns:

Tensor of shape (steps, num_nodes, 3) representing S, I, R over time.

Return type:

torch.Tensor

class epilearn.models.Temporal.Compartmental.NetworkSIS(num_nodes, horizon=None, infection_rate=0.01, recovery_rate=0.038, population=None)

Network-based SIS (Susceptible-Infected-Susceptible) Model with spatial coupling.

Parameters:

num_nodes (int) – Number of nodes in the network.
horizon (int, optional) – Number of future time steps to simulate.
infection_rate (float, optional) – Initial infection rate parameter. Default: 0.01.
recovery_rate (float, optional) – Initial recovery rate parameter. Default: 0.038.
population (int, optional) – Total population. If None, computed from initial conditions.

forward(x, adj, steps=1)

Parameters:

x (torch.Tensor) – Initial state tensor with shape (num_nodes, 2) representing S, I for each node.
adj (torch.Tensor) – Adjacency matrix with shape (num_nodes, num_nodes).
steps (int, optional) – Number of future time steps to simulate.

Returns:

Tensor of shape (steps, num_nodes, 2) representing S, I over time.

Return type:

torch.Tensor

class epilearn.models.Temporal.Compartmental.NetworkSEIR(num_nodes, horizon=None, infection_rate=0.01, recovery_rate=0.038, cure_rate=0.01, latency=0.1, population=None)

Network-based SEIR (Susceptible-Exposed-Infected-Recovered) Model with spatial coupling.

Parameters:

num_nodes (int) – Number of nodes in the network.
horizon (int, optional) – Number of future time steps to simulate.
infection_rate (float, optional) – Initial infection rate parameter. Default: 0.01.
recovery_rate (float, optional) – Initial recovery rate parameter. Default: 0.038.
cure_rate (float, optional) – Natural immunity loss rate. Default: 0.01.
latency (float, optional) – Rate of progression from exposed to infected. Default: 0.1.
population (int, optional) – Total population. If None, computed from initial conditions.

forward(x, adj, steps=1)

Parameters:

x (torch.Tensor) – Initial state tensor with shape (num_nodes, 4) representing S, E, I, R for each node.
adj (torch.Tensor) – Adjacency matrix with shape (num_nodes, num_nodes).
steps (int, optional) – Number of future time steps to simulate.

Returns:

Tensor of shape (steps, num_nodes, 4) representing S, E, I, R over time.

Return type:

torch.Tensor

Note

epilearn.models.Temporal.SIR still imports SIR, SIS and SEIR for backwards compatibility, but the module moved to epilearn.models.Temporal.Compartmental in 0.1.0. Prefer the new path.

Code Example

Every model can also be driven directly, without a task, through .fit():

import torch
from epilearn.models.Temporal.GRU import GRUModel

num_features = 1
lookback = 16 # inputs size
horizon = 3 # predicts size

features = torch.round(torch.rand((10, lookback, num_features)))
node_target = torch.round(torch.rand((10, horizon)))

model=GRUModel(num_features=num_features, num_timesteps_input=lookback, num_timesteps_output=horizon, device='cpu')
model.fit(
        train_input=features,
        train_target=node_target,
        val_input=None,
        val_target=None,
        val_graph=None,
        epochs=20,
        loss='mse'
        )

Spatial Models

Graph neural networks that classify nodes from a single snapshot; the Detection task uses them for outbreak detection. They take num_features, num_classes and device instead of a lookback/horizon pair.

class epilearn.models.Spatial.GCN.GCN(num_features, hidden_dim=16, num_classes=2, nlayers=2, dropout=0.5, with_bn=False, with_bias=True, device='cpu', **kwargs)

Graph Convolutional Network (GCN)

Parameters:

num_features (int) – Number of input features per node.
hidden_dim (int, optional) – Dimension of hidden layers. Default: 16.
num_classes (int, optional) – Number of output classes for each node. Default: 2.
nlayers (int, optional) – Number of layers in the GCN. Default: 2.
dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.5.
with_bn (bool, optional) – Specifies whether batch normalization should be included. Default: False.
with_bias (bool, optional) – Specifies whether to include bias parameters in the GCN layers. Default: True.
device (str) – The device (cpu or gpu) on which the model will be run.

Returns:

A tensor of shape (batch_size, num_nodes, num_classes), representing the predicted outcomes for each node after passing through the GCN.

Return type:

torch.Tensor

forward(X, edge_index, edge_weight=None)

Parameters: X : torch.Tensor

The input features tensor with shape (batch_size, num_nodes, num_features).

edge_indextorch.Tensor: The edge indices in COO format with shape (2, num_edges).

Returns: torch.Tensor

The output predictions for each node with shape (batch_size * num_nodes, num_classes).

class epilearn.models.Spatial.GAT.GAT(num_features, hidden_dim, num_classes, nlayers=2, nheads=[2, 2], dropout=0.5, with_bn=False, with_bias=True, device=None, concat=False)

Graph Attention Network (GAT)

Parameters:

num_features (int) – Number of input features per node.
hidden_dim (int) – Dimension of hidden layers.
num_classes (int) – Number of output features per node.
nlayers (int, optional) – Number of layers in the GAT. Default: 2.
nheads (list of int) – Number of attention heads in each GAT layer. Length must match nlayers.
dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.5.
with_bn (bool, optional) – Specifies whether batch normalization should be included. Default: False.
with_bias (bool, optional) – Specifies whether to include bias parameters in the attention calculations. Default: True.
device (torch.device) – The device (cpu or gpu) on which the model will be run.
concat (bool, optional) – Specifies whether to concatenate the outputs of the attention heads instead of averaging them. Default: False.

Returns:

A tensor of shape (batch_size, num_nodes, output_dim), representing the predicted values for each node over future timesteps.

Return type:

torch.Tensor

forward(X, edge_index, edge_weight=None)

Parameters:

X (torch.Tensor) – Input features tensor with shape (num_nodes, num_features).
edge_index (torch.Tensor) – Tensor defining the edges of the graph with shape (2, num_edges), where each column represents an edge as a pair of node indices.
edge_weight (torch.Tensor, optional) – Edge weights with shape (num_edges,). Default is None.

Returns:

Output tensor of shape (num_nodes, num_classes), representing the predicted values for each node.

Return type:

torch.Tensor

class epilearn.models.Spatial.GIN.GIN(num_features, num_classes, hidden_dim=16, nlayers=2, dropout=0.5, with_bias=True, device=None)

Graph Isomorphism Network (GIN)

Parameters:

num_features (int) – Number of input features per node.
hidden_dim (int) – Dimension of hidden layers.
num_classes (int) – Number of output features per node.
nlayers (int, optional) – Number of layers in the GIN. Default: 2.
dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.5.
with_bias (bool, optional) – Specifies whether to include bias parameters in the MLP layers. Default: True.
device (str) – The device (cpu or gpu) on which the model will be run. Must be specified.

Returns:

A tensor of shape (batch_size, num_nodes, output_dim), representing the predicted outcomes for each node after passing through the GIN.

Return type:

torch.Tensor

forward(X, edge_index, edge_weight=None)

Parameters: X : torch.Tensor

Node feature matrix with shape (batch_size, num_nodes, num_features) or (num_nodes, num_features).

edge_indextorch.Tensor: Edge index in COO format with shape (2, num_edges).

Returns: torch.Tensor

Output from the network with shape (batch_size * num_nodes, num_classes) or (num_nodes, num_classes).

class epilearn.models.Spatial.SAGE.SAGE(num_features, hidden_dim, num_classes, nlayers=2, dropout=0.5, with_bn=False, with_bias=True, device=None, aggr='mean')

Graph Sample and Aggregate (SAGE)

Parameters:

num_features (int) – Number of input features per node.
hidden_dim (int) – Dimension of hidden layers.
num_classes (int) – Number of output features per node.
nlayers (int, optional) – Number of layers in the GraphSAGE model. Default: 2.
dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.5.
with_bn (bool, optional) – Specifies whether batch normalization should be included. Default: False.
with_bias (bool, optional) – Specifies whether to include bias parameters in the GraphSAGE layers. Default: True.
device (str) – The device (cpu or gpu) on which the model will be run. Must be specified.
aggr (str or callable, optional) – The aggregation function to use (‘mean’, ‘sum’, ‘max’, etc.), or a callable that returns a custom aggregation function. Default: ‘mean’.

Returns:

A tensor of shape (batch_size, num_nodes, output_dim), representing the predicted outcomes for each node after passing through the GraphSAGE model.

Return type:

torch.Tensor

forward(X, edge_index, edge_weight=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Note

Only GCN accepts the extra num_nodes keyword that Detection passes when it builds the model. GAT, SAGE and GIN currently have to be constructed and fitted directly, as below.

Code Example

import torch
from epilearn.models.Spatial import GCN

num_features = 4
num_classes = 2
lookback = 1 # inputs size
horizon = 2 # predicts size

graph = torch.round(torch.rand((47,47)))
features = torch.round(torch.rand((10,47,1,4)))
node_target = torch.round(torch.rand((10,47)))

model=GCN(num_features=num_features, num_classes=horizon, device='cpu')
model.fit(
        train_input=features,
        train_target=node_target,
        train_graph=graph,
        val_input=None,
        val_target=None,
        val_graph=None,
        epochs=20,
        loss='ce'
        )

Spatial-Temporal Models

Models that consume a window of features and a graph, and predict a horizon for every node jointly. All of them take num_nodes, num_features, num_timesteps_input, num_timesteps_output, device, and their forward accepts (X, adj, states, dynamic_adj) so a model can use node states and a time-varying graph if the dataset provides them.

Spatio-temporal graph networks

General-purpose architectures that combine temporal convolution or recurrence with graph propagation: STGCN (spatio-temporal graph convolution), DSTGCN (STGCN plus a GraphLearningLayer that learns the adjacency instead of taking it as given), DCRNN (diffusion convolutional recurrent network), GraphWaveNet (dilated temporal convolution with an optional adaptive adjacency) and ATMGNN (attention-based temporal multiresolution GNN).

class epilearn.models.SpatialTemporal.STGCN.STGCN(num_nodes, num_features, num_timesteps_input, num_timesteps_output, nhids=128, device='cpu', **kwargs)

Spatio-temporal graph convolutional network as described in https://arxiv.org/abs/1709.04875v3 by Yu et al. Input should have shape (batch_size, num_nodes, num_input_time_steps, num_features).

forward(X, adj, states=None, dynamic_adj=None, **kargs)

Parameters:

X (torch.Tensor) – Input from task: Shape (batch_size, num_timesteps_input, num_nodes, num_features)
adj (torch.Tensor) – Shape (num_nodes, num_nodes)

Returns:

Output shape (batch_size, num_nodes, num_timesteps_output)

Return type:

torch.Tensor

class epilearn.models.SpatialTemporal.DSTGCN.DSTGCN(num_nodes, num_features, num_timesteps_input, num_timesteps_output, nhids=128, gat_heads=1, **kwargs)

Spatio-temporal graph convolutional network as described in https://arxiv.org/abs/1709.04875v3 by Yu et al. Input should have shape (batch_size, num_nodes, num_input_time_steps, num_features).

forward(X, adj=None, states=None, dynamic_adj=None, debug=False, y=None, **kargs)

Parameters:

X (torch.Tensor) – Shape (batch_size, num_nodes, num_timesteps_input, num_features)
adj (torch.Tensor) – Shape (num_nodes, num_nodes)

Returns:

Output shape (batch_size, num_timesteps_output, num_nodes)

Return type:

torch.Tensor

class epilearn.models.SpatialTemporal.DCRNN.DCRNN(num_features=1, num_timesteps_input=5, num_classes=1, num_timesteps_output=1, max_diffusion_step=2, filter_type='laplacian', num_rnn_layers=1, rnn_units=1, nonlinearity='tanh', dropout=0, device='cpu', num_nodes=None, **kwargs)

Diffusion Convolutional Recurrent Neural Network (DCRNN)

Parameters:

num_features (int) – Number of input features per node.
num_timesteps_input (int) – Number of past time steps used as input by the network.
num_classes (int) – Number of output features per node.
num_timesteps_output (int) – Number of future time steps to predict.
max_diffusion_step (int) – Maximum number of diffusion steps in the graph convolution operations. Default: 2.
filter_type (str) – Type of filter used in graph convolutions, e.g., ‘laplacian’. Default: “laplacian”.
num_rnn_layers (int) – Number of recurrent neural network layers. Default: 1.
rnn_units (int) – Number of units per recurrent layer. Default: 1.
nonlinearity (str) – Type of nonlinearity function used in RNN. Default: “tanh”.
dropout (float) – Dropout rate applied in the network to prevent overfitting. Default: 0.
device (str, optional) – The device (cpu or gpu) on which the model will be run. Default: ‘cpu’.
num_nodes (int, optional) – Number of nodes in the graph. This parameter is accepted for API compatibility but not used in initialization. Default: None.

Returns:

A tensor of shape (batch_size, num_nodes, horizon), representing the predicted values for each node over future timesteps.

Return type:

torch.Tensor

decoder(encoder_hidden_state, adj_mx, num_nodes): Decoder forward pass :param encoder_hidden_state: (num_layers, batch_size, self.hidden_state_size) :param labels: (self.num_timesteps_output, batch_size, self.num_nodes * self.num_classes) [optional, not exist for inference] :param batches_seen: global step [optional, not exist for inference] :return: output: (self.num_timesteps_output, batch_size, self.num_nodes * self.num_classes)

encoder(inputs, adj_mx, num_nodes): encoder forward pass on t time steps :param inputs: shape (num_timesteps_input, batch_size, num_sensor * num_features) :return: encoder_hidden_state: (num_layers, batch_size, self.hidden_state_size)

forward(X_batch, graph, X_states, batch_graph)

Parameters:

X_batch (torch.Tensor) – Input tensor with shape (batch_size, num_nodes, num_timesteps_input, num_features), representing the input features over multiple timesteps for each node.
graph (torch.Tensor) – Static adjacency matrix with shape (num_nodes, num_nodes), representing the fixed connections between nodes.
X_states (torch.Tensor, optional) – States of the nodes if available, with the same shape as X_batch. Used for models that incorporate node states over time. Default: None.
batch_graph (torch.Tensor, optional) – Dynamic adjacency matrix if available, with shape similar to graph but possibly varying over time. Used for models that account for changing graph structures. Default: None.

Returns:

The output tensor of shape (batch_size, num_nodes, num_timesteps_output), representing the predicted values for each node over the specified output timesteps.

Return type:

torch.Tensor

class epilearn.models.SpatialTemporal.GraphWaveNet.GraphWaveNet(device='cpu', dropout=0.3, gcn_bool=True, addaptadj=True, aptinit=None, num_timesteps_input=2, num_timesteps_output=12, residual_channels=32, dilation_channels=32, skip_channels=256, end_channels=512, kernel_size=2, blocks=4, nlayers=2, adj_m=None, num_nodes=None, num_features=1, **kwargs)

Graph Convolutional Wave Network (GraphWaveNet)

Parameters:

device (str) – The device (cpu or gpu) on which the model will be run.
dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.3.
gcn_bool (bool, optional) – Indicates whether to include graph convolution layers. Default: True.
addaptadj (bool, optional) – Indicates whether to include an adaptive adjacency matrix. Default: True.
aptinit (tensor, optional) – Initial tensor for adaptive adjacency matrix. Default: None.
num_timesteps_input (int) – Number of input timesteps per node.
num_timesteps_output (int) – Number of output timesteps per node.
residual_channels (int) – Number of channels in residual layers. Default: 32.
dilation_channels (int) – Number of channels in dilation layers. Default: 32.
skip_channels (int) – Number of channels in skip connection layers. Default: 256.
end_channels (int) – Number of channels in the final convolution layers. Default: 512.
kernel_size (int) – Kernel size for the convolution layers. Default: 2.
blocks (int) – Number of blocks in the WaveNet architecture. Default: 4.
nlayers (int) – Number of layers in each block. Default: 2.
adj_m (tensor) – Initial adjacency matrix if static graph structure is used. Default: None.

Returns:

A tensor of shape (batch_size, num_nodes, output_dim), representing the predicted values for each node over future timesteps. Each slice along the second dimension corresponds to a timestep, with each column representing a node.

Return type:

torch.Tensor

forward(X_batch, graph, X_states, batch_graph)

Parameters:

X_batch (torch.Tensor) – Input features tensor with shape (batch_size, num_timesteps_input, num_nodes, num_features).
adj (torch.Tensor) – Static adjacency matrix of the graph with shape (num_nodes, num_nodes).
states (torch.Tensor, optional) – States of the nodes if available, with the same shape as x. Default: None.
dynamic_adj (torch.Tensor, optional) – Dynamic adjacency matrix if available, with shape similar to adj but possibly varying over time. Default: None.

Returns:

The output tensor of shape (batch_size, num_timesteps_output, num_nodes), representing the predicted values for each node over the specified output timesteps.

Return type:

torch.Tensor

class epilearn.models.SpatialTemporal.ATMGNN.ATMGNN(num_nodes, num_features, num_timesteps_input, num_timesteps_output, nhid=256, dropout=0.5, nhead=1, num_clusters=[10, 5], use_norm=False, device='cpu', **kwargs)

Attention-based Temporal Multiresolution Graph Neural Network (ATMGNN)

Parameters:

num_nodes (int) – Number of nodes in the graph.
num_features (int) – Number of features per node per timestep.
num_timesteps_input (int) – Number of timesteps considered for each input sample (window size).
num_timesteps_output (int) – Number of output timesteps to predict.
nhid (int, optional) – Number of hidden units in the network and the output size of graph convolution layers. Default: 256.
dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.5.
nhead (int, optional) – Number of heads in the multi-head attention mechanism. Default: 1.
num_clusters (list, optional) – List of integers defining the number of clusters for each multiresolution layer. Default: [10, 5].
use_norm (bool, optional) – Whether to use normalization on outputs from each layer. Default: False.

Returns:

A tensor of shape (batch_size, num_timesteps_output, num_nodes), representing the predicted values for each node over future timesteps. Each slice along the second dimension corresponds to a timestep, with each column representing a node.

Return type:

torch.Tensor

forward(x, adj, states=None, dynamic_adj=None, **kargs)

Parameters:

x (torch.Tensor) – Input features tensor with shape (batch_size, num_timestamps, num_nodes, num_features).
adj (torch.Tensor) – Static adjacency matrix of the graph with shape (num_nodes, num_nodes).
states (torch.Tensor, optional) – States of the nodes if available, with the same shape as x. Default: None.
dynamic_adj (torch.Tensor, optional) – Dynamic adjacency matrix if available, with shape similar to adj but possibly varying over time. Default: None.

Returns:

The output tensor of shape (batch_size, num_timesteps_output, num_nodes), representing the predicted values for each node over the specified output timesteps.

Return type:

torch.Tensor

Epidemic-specific graph models

Architectures published for multi-region epidemic forecasting: ColaGNN (cross-location attention graph network), its epidemiological variant EpiColaGNN, EpiGNN (epidemiological graph network) and CNNRNN_Res (convolution + recurrence with residual connections).

class epilearn.models.SpatialTemporal.ColaGNN.ColaGNN(num_nodes, num_features, num_timesteps_input, num_timesteps_output, nhid=32, n_channels=1, n_spatial=None, rnn_model='GRU', n_layer=1, bidirect=False, dropout=0.5, device='cpu', **kwargs)

Convolutional-Layer Graph Neural Network (ColaGNN)

Parameters:

num_nodes (int) – Number of nodes in the graph.
num_features (int) – Number of features per node per timestep.
num_timesteps_input (int) – Number of timesteps considered for each input sample.
num_timesteps_output (int) – Number of output timesteps to predict.
nhid (int, optional) – Number of hidden units in the RNN and GNN layers. Default: 32.
n_channels (int, optional) – Number of channels for the convolution layers. Default: 1.
n_spatial (int, optional) – Number of spatial features from graph convolutions. Default: max(10, nhid // 2).
rnn_model (str, optional) – Type of RNN model to use (‘LSTM’, ‘GRU’, ‘RNN’). Default: ‘GRU’.
n_layer (int, optional) – Number of layers in the RNN model. Default: 1.
bidirect (bool, optional) – Whether the RNN layers are bidirectional. Default: False.
dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.5.
device (str, optional) – The device (cpu or gpu) on which the model will be run. Default: ‘cpu’.

Returns:

A tensor of shape (batch_size, num_timesteps_output, num_nodes), representing the predicted values for each node over future timesteps. Each slice along the second dimension corresponds to a timestep, with each column representing a node.

Return type:

torch.Tensor

forward(x, adj, states=None, dynamic_adj=None)

Parameters:

x (torch.Tensor) – Input features tensor with shape (batch_size, num_timesteps_input, num_nodes, num_features).
adj (torch.Tensor) – Static adjacency matrix of the graph with shape (num_nodes, num_nodes).
states (torch.Tensor, optional) – States of the nodes if available, with the same shape as x. Default: None.
dynamic_adj (torch.Tensor, optional) – Dynamic adjacency matrix if available, with shape similar to adj but possibly varying over time. Default: None.

Returns:

The output tensor of shape (batch_size, num_timesteps_output, num_nodes), representing the predicted values for each node over the specified output timesteps.

Return type:

torch.Tensor

class epilearn.models.SpatialTemporal.EpiColaGNN.EpiColaGNN(num_nodes, num_features, num_timesteps_input, num_timesteps_output, nhid=32, rnn_model='GRU', n_layer=1, bidirect=False, target_idx=0, dropout=0.5, device='cpu')

Epidemiological Convolutional-Layer Graph Neural Network (EpiColaGNN)

Parameters:

num_nodes (int) – Number of nodes in the graph.
num_features (int) – Number of features per node per timestep.
num_timesteps_input (int) – Number of timesteps considered for each input sample.
num_timesteps_output (int) – Number of output timesteps to predict.
nhid (int, optional) – Number of hidden units in the RNN and GNN layers. Default: 32.
rnn_model (str, optional) – Type of RNN model to use (‘LSTM’, ‘GRU’, ‘RNN’). Default: ‘GRU’.
n_layer (int, optional) – Number of layers in the RNN model. Default: 1.
bidirect (bool, optional) – Whether the RNN layers are bidirectional. Default: False.
target_idx (int, optional) – Index of the target variable in the feature set. Default: 0.
dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.5.
device (str, optional) – The device (cpu or gpu) on which the model will be run. Default: ‘cpu’.

Returns:

A tensor of shape (batch_size, num_timesteps_output, num_nodes), representing the predicted values for each node over future timesteps. Each slice along the second dimension corresponds to a timestep, with each column representing a node.

Return type:

torch.Tensor

forward(x, adj, states=None, dynamic_adj=None)

Parameters:

x (torch.Tensor) – Input features tensor with shape (batch_size, num_timesteps_input, num_nodes, num_features).
adj (torch.Tensor) – Static adjacency matrix of the graph with shape (num_nodes, num_nodes).
states (torch.Tensor, optional) – States of the nodes if available, with the same shape as x. Default: None.
dynamic_adj (torch.Tensor, optional) – Dynamic adjacency matrix if available, with shape similar to adj but possibly varying over time. Default: None.

Returns:

The output tensor of shape (batch_size, num_timesteps_output, num_nodes), representing the predicted values for each node over the specified output timesteps.

Return type:

torch.Tensor

class epilearn.models.SpatialTemporal.EpiGNN.EpiGNN(num_nodes, num_features, num_timesteps_input, num_timesteps_output, k=8, hidA=128, hidR=32, hidP=1, n_layer=2, dropout=0, nhids=None, device='cpu', **kwargs)

Epidemiological Graph Neural Network (EpiGNN)

Parameters:

num_nodes (int) – Number of nodes in the graph.
num_features (int) – Number of features per node per timestep.
num_timesteps_input (int) – Number of timesteps considered for each input sample.
num_timesteps_output (int) – Number of output timesteps to predict.
k (int, optional) – Number of local neighborhoods to consider in the graph learning layer. Default: 8.
hidA (int, optional) – Dimension of attention in the model. Default: 64.
hidR (int, optional) – Dimension of hidden layers in the recurrent neural network part. Default: 40.
hidP (int, optional) – Dimension of positional encoding in the model. Default: 1.
n_layer (int, optional) – Number of layers in the graph neural network. Default: 2.
dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.5.
device (str, optional) – The device (cpu or gpu) on which the model will be run. Default: ‘cpu’.

Returns:

A tensor of shape (batch_size, num_nodes, num_timesteps_output), representing the predicted values for each node over future timesteps.

Return type:

torch.Tensor

forward(X, adj, states=None, dynamic_adj=None, index=None)

Parameters:

X (torch.Tensor) – Input features tensor with shape (batch_size, num_timesteps_input, num_nodes, num_features).
adj (torch.Tensor) – Static adjacency matrix of the graph with shape (num_nodes, num_nodes).
states (torch.Tensor, optional) – States of the nodes if available, with the same shape as x. Default: None.
dynamic_adj (torch.Tensor, optional) – Dynamic adjacency matrix if available, with shape similar to adj but possibly varying over time. Default: None.

Returns:

The output tensor of shape (batch_size, num_timesteps_output, num_nodes), representing the predicted values for each node over the specified output timesteps.

Return type:

torch.Tensor

class epilearn.models.SpatialTemporal.CNNRNN_Res.CNNRNN_Res(num_nodes, num_features, num_timesteps_input, num_timesteps_output, nhid=32, residual_ratio=0, residual_window=0, dropout=0.5, device='cpu')

Combined Convolutional Neural Network and Recurrent Neural Network with Residual Connections (CNNRNN_Res)

Parameters:

num_nodes (int) – Number of nodes in the graph.
num_features (int) – Number of features per node per timestep.
num_timesteps_input (int) – Number of timesteps considered for each input sample (window size).
num_timesteps_output (int) – Number of output timesteps to predict.
nhid (int, optional) – Number of hidden units in the GRU layer. Default: 32.
residual_ratio (float, optional) – Proportion of the residual connection compared to the GRU output. Default: 0.
residual_window (int, optional) – Number of timesteps to include in the residual connection. Default: 0.
dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.5.
device (str, optional) – The device (cpu or gpu) on which the model will be run. Default: ‘cpu’.

Returns:

A tensor of shape (batch_size, num_timesteps_output, num_nodes), representing the predicted values for each node over future timesteps. Each slice along the second dimension corresponds to a timestep, with each column representing a node.

Return type:

torch.Tensor

forward(x, adj, states=None, dynamic_adj=None, **kargs)

Parameters:

x (torch.Tensor) – Input features tensor with shape (batch_size, num_timesteps_input, num_nodes, num_features).
adj (torch.Tensor) – Static adjacency matrix of the graph with shape (num_nodes, num_nodes).
states (torch.Tensor, optional) – States of the nodes if available, with the same shape as x. Default: None.
dynamic_adj (torch.Tensor, optional) – Dynamic adjacency matrix if available, with shape similar to adj but possibly varying over time. Default: None.

Returns:

The output tensor of shape (batch_size, num_timesteps_output, num_nodes), representing the predicted values for each node over the specified output timesteps.

Return type:

torch.Tensor

Network simulators

Mechanistic spreading models on a network. Like the Temporal compartmental blocks, these have no fit() and are not task prototypes – they are simulators.

class epilearn.models.SpatialTemporal.NetworkSIR.NetSIR(num_nodes=None, horizon=None, infection_rate=0.01, recovery_rate=0.038, population=None)

Network-based SIR (Susceptible-Infected-Recovered)

Parameters:

num_nodes (int, optional) – Number of nodes in the graph representing individuals or groups. Default: None.
horizon (int, optional) – Number of future time steps to simulate. If None, a single step is simulated unless overridden in the forward method.
infection_rate (float, optional) – Initial infection rate parameter, representing the rate at which susceptible individuals become infected. Default: 0.01.
recovery_rate (float, optional) – Initial recovery rate parameter, representing the rate at which infected individuals recover. Default: 0.038.
population (int, optional) – Total population considered in the model. If None, the sum of the initial conditions (susceptible, infected, recovered) is used as the total population.

Returns:

A tensor of shape (time_step, num_nodes, 3), representing the predicted number of susceptible, infected, and recovered individuals at each timestep for each node. Each row corresponds to a timestep, with the columns representing the susceptible, infected, and recovered counts respectively for each node.

Return type:

torch.Tensor

forward(x, adj, steps=1)

Parameters:

x (torch.Tensor) – Input features tensor with shape (n_nodes, one-hot encoding of states).
adj (torch.Tensor) – Static adjacency matrix of the graph with shape (num_nodes, num_nodes).
states (torch.Tensor, optional) – States of the nodes if available, with the same shape as x. Default: None.
dynamic_adj (torch.Tensor, optional) – Dynamic adjacency matrix if available, with shape similar to adj but possibly varying over time. Default: None.

Returns:

The output tensor of shape (time_step, n_nodes, probability of states), representing the predicted values for each node over the specified output timesteps.

Return type:

torch.Tensor

class epilearn.models.SpatialTemporal.DMP.DMP(num_nodes, recover_rate=None, horizon=1, seed_list=[14, 8], device='cpu')

forward(x, adj)

Parameters:

x (torch.Tensor) – Expected shape (batch_size, num_timesteps_input, num_nodes, num_features).
adj (torch.Tensor) – Adjacency matrix of the graph with shape (num_nodes, num_nodes), indicating connections between nodes.

Returns:

The output tensor of shape (horizon, num_nodes, 3), representing the predicted number of susceptible, infected, and recovered individuals at each timestep.

Return type:

torch.Tensor

Models with known limitations

Four exported spatial-temporal classes need care before they run unchanged through Forecast.rolling_train on a dataset built by Dataset.generate_dataset. They are documented here so the failure mode is not a surprise; the alternative is listed for each. MepoGNN is the mild case – it is a fully supported benchmark model that only trips on datasets carrying a dynamic graph.

class epilearn.models.SpatialTemporal.STGCN.STGCN_c(num_nodes, num_features, num_timesteps_input, num_timesteps_output, nhids=128, device='cpu', **kwargs)

Spatio-temporal graph convolutional network as described in https://arxiv.org/abs/1709.04875v3 by Yu et al. Input should have shape (batch_size, num_nodes, num_input_time_steps, num_features).

forward(X, adj, states=None, dynamic_adj=None, **kargs)

Parameters:

X (torch.Tensor) – Shape (batch_size, num_nodes, num_timesteps_input, num_features)
adj (torch.Tensor) – Shape (num_nodes, num_nodes)

Returns:

Output shape (batch_size, num_nodes) for travel_rate task

Return type:

torch.Tensor

class epilearn.models.SpatialTemporal.DASTGN.DASTGN(num_nodes, num_features, num_timesteps_input, num_timesteps_output, GNN_layers=2, nhids=None, device='cpu')

Dynamic and Adaptive Spatio-Temporal Graph Network (DASTGN)

Parameters:

num_nodes (int) – Number of nodes in the graph.
num_features (int) – Number of features per node per timestep.
num_timesteps_input (int) – Number of timesteps considered for each input sample.
num_timesteps_output (int) – Number of output timesteps to predict.
GNN_layers (int, optional) – Number of Graph Neural Network layers to use. Default: 2.
device (str, optional) – The device (cpu or gpu) on which the model will be run. Default: ‘cpu’.

Returns:

A tensor of shape (batch_size, num_timesteps_output, num_nodes), representing the predicted values for each node over future timesteps. Each slice along the second dimension corresponds to a timestep, with each column representing a node.

Return type:

torch.Tensor

forward(x, adj, states=None, dynamic_adj=None, **kwargs)

Parameters:

x (torch.Tensor) – Input features tensor with shape (batch_size, num_timesteps_input, num_nodes, num_features).
adj (torch.Tensor) – Static adjacency matrix of the graph with shape (num_nodes, num_nodes).
states (torch.Tensor, optional) – States of the nodes if available, with the same shape as x. Default: None.
dynamic_adj (torch.Tensor, optional) – Dynamic adjacency matrix if available, with shape similar to adj but possibly varying over time. Default: None.

Returns:

The output tensor of shape (batch_size, num_timesteps_output, num_nodes), representing the predicted values for each node over the specified output timesteps.

Return type:

torch.Tensor

class epilearn.models.SpatialTemporal.MepoGNN.MepoGNN(num_nodes, num_features, num_timesteps_input, num_timesteps_output, glm_type='Dynamic', adapt_graph=None, dropout=0.5, residual_channels=32, dilation_channels=32, skip_channels=256, end_channels=512, kernel_size=2, blocks=2, layers=3, nhids=None, device='cpu', **kwargs)

Meta-Population Graph Neural Network (MepoGNN)

Parameters:

num_nodes (int) – Number of nodes in the graph.
num_features (int) – Number of features per node per timestep.
num_timesteps_input (int) – Number of timesteps considered for each input sample.
num_timesteps_output (int) – Number of output timesteps to predict.
glm_type (str, optional) – Type of graph learning model (‘Dynamic’, ‘Adaptive’). Default: ‘Dynamic’.
adapt_graph (tensor, optional) – Initial tensor for adaptive adjacency matrix. Only needed if glm_type is ‘Adaptive’. Default: None.
dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.5.
residual_channels (int) – Number of channels in residual layers.
dilation_channels (int) – Number of channels in dilation layers.
skip_channels (int) – Number of channels in skip connection layers.
end_channels (int) – Number of channels in the final convolution layers.
kernel_size (int) – Kernel size for the convolution layers.
blocks (int) – Number of blocks in the WaveNet architecture.
layers (int) – Number of layers in each block.
device (str, optional) – The device (cpu or gpu) on which the model will be run. Default: ‘cpu’.

Returns:

A tensor of shape (batch_size, num_timesteps_output, num_nodes), representing the predicted values for each node over future timesteps. Each slice along the second dimension corresponds to a timestep, with each column representing a node.

Return type:

torch.Tensor

forward(x, adj, states, dynamic_adj, max_od=1000000.0)

Parameters:

x (torch.Tensor) – Input features tensor with shape (batch_size, num_timesteps_input, num_nodes, num_features).
adj (torch.Tensor) – Static adjacency matrix of the graph with shape (num_nodes, num_nodes).
states (torch.Tensor, optional) – States of the nodes if available, with the same shape as x. Default: None.
dynamic_adj (torch.Tensor, optional) – Dynamic adjacency matrix if available, with shape similar to adj but possibly varying over time. Default: None.

Returns:

The output tensor of shape (batch_size, num_timesteps_output, num_nodes), representing the predicted values for each node over the specified output timesteps.

Return type:

torch.Tensor

class epilearn.models.SpatialTemporal.STAN.STAN(num_nodes, num_features, num_timesteps_input, num_timesteps_output, population=10000000000.0, gat_dim1=32, gat_dim2=32, gru_dim=32, num_heads=1, nhids=None, device='cpu')

Spatio-Temporal Attention Network (STAN)

Parameters:

num_nodes (int) – Number of nodes in the graph.
num_features (int) – Number of features per node per timestep.
num_timesteps_input (int) – Number of timesteps considered for each input sample.
num_timesteps_output (int) – Number of output timesteps to predict.
population (float, optional) – Total population considered in the model. Default: 1e10.
gat_dim1 (int) – Dimension of the output space for the first GAT layer. Default: 32.
gat_dim2 (int) – Dimension of the output space for the second GAT layer. Default: 32.
gru_dim (int) – Dimension of the hidden state for the GRU layer. Default: 32.
num_heads (int) – Number of attention heads in the first GAT layer. Default: 1.
device (str, optional) – The device (cpu or gpu) on which the model will be run. Default: ‘cpu’.

Returns:

A tuple containing two tensors, each of shape (batch_size, num_timesteps_output, num_nodes, 1), representing the predicted values for newly infected and recovered individuals respectively for each node over future timesteps. The second tensor contains the physical model predictions.

Return type:

tuple of torch.Tensor

forward(X, adj, states, dynamic_adj=None, N=None, h=None)

Parameters:

X (torch.Tensor) – Input feature tensor with shape (batch_size, num_timesteps_input, num_nodes, num_features).
adj (torch.Tensor) – Static adjacency matrix with shape (num_nodes, num_nodes).
states (torch.Tensor) – States of the nodes, with the same shape as X, containing current infection and recovery data.
dynamic_adj (torch.Tensor, optional) – Dynamic adjacency matrix, with shape similar to adj but possibly varying over time. Default: None.
N (float, optional) – Total population considered in the model. Default: None.
h (torch.Tensor, optional) – Hidden states for the GRU layer, used if provided. Default: None.

Returns:

A tuple containing two tensors: 1. Predicted new infections and recoveries, shape (batch_size, num_timesteps_output, num_nodes, 2). 2. Physical model predictions based on current states, shape (batch_size, num_timesteps_output, num_nodes, 2).

Return type:

tuple of torch.Tensor

Class	What happens	Workaround
`STGCN_c`	`RuntimeError: mat1 and mat2 shapes cannot be multiplied`. Its output layer is sized for the unpadded temporal convolution of the original paper, while the shared `TimeBlock` now pads and preserves the temporal length.	Use `STGCN`.
`DASTGN`	`RuntimeError: The size of tensor a (num_nodes) must match the size of tensor b (horizon)`. `forward` returns `(batch, horizon, num_nodes)` while every other model – and `generate_dataset`’s targets – use `(batch, num_nodes, horizon)`.	Transpose the output yourself, or use another graph model.
`MepoGNN`	Only with a dynamic graph: `RuntimeError: einsum(): the number of subscripts in the equation (5) does not match the number of dimensions (4)`, because it needs a 5-D dynamic graph `(batch, lookback, num_nodes, num_nodes, 1)` while `generate_dataset` produces a 4-D one. The bundled toy `Dataset` carries one, so every fold fails there. With a static-only graph it trains normally – given `dynamic_adj=None` it derives its own 5-D mobility tensor from the static adjacency – which is why it is a supported entry in `SPATIOTEMPORAL_MODELS` and runs through the benchmark CLI on the toy CSV (see Benchmark).	Rebuild the `Dataset` without `dynamic_graph=`, or call `forward` directly with a 5-D `dynamic_adj`.
`STAN`	`AttributeError: 'tuple' object has no attribute 'size'`. `forward` returns `(predictions, physical_predictions)`, and the generic training loop feeds that tuple straight into the loss.	Call `forward` directly and combine the two heads with your own loss.

Code Example

import torch
from epilearn.models.SpatialTemporal import ColaGNN

num_nodes=47
num_features = 1
lookback = 16 # inputs size
horizon = 3 # predicts size


graph = torch.round(torch.rand((num_nodes, num_nodes)))
features = torch.round(torch.rand((10, lookback, num_nodes, num_features)))
# targets are (samples, num_nodes, horizon), the layout generate_dataset produces
node_target = torch.round(torch.rand((10, num_nodes, horizon)))

model=ColaGNN(num_nodes = num_nodes, num_features=num_features, num_timesteps_input=lookback, num_timesteps_output=horizon, device='cpu')
model.fit(
        train_input=features,
        train_target=node_target,
        train_graph=graph,
        val_input=None,
        val_target=None,
        val_graph=None,
        epochs=20,
        loss='mse'
        )