Utils

epilearn.utils groups the framework’s helper code into six modules:

Module	Contents
`epilearn.utils.utils`	Tensor/graph helpers, plus the `moving_avg` / `series_decomp` blocks used by the deep models.
`epilearn.utils.metrics`	Loss functions and error metrics (MSE/MAE/RMSE/ACC).
`epilearn.utils.transforms`	Composable preprocessing applied to a `Dataset`.
`epilearn.utils.uncertainty`	Conformal prediction strategies and interval-quality metrics. New in 0.1.0.
`epilearn.utils.compartmental_models`	SIR / SEIR / SIRS / SEIR-VI ODE models with time-varying parameters. New in 0.1.0.
`epilearn.utils.simulation`	Graph generators and the temporal / individual / regional epidemic simulators.

Every module is imported eagerly, so from epilearn.utils import uncertainty (or from epilearn import utils; utils.simulation...) works after a plain import epilearn.

Utility_Functions

Accuracy

epilearn.utils.utils.accuracy(output, labels)

Return accuracy of output compared to labels.

Parameters:

output (torch.Tensor) – output from model
labels (torch.Tensor or numpy.array) – node labels

Returns:

accuracy

Return type:

float

Normalize

epilearn.utils.utils.normalize(X)

Normalizes the input tensor X to have zero mean and unit standard deviation. Handles 3D tensors.

Parameters:

X (torch.Tensor) – The input tensor to be normalized.

Returns:

torch.Tensor – The normalized tensor.
torch.Tensor – The means of the input tensor.
torch.Tensor – The standard deviations of the input tensor.

Normalize_Adj

epilearn.utils.utils.normalize_adj(Adj)

Returns the degree normalized adjacency matrix.

Parameters:: Adj (torch.Tensor or np.array) – The input adjacency matrix.
Returns:: The degree normalized adjacency matrix.
Return type:: torch.Tensor

Diff

epilearn.utils.utils.diff(features)

Computes the discrete difference along the time dimension of the input tensor.

Parameters:: features (torch.Tensor) – The input feature tensor.
Returns:: The tensor of differences along the time dimension.
Return type:: torch.Tensor

Degree_Matrix

epilearn.utils.utils.Degree_Matrix(ST_matrix)

Computes the degree matrix for a given spatio-temporal adjacency matrix.

Parameters:: ST_matrix (torch.Tensor) – The input spatio-temporal adjacency matrix.
Returns:: The degree matrix.
Return type:: torch.Tensor

Static_Full

epilearn.utils.utils.Static_full(n, t, A)

Constructs the full spatio-temporal adjacency matrix using the binary spatio-temporal adjacency matrix method.

Parameters:

n (int) – The dimension of the spatial adjacency matrix.
t (int) – The length of periods.
A (torch.Tensor) – The spatial adjacency matrix.

Returns:

The full spatio-temporal adjacency matrix.

Return type:

torch.Tensor

Kronecker

epilearn.utils.utils.kronecker(A, B)

Constructs the spatio-temporal adjacency matrix using the Kronecker product.

Parameters:

A (torch.Tensor) – The temporal adjacency matrix.
B (torch.Tensor) – The spatial adjacency matrix.

Returns:

The adjacency matrix of one space-time neighboring block.

Return type:

torch.Tensor

Edge_to_Adj

epilearn.utils.utils.edge_to_adj(edge_index, num_nodes)

Converts edge index representation to adjacency matrix.

Parameters:

edge_index (torch.Tensor) – The edge index tensor where each column represents an edge.
num_nodes (int) – The number of nodes in the graph.

Returns:

The adjacency matrix.

Return type:

torch.Tensor

Moving_Avg

Note

moving_avg, series_decomp and series_decomp_multi live in epilearn.utils.utils; epilearn.utils.transforms.moving_avg still resolves (transforms does from .utils import *) but is not the canonical import.

class epilearn.utils.utils.moving_avg(kernel_size, stride)

Moving average block to highlight the trend of time series. This module smooths the input time series data using a moving average filter, which helps in capturing the underlying trend by averaging over a specified window.

Parameters:

kernel_size (int) – The size of the moving average window.
stride (int) – The stride of the moving average window.

forward(x)

Applies the moving average to the input tensor ‘x’, padding the time series at both ends to ensure that the moving average is computed correctly for the entire series.

Parameters:: x (torch.Tensor) – The input time-series data tensor, typically of shape (batch_size, time_steps, features).
Returns:: The tensor smoothed using the moving average, maintaining the original shape of the input.
Return type:: torch.Tensor

Series_Decomp

class epilearn.utils.utils.series_decomp(kernel_size)

A PyTorch module for series decomposition using a single moving average kernel. This module decomposes a time series into its residual and moving average components, providing a straightforward approach to trend extraction.

Parameters:: kernel_size (int) – The kernel size to be used for the moving average calculation. The kernel size defines the window for the moving average.

forward(x)

Applies the series decomposition to the input tensor ‘x’. It calculates the moving average using the specified kernel size and computes the residual by subtracting the moving average from the original series.

Parameters:: x (torch.Tensor) – The input time-series data tensor, typically of shape (batch_size, time_steps).
Returns:: A tuple containing: - The residual tensor after subtracting the moving average from the input. - The moving average tensor computed from the input.
Return type:: tuple(torch.Tensor, torch.Tensor)

Series_Decomp_Multi

class epilearn.utils.utils.series_decomp_multi(kernel_size)

A PyTorch module for series decomposition using multiple moving average kernels. This module decomposes a time series into its residual and moving average components, leveraging multiple kernel sizes for enhanced flexibility and accuracy in capturing trends.

Parameters:: kernel_size (list of int) – List of kernel sizes to be used for moving average calculations. Each kernel size defines the window for the moving average.

forward(x)

Applies the series decomposition to the input tensor ‘x’. It calculates multiple moving averages using the initialized kernels, weights them using a linear layer followed by a softmax, and computes the residual by subtracting the weighted moving average from the original series.

Parameters:: x (torch.Tensor) – The input time-series data tensor, typically of shape (batch_size, time_steps).
Returns:: A tuple containing: - The residual tensor after subtracting the weighted moving average from the input. - The weighted moving average tensor computed from the input.
Return type:: tuple(torch.Tensor, torch.Tensor)

Metrics

MSE_loss

epilearn.utils.metrics.get_loss(loss_name='mse')

Retrieves the specified loss function based on the input loss name. It supports mean squared error (MSE), a standardized loss (stan), an epidemic-collaboration specific loss (epi_cola), and cross-entropy loss.

Parameters:: loss_name (str, optional) – Name of the loss function to retrieve. Default is ‘mse’.
Returns:: The corresponding loss function as specified by loss_name.
Return type:: callable

Stan_loss

epilearn.utils.metrics.stan_loss(output, label, scale=0.5)

Calculates a combined mean squared error loss on predicted and physically informed predicted values, scaled by a given factor.

Parameters:

output (tuple of torch.Tensor) – The predicted values and the physically informed predicted values.
label (torch.Tensor) – The ground truth values.
scale (float, optional) – Scaling factor for the physical informed loss component. Default: 0.5.

Returns:

The calculated total loss as a scalar tensor.

Return type:

torch.Tensor

Epi_cola_loss

epilearn.utils.metrics.epi_cola_loss(output, label, scale=0.5)

Calculates a combined L1 and mean squared error loss on the output and an epidemiological output, scaled by a given factor.

Parameters:

output (tuple of torch.Tensor) – The primary model output and the epidemiological model output.
label (torch.Tensor) – The ground truth values.
scale (float, optional) – Scaling factor for the epidemiological loss component. Default: 0.5.

Returns:

The calculated total loss as a scalar tensor.

Return type:

torch.Tensor

Cross_entropy_loss

epilearn.utils.metrics.cross_entropy_loss(output, label)

Computes the cross-entropy loss between the logits and labels, adjusting the label tensor to fit the logits dimensions.

Parameters:

output (torch.Tensor) – The logits from the model.
label (torch.Tensor) – The ground truth labels, scaled to match the number of classes based on output dimensions.

Returns:

The cross-entropy loss as a scalar tensor.

Return type:

torch.Tensor

MAE

epilearn.utils.metrics.get_MAE(pred, target)

Calculates the Mean Absolute Error (MAE) between predictions and targets.

Parameters:

pred (torch.Tensor) – Predicted values.
target (torch.Tensor) – Ground truth values.

Returns:

The MAE value as a scalar tensor.

Return type:

torch.Tensor

RMSE

epilearn.utils.metrics.get_RMSE(pred, target)

Calculates the Root Mean Squared Error (RMSE) between predictions and targets.

Parameters:

pred (torch.Tensor) – Predicted values.
target (torch.Tensor) – Ground truth values.

Returns:

The RMSE value as a scalar tensor.

Return type:

torch.Tensor

ACC

epilearn.utils.metrics.get_ACC(pred, target)

Calculates the accuracy of predictions by comparing them to the targets.

Parameters:

pred (torch.Tensor) – Predicted labels.
target (torch.Tensor) – True labels.

Returns:

The accuracy as a scalar tensor.

Return type:

torch.Tensor

Uncertainty

New in 0.1.0. epilearn.utils.uncertainty turns point forecasts into calibrated prediction intervals. Four conformal strategies are provided, all sharing one signature and one return schema:

strategy(val_residuals, predictions, targets, target_alpha=0.1)

Strategy	Interval width
`static_conformal`	Constant (split conformal).
`compute_aci`	Varies over time; `alpha_t` adapts sequentially (Gibbs & Candès, 2021).
`locally_weighted_conformal`	Per-sample, scales with prediction magnitude.
`locally_weighted_aci`	Per-sample and time-adaptive.

Note the argument order: val_residuals (1-D absolute residuals from the calibration split) comes first, and the miscoverage keyword is target_alpha, not alpha. All four share one return schema — lower, upper, coverage, avg_width, winkler_score plus the alpha_trace / quantile_trace / coverage_per_sample diagnostics — or None when val_residuals is empty.

Note

The returned coverage is joint — a test sample counts as covered only if every horizon step is inside its interval — while task.rolling_train and compute_uncertainty_metrics() report marginal (element-wise) coverage. The two are not comparable.

Inputs are NumPy arrays in original (denormalized) scale, and the locally-weighted pair wants val_predictions= (calibration-set predictions) or falls back to a simplified, over-covering procedure. Utilities works both points through with numbers and compares the four strategies on a trained model.

Static_Conformal

epilearn.utils.uncertainty.static_conformal(val_residuals: ndarray, predictions: ndarray, targets: ndarray, target_alpha: float = 0.1) → dict

Plain split-conformal prediction intervals.

Parameters:

val_residuals (1-D array) – Absolute residuals from the calibration set.
predictions (arrays, shape (n_test, ...)) – Test-set predictions and ground truth (denormalized).
targets (arrays, shape (n_test, ...)) – Test-set predictions and ground truth (denormalized).
target_alpha (float) – Target miscoverage rate.

Return type:

dict — same schema as compute_aci().

Compute_ACI

epilearn.utils.uncertainty.compute_aci(val_residuals: np.ndarray, predictions: np.ndarray, targets: np.ndarray, target_alpha: float = 0.1, gamma: float = 0.005) → dict | None

Adaptive Conformal Inference (Gibbs & Candès, 2021).

Processes test samples sequentially, adapting the miscoverage rate α_t based on observed coverage at each step.

Parameters:

val_residuals (1-D array, shape (n_cal,)) – Absolute residuals from the calibration set.
predictions (array, shape (n_test,) or (n_test, horizon)) – Point predictions for the test set (denormalized).
targets (array, same shape as predictions) – Ground truth values.
target_alpha (float) – Target miscoverage rate (0.1 → 90 % coverage).
gamma (float) – Learning rate for α adaptation.

Returns:

lower, upper, coverage, avg_width, winkler_score, alpha_trace, quantile_trace, coverage_per_sample.

Return type:

dict

Locally_Weighted_Conformal

epilearn.utils.uncertainty.locally_weighted_conformal(val_residuals: np.ndarray, predictions: np.ndarray, targets: np.ndarray, target_alpha: float = 0.1, *, val_predictions: np.ndarray | None = None, difficulty_fn=None, floor_quantile: float = 0.1) → dict | None

Normalised / locally-weighted split-conformal prediction.

Produces per-sample heteroscedastic intervals whose width scales with prediction magnitude — matching the natural heteroscedasticity of epidemic time-series.

When val_predictions are available the full Lei et al. (2018) / Papadopoulos et al. (2008) procedure is used:

Compute difficulty s_cal for calibration samples.
Normalise residuals: r_norm = r_i / s_cal_i.
q_norm = (1-α)-quantile of r_norm.
interval_i = pred_i ± q_norm × s_test_i.

When val_predictions are not available (the common case for post-hoc analysis of saved benchmarks), a simplified procedure is used that is equivalent to assuming calibration difficulty ≈ 1:

q = (1-α)-quantile of raw val_residuals.
interval_i = pred_i ± q × s_test_i.

Parameters:

val_residuals (1-D array) – Absolute residuals from the calibration set (original scale).
predictions (arrays, shape (n_test, ...)) – Test-set data (denormalized).
targets (arrays, shape (n_test, ...)) – Test-set data (denormalized).
target_alpha (float) – Target miscoverage rate.
val_predictions (array, optional) – Calibration-set predictions (denormalized). When provided the full normalised procedure is used.
difficulty_fn (callable, optional) – f(predictions) → difficulty (n,) array. Defaults to difficulty_from_predictions().
floor_quantile (float) – Passed to the default difficulty function.

Returns:

difficulty_test — per-sample difficulty score.

Return type:

dict — same schema as compute_aci(), plus

Locally_Weighted_ACI

epilearn.utils.uncertainty.locally_weighted_aci(val_residuals: np.ndarray, predictions: np.ndarray, targets: np.ndarray, target_alpha: float = 0.1, gamma: float = 0.005, *, val_predictions: np.ndarray | None = None, difficulty_fn=None, floor_quantile: float = 0.1) → dict | None

Locally-weighted Adaptive Conformal Inference.

Combines the per-sample difficulty scaling of locally_weighted_conformal() with the sequential α-adaptation of compute_aci(). This gives:

Per-sample heteroscedastic widths (via difficulty scores), and
Long-run coverage control (via ACI α-adaptation).

Parameters:

val_residuals (1-D array) – Absolute residuals from calibration (original scale, flattened).
predictions (arrays, shape (n_test, ...)) – Test set.
targets (arrays, shape (n_test, ...)) – Test set.
target_alpha (float) – ACI parameters.
gamma (float) – ACI parameters.
val_predictions (array, optional) – Calibration-set predictions. See locally_weighted_conformal().
difficulty_fn – See locally_weighted_conformal().
floor_quantile – See locally_weighted_conformal().

Return type:

dict — same schema as compute_aci(), plus difficulty_test.

Winkler_Score

epilearn.utils.uncertainty.winkler_score(lower: ndarray, upper: ndarray, targets: ndarray, alpha: float) → float

Winkler interval score (lower is better).

score = width + (2/α) × (penalty_below + penalty_above)

Compute_Uncertainty_Metrics

epilearn.utils.uncertainty.compute_uncertainty_metrics(predictions: ndarray, targets: ndarray, lower: ndarray, upper: ndarray, target_alpha: float = 0.1) → dict

Evaluate quality of prediction intervals.

Returns:: coverage, avg_width, winkler_score, width_abs_corr_r, width_abs_corr_rho, mse_by_uncertainty_q, per_horizon_coverage.
Return type:: dict

Difficulty_From_Predictions

epilearn.utils.uncertainty.difficulty_from_predictions(predictions: np.ndarray, *, floor: float | None = None, floor_quantile: float = 0.1, ref_floor: float | None = None, ref_median: float | None = None) → np.ndarray

Per-sample difficulty score based on prediction magnitude.

For epidemic time-series errors typically scale with the magnitude of the predicted count — high-incidence regions/periods have larger absolute errors. This function returns max(mean|pred_i|, floor) per sample, normalised so the median difficulty is 1.0.

Parameters:

predictions (ndarray, shape (n, ...)) – Denormalized point predictions. The first axis is the sample axis; remaining axes (e.g. horizons, nodes) are averaged.
floor (float, optional) – Minimum difficulty value before normalisation. If None (default) it is set to quantile(mean|pred|, floor_quantile).
floor_quantile (float) – Quantile of mean|pred| used to set the floor when floor is None. Prevents near-zero predictions from creating infinite normalised residuals.
ref_floor (float, optional) – Pre-computed floor from calibration data. When provided, overrides the floor computed from predictions, avoiding test-set leakage.
ref_median (float, optional) – Pre-computed median difficulty from calibration data. When provided, overrides the median computed from predictions, avoiding test-set leakage.

Returns:

Per-sample difficulty scores, median ≈ 1.0 (w.r.t. reference).

Return type:

ndarray, shape (n,)

Difficulty_Reference_Stats

epilearn.utils.uncertainty.difficulty_reference_stats(reference_data: ndarray, *, floor_quantile: float = 0.1, is_residual: bool = False) → tuple[float, float]

Compute difficulty normalisation constants from calibration data.

These constants (ref_floor, ref_median) should be passed to difficulty_from_predictions() when scoring test samples so that the floor and median are derived from the calibration set rather than the test batch — preventing information leakage.

Parameters:

reference_data (ndarray) – Either calibration predictions (shape (n, ...)) or calibration residuals (1-D |pred - target|). Set is_residual accordingly.
floor_quantile (float) – Quantile used for the floor.
is_residual (bool) – If True, reference_data are absolute residuals (already positive, 1-D). If False, per-sample mean absolute value is computed first.

Returns:

(ref_floor, ref_median) – Values to pass to difficulty_from_predictions(…, ref_floor=ref_floor, ref_median=ref_median).

Return type:

tuple[float, float]

Evaluate_ACI_From_Saved

epilearn.utils.uncertainty.evaluate_aci_from_saved(results_dir: str, model_name: str | None = None, timestamp: str | None = None, gamma: float = 0.005) → dict

Run ACI + locally-weighted ACI on saved benchmark predictions.

Always denormalizes predictions/targets to match the original-scale validation residuals.

Parameters:

results_dir (str) – Path to benchmark results directory.
model_name (str, optional) – Single model to evaluate. If None, evaluates all models.
timestamp (str, optional) – Specific benchmark run timestamp.
gamma (float) – ACI learning rate.

Returns:

{model_name: {'aci': {...}, 'static': {...}, ...}}.

Return type:

dict

Compartmental_Models

New in 0.1.0. epilearn.utils.compartmental_models provides deterministic ODE models that need no data at all: SIRModel / SIRSModel over ('S', 'I', 'R') (SIRSModel adds omega, waning immunity R → S), SEIRModel over ('S', 'E', 'I', 'R') (sigma = 1/latent period) and SEIRVIModel over ('S', 'E', 'I', 'R', 'V', 'Q') (V vaccinated, Q isolated). Constructors are in the class signatures below; every model exposes

compartments — the compartment names, in state-vector order;
parameters — the rate dictionary;
step(state, t, dt, method='rk4'|'euler', external_inputs=None, parameter_overrides=None) — one integration step;
simulate(initial_state, steps, dt=1.0, method='rk4', parameter_schedule=None, input_schedule=None) → {'time', 'trajectory', 'compartments'} where trajectory has shape (steps + 1, n_compartments).

initial_state is in counts, not fractions, and its length must equal len(model.compartments) or validate_state raises ValueError; project_state clamps every compartment at 0 after each step.

parameter_schedule overrides rate parameters (beta, gamma, …) while input_schedule supplies external inputs (force_of_infection, and for SEIRVIModel also vaccination_rate / isolation_rate). Both accept a callable f(step_idx, t, state) -> dict | None, a {step_idx: dict} mapping, or a sequence indexed by step; None leaves the base values in place. Utilities explains which channel an intervention belongs to (with worked lockdown and vaccination examples, and when to prefer method='euler'); Simulation drives these models at population, individual and regional level.

CompartmentalModel

class epilearn.utils.compartmental_models.CompartmentalModel(compartments, parameters, metadata=None, name=None): Base class for deterministic compartmental epidemic models.

SIRModel

class epilearn.utils.compartmental_models.SIRModel(beta, gamma, mu=0.0, birth_rate=0.0): Classical SIR model with optional demography.

SEIRModel

class epilearn.utils.compartmental_models.SEIRModel(beta, gamma, sigma, mu=0.0, birth_rate=0.0): SEIR model (Susceptible-Exposed-Infectious-Recovered).

SIRSModel

class epilearn.utils.compartmental_models.SIRSModel(beta, gamma, omega, mu=0.0, birth_rate=0.0): SIRS model with waning immunity from R back to S.

SEIRVIModel

class epilearn.utils.compartmental_models.SEIRVIModel(beta, gamma, sigma, mu=0.0, birth_rate=0.0, vaccine_efficacy=0.8, isolation_efficacy=0.9)

SEIR model with Vaccination and Isolation interventions.

Compartments: S (Susceptible), E (Exposed), I (Infectious), R (Recovered), V (Vaccinated), Q (Isolated/Quarantined)

Intervention effects: - Vaccination: Moves susceptibles to V with rate dependent on vaccination policy - Isolation: Moves infectious individuals to Q, reducing transmission

Parameters:

beta (float) – Base transmission rate.
gamma (float) – Recovery rate (1/infectious_period).
sigma (float) – Incubation rate (1/latent_period).
mu (float) – Natural death rate (default 0).
birth_rate (float) – Birth rate (default 0).
vaccine_efficacy (float) – Vaccine efficacy in preventing infection (0-1, default 0.8).
isolation_efficacy (float) – Reduction in transmission from isolated individuals (0-1, default 0.9).
schedules) (Intervention inputs (via external_inputs or)
vaccination_rate (-)
vaccination_delay (-)
isolation_rate (-)
isolation_delay (-)

create_multi_scenario_interventions(lookback: int, horizon: int, n_scenarios: int, static_vacc_rate: float, static_vacc_delay: int, static_isol_rate: float, static_isol_delay: int, dynamic_vacc_rates: list, dynamic_vacc_delays: list, dynamic_isol_rates: list, dynamic_isol_delays: list)

Create intervention scenarios tensor ensuring static/dynamic constraints.

Static interventions (t + delay < L) are applied identically. Dynamic interventions (t + delay >= L) vary per scenario.

Returns:: Tensor (L+H, n_scenarios, 2, 2)
Return type:: intervention_scenarios

generate_multi_scenario_dataset(n_samples: int, lookback: int, horizon: int, n_scenarios: int, static_vacc_rate_range: tuple = (0.0, 0.02), static_vacc_delay_range: tuple = (10, 30), static_isol_rate_range: tuple = (0.0, 0.1), static_isol_delay_range: tuple = (5, 20), dynamic_vacc_rate_range: tuple = (0.0, 0.08), dynamic_isol_rate_range: tuple = (0.0, 0.3), initial_infected_frac_range: tuple = (0.001, 0.02), population: float = 1000000.0, process_noise: float | None = None, seed: int = 42)

Generate a dataset using generate_multi_scenario_samples.

Each sample generates: - X: (L, n_scenarios, N+4) - shared history with scenario-specific interventions - Y: (H, n_scenarios, N) - different futures per scenario

Parameters:

n_samples (int) – Number of samples to generate.
lookback (int) – Historical window size L.
horizon (int) – Future projection size H.
n_scenarios (int) – Number of scenarios per sample.
static_vacc_rate_range (tuple) – Range for static vaccination rate (min, max).
static_vacc_delay_range (tuple) – Range for static vaccination delay (min, max).
static_isol_rate_range (tuple) – Range for static isolation rate (min, max).
static_isol_delay_range (tuple) – Range for static isolation delay (min, max).
dynamic_vacc_rate_range (tuple) – Range for dynamic vaccination rate (min, max).
dynamic_isol_rate_range (tuple) – Range for dynamic isolation rate (min, max).
initial_infected_frac_range (tuple) – Range for initial infected fraction (min, max).
population (float) – Total population size.
process_noise (float, optional) – Standard deviation of process noise. If None, no noise is added. Noise is applied identically during historical period and independently per scenario during horizon period.
seed (int) – Random seed for reproducibility.

Returns:

X (torch.Tensor) – Shape (n_samples, L, n_scenarios, N+4)
Y (torch.Tensor) – Shape (n_samples, H, n_scenarios, N)
metadata (list) – List of sample metadata dictionaries.

generate_multi_scenario_samples(initial_state, lookback: int, horizon: int, intervention_scenarios: Tensor, dt: float = 1.0, method: str = 'rk4', process_noise: float | None = None, seed: int | None = None)

Generate simulation samples for multiple scenarios with shared history.

This function generates: - A single historical trajectory that is SHARED across all scenarios - Multiple future projections, one per scenario

Key concept: - Static interventions: t + delay < L (take effect within historical period)

Must be IDENTICAL across all scenarios.

Dynamic interventions: t + delay >= L (take effect in horizon period) Can VARY across scenarios.

Parameters:

initial_state (array-like) – Initial state for compartments [S, E, I, R, V, Q].
lookback (int) – Lookback window size L (historical period).
horizon (int) – Horizon size H (future projection period).
intervention_scenarios (torch.Tensor) –
Intervention scenarios of shape (L+H, n_scenarios, 2, 2): - Dim 0: timesteps (L+H total) - Dim 1: scenarios - Dim 2: [vaccination, isolation] - Dim 3: [value, delay]

Static entries (where t + delay < L) MUST be identical across scenarios. Only dynamic entries (where t + delay >= L) can differ.
dt (float) – Time step size (default 1.0).
method (str) – Integration method (‘rk4’ or ‘euler’).
process_noise (float, optional) – Standard deviation of process noise (applied identically across scenarios during historical period, independently during horizon).
seed (int, optional) – Random seed for reproducibility.

Returns:

‘historical’: Tensor of shape (L, n_scenarios, N+4) Compartment values (N) are IDENTICAL across scenarios. Intervention columns (+4) may differ for dynamic interventions.
’future’: Tensor of shape (H, n_scenarios, N) Different future trajectories per scenario.
’time’: Time axis for entire simulation
’compartments’: Compartment names
’n_scenarios’: Number of scenarios
’static_mask’: Boolean mask indicating static intervention entries
’dynamic_mask’: Boolean mask indicating dynamic intervention entries

Return type:

dict

generate_scenario_samples(initial_state, lookback: int, horizon: int, intervention_trajectories: Tensor, dt: float = 1.0, method: str = 'rk4', process_noise: float | None = None, seed: int | None = None)

Generate simulation samples for scenario modeling.

This function generates both historical trajectories (with interventions applied according to their delays) and future projections.

Parameters:

initial_state (array-like) – Initial state for compartments [S, E, I, R, V, Q].
lookback (int) – Lookback window size L (historical period).
horizon (int) – Horizon size H (future projection period).
intervention_trajectories (torch.Tensor) –
Intervention trajectories of shape (L+H, 2, 2): - First dimension: timesteps - Second dimension: [vaccination, isolation] - Third dimension: [value, delay] Example: intervention_trajectories[t, 0, :] = [vaccination_rate, vaccination_delay]

intervention_trajectories[t, 1, :] = [isolation_rate, isolation_delay]
dt (float) – Time step size (default 1.0).
method (str) – Integration method (‘rk4’ or ‘euler’).
process_noise (float, optional) – Standard deviation of process noise.
seed (int, optional) – Random seed for reproducibility.

Returns:

‘historical’: Tensor of shape (L, N+4) where N is number of compartments (6), and +4 is for [vaccination_value, vaccination_delay, isolation_value, isolation_delay]
’future’: Tensor of shape (H, N) containing future compartment trajectories
’time’: Time axis for entire simulation
’compartments’: Compartment names

Return type:

dict

Simulation

Deprecated since version 0.1.0: epilearn.utils.simulation.Time_geo has been removed with no replacement. Use simulate_spatiotemporal_individual() for individual-level dynamics on a contact graph, or simulate_spatiotemporal_regions() for metapopulation dynamics with mobility-driven flows.

0.1.0 exposes three simulators, all driven by a CompartmentalModel and all returning a dict of tensors whose keys are listed per function below:

simulate_temporal_epidemic — one population; trajectory (steps+1, n_compartments). It is model.simulate() plus optional process_noise (scalar or per-compartment, with seed), added after each step and re-projected to be non-negative.
simulate_spatiotemporal_individual — individuals on a contact graph; trajectory (steps+1, n_nodes) of compartment indices.
simulate_spatiotemporal_regions — regions with mobility-driven flows; trajectory (steps+1, n_regions, n_compartments), signed net flows in dynamic_graph (directed_flow is the non-negative form) and per-region \(R_t = (\beta_t / \gamma)\,(S / N)\) in effective_reproduction_number.

Both spatiotemporal simulators take their graph as a static adjacency matrix / NetworkX graph, a (steps+1, N, N) tensor or a callable f(t, step_idx) -> adjacency, and both return a node_features tensor already shaped for a SpatialTemporal model; initial conditions come from create_initial_conditions_individual / create_initial_conditions_region, with create_regional_forcing_params for per-region seasonal forcing.

Simulation is where these are taught: it runs all three end to end, feeds the output into a Dataset, explains the Gravity_model connection-strength convention and which flow tensor to hand a model, and documents the two traps in the helpers above (initial_compartment_fractions is renormalized to sum to 1; n_initial_infected must stay below n_regions).

Simulate_Temporal_Epidemic

epilearn.utils.simulation.simulate_temporal_epidemic(model: CompartmentalModel, initial_state: Sequence[float] | Tensor | ndarray, steps: int, dt: float = 1.0, parameter_schedule: Mapping[int, Dict[str, float]] | Sequence | None = None, input_schedule: Mapping[int, Dict[str, float]] | Sequence | None = None, process_noise: float | Sequence[float] | Tensor | None = None, method: str = 'rk4', seed: int | None = None): Run a temporal (population-level) simulation for the provided compartmental model.

Simulate_Spatiotemporal_Individual

Simulate an individual-level compartmental process on a contact graph. Nodes follow the provided model (SIR/SEIR/SIRS) with infection pressure driven by neighbors. Returns per-step node features (one-hot compartment indicators) and a dynamic graph tensor of shape (time, N, N).

Parameters:

model (CompartmentalModel) – The compartmental model (e.g., SIRModel, SEIRModel, SIRSModel).
contact_graph (Union[torch.Tensor, np.ndarray, nx.Graph, Callable]) –
Contact graph specification. Can be: - Static: Binary adjacency matrix (num_nodes, num_nodes) or NetworkX graph - Time-varying: Tensor of shape (steps+1, num_nodes, num_nodes) - Dynamic: Callable function f(t, step_idx) -> adjacency_matrix that returns

the contact graph at each time step
initial_states (Union[Sequence[int], Sequence[str], torch.Tensor, np.ndarray, Dict[str, Iterable[int]]]) – Initial compartment states for each individual.
steps (int) – Number of simulation steps.
dt (float) – Time step size (default 1.0).
stochastic (bool) – Whether to use stochastic transitions (default True).
seed (int, optional) – Random seed for reproducibility.

Returns:

Simulation results with keys: - ‘time’: Time axis (steps+1,) - ‘trajectory’: Individual state history (steps+1, num_nodes) with compartment indices - ‘counts’: Compartment counts over time (steps+1, num_compartments) - ‘compartments’: Compartment names - ‘contact_graph’: Static or initial contact graph - ‘dynamic_graph’: Time-varying contact graphs (steps+1, num_nodes, num_nodes) - ‘node_features’: One-hot encoded compartment states (steps+1, num_nodes, num_compartments)

Return type:

dict

Simulate_Spatiotemporal_Regions

epilearn.utils.simulation.simulate_spatiotemporal_regions(model: CompartmentalModel, region_states: Tensor | ndarray, adjacency_graph: Tensor | ndarray | Graph | Callable, steps: int, dt: float = 1.0, forcing_params: List[Dict[str, float]] | None = None, forcing_parameter: str = 'beta', method: str = 'euler', gravity_model: Gravity_model | None = None, travel_rate: float = 1.0) → Dict[str, Tensor]

Simulate region-level epidemics with mobility-driven population flows.

This function implements a two-step process: 1. Internal dynamics update (with optional multi-frequency forcing) 2. Population flow based on mobility model (diffusive or gravity-based)

The flows are applied AFTER internal dynamics. Flow models are configured via the gravity_model parameter, which supports both diffusive and gravity-based flows.

Parameters:

model (CompartmentalModel) – Compartmental epidemic model (must be SIRS-like for infinite stability).
region_states (Union[torch.Tensor, np.ndarray]) – Initial states for each region as counts (not fractions), shape (num_regions, num_compartments). For SIRS: [S, I, R] counts for each region.
adjacency_graph (Union[torch.Tensor, np.ndarray, nx.Graph, Callable]) – Connectivity/adjacency graph specification where higher values = stronger connections. Edge weights represent connection strength (travel rates, similarity, interaction frequency). Can be: - Static: Weighted matrix (num_regions, num_regions) or NetworkX graph - Time-varying: Tensor of shape (steps+1, num_regions, num_regions) - Dynamic: Callable f(t, step_idx) -> adjacency_matrix
steps (int) – Number of simulation steps.
dt (float) – Time step size (default 1.0 for daily updates).
forcing_params (List[Dict[str, float]], optional) – Region-specific forcing parameters with keys: ‘amp1’, ‘phase1’, ‘period1’, ‘amp2’, ‘phase2’, ‘period2’, ‘amp3’, ‘phase3’, ‘period3’.
forcing_parameter (str) – Name of the parameter to apply forcing to (default ‘beta’).
method (str) – Integration method (‘euler’ recommended for stability, default ‘euler’).
gravity_model (Gravity_model, optional) – Mobility model instance. If None, defaults to diffusive model. - Diffusive: Gravity_model(rho=0, theta=0, delta=1.0, normalize=False) - Gravity: Gravity_model(rho=1.0, theta=1.0, delta=100.0, normalize=True)
travel_rate (float) – Global mobility scaling factor applied to every edge-wise flow (default 1.0). Use this to tune the overall level of inter-regional travel without modifying the adjacency weights.

Returns:

Simulation results with keys: - ‘time’: Time axis (steps+1,) - ‘trajectory’: Regional state history (steps+1, num_regions, num_compartments) - ‘compartments’: Compartment names - ‘adjacency’: Static adjacency or provided tensor (for backward compatibility) - ‘adjacency_history’: Adjacency tensor for every timestep (steps+1, num_regions, num_regions) - ‘dynamic_graph’: Signed flow graphs (steps+1, num_regions, num_regions) - ‘directed_flow’: Non-negative directed flows derived from dynamic_graph - ‘parameter_history’: Value of forcing_parameter used per region and timestep - ‘effective_reproduction_number’: R_t estimates when applicable - ‘counts’: Aggregate compartment counts summed across regions - ‘regional_totals’: Population of each region at every timestep - ‘node_features’: Alias for trajectory (kept for convenience)

Return type:

dict

Notes

Flow Models:

Both models use unified connectivity semantics (higher = stronger connection):

Diffusive (default):: F(i,j) = connectivity[i,j] × (N_i - N_j)
Gravity:: F(i,j) = [N_i^ρ × N_j^θ × exp((connectivity[i,j]-1)/δ)] × (N_i - N_j)/(N_i + N_j)

Both ensure population conservation through antisymmetric flows.

Connectivity Normalization:

Normalize edge weights to [0, 1] for best results: >>> adjacency_normalized = adjacency / adjacency.max()

Examples

Diffusive flow (default):

>>> result = simulate_spatiotemporal_regions(
...     model, states, adjacency, steps=100)

Diffusive flow (explicit):

>>> diffusive = Gravity_model(rho=0, theta=0, delta=1.0, normalize=False)
>>> result = simulate_spatiotemporal_regions(
...     model, states, adjacency, steps=100, gravity_model=diffusive)

Gravity model:

>>> # Normalize connectivity (edge weights)
>>> adjacency_norm = adjacency / adjacency.max()
>>> gravity = Gravity_model(rho=1.0, theta=1.0, delta=0.5, normalize=True)
>>> result = simulate_spatiotemporal_regions(
...     model, states, adjacency_norm, steps=100, gravity_model=gravity)

Create_Initial_Conditions_Individual

epilearn.utils.simulation.create_initial_conditions_individual(model, num_individuals=100, p_edge=0.01, initial_compartment_fractions: dict = {'I': 0.01}, seed=42)

Create initial conditions for individual-level simulations.

Parameters:

model (CompartmentalModel) – The compartmental model (e.g., SIRModel, SIRSModel, SEIRModel) that defines the disease dynamics and compartment structure
num_individuals (int) – Number of individuals in the network
p_edge (float) – Edge probability for Erdos-Renyi random graph (contact network)
initial_compartment_fractions (dict) – Dictionary mapping compartment names to initial fractions. Example: {‘I’: 0.01} means 1% start infected, rest susceptible Example: {‘S’: 0.9, ‘E’: 0.05, ‘I’: 0.05} for SEIR model
seed (int) – Random seed for reproducibility

Returns:

node_states: torch.Tensor of shape (num_individuals,) containing compartment indices
adjacency: torch.Tensor of shape (num_individuals, num_individuals) contact graph

Return type:

tuple

Create_Initial_Conditions_Region

epilearn.utils.simulation.create_initial_conditions_region(model, n_regions=100, p_edge=0.01, pop_range=(800, 20000), n_initial_infected=20, initial_infected_size=100, initial_compartment_fractions=None, ensure_connected=True, seed=42)

Create initial conditions for spatiotemporal epidemic simulation based on a compartmental model.

Parameters:

model (CompartmentalModel) – The compartmental model (e.g., SIRModel, SIRSModel, SEIRModel) that defines the disease dynamics and compartment structure
n_regions (int) – Number of regions in the network
p_edge (float) – Edge probability for Erdos-Renyi random graph
pop_range (tuple of (int, int)) – Range for random population initialization (min_pop, max_pop)
n_initial_infected (int) – Number of regions to seed with initial infections
initial_infected_size (int) – Number of individuals initially infected in each seeded region
initial_compartment_fractions (dict or None) – Optional dictionary mapping compartment names to initial fractions. If None, all individuals start in ‘S’ (susceptible). Example: {‘S’: 0.9, ‘E’: 0.05, ‘I’: 0.05} for SEIR model
ensure_connected (bool) – If True, extract largest connected component
seed (int) – Random seed for reproducibility

Returns:

Dictionary containing: - ‘adjacency’: torch.Tensor of shape (n_regions, n_regions) - ‘region_states’: torch.Tensor of shape (n_regions, n_compartments) - ‘graph’: networkx.Graph object - ‘n_regions’: int (actual number of regions after connectivity check) - ‘n_edges’: int - ‘infected_nodes’: numpy.ndarray of initially infected region indices - ‘model’: the compartmental model used - ‘compartment_names’: list of compartment names - ‘total_population’: total population across all regions

Return type:

dict

Create_Regional_Forcing_Params

epilearn.utils.simulation.create_regional_forcing_params(num_regions: int, base_amplitudes: Sequence[float] = (0.2, 0.15, 0.1), periods: Sequence[float] = (30.0, 47.0, 73.0), amplitude_noise: float = 0.1, seed: int | None = None) → List[Dict[str, List[float] | float]]

Create region-specific multi-frequency forcing parameters.

Each region gets slightly different amplitudes and random phases to desynchronize dynamics across regions.

Parameters:

num_regions (int) – Number of regions.
base_amplitudes (Sequence[float]) – Base amplitudes for each frequency component.
periods (Sequence[float]) – Periods for each frequency component (in time units).
amplitude_noise (float) – Relative noise level for amplitude perturbations (default 0.1 = 10%).
seed (int, optional) – Random seed for reproducibility.

Returns:

List of dictionaries, one per region, with keys: ‘amp1’, ‘phase1’, ‘period1’, ‘amp2’, ‘phase2’, ‘period2’, ‘amp3’, ‘phase3’, ‘period3’.

Return type:

List[Dict[str, Union[List[float], float]]]

Gravity_Model

class epilearn.utils.simulation.Gravity_model(rho: float = 0.0, theta: float = 0.0, delta: float = 1.0, normalize: bool = False)

Elegant gravity model for human mobility in epidemic simulations.

Computes flow between regions based on population attraction and connectivity strength:

\[F_{ij} = N_i^{\rho} \cdot N_j^{\theta} \cdot \exp((w_{ij} - 1) / \delta)\]

where higher w_{ij} indicates stronger connection between regions.

Parameters:

rho (float) – Source population exponent (typically 0.5-1.0).
theta (float) – Target population exponent (typically 0.5-1.0).
delta (float) – Connectivity decay parameter. Controls sensitivity to connectivity variations. Recommended: 0.2-2.0 for normalized connectivity [0, 1].
normalize (bool) – If True, normalize flows by total population (default True).

Notes

Unified Connectivity Semantics:

Connectivity values represent connection strength where: - Higher values = Stronger connection = More flow - Lower values = Weaker connection = Less flow - 0 = No connection

This is consistent across both diffusive and gravity models.

It is recommended to normalize your connectivity matrix to [0, 1] range:

>>> # Normalize connectivity matrix (edge weights)
>>> max_conn = connectivity_matrix.max()
>>> normalized_connectivity = connectivity_matrix / max_conn
>>>
>>> # Choose delta based on desired spatial spread
>>> # Sharp decay (local): delta = 0.2-0.5
>>> # Moderate (regional): delta = 0.5-1.0
>>> # Gradual (long-range): delta = 1.0-2.0

Examples

>>> # Diffusive flow (special case: rho=0, theta=0)
>>> diffusive = Gravity_model(rho=0, theta=0, delta=1.0, normalize=False)
>>>
>>> # Gravity model with normalized connectivity
>>> gravity = Gravity_model(rho=1.0, theta=1.0, delta=0.5, normalize=True)
>>>
>>> # Using with normalized edge weights
>>> weights_normalized = adjacency / adjacency.max()
>>> result = simulate_spatiotemporal_regions(
...     model, states, weights_normalized, steps=100, gravity_model=gravity)

compute_flow(pop_i: float, pop_j: float, connectivity: float) → float

Compute flow between two regions.

Parameters:

pop_i (float) – Population sizes of regions i and j.
pop_j (float) – Population sizes of regions i and j.
connectivity (float) –
Connection strength between regions (recommended range: 0-1). Higher values = stronger connection = more flow. - 0: No connection - 1: Maximum connection strength

Interpretation is unified across all models: - Diffusive: Direct multiplier on population flow - Gravity: Exponential enhancement of attraction

Returns:

Flow magnitude or edge weight for further computation.

Return type:

float

Notes

Unified Connectivity Semantics:

Connectivity always represents connection strength where higher = stronger:

Diffusive model (rho=0, theta=0):

Formula: F = connectivity × (N_i - N_j)
Range: [0.0, 1.0]
Interpretation: Fraction of population difference that flows per timestep
Example: 0.1 = 10% of population difference travels per day

Gravity model (rho>0, theta>0):

Formula: F = [N_i^ρ × N_j^θ × exp((connectivity-1)/δ)] × (N_i - N_j)/(N_i + N_j)
Range: [0.0, 1.0] normalized
connectivity=1.0 gives baseline gravity attraction
connectivity>1.0 enhances flow (if using unnormalized weights)
connectivity<1.0 reduces flow
Delta parameter controls sensitivity: * Small delta (0.2-0.5): Sharp response to connectivity differences * Large delta (1.0-2.0): Gradual response to connectivity differences

Normalization Strategy:

For any connectivity matrix (edge weights, similarity scores, etc.): 1. Find max value: max_conn = connectivity_matrix.max() 2. Normalize: connectivity_matrix = connectivity_matrix / max_conn 3. Choose delta based on desired sensitivity:

Sharp (local spread): delta = 0.2-0.5

Moderate: delta = 0.5-1.0

Gradual (long-range): delta = 1.0-2.0

compute_mobility_matrix(populations: Tensor, connectivity_matrix: Tensor) → Tensor

Vectorized computation of full mobility matrix.

Parameters:

populations (torch.Tensor) – Population of each region, shape (num_regions,).
connectivity_matrix (torch.Tensor) – Connectivity matrix, shape (num_regions, num_regions). Higher values = stronger connection = more flow (unified semantics).

Returns:

Mobility flow matrix, shape (num_regions, num_regions).

Return type:

torch.Tensor

compute_net_flow(pop_i: float, pop_j: float, connectivity: float) → float

Compute net directional flow from region i to region j.

Parameters:

pop_i (float) – Population sizes of regions i and j.
pop_j (float) – Population sizes of regions i and j.
connectivity (float) – Connection strength between regions. Higher values = stronger connection = more flow (unified across all models).

Returns:

Net flow (positive = i→j, negative = j→i).

Return type:

float

Get_Random_Graph

epilearn.utils.simulation.get_random_graph(num_nodes=None, connect_prob=None, block_sizes=None, num_edges=None, graph_type='erdos_renyi')

Generates a random static graph using one of the supported graph types: Erdos-Renyi, Stochastic Blockmodel, or Barabasi-Albert.

Parameters:

num_nodes (int) – Number of nodes in the graph.
connect_prob (float, optional) – Probability of edge creation (for Erdos-Renyi and Stochastic Blockmodel graphs).
block_sizes (list of int, optional) – Sizes of blocks (for Stochastic Blockmodel graph).
num_edges (int, optional) – Number of edges (for Barabasi-Albert graph).
graph_type (str) – Type of graph to generate. Options are ‘erdos_renyi’, ‘stochastic_blockmodel’, ‘barabasi_albert’. Default is ‘erdos_renyi’.

Returns:

Adjacency matrix of the generated graph.

Return type:

torch.Tensor

Get_Graph_From_Features

epilearn.utils.simulation.get_graph_from_features(features, adj=None, G=1)

Generate a graph from node features using cosine similarity.

This function generates a graph where each edge weight is computed based on the cosine similarity between the feature vectors of the connected nodes. If an adjacency matrix is provided, the cosine similarity is adjusted by the corresponding entry in the adjacency matrix.

Parameters:

features (torch.Tensor) – A tensor of shape (num_nodes, feat_dim) where num_nodes is the number of nodes and feat_dim is the dimensionality of the feature vectors.
adj (torch.Tensor, optional) – A tensor of shape (num_nodes, num_nodes) representing the adjacency matrix, where adj[i, j] denotes the distance or weight between node i and node j. If None, the cosine similarity is used directly as the edge weight. Default is None.

Returns:

A tensor of shape (num_nodes, num_nodes) representing the generated graph’s adjacency matrix, where each entry [i, j] contains the adjusted cosine similarity between nodes i and j.

Return type:

torch.Tensor

Transformation

Note

Compose.__call__ returns a tuple (data, process_history) in 0.1.0, so data = transformation(data) silently binds the tuple. process_history holds the fitted normalization constants — feat_mean / feat_std and target_mean / target_std — replacing the old Compose.feat_mean / Compose.feat_std attributes; it is also readable as transformation.process_history after the call.

from epilearn.data import Dataset
from epilearn.utils import transforms

dataset = Dataset(); dataset.load_toy_dataset()
transformation = transforms.Compose({
    "features": [transforms.normalize_feat()],
    "target": [transforms.normalize_target()],
    "graph": [transforms.normalize_adj()]})
data, process_history = transformation({"features": dataset.x,
                                       "target": dataset.y,
                                       "graph": dataset.graph})
print(sorted(process_history))
# ['feat_mean', 'feat_std', 'target_mean', 'target_std']

Usually the Compose is handed to the dataset instead (dataset.set_transforms(transformation, apply_now=True)) and applied per fold by the task. Keep target_mean / target_std: without them metrics and interval widths stay in normalized units (Model&Dataset Customization).

Compose

class epilearn.utils.transforms.Compose(transforms, device='cpu')

Composes several transforms together. This transform does not support torchscript. Please, see the note below.

Parameters:: transforms (list of Transform objects) – list of transforms to compose.

Example

>>> transforms.Compose([
>>>     transforms.CenterCrop(10),
>>>     transforms.ToTensor(),
>>> ])

Note

In order to script the transformations, please use torch.nn.Sequential as below.

>>> transforms = torch.nn.Sequential(
>>>     transforms.CenterCrop(10),
>>>     transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
>>> )
>>> scripted_transforms = torch.jit.script(transforms)

Make sure to use only scriptable transformations, i.e. that work with torch.Tensor, does not require lambda functions or PIL.Image.

Normalize_Feat

class epilearn.utils.transforms.normalize_feat(dim=1)

A normalization module for feature standardization in PyTorch. This module adjusts features to have zero mean and unit variance along specified dimensions, handling 3D and 4D tensors.

Parameters:: dim (int, optional) – The dimension over which to calculate the mean and standard deviation for normalization. Default: 1.

forward(X, device='cpu')

Forward pass of the normalization module that normalizes a given input tensor X.

Uses GLOBAL normalization (single mean/std across all dimensions) to properly handle time series data where temporal trends cause different distributions across train/val/test splits.

Parameters:

X (torch.Tensor) – The input tensor to be normalized. Can be 2D, 3D, or 4D tensor.
device (str, optional) – The device to which the normalized tensor is transferred. Default: ‘cpu’.

Returns:

The normalized tensor, adjusted to have zero mean and unit variance globally, and transferred to the specified device.

Return type:

torch.Tensor

Normalize_Target

class epilearn.utils.transforms.normalize_target(dim=1)

A normalization module for feature standardization in PyTorch. This module adjusts features to have zero mean and unit variance along specified dimensions, handling 3D and 4D tensors.

Parameters:: dim (int, optional) – The dimension over which to calculate the mean and standard deviation for normalization. Default: 1.

forward(X, device='cpu')

Forward pass of the normalization module that normalizes a given input tensor X.

Uses GLOBAL normalization (single mean/std across all dimensions) to properly handle time series data where temporal trends cause different distributions across train/val/test splits.

Parameters:

X (torch.Tensor) – The input tensor to be normalized. Can be 2D, 3D, or 4D tensor.
device (str, optional) – The device to which the normalized tensor is transferred. Default: ‘cpu’.

Returns:

The normalized tensor, adjusted to have zero mean and unit variance globally, and transferred to the specified device.

Return type:

torch.Tensor

Normalize_Adj

class epilearn.utils.transforms.normalize_adj(dim=0)

A PyTorch module for normalizing adjacency matrices to facilitate operations in graph neural networks. The normalization adjusts adjacency matrices to account for node degrees, enhancing the propagation of features through the network. This class can handle both batched and single adjacency matrices.

Parameters:: dim (int, optional) – The dimension over which to perform normalization (not utilized in current implementation). Default: 0.

forward(Adj, device='cpu')

Forward pass of the normalize_adj module that computes a degree-normalized adjacency matrix from the input adjacency matrix ‘Adj’.

Parameters:

Adj (torch.Tensor or np.array) – The input adjacency matrix, which can be a 2D matrix for a single graph or a 3D tensor for batch processing.
device (str, optional) – The device to which the normalized adjacency matrix is transferred. Default: ‘cpu’.

Returns:

The degree-normalized adjacency matrix, transferred to the specified device.

Return type:

torch.Tensor

Convert_To_Frequency

class epilearn.utils.transforms.convert_to_frequency(ftype='fft', n_fft=8)

A PyTorch module for transforming time-domain data into frequency-domain representations using either FFT (Fast Fourier Transform) or STFT (Short Time Fourier Transform). This module is configurable to handle different lengths of FFT and hop sizes for STFT, making it flexible for various signal processing tasks.

Parameters:

ftype (str, optional) – The type of frequency transformation to perform, either ‘fft’ for Fast Fourier Transform or ‘stft’ for Short Time Fourier Transform. Default: “fft”.
n_fft (int, optional) – The window size for the FFT or STFT. Default: 8.

forward(data, **kwarg)

Applies the configured frequency transformation (FFT or STFT) to the input data.

Parameters:

data (torch.Tensor) – The input data tensor, expected to be a time-domain signal. It can be a 3D tensor (batch, channels, time) for ‘fft’ or a 4D tensor (batch, nodes, time, features) for ‘stft’.
**kwargs (dict) – Additional keyword arguments, such as ‘device’ for specifying the computation device.

Returns:

The frequency-domain representation of the input data. The output format depends on the transformation type (‘fft’ returns real parts of the FFT, ‘stft’ returns the magnitude of the STFT as a 5D tensor).

Return type:

torch.Tensor

Add_Time_Embedding

class epilearn.utils.transforms.add_time_embedding(embedding_dim=13, fourier=False)

A PyTorch module that appends time-based embeddings to each feature vector in the dataset. It can generate embeddings using sinusoidal functions, optionally using a Fourier transform approach, enhancing the temporal aspects of data for tasks such as time series forecasting or sequence modeling.

Parameters:

embedding_dim (int, optional) – The dimensionality of the time embeddings. Default: 13.
fourier (bool, optional) – Specifies whether to use a Fourier transform-based approach for time embedding. Default: False.

forward(data, **kwarg)

Generates and appends time-based embeddings to the input data tensor.

Parameters:

data (torch.Tensor) – The input data tensor, which can vary in dimensions depending on the application (e.g., batch, nodes, time).
**kwargs (dict) – Additional keyword arguments such as ‘device’ for specifying the computation device.

Returns:

The input data tensor augmented with time-based embeddings. The resulting tensor includes an additional dimension for the embeddings, concatenated to the last dimension of the input data tensor.

Return type:

torch.Tensor

Learnable_Time_Embedding

class epilearn.utils.transforms.learnable_time_embedding(timesteps=13, embedding_dim=13)

A PyTorch module that appends learnable time embeddings to the input data, facilitating the capture of temporal dependencies in models that process sequences or time-series data. The embedding is learned during the training process, allowing it to adapt to specific temporal patterns observed in the dataset.

Parameters:

timesteps (int) – The total number of timesteps in the data sequence, which defines the number of unique time indices.
embedding_dim (int) – The dimensionality of each time embedding vector.

forward(data, **kwarg)

Appends learnable time-based embeddings to the input data tensor. The embeddings are added to each timestep, augmenting the feature dimensions of the data.

Parameters:

data (torch.Tensor) – The input data tensor, typically including dimensions for batch, nodes, and time.
**kwargs (dict) – Additional keyword arguments such as ‘device’ for specifying the computation device.

Returns:

The input data tensor augmented with time-based embeddings, expanding the last dimension to include the embeddings.

Return type:

torch.Tensor

Seasonality_And_Trend_Decompose

class epilearn.utils.transforms.seasonality_and_trend_decompose(decompose_type='dynamic', moving_avg=25, kernel_size=[4, 8, 12])

A PyTorch module designed to decompose time-series data into seasonality and trend components. It supports both dynamic and static decomposition methods. Dynamic decomposition leverages a Fourier transform approach for seasonality and convolutional filters for trend extraction. Static decomposition utilizes a moving average approach.

Parameters:

decompose_type (str, optional) – The type of decomposition to perform. “dynamic” for Fourier and convolutional methods, “static” for moving average based decomposition. Default: “dynamic”.
moving_avg (int, optional) – The window size for the moving average in static decomposition. Default: 25.
kernel_size (list of int, optional) – List of kernel sizes for convolutional filters in dynamic trend decomposition. Default: [4, 8, 12].

forward(data, **kwarg)

Decomposes the input data tensor into seasonality and trend components. The method of decomposition (dynamic or static) impacts the models and techniques used.

Parameters:

data (torch.Tensor) – The input data tensor, typically including dimensions for batch, nodes, and time.
**kwargs (dict) – Additional keyword arguments such as ‘device’ for specifying the computation device.

Returns:

A list containing the seasonality and trend components of the input data, each as a separate tensor.

Return type:

list of torch.Tensor

Calculate_DTW_Matrix

class epilearn.utils.transforms.calculate_dtw_matrix(dataset_name)

A PyTorch module designed to compute the Dynamic Time Warping (DTW) distance matrix between all pairs of time-series in a dataset. DTW is a method that calculates an optimal match between two given sequences with certain restrictions. The matrix is computed once and saved for future use to avoid redundant computations.

Parameters:: dataset_name (str) – The name of the dataset, used to save and retrieve the computed DTW matrix.

forward(data, **kwarg)

Computes the Dynamic Time Warping (DTW) matrix for the input data. If a precomputed matrix exists in the cache, it loads that matrix; otherwise, it computes a new matrix and saves it.

Parameters:

data (np.ndarray) – The input data array where each row represents a time step and each column a time-series node.
**kwargs (dict) – Additional keyword arguments not utilized in this method.

Returns:

The computed or loaded DTW distance matrix, where each element (i, j) represents the DTW distance between the i-th and j-th time-series.

Return type:

np.ndarray