From b664313e3a706ebc9f6498ee4a367aaffb9cff91 Mon Sep 17 00:00:00 2001
From: Simon Heybrock <simon.heybrock@ess.eu>
Date: Thu, 21 Aug 2025 13:06:54 +0200
Subject: [PATCH] Fix run mapping so it works also for Qxy mode

---
 src/ess/sans/workflow.py | 12 +++++++-----
 tests/loki/iofq_test.py  |  7 +++++++
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/ess/sans/workflow.py b/src/ess/sans/workflow.py
index 2459bb2c..45a3451e 100644
--- a/src/ess/sans/workflow.py
+++ b/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
 
 
diff --git a/tests/loki/iofq_test.py b/tests/loki/iofq_test.py
index 8d2af5e5..90d8549b 100644
--- a/tests/loki/iofq_test.py
+++ b/tests/loki/iofq_test.py
@@ -13,6 +13,7 @@
 from ess import loki, sans
 from ess.sans.conversions import ElasticCoordTransformGraph
 from ess.sans.types import (
+    BackgroundRun,
     BackgroundSubtractedIofQ,
     BackgroundSubtractedIofQxy,
     BeamCenter,
@@ -20,6 +21,7 @@
     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():