From ae465191e70c7d88dd3fcf21112094e56482634f Mon Sep 17 00:00:00 2001 From: Henrik Andersson Date: Sat, 16 Mar 2024 17:43:56 +0100 Subject: [PATCH] Delete duplicate _calc_element_coordinates --- mikeio/spatial/_FM_geometry.py | 65 +++++++++--------- mikeio/spatial/_FM_geometry_layered.py | 91 +++++++------------------- 2 files changed, 57 insertions(+), 99 deletions(-) diff --git a/mikeio/spatial/_FM_geometry.py b/mikeio/spatial/_FM_geometry.py index b8145f561..47f89ceda 100644 --- a/mikeio/spatial/_FM_geometry.py +++ b/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, diff --git a/mikeio/spatial/_FM_geometry_layered.py b/mikeio/spatial/_FM_geometry_layered.py index 4201c0215..36f9b7513 100644 --- a/mikeio/spatial/_FM_geometry_layered.py +++ b/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: