diff --git a/packages/essnmx/pyproject.toml b/packages/essnmx/pyproject.toml index 4f79d4231..2fe21ca8d 100644 --- a/packages/essnmx/pyproject.toml +++ b/packages/essnmx/pyproject.toml @@ -49,6 +49,7 @@ dynamic = ["version"] [project.scripts] essnmx_reduce_mcstas = "ess.nmx.mcstas.executables:main" essnmx-reduce = "ess.nmx.executables:main" +essmandi-reduce = "ess.mandi.workflows:main" [project.optional-dependencies] test = [ diff --git a/packages/essnmx/src/ess/mandi/__init__.py b/packages/essnmx/src/ess/mandi/__init__.py new file mode 100644 index 000000000..e392b3b27 --- /dev/null +++ b/packages/essnmx/src/ess/mandi/__init__.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2025 Scipp contributors (https://github.com/scipp) +# ruff: noqa: RUF100, E402, I + +import importlib.metadata + +try: + __version__ = importlib.metadata.version("essnmx") +except importlib.metadata.PackageNotFoundError: + __version__ = "0.0.0" + +del importlib + +# from .workflows import MandiWorkflow + +# __all__ = ["MandiWorkflow"] diff --git a/packages/essnmx/src/ess/mandi/_idf_helper.py b/packages/essnmx/src/ess/mandi/_idf_helper.py new file mode 100644 index 000000000..3d99bbe9c --- /dev/null +++ b/packages/essnmx/src/ess/mandi/_idf_helper.py @@ -0,0 +1,437 @@ +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2026 Scipp contributors (https://github.com/scipp) +# Mantid IDF related functions. +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from types import MappingProxyType +from typing import Protocol + +import h5py +import scipp as sc +from defusedxml.ElementTree import fromstring + +from ess.nmx.rotation import axis_angle_to_quaternion, quaternion_to_matrix +from ess.reduce.nexus.types import FilePath + +_AXISNAME_TO_UNIT_VECTOR = MappingProxyType( + { + 'x': sc.vector([1.0, 0.0, 0.0]), + 'y': sc.vector([0.0, 1.0, 0.0]), + 'z': sc.vector([0.0, 0.0, 1.0]), + } +) + + +class _XML(Protocol): + """XML element or tree type. + + Temporarily used for type hinting. + Builtin XML type is blocked by bandit security check.""" + + tag: str + attrib: dict[str, str] + + def find(self, name: str) -> '_XML | None': ... + + def findall(self, tag: str) -> 'Iterable[_XML]': ... + + def get(self, name: str, default: str | None = None) -> str: ... + + def __iter__(self) -> '_XML': ... + + def __next__(self) -> '_XML': ... + + +@dataclass +class DetectorDesc: + """Detector information extracted from McStas instrument xml description.""" + + # Name defined in the location. + name: str + id_start: int # 'idstart' + fast_axis_name: str # 'idfillbyfirst' + # From + num_x: int # 'xpixels' + num_y: int # 'ypixels' + step_x: sc.Variable # 'xstep' + step_y: sc.Variable # 'ystep' + start_x: float # 'xstart' + start_y: float # 'ystart' + # From under + position: sc.Variable # 'x', 'y', 'z' + # Calculated fields + rotation_matrix: sc.Variable + slow_axis_name: str + fast_axis: sc.Variable + slow_axis: sc.Variable + + @property + def total_pixels(self) -> int: + return self.num_x * self.num_y + + @property + def slow_step(self) -> sc.Variable: + return self.step_y if self.fast_axis_name == 'x' else self.step_x + + @property + def fast_step(self) -> sc.Variable: + return self.step_x if self.fast_axis_name == 'x' else self.step_y + + @property + def num_fast_pixels_per_row(self) -> int: + """Number of pixels in each row of the detector along the fast axis.""" + return self.num_x if self.fast_axis_name == 'x' else self.num_y + + @property + def detector_shape(self) -> tuple: + """Shape of the detector panel. (num_x, num_y)""" + return (self.num_x, self.num_y) + + @property + def pixel_ids(self) -> sc.Variable: + start, stop = ( + self.id_start, + self.id_start + self.total_pixels, + ) + return sc.arange('event_id', start, stop, unit=None) + + @property + def pixel_positions(self) -> sc.Variable: + # Assuming sample is always at 0,0,0 + pixel_idx = sc.arange('event_id', self.total_pixels) + n_col = sc.scalar(self.num_fast_pixels_per_row) + + pixel_n_slow = pixel_idx // n_col + pixel_n_fast = pixel_idx % n_col + + fast_axis_steps = self.fast_axis * self.fast_step + slow_axis_steps = self.slow_axis * self.slow_step + + return self.position + ( + (pixel_n_slow * slow_axis_steps) + + (pixel_n_fast * fast_axis_steps) + + self.rotation_matrix + * sc.vector( + [self.start_x, self.start_y, 0.0], unit=self.position.unit + ) # Detector pixel offset should also be rotated first. + ) + + def fold(self, da: sc.DataArray) -> sc.DataArray: + sizes = {'x': self.num_x, 'y': self.num_y} + axis_names = (self.fast_axis_name, self.slow_axis_name) + sizes = {f'{i}_pixel_offset': sizes[i] for i in axis_names} + return da.fold(dim='event_id', sizes=sizes) + + +@dataclass +class SampleDesc: + """Sample description extracted from McStas instrument xml description.""" + + name: str + position: sc.Variable + + def position_from_sample(self, other: sc.Variable) -> sc.Variable: + """Position of ``other`` relative to the sample. + + All positions and distance are stored relative to the sample position. + + Parameters + ---------- + other: + Position of the other object in 3D vector. + + """ + return other - self.position + + +@dataclass +class SourceDesc: + """Source description extracted from IDF.""" + + # From + name: str + # From under + position: sc.Variable + + +@dataclass +class MonitorDesc: + """Monitor description extracted from IDF.""" + + name: str + position: sc.Variable + + +@dataclass +class MandiInstrument: + instrument_definition: str + detectors: tuple[DetectorDesc, ...] + monitors: tuple[MonitorDesc, ...] + source: SourceDesc + sample: SampleDesc + + +@dataclass(frozen=True) +class ReferenceFrame: + along_beam_axis: str + pointing_up_axis: str + handedness: str + + +@dataclass(frozen=True) +class DefaultSettings: + length_unit: str + angle_unit: str + reference_frame: ReferenceFrame + default_view: str + + +def _retrieve_default_settings(tree: _XML, find: Callable) -> DefaultSettings: + default_settings_xml = find(tree, "defaults") + reference_frame_xml = find(default_settings_xml, "reference-frame") + reference_frame = ReferenceFrame( + along_beam_axis=find(reference_frame_xml, "along-beam").get("axis"), + pointing_up_axis=find(reference_frame_xml, "pointing-up").get("axis"), + handedness=find(reference_frame_xml, "handedness").get("val"), + ) + default_settings = DefaultSettings( + length_unit=find(default_settings_xml, "length").get("unit"), + angle_unit=find(default_settings_xml, "angle").get("unit"), + default_view=find(default_settings_xml, "default-view").get("view"), + reference_frame=reference_frame, + ) + return default_settings + + +def _retrieve_location(tree: _XML, find: Callable, length_unit: str) -> sc.Variable: + # Sometimes one component contains multiple locations + location_xml = find(tree, "location") if not tree.tag.endswith("location") else tree + xyz = [float(location_xml.get(i, 0)) for i in "xyz"] + return sc.vector(xyz, unit=length_unit) + + +def _retrieve_source( + *, + all_types: list[_XML], + all_components: list[_XML], + find: Callable, + length_unit: str, +) -> SourceDesc: + # Use the first one assuming there is single source. + source_type = next( + type_xml for type_xml in all_types if type_xml.get('is') == "Source" + ) + source_type_name = source_type.get("name") or "" + source_component = next( + comp for comp in all_components if comp.get("type") == source_type_name + ) + location = _retrieve_location(source_component, find=find, length_unit=length_unit) + + return SourceDesc(name=source_type_name, position=location) + + +def _retrieve_sample( + *, + all_types: list[_XML], + all_components: list[_XML], + find: Callable, + length_unit: str, +) -> SampleDesc: + # Use the first one assuming there is single sample. + sample_type = next( + type_xml for type_xml in all_types if type_xml.get('is') == "SamplePos" + ) + sample_type_name = sample_type.get("name") or "" + sample_component = next( + comp for comp in all_components if comp.get("type") == sample_type_name + ) + location = _retrieve_location(sample_component, find=find, length_unit=length_unit) + + return SampleDesc(name=sample_type_name, position=location) + + +def _retrieve_monitors( + *, + all_types: list[_XML], + find: Callable, + findall: Callable, + length_unit: str, +) -> list[MonitorDesc]: + # All monitors are defined under single "monitors" type. + monitors = next( + type_xml for type_xml in all_types if type_xml.get("name") == "monitors" + ) + monitor = find(monitors, "component") + locations = findall(monitor, "location") + return [ + MonitorDesc( + name=loc.get("name"), + position=_retrieve_location(loc, find=find, length_unit=length_unit), + ) + for loc in locations + ] + + +def _resolve_rotation_chain( + loc_tree: _XML, + *, + find: Callable, + angle_unit: str, + cur_matrix: sc.Variable | None = None, + handedness: str, +) -> sc.Variable: + """Resolve nested rotation chain. + + In the IDF, rotations can be nested. + MANDI does not have any translation/rotation chains + so this helper only resolves nested rotations. + + Returns + ------- + : + Rotation matrix. + + """ + try: + rot = find(loc_tree, "rot") + except KeyError: + return cur_matrix + + theta = sc.scalar(float(rot.get("val", 0)), unit=angle_unit) + if handedness == "right": + theta = -theta + + x, y, z, w = axis_angle_to_quaternion( + x=float(rot.get("axis-x", 0)), + y=float(rot.get("axis-y", 0)), + z=float(rot.get("axis-z", 1)), + theta=theta, + ) + new_matrix = quaternion_to_matrix(x=x, y=y, z=z, w=w) + + if cur_matrix is not None: + cur_matrix = cur_matrix * new_matrix + else: + cur_matrix = new_matrix + + return _resolve_rotation_chain( + loc_tree=rot, + find=find, + angle_unit=angle_unit, + cur_matrix=cur_matrix, + handedness=handedness, + ) + + +def _retrieve_detectors( + *, + all_types: list[_XML], + all_components: list[_XML], + find: Callable, + length_unit: str, + angle_unit: str, + handedness: str, +) -> list[DetectorDesc]: + detector_types = { + type_xml.get("name"): type_xml + for type_xml in all_types + if type_xml.get('is') == "rectangular_detector" + } + detector_type_names = set(detector_types.keys()) + detector_components = [ + comp for comp in all_components if comp.get("type") in detector_type_names + ] + detectors = [] + for comp in detector_components: + type_def = detector_types[comp.get("type")] + location_xml = find(comp, "location") + fast_axis_name = comp.get("idfillbyfirst") + slow_axis_name = 'x' if fast_axis_name == 'y' else 'y' + step_x = sc.scalar(float(type_def.get("xstep")), unit=length_unit) + step_y = sc.scalar(float(type_def.get("ystep")), unit=length_unit) + start_x = float(type_def.get("xstart")) + start_y = float(type_def.get("ystart")) + + position = _retrieve_location(location_xml, find=find, length_unit=length_unit) + rotation_matrix = _resolve_rotation_chain( + location_xml, find=find, angle_unit=angle_unit, handedness=handedness + ) + + cur_det = DetectorDesc( + name=location_xml.get("name"), + id_start=int(comp.get("idstart")), + fast_axis_name=fast_axis_name, + slow_axis_name=slow_axis_name, + num_x=int(type_def.get("xpixels")), + num_y=int(type_def.get("ypixels")), + step_x=step_x, + step_y=step_y, + start_x=start_x, + start_y=start_y, + position=position, + rotation_matrix=rotation_matrix, + fast_axis=rotation_matrix * _AXISNAME_TO_UNIT_VECTOR[fast_axis_name], + slow_axis=rotation_matrix * _AXISNAME_TO_UNIT_VECTOR[slow_axis_name], + ) + detectors.append(cur_det) + + return detectors + + +def read_mandi_geometry_xml(file_path: FilePath) -> MandiInstrument: + """Retrieve geometry parameters from Mandi file.""" + instrument_xml_path = 'entry/instrument/instrument_xml/data' + with h5py.File(file_path) as file: + idf_str: str = file[instrument_xml_path][...][0].decode() + tree = fromstring(idf_str) + + # Probably better way to retrieve the namespace... + namespace = tree.tag.removesuffix("instrument") + + def find(tree: _XML, tag) -> _XML: + elem = tree.find(f"{namespace}{tag}") + if elem is None: + raise KeyError(f"{tag=} not found in {elem=}") + return elem + + def findall(tree: _XML, tag) -> Iterable[_XML]: + return tree.findall(f"{namespace}{tag}") + + default_settings = _retrieve_default_settings(tree, find) + all_types = list(findall(tree, "type")) + all_components = list(findall(tree, "component")) + + source = _retrieve_source( + all_types=all_types, + all_components=all_components, + find=find, + length_unit=default_settings.length_unit, + ) + sample = _retrieve_sample( + all_types=all_types, + all_components=all_components, + find=find, + length_unit=default_settings.length_unit, + ) + monitors = _retrieve_monitors( + all_types=all_types, + find=find, + findall=findall, + length_unit=default_settings.length_unit, + ) + detectors = _retrieve_detectors( + all_types=all_types, + all_components=all_components, + find=find, + length_unit=default_settings.length_unit, + angle_unit=default_settings.angle_unit, + handedness=default_settings.reference_frame.handedness, + ) + + return MandiInstrument( + instrument_definition=idf_str, + detectors=tuple(detectors), + monitors=tuple(monitors), + sample=sample, + source=source, + ) diff --git a/packages/essnmx/src/ess/mandi/configurations.py b/packages/essnmx/src/ess/mandi/configurations.py new file mode 100644 index 000000000..c6550ac88 --- /dev/null +++ b/packages/essnmx/src/ess/mandi/configurations.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2025 Scipp contributors (https://github.com/scipp) +import enum + +from pydantic import BaseModel, Field, model_validator + +from ess.nmx.configurations import ( + AuxiliaryOutputConfig as NMXAuxiliaryOutputConfig, +) +from ess.nmx.configurations import ( + OutputConfig as NMXOutputConfig, +) + + +class InputConfig(BaseModel): + # Add title of the basemodel + model_config = {"title": "Input Configuration"} + # File IO + input_file: str = Field(title="Input File", description="Path to the input file.") + swmr: bool = Field( + title="SWMR Mode", + description="Open the input file in SWMR mode", + default=False, + ) + # Detector selection + ignore_list: list[str] = Field( + title="Detector names to be excluded.", + description="Detector indices to process", + default=["bank_error", "bank_unmapped"], + ) + + +class TimeBinUnit(enum.StrEnum): + ms = 'ms' + us = 'us' + ns = 'ns' + + +class _NotSet: ... + + +_notset = _NotSet() + + +class WorkflowConfig(BaseModel): + # Add title of the basemodel + @model_validator(mode='after') + def nbins_or_time_bin_width(self): + if self.time_bin_width is not None and self.nbins is not None: + raise ValueError( + "Either `nbins` or `time_bin_width` should be set. " + "They cannot be set at the same time. " + "It is allowed not setting any of them. " + "Then 300 [us] of `time_bin_width` will be used." + ) + return self + + @model_validator(mode='after') + def positive_time_bin_width(self): + if self.time_bin_width is not None and self.time_bin_width <= 0: + raise ValueError("`time_bin_width` should be a positive number.") + return self + + @model_validator(mode='after') + def positive_nbins(self): + if self.nbins is not None and self.nbins <= 0: + raise ValueError("`nbins` should be a positive integer.") + return self + + model_config = {"title": "Workflow Configuration"} + time_bin_width: int | None = Field( + title="Time Bin Width", + description="Width(Length) of each Time Bin in [time_bin_unit]. " + "If none of `time_bin_width` or `nbins` is given, " + "300 [us] of `time_bin_width` will be used.", + default=None, + ) + nbins: int | None = Field( + title="Number of Time Bins", + description="Number of Time bins. ", + default=None, + ) + min_time_bin: int | None = Field( + title="Minimum Time", + description="Minimum time edge of [time_bin_coordinate] in [time_bin_unit].", + default=None, + ) + max_time_bin: int | None = Field( + title="Maximum Time", + description="Maximum time edge of [time_bin_coordinate] in [time_bin_unit].", + default=None, + ) + time_bin_unit: TimeBinUnit = Field( + title="Unit of Time Bins", + description="Unit of time bins.", + default=TimeBinUnit.us, + ) + result_time_bin_unit: TimeBinUnit = Field( + title="Output Time Bin Unit", + description="Time bin unit of the histogram after reduction. " + "If the input time bin is different from the result time bin unit, " + "the unit will be converted to the result time bin " + "before the result is returned.", + default=TimeBinUnit.ns, + # DIALS expects [ns] by default. + ) + + +class AuxiliaryOutputConfig(NMXAuxiliaryOutputConfig, BaseModel): ... + + +class OutputConfig(NMXOutputConfig, BaseModel): + output_file: str = Field( + title="Output File", + description="Path to the output file. " + "It will be overwritten if ``overwrite`` is True.", + default="scipp_mandi_output.h5", + ) + + +class ReductionConfig(BaseModel): + """Container for all reduction configurations.""" + + inputs: InputConfig + workflow: WorkflowConfig = Field(default_factory=WorkflowConfig) + output: OutputConfig = Field(default_factory=OutputConfig) + aux: AuxiliaryOutputConfig = Field(default_factory=AuxiliaryOutputConfig) + + @property + def _children(self) -> list[BaseModel]: + return [self.inputs, self.workflow, self.output, self.aux] diff --git a/packages/essnmx/src/ess/mandi/workflows.py b/packages/essnmx/src/ess/mandi/workflows.py new file mode 100644 index 000000000..781ec3574 --- /dev/null +++ b/packages/essnmx/src/ess/mandi/workflows.py @@ -0,0 +1,303 @@ +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2026 Scipp contributors (https://github.com/scipp) +import argparse +import logging +import warnings +from collections.abc import Callable, Iterable + +import scipp as sc +import scippnexus as snx + +from ess.nmx._executable_helper import ( + add_args_from_pydantic_model, + build_logger, + from_args, +) +from ess.nmx.nexus import _check_file +from ess.nmx.types import ( + NMXDetectorMetadata, + NMXInstrument, + NMXLauetof, + NMXMonitorMetadata, + NMXReducedDetector, + NMXSampleMetadata, + NMXSourceMetadata, +) + +from ._idf_helper import read_mandi_geometry_xml +from .configurations import ( + AuxiliaryOutputConfig, + InputConfig, + OutputConfig, + ReductionConfig, + WorkflowConfig, +) + + +def build_reduction_argument_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Command line arguments for the Mandi reduction. " + "It assumes 60 Hz pulse speed." + ) + parser = add_args_from_pydantic_model(model_cls=InputConfig, parser=parser) + parser = add_args_from_pydantic_model(model_cls=WorkflowConfig, parser=parser) + parser = add_args_from_pydantic_model(model_cls=OutputConfig, parser=parser) + parser = add_args_from_pydantic_model( + model_cls=AuxiliaryOutputConfig, parser=parser + ) + return parser + + +def reduction_config_from_args(args: argparse.Namespace) -> ReductionConfig: + return ReductionConfig( + inputs=from_args(InputConfig, args), + workflow=from_args(WorkflowConfig, args), + output=from_args(OutputConfig, args), + aux=from_args(AuxiliaryOutputConfig, args), + ) + + +def _normalize_vector(vec: sc.Variable) -> sc.Variable: + return vec / sc.norm(vec) + + +def _build_mandi_time_bin_edges( + *, wf_config: WorkflowConfig, das: dict[str, sc.DataArray] +) -> sc.Variable: + """Build time bin edges. + + Mostly copied from ess.nmx.executables module. + However, for Mandi, we can build the time bin edges + before we group/bin the event data. + Therefore the helper function is slightly different. + """ + import numpy as np + + from ess.nmx.executables import _warn_bin_edge_out_of_range + + t_coord_name = "tof" + da_min_t = min(da.coords[t_coord_name].nanmin() for da in das.values()) + da_max_t = max(da.coords[t_coord_name].nanmax() for da in das.values()) + + # Use the user-set parameters if available + # and validate them according to the data. + # Lower Time Bin Edge + if wf_config.min_time_bin is not None: + min_t = sc.scalar(wf_config.min_time_bin, unit=wf_config.time_bin_unit) + min_t = min_t.to(unit=da_min_t.unit, dtype=da_min_t.dtype) + # If the user-set minimum time bin value + # is bigger than all time-bin-coordinate values. + if min_t > da_max_t: + _warn_bin_edge_out_of_range( + edge=min_t, coord_name=t_coord_name, desc='bigger' + ) + else: + min_t = da_min_t + + # Upper Time Bin Edge + if wf_config.max_time_bin is not None: + max_t = sc.scalar(wf_config.max_time_bin, unit=wf_config.time_bin_unit) + max_t = max_t.to(unit=da_max_t.unit, dtype=da_max_t.dtype) + # If the user-set maximum time bin value + # is smaller than all time-bin-coordinate values. + if max_t <= da_min_t: + _warn_bin_edge_out_of_range( + edge=max_t, coord_name=t_coord_name, desc='smaller' + ) + else: + max_t = da_max_t + + # Validate the results. + if min_t >= max_t: + raise ValueError( + f"Minimum time bin edge, {min_t} " + "is bigger than or equal to the " + f"maximum time bin edge, {max_t}.\n" + "Cannot build a time bin edges coordinate.\n" + "Please check your configurations again." + ) + + # If either min/max were manually selected and bin width is set. + if wf_config.nbins is None: + if wf_config.time_bin_width is None: + time_bin_width = sc.scalar(300, unit='us').to(unit=wf_config.time_bin_unit) + else: + time_bin_width = sc.scalar( + wf_config.time_bin_width, unit=wf_config.time_bin_unit + ) + # We do not return a scalar bin width since we histogram + # detector panels individually + # and all histograms should have the same bin edges. + min_t = min_t.to(unit=wf_config.time_bin_unit) + max_t = max_t.to(unit=wf_config.time_bin_unit) + bin_edges = sc.arange( + dim=t_coord_name, start=min_t, stop=max_t, step=time_bin_width + ) + # If the last bin edge is smaller than `max_t` + if bin_edges[t_coord_name, -1] <= max_t: + # Need to append one more edge to cover the whole range. + true_last_bin_edge = bin_edges[t_coord_name, -1] + time_bin_width + bin_edges = sc.concat([bin_edges, true_last_bin_edge], dim=t_coord_name) + + return bin_edges.to(dtype=float) + + else: # Number of bin edges are given but not the bin width. + n_edges = wf_config.nbins + 1 + if min_t.unit != max_t.unit: + min_t = min_t.to(unit=wf_config.time_bin_unit) + max_t = max_t.to(unit=wf_config.time_bin_unit) + + # Avoid dropping the event that has the exact same + # `event_time_offset`` or `tof` value as the upper bin edge. + max_t.value = np.nextafter(max_t.value, np.inf) + return sc.linspace( + dim=t_coord_name, start=min_t, stop=max_t, num=n_edges, dtype=float + ) + + +def _retrieve_display( + logger: logging.Logger | None, display: Callable | None +) -> Callable: + if display is not None: + return display + elif logger is not None: + return logger.info + else: + return logging.getLogger(__name__).info + + +def _sort_mandi_detbank_names(bank_names: Iterable[str]) -> list[str]: + """Sort the bank names. + + Bank names are sorted as bank1, bank2, bank22, bank5, bank51, + when they are loaded, + so this function helps to sort them by the number prefix instead. + """ + char_sorted = sorted(bank_names) + return sorted(char_sorted, key=lambda x: len(x)) + + +def reduction( + *, + config: ReductionConfig, + logger: logging.Logger | None = None, + display: Callable | None = None, +) -> NMXLauetof: + from ess.nmx.executables import save_results + + if not config.output.skip_file_output: + _check_file(config.output.output_file, config.output.overwrite) + config.aux.check_output_dir() + + display = _retrieve_display(logger, display) + + warnings.filterwarnings("ignore", category=UserWarning) + # Loading + with snx.File(config.inputs.input_file) as file: + detectors = dict( + filter( + lambda kv: kv[0] not in config.inputs.ignore_list, + file['entry/instrument'][snx.NXdetector].items(), + ) + ) + total_detectors = len(detectors) + banks = {} + # Sort the bank names + detectors = { + name: detectors[name] + for name in _sort_mandi_detbank_names(detectors.keys()) + } + for idet, (name, det) in enumerate(detectors.items()): + da = det[()]['events'].bins.concat().value.copy() + # Mandi files' event_time_offset is time-of-flight + da.coords['tof'] = da.coords.pop('event_time_offset') + banks[name] = da + display(f"{idet + 1}/{total_detectors} detector bank {name=} loaded.") + + mandi_geo = read_mandi_geometry_xml(config.inputs.input_file) + detector_dict = {det.name: det for det in mandi_geo.detectors} + tof_bin_edges = _build_mandi_time_bin_edges(wf_config=config.workflow, das=banks) + det_hists = {} + source_position = mandi_geo.source.position + sample_position = mandi_geo.sample.position + monitor_metadata = NMXMonitorMetadata( + tof_bin_coord='tof', + # TODO: Use real monitor data + data=sc.DataArray( + coords={'tof': tof_bin_edges}, + data=sc.ones_like(tof_bin_edges), + ), + ) + + sample_meta = NMXSampleMetadata( + # TODO: retrieve crystal rotation from the file correctly. + crystal_rotation=sc.vector([0.0, 0.0, 0.0], unit='deg'), + name=mandi_geo.sample.name, + position=sample_position, + ) + + for ibank, (name, bank) in enumerate(banks.items()): + if name not in detector_dict: + warnings.warn(f"Detector {name=} not found in the IDF.", stacklevel=2) + continue + + det_geo = detector_dict[name] + binned = bank.group(det_geo.pixel_ids) + hist = binned.hist(tof=tof_bin_edges.to(unit=bank.coords['tof'].unit)) + hist.coords['tof'] = hist.coords['tof'].to( + unit=config.workflow.result_time_bin_unit + ) + pixel_positions = det_geo.pixel_positions + origin = pixel_positions.mean() + distance = sc.norm(origin - source_position.to(unit=origin.unit)) + hist.coords['position'] = pixel_positions + # We save the first pixel position so that DIALS can read use it. + # first_pixel_position should be retrieved before folding. + first_pixel_number = hist.coords['event_id'].min() + first_pixel_position = hist['event_id', first_pixel_number].coords['position'] + first_pixel_position_from_sample = first_pixel_position - sample_position + + hist = det_geo.fold(hist) + detector_meta = NMXDetectorMetadata( + detector_name=name, + x_pixel_size=det_geo.step_x, + y_pixel_size=det_geo.step_y, + origin=origin, + fast_axis=_normalize_vector(det_geo.fast_axis), + fast_axis_dim=det_geo.fast_axis_name + '_pixel_offset', + slow_axis=_normalize_vector(det_geo.slow_axis), + slow_axis_dim=det_geo.slow_axis_name + '_pixel_offset', + distance=distance, + first_pixel_position=first_pixel_position_from_sample, + ) + det_hists[name] = NMXReducedDetector(data=hist, metadata=detector_meta) + display(f"{ibank + 1}/{total_detectors} bank {name=} reduced") + display(hist) + + instrument = NMXInstrument( + instrument_definition=mandi_geo.instrument_definition, + detectors=sc.DataGroup(det_hists), + name="MANDI", + source=NMXSourceMetadata(position=source_position), + ) + results = NMXLauetof( + control=monitor_metadata, + instrument=instrument, + sample=sample_meta, + ) + if not config.output.skip_file_output: + save_results( + results=results, + output_config=config.output, + aux_config=config.aux, + display=display, + ) + return results + + +def main() -> None: + parser = build_reduction_argument_parser() + config = reduction_config_from_args(parser.parse_args()) + logger = build_logger(config.output) + + reduction(config=config, logger=logger) diff --git a/packages/essnmx/src/ess/nmx/_display_helper.py b/packages/essnmx/src/ess/nmx/_display_helper.py index 676c3b9e5..dbac7bafa 100644 --- a/packages/essnmx/src/ess/nmx/_display_helper.py +++ b/packages/essnmx/src/ess/nmx/_display_helper.py @@ -9,21 +9,25 @@ def _is_nested(obj) -> bool: return is_dataclass(obj) or isinstance(obj, sc.DataGroup | dict) -def to_datagroup(obj) -> sc.DataGroup: +def to_datagroup(obj, *, drop_nones: bool = True) -> sc.DataGroup: if is_dataclass(obj): return sc.DataGroup( { - field.name: to_datagroup(value) - if _is_nested(value := getattr(obj, field.name)) + field.name: to_datagroup(value, drop_nones=drop_nones) + if _is_nested(value) else value for field in fields(obj) + if (value := getattr(obj, field.name)) is not None and drop_nones } ) elif isinstance(obj, sc.DataGroup | dict): return sc.DataGroup( { - name: to_datagroup(value) if _is_nested(value) else value + name: to_datagroup(value, drop_nones=drop_nones) + if _is_nested(value) + else value for name, value in obj.items() + if value is not None and drop_nones } ) else: diff --git a/packages/essnmx/src/ess/nmx/_nxlauetof_io.py b/packages/essnmx/src/ess/nmx/_nxlauetof_io.py index bd5bbfde6..91a83e134 100644 --- a/packages/essnmx/src/ess/nmx/_nxlauetof_io.py +++ b/packages/essnmx/src/ess/nmx/_nxlauetof_io.py @@ -194,6 +194,12 @@ def load_essnmx_nxlauetof(file: str | FilePath | NeXusFile) -> sc.DataGroup: _handle_sample(dg['entry']['sample'], entry['sample']) _handle_monitor(dg['entry']['control'], entry['control']) _handle_source(dg['entry']['instrument'], entry['instrument']) + # handle instrument definition - only needed for MANDI + if ( + 'instrument_definition' in dg['entry']['instrument'] + and dg['entry']['instrument']['instrument_definition'] is None + ): + dg['entry']['instrument'].pop('instrument_definition') _handle_detector_data(dg['entry']['instrument'], entry['instrument']) return dg['entry'] diff --git a/packages/essnmx/src/ess/nmx/executables.py b/packages/essnmx/src/ess/nmx/executables.py index dc1aa2888..00cac9563 100644 --- a/packages/essnmx/src/ess/nmx/executables.py +++ b/packages/essnmx/src/ess/nmx/executables.py @@ -355,7 +355,7 @@ def save_results( # Validate if results have expected fields export_static_metadata_as_nxlauetof( sample_metadata=results.sample, - source_metadata=results.instrument.source, + instrument_metadata=results.instrument, program=results.reducer, output_file=output_config.output_file, overwrite=output_config.overwrite, diff --git a/packages/essnmx/src/ess/nmx/nexus.py b/packages/essnmx/src/ess/nmx/nexus.py index 1a3bc6afb..f0c9a533d 100644 --- a/packages/essnmx/src/ess/nmx/nexus.py +++ b/packages/essnmx/src/ess/nmx/nexus.py @@ -13,10 +13,10 @@ from .configurations import Compression from .types import ( NMXDetectorMetadata, + NMXInstrument, NMXMonitorMetadata, NMXProgram, NMXSampleMetadata, - NMXSourceMetadata, ) @@ -146,7 +146,7 @@ def _set_default_instrument(nx_entry: snx.Group) -> snx.Group: def export_static_metadata_as_nxlauetof( *, sample_metadata: NMXSampleMetadata, - source_metadata: NMXSourceMetadata, + instrument_metadata: NMXInstrument, program: NMXProgram, output_file: str | pathlib.Path | io.BytesIO, overwrite: bool = False, @@ -166,6 +166,8 @@ def export_static_metadata_as_nxlauetof( Sample metadata object. source_metadata: Source metadata object. + instrument_metadata: + Instrument metadata object. monitor_metadata: Monitor metadata object. output_file: @@ -183,7 +185,12 @@ def export_static_metadata_as_nxlauetof( nx_entry['reducer'] = program nx_instrument = _set_default_instrument(nx_entry) - nx_instrument['source'] = source_metadata + nx_instrument['source'] = instrument_metadata.source + if instrument_metadata.instrument_definition is not None: + idf = nx_instrument.create_field( + 'IDF', value=instrument_metadata.instrument_definition + ) + idf.attrs['long_name'] = 'instrument definition xml' _add_arbitrary_metadata(nx_entry._group, **arbitrary_metadata) diff --git a/packages/essnmx/src/ess/nmx/types.py b/packages/essnmx/src/ess/nmx/types.py index 077e4ed37..c3e7450b3 100644 --- a/packages/essnmx/src/ess/nmx/types.py +++ b/packages/essnmx/src/ess/nmx/types.py @@ -249,6 +249,8 @@ class NMXReducedDetector: class NMXInstrument: nx_class = snx.NXinstrument + instrument_definition: str | None = None + """Instrument definition xml string.""" detectors: sc.DataGroup[NMXReducedDetector] name: str = "NMX" source: NMXSourceMetadata diff --git a/packages/essnmx/tests/nxlauetof_io_helper_test.py b/packages/essnmx/tests/nxlauetof_io_helper_test.py index b53bbaa82..bd66b1982 100644 --- a/packages/essnmx/tests/nxlauetof_io_helper_test.py +++ b/packages/essnmx/tests/nxlauetof_io_helper_test.py @@ -50,8 +50,6 @@ def test_loaded_data_same_as_in_memory_result( result = reduction(config=reduction_config) original_result_dg = result.to_datagroup() - # Adjust original result to be same as expected loaded data group. - original_result_dg.pop('lookup_table') original_positions = {} detectors = original_result_dg['instrument']['detectors'] for det_name, det in detectors.items():