From a42ac927ccc227e090bac96f73d854796b2241a2 Mon Sep 17 00:00:00 2001 From: YooSunyoung Date: Tue, 29 Oct 2024 16:39:46 +0100 Subject: [PATCH] McStas 3.5 version workaround. --- src/ess/nmx/mcstas/__init__.py | 31 ++++++++- src/ess/nmx/mcstas/{xml.py => geometry.py} | 76 +++++++++++++++++++++- src/ess/nmx/mcstas/load.py | 57 +++++++++++++++- src/ess/nmx/reduction.py | 2 +- 4 files changed, 162 insertions(+), 4 deletions(-) rename src/ess/nmx/mcstas/{xml.py => geometry.py} (80%) diff --git a/src/ess/nmx/mcstas/__init__.py b/src/ess/nmx/mcstas/__init__.py index 24397a5..69fdd00 100644 --- a/src/ess/nmx/mcstas/__init__.py +++ b/src/ess/nmx/mcstas/__init__.py @@ -8,8 +8,37 @@ def McStasWorkflow(): from ess.nmx.reduction import bin_time_of_arrival from .load import providers as loader_providers - from .xml import read_mcstas_geometry_xml + from .geometry import read_mcstas_geometry_xml return sl.Pipeline( (*loader_providers, read_mcstas_geometry_xml, bin_time_of_arrival) ) + + +def McStas35Workflow(): + import sciline as sl + + from ess.nmx.reduction import bin_time_of_arrival + + from .load import providers as loader_providers + from .load import ( + load_raw_event_data_mcstas_35, + load_raw_event_data, + load_crystal_rotation_mcstas_35, + load_crystal_rotation, + ) + from .geometry import load_mcstas_geometry + + loader_providers = list(loader_providers) + loader_providers.remove(load_raw_event_data) + loader_providers.remove(load_crystal_rotation) + + return sl.Pipeline( + ( + *loader_providers, + load_mcstas_geometry, + bin_time_of_arrival, + load_raw_event_data_mcstas_35, + load_crystal_rotation_mcstas_35, + ) + ) diff --git a/src/ess/nmx/mcstas/xml.py b/src/ess/nmx/mcstas/geometry.py similarity index 80% rename from src/ess/nmx/mcstas/xml.py rename to src/ess/nmx/mcstas/geometry.py index c872e68..4bee90a 100644 --- a/src/ess/nmx/mcstas/xml.py +++ b/src/ess/nmx/mcstas/geometry.py @@ -11,7 +11,7 @@ from defusedxml.ElementTree import fromstring from ..rotation import axis_angle_to_quaternion, quaternion_to_matrix -from ..types import FilePath +from ..types import DetectorName, FilePath T = TypeVar('T') @@ -410,3 +410,77 @@ def read_mcstas_geometry_xml(file_path: FilePath) -> McStasInstrument: with h5py.File(file_path) as file: tree = fromstring(file[instrument_xml_path][...][0]) return McStasInstrument.from_xml(tree) + + +def load_mcstas_geometry( + file_path: FilePath, component_name: DetectorName +) -> McStasInstrument: + """Retrieve geometric information from mcstas file""" + import scippnexus as snx + from scipy.spatial.transform import Rotation + + with snx.File(file_path, 'r') as f: + root = f[f"entry1/instrument/components/{component_name}"] + position_x_offset = root['output/BINS/x__m_'][()] + position_y_offset = root['output/BINS/y__m_'][()] + step_x = position_x_offset['x__m_', 1] - position_x_offset['x__m_', 0] + step_x.unit = 'm' + step_y = position_y_offset['y__m_', 1] - position_y_offset['y__m_', 0] + step_y.unit = 'm' + position = root['Position'][()] + origin_position_vector = sc.vector([*position.values], unit='m') + det_rotation = root['Rotation'][()] + det_rotation_matrix = sc.spatial.rotations_from_rotvecs( + rotation_vectors=sc.vector( + Rotation.from_matrix(det_rotation.values).as_rotvec(degrees=False), + unit='rad', + ) + ) # Not sure if it is the correct way... + pixel_ids = root['output/BINS/pixels'][()] + sim_settings = SimulationSettings( + length_unit='m', angle_unit='degree', beam_axis='z', handedness='right' + ) # Hard-coded for now. + sample_position = f["entry1/instrument/components/sampleMantid/Position"][()] + sample_rotation = f["entry1/instrument/components/sampleMantid/Rotation"][()] + sample_rotation_matrix = sc.spatial.rotations_from_rotvecs( + rotation_vectors=sc.vector( + Rotation.from_matrix(sample_rotation.values).as_rotvec(degrees=False), + unit='rad', + ) + ) # Not sure if it is the correct way... + sample_desc = SampleDesc( + component_type='sampleMantid', + name='sample', + position=sc.vector([*sample_position.values], unit='m'), + rotation_matrix=sample_rotation_matrix, + ) + source_position = f["entry1/instrument/components/sourceMantid/Position"][()] + source_desc = SourceDesc( + component_type='Source', + name='source', + position=sc.vector([*source_position.values], unit='m'), + ) + return McStasInstrument( + simulation_settings=sim_settings, + detectors=( + DetectorDesc( + component_type='MonNDtype', + name=component_name, + id_start=pixel_ids.min().value, + fast_axis_name='x', + num_x=pixel_ids.shape[1], + num_y=pixel_ids.shape[0], + step_x=step_x, + step_y=step_y, + start_x=position_x_offset.min().value, + start_y=position_y_offset.min().value, + rotation_matrix=det_rotation_matrix, + slow_axis_name='y', + fast_axis=det_rotation_matrix * _AXISNAME_TO_UNIT_VECTOR['x'], + slow_axis=det_rotation_matrix * _AXISNAME_TO_UNIT_VECTOR['y'], + position=origin_position_vector, + ), + ), + source=source_desc, + sample=sample_desc, + ) diff --git a/src/ess/nmx/mcstas/load.py b/src/ess/nmx/mcstas/load.py index 02ed179..1dcf690 100644 --- a/src/ess/nmx/mcstas/load.py +++ b/src/ess/nmx/mcstas/load.py @@ -17,7 +17,7 @@ ProtonCharge, RawEventData, ) -from .xml import McStasInstrument, read_mcstas_geometry_xml +from .geometry import McStasInstrument, read_mcstas_geometry_xml def detector_name_from_index(index: DetectorIndex) -> DetectorName: @@ -176,6 +176,61 @@ def load_mcstas( ) +def load_crystal_rotation_mcstas_35(file_path: FilePath) -> CrystalRotation: + """Retrieve crystal rotation from the file. + + Raises + ------ + KeyError + If the crystal rotation is not found in the file. + + """ + from scipy.spatial.transform import Rotation + + with snx.File(file_path, 'r') as f: + root = f["entry1/instrument/components/sampleMantid"] + rotation = root['Rotation'][()] + rotation_matrix = sc.spatial.rotations_from_rotvecs( + rotation_vectors=sc.vector( + Rotation.from_matrix(rotation.values).as_rotvec(degrees=False), + unit='rad', + ) + ) # Not sure if it is the correct way... + return CrystalRotation(rotation_matrix) + + +def load_raw_event_data_mcstas_35( + file_path: FilePath, + bank_prefix: DetectorBankPrefix, + detector_name: DetectorName, +) -> RawEventData: + """Retrieve events from the nexus file.""" + bank_name = f'{bank_prefix}_dat_list_p_x_y_n_id_t' + with snx.File(file_path, 'r') as f: + root = f["entry1/data"] + pixel_ids: sc.Variable = ( + f[f"entry1/instrument/components/{detector_name}/output/BINS/pixels"][()] + .astype(int) + .flatten(to='id') + ) + (bank_name,) = (name for name in root.keys() if bank_name in name) + data = root[bank_name]["events"][()].rename_dims({'dim_0': 'event'}) + return sc.DataArray( + coords={ + 'id': sc.array( + dims=['event'], + values=data['dim_1', 4].values, + dtype='int64', + unit=None, + ), + 't': sc.array(dims=['event'], values=data['dim_1', 5].values, unit='s'), + }, + data=sc.array( + dims=['event'], values=data['dim_1', 0].values, unit='counts' + ), + ).group(pixel_ids) + + providers = ( read_mcstas_geometry_xml, detector_name_from_index, diff --git a/src/ess/nmx/reduction.py b/src/ess/nmx/reduction.py index 1bb3df4..8128303 100644 --- a/src/ess/nmx/reduction.py +++ b/src/ess/nmx/reduction.py @@ -4,7 +4,7 @@ import scipp as sc -from .mcstas.xml import McStasInstrument +from .mcstas.geometry import McStasInstrument from .types import DetectorName, TimeBinSteps NMXData = NewType("NMXData", sc.DataGroup)