Assign sumo coord in load_geant4_csv

scipp · Sep 18, 2024 · 8ce50f6 · 8ce50f6
1 parent fa65af7
commit 8ce50f6
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 1 deletion.
diff --git a/src/ess/dream/io/geant4.py b/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()}

diff --git a/tests/dream/geant4_reduction_test.py b/tests/dream/geant4_reduction_test.py
@@ -95,7 +95,7 @@ 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'}
+    assert set(result.dims) == {'sumo', 'segment', 'wire', 'counter', 'strip', 'module'}
 
 
 def test_pipeline_group_by_two_theta(workflow):

diff --git a/tests/dream/io/geant4_test.py b/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