graphnet-team / graphnet

A Deep learning library for neutrino telescopes
https://graphnet-team.github.io/graphnet/
Apache License 2.0
94 stars 94 forks source link

Add pulse merging functionality to `GraphDefinition` #748

Open RasmusOrsoe opened 2 months ago

RasmusOrsoe commented 2 months ago

This PR adds new functionality to GraphDefinition that allows it to merge coincident pulses/photons on the same PMT within some merge_window if the merge_coincident argument is set to True (defaults to False). The merging is charge-weighted if charge is available in the input data.

To avoid having to pass arguments to the GraphDefinition, the names of the time and charge columns are created as Detector properties, and this PR adds the property for all existing Detectors.

Here's an example:

import numpy as np

from graphnet.models.graphs import KNNGraph
from graphnet.models.detector.prometheus import ORCA150

input_feature_names = ['sensor_pos_x',
                        'sensor_pos_y',
                        'sensor_pos_z',
                        't']

input_features = np.array([[1,2,3,0.1],
                            [1,2,3,2],
                            [2,2,1,0.1],
                            [1,2,1,0.1],
                            [1,2,1,7]])
print(input_features)
graph_definition = KNNGraph(detector = ORCA150(),
                merge_coincident = True,
                merge_window = 4.5, # nanoseconds
                input_feature_names=input_feature_names)
y = graph_definition(input_features = input_features,
                input_feature_names = input_feature_names)

print(y.x.numpy())

The corresponding output is

# Original data
[[1.  2.  3.  0.1] #same-pmt
 [1.  2.  3.  2. ]  #same-pmt
 [2.  2.  1.  0.1]
 [1.  2.  1.  0.1] #same-pmt
 [1.  2.  1.  7. ]] #same-pmt 

# merged (standardization removed here to make it easier to see)
[[2.  2.  1.  0.1]
 [1.  2.  1.  0.1]
 [1.  2.  1.  7. ]
 [2.  4.  6.  2.1]] 

As you can see, the two first rows and the last two rows in the original data is same-pmt pulses (first is less than 2 ns apart and second is more than 6 ns apart). Because we chose a merging window of 4.5ns, only the first same-pmt pulses are merged.

pweigel commented 2 months ago

Just a minor comment/warning in _find_photons_for_merging: the time windows are assumed to begin at the start of the first pulse/photon, which can lead to some bias. In data, the start time can be between 0 and merge_window ns before the first pulse because the DAQ is always sampling (well, for our purposes). A second pulse that comes soon after may not be merged in data, but if the window start time is set to the time of the first pulse it may always be merged.