Fix isel on Dfsu 3d single timestep. (#828)

* Failing test

* Fix error (+ minor refactoring)

* Ignore/assert types

* Not needed

* Remove useless comment

Author: ecomodeller

mikeio/dataset/_dataarray.py

@@ -804,12 +804,15 @@ def isel(
                 spatial_axis = axis - 1 if self.dims[0] == "time" else axis
                 geometry = self.geometry.isel(idx, axis=spatial_axis)
 
-            # TODO this is ugly
             if isinstance(geometry, _GeometryFMLayered):
                 node_ids, _ = self.geometry._get_nodes_and_table_for_elements(
                     idx, node_layers="all"
                 )
-                zn = self._zn[:, node_ids]  # type: ignore
+                assert isinstance(self._zn, np.ndarray)
+                if self._zn.ndim == 2:
+                    zn = self._zn[:, node_ids]
+                else:
+                    zn = self._zn[node_ids]
 
         # reduce dims only if singleton idx
         dims = (

mikeio/spatial/_FM_geometry_layered.py

@@ -206,19 +206,18 @@ def _get_nodes_and_table_for_elements(
 
         """
         elem_tbl = np.empty(len(elements), dtype=np.dtype("O"))
-        if (node_layers == "all") or self.is_2d:
+        if node_layers == "all":
             for j, eid in enumerate(elements):
-                elem_tbl[j] = np.asarray(self.element_table[eid])
+                elem_tbl[j] = self.element_table[eid]
 
         else:
             # 3D => 2D
             for j, eid in enumerate(elements):
-                elem_nodes = np.asarray(self.element_table[eid])
-                nn = len(elem_nodes)
-                halfn = int(nn / 2)
+                elem_nodes = self.element_table[eid]
+                halfn = len(elem_nodes) // 2
                 if node_layers == "bottom":
                     elem_nodes = elem_nodes[:halfn]
-                if node_layers == "top":
+                elif node_layers == "top":
                     elem_nodes = elem_nodes[halfn:]
                 elem_tbl[j] = elem_nodes
 

tests/test_dfsu3d.py

@@ -662,3 +662,14 @@ def test_write_3d_non_equidistant(tmp_path: Path) -> None:
 
     # but getting the end time is not that expensive
     assert dfs.end_time == pd.Timestamp("2000-01-10")
+
+
+def test_isel_3d_single_time() -> None:
+    ds = mikeio.Dfsu3D("tests/testdata/basin_3d.dfsu").read()
+    ds1 = ds.isel(element=[0, 1])
+    assert ds1.geometry.n_elements == 2
+
+    ds2 = ds.isel(time=-1)
+    assert "time" not in ds2.dims
+    ds3 = ds2.isel(element=[0, 1])
+    assert ds3.geometry.n_elements == 2