Merge branch 'main' into remove_microphysics_context

fuse/common.py

@@ -204,25 +204,6 @@ def ak_num(array, **kwargs):
     return ak.num(array, **kwargs)
 
 
-@numba.njit
-def offset_range(offsets):
-    """Computes range of constant event ids while in same offset. E.g. for an
-    array [1], [1,2,3], [5] this function yields [0, 1, 1, 1, 2].
-
-    Args:
-        offsets (ak.array): jagged array offsets.
-
-    Returns:
-        np.array: Indicies.
-    """
-    res = np.zeros(np.sum(offsets), dtype=np.int32)
-    i = 0
-    for ind, o in enumerate(offsets):
-        res[i : i + o] = ind
-        i += o
-    return res
-
-
 # Code shared between S1 and S2 photon propagation
 def init_spe_scaling_factor_distributions(spe_shapes):
     # Create a converter array from uniform random numbers to SPE gains

fuse/context.py

@@ -87,6 +87,11 @@
     fuse.truth_information.ClusterTagging,
 ]
 
+# Plugins to override the default processing plugins in straxen
+processing_plugins = [
+    fuse.processing.CorrectedAreasMC,
+]
+
 
 def full_chain_context(
     output_folder="./fuse_data",
@@ -104,6 +109,9 @@ def full_chain_context(
 ):
     """Function to create a fuse full chain simulation context."""
 
+    # Lets go for info level logging when working with fuse
+    log.setLevel("INFO")
+
     if corrections_run_id is None:
         raise ValueError("Specify a corrections_run_id to load the corrections")
     if (corrections_version is None) & (not run_without_proper_corrections):
@@ -168,6 +176,12 @@ def full_chain_context(
     for plugin in truth_information_plugins:
         st.register(plugin)
 
+    # Register processing plugins
+    log.info("Overriding processing plugins:")
+    for plugin in processing_plugins:
+        log.info(f"Registering {plugin}")
+        st.register(plugin)
+
     if corrections_version is not None:
         st.apply_xedocs_configs(version=corrections_version)
 
@@ -195,6 +209,9 @@ def full_chain_context(
     # Deregister plugins with missing dependencies
     st.deregister_plugins_with_missing_dependencies()
 
+    # Purge unused configs
+    st.purge_unused_configs()
+
     return st
 
 

fuse/plugins/__init__.py

@@ -9,3 +9,6 @@
 
 from . import truth_information
 from .truth_information import *
+
+from . import processing
+from .processing import *

fuse/plugins/micro_physics/__init__.py

@@ -16,9 +16,6 @@
 from . import yields
 from .yields import *
 
-from . import wfsim_connection
-from .wfsim_connection import *
-
 from . import detector_volumes
 from .detector_volumes import *
 

fuse/plugins/processing/__init__.py

@@ -0,0 +1,2 @@
+from . import corrected_areas
+from .corrected_areas import *

fuse/plugins/processing/corrected_areas.py

@@ -0,0 +1,32 @@
+import strax
+import straxen
+
+export, __all__ = strax.exporter()
+
+
+@export
+class CorrectedAreasMC(straxen.CorrectedAreas):
+    """Corrected areas plugin for MC data.
+
+    This plugin overwrites the (alt_)cs1 and (alt_)cs2 fields with the
+    not-time- corrected values as the effects of time dependent single-
+    electron gain, extraction efficiency and photo ionization are not
+    simulated in fuse.
+
+    Applying corrections for these effects to simulated data would lead
+    to incorrect results. This plugins should only be used for simulated
+    data!
+    """
+
+    __version__ = "0.0.1"
+    child_plugin = True
+
+    def compute(self, events):
+        result = super().compute(events)
+        result["cs1"] = result["cs1_wo_timecorr"]
+        result["cs2"] = result["cs2_wo_timecorr"]
+
+        result["alt_cs1"] = result["alt_cs1_wo_timecorr"]
+        result["alt_cs2"] = result["alt_cs2_wo_timecorr"]
+
+        return result

pyproject.toml

@@ -39,6 +39,7 @@ awkward = ">=2.5.1"
 uproot = ">=5.2.1"
 strax = ">=1.6.0"
 straxen = ">=2.2.3"
+utilix = ">=0.11.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.8", "setuptools"]

tests/test_FullChain.py

@@ -4,7 +4,7 @@
 import tempfile
 import timeout_decorator
 import fuse
-import straxen
+import utilix
 from _utils import test_root_file_name
 
 TIMEOUT = 240
@@ -36,7 +36,7 @@ def tearDownClass(cls):
         cls.temp_dir.cleanup()
 
     def setUp(self):
-        downloader = straxen.MongoDownloader(store_files_at=(self.temp_dir.name,))
+        downloader = utilix.mongo_storage.MongoDownloader(store_files_at=(self.temp_dir.name,))
         downloader.download_single(test_root_file_name, human_readable_file_name=True)
 
         assert os.path.exists(os.path.join(self.temp_dir.name, test_root_file_name))

tests/test_FullChain_w_DelayedElectrons.py

@@ -4,7 +4,7 @@
 import tempfile
 import timeout_decorator
 import fuse
-import straxen
+import utilix
 from _utils import test_root_file_name
 
 TIMEOUT = 240
@@ -39,7 +39,7 @@ def tearDownClass(cls):
         cls.temp_dir.cleanup()
 
     def setUp(self):
-        downloader = straxen.MongoDownloader(store_files_at=(self.temp_dir.name,))
+        downloader = utilix.mongo_storage.MongoDownloader(store_files_at=(self.temp_dir.name,))
         downloader.download_single(test_root_file_name, human_readable_file_name=True)
 
         assert os.path.exists(os.path.join(self.temp_dir.name, test_root_file_name))

tests/test_MicroPhysics.py

@@ -4,7 +4,7 @@
 import tempfile
 import timeout_decorator
 import fuse
-import straxen
+import utilix
 from _utils import test_root_file_name
 
 TIMEOUT = 60
@@ -35,7 +35,7 @@ def tearDownClass(cls):
         cls.temp_dir.cleanup()
 
     def setUp(self):
-        downloader = straxen.MongoDownloader(store_files_at=(self.temp_dir.name,))
+        downloader = utilix.mongo_storage.MongoDownloader(store_files_at=(self.temp_dir.name,))
         downloader.download_single(test_root_file_name, human_readable_file_name=True)
 
         assert os.path.exists(os.path.join(self.temp_dir.name, test_root_file_name))
@@ -90,11 +90,6 @@ def test_BBFYields(self):
         self.test_context.register(fuse.plugins.BBFYields)
         self.test_context.make(self.run_number, "quanta")
 
-    @timeout_decorator.timeout(TIMEOUT, exception_message="WFSim connection timed out")
-    def test_WFSimConnection(self):
-        self.test_context.register(fuse.plugins.output_plugin)
-        self.test_context.make(self.run_number, "wfsim_instructions")
-
     @timeout_decorator.timeout(TIMEOUT, exception_message="GasPhasePlugin timed out")
     def test_GasPhasePlugin(self):
         self.test_context.register(fuse.plugins.XENONnT_GasPhase)

tests/test_MicroPhysics_with_lineage_clustering.py

@@ -4,7 +4,7 @@
 import tempfile
 import timeout_decorator
 import fuse
-import straxen
+import utilix
 from _utils import test_root_file_name
 
 TIMEOUT = 240
@@ -36,7 +36,7 @@ def tearDownClass(cls):
         cls.temp_dir.cleanup()
 
     def setUp(self):
-        downloader = straxen.MongoDownloader(store_files_at=(self.temp_dir.name,))
+        downloader = utilix.mongo_storage.MongoDownloader(store_files_at=(self.temp_dir.name,))
         downloader.download_single(test_root_file_name, human_readable_file_name=True)
 
         assert os.path.exists(os.path.join(self.temp_dir.name, test_root_file_name))

tests/test_deterministic_seed.py

@@ -3,7 +3,7 @@
 import tempfile
 import timeout_decorator
 import fuse
-import straxen
+import utilix
 from numpy.testing import assert_array_equal, assert_raises
 from _utils import test_root_file_name
 
@@ -16,7 +16,7 @@ def setUp(self):
         self.temp_dir_1 = tempfile.TemporaryDirectory()
 
         for temp_dir in [self.temp_dir_0, self.temp_dir_1]:
-            downloader = straxen.MongoDownloader(store_files_at=(temp_dir.name,))
+            downloader = utilix.mongo_storage.MongoDownloader(store_files_at=(temp_dir.name,))
             downloader.download_single(test_root_file_name, human_readable_file_name=True)
             assert os.path.exists(os.path.join(temp_dir.name, test_root_file_name))
 

tests/test_input.py

@@ -4,7 +4,7 @@
 import tempfile
 import timeout_decorator
 import fuse
-import straxen
+import utilix
 import numpy as np
 from _utils import test_root_file_name
 
@@ -22,7 +22,7 @@ def tearDownClass(cls):
         cls.temp_dir.cleanup()
 
     def setUp(self):
-        downloader = straxen.MongoDownloader(store_files_at=(self.temp_dir.name,))
+        downloader = utilix.mongo_storage.MongoDownloader(store_files_at=(self.temp_dir.name,))
         downloader.download_single(test_root_file_name, human_readable_file_name=True)
 
         assert os.path.exists(os.path.join(self.temp_dir.name, test_root_file_name))

tests/test_plugin_random_seed.py

@@ -4,7 +4,7 @@
 import tempfile
 import timeout_decorator
 import fuse
-import straxen
+import utilix
 from _utils import test_root_file_name
 
 TIMEOUT = 60
@@ -39,7 +39,7 @@ def tearDownClass(cls):
         cls.temp_dir.cleanup()
 
     def setUp(self):
-        downloader = straxen.MongoDownloader(store_files_at=(self.temp_dir.name,))
+        downloader = utilix.mongo_storage.MongoDownloader(store_files_at=(self.temp_dir.name,))
         downloader.download_single(test_root_file_name, human_readable_file_name=True)
 
         assert os.path.exists(os.path.join(self.temp_dir.name, test_root_file_name))

tests/test_truth_plugins.py

@@ -4,7 +4,7 @@
 import tempfile
 import timeout_decorator
 import fuse
-import straxen
+import utilix
 from _utils import test_root_file_name
 
 TIMEOUT = 240
@@ -36,7 +36,7 @@ def tearDownClass(cls):
         cls.temp_dir.cleanup()
 
     def setUp(self):
-        downloader = straxen.MongoDownloader(store_files_at=(self.temp_dir.name,))
+        downloader = utilix.mongo_storage.MongoDownloader(store_files_at=(self.temp_dir.name,))
         downloader.download_single(test_root_file_name, human_readable_file_name=True)
 
         assert os.path.exists(os.path.join(self.temp_dir.name, test_root_file_name))