lloca.utils.utils

Tools for graph construction and manipulation.

Functions

get_batch_from_ptr(ptr)	Reconstruct batch indices (batch) from pointer (ptr).
get_edge_attr(fourmomenta, edge_index[, ...])	Calculate edge attributes based on the squared Lorentz norm of the sum of four-momenta.
get_edge_index_from_ptr(ptr, shape[, ...])	Construct edge index of fully connected graph from pointer (ptr).
get_edge_index_from_shape(shape, device[, ...])	Construct edge index of fully connected graph from the shape of a corresponding dense tensor.
get_node_to_edge_ptr_fully_connected(ptr, batch)	Get pointer (ptr) mapping nodes to edges in a fully connected graph.
get_ptr_from_batch(batch)	Reconstruct pointer (ptr) from batch indices (batch).

lloca.utils.utils.get_batch_from_ptr(ptr)

Reconstruct batch indices (batch) from pointer (ptr).

Parameters:

ptr (torch.Tensor) – Pointer tensor indicating the start of each batch. Tensor of shape (B+1,) where B is the number of batches.

Returns:

A tensor where each element indicates the batch index for each item. Tensor of shape (N,) where N is the total number of items across all batches.

Return type:

torch.Tensor

lloca.utils.utils.get_edge_attr(fourmomenta, edge_index, eps=1e-10, use_float64=True)

Calculate edge attributes based on the squared Lorentz norm of the sum of four-momenta.

Parameters:

• fourmomenta (torch.Tensor) – A tensor of shape (B, N, 4) representing the four-momenta of particles.

• edge_index (torch.Tensor) – A tensor of shape (2, E) representing the edge index of the graph.

• eps (float, optional) – A small value to avoid log(0) issues, by default 1e-10.

• use_float64 (bool, optional) – Whether to use float64 precision for calculations, by default True.

Returns:

A tensor of shape (E,) representing the edge attributes, which are the logarithm of the squared Lorentz norm.

Return type:

torch.Tensor

lloca.utils.utils.get_edge_index_from_ptr(ptr, shape, remove_self_loops=True)

Construct edge index of fully connected graph from pointer (ptr). This function should be used for graphs represented by sparse tensors, i.e. graphs where the number of nodes per graph can vary.

Parameters:

• ptr (torch.Tensor) – Pointer tensor indicating the start of each batch. Tensor of shape (B+1,) where B is the number of batches.

• shape (torch.Size) – Shape of the node tensor, expected to be (N, C)

• remove_self_loops (bool, optional) – Whether to remove self-loops from the edge index, by default True.

Returns:

A tensor of shape (2, E) where E is the number of edges, representing the edge index.

Return type:

torch.Tensor

lloca.utils.utils.get_edge_index_from_shape(shape, device, remove_self_loops=True)

Construct edge index of fully connected graph from the shape of a corresponding dense tensor. This function should be used for graphs represented by dense tensors, i.e. graphs where the number of nodes per graph is fixed.

Parameters:

• shape (torch.Size) – Shape of the dense node tensor, expected to be (B, N, C) where B is the batch size and N is the number of nodes.

• device (torch.device) – Device on which the tensors are allocated.

• remove_self_loops (bool, optional) – Whether to remove self-loops from the edge index, by default True.

Returns:

A tensor of shape (2, E) where E is the number of edges, representing the edge index.

Return type:

torch.Tensor

lloca.utils.utils.get_node_to_edge_ptr_fully_connected(ptr, batch, remove_self_loops=True)

Get pointer (ptr) mapping nodes to edges in a fully connected graph.

Parameters:

• ptr (torch.Tensor) – Pointer tensor indicating the start of each batch. Tensor of shape (B+1,) where B is the number of batches.

• batch (torch.Tensor) – A tensor where each element indicates the batch index for each node. Tensor of shape (N,) where N is the total number of nodes across all batches.

• remove_self_loops (bool) – Whether self-loops were removed when constructing the edge index, by default True.

Returns:

A pointer tensor mapping nodes to edges in a fully connected graph. Tensor of shape (N,) where N is the total number of nodes across all batches.

Return type:

torch.Tensor

lloca.utils.utils.get_ptr_from_batch(batch)

Reconstruct pointer (ptr) from batch indices (batch).

Parameters:

batch (torch.Tensor) – A tensor where each element indicates the batch index for each item. Tensor of shape (N,) where N is the total number of items across all batches.

Returns:

A pointer tensor indicating the start of each batch. Tensor of shape (B+1,) where B is the number of batches.

Return type:

torch.Tensor