diff --git a/requirements.txt b/requirements.txt index f9f2bba..fdf94ff 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ strax straxen -wfsim +git+https://github.com/XENONnT/repository.git@8180bda7a18d7e6c4f4217cfb314154d016fcd10#egg=wfsim scipy numpy pandas diff --git a/saltax/README.md b/saltax/README.md index 27bb46e..b022445 100644 --- a/saltax/README.md +++ b/saltax/README.md @@ -2,7 +2,7 @@ ## Scope Core functions of `saltax`. ## Structure -- `instructions`: `wfsim` instruction related, including event generators. +- `instructions`: `wfsim` or `fuse` CSV instruction related, including event generators. - `match`: tools to match salted truth and reconstucted signals. - `plugins`: `straxen`-like plugins for salting, equivalently from `raw_records` to `peaklets`. - `contexts.py`: `strax` context for `saltax`. You will want to use `sxenonnt` as default context. diff --git a/saltax/__init__.py b/saltax/__init__.py index a7f7677..ef2f1aa 100644 --- a/saltax/__init__.py +++ b/saltax/__init__.py @@ -1,4 +1,4 @@ -__version__ = "0.0.2" +__version__ = "0.1.0" from . import instructions from .instructions import * diff --git a/saltax/contexts.py b/saltax/contexts.py index 156f016..3ad09d7 100644 --- a/saltax/contexts.py +++ b/saltax/contexts.py @@ -3,29 +3,86 @@ import cutax import strax from immutabledict import immutabledict +import fuse +import logging # import pema import pandas as pd +logging.basicConfig(handlers=[logging.StreamHandler()]) +log = logging.getLogger("fuse.context") + + +# fuse plugins +# Plugins to simulate microphysics +MICROPHYSICS_PLUGINS = fuse.context.microphysics_plugins +# Plugins to simulate S1 signals +S1_SIMULATION_PLUGINS = fuse.context.s1_simulation_plugins +# Plugins to simulate S2 signals +S2_SIMULATION_PLUGINS = fuse.context.s2_simulation_plugins +# Plugins to simulate delayed electrons +DELAYED_ELECTRON_SIMULATION_PLUGINS = fuse.context.delayed_electron_simulation_plugins +# Plugins to merge delayed and regular electrons +DELAYED_ELECTRON_MERGER_PLUGINS = fuse.context.delayed_electron_merger_plugins +# Plugins to simulate PMTs and DAQ +PMT_AND_DAQ_PLUGINS = fuse.context.pmt_and_daq_plugins +# Plugins to get truth information +TRUTH_INFORMATION_PLUGINS = fuse.context.truth_information_plugins +# All plugins with fuse +FUSED_PLUGINS = [ + MICROPHYSICS_PLUGINS, + S1_SIMULATION_PLUGINS, + S2_SIMULATION_PLUGINS, + DELAYED_ELECTRON_SIMULATION_PLUGINS, + DELAYED_ELECTRON_MERGER_PLUGINS, + PMT_AND_DAQ_PLUGINS, + TRUTH_INFORMATION_PLUGINS, +] +FUSE_DONT_REGISTER = [ + fuse.plugins.micro_physics.microphysics_summary.MicroPhysicsSummary, + fuse.plugins.pmt_and_daq.pmt_response_and_daq.PMTResponseAndDAQ, +] + +# ~Infinite raw_records file size to avoid downchunking +MAX_RAW_RECORDS_FILE_SIZE_MB = 1e9 + +# fuse placeholder parameters +CORRECTION_RUN_ID_DEFAULT = "046477" # straxen XENONnT options/configuration XNT_COMMON_OPTS = straxen.contexts.xnt_common_opts.copy() XNT_COMMON_CONFIG = straxen.contexts.xnt_common_config.copy() XNT_SIMULATION_CONFIG = straxen.contexts.xnt_simulation_config.copy() -# saltax options overrides +# wfsim based saltax options overrides SXNT_COMMON_OPTS_REGISTER = XNT_COMMON_OPTS["register"].copy() +SXNT_COMMON_OPTS_REGISTER.remove(straxen.Peaklets) SXNT_COMMON_OPTS_REGISTER.remove(straxen.PulseProcessing) SXNT_COMMON_OPTS_REGISTER = [ saltax.SPulseProcessing, saltax.SRawRecordsFromFaxNT, ] + SXNT_COMMON_OPTS_REGISTER -XNT_COMMON_OPTS_OVERRIDE = dict( +SXNT_COMMON_OPTS_OVERRIDE = dict( register=SXNT_COMMON_OPTS_REGISTER, ) SXNT_COMMON_OPTS = XNT_COMMON_OPTS.copy() -SXNT_COMMON_OPTS["register"] = XNT_COMMON_OPTS_OVERRIDE["register"] +SXNT_COMMON_OPTS["register"] = SXNT_COMMON_OPTS_OVERRIDE["register"] +# fuse based saltax options overrides +FXNT_COMMON_OPTS_REGISTER = XNT_COMMON_OPTS["register"].copy() +FXNT_COMMON_OPTS_REGISTER.remove(straxen.Peaklets) +FXNT_COMMON_OPTS_REGISTER.remove(straxen.PulseProcessing) +FXNT_COMMON_OPTS_REGISTER = [ + saltax.SChunkCsvInput, + saltax.SPeaklets, + saltax.SPulseProcessing, + saltax.SPMTResponseAndDAQ, +] + FXNT_COMMON_OPTS_REGISTER +FXNT_COMMON_OPTS_OVERRIDE = dict( + register=FXNT_COMMON_OPTS_REGISTER, +) +FXNT_COMMON_OPTS = XNT_COMMON_OPTS.copy() +FXNT_COMMON_OPTS["register"] = FXNT_COMMON_OPTS_OVERRIDE["register"] # saltax configuration overrides SCHANNEL_STARTS_AT = 3000 @@ -48,6 +105,7 @@ ) SXNT_COMMON_CONFIG = XNT_COMMON_CONFIG.copy() SXNT_COMMON_CONFIG["channel_map"] = XNT_COMMON_CONFIG_OVERRIDE["channel_map"] +FXNT_COMMON_CONFIG = SXNT_COMMON_CONFIG DEFAULT_XEDOCS_VERSION = cutax.contexts.DEFAULT_XEDOCS_VERSION # saltax modes supported @@ -64,18 +122,161 @@ def get_generator(generator_name): return generator_func -def xenonnt_salted( +def xenonnt_salted_fuse( + runid=None, + saltax_mode="salt", + output_folder="./fuse_data", + cut_list=cutax.BasicCuts, + corrections_version=DEFAULT_XEDOCS_VERSION, + simulation_config_file="fuse_config_nt_sr1_dev.json", + run_id_specific_config={ + "gain_model_mc": "gain_model", + "electron_lifetime_liquid": "elife", + "drift_velocity_liquid": "electron_drift_velocity", + "drift_time_gate": "electron_drift_time_gate", + }, + run_without_proper_corrections=False, + generator_name="flat", + recoil=8, + simu_mode="all", + **kwargs, +): + """Return a strax context for XENONnT data analysis with saltax. + + :param runid: run number in integer. Must exist in RunDB if you use + this context to compute raw_records_simu, or use None for data- + loading only. + :param saltax_mode: 'data', 'simu', or 'salt'. + :param output_folder: Directory where data will be stored, defaults + to ./strax_data + :param corrections_version: XENONnT documentation version to use, + defaults to DEFAULT_XEDOCS_VERSION + :param cut_list: Cut list to use, defaults to cutax.BasicCuts + :param simulation_config_file: File containing simulation + configuration + :param run_id_specific_config: Mapping of run_id specific config + :param run_without_proper_corrections: Whether to run without proper + corrections, defaults to False + :param generator_name: Instruction mode to use, defaults to 'flat' + :param recoil: NEST recoil type, defaults to 8 + :param simu_mode: 's1', 's2', or 'all'. Defaults to 'all' + :param kwargs: Extra options to pass to strax.Context or generator + :return: strax context + """ + # Manually assign a correction_run_id if runid is None + if runid is None: + corrections_run_id = CORRECTION_RUN_ID_DEFAULT + else: + corrections_run_id = runid + if (corrections_version is None) & (not run_without_proper_corrections): + raise ValueError( + "Specify a corrections_version. If you want to run without proper " + "corrections for testing or just trying out fuse, " + "set run_without_proper_corrections to True" + ) + if simulation_config_file is None: + raise ValueError("Specify a simulation configuration file") + + if run_without_proper_corrections: + log.warning( + "Running without proper correction version. This is not recommended for production use." + "Take the context defined in cutax if you want to run XENONnT simulations." + ) + + if kwargs is not None: + context_options = dict(**FXNT_COMMON_OPTS, **kwargs) + else: + context_options = FXNT_COMMON_OPTS.copy() + context_config = dict( + detector="XENONnT", + check_raw_record_overlaps=True, + **FXNT_COMMON_CONFIG, + ) + st = strax.Context( + storage=strax.DataDirectory(output_folder), config=context_config, **context_options + ) + + for plugin_list in FUSED_PLUGINS: + for plugin in plugin_list: + if plugin not in FUSE_DONT_REGISTER: + st.register(plugin) + + if corrections_version is not None: + st.apply_xedocs_configs(version=corrections_version) + + fuse.context.set_simulation_config_file(st, simulation_config_file) + + local_versions = st.config + for config_name, url_config in local_versions.items(): + if isinstance(url_config, str): + if "run_id" in url_config: + local_versions[config_name] = straxen.URLConfig.format_url_kwargs( + url_config, run_id=corrections_run_id + ) + st.config = local_versions + + # Update some run specific config + for mc_config, processing_config in run_id_specific_config.items(): + if processing_config in st.config: + st.config[mc_config] = st.config[processing_config] + else: + print(f"Warning! {processing_config} not in context config, skipping...") + + # No blinding in simulations + st.config["event_info_function"] = "disabled" + + # Deregister plugins with missing dependencies + st.deregister_plugins_with_missing_dependencies() + + # Add saltax mode + st.set_config(dict(saltax_mode=saltax_mode)) + + # Register cuts plugins + if cut_list is not None: + st.register_cut_list(cut_list) + + # Get salt generator + generator_func = get_generator(generator_name) + + # Specify simulation instructions + instr_file_name = saltax.instr_file_name( + runid=runid, recoil=recoil, generator_name=generator_name, mode=simu_mode, **kwargs + ) + # If runid is not None, then we need to either load instruction or generate it + if runid is not None: + # Try to load instruction from file and generate if not found + try: + instr = pd.read_csv(instr_file_name) + print("Loaded instructions from file", instr_file_name) + except: + print(f"Instruction file {instr_file_name} not found. Generating instructions...") + instr = generator_func(runid=runid, **kwargs) + pd.DataFrame(instr).to_csv(instr_file_name, index=False) + print(f"Instructions saved to {instr_file_name}") + + # Load instructions into config + st.set_config( + { + "input_file": instr_file_name, + "raw_records_file_size_target": MAX_RAW_RECORDS_FILE_SIZE_MB, + } + ) + + return st + + +def xenonnt_salted_wfsim( runid=None, saltax_mode="salt", output_folder="./strax_data", - xedocs_version=DEFAULT_XEDOCS_VERSION, + corrections_version=DEFAULT_XEDOCS_VERSION, cut_list=cutax.BasicCuts, - auto_register=True, + auto_register_cuts=True, faxconf_version="sr0_v4", cmt_version="global_v9", cmt_run_id="026000", generator_name="flat", - recoil=7, + recoil=8, simu_mode="all", **kwargs, ): @@ -87,10 +288,10 @@ def xenonnt_salted( :param saltax_mode: 'data', 'simu', or 'salt'. :param output_folder: Directory where data will be stored, defaults to ./strax_data - :param xedocs_version: XENONnT documentation version to use, + :param corrections_version: XENONnT documentation version to use, defaults to DEFAULT_XEDOCS_VERSION :param cut_list: Cut list to use, defaults to cutax.BasicCuts - :param auto_register: Whether to automatically register cuts, + :param auto_register_cuts: Whether to automatically register cuts, defaults to True :param faxconf_version: (for simulation) fax configuration version to use, defaults to "sr0_v4" @@ -147,8 +348,8 @@ def xenonnt_salted( # Based on straxen.contexts.xenonnt() # st.apply_cmt_version(cmt_version) - if xedocs_version is not None: - st.apply_xedocs_configs(version=xedocs_version, **kwargs) + if corrections_version is not None: + st.apply_xedocs_configs(version=corrections_version, **kwargs) # Based on cutax.xenonnt_offline() # extra plugins to register @@ -160,7 +361,7 @@ def xenonnt_salted( "g2": "bodega://g2?bodega_version=v5", } ) - if auto_register: + if auto_register_cuts: st.register_cuts() if cut_list is not None: st.register_cut_list(cut_list) @@ -215,13 +416,84 @@ def xenonnt_salted( return st +def fxenonnt( + runid=None, + saltax_mode="salt", + output_folder="./fuse_data", + cut_list=cutax.BasicCuts, + corrections_version=DEFAULT_XEDOCS_VERSION, + simulation_config_file="fuse_config_nt_sr1_dev.json", + run_id_specific_config={ + "gain_model_mc": "gain_model", + "electron_lifetime_liquid": "elife", + "drift_velocity_liquid": "electron_drift_velocity", + "drift_time_gate": "electron_drift_time_gate", + }, + run_without_proper_corrections=False, + generator_name="flat", + recoil=8, + simu_mode="all", + **kwargs, +): + """United strax context for XENONnT data, simulation, or salted data. Based + on fuse. + + :param runid: run number in integer. Must exist in RunDB if you use + this context to compute raw_records_simu, or use None for data- + loading only. + :param saltax_mode: 'data', 'simu', or 'salt' + :param output_folder: Output folder for strax data, default + './strax_data' + :param corrections_version: xedocs version to use, default is synced + with cutax latest + :param cut_list: List of cuts to register, default is + cutax.BasicCuts + :param auto_register_cuts: Whether to auto register cuts, default + True + :param faxconf_version: fax config version to use, default is synced + with cutax latest + :param cmt_version: cmt version to use, default is synced with cutax + latest + :param cmt_run_id: cmt run id to use, default is synced with cutax + :param generator_name: (for simulation) Instruction mode to use, + defaults to 'flat' + :param recoil: (for simulation) NEST recoil type, defaults to 7 + (beta ER) + :param simu_mode: 's1', 's2', or 'all'. Defaults to 'all' + :param kwargs: Additional kwargs to pass + :return: strax context + """ + assert saltax_mode in SALTAX_MODES, "saltax_mode must be one of %s" % (SALTAX_MODES) + if runid is None: + print( + "Since you specified runid=None, this context will not be able to compute raw_records_simu." + ) + print("Welcome to data-loading only mode!") + else: + print("Welcome to computation mode which only works for run %s!" % (runid)) + + return xenonnt_salted_fuse( + runid=runid, + output_folder=output_folder, + corrections_version=corrections_version, + cut_list=cut_list, + simulation_config_file=simulation_config_file, + run_id_specific_config=run_id_specific_config, + run_without_proper_corrections=run_without_proper_corrections, + generator_name=generator_name, + recoil=recoil, + simu_mode=simu_mode, + **kwargs, + ) + + def sxenonnt( runid=None, saltax_mode="salt", output_folder="./strax_data", - xedocs_version=DEFAULT_XEDOCS_VERSION, + corrections_version=DEFAULT_XEDOCS_VERSION, cut_list=cutax.BasicCuts, - auto_register=True, + auto_register_cuts=True, faxconf_version="sr0_v4", cmt_version="global_v9", cmt_run_id="026000", @@ -230,7 +502,8 @@ def sxenonnt( simu_mode="all", **kwargs, ): - """United strax context for XENONnT data, simulation, or salted data. + """United strax context for XENONnT data, simulation, or salted data. Based + on wfsim. :param runid: run number in integer. Must exist in RunDB if you use this context to compute raw_records_simu, or use None for data- @@ -238,11 +511,12 @@ def sxenonnt( :param saltax_mode: 'data', 'simu', or 'salt' :param output_folder: Output folder for strax data, default './strax_data' - :param xedocs_version: xedocs version to use, default is synced with - cutax latest + :param corrections_version: xedocs version to use, default is synced + with cutax latest :param cut_list: List of cuts to register, default is cutax.BasicCuts - :param auto_register: Whether to auto register cuts, default True + :param auto_register_cuts: Whether to auto register cuts, default + True :param faxconf_version: fax config version to use, default is synced with cutax latest :param cmt_version: cmt version to use, default is synced with cutax @@ -265,12 +539,12 @@ def sxenonnt( else: print("Welcome to computation mode which only works for run %s!" % (runid)) - return xenonnt_salted( + return xenonnt_salted_wfsim( runid=runid, output_folder=output_folder, - xedocs_version=xedocs_version, + corrections_version=corrections_version, cut_list=cut_list, - auto_register=auto_register, + auto_register_cuts=auto_register_cuts, faxconf_version=faxconf_version, cmt_version=cmt_version, cmt_run_id=cmt_run_id, diff --git a/saltax/instructions/__init__.py b/saltax/instructions/__init__.py index eb1b5a8..f68c943 100644 --- a/saltax/instructions/__init__.py +++ b/saltax/instructions/__init__.py @@ -1,3 +1,3 @@ -__version__ = "0.0.2" +__version__ = "0.1.0" from . import generator from .generator import * diff --git a/saltax/match/__init__.py b/saltax/match/__init__.py index b147dd4..eda09cd 100644 --- a/saltax/match/__init__.py +++ b/saltax/match/__init__.py @@ -1,4 +1,4 @@ -__version__ = "0.0.2" +__version__ = "0.1.0" from . import match from .match import * from . import visual diff --git a/saltax/plugins/README.md b/saltax/plugins/README.md index 89adcf3..49fee44 100644 --- a/saltax/plugins/README.md +++ b/saltax/plugins/README.md @@ -3,8 +3,8 @@ Modified `straxen` [data structure](https://straxen.readthedocs.io/en/latest/reference/datastructure_nT.html) up to `peaklets`. Everything above is designed to be exactly the same as `straxen`. ## Structure - `s_raw_records.py`: Modified `raw_records` specifically for simulated `raw_records`, which determine chunking time range based on the ones' from `raw_reocrds` -- `s_records.py`: Modified `records`, which combines `raw_reords` and `raw_records_simu` together. The latter one's channel starts at constant variable `SCHANNEL_STARTS_AT`. -- `s_peaklets.py`: Modified `peaklets`, which build peaks as there are `2*n_tpc_pmts` channels, but sum up per-channel information (`area_per_channel` and top-bottom specific fields) as if there are only `n_tpc_pmts`. +- `records.py`: Modified `records`, which combines `raw_reords` and `raw_records_simu` together. The latter one's channel starts at constant variable `SCHANNEL_STARTS_AT`. +- `peaklets.py`: Modified `peaklets`, which build peaks as there are `2*n_tpc_pmts` channels, but sum up per-channel information (`area_per_channel` and top-bottom specific fields) as if there are only `n_tpc_pmts`. ## Data Structure image diff --git a/saltax/plugins/__init__.py b/saltax/plugins/__init__.py index 0985e22..50bab23 100644 --- a/saltax/plugins/__init__.py +++ b/saltax/plugins/__init__.py @@ -1,8 +1,11 @@ +from . import f_raw_records +from .f_raw_records import * + from . import s_raw_records from .s_raw_records import * -from . import s_records -from .s_records import * +from . import records +from .records import * -from . import s_peaklets -from .s_peaklets import * +from . import peaklets +from .peaklets import * diff --git a/saltax/plugins/f_raw_records.py b/saltax/plugins/f_raw_records.py new file mode 100644 index 0000000..0155c87 --- /dev/null +++ b/saltax/plugins/f_raw_records.py @@ -0,0 +1,322 @@ +import logging +from typing import Tuple + +import numpy as np +import pandas as pd +import strax +import straxen + +from fuse.plugin import FuseBasePlugin +from fuse.plugins.pmt_and_daq.pmt_response_and_daq import PMTResponseAndDAQ +from .s_raw_records import NS_NO_INSTRUCTION_BEFORE_CHUNK_END, NS_NO_INSTRUCTION_AFTER_CHUNK_START + +export, __all__ = strax.exporter() + +logging.basicConfig(handlers=[logging.StreamHandler()]) +log = logging.getLogger("fuse.detector_physics.csv_input") + + +@export +class SPMTResponseAndDAQ(PMTResponseAndDAQ): + __version__ = "0.0.0" + provides = "raw_records_simu" + data_kind = "raw_records_simu" + + +@export +class SChunkCsvInput(FuseBasePlugin): + """Plugin which reads a CSV file containing instructions for the detector + physics simulation (in wfsim format) and returns the data in chunks. + + Modified from csv_file_loader: https://github.com/XENONnT/fuse/blob/e538d32a5a0735757a77b1dce31d6f95a379bf4e/fuse/plugins/detector_physics/csv_input.py#L118 + Similar to the case event_rate=0 there: we don't reassign times. + """ + + __version__ = "0.0.0" + + depends_on: Tuple = "raw_records" + provides = "microphysics_summary" + data_kind = "interactions_in_roi" + + save_when = strax.SaveWhen.TARGET + + # source_done = False + + dtype = [ + (("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 photons at interaction position", "photons"), np.int32), + (("Number of electrons at interaction position", "electrons"), np.int32), + (("Number of excitons at interaction position", "excitons"), np.int32), + (("Electric field value at the cluster position [V/cm]", "e_field"), np.float32), + (("Energy of the cluster [keV]", "ed"), np.float32), + (("NEST interaction type", "nestid"), np.int8), + (("ID of the cluster", "cluster_id"), np.int32), + ] + dtype = dtype + strax.time_fields + + # Config options + input_file = straxen.URLConfig( + track=False, + infer_type=False, + help="CSV file (in wfsim format) to read.", + ) + ns_no_instruction_after_chunk_start = straxen.URLConfig( + default=5e7, + track=False, + type=(int), + help="No instruction this amount of time after a chunk starts will be used, " + "as a safeguard for not getting raw_records_simu out of raw_records chunk time range. ", + ) + ns_no_instruction_before_chunk_end = straxen.URLConfig( + default=5e7, + track=False, + type=(int), + help="No instruction this amount of time before a chunk ends will be used, " + "as a safeguard for not getting raw_records_simu out of raw_records chunk time range. ", + ) + + def setup(self): + super().setup() + self.csv_file_reader = SCsvFileLoader( + input_file=self.input_file, + random_number_generator=self.rng, + ns_no_instruction_before_chunk_end=self.ns_no_instruction_before_chunk_end, + ns_no_instruction_after_chunk_start=self.ns_no_instruction_after_chunk_start, + debug=self.debug, + ) + + def compute(self, raw_records, start, end): + try: + chunk_data, source_done = self.csv_file_reader.output_chunk(start, end) + chunk_data["time"] = chunk_data["t"] + chunk_data["endtime"] = chunk_data["time"] + chunk_data = chunk_data.to_records(index=False) + data = np.zeros(len(chunk_data), dtype=self.dtype) + strax.copy_to_buffer(chunk_data, data, "_bring_data_into_correct_format") + + self.source_done = source_done + + # Stick rigorously with raw_records time range + return self.chunk(start=start, end=end, data=data, data_type="geant4_interactions") + + except StopIteration: + raise RuntimeError("Bug in chunk building!") + + def source_finished(self): + return self.source_done + + def is_ready(self, chunk_i): + """Overwritten to mimic online input plugin. + + Returns False to check source finished; Returns True to get next + chunk. + """ + if "ready" not in self.__dict__: + self.ready = False + self.ready ^= True # Flip + return self.ready + + +class SCsvFileLoader: + """Class to load a CSV file (in wfsim format) with detector simulation + instructions.""" + + def __init__( + self, + input_file, + random_number_generator, + ns_no_instruction_before_chunk_end=NS_NO_INSTRUCTION_BEFORE_CHUNK_END, + ns_no_instruction_after_chunk_start=NS_NO_INSTRUCTION_AFTER_CHUNK_START, + debug=False, + ): + self.input_file = input_file + self.rng = random_number_generator + self.ns_no_instruction_before_chunk_end = ns_no_instruction_before_chunk_end + self.ns_no_instruction_after_chunk_start = ns_no_instruction_after_chunk_start + self.debug = debug + + self.dtype = [ + (("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 photons at interaction position", "photons"), np.int32), + (("Number of electrons at interaction position", "electrons"), np.int32), + (("Number of excitons at interaction position", "excitons"), np.int32), + (("Electric field value at the cluster position [V/cm]", "e_field"), np.float32), + (("Energy of the cluster [keV]", "ed"), np.float32), + (("NEST interaction type", "nestid"), np.int8), + (("ID of the cluster", "cluster_id"), np.int32), + (("Time of the interaction", "t"), np.int64), + (("Geant4 event ID", "eventid"), np.int32), + ] + self.dtype = self.dtype + strax.time_fields + + # The csv file needs to have these columns: + self.columns = [ + "x", + "y", + "z", + "photons", + "electrons", + "excitons", + "e_field", + "ed", + "nestid", + "t", + "eventid", + "cluster_id", + ] + + # Translator to translate the wfsim instructions to the fuse format + self.translator = InstrTranslator(input_format="wfsim", output_format="fuse") + + def output_chunk(self, chunk_start, chunk_end): + """Load the simulation instructions from the csv file in wfsim format + and then translate them to the fuse format. + + Truncate the instructions to the chunk time range. + """ + # Load the csv file in wfsim format + log.debug("Loaded detector simulation instructions from a csv file in wfsim format!") + instructions = self._load_csv_file() + + # Translate the wfsim instructions to the fuse format + log.debug("Translating the wfsim instructions to the fuse format!") + instructions = self.translator.translate(instructions) + log.debug("Instructions translated to the fuse format!") + + # truncate instructions to the chunk time range + log.debug("Truncating instructions to the chunk time range!") + log.debug( + "We will further truncate the instructions to the range [%d, %d]", + chunk_start + self.ns_no_instruction_after_chunk_start, + chunk_end - self.ns_no_instruction_before_chunk_end, + ) + + # See if we have any instructions after the chunk end + mask_next = instructions["t"] > chunk_end + if np.any(mask_next): + source_done = False + else: + log.debug("This is the last chunk! No more instructions available!") + source_done = True + + mask = instructions["t"] >= chunk_start + self.ns_no_instruction_after_chunk_start + mask &= instructions["t"] < chunk_end - self.ns_no_instruction_before_chunk_end + instructions = instructions[mask] + + return instructions, source_done + + def _load_csv_file(self): + """Load the simulation instructions from a csv file in wfsim format.""" + log.debug("Loading detector simulation instructions from a csv file in wfsim format!") + df = pd.read_csv(self.input_file) + + return df + + +class InstrTranslator: + """Class to translate instructions between wfsim and fuse formats.""" + + def __init__(self, input_format="wfsim", output_format="fuse"): + self.input_format = input_format + self.output_format = output_format + + assert self.input_format in [ + "wfsim", + "fuse", + ], "Unknown input format! Choose 'wfsim' or 'fuse'!" + assert self.output_format in [ + "wfsim", + "fuse", + ], "Unknown output format! Choose 'wfsim' or 'fuse'!" + + log.debug("Translating instructions from %s to %s", self.input_format, self.output_format) + self.translator = self.translator() + + def translator(self): + if self.input_format == "wfsim" and self.output_format == "fuse": + return self.translate_wfsim_to_fuse + elif self.input_format == "fuse" and self.output_format == "wfsim": + return self.translate_fuse_to_wfsim + else: + raise NotImplementedError( + "Translation from {} to {} is not implemented yet!".format( + self.input_format, self.output_format + ) + ) + + def translate(self, instructions): + return self.translator(instructions) + + def translate_wfsim_to_fuse(self, instructions): + """Translate the wfsim instructions to the fuse format.""" + # Sort time for sanity + instructions = instructions.sort_values(by="time") + + # Find cluster and events row by row + previous_time = 0 + previous_event_number = -1 + cluster_id = 0 + event_id = 0 + first_row = True + for row in instructions.itertuples(): + new_cluster = False + new_event = False + + # Check if we have a new cluster or event + if row.time > previous_time: + new_cluster = True + if row.event_number != previous_event_number: + new_event = True + if new_event and (not new_cluster): + raise ValueError("New event without new cluster at time %s!?" % (row.time)) + previous_time = row.time + previous_event_number = row.event_number + + # Update the previous event number + if new_event: + event_id += 1 + + # Update the cluster as a new row + if new_cluster: + cluster_id += 1 + new_row = { + "x": np.float32(row.x), + "y": np.float32(row.y), + "z": np.float32(row.z), + "e_field": np.float32(row.local_field), + "ed": np.float32(row.e_dep), + "nestid": np.int8(row.recoil), + "t": np.int64(row.time), + "cluster_id": np.int32(cluster_id), + "eventid": np.int32(event_id), + "photons": np.int32(0), + "electrons": np.int32(0), + "excitons": np.int32(0), + } + last_row = new_row + + # Assign the number of photons, excitons and electrons to the last row + if row.type == 1: + last_row["photons"] = np.int32(row.amp) + last_row["excitons"] = np.int32(row.n_excitons) + elif row.type == 2: + last_row["electrons"] = np.int32(row.amp) + else: + raise ValueError("Unknown type %s!" % (row.type)) + + # Concatenate the last row to the new instructions + if first_row: + rows = [last_row] + first_row = False + elif new_cluster: + rows.append(last_row) + + return pd.DataFrame(rows) + + def translate_fuse_to_wfsim(self, instructions): + """Translate the fuse instructions to the wfsim format.""" + raise NotImplementedError("Not implemented yet!") diff --git a/saltax/plugins/s_peaklets.py b/saltax/plugins/peaklets.py similarity index 99% rename from saltax/plugins/s_peaklets.py rename to saltax/plugins/peaklets.py index 544d309..4f6db9a 100644 --- a/saltax/plugins/s_peaklets.py +++ b/saltax/plugins/peaklets.py @@ -6,7 +6,7 @@ from immutabledict import immutabledict from strax.processing.general import _touching_windows from strax.dtypes import DIGITAL_SUM_WAVEFORM_CHANNEL -from .s_records import SCHANNEL_STARTS_AT +from .records import SCHANNEL_STARTS_AT import straxen from straxen.plugins.peaklets.peaklets import hit_max_sample, get_tight_coin, drop_data_top_field diff --git a/saltax/plugins/s_records.py b/saltax/plugins/records.py similarity index 100% rename from saltax/plugins/s_records.py rename to saltax/plugins/records.py diff --git a/setup.py b/setup.py index 6ddcb24..f09cec4 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,7 @@ author="Lanqing Yuan", description="Salting analysis framework for XENONnT.", long_description=readme, - version="0.0.2", + version="0.1.0", install_requires=requires, setup_requires=["pytest-runner"], tests_require=requires + ["pytest", "hypothesis", "boltons"],