Isel should respect order

mikeio/dfsu/_layered.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 from pathlib import Path
-from typing import Any, Collection, Sequence, Tuple, TYPE_CHECKING
+from typing import Any, Sequence, Tuple, TYPE_CHECKING
 
 import numpy as np
 from mikecore.DfsuFile import DfsuFile, DfsuFileType
@@ -178,7 +178,7 @@ def read(
         *,
         items: str | int | Sequence[str | int] | None = None,
         time: int | str | slice | None = None,
-        elements: Collection[int] | None = None,
+        elements: Sequence[int] | None = None,
         area: Tuple[float, float, float, float] | None = None,
         x: float | None = None,
         y: float | None = None,
@@ -246,7 +246,7 @@ def read(
                 elements = self.geometry.find_index(  # type: ignore
                     x=x, y=y, z=z, area=area, layers=layers
                 )
-                if len(elements) == 0:
+                if len(elements) == 0:  # type: ignore
                     raise ValueError("No elements in selection!")
 
         geometry = (

mikeio/spatial/_FM_geometry.py

@@ -3,14 +3,14 @@
 from functools import cached_property
 from pathlib import Path
 from typing import (
-    Collection,
     List,
     Any,
     Literal,
     Sequence,
     Sized,
     Tuple,
     TYPE_CHECKING,
+    overload,
 )
 
 
@@ -943,8 +943,18 @@ def _get_boundary_faces(self) -> np.ndarray:
         bnd_face_id = face_counts == 1
         return all_faces[uf_id[bnd_face_id]]
 
+    @overload
     def isel(
-        self, idx: Collection[int], keepdims: bool = False, **kwargs: Any
+        self, idx: int, keepdims: bool = False, **kwargs: Any
+    ) -> GeometryPoint2D: ...
+
+    @overload
+    def isel(
+        self, idx: Sequence[int], keepdims: bool = False, **kwargs: Any
+    ) -> "GeometryFM2D": ...
+
+    def isel(
+        self, idx: int | Sequence[int], keepdims=False, **kwargs: Any
     ) -> "GeometryFM2D" | GeometryPoint2D:
         """export a selection of elements to a new geometry
 
@@ -953,12 +963,14 @@ def isel(
 
         Parameters
         ----------
-        idx : collection(int)
-            collection of element indicies
+        idx : list(int)
+            element ids to be selected
         keepdims : bool, optional
             Should the original Geometry type be kept (keepdims=True)
             or should it be reduced e.g. to a GeometryPoint2D if possible
             (keepdims=False), by default False
+        **kwargs: optional
+            Not used
 
         Returns
         -------
@@ -1073,7 +1085,7 @@ def _elements_in_area(
             raise ValueError("'area' must be bbox [x0,y0,x1,y1] or polygon")
 
     def _nodes_to_geometry(
-        self, nodes: Collection[int]
+        self, nodes: Sequence[int] | int
     ) -> "GeometryFM2D" | GeometryPoint2D:
         """export a selection of nodes to new flexible file geometry
 
@@ -1100,7 +1112,9 @@ def _nodes_to_geometry(
                 elements.append(j)
 
         assert len(elements) > 0, "no elements found"
-        elements = np.sort(elements)  # make sure elements are sorted!
+        elements = sorted(
+            elements
+        )  # make sure elements are sorted! # TODO: should this be here?
 
         node_ids, elem_tbl = self._get_nodes_and_table_for_elements(elements)
         node_coords = self.node_coordinates[node_ids]
@@ -1118,7 +1132,7 @@ def _nodes_to_geometry(
         )
 
     def elements_to_geometry(
-        self, elements: int | Collection[int], keepdims: bool = False
+        self, elements: int | Sequence[int], keepdims: bool = False
     ) -> "GeometryFM2D" | GeometryPoint2D:
         if isinstance(elements, (int, np.integer)):
             sel_elements: List[int] = [elements]
@@ -1129,16 +1143,12 @@ def elements_to_geometry(
 
             return GeometryPoint2D(x=x, y=y, projection=self.projection)
 
-        sorted_elements = np.sort(
-            sel_elements
-        )  # make sure elements are sorted! # TODO is this necessary? If so, should be done in the initialiser
-
         # extract information for selected elements
 
-        node_ids, elem_tbl = self._get_nodes_and_table_for_elements(sorted_elements)
+        node_ids, elem_tbl = self._get_nodes_and_table_for_elements(sel_elements)
         node_coords = self.node_coordinates[node_ids]
         codes = self.codes[node_ids]
-        elem_ids = self.element_ids[sorted_elements]
+        elem_ids = self.element_ids[sel_elements]
 
         return GeometryFM2D(
             node_coordinates=node_coords,
@@ -1152,8 +1162,8 @@ def elements_to_geometry(
         )
 
     def _get_nodes_and_table_for_elements(
-        self, elements: np.ndarray | List[int]
-    ) -> tuple[Any, Any]:
+        self, elements: Sequence[int] | np.ndarray
+    ) -> tuple[np.ndarray, list[np.ndarray]]:
         """list of nodes and element table for a list of elements
 
         Parameters
@@ -1165,15 +1175,12 @@ def _get_nodes_and_table_for_elements(
         -------
         np.array(int)
             array of node ids (unique)
-        list(list(int))
+        list(np.array(int))
             element table with a list of nodes for each element
         """
-        elem_tbl = np.empty(len(elements), dtype=np.dtype("O"))
-
-        for j, eid in enumerate(elements):
-            elem_tbl[j] = np.asarray(self.element_table[eid])
+        elem_tbl = [np.asarray(self.element_table[eid]) for eid in elements]
 
-        nodes = np.unique(np.hstack(elem_tbl))  # type: ignore
+        nodes = np.unique(np.hstack(elem_tbl))
         return nodes, elem_tbl
 
     def get_node_centered_data(

mikeio/spatial/_FM_geometry_layered.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 from functools import cached_property
-from typing import Any, Collection, Iterable, Literal, Sequence, List, Tuple
+from typing import Any, Iterable, Literal, Sequence, List, Tuple, overload
 
 import numpy as np
 from mikecore.DfsuFile import DfsuFileType
@@ -32,7 +32,6 @@ def __init__(
         validate: bool = True,
         reindex: bool = False,
     ) -> None:
-
         super().__init__(
             node_coordinates=node_coordinates,
             element_table=element_table,
@@ -68,18 +67,24 @@ def __repr__(self) -> str:
     def geometry2d(self) -> GeometryFM2D:
         return self.to_2d_geometry()
 
+    @overload
     def isel(
-        self, idx: Collection[int], keepdims: bool = False, **kwargs: Any
-    ) -> GeometryFM3D | GeometryPoint3D | GeometryFM2D | GeometryFMVerticalColumn:
+        self, idx: int, keepdims: bool = False, **kwargs: Any
+    ) -> GeometryPoint3D: ...
 
-        return self.elements_to_geometry(elements=idx, keepdims=keepdims)
+    @overload
+    def isel(
+        self, idx: Sequence[int], keepdims: bool = False, **kwargs: Any
+    ) -> _GeometryFMLayered: ...
+
+    def isel(self, idx: int | Sequence[int], keepdims=False, **kwargs):
+        return self.elements_to_geometry(
+            elements=idx, node_layers=None, keepdims=keepdims
+        )
 
     def elements_to_geometry(
-        self,
-        elements: int | Collection[int],
-        node_layers: Layer = "all",
-        keepdims: bool = False,
-    ) -> GeometryFM3D | GeometryPoint3D | GeometryFM2D | GeometryFMVerticalColumn:
+        self, elements: int | Sequence[int], node_layers="all", keepdims=False
+    ):
         sel_elements: List[int]
 
         if isinstance(elements, (int, np.integer)):
@@ -91,21 +96,16 @@ def elements_to_geometry(
 
             return GeometryPoint3D(x=x, y=y, z=z, projection=self.projection)
 
-        sorted_elements = np.sort(
-            sel_elements
-        )  # make sure elements are sorted! # TODO is this necessary?
-
         # create new geometry
         new_type = self._type
 
-        layers_used = self.layer_ids[sorted_elements]
+        layers_used = self.layer_ids[sel_elements]
         unique_layer_ids = np.unique(layers_used)
         n_layers = len(unique_layer_ids)
 
         if n_layers > 1:
-            bottom: Layer = "bottom"
-            elem_bot = self.get_layer_elements(layers=bottom)
-            if np.all(np.in1d(sorted_elements, elem_bot)):
+            elem_bot = self.get_layer_elements("bottom")
+            if np.all(np.in1d(sel_elements, elem_bot)):
                 n_layers = 1
 
         if (
@@ -119,7 +119,7 @@ def elements_to_geometry(
 
         # extract information for selected elements
         if n_layers == 1:
-            elem2d = self.elem2d_ids[sorted_elements]
+            elem2d = self.elem2d_ids[sel_elements]
             geom2d = self.geometry2d
             node_ids, elem_tbl = geom2d._get_nodes_and_table_for_elements(elem2d)
             assert len(elem_tbl[0]) <= 4, "Not a 2D element"
@@ -128,11 +128,11 @@ def elements_to_geometry(
             elem_ids = self._element_ids[elem2d]
         else:
             node_ids, elem_tbl = self._get_nodes_and_table_for_elements(
-                sorted_elements, node_layers=node_layers
+                sel_elements, node_layers=node_layers
             )
             node_coords = self.node_coordinates[node_ids]
             codes = self.codes[node_ids]
-            elem_ids = self._element_ids[sorted_elements]
+            elem_ids = self._element_ids[sel_elements]
 
         if new_type == DfsuFileType.Dfsu2D:
             return GeometryFM2D(
@@ -163,7 +163,7 @@ def elements_to_geometry(
                 )
             else:
                 klass = self.__class__
-                return klass(  # type: ignore
+                return klass(
                     node_coordinates=node_coords,
                     codes=codes,
                     node_ids=node_ids,
@@ -183,7 +183,7 @@ def element_coordinates(self) -> np.ndarray:
 
     def _get_nodes_and_table_for_elements(
         self,
-        elements: Collection[int] | np.ndarray,
+        elements: Sequence[int] | np.ndarray,
         node_layers: Layer = "all",
     ) -> Tuple[Any, Any]:
         """list of nodes and element table for a list of elements
@@ -760,7 +760,6 @@ def get_nearest_relative_distance(self, coords: Tuple[float, float]) -> float:
         return self.relative_element_distance[idx]
 
     def find_index(self, x=None, y=None, z=None, coords=None, layers=None):
-
         if layers is not None:
             idx = self.get_layer_elements(layers)
         else:
@@ -893,7 +892,6 @@ def __call__(self, ax=None, figsize=None, **kwargs):
         return ax
 
     def mesh(self, title="Mesh", edge_color="0.5", **kwargs):
-
         v = np.full_like(self.g.element_coordinates[:, 0], np.nan)
         return _plot_vertical_profile(
             node_coordinates=self.g.node_coordinates,

tests/test_dataarray.py

@@ -696,6 +696,27 @@ def test_da_sel_area_dfsu2d():
     assert da1.geometry.n_elements == 14
 
 
+def test_da_isel_order_is_important_dfsu2d():
+    filename = "tests/testdata/FakeLake.dfsu"
+    da = mikeio.read(filename, items=0, time=0)[0]
+
+    # select elements sorted
+    da1 = da.isel(element=[0, 1])
+    assert da1.values[0] == pytest.approx(-3.2252840995788574)
+    assert da1.geometry.element_coordinates[0, 0] == pytest.approx(-0.61049269425)
+
+    # select elements in arbitrary order
+    da2 = da.isel(element=[1, 0])
+    assert da2.values[1] == pytest.approx(-3.2252840995788574)
+    assert da2.geometry.element_coordinates[1, 0] == pytest.approx(-0.61049269425)
+
+    # select same elements multiple times, not sure why, but consistent with NumPy, xarray
+    da3 = da.isel(element=[1, 0, 1])
+    assert da3.values[1] == pytest.approx(-3.2252840995788574)
+    assert da3.geometry.element_coordinates[1, 0] == pytest.approx(-0.61049269425)
+    assert len(da3.geometry.element_coordinates) == 3
+
+
 def test_da_sel_area_grid2d():
     filename = "tests/testdata/gebco_sound.dfs2"
     da = mikeio.read(filename, items=0)[0]