feat: add workflow helper for combining files

Author: jokasimr

src/ess/reflectometry/conversions.py

@@ -9,6 +9,7 @@
 from .types import (
     BeamDivergenceLimits,
     DataWithScatteringCoordinates,
+    DetectorRotation,
     Gravity,
     IncidentBeam,
     MaskedData,
@@ -109,6 +110,7 @@ def specular_reflection(
     incident_beam: IncidentBeam[RunType],
     sample_position: SamplePosition[RunType],
     sample_rotation: SampleRotation[RunType],
+    detector_rotation: DetectorRotation[RunType],
     gravity: Gravity,
 ) -> SpecularReflectionCoordTransformGraph[RunType]:
     """
@@ -127,6 +129,7 @@ def specular_reflection(
         "incident_beam": lambda: incident_beam,
         "sample_position": lambda: sample_position,
         "sample_rotation": lambda: sample_rotation,
+        "detector_rotation": lambda: detector_rotation,
         "gravity": lambda: gravity,
     }
     return SpecularReflectionCoordTransformGraph(graph)
@@ -136,7 +139,10 @@ def add_coords(
     da: ReducibleDetectorData[RunType],
     graph: SpecularReflectionCoordTransformGraph[RunType],
 ) -> DataWithScatteringCoordinates[RunType]:
-    da = da.transform_coords(["theta", "wavelength", "Q"], graph=graph)
+    da = da.transform_coords(
+        ["theta", "wavelength", "Q", "detector_rotation"], graph=graph
+    )
+    da.coords.set_aligned('detector_rotation', False)
     da.coords["z_index"] = sc.arange(
         "row", 0, da.sizes["blade"] * da.sizes["wire"], unit=None
     ).fold("row", sizes={dim: da.sizes[dim] for dim in ("blade", "wire")})

src/ess/reflectometry/orso.py

@@ -59,6 +59,9 @@
 OrsoSample = NewType("OrsoSample", data_source.Sample)
 """ORSO sample."""
 
+OrsoSampleFilenames = NewType("OrsoSampleFilenames", list[orso_base.File])
+"""Collection of filenames used to create the ORSO file"""
+
 
 def parse_orso_experiment(filename: Filename[SampleRun]) -> OrsoExperiment:
     """Parse ORSO experiment metadata from raw NeXus data."""
@@ -107,8 +110,13 @@ def parse_orso_sample(filename: Filename[SampleRun]) -> OrsoSample:
     )
 
 
+def orso_data_files(filename: Filename[SampleRun]) -> OrsoSampleFilenames:
+    '''Collects names of files used in the experiment'''
+    return [orso_base.File(file=os.path.basename(filename))]
+
+
 def build_orso_measurement(
-    sample_filename: Filename[SampleRun],
+    sample_filenames: OrsoSampleFilenames,
     reference_filename: Filename[ReferenceRun],
     instrument: OrsoInstrument,
 ) -> OrsoMeasurement:
@@ -127,7 +135,7 @@ def build_orso_measurement(
     return OrsoMeasurement(
         data_source.Measurement(
             instrument_settings=instrument,
-            data_files=[orso_base.File(file=os.path.basename(sample_filename))],
+            data_files=sample_filenames,
             additional_files=additional_files,
         )
     )
@@ -220,4 +228,5 @@ def find_corrections(task_graph: TaskGraph) -> list[str]:
     parse_orso_experiment,
     parse_orso_owner,
     parse_orso_sample,
+    orso_data_files,
 )

src/ess/reflectometry/workflow.py

@@ -0,0 +1,76 @@
+from collections.abc import Hashable, Sequence
+from itertools import chain
+
+import pandas as pd
+import sciline
+import scipp as sc
+
+from ess.amor.types import RawChopper
+from ess.reflectometry.orso import (
+    OrsoExperiment,
+    OrsoOwner,
+    OrsoSample,
+    OrsoSampleFilenames,
+)
+from ess.reflectometry.types import (
+    Filename,
+    FootprintCorrectedData,
+    RunType,
+    SampleRotation,
+    SampleRun,
+)
+
+
+def _concatenate_event_lists(*das):
+    return (
+        sc.reduce(das)
+        .bins.concat()
+        .assign_coords(
+            {
+                name: das[0].coords[name]
+                for name in ('position', 'sample_rotation', 'detector_rotation')
+            }
+        )
+    )
+
+
+def _any_value(x, *_):
+    return x
+
+
+def _concatenate_lists(*x):
+    return list(chain(*x))
+
+
+def with_filenames(
+    workflow, runtype: Hashable, runs: Sequence[Filename[RunType]]
+) -> sciline.Pipeline:
+    axis_name = f'{str(runtype).lower()}_runs'
+    df = pd.DataFrame({Filename[runtype]: runs}).rename_axis(axis_name)
+    wf = workflow.copy()
+
+    mapped = wf.map(df)
+
+    wf[FootprintCorrectedData[runtype]] = mapped[
+        FootprintCorrectedData[runtype]
+    ].reduce(index=axis_name, func=_concatenate_event_lists)
+    wf[RawChopper[runtype]] = mapped[RawChopper[runtype]].reduce(
+        index=axis_name, func=_any_value
+    )
+    wf[SampleRotation[runtype]] = mapped[SampleRotation[runtype]].reduce(
+        index=axis_name, func=_any_value
+    )
+
+    if runtype is SampleRun:
+        wf[OrsoSample] = mapped[OrsoSample].reduce(index=axis_name, func=_any_value)
+        wf[OrsoExperiment] = mapped[OrsoExperiment].reduce(
+            index=axis_name, func=_any_value
+        )
+        wf[OrsoOwner] = mapped[OrsoOwner].reduce(index=axis_name, func=lambda x, *_: x)
+        wf[OrsoSampleFilenames] = mapped[OrsoSampleFilenames].reduce(
+            # When we don't map over filenames
+            # each OrsoSampleFilenames is a list with a single entry.
+            index=axis_name,
+            func=_concatenate_lists,
+        )
+    return wf

tests/amor/pipeline_test.py

@@ -7,6 +7,7 @@
 import sciline
 import scipp as sc
 from orsopy import fileio
+from scipp.testing import assert_allclose
 
 from ess import amor
 from ess.amor import data  # noqa: F401
@@ -23,6 +24,7 @@
     YIndexLimits,
     ZIndexLimits,
 )
+from ess.reflectometry.workflow import with_filenames
 
 
 @pytest.fixture
@@ -34,6 +36,9 @@ def amor_pipeline() -> sciline.Pipeline:
     pl[WavelengthBins] = sc.geomspace("wavelength", 2.8, 12, 300, unit="angstrom")
     pl[YIndexLimits] = sc.scalar(11, unit=None), sc.scalar(41, unit=None)
     pl[ZIndexLimits] = sc.scalar(80, unit=None), sc.scalar(370, unit=None)
+    pl[QBins] = sc.geomspace(
+        dim="Q", start=0.005, stop=0.115, num=391, unit="1/angstrom"
+    )
 
     # The sample rotation value in the file is slightly off, so we set it manually
     pl[SampleRotation[ReferenceRun]] = sc.scalar(0.65, unit="deg")
@@ -46,19 +51,15 @@ def amor_pipeline() -> sciline.Pipeline:
             contact="[email protected]",
         )
     )
-
-    # The sample rotation value in the file is slightly off, so we set it manually
-    pl[SampleRotation[SampleRun]] = sc.scalar(0.85, unit="deg")
-    pl[Filename[SampleRun]] = amor.data.amor_sample_run(608)
-    pl[QBins] = sc.geomspace(
-        dim="Q", start=0.005, stop=0.115, num=391, unit="1/angstrom"
-    )
     return pl
 
 
 @pytest.mark.filterwarnings("ignore:Failed to convert .* into a transformation")
 @pytest.mark.filterwarnings("ignore:Invalid transformation, missing attribute")
 def test_run_data_pipeline(amor_pipeline: sciline.Pipeline):
+    # The sample rotation value in the file is slightly off, so we set it manually
+    amor_pipeline[SampleRotation[SampleRun]] = sc.scalar(0.85, unit="deg")
+    amor_pipeline[Filename[SampleRun]] = amor.data.amor_sample_run(608)
     res = amor_pipeline.compute(ReflectivityOverQ)
     assert "Q" in res.coords
     assert "Q_resolution" in res.coords
@@ -67,6 +68,8 @@ def test_run_data_pipeline(amor_pipeline: sciline.Pipeline):
 @pytest.mark.filterwarnings("ignore:Failed to convert .* into a transformation")
 @pytest.mark.filterwarnings("ignore:Invalid transformation, missing attribute")
 def test_run_full_pipeline(amor_pipeline: sciline.Pipeline):
+    amor_pipeline[SampleRotation[SampleRun]] = sc.scalar(0.85, unit="deg")
+    amor_pipeline[Filename[SampleRun]] = amor.data.amor_sample_run(608)
     res = amor_pipeline.compute(orso.OrsoIofQDataset)
     assert res.info.data_source.experiment.instrument == "Amor"
     assert res.info.reduction.software.name == "ess.reflectometry"
@@ -75,7 +78,53 @@ def test_run_full_pipeline(amor_pipeline: sciline.Pipeline):
     assert np.all(res.data[:, 1] >= 0)
 
 
+@pytest.mark.filterwarnings("ignore:Failed to convert .* into a transformation")
+@pytest.mark.filterwarnings("ignore:Invalid transformation, missing attribute")
+def test_pipeline_can_compute_reflectivity_merging_events_from_multiple_runs(
+    amor_pipeline: sciline.Pipeline,
+):
+    sample_runs = [
+        amor.data.amor_sample_run(608),
+        amor.data.amor_sample_run(609),
+    ]
+    pipeline = with_filenames(amor_pipeline, SampleRun, sample_runs)
+    pipeline[SampleRotation[SampleRun]] = pipeline.compute(
+        SampleRotation[SampleRun]
+    ) + sc.scalar(0.05, unit="deg")
+    result = pipeline.compute(ReflectivityOverQ)
+    assert result.dims == ('Q',)
+
+
+@pytest.mark.filterwarnings("ignore:Failed to convert .* into a transformation")
+@pytest.mark.filterwarnings("ignore:Invalid transformation, missing attribute")
+def test_pipeline_merging_events_result_unchanged(amor_pipeline: sciline.Pipeline):
+    sample_runs = [
+        amor.data.amor_sample_run(608),
+    ]
+    pipeline = with_filenames(amor_pipeline, SampleRun, sample_runs)
+    pipeline[SampleRotation[SampleRun]] = pipeline.compute(
+        SampleRotation[SampleRun]
+    ) + sc.scalar(0.05, unit="deg")
+    result = pipeline.compute(ReflectivityOverQ).hist()
+    sample_runs = [
+        amor.data.amor_sample_run(608),
+        amor.data.amor_sample_run(608),
+    ]
+    pipeline = with_filenames(amor_pipeline, SampleRun, sample_runs)
+    pipeline[SampleRotation[SampleRun]] = pipeline.compute(
+        SampleRotation[SampleRun]
+    ) + sc.scalar(0.05, unit="deg")
+    result2 = pipeline.compute(ReflectivityOverQ).hist()
+    assert_allclose(
+        2 * sc.values(result.data), sc.values(result2.data), rtol=sc.scalar(1e-6)
+    )
+    assert_allclose(
+        2 * sc.variances(result.data), sc.variances(result2.data), rtol=sc.scalar(1e-6)
+    )
+
+
 def test_find_corrections(amor_pipeline: sciline.Pipeline):
+    amor_pipeline[Filename[SampleRun]] = amor.data.amor_sample_run(608)
     graph = amor_pipeline.get(orso.OrsoIofQDataset)
     # In topological order
     assert orso.find_corrections(graph) == [