Merge pull request #786 from DHI/isel

Use named arguments to isel/sel

Author: ecomodeller

mikeio/dataset/_dataarray.py

@@ -81,6 +81,8 @@
     Grid3D,
 ]
 
+IndexType = Union[int, slice, Sequence[int], np.ndarray, None]
+
 
 class _DataArraySpectrumToHm0:
     def __init__(self, da: "DataArray") -> None:
@@ -584,9 +586,18 @@ def __setitem__(self, key: Any, value: np.ndarray) -> None:
 
     def isel(
         self,
-        idx: int | Sequence[int] | slice | None = None,
+        idx: IndexType = None,
+        *,
+        time: IndexType = None,
+        x: IndexType = None,
+        y: IndexType = None,
+        z: IndexType = None,
+        element: IndexType = None,
+        node: IndexType = None,
+        layer: IndexType = None,
+        frequency: IndexType = None,
+        direction: IndexType = None,
         axis: int | str = 0,
-        **kwargs: Any,
     ) -> "DataArray":
         """Return a new DataArray whose data is given by
         integer indexing along the specified dimension(s).
@@ -617,11 +628,17 @@ def isel(
             y index, by default None
         z : int, optional
             z index, by default None
+        layer: int, optional
+            layer index, only used in dfsu 3d
+        direction: int, optional
+            direction index, only used in sprectra
+        frequency: int, optional
+            frequencey index, only used in spectra
+        node: int, optional
+            node index, only used in spectra
         element : int, optional
             Bounding box of coordinates (left lower and right upper)
             to be selected, by default None
-        **kwargs: Any
-            Not used
 
         Returns
         -------
@@ -654,10 +671,23 @@ def isel(
         ```
 
         """
-        if isinstance(self.geometry, Grid2D) and ("x" in kwargs and "y" in kwargs):
-            idx_x = kwargs["x"]
-            idx_y = kwargs["y"]
-            return self.isel(x=idx_x).isel(y=idx_y)
+        if isinstance(self.geometry, Grid2D) and (x is not None and y is not None):
+            return self.isel(x=x).isel(y=y)
+        kwargs = {
+            k: v
+            for k, v in dict(
+                time=time,
+                x=x,
+                y=y,
+                z=z,
+                element=element,
+                node=node,
+                layer=layer,
+                frequency=frequency,
+                direction=direction,
+            ).items()
+            if v is not None
+        }
         for dim in kwargs:
             if dim in self.dims:
                 axis = dim
@@ -698,7 +728,7 @@ def isel(
                 spatial_axis = axis - 1 if self.dims[0] == "time" else axis
                 geometry = self.geometry.isel(idx, axis=spatial_axis)
 
-            # TOOD this is ugly
+            # TODO this is ugly
             if isinstance(geometry, _GeometryFMLayered):
                 node_ids, _ = self.geometry._get_nodes_and_table_for_elements(
                     idx, node_layers="all"
@@ -741,7 +771,12 @@ def sel(
         self,
         *,
         time: str | pd.DatetimeIndex | "DataArray" | None = None,
-        **kwargs: Any,
+        x: float | slice | None = None,
+        y: float | slice | None = None,
+        z: float | slice | None = None,
+        coords: np.ndarray | None = None,
+        area: tuple[float, float, float, float] | None = None,
+        layers: int | str | Sequence[int | str] | None = None,
     ) -> "DataArray":
         """Return a new DataArray whose data is given by
         selecting index labels along the specified dimension(s).
@@ -780,8 +815,6 @@ def sel(
             layer(s) to be selected: "top", "bottom" or layer number
             from bottom 0,1,2,... or from the top -1,-2,... or as
             list of these; only for layered dfsu, by default None
-        **kwargs: Any
-            Additional keyword arguments
 
         Returns
         -------
@@ -823,24 +856,32 @@ def sel(
         ```
 
         """
+        # time is not part of kwargs
+        kwargs = {
+            k: v
+            for k, v in dict(
+                x=x, y=y, z=z, area=area, coords=coords, layers=layers
+            ).items()
+            if v is not None
+        }
         if any([isinstance(v, slice) for v in kwargs.values()]):
-            return self._sel_with_slice(kwargs)
+            return self._sel_with_slice(kwargs)  # type: ignore
 
         da = self
 
         # select in space
         if len(kwargs) > 0:
             idx = self.geometry.find_index(**kwargs)
+
+            # TODO this seems fragile
             if isinstance(idx, tuple):
                 # TODO: support for dfs3
                 assert len(idx) == 2
-                t_ax_offset = 1 if self._has_time_axis else 0
                 ii, jj = idx
                 if jj is not None:
-                    da = da.isel(idx=jj, axis=(0 + t_ax_offset))
+                    da = da.isel(y=jj)
                 if ii is not None:
-                    sp_axis = 0 if (jj is not None and len(jj) == 1) else 1
-                    da = da.isel(idx=ii, axis=(sp_axis + t_ax_offset))
+                    da = da.isel(x=ii)
             else:
                 da = da.isel(idx, axis="space")
 
@@ -1866,7 +1907,7 @@ def to_dfs(self, filename: str | Path, **kwargs: Any) -> None:
             Dfs0 only: set the dfs data type of the written data
             to e.g. np.float64, by default: DfsSimpleType.Float (=np.float32)
         **kwargs: Any
-            Additional keyword arguments, e.g. dtype for dfs0
+            additional arguments passed to the writing function, e.g. dtype for dfs0
 
         """
         self._to_dataset().to_dfs(filename, **kwargs)

mikeio/dataset/_dataset.py

@@ -44,6 +44,8 @@
 
 from ._data_plot import _DatasetPlotter
 
+from ._dataarray import IndexType
+
 
 class Dataset:
     """Dataset containing one or more DataArrays with common geometry and time.
@@ -431,7 +433,7 @@ def dropna(self) -> "Dataset":
             else:
                 all_index = list(np.intersect1d(all_index, idx))
 
-        return self.isel(all_index, axis=0)
+        return self.isel(time=all_index)
 
     def flipud(self) -> "Dataset":
         """Flip data upside down (on first non-time axis)."""
@@ -476,7 +478,7 @@ def create_data_array(
     def create_empty_data(
         n_items: int = 1,
         n_timesteps: int = 1,
-        n_elements: int | None = None,
+        n_elements: IndexType = None,
         shape: tuple[int, ...] | None = None,
     ) -> list:
         data = []
@@ -641,15 +643,15 @@ def __getitem__(self, key: Any) -> DataArray | "Dataset":
             time_steps = _get_time_idx_list(self.time, key)
             if _n_selected_timesteps(self.time, time_steps) == 0:
                 raise IndexError("No timesteps found!")
-            return self.isel(time_steps, axis=0)
+            return self.isel(time=time_steps)
         if isinstance(key, slice):
             if self._is_slice_time_slice(key):
                 try:
                     s = self.time.slice_indexer(key.start, key.stop)
                     time_steps = list(range(s.start, s.stop))
                 except ValueError:
                     time_steps = list(range(*key.indices(len(self.time))))
-                return self.isel(time_steps, axis=0)
+                return self.isel(time=time_steps)
 
         if self._multi_indexing_attempted(key):
             raise TypeError(
@@ -757,9 +759,18 @@ def __delitem__(self, key: Hashable | int) -> None:
 
     def isel(
         self,
-        idx: int | Sequence[int] | slice | None = None,
+        idx: IndexType = None,
+        *,
+        time: IndexType = None,
+        x: IndexType = None,
+        y: IndexType = None,
+        z: IndexType = None,
+        element: IndexType = None,
+        node: IndexType = None,
+        layer: IndexType = None,
+        frequency: IndexType = None,
+        direction: IndexType = None,
         axis: int | str = 0,
-        **kwargs: Any,
     ) -> "Dataset":
         """Return a new Dataset whose data is given by
         integer indexing along the specified dimension(s).
@@ -789,8 +800,14 @@ def isel(
         element : int, optional
             Bounding box of coordinates (left lower and right upper)
             to be selected, by default None
-        **kwargs: Any
-            Not used
+        layer: int, optional
+            layer index, only used in dfsu 3d
+        direction: int, optional
+            direction index, only used in sprectra
+        frequency: int, optional
+            frequencey index, only used in spectra
+        node: int, optional
+            node index, only used in spectra
 
         Returns
         -------
@@ -809,12 +826,36 @@ def isel(
         >>> ds3 = ds2.isel(elements=[100,200])
 
         """
-        res = [da.isel(idx=idx, axis=axis, **kwargs) for da in self]
+        # TODO deprecate idx, axis to prefer x= instead
+
+        res = [
+            da.isel(
+                idx=idx,
+                axis=axis,
+                time=time,
+                x=x,
+                y=y,
+                z=z,
+                element=element,
+                node=node,
+                frequency=frequency,
+                direction=direction,
+                layer=layer,
+            )
+            for da in self
+        ]
         return Dataset(data=res, validate=False)
 
     def sel(
         self,
-        **kwargs: Any,
+        *,
+        time: Any = None,
+        x: float | None = None,
+        y: float | None = None,
+        z: float | None = None,
+        coords: np.ndarray | None = None,
+        area: tuple[float, float, float, float] | None = None,
+        layers: int | str | Sequence[int | str] | None = None,
     ) -> "Dataset":
         """Return a new Dataset whose data is given by
         selecting index labels along the specified dimension(s).
@@ -853,8 +894,6 @@ def sel(
             layer(s) to be selected: "top", "bottom" or layer number
             from bottom 0,1,2,... or from the top -1,-2,... or as
             list of these; only for layered dfsu, by default None
-        **kwargs: Any
-            Not used
 
         Returns
         -------
@@ -878,7 +917,10 @@ def sel(
         >>> ds.sel(layers="bottom")
 
         """
-        res = [da.sel(**kwargs) for da in self]
+        res = [
+            da.sel(time=time, x=x, y=y, z=z, coords=coords, area=area, layers=layers)
+            for da in self
+        ]
         return Dataset(data=res, validate=False)
 
     def interp(
@@ -1849,15 +1891,16 @@ def to_dfs(self, filename: str | Path, **kwargs: Any) -> None:
         """
         filename = str(filename)
 
+        # TODO is this a candidate for match/case?
         if isinstance(
             self.geometry, (GeometryPoint2D, GeometryPoint3D, GeometryUndefined)
         ):
             if self.ndim == 0:  # Not very common, but still...
                 self._validate_extension(filename, ".dfs0")
-                self._to_dfs0(filename, **kwargs)
+                self._to_dfs0(filename=filename, **kwargs)
             elif self.ndim == 1 and self[0]._has_time_axis:
                 self._validate_extension(filename, ".dfs0")
-                self._to_dfs0(filename, **kwargs)
+                self._to_dfs0(filename=filename, **kwargs)
             else:
                 raise ValueError("Cannot write Dataset with no geometry to file!")
         elif isinstance(self.geometry, Grid2D):
@@ -1885,12 +1928,15 @@ def _validate_extension(filename: str | Path, valid_extension: str) -> None:
         if ext != valid_extension:
             raise ValueError(f"File extension must be {valid_extension}")
 
-    def _to_dfs0(self, filename: str | Path, **kwargs: Any) -> None:
+    def _to_dfs0(
+        self,
+        filename: str | Path,
+        dtype: DfsSimpleType = DfsSimpleType.Float,
+        title: str = "",
+    ) -> None:
         from ..dfs._dfs0 import _write_dfs0
 
-        dtype = kwargs.get("dtype", DfsSimpleType.Float)
-
-        _write_dfs0(filename, self, dtype=dtype)
+        _write_dfs0(filename, self, dtype=dtype, title=title)
 
     def _to_dfs2(self, filename: str | Path) -> None:
         # assumes Grid2D geometry

mikeio/dfs/_dfs0.py

@@ -189,7 +189,7 @@ def read(
                 _, time_steps = _valid_timesteps(dfs.FileInfo, time)
 
         if time_steps:
-            ds = ds.isel(time_steps, axis=0)
+            ds = ds.isel(time=time_steps)
 
         if sel_time_step_str:
             parts = sel_time_step_str.split(",")

tests/test_consistency.py

@@ -33,7 +33,7 @@ def test_read_dfs1():
 def test_dfs1_isel_t():
     ds = mikeio.read("tests/testdata/consistency/oresundHD.dfs1")
 
-    ds1 = ds.isel([0, 1], axis="t")
+    ds1 = ds.isel(time=[0, 1])
     assert ds1.dims == ("time", "x")
     assert isinstance(ds1.geometry, type(ds.geometry))
     assert ds1[0].values[0, 8] == pytest.approx(0.203246)
@@ -42,10 +42,10 @@ def test_dfs1_isel_t():
 def test_dfs1_isel_x():
     ds = mikeio.read("tests/testdata/consistency/oresundHD.dfs1")
 
-    ds1 = ds.isel(8, axis="x")
+    ds1 = ds.isel(x=8)
     assert ds1.dims == ("time",)
     assert isinstance(ds1.geometry, GeometryUndefined)
-    assert ds1[0].isel(0, axis="time").values == pytest.approx(0.203246)
+    assert ds1[0].isel(time=0).values == pytest.approx(0.203246)
 
 
 def test_dfs1_sel_t():
@@ -63,13 +63,13 @@ def test_dfs1_sel_x():
     ds1 = ds.sel(x=7.8)
     assert ds1.dims == ("time",)
     assert isinstance(ds1.geometry, GeometryUndefined)
-    assert ds1[0].isel(0, axis="time").values == pytest.approx(0.203246)
+    assert ds1[0].isel(time=0).values == pytest.approx(0.203246)
 
     da: DataArray = ds[0]
     da1 = da.sel(x=7.8)
     assert da1.dims == ("time",)
     assert isinstance(ds1.geometry, GeometryUndefined)
-    assert da1.isel(0, axis="time").values == pytest.approx(0.203246)
+    assert da1.isel(time=0).values == pytest.approx(0.203246)
 
 
 def test_dfs1_interp_x():
@@ -78,7 +78,7 @@ def test_dfs1_interp_x():
     ds1 = ds.interp(x=7.75)
     assert ds1.dims == ("time",)
     assert isinstance(ds1.geometry, GeometryUndefined)
-    assert ds1[0].isel(0, axis="time").values == pytest.approx(0.20202248)
+    assert ds1[0].isel(time=0).values == pytest.approx(0.20202248)
 
 
 # Nice to have...