Fix run mapping so it works also for Qxy mode

src/ess/sans/workflow.py

@@ -30,6 +30,7 @@
     WavelengthBands,
     WavelengthMask,
     WavelengthScaledQ,
+    WavelengthScaledQxy,
 )
 
 
@@ -99,11 +100,12 @@ def _set_runs(
     pipeline = pipeline.copy()
     runs = pd.DataFrame({Filename[key]: runs}).rename_axis(axis_name)
     for part in (Numerator, Denominator):
-        pipeline[WavelengthScaledQ[key, part]] = (
-            pipeline[WavelengthScaledQ[key, part]]
-            .map(runs)
-            .reduce(index=axis_name, func=merge_contributions)
-        )
+        for qtype in (WavelengthScaledQ, WavelengthScaledQxy):
+            pipeline[qtype[key, part]] = (
+                pipeline[qtype[key, part]]
+                .map(runs)
+                .reduce(index=axis_name, func=merge_contributions)
+            )
     return pipeline
 
 

tests/loki/iofq_test.py

@@ -13,13 +13,15 @@
 from ess import loki, sans
 from ess.sans.conversions import ElasticCoordTransformGraph
 from ess.sans.types import (
+    BackgroundRun,
     BackgroundSubtractedIofQ,
     BackgroundSubtractedIofQxy,
     BeamCenter,
     CleanWavelength,
     CorrectForGravity,
     Denominator,
     DimsToKeep,
+    Filename,
     IofQ,
     IofQxy,
     MaskedData,
@@ -192,11 +194,16 @@ def test_pipeline_can_compute_IofQ_merging_events_from_multiple_runs():
     pipeline = make_workflow()
     pipeline[BeamCenter] = _compute_beam_center()
 
+    # Remove previously set runs so we can be sure that below we use the mapped ones
+    pipeline[Filename[SampleRun]] = None
+    pipeline[Filename[BackgroundRun]] = None
     pipeline = sans.with_sample_runs(pipeline, runs=sample_runs)
     pipeline = sans.with_background_runs(pipeline, runs=background_runs)
 
     result = pipeline.compute(BackgroundSubtractedIofQ)
     assert result.dims == ('Q',)
+    result = pipeline.compute(BackgroundSubtractedIofQxy)
+    assert result.dims == ('Qy', 'Qx')
 
 
 def test_pipeline_can_compute_IofQ_by_bank():