Dataset

In EpiLearn, we use Dataset to load preprocessed datasets. For customized data, we can simply initialize the Dataset given features, graphs, and states.

Note

Dataset was named UniversalDataset before version 0.1.0. The old name is kept as a deprecated alias (from epilearn.data import UniversalDataset) and the constructor signature is unchanged, so existing scripts keep working – but new code should use Dataset.

Dataset class

Dataset class for time series data with proper indexing and slicing.

This class provides: - Loading from various sources (CSV, PT, NPY, or direct numpy/tensor) - Timestamp and region-based slicing - Train/val/test splitting with proper normalization - Sliding window dataset generation for training - Backward compatibility with existing code

Parameters:

name (str, optional) – Name of built-in dataset to load
root (str, optional) – Root directory for downloads
x (torch.Tensor or numpy.ndarray, optional) – Features of shape (T, N, F) for spatiotemporal or (T, F) for temporal
y (torch.Tensor or numpy.ndarray, optional) – Targets
graph (torch.Tensor or numpy.ndarray, optional) – Static adjacency matrix
dynamic_graph (torch.Tensor or numpy.ndarray, optional) – Dynamic graph over time
states (torch.Tensor or numpy.ndarray, optional) – State variables (e.g., SIR compartments)
timestamps (list, optional) – List of timestamp values
regions (list, optional) – List of region identifiers
feature_names (list, optional) – Names of feature columns
target_names (list, optional) – Names of target columns

apply_transforms(inplace: bool = True)

Apply the stored transformations to the dataset.

Parameters:: inplace (bool, default True) – If True, modify the dataset in place. If False, return a new dataset with transformations applied.
Returns:: If inplace=True, returns self. If inplace=False, returns (transformed_data_dict, process_history).
Return type:: Dataset or tuple

download(): Download dataset (to be implemented by subclasses).

classmethod from_csv(file_path: str, timestamp_col: str, feature_cols: List[str], target_cols: List[str] | None = None, region_col: str | None = None, graph_file: str | None = None, **kwargs) → Dataset

Load dataset from CSV file with flexible column mapping.

Parameters:

file_path – Path to CSV file
timestamp_col – Column name for timestamps
feature_cols – List of feature column names
target_cols – List of target column names. If None, uses the first feature column as target (common for forecasting tasks)
region_col – Column name for regions (for spatiotemporal data). If provided, creates spatiotemporal dataset.
graph_file – Optional path to CSV with graph edges
**kwargs – Additional arguments for CSVLoader

Returns:

Dataset instance

Example

# Minimal usage - timestamp, region, features (first feature = target) dataset = Dataset.from_csv(

“data.csv”, timestamp_col=”date”, feature_cols=[“cases”, “deaths”], # ‘cases’ will be target region_col=”state”

)

# Explicit target dataset = Dataset.from_csv(

“data.csv”, timestamp_col=”date”, feature_cols=[“cases”, “deaths”, “tests”], target_cols=[“cases”], region_col=”state”, graph_file=”edges.csv”

)

# Temporal only (no region) dataset = Dataset.from_csv(

“data.csv”, timestamp_col=”date”, feature_cols=[“value”]

)

classmethod from_numpy(file_path: str, **kwargs) → Dataset: Load dataset from NumPy file.

classmethod from_tensor(file_path: str, **kwargs) → Dataset: Load dataset from PyTorch tensor file.

classmethod from_time_series_data(ts_data: TimeSeriesData) → Dataset: Create Dataset from TimeSeriesData object.

generate_dataset(X=None, Y=None, states=None, dynamic_adj=None, adj=None, lookback_window_size=1, horizon_size=1, interval=None, ahead=0, permute=False, region_idx=None)

Generate sliding window dataset for training/evaluation.

Parameters:

X – Features tensor of shape (T, N, F) or (T, F)
Y – Targets tensor
states – State variables (e.g., SIR compartments)
dynamic_adj – Dynamic adjacency matrices
adj – Static adjacency matrix
lookback_window_size – Number of timesteps in input window
horizon_size – Number of timesteps to predict
interval – Offset to align predictions across different lookback sizes. When comparing models with different lookbacks, set this to (max_lookback - current_lookback) so all models predict the same time points. Default None means no offset.
ahead – Gap between input window and prediction window
permute – Whether to permute dimensions
region_idx – Optional region indices to select

Returns:

Dictionary with ‘features’, ‘targets’, ‘states’, ‘dynamic_graph’, ‘graph’

Example - Aligning different lookbacks:

# With max_lookback=21, both produce samples predicting same time points split_21 = ds.generate_dataset(lookback_window_size=21, interval=0) split_7 = ds.generate_dataset(lookback_window_size=7, interval=14)

# split_21 sample 0: input [0:21], predict [21:28] # split_7 sample 0: input [14:21], predict [21:28] <- same prediction!

get_process_history()

Get the processing history (normalization statistics).

Returns:: Dictionary containing normalization statistics like ‘feat_mean’, ‘feat_std’, ‘target_mean’, ‘target_std’.
Return type:: dict

get_region_idx(region: Any) → int: Get index for a region identifier.

get_region_indices(regions: List[Any] | None = None) → List[int] | None: Get indices for a list of regions.

Get a slice of the dataset for a specific timestamp range and regions.

This method supports both timestamp values and rate-based slicing, making it ideal for rolling evaluation where you progressively expand or slide the training/validation/test windows.

Timestamp values should match the format in your data file: - Integer indices: 0, 1, 2, 3, … - Date strings: ‘2020-01-01’, ‘2020-01-02’, … - Pandas Timestamps: pd.Timestamp(‘2020-01-01’), …

Parameters:

start – Start timestamp value (inclusive). Use actual values from your data, not indices. If None, starts from beginning.
end – End timestamp value. By default exclusive (like Python slicing). Set end_inclusive=True to include this timestamp.
start_rate – Alternative to ‘start’ - proportion of total timesteps (0.0-1.0)
end_rate – Alternative to ‘end’ - proportion of total timesteps (0.0-1.0)
end_inclusive – If True, include the ‘end’ timestamp in the slice.
regions – List of region identifiers to include (None = all regions)

Returns:

New Dataset object with sliced data

Example - Rolling Evaluation with timestamp values:

# Timestamps are integers: 1, 2, 3, …, 100

# Train on timestamps 1-70 (inclusive) train = dataset.get_slice(start=1, end=70, end_inclusive=True)

# Validate on 71-85 val = dataset.get_slice(start=71, end=85, end_inclusive=True)

# Test on 86-100 test = dataset.get_slice(start=86, end=100, end_inclusive=True)

Example - Rolling Evaluation with date timestamps:

# Timestamps are dates: ‘2020-01-01’, ‘2020-01-02’, …

train = dataset.get_slice(start=’2020-01-01’, end=’2020-07-01’) val = dataset.get_slice(start=’2020-07-01’, end=’2020-09-01’) test = dataset.get_slice(start=’2020-09-01’) # To the end

Example - Rolling Evaluation with rates:

train = dataset.get_slice(end_rate=0.7) val = dataset.get_slice(start_rate=0.7, end_rate=0.85) test = dataset.get_slice(start_rate=0.85)

Example - Region subsetting:

subset = dataset.get_slice(regions=[‘RegionA’, ‘RegionB’])

Example - Combined timestamp and region slicing:

subset = dataset.get_slice(start=50, end=100, regions=[‘RegionA’])

get_timestamp_idx(timestamp: Any) → int: Get index for a timestamp value.

get_timestamp_range_indices(start: Any | None = None, end: Any | None = None, start_rate: float | None = None, end_rate: float | None = None, end_inclusive: bool = False) → Tuple[int, int]

Get start and end indices for a timestamp range.

Supports two modes: - Timestamp values: Use actual values from the data (dates, integers, etc.) - Rate-based: Proportions (0.0 to 1.0) of the total timesteps

Timestamp values take precedence over rates if both are provided.

Parameters:

start – Start timestamp value (inclusive). Format should match timestamps in the dataset (e.g., integers 0,1,2… or dates ‘2020-01-01’).
end – End timestamp value. By default exclusive, set end_inclusive=True to include the end timestamp.
start_rate – Start as proportion of total timesteps (0.0 to 1.0)
end_rate – End as proportion of total timesteps (0.0 to 1.0)
end_inclusive – If True, include the end timestamp in the slice. Only applies when ‘end’ is a timestamp value, not rate.

Returns:

Tuple of (start_idx, end_idx) where end_idx is exclusive

Example

# Using timestamp values (integers) start_idx, end_idx = dataset.get_timestamp_range_indices(start=10, end=50)

# Using timestamp values (dates) start_idx, end_idx = dataset.get_timestamp_range_indices(

start=’2020-01-01’, end=’2020-06-01’

)

# Using rates start_idx, end_idx = dataset.get_timestamp_range_indices(

start_rate=0.0, end_rate=0.7

)

get_transforms()

Get the current transformation configuration.

Returns:: The current transforms object, or None if not set.
Return type:: Compose or None

property is_spatiotemporal: bool: Check if data is spatiotemporal (has region dimension).

classmethod load(path: str) → Dataset: Load dataset from file.

classmethod load_from_csv(file_path: str, timestamp_col: str, feature_cols: List[str], target_cols: List[str] | None = None, region_col: str | None = None, graph_file: str | None = None, **kwargs) → Dataset

Load dataset from CSV file with flexible column mapping.

Parameters:

file_path – Path to CSV file
timestamp_col – Column name for timestamps
feature_cols – List of feature column names
target_cols – List of target column names. If None, uses the first feature column as target (common for forecasting tasks)
region_col – Column name for regions (for spatiotemporal data). If provided, creates spatiotemporal dataset.
graph_file – Optional path to CSV with graph edges
**kwargs – Additional arguments for CSVLoader

Returns:

Dataset instance

Example

# Minimal usage - timestamp, region, features (first feature = target) dataset = Dataset.from_csv(

“data.csv”, timestamp_col=”date”, feature_cols=[“cases”, “deaths”], # ‘cases’ will be target region_col=”state”

)

# Explicit target dataset = Dataset.from_csv(

“data.csv”, timestamp_col=”date”, feature_cols=[“cases”, “deaths”, “tests”], target_cols=[“cases”], region_col=”state”, graph_file=”edges.csv”

)

# Temporal only (no region) dataset = Dataset.from_csv(

“data.csv”, timestamp_col=”date”, feature_cols=[“value”]

)

load_toy_dataset(): Load toy dataset for testing.

property n_features: int: Number of features.

property n_regions: int: Number of regions.

property n_timesteps: int: Number of timesteps.

rolling_splits(train_size: int, test_size: int, step_size: int | None = None, val_size: int = 0, expanding: bool = True, regions: List[Any] | None = None) → Iterator[Tuple[Dataset, Dataset | None, Dataset]]

Generate rolling train/val/test splits for time series cross-validation.

This is the recommended way to evaluate forecasting models, as it respects the temporal ordering of the data and avoids data leakage.

Two modes are supported: - Expanding window: Training window grows, test window slides - Sliding window: Both train and test windows slide (fixed train size)

Parameters:

train_size – Initial training window size (number of timesteps)
test_size – Test window size (number of timesteps)
step_size – How many timesteps to advance each iteration. Defaults to test_size (non-overlapping test windows).
val_size – Validation window size between train and test. Default 0.
expanding – If True, training window expands. If False, slides.
regions – Optional list of regions to include in all splits.

Yields:

Tuple of (train_dataset, val_dataset, test_dataset) val_dataset is None if val_size=0

Example - Basic rolling evaluation:

for train, val, test in dataset.rolling_splits(: train_size=100, test_size=20, val_size=10
):: # Train on train, validate on val, evaluate on test model.fit(train) model.evaluate(test)

Example - Sliding window (fixed train size):

for train, _, test in dataset.rolling_splits(: train_size=100, test_size=20, expanding=False
):: model.fit(train) model.evaluate(test)

Example - Walk-forward validation:

results = [] for train, val, test in dataset.rolling_splits(

train_size=200, test_size=7, val_size=14, step_size=7

):: # train_size expands, we step forward 7 days each iteration metrics = train_and_evaluate(train, val, test) results.append(metrics)

save(path: str): Save dataset to file.

set_transforms(transforms, apply_now: bool = False)

Set preprocessing transformations for the dataset.

This method properly sets transformation configurations that will be applied during training/evaluation. Optionally applies transformations immediately.

Parameters:

transforms (Compose) –
A Compose object containing transformations for different data components. Expected format:

transforms.Compose({
“target”: [transforms.normalize_target()], “features”: [transforms.normalize_feat()], “graph”: [transforms.normalize_adj()], “dynamic_graph”: [transforms.normalize_adj()],

})
apply_now (bool, default False) – If True, immediately apply transformations to the current data. If False, transformations are stored and applied during training/splits.

Returns:

self – Returns self for method chaining.

Return type:

Dataset

Examples

>>> from epilearn.utils import transforms
>>> transformation = transforms.Compose({
...     "target": [transforms.normalize_target()],
...     "features": [transforms.normalize_feat()],
...     "graph": [transforms.normalize_adj()],
... })
>>> dataset.set_transforms(transformation)
>>>
>>> # Or apply immediately
>>> dataset.set_transforms(transformation, apply_now=True)

to_time_series_data() → TimeSeriesData: Convert to TimeSeriesData object.

Preprocessed Datasets

We collect epidemic data from various sources including the followings:

The name in code font is the value to pass as name=.

Temporal Data

Tycho_v1 – Tycho v1.0.0: Including eight diseases collected across 50 US states and 122 US cities from 1916 to 2009.

Measles – Measles: Contains measles infections in England and Wales across 954 urban centers (cities and towns) from 1944 to 1964.

JHU_covid – Johns Hopkins University global COVID-19 case counts.

Spatial&Temporal Data

Covid_<Country> with a static graph – Covid_China, Covid_Brazil, Covid_Austria: covid infections with static graph. [1]

Covid_<Country> with a dynamic graph – Covid_England, Covid_France, Covid_Italy, Covid_NewZealand, Covid_Spain: covid infections with dynamic graph. [2] [3]

covid_static.pt / covid_dynamic.pt are the two downloaded archives that back those eight names – they are not themselves values you can pass as name=.

Dataset Loading

Passing name= downloads the archive into root on first use and reuses the local copy afterwards. Auxiliary tables that used to live on the dataset object (anual_population, coordinates, index) are now collected under dataset.metadata.

from epilearn.data import Dataset

measle_dataset = Dataset(name='Measles', root='./tmp/')
print(measle_dataset)                       # Dataset(Temporal, x=(946, 1108))
print(measle_dataset.feature_names[:3])     # ['Abingdon', 'Abram', 'Accrington']
print(measle_dataset.metadata.keys())       # anual_population, anual_birth, coordinates

jhu_dataset = Dataset(name='JHU_covid', root='./tmp/')
print(jhu_dataset)                          # Dataset(Temporal, x=(3342, 1143))

For other countries, please use ‘Covid_’+’country’ to acquire the correspnding covid dataset. Three countries come with a static graph – China, Brazil and Austria – and load as spatiotemporal data. Five more come with a dynamic graph: England, France, Italy, NewZealand and Spain (note the capital Z); their x is 2-D (T, N), so they print as Temporal even though dynamic_graph is populated.

covid_dataset = Dataset(name='Covid_Brazil', root='./tmp/')
print(covid_dataset)
# Dataset(Spatiotemporal, x=(122, 27, 3), graph=(27, 27))

covid_spain = Dataset(name='Covid_Spain', root='./tmp/')
print(covid_spain.dynamic_graph.shape)      # torch.Size([62, 52, 52])

An unsupported country name raises ValueError listing the ones that are available. In a source checkout, pass root='./datasets' instead of './tmp/' to use the copies already in the repository and skip the download entirely.

Warning

The named loaders fill in x (and the graph), but leave y as None, so a built-in dataset cannot go straight into a task: rolling_train dies with TypeError: 'NoneType' object is not subscriptable. Choose a target first:

m = Dataset(name='Measles', root='./datasets')
dataset = Dataset(x=m.x.unsqueeze(-1), y=m.x)   # forecast each city's own series

Rebuilding the Dataset like this also drops timestamps, which is what you want: JHU_covid carries 1144 timestamp labels for 3342 timesteps and the dynamic-graph countries 62 for 122, and rolling_splits raises on that mismatch (IndexError, or ValueError: The truth value of a Index is ambiguous).

The Tycho archive stores one variable-length series per disease, so it has to be indexed by disease before it becomes a Dataset:

import torch
from epilearn.data import Dataset

try:                                         # downloads ./tmp/Tycho_v1.pt, then raises
    Dataset(name='Tycho_v1', root='./tmp/')
except ValueError:
    pass

raw = torch.load('./tmp/Tycho_v1.pt', weights_only=False)
print(list(raw.keys()))
# ['DIPHTHERIA', 'HEPATITIS A', 'MEASLES', 'MUMPS', 'SMALLPOX']
series = raw['MEASLES']                      # torch.Size([3772])

tycho_measles = Dataset(x=series.unsqueeze(-1), y=series)
print(tycho_measles)                         # Dataset(Temporal, x=(3772, 1), y=(3772,))

Warning

Dataset(name='Tycho_v1', ...) raises ValueError: only one element tensors can be converted to Python scalars because the diseases have different lengths and cannot be stacked into a single tensor. The call still downloads the file, which is why the snippet above swallows the error and then reads ./tmp/Tycho_v1.pt itself.

Customize Your Own Dataset

For your own data, form a dictionary with keys features, graph, dynamic_graph, targets and states, then pass the arrays to Dataset. Not every argument is required – see Dataset class for the full signature. examples/example.pt in the repository is exactly such a file:

import torch
from epilearn.data import Dataset

data = torch.load("examples/example.pt", weights_only=False)
print(data.keys())
# dict_keys(['features', 'graph', 'dynamic_graph', 'targets', 'states'])

dataset = Dataset(x=data['features'],                    # (time, nodes, channels)
                  y=data['targets'],                     # (time, nodes)
                  states=data['states'],                 # e.g. SIR states per node
                  graph=torch.Tensor(data['graph']),     # (nodes, nodes); edge_index also works
                  dynamic_graph=data['dynamic_graph'])   # (time, nodes, nodes)

print(dataset)
# Dataset(Spatiotemporal, x=(539, 47, 4), y=(539, 47), graph=(47, 47))
print(dataset.edge_index.shape)              # torch.Size([2, 2189])
print(dataset.n_timesteps, dataset.n_regions, dataset.n_features)   # 539 47 4

graph is converted to edge_index / edge_weight automatically, and timestamps / regions default to range(T) / range(N) when you do not pass them. A dictionary-style .pt file like this one can also be handed to Dataset.from_tensor directly, since the keys are auto-detected (see Loading from tensor and numpy files). For more sample code in a real training process, refer to examples/dataset_customization.ipynb on the github page.

Loading from a CSV file

Dataset.from_csv reads long-format data: one row per (timestamp, region), with the features and targets as columns. This is the usual shape of public surveillance exports.

time,node,f0,f1,f2,f3,y
1,0,5.0,-0.10555339604616165,2.0,0.09615384787321091,5.0
1,1,0.0,-0.13394306600093842,2.0,0.0,0.0
2,0,3.0,-0.05203865468502045,3.0,0.06382978707551956,3.0
...

Regions are named by region_col; the loader pivots the frame into an (T, N, F) feature tensor. An optional graph_file is a second CSV of edges with source / target columns (and an optional weight column via graph_weight_col), which becomes an (N, N) adjacency matrix aligned to the sorted region order.

from epilearn.data import Dataset

dataset = Dataset.from_csv(
    file_path="datasets/toy_features.csv",
    timestamp_col="time",
    feature_cols=["f0", "f1", "f2", "f3"],
    target_cols=["y"],
    region_col="node",
    graph_file="datasets/toy_edges.csv",
)

print(dataset)
# Dataset(Spatiotemporal, x=(539, 47, 4), y=(539, 47), graph=(47, 47))
print(dataset.timestamps[:3], dataset.regions[:3])   # [1, 2, 3] [0, 1, 2]
print(dataset.feature_names, dataset.target_names)   # ['f0','f1','f2','f3'] ['y']

Drop region_col for a single series – the result is a Temporal dataset. target_cols is optional: omit it and the first entry of feature_cols becomes the target, which is the common forecasting setup. The result is a normal Dataset, so it goes straight into a task: attach transforms and call rolling_train as in the Quickstart for EpiLearn.

Note

The timestamp column is date-parsed automatically when it looks like dates, so dataset.timestamps holds pandas.Timestamp objects rather than strings. Index with pd.Timestamp('2020-01-01'), or pass parse_dates=False to keep the raw strings. Other keyword arguments are forwarded to CSVLoader (and from there to pandas.read_csv): fillna_method, fillna_value, strict_numeric, sort_timestamps, date_format.

Loading from tensor and numpy files

.pt and .npy / .npz files are read by Dataset.from_tensor and Dataset.from_numpy. For dictionary-style .pt files the keys are auto-detected from a list of aliases (features/x/inputs…, targets/y/labels…, graph/adj…, dynamic_graph/od…, states/SIR…), so the file written by Dataset.save round-trips without any configuration.

dataset.save("my_dataset.pt")               # save() now requires an explicit path

reloaded = Dataset.from_tensor("my_dataset.pt")
print(reloaded)
# Dataset(Spatiotemporal, x=(539, 47, 4), y=(539, 47), graph=(47, 47))

same = Dataset.load("my_dataset.pt")        # equivalent for files written by save()

# a bare .npy array is taken as the features
arrays = Dataset.from_numpy("features.npy")  # file holding a (50, 6, 2) array
print(arrays)                               # Dataset(Spatiotemporal, x=(50, 6, 2))

Loaders

The loaders behind the from_* constructors are also usable directly. They return a TimeSeriesData (or a LoadedData for load_from_file, which picks the loader from the file extension) rather than a Dataset; wrap it with Dataset.from_time_series_data when you need the full dataset API. This is the hook to use if you want to inspect or patch the parsed arrays before building the dataset.

from epilearn.data import (Dataset, load_csv, load_tensor, load_numpy,
                           load_from_file, DataLoaderRegistry)

print(DataLoaderRegistry.supported_extensions())
# ['csv', 'pt', 'pth', 'npy', 'npz']

ts = load_csv("datasets/toy_features.csv",
              timestamp_col="time",
              feature_cols=["f0", "f1", "f2", "f3"],
              target_cols=["y"],
              region_col="node")
print(ts)
# TimeSeriesData(Spatiotemporal, features=(539, 47, 4), targets=(539, 47), T=539, N=47)

loader_dataset = Dataset.from_time_series_data(ts)

ts_pt  = load_tensor("my_dataset.pt")
ts_npy = load_numpy("features.npy")

loaded = load_from_file("datasets/toy_features.csv", timestamp_col="time",
                        feature_cols=["f0", "f1", "f2", "f3"], target_cols=["y"],
                        region_col="node")
same_dataset = Dataset.from_time_series_data(loaded.to_time_series_data())

epilearn.data.loaders.load_csv(file_path, timestamp_col, feature_cols, target_cols=None, region_col=None, **kwargs)

Load time series data from a CSV file.

Parameters:

file_path – Path to CSV file
timestamp_col – Column name for timestamps
feature_cols – List of feature column names
target_cols – List of target column names (optional)
region_col – Column name for regions (for spatiotemporal data)
**kwargs – Additional arguments for CSVLoader

Returns:

TimeSeriesData object

epilearn.data.loaders.load_tensor(file_path, **kwargs)

Load time series data from a PyTorch tensor file (.pt).

Parameters:

file_path – Path to .pt file
**kwargs – Additional arguments for TensorLoader

Returns:

TimeSeriesData object

epilearn.data.loaders.load_numpy(file_path, **kwargs)

Load time series data from a NumPy file (.npy or .npz).

Parameters:

file_path – Path to .npy or .npz file
**kwargs – Additional arguments for NumpyLoader

Returns:

TimeSeriesData object

epilearn.data.loaders.load_from_file(file_path, **kwargs)

Convenience function to load data from a file using the appropriate loader.

Parameters:

file_path – Path to the data file
**kwargs – Additional arguments passed to the loader

Returns:

LoadedData object

Example

# Load CSV with column specifications data = load_from_file(

“data.csv”, timestamp_col=”date”, feature_cols=[“cases”, “deaths”], region_col=”state”

)

# Load tensor file data = load_from_file(“data.pt”)

class epilearn.data.loaders.csv_loader.CSVLoader(file_path: str | Path, timestamp_col: str, feature_cols: List[str], target_cols: List[str] | None = None, region_col: str | None = None, graph_file: str | Path | None = None, graph_source_col: str = 'source', graph_target_col: str = 'target', graph_weight_col: str | None = None, parse_dates: bool | str = 'auto', date_format: str | None = None, fillna_method: str | None = None, fillna_value: float | None = None, strict_numeric: bool = True, sort_timestamps: bool = True, **read_csv_kwargs)

Load time series data from CSV files with flexible column mapping.

Supports both temporal (single time series) and spatiotemporal (multiple regions) data formats. Users specify which columns contain timestamps, regions, features, and targets.

Example CSV format for spatiotemporal data:

date,region,cases,deaths,population 2020-01-01,RegionA,10,1,100000 2020-01-01,RegionB,5,0,50000 2020-01-02,RegionA,15,2,100000 2020-01-02,RegionB,8,1,50000

Example usage:

loader = CSVLoader(: file_path=”data.csv”, timestamp_col=”date”, region_col=”region”, feature_cols=[“cases”, “deaths”, “population”], target_cols=[“cases”]

) data = loader.load()

load() → LoadedData

Load and process the CSV file.

Returns:: LoadedData object with processed data

Load time series data from PyTorch tensor files (.pt).

Supports various data formats stored in .pt files, including: - Single tensor (interpreted as features) - Dictionary with ‘features’, ‘targets’, etc. keys - Dictionary with custom key mappings

Example usage:

loader = TensorLoader(: file_path=”data.pt”, feature_key=”features”, target_key=”targets”, graph_key=”adjacency”

) data = loader.load()

load() → LoadedData: Load and process the tensor file.

class epilearn.data.loaders.tensor_loader.NumpyLoader(file_path: str | Path, feature_key: str = 'features', target_key: str = 'targets', graph_key: str = 'graph', dynamic_graph_key: str = 'dynamic_graph', states_key: str = 'states', allow_pickle: bool = True, **kwargs)

Load time series data from NumPy files (.npy or .npz).

For .npy files, the array is treated as features. For .npz files, keys are used to identify features, targets, etc.

load() → LoadedData: Load and process the NumPy file.

class epilearn.data.core.TimeSeriesData(features: ~torch.Tensor, targets: ~torch.Tensor | None = None, timestamps: ~typing.List[~typing.Any] | None = None, regions: ~typing.List[~typing.Any] | None = None, graph: ~torch.Tensor | None = None, dynamic_graph: ~torch.Tensor | None = None, states: ~torch.Tensor | None = None, feature_names: ~typing.List[str] = <factory>, target_names: ~typing.List[str] = <factory>, metadata: ~typing.Dict[str, ~typing.Any] = <factory>)

Core container for time series data with proper indexing.

This class stores numerical data along with timestamp and region indices, supporting both temporal-only and spatiotemporal data formats.

features

Tensor of shape (T, N, F) for spatiotemporal or (T, F) for temporal where T=timesteps, N=regions/nodes, F=features

Type:: torch.Tensor

targets

Tensor of shape (T, N, T_out) or (T, T_out) for targets

Type:: torch.Tensor | None

timestamps

List of timestamp values (can be datetime, int, str)

Type:: List[Any] | None

regions

Optional list of region identifiers (for spatiotemporal data)

Type:: List[Any] | None

graph

Optional static adjacency matrix of shape (N, N)

Type:: torch.Tensor | None

dynamic_graph

Optional dynamic graph of shape (T, N, N)

Type:: torch.Tensor | None

states

Optional state variables (e.g., SIR states) of shape matching features

Type:: torch.Tensor | None

feature_names

List of feature column names

Type:: List[str]

target_names

List of target column names

Type:: List[str]

metadata

Additional metadata dictionary

Type:: Dict[str, Any]

class epilearn.data.core.LoadedData(features: ~numpy.ndarray, targets: ~numpy.ndarray | None = None, timestamps: ~typing.List[~typing.Any] | None = None, regions: ~typing.List[~typing.Any] | None = None, graph: ~numpy.ndarray | None = None, dynamic_graph: ~numpy.ndarray | None = None, states: ~numpy.ndarray | None = None, feature_names: ~typing.List[str] = <factory>, target_names: ~typing.List[str] = <factory>, metadata: ~typing.Dict[str, ~typing.Any] = <factory>)

Container for data loaded from external sources before creating TimeSeriesData.

This is an intermediate representation used by loaders.

to_time_series_data() → TimeSeriesData: Convert to TimeSeriesData.

Transformations

Transformations are attached with set_transforms. By default they are stored and applied by the task during training; pass apply_now=True to rewrite the dataset immediately. get_process_history returns the statistics that were used, which is what you need to map predictions back to the original scale.

from epilearn.utils import transforms

transformation = transforms.Compose({
    "features": [transforms.normalize_feat()],
    "target": [transforms.normalize_target()],
    "graph": [transforms.normalize_adj()]})

dataset.set_transforms(transformation, apply_now=True)

print(dataset.x.mean(), dataset.x.std())        # ~0.0, ~1.0
print(sorted(dataset.get_process_history()))
# ['feat_mean', 'feat_std', 'target_mean', 'target_std']

# equivalent, in two steps
dataset.set_transforms(transformation)
dataset.apply_transforms()                      # in place; inplace=False returns
                                                # (transformed_dict, process_history)

Note

Dataset.get_transformed() was removed in 0.1.0; use set_transforms / apply_transforms. Note also that metrics are not automatically denormalized: with normalize_target() in the pipeline, reported errors are in normalized units unless you ask for inverse_normalize=True when evaluating.

Slicing and rolling splits

get_slice returns a new Dataset restricted to a timestamp range and/or a subset of regions. The range can be given as timestamp values (start / end, exclusive unless end_inclusive=True; pandas.Timestamp for date-typed series) or as proportions of the series (start_rate / end_rate). Slicing a region subset also slices graph and dynamic_graph, and carries the transforms and process_history over.

train = dataset.get_slice(end_rate=0.7)
test  = dataset.get_slice(start_rate=0.7)
print(train, test)
# Dataset(Spatiotemporal, x=(377, 47, 4), ...) Dataset(Spatiotemporal, x=(162, 47, 4), ...)

# timestamp values, inclusive end, and a region subset
window = dataset.get_slice(start=10, end=20, end_inclusive=True, regions=[0, 1, 2])
print(window)                # Dataset(Spatiotemporal, x=(11, 3, 4), y=(11, 3), graph=(3, 3))

rolling_splits yields the (train, val, test) triples used for rolling-origin evaluation – the protocol that rolling_train runs internally. expanding=True grows the training window each fold; expanding=False slides a fixed-size one. val_size=0 yields None for the validation dataset. step_size (default: test_size, i.e. non-overlapping test windows) controls how far the windows advance per fold, and regions restricts every split to a subset of regions.

for train, val, test in dataset.rolling_splits(train_size=300, val_size=50, test_size=50):
    print(train.n_timesteps, val.n_timesteps, test.n_timesteps,
          test.timestamps[0], test.timestamps[-1])

# 300 50 50 351 400
# 350 50 50 401 450
# 400 50 50 451 500
# 450 50 39 501 539     <- trailing partial fold, kept when >= test_size // 2 remains

Sliding windows: `generate_dataset`

generate_dataset turns a dataset into the supervised (lookback, horizon) samples a model consumes. It returns a dict with keys features, targets, states, dynamic_graph and graph – in 0.0.x this was a tuple.

split = dataset.generate_dataset(X=dataset.x,
                                 Y=dataset.y,
                                 adj=dataset.graph,
                                 lookback_window_size=12,
                                 horizon_size=3)

print(list(split.keys()))
# ['features', 'targets', 'states', 'dynamic_graph', 'graph']
print(split['features'].shape)      # torch.Size([525, 12, 47, 4])
print(split['targets'].shape)       # torch.Size([525, 47, 3])
print(split['graph'].shape)         # torch.Size([47, 47])

Warning

Pass adj=dataset.graph explicitly. generate_dataset does not fall back to the dataset’s own graph, so omitting it silently produces graph=None and graph-based models will fail. The same holds for states= and dynamic_adj=.

These dicts are what the low-level task.train_model(train_split=..., val_split=..., test_split=...) and task.evaluate_model(dataset=<split dict>) expect. Use interval= (set to max_lookback - lookback) when you compare models with different lookbacks and want them to predict the same time points.

Dataset

Dataset class

Preprocessed Datasets

Customize Your Own Dataset

Loading from a CSV file

Loading from tensor and numpy files

Loaders

Transformations

Slicing and rolling splits

Sliding windows: generate_dataset

Sliding windows: `generate_dataset`