diff --git a/src/ess/reduce/nexus/__init__.py b/src/ess/reduce/nexus/__init__.py index 7785bec0..22eefad2 100644 --- a/src/ess/reduce/nexus/__init__.py +++ b/src/ess/reduce/nexus/__init__.py @@ -20,6 +20,8 @@ load_all_components, load_component, load_data, + load_field, + load_group, open_component_group, open_nexus_file, ) @@ -33,6 +35,8 @@ 'load_all_components', 'load_component', 'load_data', + 'load_field', + 'load_group', 'open_component_group', 'open_nexus_file', 'types', diff --git a/src/ess/reduce/nexus/_nexus_loader.py b/src/ess/reduce/nexus/_nexus_loader.py index b4375f74..596d9ee8 100644 --- a/src/ess/reduce/nexus/_nexus_loader.py +++ b/src/ess/reduce/nexus/_nexus_loader.py @@ -21,6 +21,7 @@ NeXusAllLocationSpec, NeXusEntryName, NeXusFile, + NeXusFileSpec, NeXusGroup, NeXusLocationSpec, ) @@ -42,6 +43,58 @@ def __repr__(self) -> str: NoLockingIfNeeded = NoLockingIfNeededType() +def load_field( + filename: NeXusFileSpec, + field_path: str, + selection: snx.typing.ScippIndex | slice = (), +) -> sc.Variable | sc.DataArray: + """Load a single field from a NeXus file. + + Parameters + ---------- + filename: + Path of the file to load from. + field_path: + Path of the field within the NeXus file. + selection: + Selection to apply to the field. + + Returns + ------- + : + The loaded field as a variable or data array. + """ + with open_nexus_file(filename.value) as f: + field = f[field_path] + return cast(sc.Variable | sc.DataArray, field[selection]) + + +def load_group( + filename: NeXusFileSpec, + group_path: str, + selection: snx.typing.ScippIndex | slice = (), +) -> sc.DataGroup: + """Load a single group from a NeXus file. + + Parameters + ---------- + filename: + Path of the file to load from. + group_path: + Path of the group within the NeXus file. + selection: + Selection to apply to the group. + + Returns + ------- + : + The loaded group as a data group. + """ + with open_nexus_file(filename.value) as f: + group = f[group_path] + return cast(sc.DataGroup, group[selection]) + + def load_component( location: NeXusLocationSpec, *, diff --git a/src/ess/reduce/nexus/workflow.py b/src/ess/reduce/nexus/workflow.py index 03a4ae75..81142865 100644 --- a/src/ess/reduce/nexus/workflow.py +++ b/src/ess/reduce/nexus/workflow.py @@ -691,16 +691,12 @@ def LoadMonitorWorkflow( def LoadDetectorWorkflow( - *, - run_types: Iterable[sciline.typing.Key], - monitor_types: Iterable[sciline.typing.Key], + *, run_types: Iterable[sciline.typing.Key] ) -> sciline.Pipeline: """Generic workflow for loading detector data from a NeXus file.""" wf = sciline.Pipeline( (*_common_providers, *_detector_providers), - constraints=_gather_constraints( - run_types=run_types, monitor_types=monitor_types - ), + constraints=_gather_constraints(run_types=run_types, monitor_types=[]), ) wf[DetectorBankSizes] = DetectorBankSizes({}) wf[PreopenNeXusFile] = PreopenNeXusFile(False) diff --git a/tests/nexus/workflow_test.py b/tests/nexus/workflow_test.py index 3b75dbe7..7b52568d 100644 --- a/tests/nexus/workflow_test.py +++ b/tests/nexus/workflow_test.py @@ -4,11 +4,17 @@ from pathlib import Path import pytest +import sciline as sl import scipp as sc import scippnexus as snx from scipp.testing import assert_identical -from ess.reduce.nexus import compute_component_position, workflow +from ess.reduce.nexus import ( + compute_component_position, + load_field, + load_group, + workflow, +) from ess.reduce.nexus.types import ( BackgroundRun, Beamline, @@ -21,6 +27,7 @@ Measurement, MonitorType, NeXusComponentLocationSpec, + NeXusFileSpec, NeXusName, NeXusTransformation, PreopenNeXusFile, @@ -577,7 +584,7 @@ def test_load_histogram_monitor_workflow(dream_coda_test_file: Path) -> None: def test_load_detector_workflow(loki_tutorial_sample_run_60250: Path) -> None: - wf = LoadDetectorWorkflow(run_types=[SampleRun], monitor_types=[]) + wf = LoadDetectorWorkflow(run_types=[SampleRun]) wf[Filename[SampleRun]] = loki_tutorial_sample_run_60250 wf[NeXusName[snx.NXdetector]] = 'larmor_detector' da = wf.compute(RawDetector[SampleRun]) @@ -587,7 +594,7 @@ def test_load_detector_workflow(loki_tutorial_sample_run_60250: Path) -> None: def test_load_histogram_detector_workflow(tbl_commissioning_orca_file: Path) -> None: - wf = LoadDetectorWorkflow(run_types=[SampleRun], monitor_types=[]) + wf = LoadDetectorWorkflow(run_types=[SampleRun]) wf[Filename[SampleRun]] = tbl_commissioning_orca_file wf[NeXusName[snx.NXdetector]] = 'orca_detector' da = wf.compute(RawDetector[SampleRun]) @@ -600,7 +607,7 @@ def test_load_histogram_detector_workflow(tbl_commissioning_orca_file: Path) -> def test_load_empty_histogram_detector_workflow( tbl_commissioning_orca_file: Path, ) -> None: - wf = LoadDetectorWorkflow(run_types=[SampleRun], monitor_types=[]) + wf = LoadDetectorWorkflow(run_types=[SampleRun]) wf[Filename[SampleRun]] = tbl_commissioning_orca_file wf[NeXusName[snx.NXdetector]] = 'orca_detector' da = wf.compute(EmptyDetector[SampleRun]) @@ -776,3 +783,44 @@ def assert_not_contains_type_arg(node: object, excluded: set[type]) -> None: assert not any( arg in excluded for arg in getattr(node, "__args__", ()) ), f"Node {node} contains one of {excluded!r}" + + +def test_generic_nexus_workflow_load_custom_field_user_affiliation( + loki_tutorial_sample_run_60250: Path, +) -> None: + class UserAffiliation(sl.Scope[RunType, str], str): + """User affiliation.""" + + def load_user_affiliation( + file: NeXusFileSpec[RunType], + path: NeXusName[UserAffiliation[RunType]], + ) -> UserAffiliation[RunType]: + return UserAffiliation[RunType](load_field(file, path)) + + wf = GenericNeXusWorkflow(run_types=[SampleRun], monitor_types=[]) + wf.insert(load_user_affiliation) + wf[Filename[SampleRun]] = loki_tutorial_sample_run_60250 + wf[NeXusName[UserAffiliation[SampleRun]]] = '/entry/user_0/affiliation' + affiliation = wf.compute(UserAffiliation[SampleRun]) + assert affiliation == 'ESS' + + +def test_generic_nexus_workflow_load_custom_group_user( + loki_tutorial_sample_run_60250: Path, +) -> None: + class UserInfo(sl.Scope[RunType, str], str): + """User info.""" + + def load_user_info( + file: NeXusFileSpec[RunType], + path: NeXusName[UserInfo[RunType]], + ) -> UserInfo[RunType]: + return UserInfo[RunType](load_group(file, path)) + + wf = GenericNeXusWorkflow(run_types=[SampleRun], monitor_types=[]) + wf.insert(load_user_info) + wf[Filename[SampleRun]] = loki_tutorial_sample_run_60250 + wf[NeXusName[UserInfo]] = '/entry/user_0' + user_info = wf.compute(UserInfo[SampleRun]) + assert user_info['affiliation'] == 'ESS' + assert user_info['name'] == 'John Doe'