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.