graphnet.models.utils module

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])