Tasks

A task pairs a model prototype with an evaluation protocol: it slices the data into windows, initializes the model from the shapes it finds, trains it, scores it and calibrates prediction intervals. The model is a parameter of the task, so moving a model between tasks means swapping the task class, not the model.

Task	Question it answers	Default metrics
Forecast	What do the next horizon steps look like?	MSE / MAE / RMSE
NowcastTask	The last few days are still being reported – what will the final counts be?	MSE / MAE / RMSE
ScenarioTask	What would have happened under a different intervention?	PEHE / ATE error
Detection	Which nodes are the outbreak sources?	accuracy / macro-F1

All four are importable from epilearn.tasks; the two new tasks also have short aliases:

from epilearn.tasks import Forecast, Detection, NowcastTask, ScenarioTask
from epilearn.tasks import Nowcast, Scenario   # aliases of the two above

This page is the reference: keywords, return shapes, and the trap each task carries. The runnable scripts live elsewhere – Quickstart for EpiLearn for a complete forecasting pipeline, Pipeline for Epidemic Modeling for a step-by-step walkthrough of all four tasks with the output each one prints.

Constructing a task

task = Forecast(prototype=STGCN, lookback=12, horizon=3, device='cpu')

Argument	Meaning
prototype / model	The model class, re-instantiated per fold with num_features, num_timesteps_input, num_timesteps_output, device and (when a graph is present) num_nodes filled in from the data; anything else goes through model_args. model= passes an already-built model instead, which is not re-initialized between folds.
lookback / horizon / ahead	Input window length, steps to predict, and the gap between the two (0 = the next step; ahead is on Forecast / Detection only). On Detection, horizon is the class count, not a step count.
dataset / device	dataset is optional on Forecast / Detection and absent from NowcastTask / ScenarioTask, which take it in rolling_train instead; device is 'cpu' or 'cuda'. NowcastTask additionally takes min_delay / max_delay, and ScenarioTask n_scenarios, baseline_scenario_idx, target_compartment and compartmental_model.

Rolling-window training

rolling_train is the primary entry point for every task, and it replaces the old train_model(dataset=..., train_rate=..., val_rate=...) call, which no longer splits the data for you (see MIGRATION.md). One call walks a rolling origin across the series, training a fresh model per fold, calibrating a conformal quantile on that fold’s validation window and scoring its test window:

fold 1: |--- train ---|- val -|- test -|
fold 2: |------ train ------|- val -|- test -|      (expanding=True)
fold 3: |-------- train --------|- val -|- test -|

result = task.rolling_train(dataset=dataset, train_size=400, val_size=50,
                            test_size=50, epochs=50, batch_size=5)
print(result['aggregate_metrics'])

Argument	Meaning
train_size / val_size / test_size	Timesteps in the first training window, in the window held out for early stopping and conformal calibration, and scored per fold. Keep val_size > 0: with 0 the quantile falls back to training residuals and the intervals come out optimistically narrow.
step_size / expanding / max_folds	How far the origin advances per fold (default test_size, i.e. non-overlapping test windows); whether the training window grows (True, default) or slides at fixed width; and a cap on the number of folds (None = every fold the series allows).
conformal_alpha	Miscoverage target, 0.1 (default) = 90% intervals.
use_optuna / n_trials	Tune hyperparameters per fold with Optuna. Search ranges come from optuna_model_args and optimizer_params; put 'lookback' in optuna_model_args to tune the input window too.
report_metrics	Metric names to compute, default ['mse', 'mae', 'rmse']. Accepts 'mape', 'r2', 'acc' or your own callables. ScenarioTask needs ['pehe', 'ate_error'].

The return value is a dictionary with seven entries: fold_results (per fold: its metrics, conformal_quantile, coverage, interval_width, val_residuals, test_split, process_history, best_params), aggregate_metrics (n_folds plus <metric>_mean / <metric>_std across folds), conformal_alpha, all_predictions, all_targets, conformal_intervals and runtime.

Warning

Seven entries only on success. A single failing fold is caught and reported, but if every fold fails the return value degrades to {'fold_results': [], 'aggregate_metrics': None}, so aggregate_metrics['mse_mean'] raises TypeError and result['all_predictions'] raises KeyError.

Note

Split-conformal intervals come for free – there is no separate calibration step – but the two coverage numbers EpiLearn reports are not the same quantity: the per-fold coverage here is marginal (fraction of individual (sample, region, horizon-step) entries inside their interval) while epilearn.utils.uncertainty.static_conformal()['coverage'] is joint (fraction of samples whose whole horizon is inside the band), and joint coverage is always the smaller number.

Note

Metrics are computed on whatever scale the model trained on: with transforms.normalize_target() attached they are in normalized units, and only evaluate_model(inverse_normalize=True) reports original units.

BaseTask.rolling_train(dataset, train_size: int, test_size: int, val_size: int = 0, step_size: int | None = None, expanding: bool = True, max_folds: int | None = None, train_loss='mse', val_loss='mse', epochs=100, batch_size=32, lr=0.001, weight_decay=0, patience=50, verbose=False, model_args={}, conformal_alpha=0.1, regions=None, device=None, use_optuna=False, n_trials=10, optimizer_params=None, optuna_model_args=None, residual_type=None, report_metrics=None)

Rolling/walk-forward evaluation for time series forecasting.

Trains and evaluates using a rolling window approach to avoid data leakage. Optionally uses Optuna for per-fold hyperparameter tuning.

For each fold: 1. Train model on training window 2. Compute predictions on validation set → calibrate conformal quantile 3. Compute predictions on test set → apply conformal intervals

This gives per-fold uncertainty estimation using the validation data that is temporally closest to the test data.

Parameters:

• dataset – Dataset object with timestamps

• train_size – Initial training window size

• test_size – Test window size for each fold

• val_size – Validation window size (required for conformal prediction)

• step_size – Timesteps to advance each fold (default = test_size)

• expanding – If True, training window expands; if False, slides

• max_folds – Maximum number of folds (None = all possible)

• train_loss – Loss functions (str like ‘mse’ or nn.Module)

• val_loss – Loss functions (str like ‘mse’ or nn.Module)

• epochs – Training params

• batch_size – Training params

• lr – Training params

• weight_decay – Training params

• patience – Training params

• verbose – Print training progress

• model_args – Fixed model arguments

• conformal_alpha – Significance level for conformal prediction (default 0.1 = 90% coverage)

• regions – Optional list of regions

• device – Training device

• use_optuna – Enable Optuna tuning per fold

• n_trials – Optuna trials per fold

• optimizer_params – Optimizer hyperparameter ranges (can include custom loss)

• optuna_model_args – Model hyperparameter ranges (can include ‘lookback’)

• residual_type – How to compute residuals for conformal prediction: - None: Standard |preds - targets| - callable: Custom function (preds, targets) -> residuals

• report_metrics – List of metric names to compute and report (default: [‘mse’, ‘mae’, ‘rmse’]) Available metrics: ‘mse’, ‘mae’, ‘rmse’, ‘mape’, ‘r2’, ‘acc’ or custom callables

Returns:

• fold_results: Per-fold metrics + conformal results

• aggregate_metrics: Mean/std of specified metrics + coverage

• conformal_intervals: Per-fold prediction intervals

• all_predictions/all_targets: Concatenated predictions and targets

Return type:

Dictionary with

Forecast

Predicts the next horizon steps from the previous lookback steps; see Quickstart for EpiLearn for the full script.

class epilearn.tasks.forecast.Forecast(prototype=None, model=None, dataset=None, lookback=None, horizon=None, ahead=0, device='cpu')

The Forecast class extends the BaseTask class, focusing on the training and evaluation of forecast models for time-series prediction tasks. It includes functionalities specific to handling time-series data, especially in settings that involve spatial-temporal dynamics. The class supports model initialization, training, evaluation, and preprocessing, facilitating the application of various neural network architectures and configurations.

evaluate_model(model=None, dataset=None, process_history=None, use_conformal=True, conformal_quantile=None, inverse_normalize=False, residue_func=None)

Evaluate the trained model and compute metrics with adaptive conformal prediction intervals.

Parameters:

• model – Model to evaluate (uses self.model if None)

• dataset – Dataset dictionary with ‘features’, ‘targets’, ‘graph’, etc.

• process_history – Dict with ‘target_mean’, ‘target_std’ for inverse normalization

• use_conformal – Whether to compute conformal prediction intervals

• conformal_quantile – Conformal quantile (uses self.conformal_quantile if None)

• inverse_normalize – Whether to inverse normalize predictions/targets

• residue_func – Custom residual function for misaligned predictions

Returns:

• Point metrics: mse, mae, rmse, mape, r2

• Residual statistics: residual_mean, residual_std

• Raw outputs: predictions, targets, residuals

• Conformal results: adaptive_lower, adaptive_upper, coverage stats (if use_conformal=True)

Return type:

Dictionary containing

inverse_norm(data, mean, std, node_indices=None)

Apply inverse normalization to data.

Handles multiple cases: - Global normalization (scalar mean/std) - Per-node normalization with data that still has node dimension - Per-node normalization with flattened data and node_indices - Per-node normalization with flattened data without node_indices (uses global average)

Parameters:

• data – Tensor to denormalize

• mean – Either scalar or array of means (one per node)

• std – Either scalar or array of stds (one per node)

• node_indices – Optional tensor of node indices for flattened data

Returns:

Denormalized tensor

plot_preds(eval_results, n_show=None, figsize=(15, 7), save_path=None, backend='matplotlib', region_idx=0, horizon_idx=-1, interactive=False)

Plot predictions with adaptive uncertainty intervals in a clean white style.

Parameters:

• eval_results – Dictionary returned from evaluate_model containing predictions, targets, and uncertainty estimates

• n_show – Number of time samples to display (None plots all samples, default: None)

• figsize – Figure size as (width, height) tuple for matplotlib (default: (15, 7))

• save_path – Optional path to save the figure (e.g., ‘plot.png’ or ‘plot.html’)

• backend – Plotting backend - ‘matplotlib’ or ‘plotly’ (default: ‘matplotlib’)

• region_idx – Index of region to plot (default: 0)

• horizon_idx – Index of horizon step to plot (default: -1, last step)

• interactive – Whether to make plotly plots interactive (default: False)

Returns:

fig, ax For plotly: fig

Return type:

For matplotlib

dataset.transforms is fitted on each fold’s training window only and applied from there to that fold’s validation and test windows, so the statistics never leak forward in time. Each fold returns its fitted constants as process_history, which is what evaluate_model needs to undo them:

last_fold = result['fold_results'][-1]
evaluation = task.evaluate_model(dataset=last_fold['test_split'],        # split dict
                                 process_history=last_fold['process_history'],
                                 inverse_normalize=True)
task.plot_preds(evaluation, region_idx=0, horizon_idx=-1, save_path='f.png')

evaluate_model returns mse, mae, rmse, mape, r2, median_ae, max_error, residual_mean, residual_std and the raw predictions / targets / residuals tensors. Its dataset= is a split dict (features, targets, graph, dynamic_graph, states) from Dataset.generate_dataset or fold_results[i]['test_split']; hand it a Dataset, or nothing at all, and it fails on the lookup instead.

Warning

Do not pass conformal_quantile= to evaluate_model: the helper it dispatched to was removed and the call now raises NotImplementedError. Read the calibrated value from fold_results[i]['conformal_quantile'], or apply the strategies in Utils to fold_results[i]['val_residuals'].

Detection

Classifies nodes – typically outbreak source detection – from a spatial snapshot, so it pairs with the Spatial models.

Warning

horizon here is the number of classes, not a number of future steps: it reaches the model as num_classes, and the target window length is set separately by horizon_size=1 in generate_dataset. That also rules out rolling_train, which builds windows with horizon_size=self.horizon: the targets come out the wrong shape and every fold dies in the cross-entropy loss with ValueError: Expected input batch_size (...) to match target batch_size (...). Use the single-split train_model instead.

class epilearn.tasks.detection.Detection(prototype=None, model=None, dataset=None, lookback=None, horizon=None, ahead=0, device='cpu')

The Detection class extends the BaseTask class, focusing on the training and evaluation of source detection models for spatial-temporal prediction tasks. It predicts the source node given future spatial-temporal observations. The class supports model initialization, training, evaluation with conformal prediction, and comprehensive metrics for classification tasks including accuracy, precision, recall, and F1 score.

evaluate_model(model=None, dataset=None, process_history=None, use_conformal=True, conformal_quantile=None, n_bootstrap=100, compute_bootstrap_ci=True)

Comprehensive evaluation of the trained detection model with uncertainty quantification.

Parameters:

• model – Model to evaluate (uses self.model if None)

• dataset – Dataset dict with ‘features’, ‘targets’, ‘graph’, etc.

• process_history – Processing history (not used for detection, kept for API consistency)

• use_conformal – Whether to compute conformal prediction intervals

• conformal_quantile – Conformal quantile (uses self.conformal_quantile if None)

• n_bootstrap – Number of bootstrap samples for uncertainty estimation

• compute_bootstrap_ci – Whether to compute bootstrap confidence intervals

Returns:

Dictionary containing comprehensive evaluation metrics and uncertainty estimates

plot_preds(eval_results, n_show=None, figsize=(15, 7), save_path=None, backend='matplotlib', sample_idx=0, interactive=False)

Plot source detection predictions with confidence scores.

Parameters:

• eval_results – Dictionary returned from evaluate_model containing predictions, targets, and probabilities

• n_show – Number of nodes to display (None plots all nodes, default: None)

• figsize – Figure size as (width, height) tuple for matplotlib (default: (15, 7))

• save_path – Optional path to save the figure (e.g., ‘plot.png’ or ‘plot.html’)

• backend – Plotting backend - ‘matplotlib’ or ‘plotly’ (default: ‘matplotlib’)

• sample_idx – Index of sample to plot (default: 0)

• interactive – Whether to make plotly plots interactive (default: False)

Returns:

fig, ax For plotly: fig

Return type:

For matplotlib

task = Detection(prototype=GCN, dataset=dataset, lookback=1, horizon=2)  # 2 classes
task.train_model(train_split=train_split, val_split=val_split,
                 test_split=test_split, train_loss='ce', val_loss='ce')
evaluation = task.evaluate_model(dataset=test_split)   # a split dict, not a Dataset

Build those splits with generate_dataset(..., horizon_size=1), passing adj= or the graph reaching the model is silently None (Pipeline for Epidemic Modeling has the script); train_model needs all three splits, and it catches training exceptions, prints them and returns None. evaluate_model returns accuracy, macro_precision, macro_recall, macro_f1, per-class precision_per_class / recall_per_class / f1_per_class, mean_confidence, predictions / targets / probabilities, and a bootstrap confidence interval per metric (turn the resampling off with compute_bootstrap_ci=False).

NowcastTask

Nowcasting corrects for reporting delay: recent counts are still incomplete, and the task learns how much each day will be revised upward. The input is a reporting triangle – row t holds the counts known about day t after delays d = min_delay ... max_delay, not-yet-observable entries marked -1. lookback is how many past days of the triangle the model sees, horizon how many recent incomplete days it nowcasts per sample, and n_delays = max_delay - min_delay + 1 becomes the model’s feature dimension.

class epilearn.tasks.nowcast.NowcastTask(prototype=None, model=None, lookback: int = 30, horizon: int = 7, min_delay: int = 3, max_delay: int | None = None, device: str = 'cpu')

Nowcasting task for epidemiological reporting triangles.

Extends BaseTask to use the same model interface as Forecast. Supports ALL models from epilearn/models: - Temporal: GRU, LSTM, MLP, Transformer, etc. - Statistical: ARIMA, etc.

Data format: - Features: (batch, lookback, n_regions, n_delays)

• At each time step, input the n_delays vector

• Unobserved values are marked as -1

• n_delays = max_delay - min_delay + 1

• Targets: (batch, n_regions, horizon) - Final counts for the last horizon days

Usage:

from epilearn.tasks import Nowcast from epilearn.models.Temporal import GRUModel

# Create task with any model task = NowcastTask(prototype=GRUModel, lookback=30, horizon=7,

min_delay=3, max_delay=30)

# Load triangle data and create dataset data = task.load_triangle(‘data.npz’) dataset = task.create_dataset(data[‘triangle’], data[‘final_counts’])

# Run training with rolling evaluation (uses base.py’s rolling_train) results = task.rolling_train(

dataset=dataset, train_size=500, test_size=100, val_size=100,

)

compute_naive_baseline(dataset: Dataset) → Dict

Compute naive baseline metrics for comparison.

The naive baseline uses the latest available observation as the prediction.

create_dataset(triangle: ndarray, final_counts: ndarray, delays: ndarray | None = None) → Dataset

Convert reporting triangle to Dataset for use with rolling_train.

Parameters:

• triangle – (n_days, total_delays) reporting triangle

• final_counts – (n_days,) final values for each day

• delays – Optional array of delay values (e.g., [3, 4, 5, …, 60]). Used to map min_delay/max_delay to column indices.

Returns:

• x: (n_samples, lookback, n_regions=1, n_delays) - features

• y: (n_samples, n_regions=1, horizon) - targets

Return type:

Dataset with

Note

• n_delays = max_delay - min_delay + 1 (or all columns if max_delay not set)

• Unobserved values in features are marked as -1

• Models should handle -1 appropriately (mask or replace)

static load_triangle(path: str) → Dict[str, ndarray]

Load reporting triangle from .npz file.

create_dataset shapes features as (n_samples, lookback, 1, n_delays) and targets as (n_samples, 1, horizon), ready for rolling_train; load_triangle expects a .npz holding triangle, final_counts, delays and time_values. compute_naive_baseline returns {'naive_mae': ..., 'n_samples': ...} for the “just trust the latest report” baseline – the number a nowcast has to beat – computed over every sample in the dataset rather than the rolling test windows, so read it as a reference level and not a fold-matched comparison.

Note

Neither NowcastTask nor BaseTask defines evaluate_model, so task.evaluate_model(...) raises AttributeError: 'NowcastTask' object has no attribute 'evaluate_model'. Score nowcasts from the rolling_train return value: aggregate_metrics, fold_results[i], or all_predictions / all_targets for raw values.

ScenarioTask

Scenario modeling answers counterfactual questions: several intervention policies share the same observed history, then diverge over the projection horizon. Each sample holds n_scenarios futures, one of which (baseline_scenario_idx, default 0) is the no-intervention control, and the task is scored on the difference from that control, measured on target_compartment ('I' by default): PEHE (sqrt(mean((tau_pred - tau_true)**2)), error on the per-sample treatment effect tau = Y_intervention - Y_baseline) and ATE error (|mean(tau_pred) - mean(tau_true)|, error on its average).

class epilearn.tasks.scenario_modeling.ScenarioTask(prototype=None, model=None, lookback: int = 60, horizon: int = 30, n_scenarios: int = 4, compartmental_model: SEIRVIModel | None = None, baseline_scenario_idx: int = 0, target_compartment: str = 'I', device: str = 'cpu')

Scenario Modeling task for epidemiological counterfactual analysis.

Extends BaseTask to support multi-scenario predictions with shared history.

Data format: - Features X: (batch, lookback, n_scenarios, N+4)

• N compartments + 4 intervention features (vacc_rate, vacc_delay, isol_rate, isol_delay)

• Compartments are IDENTICAL across scenarios (shared history)

• Intervention features can differ for dynamic interventions

• Targets Y: (batch, horizon, n_scenarios, N) - Future compartment values per scenario - Scenarios diverge based on different intervention policies

Metrics (Treatment Effect): - PEHE (Precision in Estimation of Heterogeneous Effects):

Measures accuracy in predicting the difference between intervention and baseline

• ATE Error (Average Treatment Effect Error): Measures accuracy in predicting the average effect of intervention

Usage:

from epilearn.tasks import ScenarioTask from epilearn.models.Temporal import GRUModel

# Create task with any model task = ScenarioTask(prototype=GRUModel, lookback=60, horizon=30, n_scenarios=4)

# Generate dataset from compartmental model dataset = task.generate_dataset(n_samples=100)

# Run training with rolling evaluation results = task.rolling_train(dataset=dataset, …)

evaluate_model(model=None, dataset=None, baseline_scenario_idx: int | None = None, target_compartment: str | None = None, **kwargs) → Dict

Evaluate scenario model with treatment effect metrics.

Unlike standard forecasting metrics (MSE, MAE), scenario modeling uses metrics that measure how well the model captures the EFFECT of interventions.

Metrics: 1. PEHE (Precision in Estimation of Heterogeneous Effects):

sqrt(E[(tau_pred - tau_true)^2]) where tau = Y_intervention - Y_baseline

2. ATE Error (Average Treatment Effect Error): |E[tau_pred] - E[tau_true]|

Parameters:

• model – Model to evaluate

• dataset – Dataset dict with ‘features’, ‘targets’

• baseline_scenario_idx – Index of baseline scenario (default: uses self.baseline_scenario_idx)

• target_compartment – Compartment name or index (default: uses self.target_compartment)

Returns:

Dictionary with PEHE, ATE Error, and predictions

generate_dataset(n_samples: int = 100, population: float = 1000000.0, process_noise: float | None = None, seed: int = 42, **kwargs) → Dataset

Generate scenario modeling dataset using compartmental model.

Parameters:

• n_samples – Number of samples to generate

• population – Population size

• process_noise – Optional noise in simulation

• seed – Random seed

• **kwargs – Additional parameters for generate_multi_scenario_dataset

Returns:

• x: (n_samples, lookback, n_scenarios, N+4)

• y: (n_samples, horizon, n_scenarios, N)

Return type:

Dataset with

plot_scenario_comparison(eval_results: Dict, sample_idx: int = 0, compartment_idx: int = 2, scenario_names: List[str] | None = None, figsize: Tuple[int, int] = (14, 5))

Plot scenario comparison for a single sample.

Parameters:

• eval_results – Results from evaluate_model

• sample_idx – Which sample to plot

• compartment_idx – Which compartment to plot (default: 2 = Infectious)

• scenario_names – Names for each scenario

• figsize – Figure size

task = ScenarioTask(prototype=GRUModel, lookback=20, horizon=10,
                    n_scenarios=4, target_compartment='I', device='cpu')
dataset = task.generate_dataset(n_samples=150, population=1e6, seed=42)
result = task.rolling_train(dataset, train_size=60, val_size=45, test_size=45,
                            report_metrics=['pehe', 'ate_error'])

ScenarioTask needs no data of your own: generate_dataset simulates trajectories with the built-in SEIRVIModel (epilearn.utils.compartmental_models, replaceable via compartmental_model=), sampling vaccination and isolation rates/delays per scenario. Features come out as (n_samples, lookback, n_scenarios, N + 4) – N compartments plus four intervention descriptors – and targets as (n_samples, horizon, n_scenarios, N). Those “timesteps” are independent simulated samples, so train_size / val_size / test_size count samples rather than days, and both metrics carry the units of the simulation (people out of population). Passing a single split dict to evaluate_model returns pehe, ate_error, mse, predictions / targets and the treatment-effect tensors tau_pred / tau_true of shape (n_samples, n_scenarios - 1, horizon), which plot_scenario_comparison draws.

Important

PEHE and ATE error are the only metrics ScenarioTask computes, and they are not in the default report_metrics. Pass report_metrics=['pehe', 'ate_error'] to rolling_train; otherwise it goes looking for 'mse', finds nothing, and aggregate_metrics comes back holding only the conformal numbers. The per-fold pehe / ate_error are in fold_results either way – it is the aggregation and the printed summary that go missing.