utils

Utility functions for graphnet.models.

graphnet.models.utils.calculate_xyzt_homophily(x, edge_index, batch)

Calculate xyzt-homophily from a batch of graphs.

Homophily is a graph scalar quantity that measures the likeness of variables in nodes. Notice that this calculator assumes a special order of input features in x.

Return type:

Tuple[Tensor, Tensor, Tensor, Tensor]

Returns:

Tuple, each element with shape [batch_size,1].

Parameters:

• x (Tensor)

• edge_index (LongTensor)

• batch (Batch)

graphnet.models.utils.calculate_distance_matrix(xyz_coords)

Calculate the matrix of pairwise distances between pulses.

Parameters:

xyz_coords (Tensor) – (x,y,z)-coordinates of pulses, of shape [nb_doms, 3].

Return type:

Tensor

Returns:

Matrix of pairwise distances, of shape [nb_doms, nb_doms]

graphnet.models.utils.knn_graph_batch(batch, k, columns)

Calculate k-nearest-neighbours with individual k for each batch event.

Parameters:

• batch (Batch) – Batch of events.

• k (List[int]) – A list of k’s.

• columns (List[int]) – The columns of Data.x used for computing the distances. E.g., Data.x[:,[0,1,2]]

Return type:

Batch

Returns:

Returns the same batch of events, but with updated edges.

graphnet.models.utils.array_to_sequence(x, batch_idx, padding_value, excluding_value)

Convert x of shape [n,d] into a padded sequence of shape [B, L, D].

Where B is the batch size, L is the sequence length and D is the features for each time step.

Parameters:

• x (Tensor) – array-like tensor with shape [n,d] where n is the total number

• features. (of pulses in the batch and d is the number of node)

• batch_idx (LongTensor) – a LongTensor identifying which row in x belongs to which training example. E.g. torch_geometric.data.Batch.batch.

• padding_value (Any, default: 0) – The value to use for padding.

• excluding_value (Any, default: inf) – This parameter represents a unique value that should not be present in the input tensor ‘x’

Returns:

Padded sequence with dimensions [B, L, D]. mask: A tensor that identifies masked entries in x.

E.g. : masked_entries = x[mask]

seq_length: A tensor containing the number of pulses in each event.

Return type:

x

graphnet.models.utils.get_fields(data, fields)

Extract named fields in Data object.

Return type:

Tensor

Parameters:

• data (Data | List[Data])

• fields (List[str])

graphnet.models.utils.full_edge_index(edge_index, batch)

Return the full batched sparse adjacency matrices given by edge indices.

Return batched sparse adjacency matrices with exactly those edges that are not in the input edge_index while ignoring self-loops. Implementation inspired by torch_geometric.utils.to_dense_adj.

Original code: https://github.com/LiamMa/GRIT/blob/main/grit/encoder/rrwp_encoder.py

Parameters:

• edge_index (Tensor) – The edge indices.

• batch (Optional[Tensor], default: None) – Batch vector, which assigns each node to a specific example.

Return type:

Tensor

Returns:

Complementary edge index.

graphnet.models.utils.add_full_rrwp(data, walk_length, attr_name_abs, attr_name_rel, add_identity, spd)

Add relative random walk probabilities.

Original code: https://github.com/LiamMa/GRIT/blob/main/grit/transform/rrwp.py

Parameters:

• data (Data) – Input data.

• walk_length (int, default: 8) – Number of random walks for encoding.

• attr_name_abs (str, default: 'rrwp') – Absolute position encoding name.

• attr_name_rel (str, default: 'rrwp') – Relative position encoding name.

• add_identity (bool, default: True) – Add identity matrix to position encoding.

• spd (bool, default: False) – Use shortest path distances.

Return type:

Data

graphnet.models.utils.get_rw_landing_probs(ksteps, edge_index, edge_weight, num_nodes, space_dim)

Compute Random Walk landing probabilities for given list of K steps.

Original code: https://github.com/ETH-DISCO/Benchmarking-PEs :type ksteps: List :param ksteps: List of k-steps for which to compute the RW landings :type edge_index: Tensor :param edge_index: PyG sparse representation of the graph :type edge_weight: Optional[Tensor], default: None :param edge_weight: (optional) Edge weights :type num_nodes: Optional[int], default: None :param num_nodes: (optional) Number of nodes in the graph :type space_dim: int, default: 0 :param space_dim: (optional) Estimated dimensionality of the space. Used to

correct the random-walk diagonal by a factor k^(space_dim/2). In euclidean space, this correction means that the height of the gaussian distribution stays almost constant across the number of steps, if space_dim is the dimension of the euclidean space.

Return type:

Tensor

Returns:

2D Tensor with shape (num_nodes, len(ksteps)) with RW landing probs

Parameters:

• ksteps (List)

• edge_index (Tensor)

• edge_weight (Tensor | None)

• num_nodes (int | None)

• space_dim (int)