Source code for geoh5py.objects.octree

#  Copyright (c) 2024 Mira Geoscience Ltd.
#
#  This file is part of geoh5py.
#
#  geoh5py is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  geoh5py is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public License
#  along with geoh5py.  If not, see <https://www.gnu.org/licenses/>.

from __future__ import annotations

import uuid
from typing import TYPE_CHECKING

import numpy as np

from .grid_object import GridObject

if TYPE_CHECKING:
    from geoh5py.objects import ObjectType


[docs] class Octree(GridObject): """ Octree mesh class that uses a tree structure such that cells can be subdivided it into eight octants. """ __TYPE_UID = uuid.UUID( fields=(0x4EA87376, 0x3ECE, 0x438B, 0xBF, 0x12, 0x3479733DED46) ) _attribute_map: dict = GridObject._attribute_map.copy() _attribute_map.update( { "NU": "u_count", "NV": "v_count", "NW": "w_count", "Origin": "origin", "Rotation": "rotation", "U Cell Size": "u_cell_size", "V Cell Size": "v_cell_size", "W Cell Size": "w_cell_size", } ) def __init__(self, object_type: ObjectType, **kwargs): self._origin: np.ndarray = np.zeros(3) self._rotation: float = 0.0 self._u_count: int = 0 self._v_count: int = 0 self._w_count: int = 0 self._u_cell_size: float | None = None self._v_cell_size: float | None = None self._w_cell_size: float | None = None self._octree_cells: np.ndarray | None = None super().__init__(object_type, **kwargs)
[docs] def base_refine(self): """ Refine the mesh to its base octree level resulting in a single cell along the shortest dimension. """ assert ( self._octree_cells is None ), "'base_refine' function only implemented if 'octree_cells' is None " assert self.u_count is not None assert self.v_count is not None assert self.w_count is not None # Number of octree levels allowed on each dimension level_u = np.log2(self.u_count) level_v = np.log2(self.v_count) level_w = np.log2(self.w_count) min_level = np.min([level_u, level_v, level_w]) # Check that the refine level doesn't exceed the shortest dimension level = np.min([0, min_level]) # Number of additional break to account for variable dimensions add_u = int(level_u - min_level) add_v = int(level_v - min_level) add_w = int(level_w - min_level) j, k, i = np.meshgrid( np.arange(0, self.v_count, 2 ** (level_v - add_v - level)), np.arange(0, self.w_count, 2 ** (level_w - add_w - level)), np.arange(0, self.u_count, 2 ** (level_u - add_u - level)), ) octree_cells = np.c_[ i.flatten(), j.flatten(), k.flatten(), np.ones_like(i.flatten()) * 2 ** (min_level - level), ] self._octree_cells = np.rec.fromarrays( octree_cells.T, names=["I", "J", "K", "NCells"], formats=["<i4", "<i4", "<i4", "<i4"], )
@property def centroids(self): """ :obj:`numpy.array` of :obj:`float`, shape (:obj:`~geoh5py.objects.octree.Octree.n_cells`, 3): Cell center locations in world coordinates. .. code-block:: python centroids = [ [x_1, y_1, z_1], ..., [x_N, y_N, z_N] ] """ if getattr(self, "_centroids", None) is None: assert self.octree_cells is not None, "octree_cells must be set" assert self.u_cell_size is not None, "u_cell_size must be set" assert self.v_cell_size is not None, "v_cell_size must be set" assert self.w_cell_size is not None, "w_cell_size must be set" angle = np.deg2rad(self.rotation) rot = np.r_[ np.c_[np.cos(angle), -np.sin(angle), 0], np.c_[np.sin(angle), np.cos(angle), 0], np.c_[0, 0, 1], ] u_grid = ( self.octree_cells["I"] + self.octree_cells["NCells"] / 2.0 ) * self.u_cell_size v_grid = ( self.octree_cells["J"] + self.octree_cells["NCells"] / 2.0 ) * self.v_cell_size w_grid = ( self.octree_cells["K"] + self.octree_cells["NCells"] / 2.0 ) * self.w_cell_size xyz = np.c_[u_grid, v_grid, w_grid] self._centroids = np.dot(rot, xyz.T).T assert self._centroids is not None for ind, axis in enumerate(["x", "y", "z"]): self._centroids[:, ind] += self.origin[axis] return self._centroids
[docs] @classmethod def default_type_uid(cls) -> uuid.UUID: return cls.__TYPE_UID
@property def n_cells(self) -> int | None: """ :obj:`int`: Total number of cells in the mesh """ if self.octree_cells is not None: return self.octree_cells.shape[0] return None @property def octree_cells(self) -> np.ndarray | None: """ :obj:`numpy.ndarray` of :obj:`int`, shape (:obj:`~geoh5py.objects.octree.Octree.n_cells`, 4): Array defining the i, j, k position and size of each cell. The size defines the width of a cell in number of base cells. .. code-block:: python cells = [ [i_1, j_1, k_1, size_1], ..., [i_N, j_N, k_N, size_N] ] """ if getattr(self, "_octree_cells", None) is None: if self.on_file: octree_cells = self.workspace.fetch_array_attribute( self, "octree_cells" ) self._octree_cells = octree_cells else: self.base_refine() return self._octree_cells @octree_cells.setter def octree_cells(self, value): if value is not None: dtypes = [("I", "<i4"), ("J", "<i4"), ("K", "<i4"), ("NCells", "<i4")] if len(value.dtype) > 1: dtype = np.dtype(dtypes) assert ( value.dtype == dtype ), f"Input of type {np.ndarray} must be of {dtype}" self._octree_cells = value else: value = np.vstack(value) assert ( value.shape[1] == 4 ), "'octree_cells' requires an ndarray of shape (*, 4)" self._centroids = None self._octree_cells = np.asarray( np.core.records.fromarrays( value.T, names="I, J, K, NCells", formats="<i4, <i4, <i4, <i4" ) ) self.workspace.update_attribute(self, "octree_cells") @property def origin(self): """ :obj:`numpy.array` of :obj:`float`, shape (3, ): Coordinates of the origin """ return self._origin @origin.setter def origin(self, value): if value is not None: if isinstance(value, np.ndarray): value = value.tolist() assert len(value) == 3, "Origin must be a list or numpy array of shape (3,)" self.workspace.update_attribute(self, "attributes") self._centroids = None value = np.asarray( tuple(value), dtype=[("x", float), ("y", float), ("z", float)] ) self._origin = value @property def rotation(self) -> float: """ :obj:`float`: Clockwise rotation angle (degree) about the vertical axis. """ return self._rotation @rotation.setter def rotation(self, value): if value is not None: value = np.r_[value] assert len(value) == 1, "Rotation angle must be a float of shape (1,)" self._centroids = None self._rotation = value.astype(float).item() self.workspace.update_attribute(self, "attributes") @property def shape(self) -> tuple | None: """ :obj:`list` of :obj:`int`, len (3, ): Number of cells along the u, v and w-axis. """ if ( self.u_count is not None and self.v_count is not None and self.w_count is not None ): return self.u_count, self.v_count, self.w_count return None @property def u_cell_size(self) -> float | None: """ :obj:`float`: Base cell size along the u-axis. """ return self._u_cell_size @u_cell_size.setter def u_cell_size(self, value: float | np.ndarray): if not isinstance(value, (float, np.ndarray)): raise TypeError("Attribute 'u_cell_size' must be type(float).") self._centroids = None if isinstance(value, np.ndarray): assert len(value) == 1, "u_cell_size must be a float of shape (1,)" self._u_cell_size = np.r_[value].astype(float).item() else: self._u_cell_size = value self.workspace.update_attribute(self, "attributes") @property def u_count(self) -> int | None: """ :obj:`int`: Number of cells along u-axis. """ return self._u_count @u_count.setter def u_count(self, value: int): value = np.int32(value).item() if not isinstance(value, (float, np.int32, int)) or np.log2(value) % 1.0 != 0: raise TypeError("Attribute 'u_count' must be type(int) in power of 2.") self._centroids = None self._u_count = np.int32(value).item() self.workspace.update_attribute(self, "attributes") @property def v_cell_size(self) -> float | None: """ :obj:`float`: Base cell size along the v-axis. """ return self._v_cell_size @v_cell_size.setter def v_cell_size(self, value: float | np.ndarray): if not isinstance(value, (float, np.ndarray)): raise TypeError("Attribute 'v_cell_size' must be type(float).") self._centroids = None if isinstance(value, np.ndarray): assert len(value) == 1, "v_cell_size must be a float of shape (1,)" self._v_cell_size = np.r_[value].astype(float).item() else: self._v_cell_size = value self.workspace.update_attribute(self, "attributes") @property def v_count(self) -> int | None: """ :obj:`int`: Number of cells along v-axis. """ return self._v_count @v_count.setter def v_count(self, value: int): value = np.int32(value).item() if not isinstance(value, (float, np.int32, int)) or np.log2(value) % 1.0 != 0: raise TypeError("Attribute 'v_count' must be type(int) in power of 2.") self._centroids = None self._v_count = np.int32(value).item() self.workspace.update_attribute(self, "attributes") @property def w_cell_size(self) -> float | None: """ :obj:`float`: Base cell size along the w-axis. """ return self._w_cell_size @w_cell_size.setter def w_cell_size(self, value: float | np.ndarray): if not isinstance(value, (float, np.ndarray)): raise TypeError("Attribute 'w_cell_size' must be type(float).") self._centroids = None if isinstance(value, np.ndarray): assert len(value) == 1, "w_cell_size must be a float of shape (1,)" self._w_cell_size = np.r_[value].astype(float).item() else: self._w_cell_size = value self.workspace.update_attribute(self, "attributes") @property def w_count(self) -> int | None: """ :obj:`int`: Number of cells along w-axis. """ return self._w_count @w_count.setter def w_count(self, value: int): value = np.int32(value).item() if not isinstance(value, (float, np.int32, int)) or np.log2(value) % 1.0 != 0: raise TypeError("Attribute 'w_count' must be type(int) in power of 2.") self._centroids = None self._w_count = np.int32(value).item() self.workspace.update_attribute(self, "attributes")