Selecting a layer (Dfsu 3d) uses vertical coordinates for the selected layer

mikeio/spatial/_FM_geometry_layered.py

@@ -117,22 +117,12 @@ def elements_to_geometry(
         if n_layers == 1 and node_layers == "all":
             node_layers = "bottom"
 
-        # extract information for selected elements
-        if n_layers == 1:
-            elem2d = self.elem2d_ids[sorted_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"
-            node_coords = geom2d.node_coordinates[node_ids]
-            codes = geom2d.codes[node_ids]
-            elem_ids = self._element_ids[elem2d]
-        else:
-            node_ids, elem_tbl = self._get_nodes_and_table_for_elements(
-                sorted_elements, node_layers=node_layers
-            )
-            node_coords = self.node_coordinates[node_ids]
-            codes = self.codes[node_ids]
-            elem_ids = self._element_ids[sorted_elements]
+        node_ids, elem_tbl = self._get_nodes_and_table_for_elements(
+            sorted_elements, node_layers=node_layers
+        )
+        node_coords = self.node_coordinates[node_ids]
+        codes = self.codes[node_ids]
+        elem_ids = self._element_ids[sorted_elements]
 
         if new_type == DfsuFileType.Dfsu2D:
             return GeometryFM2D(
@@ -256,6 +246,7 @@ def to_2d_geometry(self) -> GeometryFM2D:
             reindex=True,
         )
 
+        # TODO move this to before creating the geometry
         # Fix z-coordinate for sigma-z:
         if self._type == DfsuFileType.Dfsu3DSigmaZ:
             zn = geom.node_coordinates[:, 2].copy()

tests/test_geometry_fm.py

@@ -6,20 +6,20 @@
 
 @pytest.fixture
 def simple_3d_geom():
-    #     x     y    z
-    nc = [
-        (0.0, 0.0, 0.0),
-        (1.0, 0.0, 0.0),
-        (1.0, 1.0, 0.0),
-        (0.0, 0.0, -1.0),
-        (1.0, 0.0, -1.0),
-        (1.0, 1.0, -1.0),
-        (0.0, 0.0, -2.0),
-        (1.0, 0.0, -2.0),
-        (1.0, 1.0, -2.0),
-    ]
 
-    el = [(0, 1, 2, 3, 4, 5), (3, 4, 5, 6, 7, 8)]
+    nc_2d = [(0.0, 0.0), (1.0, 0.0), (1.0, 1.0)]
+    # bottom first approach
+    z = [-2.0, -1.0, 0.0]
+
+    nc = []
+
+    for col in nc_2d:
+        for depth in z:
+            nc.append((col[0], col[1], depth))
+
+    # TODO constructing a 3d mesh following the convention needs a helper function
+    # Arbitrary 3d element tables doesn't work
+    el = [(0, 3, 6, 1, 4, 7), (1, 4, 7, 2, 5, 8)]
 
     g = GeometryFM3D(
         node_coordinates=nc,
@@ -183,7 +183,7 @@ def test_plot_mesh():
     g.plot.mesh()
 
 
-def test_layered(simple_3d_geom: GeometryFM3D):
+def test_layered(simple_3d_geom: GeometryFM3D) -> None:
 
     g = simple_3d_geom
 
@@ -203,3 +203,21 @@ def test_layered(simple_3d_geom: GeometryFM3D):
 
     assert "elements: 2" in repr(g2)
     assert "layers: 2" in repr(g2)
+
+
+def test_select_single_layer_preserved_vertical_coordinates(
+    simple_3d_geom: GeometryFM3D,
+) -> None:
+    g = simple_3d_geom
+
+    bot_el = g.bottom_elements
+    gb = g.elements_to_geometry(bot_el, keepdims=True)
+    assert isinstance(gb, GeometryFM2D)
+    assert gb.node_coordinates[0][2] == -2.0
+
+    top_el = g.top_elements
+
+    gt = g.elements_to_geometry(top_el, keepdims=True)
+    assert isinstance(gt, GeometryFM2D)
+    assert gt.n_elements == 1
+    assert gt.node_coordinates[0][2] == -1.0