Delete duplicate _calc_element_coordinates

mikeio/spatial/_FM_geometry.py

@@ -261,6 +261,38 @@ def __init__(
         if reindex:
             self._reindex()
 
+    def _calc_element_coordinates(self, maxnodes: int = 4) -> np.ndarray:
+        element_table = self.element_table
+
+        n_elements = len(element_table)
+        ec = np.empty([n_elements, 3])
+
+        # pre-allocate for speed
+        idx = np.zeros(maxnodes, dtype=int)
+        xcoords = np.zeros([maxnodes, n_elements])
+        ycoords = np.zeros([maxnodes, n_elements])
+        zcoords = np.zeros([maxnodes, n_elements])
+        nnodes_per_elem = np.zeros(n_elements)
+
+        for j in range(n_elements):
+            nodes = element_table[j]
+            nnodes = len(nodes)
+            nnodes_per_elem[j] = nnodes
+            for i in range(nnodes):
+                idx[i] = nodes[i]  # - 1
+
+            x, y, z = self.node_coordinates[idx[:nnodes]].T
+
+            xcoords[:nnodes, j] = x
+            ycoords[:nnodes, j] = y
+            zcoords[:nnodes, j] = z
+
+        ec[:, 0] = np.sum(xcoords, axis=0) / nnodes_per_elem
+        ec[:, 1] = np.sum(ycoords, axis=0) / nnodes_per_elem
+        ec[:, 2] = np.sum(zcoords, axis=0) / nnodes_per_elem
+
+        return ec
+
     def _check_elements(
         self,
         element_table: np.ndarray | List[Sequence[int]] | List[np.ndarray],
@@ -463,39 +495,6 @@ def _tree2d(self) -> cKDTree:
         xy = self.element_coordinates[:, :2]
         return cKDTree(xy)
 
-    def _calc_element_coordinates(self) -> np.ndarray:
-        element_table = self.element_table
-
-        n_elements = len(element_table)
-        ec = np.empty([n_elements, 3])
-
-        # pre-allocate for speed
-        maxnodes = 4
-        idx = np.zeros(maxnodes, dtype=int)
-        xcoords = np.zeros([maxnodes, n_elements])
-        ycoords = np.zeros([maxnodes, n_elements])
-        zcoords = np.zeros([maxnodes, n_elements])
-        nnodes_per_elem = np.zeros(n_elements)
-
-        for j in range(n_elements):
-            nodes = element_table[j]
-            nnodes = len(nodes)
-            nnodes_per_elem[j] = nnodes
-            for i in range(nnodes):
-                idx[i] = nodes[i]  # - 1
-
-            x, y, z = self.node_coordinates[idx[:nnodes]].T
-
-            xcoords[:nnodes, j] = x
-            ycoords[:nnodes, j] = y
-            zcoords[:nnodes, j] = z
-
-        ec[:, 0] = np.sum(xcoords, axis=0) / nnodes_per_elem
-        ec[:, 1] = np.sum(ycoords, axis=0) / nnodes_per_elem
-        ec[:, 2] = np.sum(zcoords, axis=0) / nnodes_per_elem
-
-        return ec
-
     def find_nearest_elements(
         self,
         x: float | np.ndarray,

mikeio/spatial/_FM_geometry_layered.py

@@ -3,7 +3,7 @@
 from typing import Any, Collection, Iterable, Literal, Sequence, List, Tuple
 
 import numpy as np
-from mikecore.DfsuFile import DfsuFileType  # type: ignore
+from mikecore.DfsuFile import DfsuFileType
 
 
 from ._FM_geometry import GeometryFM2D, _GeometryFM, _GeometryFMPlotter
@@ -46,13 +46,13 @@ def __init__(
         )
 
         self._n_layers = n_layers
-        self._n_sigma: int = n_sigma if n_sigma is not None else n_layers
+        self._n_sigma = n_sigma if n_sigma is not None else n_layers
 
         # Lazy properties
-        self._bot_elems = None
-        self._e2_e3_table = None
-        self._2d_ids = None
-        self._layer_ids = None
+        self._bot_elems: np.ndarray | None = None
+        self._e2_e3_table: np.ndarray | None = None
+        self._2d_ids: np.ndarray | None = None
+        self._layer_ids: np.ndarray | None = None
 
     def __repr__(self) -> str:
         return (
@@ -65,7 +65,7 @@ def __repr__(self) -> str:
         )
 
     @cached_property
-    def geometry2d(self):
+    def geometry2d(self) -> GeometryFM2D:
         return self.to_2d_geometry()
 
     def isel(
@@ -177,43 +177,9 @@ def elements_to_geometry(
                 )
 
     @cached_property
-    def element_coordinates(self):
+    def element_coordinates(self) -> np.ndarray:
         """Center coordinates of each element"""
-        return self._calc_element_coordinates()
-
-    def _calc_element_coordinates(self):
-
-        node_coordinates = self.node_coordinates
-
-        element_table = self.element_table
-
-        n_elements = len(element_table)
-        ec = np.empty([n_elements, 3])
-
-        # pre-allocate for speed
-        maxnodes = 4 if self.is_2d else 8
-        idx = np.zeros(maxnodes, dtype=int)
-        xcoords = np.zeros([maxnodes, n_elements])
-        ycoords = np.zeros([maxnodes, n_elements])
-        zcoords = np.zeros([maxnodes, n_elements])
-        nnodes_per_elem = np.zeros(n_elements)
-
-        for j in range(n_elements):
-            nodes = element_table[j]
-            nnodes = len(nodes)
-            nnodes_per_elem[j] = nnodes
-            for i in range(nnodes):
-                idx[i] = nodes[i]  # - 1
-
-            xcoords[:nnodes, j] = node_coordinates[idx[:nnodes], 0]
-            ycoords[:nnodes, j] = node_coordinates[idx[:nnodes], 1]
-            zcoords[:nnodes, j] = node_coordinates[idx[:nnodes], 2]
-
-        ec[:, 0] = np.sum(xcoords, axis=0) / nnodes_per_elem
-        ec[:, 1] = np.sum(ycoords, axis=0) / nnodes_per_elem
-        ec[:, 2] = np.sum(zcoords, axis=0) / nnodes_per_elem
-
-        return ec
+        return self._calc_element_coordinates(maxnodes=8)
 
     def _get_nodes_and_table_for_elements(
         self,
@@ -244,8 +210,6 @@ def _get_nodes_and_table_for_elements(
 
         else:
             # 3D => 2D
-            if (node_layers != "bottom") and (node_layers != "top"):
-                raise ValueError("node_layers must be either all, bottom or top")
             for j, eid in enumerate(elements):
                 elem_nodes = np.asarray(self.element_table[eid])
                 nn = len(elem_nodes)
@@ -375,7 +339,7 @@ def _elements_in_area(self, area):
             return np.array([], dtype=int)
 
     @staticmethod
-    def _find_top_layer_elements(elementTable):
+    def _find_top_layer_elements(elementTable: np.ndarray) -> np.ndarray:
         """
         Find element indices (zero based) of the elements being the upper-most element
         in its column.
@@ -407,13 +371,6 @@ def _find_top_layer_elements(elementTable):
                 topLayerElments.append(i)
                 continue
 
-            if len(elmt1) % 2 != 0:
-                raise Exception(
-                    "In a layered mesh, each element must have an even number of elements (element index "
-                    + i
-                    + ")"
-                )
-
             # Number of nodes in a 2D element
             elmt2DSize = len(elmt1) // 2
 
@@ -439,7 +396,7 @@ def _find_top_layer_elements(elementTable):
         return np.array(topLayerElments, dtype=np.int32)
 
     @cached_property
-    def n_layers_per_column(self):
+    def n_layers_per_column(self) -> np.ndarray:
         """List of number of layers for each column"""
 
         top_elems = self.top_elements
@@ -451,7 +408,7 @@ def n_layers_per_column(self):
         return n_layers_column
 
     @cached_property
-    def bottom_elements(self):
+    def bottom_elements(self) -> np.ndarray:
         """List of 3d element ids of bottom layer"""
         return self.top_elements - self.n_layers_per_column + 1
 
@@ -502,7 +459,7 @@ def get_layer_elements(self, layers: int | Layer | Iterable[int]) -> np.ndarray:
         return self._element_ids[self.layer_ids == layers]
 
     @property
-    def e2_e3_table(self):
+    def e2_e3_table(self) -> np.ndarray:
         """The 2d-to-3d element connectivity table for a 3d object"""
 
         # e2_e3, 2d_ids and layer_ids are all set at the same time
@@ -515,7 +472,7 @@ def e2_e3_table(self):
         return self._e2_e3_table
 
     @property
-    def elem2d_ids(self):
+    def elem2d_ids(self) -> np.ndarray:
         """The associated 2d element id for each 3d element"""
         if self._2d_ids is None:
             res = self._get_2d_to_3d_association()
@@ -524,7 +481,7 @@ def elem2d_ids(self):
             self._layer_ids = res[2]
         return self._2d_ids
 
-    def _get_2d_to_3d_association(self):
+    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
@@ -692,7 +649,7 @@ def __init__(
     def boundary_polylines(self):
         return self.geometry2d.boundary_polylines
 
-    def contains(self, points) -> Sequence[bool]:
+    def contains(self, points) -> np.ndarray:
         return self.geometry2d.contains(points)
 
     def to_mesh(self, outfilename):
@@ -841,17 +798,17 @@ class GeometryFMVerticalColumn(GeometryFM3D):
 
     # TODO: add plotter
 
-    def calc_ze(self, zn=None):
+    def calc_ze(self, zn: np.ndarray | None = None) -> np.ndarray:
         if zn is None:
             zn = self.node_coordinates[:, 2]
         return self._calc_z_using_idx(zn, self._idx_e)
 
-    def calc_zf(self, zn=None):
+    def calc_zf(self, zn: np.ndarray | None = None) -> np.ndarray:
         if zn is None:
             zn = self.node_coordinates[:, 2]
         return self._calc_z_using_idx(zn, self._idx_f)
 
-    def _calc_zee(self, zn=None):
+    def _calc_zee(self, zn: np.ndarray | None = None) -> np.ndarray:
         ze = self.calc_ze(zn)
         zf = self.calc_zf(zn)
         if ze.ndim == 1:
@@ -862,7 +819,9 @@ def _calc_zee(self, zn=None):
                 (zf[:, 0].reshape((-1, 1)), ze, zf[:, -1].reshape((-1, 1)))
             )
 
-    def _interp_values(self, zn, data, z):
+    def _interp_values(
+        self, zn: np.ndarray, data: np.ndarray, z: np.ndarray
+    ) -> np.ndarray:
         """Interpolate to other z values, allow linear extrapolation"""
         from scipy.interpolate import interp1d  # type: ignore
 
@@ -878,7 +837,7 @@ def _interp_values(self, zn, data, z):
         return dati
 
     @property
-    def _idx_f(self):
+    def _idx_f(self) -> np.ndarray:
         nnodes_half = int(len(self.element_table[0]) / 2)
         n_vfaces = self.n_elements + 1
         idx_f = np.zeros((n_vfaces, nnodes_half), dtype=int)
@@ -888,7 +847,7 @@ def _idx_f(self):
         return idx_f
 
     @property
-    def _idx_e(self):
+    def _idx_e(self) -> np.ndarray:
         nnodes_per_elem = len(self.element_table[0])
         idx_e = np.zeros((self.n_elements, nnodes_per_elem), dtype=int)
 
@@ -899,7 +858,7 @@ def _idx_e(self):
 
         return idx_e
 
-    def _calc_z_using_idx(self, zn, idx):
+    def _calc_z_using_idx(self, zn: np.ndarray, idx: np.ndarray) -> np.ndarray:
         if zn.ndim == 1:
             zf = zn[idx].mean(axis=1)
         elif zn.ndim == 2: