spine.data.batch.EdgeIndexBatch

class spine.data.batch.EdgeIndexBatch(data: Any, counts: Sequence[int] | Any, spans: Sequence[int] | Any, directed: bool)[source]

Batched edge index with the necessary methods to slice it.

spans

Per-entry parent spans used to build the batch offsets.

Type:: Union[np.ndarray, torch.Tensor]

offsets

(B) Offsets between successive indexes in the batch, computed from the cumulative sum of spans.

Type:: Union[np.ndarray, torch.Tensor]

directed

Whether the edge index is directed or undirected

Type:: bool

Attributes:

batch_ids: Returns the batch ID of each element in the full index list.
directed_batch_ids: Returns the batch ID of each element in the directed index.
directed_counts: Returns the number of edges per entry, counting edges once even if they are bidirectional.
directed_index: Index of the directed graph.
directed_index_t: Index of the directed graph, transposed.
index: Alias for the underlying data stored.
index_t: Alias for the underlying data stored, transposed
shape: Shape of the underlying data.
splits: Boundaries needed to split the data into its constituents.

Methods

`get_counts`(batch_ids, batch_size)	Finds the number of elements in each entry, provided a batch ID list.
`get_edges`(counts)	Finds the edges between successive entries in the batch.
`split`()	Breaks up the index batch into its constituents.
`to_numpy`()	Cast underlying index to a np.ndarray and return a new instance.
`to_tensor`([dtype, device])	Cast underlying index to a torch.tensor and return a new instance.

__init__(data: Any, counts: Sequence[int] | Any, spans: Sequence[int] | Any, directed: bool) → None[source]

Initialize the attributes of the class.

If the edge index corresponds to an undirected graph, each edge should have its reciprocal edge immediately after, e.g.

[[0,1,0,2,0,3,...],
 [1,0,2,0,3,0,...]]

Parameters:

data (Union[np.ndarray, torch.Tensor]) – (2, E) Batched edge index
counts (Union[List[int], np.ndarray, torch.Tensor]) –
1. Number of index elements per entry in the batch
spans (Union[List[int], np.ndarray, torch.Tensor]) –
1. Per-entry parent spans used to derive offsets.
directed (bool) – Whether the edge index is directed or undirected

Methods

`__init__`(data, counts, spans, directed)	Initialize the attributes of the class.
`get_counts`(batch_ids, batch_size)	Finds the number of elements in each entry, provided a batch ID list.
`get_edges`(counts)	Finds the edges between successive entries in the batch.
`split`()	Breaks up the index batch into its constituents.
`to_numpy`()	Cast underlying index to a np.ndarray and return a new instance.
`to_tensor`([dtype, device])	Cast underlying index to a torch.tensor and return a new instance.

Attributes

`batch_ids`	Returns the batch ID of each element in the full index list.
`directed_batch_ids`	Returns the batch ID of each element in the directed index.
`directed_counts`	Returns the number of edges per entry, counting edges once even if they are bidirectional.
`directed_index`	Index of the directed graph.
`directed_index_t`	Index of the directed graph, transposed.
`index`	Alias for the underlying data stored.
`index_t`	Alias for the underlying data stored, transposed
`shape`	Shape of the underlying data.
`splits`	Boundaries needed to split the data into its constituents.
`data`
`counts`
`edges`
`batch_size`
`spans`
`offsets`
`directed`

data: Any

counts: Any

spans: Any

edges: Any

offsets: Any

directed: bool

batch_size: int

property index: Any

Alias for the underlying data stored.

Returns:: (2, E) Underlying batch of edge indexes
Return type:: Union[np.ndarray, torch.Tensor]

property index_t: Any

Alias for the underlying data stored, transposed

Returns:: (E, 2) Underlying batch of edge indexes, transposed
Return type:: Union[np.ndarray, torch.Tensor]

property batch_ids: Any

Returns the batch ID of each element in the full index list.

Returns:

Complete batch ID array, one per element

Return type:

Union[np.ndarray, torch.Tensor]

property directed_index: Any

Index of the directed graph. If a graph is undirected, it only returns one of the two edges corresponding to a connection.

Returns:: (2, E//2) Underlying batch of edge indexes
Return type:: Union[np.ndarray, torch.Tensor]

property directed_index_t: Any

Index of the directed graph, transposed. If the graph is undirected, it only returns one of the two edges corresponding to a connection.

Returns:: (E//2, 2) Underlying batch of edge indexes, transposed
Return type:: Union[np.ndarray, torch.Tensor]

property directed_counts: Any

Returns the number of edges per entry, counting edges once even if they are bidirectional.

Returns:

Complete batch ID array, one per element

Return type:

Union[np.ndarray, torch.Tensor]

property directed_batch_ids: Any

Returns the batch ID of each element in the directed index.

Returns:

Complete batch ID array, one per element

Return type:

Union[np.ndarray, torch.Tensor]

split() → list[Any][source]

Breaks up the index batch into its constituents.

Returns:: List of one index per entry in the batch
Return type:: List[Union[np.ndarray, torch.Tensor]]

to_numpy() → EdgeIndexBatch[source]

Cast underlying index to a np.ndarray and return a new instance.

Returns:: New TensorBatch object with an underlying np.ndarray tensor.
Return type:: TensorBatch

to_tensor(dtype: Any = None, device: Any = None) → EdgeIndexBatch[source]

Cast underlying index to a torch.tensor and return a new instance.

Parameters:

dtype (torch.dtype, optional) – Data type of the tensor to create
device (torch.device, optional) – Device on which to put the tensor

Returns:

New TensorBatch object with an underlying np.ndarray tensor.

Return type:

TensorBatch