pipeline.GNN package

This package define the Graph Neural Network which allows to classify the edges and triplets in the graph.

pipeline.GNN.gnn_plots module

pipeline.GNN.gnn_plots.plot_best_performances_score_cut(model, partition, edge_score_cuts, builder='default', n_events=None, seed=None, identifier=None, path_or_config=None, step='gnn', **kwargs)

Return type:

Tuple[Figure | npt.NDArray, List[Axes], Dict[float, Dict[Tuple[str | None, str], float]]]

pipeline.GNN.gnn_plots.plot_best_performances_score_cut_triplets(model, partition, edge_score_cut, triplet_score_cuts, n_events=None, seed=None, identifier=None, path_or_config=None, step='gnn', **kwargs)

Return type:

Tuple[Figure | npt.NDArray, List[Axes], Dict[float, Dict[Tuple[str | None, str], float]]]

pipeline.GNN.gnn_validation module

class pipeline.GNN.gnn_validation.GNNScoreCutExplorer(model, builder='default')

Bases: ParamExplorer

A class that allows to vary the score cut after the GNN, and compare the metric performances of track finding.

property default_step: str

Name of the temp to fall back to if not provided.

get_tracks(value, batches)

Get the dataframe of tracks from the inferred batches.

Parameters:

• value (float) – current value of the parameter that is explored

• batches (List[Data]) – list of inferred batches

Returns:

Dataframe of tracks, with columns track_id and hit_id

run_inference(batches)

Run the inference on a batch.

Parameters:

batches (Union[List[Data], LazyDatasetBase]) – List of batches

Returns:

List of inferred batches

class pipeline.GNN.gnn_validation.TripletGNNScoreCutExplorer(model)

Bases: ParamExplorer

A class that allows to vary the score cut after the GNN, and compare the metric performances of track finding.

add_lhcb_text(ax, metric_name)

get_tracks(value, batches)

Get the dataframe of tracks from the inferred batches.

Parameters:

• value (float) – current value of the parameter that is explored

• batches (List[Data]) – list of inferred batches

Returns:

Dataframe of tracks, with columns track_id and hit_id

run_inference(batches)

Run the inference on a batch.

Parameters:

batches (Union[List[Data], LazyDatasetBase]) – List of batches

Returns:

List of inferred batches

pipeline.GNN.perfect_gnn module

Replace the GNN by a perfect inference in order to understand what is the best result that can be obtained with the current pipeline.

class pipeline.GNN.perfect_gnn.PerfectInferenceBuilder

Bases: BuilderBase

Generate perfect inference, that is, the edge score is equal to the truth.

construct_downstream(batch, pid=False)

Run the inference on a PyTorch Data. In-place.

class pipeline.GNN.perfect_gnn.PerfectTripletInferenceBuilder

Bases: BuilderBase

construct_downstream(batch)

Run the inference on a PyTorch Data. In-place.

load_batch(input_path)

Load a PyTorch Data object from its path. Might apply necessary pre-processing.

Return type:

Data

pipeline.GNN.triplet_gnn_base module

A module that define TripletGNNBase, the base class of all triplet-based GNNs in this repository.

class pipeline.GNN.triplet_gnn_base.TripletGNNBase(*args: Any, **kwargs: Any)

Bases: ModelBase

The base class for triplet-base models, that first classify edges, then triplets.

common_training_validation_step(batch, edge_score_cut=None, with_triplets=None, compute_loss=False)

Common forward step and loss computation for the training and validation steps.

Parameters:

• batch (Data) – event graph

• with_triplets (Optional[bool]) – whether to include the forward step on triplets

• edge_score_cut (Optional[float]) – minimal edge score the edges are required to have

Return type:

Dict[str, Any]

Returns:

Output of the forward step and loss computation.

compute_normalised_loss(output, truth)

Compute typical weighted focal loss for given output and truth.

Parameters:

• output (Tensor) – logits

• truth (Tensor) – targets

Return type:

Tensor

Returns:

Normalised sigmoid focal loss.

property edge_checkpointing: bool

filter_edges(edge_index, edge_score, edge_score_cut=None)

Return type:

Tuple[Tensor, Tensor]

forward(x, edge_index, edge_score_cut=None, with_triplets=True)

Forward step of the triplet-based Neural Network.

1. forward_edges() method is called, and outputs the edge logits edge_output, with possibly other tensors that can be used for the triplets.

2. The edges are filtered using the self.filter_edges_for_triplets() method.

3. The triplets are built using utils.graphutils.tripletbuilding.from_edge_index_to_triplet_indices()

4. forward_triplets() method is called, and outputs the triplets logits triplet_outputs.

Parameters:

• x (Tensor) – node features

• edge_index (Tensor) – tensor with shape (2, n_edges) of the edge indices

• edge_score_cut (Optional[float]) – Minal edge score

• with_triplets (bool) – whether to include the triplet inference

Return type:

Dict[str, Any]

forward_edges(x, start, end)

Forward step for edge classification.

Parameters:

• x (Tensor) – Hit features

• start (Tensor) – tensor of start indices

• end (Tensor) – tensor of edge indices

Return type:

Dict[str, Tensor]

Returns:

A dictionary of tensors. Should at least contain edge_output, the logits of each edges.

forward_triplets(dict_triplet_indices, *args, **kwargs)

Forward step for triplet building and classification.

Parameters:

• dict_triplet_indices (Dict[str, Tensor]) – associates articulation, elbow_left and elbow_right with the corresponding triplet indices.

• args – Other arguments to pass to the triplet output step.

• kwargs – Other arguments to pass to the triplet output step.

Return type:

Dict[str, Tensor]

Returns:

A dictionary that associates articulation, elbow_left and elbow_right with the logits of the corresponding triplets.

get_lazy_dataset(*args, **kwargs)

Get the lazy dataset object.

Parameters:

• input_dir – input directory

• n_events – number of events to load

• shuffle – whether to shuffle the input paths (applied before selected the first n_events)

• seed – seed for the shuffling

• **kwargs – Other keyword arguments passed to the utils.loaderutils.dataiterator.LazyDatasetBase constructor.

Return type:

TripletGNNLazyDataset

Returns:

utils.loaderutils.dataiterator.LazyDatasetBase object

get_lazy_dataset_partition(partition, *args, **kwargs)

Get the lazy dataset of a partition.

Parameters:

• partition (str) – train, val or name of the test dataset

• n_events – number of events to load

• shuffle – whether to shuffle the input paths (applied before selected the first n_events)

• seed – seed for the shuffling

• **kwargs – Other keyword arguments passed to ModelBase.get_lazy_dataset()

Return type:

LazyDatasetBase

Returns:

Lazy dataset of the partition

inference(batch, with_triplets=True, with_triplet_truths=False, edge_score_cut=None)

Run inference (without loss computation).

Parameters:

• batch (Data) – event graph

• with_triplets (bool) – whether to include the forward step on triplets

• edge_score_cut (Optional[float]) – minimal edge score the edges are required to have

Return type:

Dict[str, Any]

Returns:

Output of the forward step.

property input_kwargs: Dict[str, Any]

Associates an input name with a dictionary corresponding to the keyword arguments used to build a dummy tensor representing the input. This dictionary basically gives the size and dtype of the tensor.

property input_to_dynamic_axes

A dictionary that associates an input name with the dynamic axis specification.

log_metrics_gen(loss, scores, predictions, truths, suffix='')

Add entry to the log.

Parameters:

• loss (Tensor) – overall loss

• scores (Tensor) – edge or triplet scores. Used to compute the AUC

• predictions (Tensor) – edge or triplet predicted targets

• truths (Tensor) – edge or triplet targets

• suffix (str) – optional suffix, e.g., _edge or _triplet

Return type:

None

property loss: str

property n_hiddens: int

Number of hidden units

shared_evaluation(batch, log=False, with_triplets=None)

Evaluation step. Can be used for validation and test.

Parameters:

• batch (Data) – event graph

• log (bool) – whether to add an entry to the log

• with_triplets (Optional[bool]) – whether to include the triplet inference

property subnetwork_to_outputs: Dict[str, List[str]]

A dictionary that associates a subnetwork name with the list of its output names.

test_step(batch, batch_idx)

to_onnx(outpath, mode=None, options=None)

Export the model to ONNX

Parameters:

• outpath (str) – path to the ONNX output file

• mode (Optional[str]) – subnetwork to save

Return type:

None

training_step(batch, batch_idx)

Training step.

property triplet_checkpointing: bool

triplet_output_step(dict_triplet_indices, *args, **kwargs)

Return type:

Dict[str, Tensor]

triplet_output_step_articulation(triplet_indices, *args, **kwargs)

triplet_output_step_elbow_left(triplet_indices, *args, **kwargs)

triplet_output_step_elbow_right(triplet_indices, *args, **kwargs)

validation_step(batch, batch_idx)

property with_triplets: bool

class pipeline.GNN.triplet_gnn_base.TripletGNNLazyDataset(input_dir, n_events=None, shuffle=False, seed=None, **kwargs)

Bases: LazyDatasetBase

fetch_dataset(input_path, **kwargs)

Load and process one PyTorch DataSet.

Parameters:

• input_path (str) – path to the PyTorch dataset

• map_location – location where to load the dataset

• **kwargs – Other keyword arguments passed to torch.load()

Return type:

Data

Returns:

Load PyTorch data object

pipeline.GNN.triplet_gnn_base.get_df_edges_from_batch_only(batch)

Get the dataframe of edges with particle ID information.

Return type:

TypeVar(DataFrame, DataFrame, DataFrame)

pipeline.GNN.models package

A package that defines various triplet-based GNNs.

pipeline.GNN.models.get_model(model_type=None)

Get a GNN model class from its name.

Parameters:

model_type (Optional[str]) – GNN type. Available are triplet_interaction, edge_based and scifi_triplet_interaction.

Return type:

Type[TripletGNNBase]

Returns:

The GNN model class that can be instantiated.

pipeline.GNN.models.edge_based_gnn module

class pipeline.GNN.models.edge_based_gnn.EdgeBasedGNN(hparams)

Bases: TripletGNNBase

A triplet-based GNN without using node encodings.

forward_edges(x, start, end)

Forward step for edge classification.

Parameters:

• x (Tensor) – Hit features

• edge_index – Torch tensor with 2 rows that define the edges.

Returns:

the hit encodings and edge encodings after message passing, and the edge classifier output.

Return type:

A tuple of 3 tensors

forward_triplets(e, filtered_edge_index, edge_mask, dict_triplet_indices, **kwargs)

Forward step for triplet classification.

Parameters:

• e (Tensor) – Edge encodings after the edge forward step

• filtered_edge_index (Tensor) – edge index after requiring the minimal edge score

• edge_mask (Tensor) – edge mask of edges kept after filtering

• edge_score_cut – custom edge score cut. If not specified, the value is taken from the internal hparams dictionary.

Return type:

Dict[str, Tensor]

Returns:

A dictionary that associates articulation, elbow_left and elbow_right with the logits of the corresponding triplets.

property input_kwargs: Dict[str, Any]

Associates an input name with a dictionary corresponding to the keyword arguments used to build a dummy tensor representing the input. This dictionary basically gives the size and dtype of the tensor.

property input_to_dynamic_axes

A dictionary that associates an input name with the dynamic axis specification.

message_step(e, start, end, dim_size)

Return type:

Tensor

property subnetwork_groups: Dict[str, List[str]]

A dictionary that associates a subnetwork actually corresponding to a list of subnetworks, with this list of subnetworks.

property subnetwork_to_outputs: Dict[str, List[str]]

A dictionary that associates a subnetwork name with the list of its output names.

property subnetworks: List[str]

List of subnetworks available. It is derived from subnetwork_to_outputs.

triplet_output_step_articulation(e, edge_indices, triplet_indices)

triplet_output_step_elbow_left(e, edge_indices, triplet_indices)

triplet_output_step_elbow_right(e, edge_indices, triplet_indices)

pipeline.GNN.models.scifi_triplet_interaction_gnn module

pipeline.GNN.models.triplet_interaction_gnn module

class pipeline.GNN.models.triplet_interaction_gnn.TripletInteractionGNN(*args: Any, **kwargs: Any)

Bases: TripletGNNBase

A triplet-based interaction network.

encoding_step(x, start, end)

Initial encoding step of the GNN.

Parameters:

• x (Tensor) – hit input features

• start (Tensor) – start indices of the edges

• end (Tensor) – end indices of the edges

Return type:

Tuple[Tensor, Tensor]

Returns:

The node encodings h and edge encodings e

forward_edges(x, start, end)

Forwrd step for edge classification.

Parameters:

• x (Tensor) – hit input features

• start (Tensor) – start indices of the edges

• end (Tensor) – end indices of the edges

Returns:

the hit encodings and edge encodings after message passing, and the edge classifier output.

Return type:

A tuple of 3 tensors

forward_triplets(h, e, filtered_edge_index, edge_mask, dict_triplet_indices, **kwargs)

Forward step for triplet classification.

Parameters:

• h (Tensor) – Hit encodings after the edge forward step

• e (Tensor) – Edge encodings after the edge forward step

• filtered_edge_index (Tensor) – edge index after requiring the minimal edge score

• edge_mask (Tensor) – edge mask of edges kept after filtering

• edge_score_cut – custom edge score cut. If not specified, the value is taken from the internal hparams dictionary.

Return type:

Dict[str, Tensor]

Returns:

A dictionary that associates articulation, elbow_left and elbow_right with the logits of the corresponding triplets.

property input_kwargs: Dict[str, Any]

Associates an input name with a dictionary corresponding to the keyword arguments used to build a dummy tensor representing the input. This dictionary basically gives the size and dtype of the tensor.

property input_to_dynamic_axes

A dictionary that associates an input name with the dynamic axis specification.

message_step(h, start, end, e, step)

Apply one step of message-passing that updates the node and edge encodings.

Return type:

Tuple[Tensor, Tensor]

property n_graph_iters: int

Number of message-passing iterations

property only_e: bool

output_step(h, start, end, e)

Apply the edge output classifier to edges to get edge logits.

Return type:

Tensor

property recursive: bool

Whether the GNN is recursive, i.e., the networks involved in the message step are the same.

scatter_add(source, index, h)

Scatter add operation. In ONNX export mode for TensorRT, the operation is replaced by a fake operator that is implemented through a plugin in TensorRT.

Return type:

Tensor

scatter_max(source, index, h)

property subnetwork_groups: Dict[str, List[str]]

A dictionary that associates a subnetwork actually corresponding to a list of subnetworks, with this list of subnetworks.

property subnetwork_to_outputs: Dict[str, List[str]]

A dictionary that associates a subnetwork name with the list of its output names.

to_onnx(outpath, mode=None, options=None)

Export the model to ONNX

Parameters:

• outpath (str) – path to the ONNX output file

• mode (Optional[str]) – subnetwork to save

Return type:

None

triplet_output_step_articulation(h, e, edge_indices, triplet_indices)

triplet_output_step_elbow_left(h, e, edge_indices, triplet_indices)

triplet_output_step_elbow_right(h, e, edge_indices, triplet_indices)