diff --git a/eos/__init__.py b/eos/__init__.py index 2205db2..ba38f85 100644 --- a/eos/__init__.py +++ b/eos/__init__.py @@ -2,5 +2,5 @@ Package to handle data redction at AMOR instrument to be used by __main__.py script. """ -__version__ = '3.2.5' +__version__ = '3.2.6' __date__ = '2026-03-10' diff --git a/eos/event_analysis.py b/eos/event_analysis.py index 3d96af1..d7c2d3a 100644 --- a/eos/event_analysis.py +++ b/eos/event_analysis.py @@ -9,7 +9,7 @@ from typing import Tuple from . import const from .event_data_types import EventDataAction, EventDatasetProtocol, append_fields, EVENT_BITMASKS -from .helpers import filter_project_x, merge_frames, extract_walltime, add_log_to_pulses +from .helpers import filter_project_x, merge_frames, extract_walltime, add_log_to_pulses, merge_frames_w_index from .instrument import Detector from .options import IncidentAngle from .header import Header @@ -25,8 +25,9 @@ class ExtractWalltime(EventDataAction): dataset.data.events = new_events class MergeFrames(EventDataAction): - def __init__(self, lamdaCut=None): + def __init__(self, lamdaCut=None, extractNeutronPulses=False): self.lamdaCut=lamdaCut + self.extractNeutronPulses=extractNeutronPulses def perform_action(self, dataset: EventDatasetProtocol)->None: if self.lamdaCut is None: @@ -36,7 +37,34 @@ class MergeFrames(EventDataAction): tofCut = lamdaCut*dataset.geometry.chopperDetectorDistance/const.hdm*1e-13 total_offset = (tofCut + dataset.timing.tau * (dataset.timing.ch1TriggerPhase + dataset.timing.chopperPhase/2)/180) - dataset.data.events.tof = merge_frames(dataset.data.events.tof, tofCut, dataset.timing.tau, total_offset) + if self.extractNeutronPulses and 'wallTime' in dataset.data.events.dtype.names: + 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) + tauns = dataset.timing.tau*1e9 + inter_times[0] = utimes[0]-tauns + inter_times[1::2] = utimes + inter_times[2::2] = utimes+tauns + # use subframe indices to sort existing events into new times + d.events.wallTime = inter_times[2*uidxs+subframes+1] + # expand pulses array with additional sub-frames + new_pulses = np.recarray(2*d.pulses.shape[0]+1, dtype=d.pulses.dtype) + new_pulses[0] = d.pulses[0] + new_pulses[1::2] = d.pulses + new_pulses[2::2] = d.pulses + new_pulses.time[0] = d.pulses.time[0]-tauns + new_pulses.time[2::2] = d.pulses.time+tauns + new_pulses.monitor /= 2.0 # ~preserve total monitor counts + d.pulses = new_pulses + else: + if self.extractNeutronPulses: + logging.error(" Trying to separate neutron pulses while wallTime is not extracted, yet!") + dataset.data.events.tof = merge_frames(dataset.data.events.tof, tofCut, dataset.timing.tau, total_offset) + class AnalyzePixelIDs(EventDataAction): diff --git a/eos/event_data_types.py b/eos/event_data_types.py index 6d1b9ca..2ccfbf6 100644 --- a/eos/event_data_types.py +++ b/eos/event_data_types.py @@ -46,6 +46,8 @@ EVENT_BITMASKS = { } def append_fields(input: np.recarray, new_fields: List[Tuple[str, np.dtype]]): + # TODO: This action is used often and time consuming as it runs len(flds) times over all indices. + # Could only be faster if array is allocated in the beginning with all fields, less flexible. # add one ore more fields to a recarray, numpy functions seems to fail flds = [(name, dtypei[0]) for name, dtypei in input.dtype.fields.items()] flds += new_fields diff --git a/eos/event_handling.py b/eos/event_handling.py index bb5b9b7..84c70c3 100644 --- a/eos/event_handling.py +++ b/eos/event_handling.py @@ -165,7 +165,9 @@ class ApplyMask(EventDataAction): self.bitmask_filter = bitmask_filter def perform_action(self, dataset: EventDatasetProtocol) ->None: - # TODO: why is this action time consuming? + # TODO: Most time in test examples is spend here. + # While the actions here are very simple, they act on a large array, + # so even just comparison and indexing become time consuming. d = dataset.data pre_filter = d.events.shape[0] if logging.getLogger().level <= logging.DEBUG: diff --git a/eos/helpers.py b/eos/helpers.py index a93e280..d8cad01 100644 --- a/eos/helpers.py +++ b/eos/helpers.py @@ -6,9 +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: - from .helpers_fallback import merge_frames, extract_walltime, filter_project_x, calculate_derived_properties_focussing + 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, 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 969510e..c39e884 100644 --- a/eos/helpers_fallback.py +++ b/eos/helpers_fallback.py @@ -8,6 +8,17 @@ def merge_frames(tof_e, tofCut, tau, total_offset): # tof shifted to 1 frame return np.remainder(tof_e-(tofCut-tau), tau)+total_offset +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 + """ + new_tof = merge_frames(tof_e, tofCut, tau, total_offset) + frame_idx = np.floor_divide(tof_e-tofCut, tau) + return new_tof, frame_idx + def extract_walltime(tof_e, dataPacket_p, dataPacketTime_p): output = np.empty(np.shape(tof_e)[0], dtype=np.int64) for i in range(len(dataPacket_p)-1): 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): diff --git a/eos/options.py b/eos/options.py index 7f56a55..4eabd24 100644 --- a/eos/options.py +++ b/eos/options.py @@ -480,6 +480,16 @@ class ReflectivityReductionConfig(ArgParsable): }, ) + extractNeutronPulses: bool = field( + default=False, + metadata={ + 'short': 'np', + 'group': 'data manicure', + 'help': 're-assign events to actual neutron pulses => ' + '2xbetter filter resolution but extra computation (experimental)', + }, + ) + class OutputFomatOption(StrEnum): Rqz_ort = "Rqz.ort" diff --git a/eos/reduction_reflectivity.py b/eos/reduction_reflectivity.py index 94e222d..ce8fc3b 100644 --- a/eos/reduction_reflectivity.py +++ b/eos/reduction_reflectivity.py @@ -65,7 +65,7 @@ class ReflectivityReduction: # the filtering only makes sense if using actual monitor data, not time self.dataevent_actions |= eh.FilterMonitorThreshold(self.config.experiment.lowCurrentThreshold) self.dataevent_actions |= eh.FilterStrangeTimes() - self.dataevent_actions |= ea.MergeFrames() + self.dataevent_actions |= ea.MergeFrames(extractNeutronPulses=self.config.reduction.extractNeutronPulses) self.dataevent_actions |= ea.AnalyzePixelIDs(self.config.experiment.yRange) self.dataevent_actions |= eh.TofTimeCorrection(self.config.experiment.incidentAngle==IncidentAngle.alphaF) self.dataevent_actions |= ea.CalculateWavelength(self.config.experiment.lambdaRange) diff --git a/tests/test_event_handling.py b/tests/test_event_handling.py index e7bf59b..ae80caa 100644 --- a/tests/test_event_handling.py +++ b/tests/test_event_handling.py @@ -46,7 +46,7 @@ class MockEventData: # list of data packates containing previous events packets = np.recarray((1000,), dtype=PACKET_TYPE) packets.start_index = np.linspace(0, events.shape[0]-1, packets.shape[0], dtype=np.uint32) - packets.time = np.linspace(1700000000000000000, 1700000000000000000+3_600_000_000, + packets.time = np.linspace(1700000000000000000, 1700000000000000000+3_600_000_000_000, packets.shape[0], dtype=np.int64) # chopper pulses within the measurement time @@ -58,7 +58,7 @@ class MockEventData: proton_current = np.recarray((50,), dtype=PC_TYPE) proton_current.current = 1500.0 proton_current[np.random.randint(0, proton_current.shape[0]-1, 10)] = 0. # random time with no current - proton_current.time = np.linspace(1700000000000000300, 1700000000000000000+3_600_000_000, + proton_current.time = np.linspace(1700000000000000300, 1700000000000000000+3_600_000_000_000, proton_current.shape[0], dtype=np.int64) self.data = AmorEventStream(events, packets, pulses, proton_current) @@ -420,20 +420,30 @@ class TestSimpleActions(TestCase): dtype=np.int32)) def test_merge_frames(self): - action = MergeFrames(lamdaCut=0.0) + action = MergeFrames(lamdaCut=0.0, extractNeutronPulses=False) action.perform_action(self.d) self.assertEqual(self.d.data.events.tof.shape, self.d.orig_data.events.tof.shape) np.testing.assert_array_compare(lambda x,y: x<=y, self.d.data.events.tof, self.d.orig_data.events.tof) self.assertTrue((-self.d.timing.tau<=self.d.data.events.tof).all()) np.testing.assert_array_less(self.d.data.events.tof, self.d.timing.tau) - action = MergeFrames(lamdaCut=2.0) + action = MergeFrames(lamdaCut=2.0, extractNeutronPulses=False) self.d.data.events.tof = self.d.orig_data.events.tof[:] action.perform_action(self.d) tofCut = 2.0*self.d.geometry.chopperDetectorDistance/const.hdm*1e-13 self.assertTrue((tofCut-self.d.timing.tau<=self.d.data.events.tof).all()) self.assertTrue((self.d.data.events.tof<=tofCut+self.d.timing.tau).all()) + def test_merge_frames_splitting(self): + action = MergeFrames(lamdaCut=0.0, extractNeutronPulses=True) + self._extract_walltime() + action.perform_action(self.d) + self.assertEqual(self.d.data.events.tof.shape, self.d.orig_data.events.tof.shape) + np.testing.assert_array_compare(lambda x,y: x<=y, self.d.data.events.tof, self.d.orig_data.events.tof) + self.assertEqual(self.d.data.pulses.shape[0], self.d.orig_data.pulses.shape[0]*2+1) + np.testing.assert_array_less(self.d.orig_data.pulses.time[:-1], self.d.orig_data.pulses.time[1:]) + np.testing.assert_array_less(self.d.data.pulses.time[:-1], self.d.data.pulses.time[1:]) + def test_analyze_pixel_ids(self): action = AnalyzePixelIDs((1000, 1001)) action.perform_action(self.d)