Models

Spatial Models

GAT

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

Graph Attention Network (GAT)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

forward(x, edge_index, edge_weight)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

GCN

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

Graph Convolutional Network (GCN)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

forward(x, edge_index, edge_weight)

Parameters: x : torch.Tensor

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

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

Returns: torch.Tensor

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

GIN

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

Graph Isomorphism Network (GIN)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

forward(x, edge_index, edge_weight)

Parameters: x : torch.Tensor

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

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

Returns: torch.Tensor

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

SAGE

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

Graph Sample and Aggregate (SAGE)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

forward(x, edge_index, edge_weight)

Parameters: x : torch.Tensor

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

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

Returns: torch.Tensor

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

Temporal Models

ARIMA

class epilearn.models.Temporal.ARIMA.VARMAXModel(num_features, num_timesteps_input, num_timesteps_output)

Vector Autoregression Moving-Average with eXogenous variables (VARMAX) Model

Parameters:

num_features (int) – Number of features in each timestep of the input data.
num_timesteps_input (int) – Number of timesteps considered for each input sample.
num_timesteps_output (int) – Number of output timesteps to predict.

Returns:

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

Return type:

torch.Tensor

fit(train_input, train_target, val_input=None, val_target=None, epochs=1000, batch_size=10, verbose=False, patience=100)

Parameters:

train_input (numpy.ndarray) – The input training data array, expected to be in the shape (batch_size, num_timesteps_input, num_features), where: - batch_size is the number of training samples, - num_timesteps_input is the number of timesteps used as input, - num_features is the number of features per timestep.
train_target (numpy.ndarray) – The target training data array corresponding to the inputs, shaped (batch_size, num_timesteps_output).
val_input (numpy.ndarray, optional) – The input validation data array, following the same format as train_input.
val_target (numpy.ndarray, optional) – The target validation data array, following the same format as train_target.
epochs (int, optional) – The number of epochs to train the model for. Default is 1000.
batch_size (int, optional) – The size of batches to use when training the model. Default is 10.
verbose (bool, optional) – If True, the model will print out progress during training. Default is False.
patience (int, optional) – Number of epochs with no improvement after which training will be stopped. Default is 100.

Returns:

A list of forecasted values for each batch in the training data, each forecast corresponding to the future timesteps as defined by num_timesteps_output.

Return type:

list

Notes

This method internally converts the input PyTorch Tensors to numpy arrays if not already provided in that format, fits a VARMAX model for each subset of data, and performs forecasting. It is designed to handle smaller subsets of the data (defined by subset_samples) to demonstrate the model’s fitting capability.

DLINEAR

class epilearn.models.Temporal.Dlinear.DlinearModel(num_features, num_timesteps_input, num_timesteps_output)

Dynamic Linear Model

Parameters:

num_features (int) – Number of features in each timestep of the input data.
num_timesteps_input (int) – Number of timesteps considered for each input sample.
num_timesteps_output (int) – Number of output timesteps to predict.

decomposition

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

Type:: function

Linear_Transform

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

Type:: torch.nn.Linear

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

forward(x)

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

GRU

class epilearn.models.Temporal.GRU.GRUModel(num_features, num_timesteps_input, num_timesteps_output, nhid=256, dropout=0.5, use_norm=False)

Single-layer Gated Recurrent Unit (GRU) Network

Parameters:

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

Returns:

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

Return type:

torch.Tensor

forward(x)

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

LSTM

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

Long Short-Term Memory (LSTM) Model

Parameters:

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

Returns:

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

Return type:

torch.Tensor

forward(x)

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

SIR

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

Susceptible-Infected-Recovered (SIR) Model

Parameters:

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

beta

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

Type:: torch.nn.Linear

gamma

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

Type:: torch.nn.Linear

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

forward(x, steps=1)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

SIS

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

Susceptible-Infected-Susceptible (SIS) Model

Parameters:

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

beta

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

Type:: torch.nn.Linear

gamma

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

Type:: torch.nn.Linear

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

forward(x, steps=1)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

SEIR

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

Susceptible-Exposed-Infected-Recovered (SEIR) Model

Parameters:

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

beta

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

Type:: torch.nn.Linear

gamma

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

Type:: torch.nn.Linear

mu

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

Type:: torch.nn.Linear

a

Linear layer with no bias to model the latency rate.

Type:: torch.nn.Linear

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

forward(x, steps=1)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

XGBOOST

class epilearn.models.Temporal.XGB.XGBModel(num_features, num_timesteps_input, num_timesteps_output, n_estimators=40, learning_rate=0.1, max_depth=5, reg_lambda=1.0, reg_alpha=0.1)

Extreme Gradient Boosting (XGBoost) Model

Parameters:

num_features (int) – Number of features in each timestep of the input data.
num_timesteps_input (int) – Number of timesteps considered for each input sample.
num_timesteps_output (int) – Number of output timesteps to predict.
n_estimators (int, optional) – Number of gradient boosted trees. Equivalent to the number of boosting rounds. Default: 40.
learning_rate (float, optional) – Step size shrinkage used in update to prevents overfitting. Default: 0.1.
max_depth (int, optional) – Maximum depth of a tree. Increasing this value will make the model more complex and more likely to overfit. Default: 5.
reg_lambda (float, optional) – L2 regularization term on weights. Increasing this value will make model more conservative. Default: 1.0.
reg_alpha (float, optional) – L1 regularization term on weights. Increasing this value will make model more conservative. Default: 0.1.

Returns:

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

Return type:

torch.Tensor

fit(train_input, train_target, val_input=None, val_target=None, epochs=1000, batch_size=10, verbose=False, patience=100)

Parameters:

train_input (torch.Tensor or numpy.ndarray) – The input training data, expected to be in the shape (batch_size, num_timesteps_input, num_features).
train_target (torch.Tensor or numpy.ndarray) – The target training data, expected to be in the shape (batch_size, num_timesteps_output).
val_input (torch.Tensor or numpy.ndarray, optional) – The input validation data, following the same format as train_input.
val_target (torch.Tensor or numpy.ndarray, optional) – The target validation data, following the same format as train_target.
epochs (int, optional) – The number of epochs to train the model for. Default is 1000.
batch_size (int, optional) – The size of batches to use when training the model. Default is 10.
verbose (bool, optional) – If True, the training process will print out progress updates and evaluation metrics. Default is False.
patience (int, optional) – Number of epochs with no improvement after which training will be stopped. Default is 100.

Notes

The method first reshapes and prepares the input data for the XGBoost regressor, converting any PyTorch tensors to numpy arrays if necessary. It configures the model with the specified hyperparameters and trains it on the provided data. If validation data is provided, it will also evaluate the model on this data after each epoch, which can be used for early stopping or parameter tuning. The method concludes by optionally printing the mean squared error of the training and validation datasets if verbose is True.

Spatial-Temporal Models

ATMGNN

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

Attention-based Temporal Multiresolution Graph Neural Network (ATMGNN)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

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

Parameters:

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

Returns:

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

Return type:

torch.Tensor

MPNN_LSTM

class epilearn.models.SpatialTemporal.ATMGNN.MPNN_LSTM(num_nodes, num_features, num_timesteps_input, num_timesteps_output, nhid=256, dropout=0.5, device='cpu')

Message Passing Neural Network combined with Long Short-Term Memory (MPNN_LSTM)

Parameters:

num_nodes (int) – Number of nodes in the graph.
num_features (int) – Number of features in each timestep of the input data.
num_timesteps_input (int) – Window size; number of timesteps considered for each input sample.
num_timesteps_output (int) – Number of output timesteps to predict.
nhid (int, optional) – Number of hidden units in the LSTM and the number of output channels in GCN layers. Default: 256.
dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.5.
device (str, optional) – The device (cpu or gpu) on which the model will be run. Default: ‘cpu’.

Returns:

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

Return type:

torch.Tensor

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

Parameters:

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

Returns:

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

Return type:

torch.Tensor

CNNRNN_Res

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

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

Parameters:

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

Returns:

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

Return type:

torch.Tensor

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

Parameters:

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

Returns:

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

Return type:

torch.Tensor

ColaGNN

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

Convolutional-Layer Graph Neural Network (ColaGNN)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

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

Parameters:

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

Returns:

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

Return type:

torch.Tensor

DASTGN

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

Dynamic and Adaptive Spatio-Temporal Graph Network (DASTGN)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

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

Parameters:

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

Returns:

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

Return type:

torch.Tensor

DCRNN

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

Diffusion Convolutional Recurrent Neural Network (DCRNN)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

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

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

forward(X_batch, graph, X_states, batch_graph)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

DMP

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

forward(x, adj)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

EpiColaGNN

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

Epidemiological Convolutional-Layer Graph Neural Network (EpiColaGNN)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

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

Parameters:

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

Returns:

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

Return type:

torch.Tensor

EpiGNN

class epilearn.models.SpatialTemporal.EpiGNN.EpiGNN(num_nodes, num_features, num_timesteps_input, num_timesteps_output, k=8, hidA=64, hidR=40, hidP=1, n_layer=2, dropout=0.5, device='cpu')

Epidemiological Graph Neural Network (EpiGNN)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

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

Parameters:

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

Returns:

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

Return type:

torch.Tensor

GraphWaveNet

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

Graph Convolutional Wave Network (GraphWaveNet)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

forward(X_batch, graph, X_states, batch_graph)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

MepoGNN

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

Meta-Population Graph Neural Network (MepoGNN)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

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

Parameters:

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

Returns:

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

Return type:

torch.Tensor

NetSIR

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

Network-based SIR (Susceptible-Infected-Recovered)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

forward(x, adj, steps=1)

Parameters:

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

Returns:

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

Return type:

torch.Tensor

STAN

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

Spatio-Temporal Attention Network (STAN)

Parameters:

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

Returns:

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

Return type:

tuple of torch.Tensor

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

Parameters:

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

Returns:

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

Return type:

tuple of torch.Tensor

STGCN

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

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

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

Parameters:

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

Returns:

Output shape (batch_size, num_timesteps_output, num_nodes)

Return type:

torch.Tensor