Fix FileManager slicing

changelog/453.bugfix.rst

@@ -0,0 +1 @@
+Minor tweak to correct indexing of >4D datasets.

dkist/conftest.py

@@ -267,6 +267,31 @@ def dataset_4d(identity_gwcs_4d, empty_meta):
     return Dataset(array, wcs=identity_gwcs_4d, meta=empty_meta, unit=u.count)
 
 
+@pytest.fixture
+def dataset_5d(identity_gwcs_5d_stokes, empty_meta):
+    shape = (4, 40, 30, 20, 10)
+    x = np.ones(shape)
+    array = da.from_array(x, tuple(shape))
+
+    identity_gwcs_4d.pixel_shape = array.shape[::-1]
+    identity_gwcs_4d.array_shape = array.shape
+
+    ds = Dataset(array, wcs=identity_gwcs_5d_stokes, meta={"inventory": {}, "headers": Table()}, unit=u.count)
+    fileuris = np.array([f"dummyfile_{i}" for i in range(np.prod(shape[:-2]))]).reshape(shape[:-2])
+    ds._file_manager = FileManager.from_parts(fileuris, 0, float, shape[-2:], loader=AstropyFITSLoader, basepath="./")
+
+    return ds
+
+
+@pytest.fixture
+def dataset_5d_dummy_filemanager_axis(dataset_5d):
+    shape = dataset_5d.data.shape
+    fileuris = np.array([f"dummyfile_{i}" for i in range(np.prod(shape[:-2]))]).reshape(shape[:-2])
+    dataset_5d._file_manager = FileManager.from_parts(fileuris, 0, float, (1, *shape[-2:]), loader=AstropyFITSLoader, basepath="./")
+
+    return dataset_5d
+
+
 @pytest.fixture
 def eit_dataset():
     eitdir = Path(rootdir) / "EIT"
@@ -350,6 +375,16 @@ def visp_dataset_no_headers(tmp_path_factory):
     return load_dataset(vispdir / "test_visp_no_headers.asdf")
 
 
+@pytest.fixture
+def large_visp_no_dummy_axis(large_visp_dataset):
+    # Slightly tweaked dataset to remove the dummy axis in the file manager array shape.
+    shape = large_visp_dataset.data.shape[:2]
+    fileuris = np.array([f"dummyfile_{i}" for i in range(np.prod(shape))]).reshape(shape)
+    large_visp_dataset._file_manager = FileManager.from_parts(fileuris, 0, float, (50, 128), loader=AstropyFITSLoader, basepath="./")
+
+    return large_visp_dataset
+
+
 @pytest.hookimpl(hookwrapper=True)
 def pytest_runtest_call(item):
     try:

dkist/dataset/tests/test_dataset.py

@@ -176,3 +176,25 @@ def test_header_slicing_3D_slice(large_visp_dataset):
 
     assert len(sliced.files.filenames) == len(sliced_headers["FILENAME"]) == len(sliced.headers)
     assert (sliced.headers["DINDEX3", "DINDEX4"] == sliced_headers["DINDEX3", "DINDEX4"]).all()
+
+
+@pytest.mark.accept_cli_dataset
+def test_file_slicing_with_dummy_axis(dataset_5d_dummy_filemanager_axis):
+    ds = dataset_5d_dummy_filemanager_axis
+    shape = ds.data.shape
+    assert len(ds.files) == np.prod(shape[:3])
+    assert len(ds[0].files) == np.prod(shape[1:3])
+    assert len(ds[0, 0].files) == np.prod(shape[2])
+    assert len(ds[0, 0, 0].files) == 1
+    assert len(ds[0, 0, 0, 0].files) == 1
+
+
+@pytest.mark.accept_cli_dataset
+def test_file_slicing_without_dummy_axis(dataset_5d):
+    ds = dataset_5d
+    shape = ds.data.shape
+    assert len(ds.files) == np.prod(shape[:3])
+    assert len(ds[0].files) == np.prod(shape[1:3])
+    assert len(ds[0, 0].files) == np.prod(shape[2])
+    assert len(ds[0, 0, 0].files) == 1
+    assert len(ds[0, 0, 0, 0].files) == 1

dkist/io/file_manager.py

@@ -242,7 +242,7 @@ def _array_slice_to_loader_slice(self, aslice):
         aslice = list(sanitize_slices(aslice, len(self.output_shape)))
         if fits_array_shape[0] == 1:
             # Insert a blank slice for the dummy dimension
-            aslice.insert(len(fits_array_shape) - 1, slice(None))
+            aslice.insert(-(len(fits_array_shape)-1), slice(None))
         # Now only use the dimensions of the slice not covered by the array axes
         aslice = aslice[:-1*len(fits_array_shape)]
         return tuple(aslice)