Merge branch 'develop' of https://github.com/ecmwf-lab/ecml-tools int…

…o develop

ecml_tools/data.py

@@ -869,6 +869,7 @@ def __init__(self, forward, mask):
         super().__init__(forward)
         assert len(forward.shape) == 4, "Grids must be 1D for now"
         self.mask = mask
+        self.axis = 3
 
     @cached_property
     def shape(self):
@@ -882,26 +883,27 @@ def latitudes(self):
     def longitudes(self):
         return self.forward.longitudes[self.mask]
 
+    @debug_indexing
     def __getitem__(self, index):
-        if isinstance(index, (int, slice)):
-            index = (index, slice(None), slice(None), slice(None))
-        return self._get_tuple(index)
+        if isinstance(index, tuple):
+            return self._get_tuple(index)
+
+        result = self.forward[index]
+        # We don't support subsetting the grid values
+        assert result.shape[-1] == len(self.mask), (result.shape, len(self.mask))
+
+        return result[..., self.mask]
 
     @debug_indexing
     @expand_list_indexing
     def _get_tuple(self, index):
-        assert self.axis >= len(index) or index[self.axis] == slice(
-            None
-        ), f"No support for selecting a subset of the 1D values {index}"
         index, changes = index_to_slices(index, self.shape)
-
-        # In case index_to_slices has changed the last slice
-        index, _ = update_tuple(index, self.axis, slice(None))
-
-        result = self.forwards[index]
-        result = result[:, :, :, self.mask]
-
-        return apply_index_to_slices_changes(result, changes)
+        index, previous = update_tuple(index, self.axis, slice(None))
+        result = self.forward[index]
+        result = result[..., self.mask]
+        result = result[..., previous]
+        result = apply_index_to_slices_changes(result, changes)
+        return result
 
 
 class Thinning(Masked):
@@ -1017,7 +1019,7 @@ def __getitem__(self, index):
     def _get_tuple(self, index):
         assert self.axis >= len(index) or index[self.axis] == slice(
             None
-        ), f"No support for selecting a subset of the 1D values {index}"
+        ), f"No support for selecting a subset of the 1D values {index} ({self.tree()})"
         index, changes = index_to_slices(index, self.shape)
 
         # In case index_to_slices has changed the last slice