Enable AmBe generator to read from fuse type input

saltax/contexts.py

@@ -260,6 +260,7 @@ def xenonnt_salted_fuse(
         st.set_config(
             {
                 "input_file": instr_file_name,
+                "input_type": saltax.generator.AMBE_INSTRUCTIONS_TYPE if generator_name=="ambe" else "wfsim",
                 "raw_records_file_size_target": MAX_RAW_RECORDS_FILE_SIZE_MB,
             }
         )

saltax/instructions/generator.py

@@ -24,10 +24,26 @@
     method="RegularGridInterpolator",
 )
 SE_INSTRUCTIONS_DIR = "/project/lgrandi/yuanlq/salt/se_instructions/"
-AMBE_INSTRUCTIONS_FILE = "/project/lgrandi/yuanlq/salt/ambe_instructions/minghao_aptinput.csv"
+AMBE_INSTRUCTIONS_FILE = "/project2/lgrandi/jjakob/AmBeSprinkling/AmBe_fuse_input.csv"
+AMBE_INSTRUCTIONS_TYPE = "fuse"
 # BASE_DIR = "/project2/lgrandi/yuanlq/shared/saltax_instr/"
 BASE_DIR = os.path.abspath(__file__)[:-12] + "../../generated/"
 
+FUSE_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),
+        ]
+
 
 def generate_vertex(r_range=R_RANGE, z_range=Z_RANGE, size=1):
     """Generate a random vertex in the TPC volume.
@@ -220,6 +236,40 @@ def instr_file_name(
 
         return filename
 
+def fill_fuse_instruction_i(i, selected_ambe,times_offset):
+    instr_i = np.zeros(len(selected_ambe), dtype=FUSE_DTYPE)
+    instr_i["t"] = times_offset[i] + selected_ambe["t"]
+    instr_i["eventid"] = i + 1
+    instr_i["x"] = selected_ambe["x"]
+    instr_i["y"] = selected_ambe["y"]
+    instr_i["z"] = selected_ambe["z"]
+    instr_i["photons"] = selected_ambe["photons"]
+    instr_i["electrons"] = selected_ambe["electrons"]
+    instr_i["excitons"] = selected_ambe["excitons"]
+    instr_i["e_field"] = selected_ambe["e_field"]
+    instr_i["ed"] = selected_ambe["ed"]
+    instr_i["nestid"] = selected_ambe["nestid"]
+    instr_i["cluster_id"] = selected_ambe["cluster_id"]
+
+    return instr_i
+
+def fill_wfsim_instruction_i(i, selected_ambe,times_offset):
+    instr_i = np.zeros(len(selected_ambe), dtype=wfsim.instruction_dtype)
+    instr_i["time"] = times_offset[i] + selected_ambe["time"]
+    instr_i["event_number"] = i + 1
+    instr_i["type"] = selected_ambe["type"]
+    instr_i["x"] = selected_ambe["x"]
+    instr_i["y"] = selected_ambe["y"]
+    instr_i["z"] = selected_ambe["z"]
+    instr_i["recoil"] = selected_ambe["recoil"]
+    instr_i["e_dep"] = selected_ambe["e_dep"]
+    instr_i["amp"] = selected_ambe["amp"]
+    instr_i["n_excitons"] = selected_ambe["n_excitons"]
+
+    # Filter out 0 amplitudes
+    instr_i = instr_i[instr_i["amp"] > 0] # probably not a good idea to do that here if there is ever an instruction that is all zero
+
+    return instr_i
 
 def generator_se(
     runid,
@@ -324,6 +374,7 @@ def generator_ambe(
     rate=1e9 / SALT_TIME_INTERVAL,
     time_mode="uniform",
     ambe_instructions_file=AMBE_INSTRUCTIONS_FILE,
+    instructions_type=AMBE_INSTRUCTIONS_TYPE,
     **kwargs
 ):
     """Generate instructions for a run with AmBe source.
@@ -353,31 +404,28 @@ def generator_ambe(
     )
 
     # assign instructions
-    instr = np.zeros(0, dtype=wfsim.instruction_dtype)
+    if instructions_type == 'fuse':
+        instr = np.zeros(0, dtype=FUSE_DTYPE)
+    if instructions_type == 'wfsim':
+        instr = np.zeros(0, dtype=wfsim.instruction_dtype)
+
     for i in tqdm(range(n_tot)):
         # bootstrapped ambe instruction
         selected_ambe = ambe_instructions[
             ambe_instructions["event_number"] == ambe_event_numbers[i]
         ]
+
         # instruction for i-th event
-        instr_i = np.zeros(len(selected_ambe), dtype=wfsim.instruction_dtype)
-        instr_i["time"] = times_offset[i] + selected_ambe["time"]
-        instr_i["event_number"] = i + 1
-        instr_i["type"] = selected_ambe["type"]
-        instr_i["x"] = selected_ambe["x"]
-        instr_i["y"] = selected_ambe["y"]
-        instr_i["z"] = selected_ambe["z"]
-        instr_i["recoil"] = selected_ambe["recoil"]
-        instr_i["e_dep"] = selected_ambe["e_dep"]
-        instr_i["amp"] = selected_ambe["amp"]
-        instr_i["n_excitons"] = selected_ambe["n_excitons"]
+        if instructions_type == 'fuse':
+            instr_i = fill_fuse_instruction_i(i, selected_ambe, times_offset)
+        elif instructions_type == 'wfsim':
+            instr_i = fill_wfsim_instruction_i(i, selected_ambe, times_offset)
+        else:
+            raise ValueError(f'Instruction type {instructions_type} is not known.')
 
         # concatenate instr
         instr = np.concatenate((instr, instr_i))
 
-    # Filter out 0 amplitudes
-    instr = instr[instr["amp"] > 0]
-
     return instr
 
 

saltax/plugins/f_raw_records.py

@@ -60,7 +60,12 @@ class SChunkCsvInput(FuseBasePlugin):
     input_file = straxen.URLConfig(
         track=False,
         infer_type=False,
-        help="CSV file (in wfsim format) to read.",
+        help="CSV file (format as in input_type) to read.",
+    )
+    input_type = straxen.URLConfig(
+        track=False,
+        infer_type=False,
+        help="Input type of CSV file ('fuse' or 'wfsim').",
     )
     ns_no_instruction_after_chunk_start = straxen.URLConfig(
         default=5e7,
@@ -81,6 +86,7 @@ def setup(self):
         super().setup()
         self.csv_file_reader = SCsvFileLoader(
             input_file=self.input_file,
+            input_type=self.input_type,
             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,
@@ -120,18 +126,20 @@ def is_ready(self, chunk_i):
 
 
 class SCsvFileLoader:
-    """Class to load a CSV file (in wfsim format) with detector simulation
+    """Class to load a CSV file (in wfsim or fuse format) with detector simulation
     instructions."""
 
     def __init__(
         self,
         input_file,
+        input_type,
         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.input_type = input_type
         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
@@ -170,22 +178,28 @@ def __init__(
         ]
 
         # Translator to translate the wfsim instructions to the fuse format
-        self.translator = InstrTranslator(input_format="wfsim", output_format="fuse")
+        if self.input_type == "wfsim":
+            self.translator = InstrTranslator(input_format="wfsim", output_format="fuse")
+        elif self.input_type == "fuse":
+            print(f"No translator defined as input is already {self.input_type} type")
+        else:
+            raise ValueError(f"Input type {self.input_type} is not defined and should not be possible")
 
     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.
+        """Load the simulation instructions from the csv file
+        and then translate them to the fuse format if neccessary.
 
         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!")
+        # Load the csv file
+        log.debug(f"Loaded detector simulation instructions from a csv file in {self.input_type} 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!")
+        if self.input_type == "wfsim":
+            # 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!")
@@ -211,7 +225,7 @@ def output_chunk(self, chunk_start, chunk_end):
 
     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!")
+        log.debug(f"Loading detector simulation instructions from a csv file in {self.input_type} format!")
         df = pd.read_csv(self.input_file)
 
         return df