import os import numpy as np import logging from unittest import TestCase from datetime import datetime from copy import deepcopy from orsopy.fileio import Person, Experiment, Sample, InstrumentSettings, Value, ValueRange, Polarization from eos import const from eos.header import Header from eos.event_data_types import EVENT_BITMASKS, AmorGeometry, AmorTiming, AmorEventStream, \ EventDataAction, EventDatasetProtocol, PACKET_TYPE, PC_TYPE, PULSE_TYPE, EVENT_TYPE, append_fields from eos.event_handling import ApplyPhaseOffset, ApplyParameterOverwrites, CorrectChopperPhase, CorrectSeriesTime, \ AssociatePulseWithMonitor, FilterMonitorThreshold, FilterStrangeTimes, TofTimeCorrection, ApplyMask from eos.event_analysis import ExtractWalltime, MergeFrames, AnalyzePixelIDs, CalculateWavelength, CalculateQ, \ FilterQzRange, FilterByLog from eos.options import MonitorType, IncidentAngle, ExperimentConfig class MockEventData: """ Simulated dataset to be used with event handling unit tests """ geometry: AmorGeometry timing: AmorTiming data: AmorEventStream def __init__(self): self.geometry = AmorGeometry(mu=2.0, nu=1.0, kap=0.5, kad=0.0, div=1.5, chopperSeparation=1000.0, detectorDistance=4000., chopperDetectorDistance=18842.) self.timing = AmorTiming( ch1TriggerPhase=-9.1, ch2TriggerPhase=6.75, chopperPhase=0.17, chopperSpeed=500., tau=0.06 ) self.create_data() def create_data(self): # list of events, here with random time of fligh and pixel location events = np.recarray((10000, ), dtype=EVENT_TYPE) events.tof = np.random.uniform(low=0., high=0.12, size=events.shape) events.pixelID = np.random.randint(0, 28671, size=events.shape) events.mask = 0 # 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_000, packets.shape[0], dtype=np.int64) # chopper pulses within the measurement time pulses = np.recarray((packets.shape[0],), dtype=PULSE_TYPE) pulses.monitor = 1.0 pulses.time = packets.time # proton current information with independent timing 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_000, proton_current.shape[0], dtype=np.int64) self.data = AmorEventStream(events, packets, pulses, proton_current) self.orig_data = deepcopy(self.data) def append(self, other): raise NotImplementedError("Just for testing, no append") def update_header(self, header:Header): # update a header with the information read from file header.owner = Person(name="test user", affiliation='PSI') header.experiment = Experiment(title='test experiment', instrument='amor', start_date=datetime.now(), probe="neutron") header.sample = Sample(name='test sample') header.measurement_instrument_settings = InstrumentSettings(incident_angle=Value(1.5, 'deg'), wavelength = ValueRange(3.0, 12.5, 'angstrom'), polarization = Polarization.unpolarized) def update_info_from_logs(self): RELEVANT_ITEMS = ['sample_temperature', 'sample_magnetic_field', 'polarization_config_label'] for key, log in self.data.device_logs.items(): if key not in RELEVANT_ITEMS: continue if log.value.dtype in [np.int8, np.int16, np.int32, np.int64]: # for integer items (flags) report the most common one value = np.bincount(log.value).argmax() if logging.getLogger().getEffectiveLevel() <= logging.DEBUG \ and np.unique(log.value).shape[0]>1: logging.debug(f' filtered values for {key} not unique, ' f'has {np.unique(log.value).shape[0]} values') else: value = log.value.mean() if key == 'polarization_config_label': self.instrument_settings.polarization = Polarization(const.polarizationConfigs[value]) elif key == 'sample_temperature': self.sample.sample_parameters['temperature'].magnitue = value elif key == 'sample_magnetic_field': self.sample.sample_parameters['magnetic_field'].magnitue = value class TestActionClass(TestCase): @classmethod def setUpClass(cls): """ Create test classes to be used """ class T1(EventDataAction): def perform_action(self, event: EventDatasetProtocol): event.data.events.mask += 1 class T2(EventDataAction): def perform_action(self, event: EventDatasetProtocol): event.data.events.mask += 2 class T4(EventDataAction): def perform_action(self, event: EventDatasetProtocol): event.data.events.mask += 4 cls.T1=T1; cls.T2=T2; cls.T4=T4 class H1(EventDataAction): def perform_action(self, event: EventDatasetProtocol): ... def update_header(self, header:Header) ->None: header.sample.name = 'h1' class H2(EventDataAction): def perform_action(self, event: EventDatasetProtocol): ... def update_header(self, header: Header) -> None: header.sample.name = 'h2' class HN(EventDataAction): def __init__(self, n): self._n = n def perform_action(self, event: EventDatasetProtocol): ... def update_header(self, header: Header) -> None: header.sample.name = self._n cls.H1=H1; cls.H2=H2; cls.HN = HN def setUp(self): self.d = MockEventData() self.header = Header() self.d.update_header(self.header) def test_individual(self): t1 = self.T1() t2 = self.T2() t4 = self.T4() np.testing.assert_array_equal(self.d.data.events.mask, 0) t1.perform_action(self.d) np.testing.assert_array_equal(self.d.data.events.mask, 1) t2.perform_action(self.d) np.testing.assert_array_equal(self.d.data.events.mask, 3) t4.perform_action(self.d) np.testing.assert_array_equal(self.d.data.events.mask, 7) def test_header(self): h1 = self.H1() h2 = self.H2() h3 = self.HN('h3') h4 = self.HN('h4') h1.update_header(self.header) self.assertEqual(self.header.sample.name, 'h1') h2.update_header(self.header) self.assertEqual(self.header.sample.name, 'h2') h3.update_header(self.header) self.assertEqual(self.header.sample.name, 'h3') h4.update_header(self.header) self.assertEqual(self.header.sample.name, 'h4') def test_combination(self): t1 = self.T1() t2 = self.T2() t4 = self.T4() t12 = t1 | t2 t24 = t2 | t4 t1224 = t1 | t2 | t2 | t4 t1224b = t12 | t24 np.testing.assert_array_equal(self.d.data.events.mask, 0) t12.perform_action(self.d) np.testing.assert_array_equal(self.d.data.events.mask, 3) t24.perform_action(self.d) np.testing.assert_array_equal(self.d.data.events.mask, 9) t1224.perform_action(self.d) np.testing.assert_array_equal(self.d.data.events.mask, 18) t1224b.perform_action(self.d) np.testing.assert_array_equal(self.d.data.events.mask, 27) def test_combine_header(self): h1 = self.H1() h2 = self.H2() h3 = self.HN('h3') h4 = self.HN('h4') (h1|h2).update_header(self.header) self.assertEqual(self.header.sample.name, 'h2') (h2|h1).update_header(self.header) self.assertEqual(self.header.sample.name, 'h1') (h3|h4).update_header(self.header) self.assertEqual(self.header.sample.name, 'h4') (h4|h3).update_header(self.header) self.assertEqual(self.header.sample.name, 'h3') def test_abstract_misssing(self): with self.assertRaises(TypeError): class E(EventDataAction): ... _ = E() def test_hash(self): """ Check that hashes of different actions are different but instances of same action have same hash """ t1 = self.T1() t1b = self.T1() t2 = self.T2() t4 = self.T4() h3 = self.HN('h3') h3b = self.HN('h3') h4 = self.HN('h4') self.assertNotEqual(t1.action_hash(), t2.action_hash()) self.assertNotEqual(t2.action_hash(), t4.action_hash()) self.assertNotEqual(t1.action_hash(), t4.action_hash()) self.assertNotEqual(h3.action_hash(), h4.action_hash()) self.assertEqual(t1.action_hash(), t1b.action_hash()) self.assertEqual(h3.action_hash(), h3b.action_hash()) class TestSimpleActions(TestCase): def setUp(self): self.d = MockEventData() self.header = Header() self.d.update_header(self.header) def test_chopper_phase(self): cp = CorrectChopperPhase() cp.perform_action(self.d) np.testing.assert_array_equal( self.d.data.events.tof, self.d.orig_data.events.tof+ self.d.timing.tau*(self.d.timing.ch1TriggerPhase-self.d.timing.chopperPhase/2)/180 ) def _extract_walltime(self): # Extract wall time for events and orig copy wt = ExtractWalltime() d = self.d.data self.d.data = self.d.orig_data wt.perform_action(self.d) self.d.data = d wt.perform_action(self.d) def test_extract_walltime(self): self._extract_walltime() # wallTime should be always a time present in the packet times np.testing.assert_array_equal(np.isin(self.d.data.events.wallTime, self.d.data.packets.time), True) # make sure extraction works on both original and copy np.testing.assert_array_equal(self.d.data.events.wallTime, self.d.orig_data.events.wallTime) def test_series_time(self): corr = 100 ct = CorrectSeriesTime(corr) with self.assertRaises(ValueError): ct.perform_action(self.d) self._extract_walltime() ct.perform_action(self.d) np.testing.assert_array_equal( self.d.data.pulses.time, self.d.orig_data.pulses.time-corr ) np.testing.assert_array_equal( self.d.data.events.wallTime, self.d.orig_data.events.wallTime-corr ) np.testing.assert_array_equal( self.d.data.proton_current.time, self.d.orig_data.proton_current.time-corr ) def test_associate_monitor(self): amPC = AssociatePulseWithMonitor(MonitorType.proton_charge) amT = AssociatePulseWithMonitor(MonitorType.time) amN = AssociatePulseWithMonitor(MonitorType.neutron_monitor) self.d.data.pulses.monitor = 13 amN.perform_action(self.d) np.testing.assert_array_equal(self.d.data.pulses.monitor, 1) self.d.data.pulses.monitor = 13 amT.perform_action(self.d) np.testing.assert_array_equal(self.d.data.pulses.monitor, np.float32(2*self.d.timing.tau)) self.d.data.pulses.monitor = 13 amPC.perform_action(self.d) pcm = self.d.data.proton_current.current *2*self.d.timing.tau*1e-3 np.testing.assert_array_equal(np.isin(self.d.data.pulses.monitor, pcm), True) def test_filter_monitor_threashold(self): amPC = AssociatePulseWithMonitor(MonitorType.proton_charge) fmt = amPC | FilterMonitorThreshold(1000.) fma = amPC | FilterMonitorThreshold(2000.) fm0 = amPC | FilterMonitorThreshold(-1.0) with self.assertRaises(ValueError): fmt.perform_action(self.d) self._extract_walltime() fm0.perform_action(self.d) self.assertEqual(self.d.data.events.mask.sum(), 0) fmt.perform_action(self.d) # calculate, which events should have 0 monitor zero_times = self.d.data.pulses.time[self.d.data.pulses.monitor==0] zero_sum = np.isin(self.d.data.events.wallTime, zero_times).sum() self.assertEqual(self.d.data.events.mask.sum(), zero_sum*EVENT_BITMASKS['MonitorThreshold']) # filter all events self.d.data.events.mask = 0 fma.perform_action(self.d) self.assertEqual(self.d.data.events.mask.sum(), self.d.data.events.shape[0]*EVENT_BITMASKS['MonitorThreshold']) def test_filter_strage_times(self): st = FilterStrangeTimes() st.perform_action(self.d) self.assertEqual(self.d.data.events.mask.sum(), 0) # half events should be strange times (outside of ToF frame) self.d.data.events.tof += self.d.timing.tau st.perform_action(self.d) self.assertEqual(self.d.data.events.mask.sum(), (self.d.data.events.tof>2*self.d.timing.tau).sum()*EVENT_BITMASKS['StrangeTimes']) def test_apply_phase_offset(self): action = ApplyPhaseOffset(12.5) action.perform_action(self.d) self.assertEqual(self.d.timing.ch1TriggerPhase, 12.5) def test_apply_parameter_overwrites(self): action = ApplyParameterOverwrites(ExperimentConfig(muOffset=0.25, mu=3.5, nu=4.5)) action.perform_action(self.d) self.assertEqual(self.d.geometry.mu, 3.5) self.assertEqual(self.d.geometry.nu, 4.5) action = ApplyParameterOverwrites(ExperimentConfig(muOffset=0.25)) action.perform_action(self.d) self.assertEqual(self.d.geometry.mu, 3.75) action = ApplyParameterOverwrites(ExperimentConfig(sampleModel='air | Si | Fe')) action.update_header(self.header) self.assertIsNotNone(self.header.sample.model) def test_apply_sample_model_file(self): if os.path.isfile('test.yaml'): os.remove('test.yaml') action = ApplyParameterOverwrites(ExperimentConfig(sampleModel='test.yaml')) action.update_header(self.header) self.assertIsNone(self.header.sample.model) with open('test.yaml', 'w') as fh: fh.write("""stack: air | Si | Fe""") try: action = ApplyParameterOverwrites(ExperimentConfig(sampleModel='test.yaml')) action.update_header(self.header) self.assertEqual(self.header.sample.model.stack, 'air | Si | Fe') finally: os.remove('test.yaml') def test_tof_time_correction(self): action = TofTimeCorrection() with self.assertRaises(ValueError): action.perform_action(self.d) new_events = append_fields(self.d.data.events, [('delta', np.float64)]) new_events.delta = 10.0 self.d.data.events = new_events tof_before = self.d.data.events.tof.copy() action.perform_action(self.d) np.testing.assert_allclose( self.d.data.events.tof, tof_before - (10.0 / 180.0) * self.d.timing.tau ) self.d.create_data() new_events = append_fields(self.d.data.events, [('delta', np.float64)]) new_events.delta = 10.0 self.d.data.events = new_events tof_before = self.d.data.events.tof.copy() action = TofTimeCorrection(correct_chopper_opening=False) action.perform_action(self.d) np.testing.assert_allclose( self.d.data.events.tof, tof_before - (self.d.geometry.kad / 180.0) * self.d.timing.tau ) def test_apply_mask(self): self.d.data.events = self.d.data.events[:6].copy() self.d.data.events.mask[:] = [0, 1, 2, 3, 4, 5] action = ApplyMask() action.perform_action(self.d) self.assertEqual(self.d.data.events.shape[0], 1) self.assertEqual(self.d.data.events.mask[0], 0) self.d.create_data() self.d.data.events = self.d.data.events[:6].copy() self.d.data.events.mask[:] = [0, 1, 2, 3, 4, 5] action = ApplyMask(bitmask_filter=EVENT_BITMASKS['MonitorThreshold']) action.perform_action(self.d) np.testing.assert_array_equal(self.d.data.events.mask, np.array([0, EVENT_BITMASKS['MonitorThreshold']], dtype=np.int32)) def test_merge_frames(self): 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, 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) self.assertIn('detZ', self.d.data.events.dtype.names) self.assertIn('detXdist', self.d.data.events.dtype.names) self.assertIn('delta', self.d.data.events.dtype.names) self.assertEqual( np.bitwise_and(self.d.data.events.mask, EVENT_BITMASKS['yRange']).astype(bool).sum(), self.d.data.events.shape[0] ) # TODO: maybe add a test actually checking correct detector-id resolution def test_calculate_wavelength(self): action = CalculateWavelength((3.0, 5.0)) with self.assertRaises(ValueError): action.perform_action(self.d) new_events = append_fields(self.d.data.events, [('detXdist', np.float64)]) new_events.detXdist = 0.0 self.d.data.events = new_events action.perform_action(self.d) self.assertIn('lamda', self.d.data.events.dtype.names) flt = self.d.data.events.mask!=EVENT_BITMASKS['LamdaRange'] # check all wavelength in range not filtered np.testing.assert_array_less(self.d.data.events.lamda[flt], 5.0) np.testing.assert_array_less(3.0, self.d.data.events.lamda[flt]) # check all wavelength out of range filtered flt = self.d.data.events.mask==EVENT_BITMASKS['LamdaRange'] self.assertTrue(((self.d.data.events.lamda[flt]<3.0)|(self.d.data.events.lamda[flt]>5.0)).all()) def test_calculate_q(self): action = CalculateQ(IncidentAngle.alphaF) with self.assertRaises(ValueError): action.perform_action(self.d) # TODO: add checks for actual resulting values new_events = append_fields(self.d.data.events, [('lamda', np.float64), ('delta', np.float64)]) new_events.lamda = 5.0 new_events.delta = 0.0 self.d.data.events = new_events action.perform_action(self.d) self.assertIn('qz', self.d.data.events.dtype.names) self.assertNotIn('qx', self.d.data.events.dtype.names) action.update_header(self.header) self.assertEqual(self.header.measurement_scheme, 'angle- and energy-dispersive') self.d.create_data() new_events = append_fields(self.d.data.events, [('lamda', np.float64), ('delta', np.float64)]) new_events.lamda = 5.0 new_events.delta = 0.0 self.d.data.events = new_events action = CalculateQ(IncidentAngle.mu) action.perform_action(self.d) self.assertIn('qz', self.d.data.events.dtype.names) self.assertIn('qx', self.d.data.events.dtype.names) action.update_header(self.header) self.assertEqual(self.header.measurement_scheme, 'energy-dispersive') self.d.create_data() new_events = append_fields(self.d.data.events, [('lamda', np.float64), ('delta', np.float64)]) new_events.lamda = 5.0 new_events.delta = 0.0 self.d.data.events = new_events action = CalculateQ(IncidentAngle.nu) action.perform_action(self.d) self.assertIn('qz', self.d.data.events.dtype.names) self.assertNotIn('qx', self.d.data.events.dtype.names) action.update_header(self.header) self.assertEqual(self.header.measurement_scheme, 'energy-dispersive') def test_filter_qz_range(self): action = FilterQzRange((0.1, 0.2)) with self.assertRaises(ValueError): action.perform_action(self.d) self.d.data.events = self.d.data.events[:5].copy() new_events = append_fields(self.d.data.events, [('qz', np.float64)]) new_events.qz = np.array([0.05, 0.1, 0.15, 0.2, 0.25]) self.d.data.events = new_events action.perform_action(self.d) np.testing.assert_array_equal( self.d.data.events.mask, np.array([1, 0, 0, 0, 1], dtype=np.int32) * EVENT_BITMASKS['qRange'] ) def test_filter_by_log(self): action = FilterByLog("test_log==0") | ApplyMask() class LogWarnError(Exception): ... def warn_raise(*args, **kwargs): raise LogWarnError() _orig_warn = logging.warning try: logging.warning = warn_raise with self.assertRaises(LogWarnError): action.perform_action(self.d) finally: logging.warning = _orig_warn self._extract_walltime() test_log = np.recarray(shape=(2,), dtype=np.dtype([('value', np.int32), ('time', np.int64)])) test_log.time = [-5, self.d.data.pulses.time[100]+123] test_log.value = [0, 1] self.d.data.device_logs['test_log'] = test_log action.perform_action(self.d) self.assertEqual(self.d.data.pulses.shape[0], 101) def test_filter_by_log_switchpulse(self): action = FilterByLog("!test_log==0") | ApplyMask() self._extract_walltime() test_log = np.recarray(shape=(2,), dtype=np.dtype([('value', np.int32), ('time', np.int64)])) test_log.time = [-5, self.d.data.pulses.time[100]+123] test_log.value = [0, 1] self.d.data.device_logs['test_log'] = test_log self.d.data.device_logs['check_log'] = test_log.copy() action.perform_action(self.d) self.assertEqual(self.d.data.pulses.shape[0], 100) np.testing.assert_array_equal( self.d.data.device_logs['test_log'], self.d.data.device_logs['check_log'], )