Strax interface

bin/run_epix

@@ -1,13 +1,7 @@
 #!/usr/bin/env python3
 
-import os
 import argparse
-import time
 import sys
-import awkward as ak
-import numpy as np
-import pandas as pd
-import warnings
 
 import epix
 
@@ -19,10 +13,13 @@ def pars_args():
     parser.add_argument('--Detector', dest='detector', type=str,
                         action='store', default='XENONnT',
                         help='Detector which should be used. Has to be defined in epix.detectors.')
-    parser.add_argument('--DetectorConfig', dest='detector_config', type=str,
+    parser.add_argument('--EpixConfigOverride', dest='epix_config_override', type=str,
                         action='store', default='',
                         help='Config file to overwrite default detector settings.')
-    parser.add_argument('--EntryStop', dest='entry_stop', type=int,
+    parser.add_argument('--EntryStart', dest='entry_stop', type=int,
+                        action='store',
+                        help='First event to be read. Defaulted is zero.'),
+    parser.add_argument('--EntryStop', dest='entry_start', type=int,
                         action='store',
                         help='Number of entries to read from first. Defaulted to all')
     parser.add_argument('--MicroSeparation', dest='micro_separation', type=float,
@@ -40,197 +37,28 @@ def pars_args():
     parser.add_argument('--MaxDelay', dest='max_delay', type=float,
                         action='store', default=1e7, #ns
                         help='Maximal time delay to first interaction which will be stored [ns]')
-    parser.add_argument('--EventRate', dest='event_rate', type=float,
-                    action='store', default=-1,
-                    help='Event rate for event separation. Use -1 for clean simulations'
-                         'or give a rate >0 to space events randomly.')
+    parser.add_argument('--SourceRate', dest='source_rate', type=float,
+                    action='store', default=0,
+                    help='Event rate for event separation. Default 0 means no time shift is applied for the different '
+                         'events. Use -1 for clean simulations or give a rate >0 to space events randomly.')
+    parser.add_argument('--JobId', dest='job_id', type=int,
+                    action='store', default=0,
+                    help='Job id in full chain simulation. Offset is computed as '
+                         '"JobId * n_simulated_events/SourceRate", n_simulated_events '
+                         'is read from file.')
     parser.add_argument('--OutputPath', dest='output_path',
                        action='store', default="",
                        help=('Optional output path. If not specified the result will be saved'
                              'in the same dir as the input file.'))
     parser.add_argument('--Debug', dest='debug',
-                       action='store_true', default="",
+                       action='store_true', default=False,
                        help=('If specifed additional information is printed to the consol.')
                         )
 
     args = parser.parse_args(sys.argv[1:])
     return args
 
-
-def main(args, return_df=False):
-    # TODO: remove setup from main for strax
-    path, file_name, detector, outer_cylinder = setup(args)
-
-    if args.debug:
-        print("epix configuration: ", args)
-        # TODO: also add memory information see starxer and change this to debug
-        # Getting time information:
-        starttime = time.time()
-        tnow = starttime
-
-    # Loading data:
-    inter = epix.loader(path, file_name, args.debug,
-                        outer_cylinder=outer_cylinder,
-                        kwargs_uproot_arrays={'entry_stop': args.entry_stop}
-                        )
-
-    if args.debug:
-        tnow = monitor_time(tnow, 'load data.')
-        print((f'Finding clusters of interactions with a dr = {args.micro_separation} mm'
-               f' and dt = {args.micro_separation_time} ns'))
-
-    # Cluster finding and clustering:
-    inter = epix.find_cluster(inter, args.micro_separation/10, args.micro_separation_time)
-
-    if args.debug:
-        tnow = monitor_time(tnow, 'cluster finding.')
-
-    result = epix.cluster(inter, args.tag_cluster_by == 'energy')
-
-    if args.debug:
-        tnow = monitor_time(tnow, 'cluster merging.')
-
-    # Add eventid again:
-    result['evtid'] = ak.broadcast_arrays(inter['evtid'][:, 0], result['ed'])[0]
-
-    # Sort detector volumes and keep interactions in selected ones:
-    if args.debug:
-        print('Removing clusters not in volumes:', *[v.name for v in detector])
-        print(f'Number of clusters before: {np.sum(ak.num(result["ed"]))}')
-
-    # Returns all interactions which are inside in one of the volumes,
-    # Checks for volume overlap, assigns Xe density and create S2 to
-    # interactions. EField comes later since interpolated maps cannot be
-    # called inside numba functions.
-    res_det = epix.in_sensitive_volume(result, detector)
-
-    # Adding new fields to result:
-    for field in res_det.fields:
-        result[field] = res_det[field]
-    m = result['vol_id'] > 0  # All volumes have an id larger zero
-    result = result[m]
-
-    # Removing now empty events as a result of the selection above:
-    m = ak.num(result['ed']) > 0
-    result = result[m]
-
-    if args.debug:
-        print(f'Number of clusters after: {np.sum(ak.num(result["ed"]))}')
-        print('Assigning electric field to clusters')
-
-    if not ak.any(m):
-        # There are not any events left so return empty array:
-        warnings.warn('No interactions left, return empty DataFrame.')
-        if return_df:
-            if args.output_path and not os.path.isdir(args.output_path):
-                os.makedirs(args.output_path)
-
-            output_path_and_name = os.path.join(args.output_path, file_name[:-5] + "_wfsim_instructions.csv")
-            df = pd.DataFrame()
-            df.to_csv(output_path_and_name, index=False)
-            return
-
-    # Add electric field to array:
-    efields = np.zeros(np.sum(ak.num(result)), np.float32)
-    # Loop over volume and assign values:
-    for volume in detector:
-        if isinstance(volume.electric_field, (float, int)):
-            ids = epix.awkward_to_flat_numpy(result['vol_id'])
-            m = ids == volume.volume_id
-            efields[m] = volume.electric_field
-        else:
-            efields = volume.electric_field(epix.awkward_to_flat_numpy(result.x),
-                                            epix.awkward_to_flat_numpy(result.y),
-                                            epix.awkward_to_flat_numpy(result.z)
-                                            )
-
-    result['e_field'] = epix.reshape_awkward(efields, ak.num(result))
-
-    # Sort in time and set first cluster to t=0, then chop all delayed
-    # events which are too far away from the rest.
-    # (This is a requirement of WFSim)
-    result = result[ak.argsort(result['t'])]
-    result['t'] = result['t'] - result['t'][:, 0]
-    result = result[result['t'] <= args.max_delay]
-
-    #Separate event in time 
-    number_of_events = len(result["t"])
-    if args.event_rate == -1:
-        dt = epix.times_for_clean_separation(number_of_events, args.max_delay)
-        if args.debug:
-            print('Clean event separation')
-    else:
-        dt = epix.times_from_fixed_rate(args.event_rate, number_of_events, args.max_delay)
-        if args.debug:
-            print(f'Fixed event rate of {args.event_rate} Hz')
-    result['t'] = result['t'][:, :] + dt
-
-    if args.debug:
-        print('Generating photons and electrons for events')
-    # Generate quanta:
-    photons, electrons = epix.quanta_from_NEST(epix.awkward_to_flat_numpy(result['ed']),
-                                               epix.awkward_to_flat_numpy(result['nestid']),
-                                               epix.awkward_to_flat_numpy(result['e_field']),
-                                               epix.awkward_to_flat_numpy(result['A']),
-                                               epix.awkward_to_flat_numpy(result['Z']),
-                                               epix.awkward_to_flat_numpy(result['create_S2']),
-                                               density=epix.awkward_to_flat_numpy(result['xe_density']))
-    result['photons'] = epix.reshape_awkward(photons, ak.num(result['ed']))
-    result['electrons'] = epix.reshape_awkward(electrons, ak.num(result['ed']))
-    if args.debug:
-        _ = monitor_time(tnow, 'get quanta.')
-
-    # Reshape instructions:
-    instructions = epix.awkward_to_wfsim_row_style(result)
-
-    # Remove entries with no quanta
-    instructions = instructions[instructions['amp'] > 0]
-    ins_df = pd.DataFrame(instructions)
-
-    if return_df:
-        if args.output_path and not os.path.isdir(args.output_path):
-            os.makedirs(args.output_path)
-
-        output_path_and_name = os.path.join(args.output_path, file_name[:-5] + "_wfsim_instructions.csv")
-        if os.path.isfile(output_path_and_name):
-            warnings.warn("Output file already exists - Overwriting")
-        ins_df.to_csv(output_path_and_name, index=False)
-
-        print('Done')
-        print('Instructions saved to ', output_path_and_name)
-        if args.debug:
-            _ = monitor_time(starttime, 'run epix.')
-
-
-def monitor_time(prev_time, task):
-    t = time.time()
-    print(f'It took {(t - prev_time):.4f} sec to {task}')
-    return t
-
-
-def setup(args):
-    """
-    Function which sets-up configurations. (for strax conversion)
-
-    :return:
-    """
-    # Getting file path and split it into directory and file name:
-    p = args.input_file
-    p = p.split('/')
-    if p[0] == "":
-        p[0] = "/"
-    path = os.path.join(*p[:-1])
-    file_name = p[-1]
-
-    # Init detector volume according to settings and get outer_cylinder
-    # for data reduction of non-relevant interactions.
-    detector = epix.init_detector(args.detector.lower(), args.detector_config)
-    outer_cylinder = getattr(epix.detectors, args.detector.lower())
-    _, outer_cylinder = outer_cylinder()
-    return path, file_name, detector, outer_cylinder
-
-
 if __name__ == "__main__":
-    args = pars_args()
-    main(args, return_df=True)
-
+    args = vars(pars_args())
+    args = epix.run_epix.setup(args)
+    epix.run_epix.main(args, return_df=True)

epix/__init__.py

@@ -8,4 +8,5 @@
 from .clustering import *
 from .ElectricFieldHandler import *
 from .event_separation import *
-from .detectors import *
+from .detectors import *
+from .run_epix import *

epix/event_separation.py

@@ -18,7 +18,7 @@ def times_for_clean_separation(n_events, MaxDelay):
 
     return dt
 
-def times_from_fixed_rate(rate, n_events, offset):
+def times_from_fixed_rate(rate, n_events, n_simulated_events, offset=0):
     """
     Function to generate event times with a fixed rate.
 
@@ -30,14 +30,16 @@ def times_from_fixed_rate(rate, n_events, offset):
 
     Args:
         rate (int or float): Mean event rate in Hz. 
-        n_events (int): Number of events
+        n_events (int): Number of events in file
+        n_events_simulated (int): True number of events which were
+            simulated.
         offset (int or float): Time offset to shift all event times to larger values.
 
     Returns:
         event_times (numpy.array): Array containing the start times of the events.
     """
 
-    simulation_time = n_events/rate
+    simulation_time = n_simulated_events/rate
     simulation_time *= 1e9
 
     event_times = np.sort(np.random.uniform(low=offset, high=simulation_time+offset, size=n_events))