diff --git a/src/ess/dream/io/geant4.py b/src/ess/dream/io/geant4.py index 7eda04f0..25c5a212 100644 --- 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 index 1cd15a4e..c57d5f4c 100644 --- 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 index d3396c3f..43e6b587 100644 --- 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