graphnet_module

Class(es) for deploying GraphNeT models in icetray as I3Modules.

class graphnet.deployment.i3modules.graphnet_module.GraphNeTI3Module(graph_definition, pulsemap, features, pulsemap_extractor, gcd_file)[source]

Bases: object

Base I3 Module for GraphNeT.

Contains methods for extracting pulsemaps, producing graphs and writing to frames.

I3Module Constructor.

Parameters:

  • graph_definition (GraphDefinition) – An instance of GraphDefinition. E.g. KNNGraph.

  • pulsemap (str) – the pulse map on which the module functions

  • features (List[str]) – the features that is used from the pulse map. E.g. [dom_x, dom_y, dom_z, charge]

  • pulsemap_extractor (Union[List[I3FeatureExtractor], I3FeatureExtractor]) – The I3FeatureExtractor used to extract the pulsemap from the I3Frames

  • gcd_file (str) – Path to the associated gcd-file.

class graphnet.deployment.i3modules.graphnet_module.I3InferenceModule(pulsemap, features, pulsemap_extractor, model_config, state_dict, model_name, gcd_file, prediction_columns)[source]

Bases: GraphNeTI3Module

General class for inference on i3 frames.

General class for inference on I3Frames (physics).

Parameters:

  • pulsemap (str) – the pulsmap that the model is expecting as input.

  • features (List[str]) – the features of the pulsemap that the model is expecting.

  • pulsemap_extractor (Union[List[I3FeatureExtractor], I3FeatureExtractor]) – The extractor used to extract the pulsemap.

  • model_config (Union[ModelConfig, str]) – The ModelConfig (or path to it) that summarizes the model used for inference.

  • state_dict (str) – Path to state_dict containing the learned weights.

  • model_name (str) – The name used for the model. Will help define the named entry in the I3Frame. E.g. “dynedge”.

  • gcd_file (str) – path to associated gcd file.

  • prediction_columns (Union[str, List[str], None], default: None) –

    column names for the predictions of the model. Will help define the named entry in the I3Frame.

    E.g. [‘energy_reco’]. Optional.

class graphnet.deployment.i3modules.graphnet_module.I3PulseCleanerModule(pulsemap, features, pulsemap_extractor, model_config, state_dict, model_name, *, gcd_file, threshold, discard_empty_events, prediction_columns)[source]

Bases: I3InferenceModule

A specialized module for pulse cleaning.

It is assumed that the model provided has been trained for this.

General class for inference on I3Frames (physics).

Parameters:

  • pulsemap (str) – the pulsmap that the model is expecting as input (the one that is being cleaned).

  • features (List[str]) – the features of the pulsemap that the model is expecting.

  • pulsemap_extractor (Union[List[I3FeatureExtractor], I3FeatureExtractor]) – The extractor used to extract the pulsemap.

  • model_config (str) – The ModelConfig (or path to it) that summarizes the model used for inference.

  • state_dict (str) – Path to state_dict containing the learned weights.

  • model_name (str) – The name used for the model. Will help define the named entry in the I3Frame. E.g. “dynedge”.

  • gcd_file (str) – path to associated gcd file.

  • threshold (float, default: 0.7) – the threshold for being considered a positive case. E.g., predictions >= threshold will be considered to be signal, all else noise.

  • discard_empty_events (bool, default: False) – When true, this flag will eliminate events whose cleaned pulse series are empty. Can be used to speed up processing especially for noise simulation, since it will not do any writing or further calculations.

  • prediction_columns (Union[str, List[str], None], default: None) – column names for the predictions of the model. Will help define the named entry in the I3Frame. E.g. [‘energy_reco’]. Optional.