Graph¶

An extension module which aims to reduce the need for writing redundant torch.scatter node and edge aggregation. Each of the functions highlighted below have the same output type and structure. These are outlined below:

clusters:
The aggregated sum of the 4-vectors
unique_sum:
A tensor holding the unique cluster masses for the predicted topology. This can be interpreted as the number of particles reconstructed in the particular graph. For instance, if the event contains n-particles, with n-signature masses, then in an ideal situation n-particles would be contained within this tensor, rather than m-constituent particles (children).
reverse_sum:
A reverse mapping of unique_sum back to each graph node.
node_sum:
A tensor holding the aggregated node’s 4-vector.

pyc.Graph.Base.unique_aggregation(clusters, node_features) → torch.tensor¶

An aggregation function used to sum node features without double summing any particular node pairs. Given the complexity of the operation in nominal PyTorch, this function is exclusive to CUDA tensors.

Parameters:

clusters (torch.tensor) – Cluster indices used to perform the summation over.
node_features (torch.tensor) – The feature vector of the node to sum.

pyc.Graph.Cartesian.edge(edge_index, prediction, px, py, pz, e, include_zero) → dict

Parameters:

edge_index (torch.tensor) – The graph topology (this is usually known as the edge_index)
prediction (torch.tensor) – The edge classification prediction of the neural network.
px (torch.tensor) – The node’s cartesian momentum component in the x-direction.
py (torch.tensor) – The node’s cartesian momentum component in the y-direction.
pz (torch.tensor) – The node’s cartesian momentum component in the z-direction.
e (torch.tensor) – The node’s energy.
include_zero (bool) – Whether to include predictions where the MLP predicts 0.

Return dict[torch.tensor]:

{ clusters : torch.tensor, unique_sum : torch.tensor, reverse_sum : torch.tensor, node_sum : torch.tensor}

pyc.Graph.Cartesian.edge(edge_index, prediction, pmc, include_zero) → dict¶

Parameters:

edge_index (torch.tensor) – The graph topology (this is usually known as the edge_index)
prediction (torch.tensor) – The edge classification prediction of the neural network.
pmc (torch.tensor) – A compact version of the particle’s 4-vector.
include_zero (bool) – Whether to include predictions where the MLP predicts 0.

Return dict[torch.tensor]:

{ clusters : torch.tensor, unique_sum : torch.tensor, reverse_sum : torch.tensor, node_sum : torch.tensor}

pyc.Graph.Cartesian.node(edge_index, prediction, px, py, pz, e, include_zero) → dict

Parameters:

edge_index (torch.tensor) – The graph topology (this is usually known as the edge_index)
prediction (torch.tensor) – The node classification prediction of the neural network.
px (torch.tensor) – The node’s cartesian momentum component in the x-direction.
py (torch.tensor) – The node’s cartesian momentum component in the y-direction.
pz (torch.tensor) – The node’s cartesian momentum component in the z-direction.
e (torch.tensor) – The node’s energy.
include_zero (bool) – Whether to include predictions where the MLP predicts 0.

Return dict[torch.tensor]:

{ clusters : torch.tensor, unique_sum : torch.tensor, reverse_sum : torch.tensor, node_sum : torch.tensor}

pyc.Graph.Cartesian.node(edge_index, prediction, pmc, include_zero) → dict¶

Parameters:

edge_index (torch.tensor) – The graph topology (this is usually known as the edge_index)
prediction (torch.tensor) – The node classification prediction of the neural network.
pmc (torch.tensor) – A compact version of the particle’s 4-vector.
include_zero (bool) – Whether to include predictions where the MLP predicts 0.

Return dict[torch.tensor]:

{ clusters : torch.tensor, unique_sum : torch.tensor, reverse_sum : torch.tensor, node_sum : torch.tensor}

pyc.Graph.Polar.edge(edge_index, prediction, pt, eta, phi, e, include_zero) → dict

Parameters:

edge_index (torch.tensor) – The graph topology (this is usually known as the edge_index)
prediction (torch.tensor) – The edge classification prediction of the neural network.
pt (torch.tensor) – The particle’s transverse momentum
eta (torch.tensor) – The rapidity of the particle
phi (torch.tensor) – The azimuthal compnent of the particle node
e (torch.tensor) – The node’s energy.
include_zero (bool) – Whether to include predictions where the MLP predicts 0.

Return dict[torch.tensor]:

{ clusters : torch.tensor, unique_sum : torch.tensor, reverse_sum : torch.tensor, node_sum : torch.tensor}

pyc.Graph.Polar.edge(edge_index, prediction, pmu, include_zero) → dict¶

Parameters:

edge_index (torch.tensor) – The graph topology (this is usually known as the edge_index)
prediction (torch.tensor) – The edge classification prediction of the neural network.
pmu (torch.tensor) – A compact version of the particle’s 4-vector.
include_zero (bool) – Whether to include predictions where the MLP predicts 0.

Return dict[torch.tensor]:

{ clusters : torch.tensor, unique_sum : torch.tensor, reverse_sum : torch.tensor, node_sum : torch.tensor}

pyc.Graph.Polar.node(edge_index, prediction, pt, eta, phi, e, include_zero) → dict

Parameters:

edge_index (torch.tensor) – The graph topology (this is usually known as the edge_index)
prediction (torch.tensor) – The node classification prediction of the neural network.
pt (torch.tensor) – The particle’s transverse momentum
eta (torch.tensor) – The rapidity of the particle
phi (torch.tensor) – The azimuthal compnent of the particle node
e (torch.tensor) – The node’s energy.
include_zero (bool) – Whether to include predictions where the MLP predicts 0.

Return dict[torch.tensor]:

{ clusters : torch.tensor, unique_sum : torch.tensor, reverse_sum : torch.tensor, node_sum : torch.tensor}

pyc.Graph.Polar.node(edge_index, prediction, pmu, include_zero) → dict¶

Parameters:

edge_index (torch.tensor) – The graph topology (this is usually known as the edge_index)
prediction (torch.tensor) – The node classification prediction of the neural network.
pmu (torch.tensor) – A compact version of the particle’s 4-vector.
include_zero (bool) – Whether to include predictions where the MLP predicts 0.

Return dict[torch.tensor]:

{ clusters : torch.tensor, unique_sum : torch.tensor, reverse_sum : torch.tensor, node_sum : torch.tensor}