Merge pull request #767 from DHI/dfsu_types

Dfsu 3d types

mikeio/dataset/_dataset.py

@@ -1307,7 +1307,7 @@ def _concat_time(
                 newdata[j][idx1] = ds[j].to_numpy()
 
         zn = None
-        if self._zn is not None:
+        if self._zn is not None and other._zn is not None:
             zshape = (len(newtime), self._zn.shape[start_dim])
             zn = np.zeros(shape=zshape, dtype=self._zn.dtype)
             if keep == "last":

mikeio/dfsu/_layered.py

@@ -203,7 +203,7 @@ def read(
         *,
         items: str | int | Sequence[str | int] | None = None,
         time: int | str | slice | None = None,
-        elements: Sequence[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/spatial/_FM_geometry_layered.py

@@ -71,7 +71,7 @@ def geometry2d(self) -> GeometryFM2D:
         return self.to_2d_geometry()
 
     def isel(
-        self, idx: Sequence[int], keepdims: bool = False, **kwargs: Any
+        self, idx: Sequence[int] | np.ndarray, keepdims: bool = False, **kwargs: Any
     ) -> GeometryFM3D | GeometryPoint3D | GeometryFM2D | GeometryFMVerticalColumn:
         return self.elements_to_geometry(elements=idx, keepdims=keepdims)
 
@@ -480,9 +480,7 @@ def elem2d_ids(self) -> np.ndarray:
         return self._2d_ids
 
     def _get_2d_to_3d_association(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        e2_to_e3 = (
-            []
-        )  # for each 2d element: the corresponding 3d element ids from bot to top
+        e2_to_e3 = []  # for each 2d element: the corresponding 3d element ids from bot to top
         index2d = []  # for each 3d element: the associated 2d element id
         layerid = []  # for each 3d element: the associated layer number
         n2d = len(self.top_elements)

tests/test_dfsu_layered.py

@@ -329,7 +329,6 @@ def test_number_of_nodes_and_elements_sigma_z():
 
 
 def test_read_and_select_single_element_dfsu_3d():
-
     filename = "tests/testdata/basin_3d.dfsu"
     dfs = mikeio.open(filename)
 
@@ -341,7 +340,6 @@ def test_read_and_select_single_element_dfsu_3d():
 
 
 def test_n_layers():
-
     filename = "tests/testdata/basin_3d.dfsu"
     dfs = mikeio.open(filename)
     assert dfs.n_layers == 10
@@ -360,7 +358,6 @@ def test_n_layers():
 
 
 def test_n_sigma_layers():
-
     filename = "tests/testdata/basin_3d.dfsu"
     dfs = mikeio.open(filename)
     assert dfs.n_sigma_layers == 10
@@ -379,7 +376,6 @@ def test_n_sigma_layers():
 
 
 def test_n_z_layers():
-
     filename = "tests/testdata/basin_3d.dfsu"
     dfs = mikeio.open(filename)
     assert dfs.n_z_layers == 0
@@ -398,7 +394,6 @@ def test_n_z_layers():
 
 
 def test_boundary_codes():
-
     filename = "tests/testdata/basin_3d.dfsu"
     dfs = mikeio.open(filename)
     assert len(dfs.geometry.boundary_codes) == 1
@@ -434,21 +429,23 @@ def test_top_elements():
     assert not hasattr(dfs, "top_elements")
 
 
-def test_top_elements_subset():
+def test_top_elements_subset() -> None:
     filename = "tests/testdata/oresund_sigma_z.dfsu"
-    g3d = mikeio.open(filename).geometry
+    g3d: GeometryFM3D = mikeio.open(filename).geometry  # type: ignore
     g2d = g3d.geometry2d
 
-    area = [356000, 6144000, 359000, 6146000]
+    area = (356000, 6144000, 359000, 6146000)
     idx2d = g2d.find_index(area=area)
     assert len(idx2d) == 6
     assert 3408 in idx2d
 
     idx3d = g3d.find_index(area=area)
-    subg = g3d.isel(idx3d)
+    subg: GeometryFM3D = g3d.isel(idx3d)  # type: ignore
 
     assert len(subg.top_elements) == 6
 
+    ds = mikeio.Dfsu3D(filename).read(elements=idx3d)
+
 
 def test_bottom_elements():
     filename = "tests/testdata/basin_3d.dfsu"
@@ -515,7 +512,6 @@ def test_n_layers_per_column():
 
 
 def test_write_from_dfsu3D(tmp_path):
-
     sourcefilename = "tests/testdata/basin_3d.dfsu"
     fp = tmp_path / "basin_3d.dfsu"
     dfs = mikeio.open(sourcefilename)
@@ -543,7 +539,6 @@ def test_extract_top_layer_to_2d(tmp_path):
 
 
 def test_modify_values_in_layer(tmp_path):
-
     ds = mikeio.read("tests/testdata/oresund_sigma_z.dfsu")
     selected_layer = 6  # Zero-based indexing!
     layer_elem_ids = ds.geometry.get_layer_elements(selected_layer)
@@ -559,7 +554,6 @@ def test_modify_values_in_layer(tmp_path):
 
 
 def test_to_mesh_3d(tmp_path):
-
     filename = "tests/testdata/oresund_sigma_z.dfsu"
 
     dfs = mikeio.open(filename)
@@ -576,7 +570,6 @@ def test_to_mesh_3d(tmp_path):
 
 
 def test_extract_surface_elevation_from_3d():
-
     dfs = mikeio.open("tests/testdata/oresund_sigma_z.dfsu")
     n_top1 = len(dfs.geometry.top_elements)
 
@@ -586,7 +579,6 @@ def test_extract_surface_elevation_from_3d():
 
 
 def test_dataset_write_dfsu3d(tmp_path):
-
     fp = tmp_path / "oresund_sigma_z.dfsu"
     ds = mikeio.read("tests/testdata/oresund_sigma_z.dfsu", time=[0, 1])
     ds.to_dfs(fp)
@@ -596,7 +588,6 @@ def test_dataset_write_dfsu3d(tmp_path):
 
 
 def test_dataset_write_dfsu3d_max(tmp_path):
-
     fp = tmp_path / "oresund_sigma_z.dfsu"
     ds = mikeio.read("tests/testdata/oresund_sigma_z.dfsu")
     assert ds._zn is not None
@@ -636,6 +627,7 @@ def test_append_dfsu_3d(tmp_path):
     assert ds3.n_timesteps == 2
     assert ds3.time[-1] == ds2.time[-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)
@@ -671,4 +663,3 @@ def test_write_3d_non_equidistant(tmp_path):
 
     # but getting the end time is not that expensive
     assert dfs.end_time == pd.Timestamp("2000-01-10")
-