Merge pull request #690 from DHI/fix/dfsu_read_elements

Dfsu read elements preserve order

mikeio/dfsu/_dfsu.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Collection, Literal, Sequence, Tuple
+
+from typing import Any, Literal, Sequence, Tuple
 
 import numpy as np
 import pandas as pd
@@ -358,7 +359,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] | np.ndarray | None = None,
         area: Tuple[float, float, float, float] | None = None,
         x: float | None = None,
         y: float | None = None,

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
 
 from matplotlib.axes import Axes
 import numpy as np
@@ -184,7 +184,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,
@@ -252,7 +252,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,7 +3,6 @@
 from functools import cached_property
 from pathlib import Path
 from typing import (
-    Collection,
     List,
     Any,
     Literal,
@@ -944,7 +943,7 @@ def _get_boundary_faces(self) -> np.ndarray:
         return all_faces[uf_id[bnd_face_id]]
 
     def isel(
-        self, idx: Collection[int], keepdims: bool = False, **kwargs: Any
+        self, idx: Sequence[int], keepdims: bool = False, **kwargs: Any
     ) -> "GeometryFM2D" | GeometryPoint2D:
         """export a selection of elements to a new geometry
 
@@ -953,7 +952,7 @@ def isel(
 
         Parameters
         ----------
-        idx : collection(int)
+        idx : list(int)
             collection of element indicies
         keepdims : bool, optional
             Should the original Geometry type be kept (keepdims=True)
@@ -970,10 +969,7 @@ def isel(
         find_index : find element indicies for points or an area
         """
 
-        if self._type == DfsuFileType.DfsuSpectral1D:
-            return self._nodes_to_geometry(nodes=idx)
-        else:
-            return self.elements_to_geometry(elements=idx, keepdims=keepdims)
+        return self.elements_to_geometry(elements=idx, keepdims=keepdims)
 
     def find_index(
         self,
@@ -1072,53 +1068,8 @@ def _elements_in_area(
         else:
             raise ValueError("'area' must be bbox [x0,y0,x1,y1] or polygon")
 
-    def _nodes_to_geometry(
-        self, nodes: Collection[int]
-    ) -> "GeometryFM2D" | GeometryPoint2D:
-        """export a selection of nodes to new flexible file geometry
-
-        Note: takes only the elements for which all nodes are selected
-
-        Parameters
-        ----------
-        nodes : list(int)
-            list of node ids
-
-        Returns
-        -------
-        UnstructuredGeometry
-            which can be used for further extraction or saved to file
-        """
-        nodes = np.atleast_1d(nodes)  # type: ignore
-        if len(nodes) == 1:
-            xy = self.node_coordinates[nodes[0], :2]
-            return GeometryPoint2D(xy[0], xy[1])
-
-        elements = []
-        for j, el_nodes in enumerate(self.element_table):
-            if np.all(np.isin(el_nodes, nodes)):
-                elements.append(j)
-
-        assert len(elements) > 0, "no elements found"
-        elements = np.sort(elements)  # make sure elements are sorted!
-
-        node_ids, elem_tbl = self._get_nodes_and_table_for_elements(elements)
-        node_coords = self.node_coordinates[node_ids]
-        codes = self.codes[node_ids]
-
-        return GeometryFM2D(
-            node_coordinates=node_coords,
-            codes=codes,
-            node_ids=node_ids,
-            dfsu_type=self._type,
-            projection=self.projection_string,
-            element_table=elem_tbl,
-            element_ids=self.element_ids[elements],
-            reindex=True,
-        )
-
     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 +1080,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,

mikeio/spatial/_FM_geometry_layered.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 from functools import cached_property
 from pathlib import Path
-from typing import Any, Collection, Iterable, Literal, Sequence, List, Tuple
+
+from typing import Any, Iterable, Literal, Sequence, List, Tuple
 
 from matplotlib.axes import Axes
 import numpy as np
@@ -71,14 +72,14 @@ def geometry2d(self) -> GeometryFM2D:
         return self.to_2d_geometry()
 
     def isel(
-        self, idx: Collection[int], keepdims: bool = False, **kwargs: Any
+        self, idx: Sequence[int], keepdims: bool = False, **kwargs: Any
     ) -> GeometryFM3D | GeometryPoint3D | GeometryFM2D | GeometryFMVerticalColumn:
 
         return self.elements_to_geometry(elements=idx, keepdims=keepdims)
 
     def elements_to_geometry(
         self,
-        elements: int | Collection[int],
+        elements: int | Sequence[int] | np.ndarray,
         node_layers: Layer = "all",
         keepdims: bool = False,
     ) -> GeometryFM3D | GeometryPoint3D | GeometryFM2D | GeometryFMVerticalColumn:
@@ -93,21 +94,17 @@ 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)):
+            if np.all(np.in1d(sel_elements, elem_bot)):
                 n_layers = 1
 
         if (
@@ -121,7 +118,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"
@@ -130,11 +127,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(
@@ -185,7 +182,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

mikeio/spatial/_FM_geometry_spectral.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Collection, Any, Sequence, Tuple
+from typing import Any, Sequence, Tuple
 
 
 import numpy as np
@@ -127,12 +127,12 @@ def directions(self) -> np.ndarray | None:
 # TODO reconsider inheritance to avoid overriding method signature
 class GeometryFMAreaSpectrum(_GeometryFMSpectrum):
     def isel(  # type: ignore
-        self, idx: Collection[int], **kwargs: Any
+        self, idx: Sequence[int], **kwargs: Any
     ) -> "GeometryFMPointSpectrum" | "GeometryFMAreaSpectrum":
         return self.elements_to_geometry(elements=idx)
 
     def elements_to_geometry(  # type: ignore
-        self, elements: Collection[int], keepdims: bool = False
+        self, elements: Sequence[int], keepdims: bool = False
     ) -> "GeometryFMPointSpectrum" | "GeometryFMAreaSpectrum":
         """export a selection of elements to new flexible file geometry
         Parameters
@@ -156,7 +156,6 @@ def elements_to_geometry(  # type: ignore
                 y=coords[1],
             )
 
-        elements = np.sort(elements)  # make sure elements are sorted!
         node_ids, elem_tbl = self._get_nodes_and_table_for_elements(elements)
         node_coords = self.node_coordinates[node_ids]
         codes = self.codes[node_ids]
@@ -213,7 +212,6 @@ def _nodes_to_geometry(  # type: ignore
                 elements.append(j)
 
         assert len(elements) > 0, "no elements found"
-        elements = np.sort(elements)  # make sure elements are sorted!
 
         node_ids, elem_tbl = self._get_nodes_and_table_for_elements(elements)
         node_coords = self.node_coordinates[node_ids]

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]

tests/test_dfsu.py

@@ -249,6 +249,16 @@ def test_read_area_polygon():
     assert subdomain.within(domain)
 
 
+def test_read_elements():
+    ds = mikeio.read(filename="tests/testdata/wind_north_sea.dfsu", elements=[0, 10])
+    assert ds.geometry.element_coordinates[0][0] == pytest.approx(1.4931853081272184)
+    assert ds.Wind_speed.to_numpy()[0, 0] == pytest.approx(9.530759811401367)
+
+    ds2 = mikeio.read(filename="tests/testdata/wind_north_sea.dfsu", elements=[10, 0])
+    assert ds2.geometry.element_coordinates[1][0] == pytest.approx(1.4931853081272184)
+    assert ds2.Wind_speed.to_numpy()[0, 1] == pytest.approx(9.530759811401367)
+
+
 def test_find_index_on_island():
     filename = "tests/testdata/FakeLake.dfsu"
     dfs = mikeio.open(filename)

tests/test_dfsu_layered.py

@@ -171,6 +171,25 @@ def test_read_dfsu3d_column():
     assert dscol2._zn.shape == (ds.n_timesteps, 5 * 3)
 
 
+def test_flip_column_upside_down():
+    filename = "tests/testdata/oresund_sigma_z.dfsu"
+    dfs = mikeio.open(filename)
+
+    (x, y) = (333934.1, 6158101.5)
+
+    ds = dfs.read()  # all data in file
+    dscol = ds.sel(x=x, y=y)
+    assert dscol.geometry.element_coordinates[0, 2] == pytest.approx(-7.0)
+    assert dscol.isel(time=-1).Temperature.values[0] == pytest.approx(17.460058)
+
+    idx = list(reversed(range(dscol.n_elements)))
+
+    dscol_ud = dscol.isel(element=idx)
+
+    assert dscol_ud.geometry.element_coordinates[-1, 2] == pytest.approx(-7.0)
+    assert dscol_ud.isel(time=-1).Temperature.values[-1] == pytest.approx(17.460058)
+
+
 def test_read_dfsu3d_column_save(tmp_path):
     filename = "tests/testdata/oresund_sigma_z.dfsu"
     dfs = mikeio.open(filename)
@@ -596,3 +615,13 @@ def test_read_wildcard_items():
     ds = dfs.read(items="Sal*")
     assert ds.items[0].name == "Salinity"
     assert ds.n_items == 1
+
+
+def test_read_elements_3d():
+    ds = mikeio.read("tests/testdata/oresund_sigma_z.dfsu", elements=[0, 10])
+    assert ds.geometry.element_coordinates[0][0] == pytest.approx(354020.46382194717)
+    assert ds.Salinity.to_numpy()[0, 0] == pytest.approx(23.18906021118164)
+
+    ds2 = mikeio.read("tests/testdata/oresund_sigma_z.dfsu", elements=[10, 0])
+    assert ds2.geometry.element_coordinates[1][0] == pytest.approx(354020.46382194717)
+    assert ds2.Salinity.to_numpy()[0, 1] == pytest.approx(23.18906021118164)

tests/test_dfsu_spectral.py

@@ -161,6 +161,16 @@ def test_read_area_spectrum_elements(dfsu_area):
     ds2 = dfs.read(elements=elems)
     assert ds2.shape[1] == len(elems)
     assert np.all(ds1[0].to_numpy()[:, elems, ...] == ds2[0].to_numpy())
+    assert ds2.geometry.element_coordinates[0, 0] == pytest.approx(2.651450863095597)
+    assert ds2["Energy density"].isel(time=-1).isel(frequency=0).isel(
+        direction=0
+    ).to_numpy()[0] == pytest.approx(1.770e-12)
+
+    ds3 = dfs.read(elements=[4, 3])
+    assert ds3.geometry.element_coordinates[1, 0] == pytest.approx(2.651450863095597)
+    assert ds3["Energy density"].isel(time=-1).isel(frequency=0).isel(
+        direction=0
+    ).to_numpy()[1] == pytest.approx(1.770e-12)
 
 
 def test_read_area_spectrum_xy(dfsu_area):