diff --git a/fuse/common.py b/fuse/common.py index 5f39e8c5..fb54da1f 100644 --- a/fuse/common.py +++ b/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 diff --git a/fuse/context.py b/fuse/context.py index dcddc614..5c43f7f0 100644 --- a/fuse/context.py +++ b/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 diff --git a/fuse/plugins/__init__.py b/fuse/plugins/__init__.py index cc298898..bebf6792 100644 --- a/fuse/plugins/__init__.py +++ b/fuse/plugins/__init__.py @@ -9,3 +9,6 @@ from . import truth_information from .truth_information import * + +from . import processing +from .processing import * diff --git a/fuse/plugins/micro_physics/__init__.py b/fuse/plugins/micro_physics/__init__.py index ada2ed9c..a489ae8a 100644 --- a/fuse/plugins/micro_physics/__init__.py +++ b/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 * diff --git a/fuse/plugins/micro_physics/wfsim_connection.py b/fuse/plugins/micro_physics/wfsim_connection.py deleted file mode 100644 index 56cbec6b..00000000 --- a/fuse/plugins/micro_physics/wfsim_connection.py +++ /dev/null @@ -1,101 +0,0 @@ -# This plugin can be used to generate output in a shape that can be read by WFSim -# We can keep this for validation of fuse but it can be removed later on - -import awkward as ak -import numpy as np -import strax - -from ...common import offset_range, reshape_awkward -from ...plugin import FuseBasePlugin - -export, __all__ = strax.exporter() - - -@export -class output_plugin(FuseBasePlugin): - __version__ = "0.2.1" - - depends_on = ("interactions_in_roi", "quanta", "electric_field_values") # Add times later - - provides = "wfsim_instructions" - data_kind = "wfsim_instructions" - - save_when = strax.SaveWhen.TARGET - - dtype = [ - (("Waveform simulator event number", "event_number"), np.int32), - (("Quanta type (S1 photons or S2 electrons)", "type"), np.int8), - (("Time of the interaction [ns]", "time"), np.int64), - (("End Time of the interaction [ns]", "endtime"), np.int64), - (("X position of the cluster [cm]", "x"), np.float32), - (("Y position of the cluster [cm]", "y"), np.float32), - (("Z position of the cluster [cm]", "z"), np.float32), - (("Number of quanta", "amp"), np.int32), - (("Recoil type of interaction", "recoil"), np.int8), - (("Energy deposit of interaction", "e_dep"), np.float32), - (("Eventid like in geant4 output rootfile", "g4id"), np.int32), - (("Volume id giving the detector subvolume", "vol_id"), np.int32), - (("Local field [ V / cm ]", "local_field"), np.float64), - (("Number of excitons", "n_excitons"), np.int32), - (("X position of the primary particle [cm]", "x_pri"), np.float32), - (("Y position of the primary particle [cm]", "y_pri"), np.float32), - (("Z position of the primary particle [cm]", "z_pri"), np.float32), - ] - - def compute(self, interactions_in_roi): - if len(interactions_in_roi) == 0: - return np.zeros(0, dtype=self.dtype) - - instructions = self.awkward_to_wfsim_row_style(interactions_in_roi) - - return instructions - - def awkward_to_wfsim_row_style(self, interactions): - """Converts awkward array instructions into instructions required by WFSim - Args: - interactions: awkward.Array containing GEANT4 simulation information - - Returns: - Structured numpy.array. Each row represents either a S1 or S2 - """ - if len(interactions) == 0: - return np.empty(0, dtype=self.dtype) - - ninteractions = len(interactions["ed"]) - res = np.zeros(2 * ninteractions, dtype=self.dtype) - - # TODO: Currently not supported rows with only electrons or photons due to - # this super odd shape - for i in range(2): - structure = np.unique(interactions["eventid"], return_counts=True)[1] - eventid = reshape_awkward(interactions["eventid"], structure) - - res["event_number"][i::2] = offset_range(ak.to_numpy(ak.num(eventid))) - res["type"][i::2] = i + 1 - res["x"][i::2] = interactions["x"] - res["y"][i::2] = interactions["y"] - res["z"][i::2] = interactions["z"] - res["x_pri"][i::2] = interactions["x_pri"] - res["y_pri"][i::2] = interactions["y_pri"] - res["z_pri"][i::2] = interactions["z_pri"] - res["g4id"][i::2] = interactions["eventid"] - res["vol_id"][i::2] = interactions["vol_id"] - res["e_dep"][i::2] = interactions["ed"] - if "local_field" in res.dtype.names: - res["local_field"][i::2] = interactions["e_field"] - - recoil = interactions["nestid"] - res["recoil"][i::2] = np.where(np.isin(recoil, [0, 6, 7, 8, 11]), recoil, 8) - - if i: - res["amp"][i::2] = interactions["electrons"] - else: - res["amp"][i::2] = interactions["photons"] - if "n_excitons" in res.dtype.names: - res["n_excitons"][i::2] = interactions["excitons"] - - res["time"][i::2] = interactions["time"] - res["endtime"][i::2] = interactions["endtime"] - # Remove entries with no quanta - res = res[res["amp"] > 0] - return res diff --git a/fuse/plugins/processing/__init__.py b/fuse/plugins/processing/__init__.py new file mode 100644 index 00000000..40e00983 --- /dev/null +++ b/fuse/plugins/processing/__init__.py @@ -0,0 +1,2 @@ +from . import corrected_areas +from .corrected_areas import * diff --git a/fuse/plugins/processing/corrected_areas.py b/fuse/plugins/processing/corrected_areas.py new file mode 100644 index 00000000..4fae1ba5 --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 2fef6653..76a84204 100644 --- a/pyproject.toml +++ b/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"] diff --git a/tests/test_FullChain.py b/tests/test_FullChain.py index 0f279f2d..f67aaf34 100644 --- a/tests/test_FullChain.py +++ b/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)) diff --git a/tests/test_FullChain_w_DelayedElectrons.py b/tests/test_FullChain_w_DelayedElectrons.py index 96df014c..09c87e5e 100644 --- a/tests/test_FullChain_w_DelayedElectrons.py +++ b/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)) diff --git a/tests/test_MicroPhysics.py b/tests/test_MicroPhysics.py index e1782c29..4c9c6042 100644 --- a/tests/test_MicroPhysics.py +++ b/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) diff --git a/tests/test_MicroPhysics_with_lineage_clustering.py b/tests/test_MicroPhysics_with_lineage_clustering.py index 3f49faae..47a2d851 100644 --- a/tests/test_MicroPhysics_with_lineage_clustering.py +++ b/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)) diff --git a/tests/test_deterministic_seed.py b/tests/test_deterministic_seed.py index 3d2e1c9a..61de7f7f 100644 --- a/tests/test_deterministic_seed.py +++ b/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)) diff --git a/tests/test_input.py b/tests/test_input.py index 23847e92..d94298ca 100644 --- a/tests/test_input.py +++ b/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)) diff --git a/tests/test_plugin_random_seed.py b/tests/test_plugin_random_seed.py index e232926e..815c2e55 100644 --- a/tests/test_plugin_random_seed.py +++ b/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)) diff --git a/tests/test_truth_plugins.py b/tests/test_truth_plugins.py index 4780adfc..d23c80bd 100644 --- a/tests/test_truth_plugins.py +++ b/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))