Add numba version of new merge w/ index function

eos/event_analysis.py

@@ -41,6 +41,7 @@ class MergeFrames(EventDataAction):
             d = dataset.data
             # put events into precise sub-frame
             d.events.tof, subframes = merge_frames_w_index(d.events.tof, tofCut, dataset.timing.tau, total_offset)
+            subframes = subframes.astype(int)
             # add a sub-pulse time 1-tau before and after each existing time
             utimes, uidxs = np.unique(d.events.wallTime, return_inverse=True)
             inter_times = np.empty(2*utimes.shape[0]+1, dtype=d.events.wallTime.dtype)

eos/helpers.py

@@ -6,12 +6,13 @@ import numpy as np
 from .event_data_types import EventDatasetProtocol, append_fields
 
 try:
-    from .helpers_numba import merge_frames, extract_walltime, filter_project_x, calculate_derived_properties_focussing
+    from .helpers_numba import merge_frames, extract_walltime, filter_project_x, \
+        calculate_derived_properties_focussing, merge_frames_w_index
 except ImportError:
     import logging
     logging.warning('Cannot import numba enhanced functions, is it installed?')
-    from .helpers_fallback import merge_frames, extract_walltime, filter_project_x, calculate_derived_properties_focussing
-from .helpers_fallback import merge_frames_w_index
+    from .helpers_fallback import merge_frames, extract_walltime, filter_project_x, \
+        calculate_derived_properties_focussing, merge_frames_w_index
 
 def add_log_to_pulses(key, dataset: EventDatasetProtocol):
     """

eos/helpers_fallback.py

@@ -16,7 +16,7 @@ def merge_frames_w_index(tof_e, tofCut, tau, total_offset):
     1  - belongs to the second neutron pulse of the original frame
     """
     new_tof = merge_frames(tof_e, tofCut, tau, total_offset)
-    frame_idx = np.floor_divide(tof_e-tofCut, tau).astype(int)
+    frame_idx = np.floor_divide(tof_e-tofCut, tau)
     return new_tof, frame_idx
 
 def extract_walltime(tof_e, dataPacket_p, dataPacketTime_p):

eos/helpers_numba.py

@@ -11,6 +11,22 @@ def merge_frames(tof_e, tofCut, tau, total_offset):
         tof_e_out[ti] = ((tof_e[ti]-dt)%tau)+total_offset  # tof shifted to 1 frame
     return tof_e_out
 
+@nb.jit(nb.float64[:,:](nb.float64[:], nb.float64, nb.float64, nb.float64),
+        nopython=True, parallel=True, cache=True)
+def merge_frames_w_index(tof_e, tofCut, tau, total_offset):
+    """
+    Version of merge frames that also returns a frame index for each pulse:
+    0  - belongs to the frame it was measured in
+    -1 - arrived in this frame but belongs to the previous neutron pulse
+    1  - belongs to the second neutron pulse of the original frame
+    """
+    tof_idx_out = np.empty((2, tof_e.shape[0]), dtype=np.float64)
+    dt = (tofCut-tau)
+    for ti in nb.prange(tof_e.shape[0]):
+        tof_idx_out[0, ti] = ((tof_e[ti]-dt)%tau)+total_offset  # tof shifted to 1 frame
+        tof_idx_out[1, ti] = ((tof_e[ti]-tofCut) // tau)  # tof shifted to 1 frame
+    return tof_idx_out
+
 @nb.jit(nb.int64[:](nb.float64[:], nb.uint32[:], nb.int64[:]),
         nopython=True, parallel=True, cache=True)
 def extract_walltime(tof_e, dataPacket_p, dataPacketTime_p):