From f7b2e00ab24ec8d94b62cd27b124b111aa573d6c Mon Sep 17 00:00:00 2001 From: Artur Glavic Date: Tue, 3 Mar 2026 00:16:09 +0100 Subject: [PATCH] Add numba version of new merge w/ index function --- eos/event_analysis.py | 1 + eos/helpers.py | 7 ++++--- eos/helpers_fallback.py | 2 +- eos/helpers_numba.py | 16 ++++++++++++++++ 4 files changed, 22 insertions(+), 4 deletions(-) diff --git a/eos/event_analysis.py b/eos/event_analysis.py index 9e65df6..c9c80d9 100644 --- a/eos/event_analysis.py +++ b/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) diff --git a/eos/helpers.py b/eos/helpers.py index d47b30e..d8cad01 100644 --- a/eos/helpers.py +++ b/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): """ diff --git a/eos/helpers_fallback.py b/eos/helpers_fallback.py index d53f62b..c39e884 100644 --- a/eos/helpers_fallback.py +++ b/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): diff --git a/eos/helpers_numba.py b/eos/helpers_numba.py index 55e02fd..2f2135c 100644 --- a/eos/helpers_numba.py +++ b/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):