Merge branch 'ybe_generator' of https://github.com/FaroutYLq/saltax i…

…nto ybe_generator

saltax/instructions/generator.py

@@ -381,43 +381,46 @@ def generator_ambe(
 
     return instr
 
+
 def generator_ybe(
-    runid, 
-    n_tot=None, 
-    rate=1e9 / SALT_TIME_INTERVAL, 
+    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. 
+    """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 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
+    :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)
+    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)
+    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]]
+        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"]