Data objects#

Data

A data object describing a spatiotemporal graph, i.e., a graph with time series of equal length associated with every node.

StaticBatch

A batch of tsl.data.Data objects for multiple spatiotemporal graphs sharing the same topology.

DisjointBatch

A data object describing a batch of graphs as one big (disconnected) graph.

class Data(input: Optional[Mapping] = None, target: Optional[Mapping] = None, edge_index: Optional[Union[Tensor, SparseTensor]] = None, edge_weight: Optional[Tensor] = None, mask: Optional[Tensor] = None, transform: Optional[Mapping] = None, pattern: Optional[Mapping] = None, **kwargs)[source]#

A data object describing a spatiotemporal graph, i.e., a graph with time series of equal length associated with every node.

The data object extends torch_geometric.data.Data, thus preserving all its functionalities (see also the accompanying tutorial).

Parameters:
  • input (Mapping, optional) – Named mapping of Tensor to be used as input to the model. (default: None)

  • target (Mapping, optional) – Named mapping of Tensor to be used as target of the task. (default: None)

  • edge_index (Adj, optional) – Graph connectivity either in COO format (a Tensor of shape [2, E]) or as a torch_sparse.SparseTensor with shape [N, N]. For dynamic graphs – with time-varying topology – can be a Python list of Tensor. (default: None)

  • edge_weight (Tensor, optional) – Weights of the edges (if edge_index is not a torch_sparse.SparseTensor). (default: None)

  • mask (Tensor, optional) – The optional mask associated with the target. (default: None)

  • transform (Mapping, optional) – Named mapping of Scaler associated with entries in input or output. (default: None)

  • pattern (Mapping, optional) – Map of the pattern of each entry in input or output. (default: None)

  • **kwargs – Any keyword argument for Data.

numpy(*args: List[str])[source]#

Transform all tensors to numpy arrays, either for all attributes or only the ones given in *args.

rearrange_element(key: str, pattern: str, **axes_lengths)[source]#

Rearrange key in Data according to the provided patter using einops.rearrange.

rearrange(patterns: Mapping)[source]#

Rearrange all keys in Data according to the provided pattern using einops.rearrange.

subgraph(subset: Tensor)[source]#

Returns the induced subgraph given by the node indices subset.

Parameters:

subset (LongTensor or BoolTensor) – The nodes to keep.

class StaticBatch(input: Optional[Mapping] = None, target: Optional[Mapping] = None, edge_index: Optional[Union[Tensor, SparseTensor]] = None, edge_weight: Optional[Tensor] = None, mask: Optional[Tensor] = None, transform: Optional[Mapping] = None, pattern: Optional[Mapping] = None, size: Optional[int] = None, **kwargs)[source]#

A batch of tsl.data.Data objects for multiple spatiotemporal graphs sharing the same topology.

The batch object extends Data, thus preserving all its functionalities.

Parameters:
  • input (Mapping, optional) – Named mapping of Tensor to be used as input to the model. (default: None)

  • target (Mapping, optional) – Named mapping of Tensor to be used as target of the task. (default: None)

  • edge_index (Adj, optional) – Shared graph connectivity, either in COO format (a Tensor of shape [2, E]) or as a torch_sparse.SparseTensor with shape [N, N]. (default: None)

  • edge_weight (Tensor, optional) – Weights of the edges (if edge_index is not a torch_sparse.SparseTensor). (default: None)

  • mask (Tensor, optional) – The optional mask associated with the target. (default: None)

  • transform (Mapping, optional) – Named mapping of Scaler associated with entries in input or output. (default: None)

  • pattern (Mapping, optional) – Map of the pattern of each entry in input or output. (default: None)

  • size (int, optional) – The batch size, i.e., the number of spatiotemporal graphs in the batch. The different samples in the batch share all the same topology, such that there is (at most) only one edge_index and edge_weight. If None, then the batch size is inferred from data (if possible). (default: None)

  • **kwargs – Any keyword argument for Data.

classmethod from_data_list(data_list: List[Data])[source]#

Constructs a Batch object from a Python list of Data representing temporal signals on a static (shared) graph.

get_example(idx: int) Data[source]#

Gets the Data or HeteroData object at index idx. The Batch object must have been created via from_data_list() in order to be able to reconstruct the initial object.

index_select(idx: Union[slice, List, Tuple, Tensor, ndarray]) List[Data][source]#

Creates a subset of Data objects from specified indices idx.

Indices idx can be a slicing object, e.g., [2:5], a list, a tuple, or a torch.Tensor or np.ndarray of type long or bool.

to_data_list() List[Data][source]#

Reconstructs the list of Data objects from the StaticBatch object.

property batch_size: int#

The batch size, i.e., the number of spatiotemporal graphs in the batch.

class DisjointBatch(*args: Any, **kwargs: Any)[source]#

A data object describing a batch of graphs as one big (disconnected) graph.

Inherits from tsl.data.Data. In addition, single graphs can be identified via the assignment vector batch, which maps each node to its respective graph identifier.

classmethod from_data_list(data_list: List[Data], force_batch: bool = False, follow_batch: Optional[List[str]] = None, exclude_keys: Optional[List[str]] = None, graph_attributes: Optional[List[str]] = None)[source]#

Constructs a DisjointBatch object from a list of Data objects.

The assignment vector batch is created on the fly. In addition, creates assignment vectors for each key in follow_batch. Will exclude any keys given in exclude_keys.

Parameters:
  • data_list (list) – The list of tsl.data.Data objects.

  • force_batch (bool) – If True, then add add dummy batch dimension for time-varying elements. (default: False)

  • follow_batch (list, optional) – Create an assignment vector for each key in follow_batch. (default: None)

  • exclude_keys (list, optional) – Exclude the keys in exclude_keys from collate. (default: None)

  • graph_attributes – Keys in graph_attributes with no node dimension will be added to the batch as graph attributes, i.e., the tensors will be stacked on a new dimension (the first one). Note that all graph attributes indexed by a key which is not in this list are repeated along a new node dimension (the second one, if the attribute is time-varying, otherwise the first one). (default: None)

get_example(idx: int) Data[source]#

Gets the Data or HeteroData object at index idx. The Batch object must have been created via from_data_list() in order to be able to reconstruct the initial object.

property batch_size: int#

Alias for num_graphs.