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, num_classes, hidden_dim=16, nlayers=2, dropout=0.5, with_bias=True, device=None)

Graph Isomorphism Network (GIN)

Parameters:

  • num_features (int) – Number of input features per node.

  • hidden_dim (int) – Dimension of hidden layers.

  • num_classes (int) – Number of output features per node.

  • nlayers (int, optional) – Number of layers in the GIN. Default: 2.

  • dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.5.

  • with_bias (bool, optional) – Specifies whether to include bias parameters in the MLP layers. Default: True.

  • device (str) – The device (cpu or gpu) on which the model will be run. Must be specified.

Returns:

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

Return type:

torch.Tensor

forward(x, edge_index, edge_weight=None)

Parameters: x : torch.Tensor

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

edge_indextorch.Tensor

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

Returns: torch.Tensor

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

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=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

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

Code Example

import torch
from epilearn.models.Spatial import GCN

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

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

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

Temporal Models

ARIMA

class epilearn.models.Temporal.ARIMA.VARMAXModel(num_features, num_timesteps_input, num_timesteps_output, order=(1, 0), trend='c')

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, train_states=None, val_input=None, val_target=None, epochs=1000, batch_size=10, verbose=False, patience=100, **kwargs)
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, device='cpu')

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, device='cpu')

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, **kwargs)

Long Short-Term Memory (LSTM) Model

Parameters:

  • num_features (int) – Number of features in each timestep of the input data.

  • num_timesteps_input (int) – Number of timesteps considered for each input sample.

  • num_timesteps_output (int) – Number of output timesteps to predict.

  • nhid (int, optional) – Number of hidden units in the LSTM layers. Default: 256.

  • dropout (float, optional) – Dropout rate for regularization during training to prevent overfitting. Default: 0.5.

  • use_norm (bool, optional) – Whether to apply layer normalization after the LSTM layers. Default: False.

Returns:

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

Return type:

torch.Tensor

forward(x)
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

Code Example

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

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

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

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

Spatial-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

DASTGN

DCRNN

DMP

EpiColaGNN

EpiGNN

GraphWaveNet

MepoGNN

NetSIR

STAN

STGCN

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

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

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

  • X (torch.Tensor) – 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

Code Example

import torch
from epilearn.models.SpatialTemporal import ColaGNN

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


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

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