Add fuse base plugins to reduce duplicate code (#112)

* Add fuseBasePlugin and fuseBaseDownChunkingPlugin * update volume and merger plugins * Update microphysics plugins * update detectorphysics plugins * update waveform building plugins * loosen save_when for Volume and Merger Plugins * Try to set the nestpy random seed in compute function * Set a timeout for the tests * Increase timeout for deterministic seed tests * Longer timeout * Use CamelCase for base plugins * Bugfix * Disable st.make in tests for secondary plugin output * Increase timeout for fullchain tests * Lock and Unlock the nestpy random seed * remove unnecessary setup functions * multi class inheritance for FuseBaseDownChunkingPlugin * remove duplicate code * Bump version: 0.2.5 → 0.3.0 --------- Co-authored-by: dachengx <[email protected]>
XENONnT · Feb 1, 2024 · 0861489 · 0861489
1 parent 5cc78d3
commit 0861489
Show file tree

Hide file tree

Showing 27 changed files with 158 additions and 592 deletions.
diff --git a/.bumpversion.cfg b/.bumpversion.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.2.5
+current_version = 0.3.0
 commit = True
 tag = True
 

diff --git a/fuse/__init__.py b/fuse/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.2.5"
+__version__ = "0.3.0"
 
 from . import plugins
 from .plugins import *

diff --git a/fuse/plugin.py b/fuse/plugin.py
@@ -0,0 +1,51 @@
+import strax
+import straxen
+import numpy as np
+import logging
+
+logging.basicConfig(handlers=[logging.StreamHandler()])
+
+class FuseBasePlugin(strax.Plugin):
+
+    """Base plugin for fuse plugins"""
+
+    #Forbid rechunking
+    rechunk_on_save = False
+
+    #Lets wait 10 minutes for the plugin to finish
+    input_timeout = 600
+
+    #Config options
+    debug = straxen.URLConfig(
+        default=False, type=bool, track=False,
+        help='Show debug informations',
+    )
+
+    deterministic_seed = straxen.URLConfig(
+        default=True, type=bool,
+        help='Set the random seed from lineage and run_id, or pull the seed from the OS.',
+    )
+
+    def setup(self):
+        super().setup()
+
+        log = logging.getLogger(f"{self.__class__.__name__}")
+
+        if self.debug:
+            log.setLevel('DEBUG')
+            log.debug(f"Running {self.__class__.__name__} version {self.__version__} in debug mode")
+        else: 
+            log.setLevel('INFO')
+
+        if self.deterministic_seed:
+            hash_string = strax.deterministic_hash((self.run_id, self.lineage))
+            self.seed = int(hash_string.encode().hex(), 16)
+            self.rng = np.random.default_rng(seed = self.seed)
+            log.debug(f"Generating random numbers from seed {self.seed}")
+        else: 
+            self.rng = np.random.default_rng()
+            log.debug(f"Generating random numbers with seed pulled from OS")
+
+class FuseBaseDownChunkingPlugin(strax.DownChunkingPlugin, FuseBasePlugin):
+
+    """Base plugin for fuse DownChunkingPlugins"""
diff --git a/fuse/plugins/detector_physics/csv_input.py b/fuse/plugins/detector_physics/csv_input.py
@@ -5,15 +5,16 @@
 import pandas as pd
 import numpy as np
 
-from ...common import dynamic_chunking, FUSE_PLUGIN_TIMEOUT
+from ...common import dynamic_chunking
+from ...plugin import FuseBasePlugin
 
 export, __all__ = strax.exporter()
 
 logging.basicConfig(handlers=[logging.StreamHandler()])
 log = logging.getLogger('fuse.detector_physics.csv_input')
 
 @export
-class ChunkCsvInput(strax.Plugin):
+class ChunkCsvInput(FuseBasePlugin):
     """
     Plugin which reads a CSV file containing instructions for the detector physics simulation
     and returns the data in chunks
@@ -24,13 +25,8 @@ class ChunkCsvInput(strax.Plugin):
     provides = "microphysics_summary"
     data_kind = "interactions_in_roi"
 
-    #Forbid rechunking
-    rechunk_on_save = False
-
     save_when = strax.SaveWhen.TARGET
 
-    input_timeout = FUSE_PLUGIN_TIMEOUT
-
     source_done = False
 
     dtype = [('x', np.float32),
@@ -48,11 +44,6 @@ class ChunkCsvInput(strax.Plugin):
     dtype = dtype + strax.time_fields
 
     #Config options
-    debug = straxen.URLConfig(
-        default=False, type=bool, track=False,
-        help='Show debug informations',
-    )
-
     input_file = straxen.URLConfig(
         track=False,
         infer_type=False,
@@ -76,27 +67,8 @@ class ChunkCsvInput(strax.Plugin):
         help='n_interactions_per_chunk',
     )
 
-    deterministic_seed = straxen.URLConfig(
-        default=True, type=bool,
-        help='Set the random seed from lineage and run_id, or pull the seed from the OS.',
-    )
-
     def setup(self):
-
-        if self.debug:
-            log.setLevel('DEBUG')
-            log.debug(f"Running ChunkCsvInput version {self.__version__} in debug mode")
-        else: 
-            log.setLevel('INFO')
-
-        if self.deterministic_seed:
-            hash_string = strax.deterministic_hash((self.run_id, self.lineage))
-            seed = int(hash_string.encode().hex(), 16)
-            self.rng = np.random.default_rng(seed = seed)
-            log.debug(f"Generating random numbers from seed {seed}")
-        else: 
-            self.rng = np.random.default_rng()
-            log.debug(f"Generating random numbers with seed pulled from OS")
+        super().setup()
 
         self.file_reader = csv_file_loader(
             input_file = self.input_file,

diff --git a/fuse/plugins/detector_physics/electron_drift.py b/fuse/plugins/detector_physics/electron_drift.py
@@ -3,15 +3,15 @@
 import straxen
 import logging
 
-from ...common import FUSE_PLUGIN_TIMEOUT
+from ...plugin import FuseBasePlugin
 
 export, __all__ = strax.exporter()
 
 logging.basicConfig(handlers=[logging.StreamHandler()])
 log = logging.getLogger('fuse.detector_physics.electron_drift')
 
 @export
-class ElectronDrift(strax.Plugin):
+class ElectronDrift(FuseBasePlugin):
 
     __version__ = "0.1.4"
 
@@ -27,20 +27,11 @@ class ElectronDrift(strax.Plugin):
              ('z_obs', np.float32),
             ]
     dtype = dtype + strax.time_fields
-
-    #Forbid rechunking
-    rechunk_on_save = False
 
     save_when = strax.SaveWhen.ALWAYS
-
-    input_timeout = FUSE_PLUGIN_TIMEOUT
 
     #Config options
-    debug = straxen.URLConfig(
-        default=False, type=bool,track=False,
-        help='Show debug informations',
-    )
-
+
     drift_velocity_liquid = straxen.URLConfig(
         default = "take://resource://"
                   "SIMULATION_CONFIG_FILE.json?&fmt=json"
@@ -131,28 +122,9 @@ class ElectronDrift(strax.Plugin):
         cache=True,
         help='fdc_map',
     )
-
-    deterministic_seed = straxen.URLConfig(
-        default=True, type=bool,
-        help='Set the random seed from lineage and run_id, or pull the seed from the OS.',
-    )
 
     def setup(self):
-
-        if self.debug:
-            log.setLevel('DEBUG')
-            log.debug(f"Running ElectronDrift version {self.__version__} in debug mode")
-        else: 
-            log.setLevel('INFO')
-
-        if self.deterministic_seed:
-            hash_string = strax.deterministic_hash((self.run_id, self.lineage))
-            seed = int(hash_string.encode().hex(), 16)
-            self.rng = np.random.default_rng(seed = seed)
-            log.debug(f"Generating random numbers from seed {seed}")
-        else: 
-            self.rng = np.random.default_rng()
-            log.debug(f"Generating random numbers with seed pulled from OS")
+        super().setup()
 
         #Can i do this scaling in the url config?
         if self.field_distortion_model == "inverse_fdc":

diff --git a/fuse/plugins/detector_physics/electron_extraction.py b/fuse/plugins/detector_physics/electron_extraction.py
@@ -4,40 +4,30 @@
 import os
 import logging
 
-from ...common import FUSE_PLUGIN_TIMEOUT
+from ...plugin import FuseBasePlugin
 
 export, __all__ = strax.exporter()
 
 logging.basicConfig(handlers=[logging.StreamHandler()])
 log = logging.getLogger('fuse.detector_physics.electron_extraction')
 
 @export
-class ElectronExtraction(strax.Plugin):
+class ElectronExtraction(FuseBasePlugin):
 
     __version__ = "0.1.3"
 
     depends_on = ("microphysics_summary", "drifted_electrons")
     provides = "extracted_electrons"
     data_kind = "interactions_in_roi"
 
-    #Forbid rechunking
-    rechunk_on_save = False
-
     save_when = strax.SaveWhen.ALWAYS
 
-    input_timeout = FUSE_PLUGIN_TIMEOUT
-
     dtype = [('n_electron_extracted', np.int32),
             ]
 
     dtype = dtype + strax.time_fields
 
     #Config options
-    debug = straxen.URLConfig(
-        default=False, type=bool,track=False,
-        help='Show debug informations',
-    )
-
     s2_secondary_sc_gain_mc = straxen.URLConfig(
         default = "take://resource://"
                   "SIMULATION_CONFIG_FILE.json?&fmt=json"
@@ -101,28 +91,6 @@ class ElectronExtraction(strax.Plugin):
         help='Map of the single electron gain',
     )
 
-    deterministic_seed = straxen.URLConfig(
-        default=True, type=bool,
-        help='Set the random seed from lineage and run_id, or pull the seed from the OS.',
-    )
-
-    def setup(self):
-
-        if self.debug:
-            log.setLevel('DEBUG')
-            log.debug(f"Running ElectronExtraction version {self.__version__} in debug mode")
-        else: 
-            log.setLevel('INFO')
-
-        if self.deterministic_seed:
-            hash_string = strax.deterministic_hash((self.run_id, self.lineage))
-            seed = int(hash_string.encode().hex(), 16)
-            self.rng = np.random.default_rng(seed = seed)
-            log.debug(f"Generating random numbers from seed {seed}")
-        else: 
-            self.rng = np.random.default_rng()
-            log.debug(f"Generating random numbers with seed pulled from OS")
-
     def compute(self, interactions_in_roi):
 
         #Just apply this to clusters with photons

diff --git a/fuse/plugins/detector_physics/electron_timing.py b/fuse/plugins/detector_physics/electron_timing.py
@@ -3,28 +3,23 @@
 import straxen
 import logging
 
-from ...common import FUSE_PLUGIN_TIMEOUT
+from ...plugin import FuseBasePlugin
 
 export, __all__ = strax.exporter()
 
 logging.basicConfig(handlers=[logging.StreamHandler()])
 log = logging.getLogger('fuse.detector_physics.electron_timing')
 
 @export
-class ElectronTiming(strax.Plugin):
+class ElectronTiming(FuseBasePlugin):
 
     __version__ = "0.1.1"
 
     depends_on = ("drifted_electrons", "extracted_electrons")
     provides = "electron_time"
     data_kind = "individual_electrons"
-
-    #Forbid rechunking
-    rechunk_on_save = False
 
     save_when = strax.SaveWhen.TARGET
-
-    input_timeout = FUSE_PLUGIN_TIMEOUT
 
     data_kind = "individual_electrons"
 
@@ -35,11 +30,6 @@ class ElectronTiming(strax.Plugin):
     dtype = dtype + strax.time_fields
 
     #Config options
-    debug = straxen.URLConfig(
-        default=False, type=bool,track=False,
-        help='Show debug informations',
-    )
-
     electron_trapping_time = straxen.URLConfig(
         default = "take://resource://"
                   "SIMULATION_CONFIG_FILE.json?&fmt=json"
@@ -49,28 +39,6 @@ class ElectronTiming(strax.Plugin):
         help='Time scale electrons are trapped at the liquid gas interface',
     )
 
-    deterministic_seed = straxen.URLConfig(
-        default=True, type=bool,
-        help='Set the random seed from lineage and run_id, or pull the seed from the OS.',
-    )
-
-    def setup(self):
-
-        if self.debug:
-            log.setLevel('DEBUG')
-            log.debug(f"Running ElectronTiming version {self.__version__} in debug mode")
-        else: 
-            log.setLevel('INFO')
-
-        if self.deterministic_seed:
-            hash_string = strax.deterministic_hash((self.run_id, self.lineage))
-            seed = int(hash_string.encode().hex(), 16)
-            self.rng = np.random.default_rng(seed = seed)
-            log.debug(f"Generating random numbers from seed {seed}")
-        else: 
-            self.rng = np.random.default_rng()
-            log.debug(f"Generating random numbers with seed pulled from OS")
-
     def compute(self, interactions_in_roi):
 
         #Just apply this to clusters with photons