graphnet.models.components.pool module

Functions for performing pooling/clustering/coarsening.

graphnet.models.components.pool.min_pool(cluster, data, transform)[source]

Perform min-pooling of Data.

Like max_pool, just negating `data.x.

Return type:

Data

Parameters:

  • cluster (LongTensor)

  • data (Data)

  • transform (Any | None)

graphnet.models.components.pool.min_pool_x(cluster, x, batch, size)[source]

Perform min-pooling of Tensor.

Like max_pool_x, just negating `x.

Return type:

Tensor

Parameters:

  • cluster (LongTensor)

  • x (Tensor)

  • batch (LongTensor)

  • size (int | None)

graphnet.models.components.pool.sum_pool_and_distribute(tensor, cluster_index, batch)[source]

Sum-pool values and distribute result to the individual nodes.

Return type:

Tensor

Parameters:

  • tensor (Tensor)

  • cluster_index (LongTensor)

  • batch (LongTensor | None)

graphnet.models.components.pool.group_by(data, keys)[source]

Group nodes in data that have identical values of keys.

This grouping is done with in each event in case of batching. This allows for, e.g., assigning the same index to all pulses on the same PMT or DOM in the same event. This can be used for coarsening graphs, e.g., from pulse- level to DOM-level by aggregating feature across each group returned by this method.

Return type:

LongTensor

Parameters:

  • data (Data | Batch)

  • keys (List[str])

Example

Given:

data.f1 = [1,1,2,2,2] data.f2 = [6,7,7,7,8]

Calls:

groupby(data, [‘f1’]) -> [0, 0, 1, 1, 1] groupby(data, [‘f2’]) -> [0, 1, 1, 1, 2] groupby(data, [‘f1’, ‘f2’]) -> [0, 1, 2, 2, 3]

graphnet.models.components.pool.group_pulses_to_dom(data)[source]

Group pulses on the same DOM, using DOM and string number.

Return type:

Data

Parameters:

data (Data)

graphnet.models.components.pool.group_pulses_to_pmt(data)[source]

Group pulses on the same PMT, using PMT, DOM, and string number.

Return type:

Data

Parameters:

data (Data)

graphnet.models.components.pool.sum_pool_x(cluster, x, batch, size)[source]

Sum-pool node features according to the clustering defined in cluster.

Parameters:

  • cluster (LongTensor) – Cluster vector \(\mathbf{c} \in \{ 0, \ldots, N - 1 \}^N\), which assigns each node to a specific cluster.

  • x (Tensor) – Node feature matrix \(\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}\).

  • batch (LongTensor) – Batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which assigns each node to a specific example.

  • size (Optional[int], default: None) – The maximum number of clusters in a single example. This property is useful to obtain a batch-wise dense representation, e.g. for applying FC layers, but should only be used if the size of the maximum number of clusters per example is known in advance.

Return type:

Tensor

graphnet.models.components.pool.std_pool_x(cluster, x, batch, size)[source]

Std-pool node features according to the clustering defined in cluster.

Parameters:

  • cluster (LongTensor) – Cluster vector \(\mathbf{c} \in \{ 0, \ldots, N - 1 \}^N\), which assigns each node to a specific cluster.

  • x (Tensor) – Node feature matrix \(\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}\).

  • batch (LongTensor) – Batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which assigns each node to a specific example.

  • size (Optional[int], default: None) – The maximum number of clusters in a single example. This property is useful to obtain a batch-wise dense representation, e.g. for applying FC layers, but should only be used if the size of the maximum number of clusters per example is known in advance.

Return type:

Tensor

graphnet.models.components.pool.sum_pool(cluster, data, transform)[source]

Pool and coarsen graph according to the clustering defined in cluster.

All nodes within the same cluster will be represented as one node. Final node features are defined by the sum of features of all nodes within the same cluster, node positions are averaged and edge indices are defined to be the union of the edge indices of all nodes within the same cluster.

Parameters:

  • cluster (LongTensor) – Cluster vector \(\mathbf{c} \in \{ 0, \ldots, N - 1 \}^N\), which assigns each node to a specific cluster.

  • data (Data) – Graph data object.

  • transform (Optional[Callable], default: None) – A function/transform that takes in the coarsened and pooled torch_geometric.data.Data object and returns a transformed version.

Return type:

Data

graphnet.models.components.pool.std_pool(cluster, data, transform)[source]

Pool and coarsen graph according to the clustering defined in cluster.

All nodes within the same cluster will be represented as one node. Final node features are defined by the std of features of all nodes within the same cluster, node positions are averaged and edge indices are defined to be the union of the edge indices of all nodes within the same cluster.

Parameters:

  • cluster (LongTensor) – Cluster vector \(\mathbf{c} \in \{ 0, \ldots, N - 1 \}^N\), which assigns each node to a specific cluster.

  • data (Data) – Graph data object.

  • transform (Optional[Callable], default: None) – A function/transform that takes in the coarsened and pooled torch_geometric.data.Data object and returns a transformed version.

Return type:

Data