diff --git a/saltax/instructions/generator.py b/saltax/instructions/generator.py index 1be3919..b2439d5 100644 --- a/saltax/instructions/generator.py +++ b/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,