Torch Extension Modules (pyc) Python Interfaces

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.Cartesian.edge(edge_index, prediction, px, py, pz, e, include_zero) → dict

Params torch.tensor edge_index:

The graph topology (this is usually known as the edge_index)

Params torch.tensor prediction:

The edge classification prediction of the neural network.

Params torch.tensor px:

The node’s cartesian momentum component in the x-direction.

Params torch.tensor py:

The node’s cartesian momentum component in the y-direction.

Params torch.tensor pz:

The node’s cartesian momentum component in the z-direction.

Params torch.tensor e:

The node’s energy.

Params bool include_zero:

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

Params torch.tensor edge_index:

The graph topology (this is usually known as the edge_index)

Params torch.tensor prediction:

The edge classification prediction of the neural network.

Params torch.tensor pmc:

A compact version of the particle’s 4-vector.

Params bool include_zero:

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

Params torch.tensor edge_index:

The graph topology (this is usually known as the edge_index)

Params torch.tensor prediction:

The node classification prediction of the neural network.

Params torch.tensor px:

The node’s cartesian momentum component in the x-direction.

Params torch.tensor py:

The node’s cartesian momentum component in the y-direction.

Params torch.tensor pz:

The node’s cartesian momentum component in the z-direction.

Params torch.tensor e:

The node’s energy.

Params bool include_zero:

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

Params torch.tensor edge_index:

The graph topology (this is usually known as the edge_index)

Params torch.tensor prediction:

The node classification prediction of the neural network.

Params torch.tensor pmc:

A compact version of the particle’s 4-vector.

Params bool include_zero:

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

Params torch.tensor edge_index:

The graph topology (this is usually known as the edge_index)

Params torch.tensor prediction:

The edge classification prediction of the neural network.

Params torch.tensor pt:

The particle’s transverse momentum

Params torch.tensor eta:

The rapidity of the particle

Params torch.tensor phi:

The azimuthal compnent of the particle node

Params torch.tensor e:

The node’s energy.

Params bool include_zero:

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

Params torch.tensor edge_index:

The graph topology (this is usually known as the edge_index)

Params torch.tensor prediction:

The edge classification prediction of the neural network.

Params torch.tensor pmu:

A compact version of the particle’s 4-vector.

Params bool include_zero:

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

Params torch.tensor edge_index:

The graph topology (this is usually known as the edge_index)

Params torch.tensor prediction:

The node classification prediction of the neural network.

Params torch.tensor pt:

The particle’s transverse momentum

Params torch.tensor eta:

The rapidity of the particle

Params torch.tensor phi:

The azimuthal compnent of the particle node

Params torch.tensor e:

The node’s energy.

Params bool include_zero:

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

Params torch.tensor edge_index:

The graph topology (this is usually known as the edge_index)

Params torch.tensor prediction:

The node classification prediction of the neural network.

Params torch.tensor pmu:

A compact version of the particle’s 4-vector.

Params bool include_zero:

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}