# 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.floating, np.integer, 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.floating, np.integer, 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.floating, np.integer, 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")