Assign sumo coord in load_geant4_csv

Author: jl-wynen

src/ess/dream/io/geant4.py

@@ -69,6 +69,7 @@ def _load_raw_events(file_path: str) -> sc.DataArray:
     table = sc.io.load_csv(
         file_path, sep="\t", header_parser="bracket", data_columns=[]
     )
+    table.coords['sumo'] = table.coords['det ID']
     table.coords.pop("lambda", None)
     table = table.rename_dims(row="event")
     return sc.DataArray(
@@ -90,10 +91,14 @@ def group(key: str, da: sc.DataArray) -> sc.DataArray:
         if key in ["high_resolution", "sans"]:
             # Only the HR and SANS detectors have sectors.
             res = da.group("sector", *elements)
+        elif key in ["endcap_backward", "endcap_forward"]:
+            # Other banks only have a single SUMO.
+            res = da.group("sumo", *elements)
         else:
             res = da.group(*elements)
         res.coords['position'] = res.bins.coords.pop('position').bins.mean()
         res.bins.coords.pop("sector", None)
+        res.bins.coords.pop("sumo", None)
         return res
 
     return {key: sc.DataGroup(events=group(key, da)) for key, da in detectors.items()}

tests/dream/geant4_reduction_test.py

@@ -94,8 +94,15 @@ def test_workflow_is_deterministic(workflow):
 
 def test_pipeline_can_compute_intermediate_results(workflow):
     workflow = powder.with_pixel_mask_filenames(workflow, [])
-    result = workflow.compute(NormalizedByProtonCharge[SampleRun])
-    assert set(result.dims) == {'segment', 'wire', 'counter', 'strip', 'module'}
+    results = workflow.compute((NormalizedByProtonCharge[SampleRun], NeXusDetectorName))
+    result = results[NormalizedByProtonCharge[SampleRun]]
+
+    detector_name = results[NeXusDetectorName]
+    expected_dims = {'segment', 'wire', 'counter', 'strip', 'module'}
+    if detector_name in ('endcap_backward', 'endcap_forward'):
+        expected_dims.add('sumo')
+
+    assert set(result.dims) == expected_dims
 
 
 def test_pipeline_group_by_two_theta(workflow):

tests/dream/io/geant4_test.py

@@ -114,6 +114,7 @@ def test_load_geant4_csv_mantle_has_expected_coords(file):
     assert_index_coord(mantle.coords["wire"], values=set(range(1, 33)))
     assert_index_coord(mantle.coords["strip"], values=set(range(1, 257)))
     assert "sector" not in mantle.coords
+    assert "sumo" not in mantle.coords
 
     assert "sector" not in mantle.bins.coords
     assert "tof" in mantle.bins.coords
@@ -127,6 +128,7 @@ def test_load_geant4_csv_endcap_backward_has_expected_coords(file):
     assert_index_coord(endcap.coords["counter"])
     assert_index_coord(endcap.coords["wire"], values=set(range(1, 17)))
     assert_index_coord(endcap.coords["strip"], values=set(range(1, 17)))
+    assert_index_coord(endcap.coords["sumo"], values=set(range(3, 7)))
     assert "sector" not in endcap.coords
 
     assert "sector" not in endcap.bins.coords
@@ -141,6 +143,7 @@ def test_load_geant4_csv_endcap_forward_has_expected_coords(file):
     assert_index_coord(endcap.coords["counter"])
     assert_index_coord(endcap.coords["wire"], values=set(range(1, 17)))
     assert_index_coord(endcap.coords["strip"], values=set(range(1, 17)))
+    assert_index_coord(endcap.coords["sumo"], values=set(range(3, 7)))
     assert "sector" not in endcap.coords
 
     assert "sector" not in endcap.bins.coords
@@ -156,6 +159,7 @@ def test_load_geant4_csv_high_resolution_has_expected_coords(file):
     assert_index_coord(hr.coords["wire"], values=set(range(1, 17)))
     assert_index_coord(hr.coords["strip"], values=set(range(1, 33)))
     assert_index_coord(hr.coords["sector"], values=set(range(1, 5)))
+    assert "sumo" not in hr.coords
 
     assert "tof" in hr.bins.coords
     assert "position" in hr.coords
@@ -172,6 +176,7 @@ def test_load_geant4_csv_sans_has_expected_coords(file):
         assert_index_coord(sans.coords["wire"], values=set(range(1, 17)))
         assert_index_coord(sans.coords["strip"], values=set(range(1, 33)))
         assert_index_coord(sans.coords["sector"], values=set(range(1, 5)))
+    assert "sumo" not in sans.coords
 
     assert "tof" in sans.bins.coords
     assert "position" in sans.coords