138 lines
6.4 KiB
Python
138 lines
6.4 KiB
Python
"""
|
|
Define an event dataformat that performs reduction actions like wavelength calculation on per-event basis.
|
|
"""
|
|
import numpy as np
|
|
import logging
|
|
|
|
from typing import Tuple
|
|
|
|
from . import const
|
|
from .event_data_types import EventDataAction, EventDatasetProtocol, EVENT_TYPE, ANA_EVENT_TYPE, FINAL_EVENT_TYPE
|
|
from .helpers import filter_project_x
|
|
from .instrument import Detector
|
|
from .options import IncidentAngle
|
|
from .header import Header
|
|
|
|
|
|
class AnalyzedEventData(EventDataAction):
|
|
def __init__(self, yRange: Tuple[int, int]):
|
|
self.yRange = yRange
|
|
|
|
def perform_action(self, dataset: EventDatasetProtocol) ->None:
|
|
d = dataset.data
|
|
if d.events.dtype != EVENT_TYPE:
|
|
raise ValueError("AnalyzeEventData only works on raw AmorEventData, this dataset has already been altered")
|
|
pixelLookUp = self.resolve_pixels()
|
|
# TODO: change numba implementation to use native pixelID type
|
|
(detZ, detXdist, delta, mask) = filter_project_x(
|
|
pixelLookUp, d.events.pixelID.astype(np.int64), self.yRange[0], self.yRange[1]
|
|
)
|
|
ana_events = np.recarray(d.events.shape, dtype=ANA_EVENT_TYPE)
|
|
# copy old data
|
|
for field in d.events.dtype.fields.keys():
|
|
ana_events[field] = d.events[field]
|
|
# add analysis per event
|
|
ana_events.detZ = detZ
|
|
ana_events.detXdist = detXdist
|
|
ana_events.delta = delta
|
|
ana_events.mask = mask
|
|
d.events = ana_events
|
|
dataset.geometry.delta_z = self.delta_z
|
|
|
|
def resolve_pixels(self):
|
|
"""determine spatial coordinats and angles from pixel number"""
|
|
nPixel = Detector.nWires * Detector.nStripes * Detector.nBlades
|
|
pixelID = np.arange(nPixel)
|
|
(bladeNr, bPixel) = np.divmod(pixelID, Detector.nWires * Detector.nStripes)
|
|
(bZi, detYi) = np.divmod(bPixel, Detector.nStripes) # z index on blade, y index on detector
|
|
detZi = bladeNr * Detector.nWires + bZi # z index on detector
|
|
detX = bZi * Detector.dX # x position in detector
|
|
# detZ = Detector.zero - bladeNr * Detector.bladeZ - bZi * Detector.dZ # z position on detector
|
|
bladeAngle = np.rad2deg( 2. * np.arcsin(0.5*Detector.bladeZ / Detector.distance) )
|
|
delta = (Detector.nBlades/2. - bladeNr) * bladeAngle \
|
|
- np.rad2deg( np.arctan(bZi*Detector.dZ / ( Detector.distance + bZi * Detector.dX) ) )
|
|
self.delta_z = delta[detYi==1]
|
|
return np.vstack((detYi.T, detZi.T, detX.T, delta.T)).T
|
|
|
|
class TofTimeCorrection(EventDataAction):
|
|
def __init__(self, correct_chopper_opening: bool = True):
|
|
self.correct_chopper_opening = correct_chopper_opening
|
|
|
|
def perform_action(self, dataset: EventDatasetProtocol) ->None:
|
|
d = dataset.data
|
|
if d.events.dtype != ANA_EVENT_TYPE:
|
|
raise ValueError("TofTimeCorrection requires dataset with analyzed events, perform AnalyzedEventData first")
|
|
|
|
if self.correct_chopper_opening:
|
|
d.events.tof -= ( d.events.delta / 180. ) * dataset.timing.tau
|
|
else:
|
|
d.events.tof -= ( dataset.geometry.kad / 180. ) * dataset.timing.tau
|
|
|
|
class WavelengthAndQ(EventDataAction):
|
|
def __init__(self, lambdaRange: Tuple[float, float], incidentAngle: IncidentAngle):
|
|
self.lambdaRange = lambdaRange
|
|
self.incidentAngle = incidentAngle
|
|
|
|
def perform_action(self, dataset: EventDatasetProtocol) ->None:
|
|
d = dataset.data
|
|
if d.events.dtype != ANA_EVENT_TYPE:
|
|
raise ValueError("WavelengthAndQ requires dataset with analyzed events, perform AnalyzedEventData first")
|
|
|
|
self.lamdaMax = const.lamdaCut+1.e13*dataset.timing.tau*const.hdm/(dataset.geometry.chopperDetectorDistance+124.)
|
|
|
|
# lambda
|
|
lamda = (1.e13*const.hdm)*d.events.tof/(dataset.geometry.chopperDetectorDistance+d.events.detXdist)
|
|
|
|
final_events = np.recarray(d.events.shape, dtype=FINAL_EVENT_TYPE)
|
|
# copy old data
|
|
for field in d.events.dtype.fields.keys():
|
|
final_events[field] = d.events[field]
|
|
# add analysis per event
|
|
final_events.lamda = lamda
|
|
final_events.mask &= (self.lambdaRange[0]<=lamda) & (lamda<=self.lambdaRange[1])
|
|
|
|
# alpha_f
|
|
# q_z
|
|
if self.incidentAngle == IncidentAngle.alphaF:
|
|
alphaF_e = dataset.geometry.nu - dataset.geometry.mu + d.events.delta
|
|
final_events.qz = 4*np.pi*(np.sin(np.deg2rad(alphaF_e))/lamda)
|
|
elif self.incidentAngle == IncidentAngle.nu:
|
|
alphaF_e = (dataset.geometry.nu + d.events.delta + dataset.geometry.kap + dataset.geometry.kad) / 2.
|
|
final_events.qz = 4*np.pi*(np.sin(np.deg2rad(alphaF_e))/lamda)
|
|
else:
|
|
alphaF_e = dataset.geometry.nu - dataset.geometry.mu + d.events.delta
|
|
alphaI = dataset.geometry.kap + dataset.geometry.kad + dataset.geometry.mu
|
|
final_events.qz = 2*np.pi * ((np.sin(np.deg2rad(alphaF_e)) + np.sin(np.deg2rad(alphaI)))/lamda)
|
|
final_events.qx = 2*np.pi * ((np.cos(np.deg2rad(alphaF_e)) - np.cos(np.deg2rad(alphaI)))/lamda)
|
|
|
|
dataset.data.events = final_events
|
|
|
|
def update_header(self, header: Header):
|
|
if self.incidentAngle == IncidentAngle.alphaF:
|
|
header.measurement_scheme = 'angle- and energy-dispersive'
|
|
else:
|
|
header.measurement_scheme = 'energy-dispersive'
|
|
|
|
class FilterQzRange(EventDataAction):
|
|
def __init__(self, qzRange: Tuple[float, float]):
|
|
self.qzRange = qzRange
|
|
|
|
def perform_action(self, dataset: EventDatasetProtocol) ->None:
|
|
d = dataset.data
|
|
if d.events.dtype != FINAL_EVENT_TYPE:
|
|
raise ValueError("FilterQzRange requires dataset with fully analyzed events, perform WavelengthAndQ first")
|
|
|
|
if self.qzRange[1]<0.5:
|
|
d.events.mask &= (self.qzRange[0]<=d.events.qz) & (d.events.qz<=self.qzRange[1])
|
|
|
|
class ApplyMask(EventDataAction):
|
|
def perform_action(self, dataset: EventDatasetProtocol) ->None:
|
|
d = dataset.data
|
|
if not 'mask' in d.events.dtype.names:
|
|
logging.debug("ApplyMask performed on dataset without mask")
|
|
return
|
|
|
|
logging.info(f' number of events: total = {d.events.shape[0]:7d}, '
|
|
f'filtered = {np.logical_not(d.events.mask).sum():7d}')
|
|
d.events = d.events[d.events.mask]
|