diff --git a/ecml_tools/data.py b/ecml_tools/data.py index de869df..316bf0c 100644 --- a/ecml_tools/data.py +++ b/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