Add lineage clustering algorithm (#190)

* Change detector physics config defaults

* Change microphysics config defaults

* Change PMT and DAQ config defaults

* Add new context functions

* cleanup

* Bump version: 0.1.1 → 0.2.0

* Update PMT Gain handling

* Update documentation

* Include some code from cutax

* Calculate S1 LCE map from patternmap

* Add _mc to s2_secondary_sc_gain config arg

* Set some default values

* Update context in tests

* Update example notebooks with new context

* bugfix in ElectronDrift

* Add S2 pattern map AFT scaling

* Update Documentation

* Cleanup

* Cleanup

* Bump version: 0.2.0 → 0.3.0

* Revert "Bump version: 0.2.0 → 0.3.0"

This reverts commit ea0df7c.

* Bump version: 0.2.0 → 0.2.1

* cache config args for better performance

* Use new strax.DownChunkingPlugin

* Always save results of fast S1 & S2 plugins

* Change default logging level to INFO

* Increase propagated_s2_photons_file_size_target

* Bump version: 0.2.1 → 0.2.2

* get n_electron from binomial process

* Update documentation

* Bump version: 0.2.2 → 0.2.3

* BugFix

* return correct shaped zero array if empty input

* BugFix

* Use np.around in waveform dtype conversion

* Bump version: 0.2.3 → 0.2.4

* Remove interact. time offset + increase delay cut

* Bump version: 0.2.4 → 0.2.5

* Sort PMT AP photons by time

* Change PMTAfterPulses version number

* Remove unused WFSim functions

* Rename __uniform_pe.. to spe_scaling_fact..

* Refactor pmt_transit_time_spread

* Update plugin versions

* Bugfixing

* Keep only photons with positive photon_gain

* update TPC volume plugins -> Remove primary positions

* Remove cuts when reading the root file

* Remove primary positions from merge clusters

* Add lineage clustering code draft

* Add periodictable to dependencies

* Add option to choose the clustering method when building the context

* Bugfix

* Disable energy cut when loading the root file

* Remove clusters with 0 energy deposit when merging the lineage

* Multiple updates to the lineage reconstruction

* Use base plugins for lineage clustering

* Restructure

* Restructure lineage algorithm

* Fix function after merge commit

* Dtype bugfixing

* Remove zero energy deposits for dbscan clustering

* input

* cluster id dtype problems fixed

* small changes

* rebase input.py

* working lineage

* add beta_brem

* remove main cluster from merged

* gammas were classified as ions

* forgot print

* make beta yields plugin work

* remove file with extra space

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix microphysics only context

* cleanup lineage cluster

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* cleanup

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* cleanup

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add periodictable to requirements

* Common place to define some dtypes and bugfixing

* Save primary positions again

* Bugfixing

* Add test for lineage clustering plugins

* set pri positions in MergeLineage

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove empty TPC events

* Remove empty events

* Fix some typo

* Add docstrings

* Bump some plugin versions

* Update plugin config arguments

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* rename g1 and g2

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove pickle import

* Add test for BetaYield plugin

* Update exception_message

* Use same defaults for g1g2 as straxen EnergyEstimates

* Remove clustered_interactions_dtype

* Update dtypes

* Bugfixing

* Debug

* Bugfix

* Bugfixing

* Even more bug fixing....

* efficient mem chunking in input plugin

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix time and t bug and dtype

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* forgot import function from common..

* fix bug

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* empty tpc events handling

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* whitespace

* allow empty chunk

* one last chance

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Bugfixing

* float to int in compare boundaries

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* better sorting on float times

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix preselection cut

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Typo

* from cs1/cs2 to n_ph/n_e and make dummy function for test

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix missing import and avoid lambda func

* function needs to be picklable

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* docstring

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix bug on cut_delay

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* do cut delay directly on interactions

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: cfuselli <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Diego Ramírez García <[email protected]>

fuse/context.py

@@ -13,10 +13,19 @@
 log = logging.getLogger("fuse.context")
 
 # Plugins to simulate microphysics
-microphysics_plugins = [
+microphysics_plugins_dbscan_clustering = [
     fuse.micro_physics.ChunkInput,
     fuse.micro_physics.FindCluster,
     fuse.micro_physics.MergeCluster,
+]
+
+microphysics_plugins_lineage_clustering = [
+    fuse.micro_physics.ChunkInput,
+    fuse.micro_physics.LineageClustering,
+    fuse.micro_physics.MergeLineage,
+]
+
+remaining_microphysics_plugins = [
     fuse.micro_physics.XENONnT_TPC,
     fuse.micro_physics.XENONnT_BelowCathode,
     fuse.micro_physics.VolumesMerger,
@@ -95,7 +104,9 @@ def microphysics_context(
     )
 
     # Register microphysics plugins
-    for plugin in microphysics_plugins:
+    for plugin in microphysics_plugins_dbscan_clustering:
+        st.register(plugin)
+    for plugin in remaining_microphysics_plugins:
         st.register(plugin)
 
     set_simulation_config_file(st, simulation_config_file)
@@ -105,6 +116,7 @@ def microphysics_context(
 
 def full_chain_context(
     output_folder="./fuse_data",
+    clustering_method="dbscan",
     corrections_version=None,
     simulation_config_file="fuse_config_nt_sr1_dev.json",
     corrections_run_id="046477",
@@ -140,15 +152,22 @@ def full_chain_context(
     )
 
     st.config.update(
-        dict(
-            # detector='XENONnT',
-            check_raw_record_overlaps=True,
-            **straxen.contexts.xnt_common_config,
+        dict(  # detector='XENONnT',
+            check_raw_record_overlaps=True, **straxen.contexts.xnt_common_config
         )
     )
 
     # Register microphysics plugins
-    for plugin in microphysics_plugins:
+    if clustering_method == "dbscan":
+        for plugin in microphysics_plugins_dbscan_clustering:
+            st.register(plugin)
+    elif clustering_method == "lineage":
+        for plugin in microphysics_plugins_lineage_clustering:
+            st.register(plugin)
+    else:
+        raise ValueError(f"Clustering method {clustering_method} not implemented!")
+
+    for plugin in remaining_microphysics_plugins:
         st.register(plugin)
 
     # Register S1 plugins

fuse/plugins/micro_physics/__init__.py

@@ -21,3 +21,9 @@
 
 from . import detector_volumes
 from .detector_volumes import *
+
+from . import lineage_cluster
+from .lineage_cluster import *
+
+from . import merge_lineage
+from .merge_lineage import *

fuse/plugins/micro_physics/detector_volumes.py

@@ -23,7 +23,7 @@ class VolumesMerger(VerticalMergerPlugin):
 
     provides = "interactions_in_roi"
     data_kind = "interactions_in_roi"
-    __version__ = "0.1.0"
+    __version__ = "0.2.0"
 
     def compute(self, **kwargs):
         return super().compute(**kwargs)

fuse/plugins/micro_physics/input.py

@@ -218,8 +218,7 @@ def output_chunk(self):
             )
 
         # Removing all events with zero energy deposit
-        m = interactions["ed"] > 0
-        interactions = interactions[m]
+        # m = interactions["ed"] > 0
 
         if self.cut_nr_only:
             log.info("'nr_only' set to True, keeping only the NR events")
@@ -232,6 +231,9 @@ def output_chunk(self):
 
         # Removing all events with no interactions:
         m = ak.num(interactions["ed"]) > 0
+        # and all events with no deposited energy
+        m = m & (ak.sum(interactions["ed"], axis=1) > 0)
+
         interactions = interactions[m]
 
         # Sort interactions in events by time and subtract time of the first interaction
@@ -240,17 +242,6 @@ def output_chunk(self):
         if self.event_rate > 0:
             interactions["t"] = interactions["t"] - interactions["t"][:, 0]
 
-        # Get the interaction times into flat numpy array
-        interaction_time = awkward_to_flat_numpy(interactions["t"])
-
-        # Remove interactions that happen way after the run ended
-        # we will apply the cut later on the times instead of t
-        delay_cut = interaction_time <= self.cut_delayed
-        log.info(
-            f"Removing {np.sum(~delay_cut)} ({np.sum(~delay_cut) / len(delay_cut):.4%}) "
-            f"interactions later than {self.cut_delayed:.2e} ns."
-        )
-
         # Adjust event times if necessary
         if self.event_rate > 0:
 
@@ -278,16 +269,17 @@ def output_chunk(self):
         else:
             raise ValueError("Source rate cannot be negative!")
 
-        # Overwrite interaction_time (based on "t") with the new event times
+        interactions = interactions[interactions["t"] < self.cut_delayed]
+
+        # Make into a flat numpy array
         interaction_time = awkward_to_flat_numpy(interactions["time"])
+
         # First caclulate sort index for the interaction times
         sort_idx = np.argsort(interaction_time)
         # and now make it an integer for strax time field
         interaction_time = interaction_time.astype(np.int64)
         # Sort the interaction times
         interaction_time = interaction_time[sort_idx]
-        # Apply the delay cut
-        interaction_time = interaction_time[delay_cut]
 
         chunk_idx = dynamic_chunking(
             interaction_time, scale=self.separation_scale, n_min=self.n_interactions_per_chunk