Merge branch 'main' of https://github.com/FaroutYLq/saltax into main

saltax/instructions/generator.py

@@ -440,6 +440,67 @@ def generator_ybe(
     return instr
 
 
+def generator_ybe(
+    runid,
+    n_tot=None,
+    rate=1e9 / SALT_TIME_INTERVAL,
+    time_mode="uniform",
+    ybe_instructions_file=YBE_INSTRUCTIONS_FILE,
+    **kwargs
+):
+    """Generate instructions for a run with YBe source.
+
+    YBe instruction was first generated by full-chain simulation, and
+    then passing the post-epix instruction to feed this function. Each
+    event with a certain event_id in the fed instructions will be
+    shifted in time based on the time_mode you specified.
+    :param runid: run number in integer
+    :param n_tot: total number of events to generate, default: None i.e.
+        generate events until end_time
+    :param rate: rate of events in Hz, default: 1e9/SALT_TIME_INTERVAL
+    :param time_mode: 'uniform' or 'realistic', default: 'uniform'
+    :param ybe_instructions_file: file containing ybe instructions,
+        default: YBE_INSTRUCTIONS_FILE
+    :return: instructions in numpy array
+    """
+    # determine time offsets to shift ybe instructions
+    start_time, end_time = get_run_start_end(runid)
+    times_offset = generate_times(start_time, end_time, size=n_tot, rate=rate, time_mode=time_mode)
+    n_tot = len(times_offset)
+
+    # bootstrap instructions
+    ybe_instructions = pd.read_csv(ybe_instructions_file)
+    ybe_event_numbers = np.random.choice(
+        np.unique(ybe_instructions.event_number), n_tot, replace=True
+    )
+
+    # assign instructions
+    instr = np.zeros(0, dtype=wfsim.instruction_dtype)
+    for i in tqdm(range(n_tot)):
+        # bootstrapped ybe instruction
+        selected_ybe = ybe_instructions[ybe_instructions["event_number"] == ybe_event_numbers[i]]
+        # instruction for i-th event
+        instr_i = np.zeros(len(selected_ybe), dtype=wfsim.instruction_dtype)
+        instr_i["time"] = times_offset[i] + selected_ybe["time"]
+        instr_i["event_number"] = i + 1
+        instr_i["type"] = selected_ybe["type"]
+        instr_i["x"] = selected_ybe["x"]
+        instr_i["y"] = selected_ybe["y"]
+        instr_i["z"] = selected_ybe["z"]
+        instr_i["recoil"] = selected_ybe["recoil"]
+        instr_i["e_dep"] = selected_ybe["e_dep"]
+        instr_i["amp"] = selected_ybe["amp"]
+        instr_i["n_excitons"] = selected_ybe["n_excitons"]
+
+        # concatenate instr
+        instr = np.concatenate((instr, instr_i))
+
+    # Filter out 0 amplitudes
+    instr = instr[instr["amp"] > 0]
+
+    return instr
+
+
 def generator_flat(
     runid,
     en_range=DEFAULT_EN_RANGE,