diff --git a/geoapps_utils/modelling/plates.py b/geoapps_utils/modelling/plates.py index bff06e0..9bba726 100644 --- a/geoapps_utils/modelling/plates.py +++ b/geoapps_utils/modelling/plates.py @@ -9,10 +9,14 @@ # ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' import numpy as np -from pydantic import BaseModel, ConfigDict +from geoh5py import Workspace +from geoh5py.objects import Octree, Surface +from geoh5py.shared.utils import fetch_active_workspace +from pydantic import BaseModel, ConfigDict, Field from geoapps_utils.utils.transformations import ( rotate_points, + rotate_xyz, x_rotation_matrix, z_rotation_matrix, ) @@ -25,20 +29,121 @@ class PlateModel(BaseModel): :param strike_length: Length of the plate in the strike direction. :param dip_length: Length of the plate in the dip direction. :param width: Width of the plate. - :param origin: Origin point of the plate in the form [x, y, z]. + :param easting: Easting of the plate center. + :param northing: Northing of the plate center. + :param elevation: Elevation of the plate center. :param direction: Dip direction of the plate in degrees from North. :param dip: Dip angle of the plate in degrees below the horizontal. """ - model_config = ConfigDict(arbitrary_types_allowed=True) + model_config = ConfigDict(arbitrary_types_allowed=True, populate_by_name=True) strike_length: float dip_length: float width: float - origin: tuple[float, float, float] = (0.0, 0.0, 0.0) - direction: float = 0.0 + easting: float + northing: float + elevation: float + direction: float = Field(default=0.0, alias="dip_direction") dip: float = 0.0 + @property + def origin(self) -> tuple[float, float, float]: + return (self.easting, self.northing, self.elevation) + + +class Plate: + """ + Plate representation for surface extraction and masking. + + :param params: Parameters describing the plate. + """ + + def __init__(self, params: PlateModel): + self.params = params + + def mask(self, mesh: Octree) -> np.ndarray: + rotations = [ + z_rotation_matrix(np.deg2rad(self.params.direction)), + x_rotation_matrix(np.deg2rad(self.params.dip)), + ] + rotated_centers = rotate_points( + mesh.centroids, origin=self.params.origin, rotations=rotations + ) + return inside_plate(rotated_centers, self.params) + + def surface(self, workspace: Workspace, name: str = "plate") -> Surface: + """ + Create a rectangular prism geoh5py.Surface representing the plate. + + :param workspace: Workspace object to save the surface in. + :param name: Name of the surface. + """ + + with fetch_active_workspace(workspace) as ws: + surface = Surface.create( + ws, + vertices=self.vertices, + cells=self.triangles, + name=name, + ) + + return surface + + @property + def triangles(self) -> np.ndarray: + """Triangulation of the block.""" + return np.vstack( + [ + [0, 2, 1], + [1, 2, 3], + [0, 1, 4], + [4, 1, 5], + [1, 3, 5], + [5, 3, 7], + [2, 6, 3], + [3, 6, 7], + [0, 4, 2], + [2, 4, 6], + [4, 5, 6], + [6, 5, 7], + ] + ) + + @property + def vertices(self) -> np.ndarray: + """Vertices for triangulation of a rectangular prism in 3D space.""" + + u_1 = self.params.origin[0] - (self.params.strike_length / 2.0) + u_2 = self.params.origin[0] + (self.params.strike_length / 2.0) + v_1 = self.params.origin[1] - (self.params.dip_length / 2.0) + v_2 = self.params.origin[1] + (self.params.dip_length / 2.0) + w_1 = self.params.origin[2] - (self.params.width / 2.0) + w_2 = self.params.origin[2] + (self.params.width / 2.0) + + vertices = np.array( + [ + [u_1, v_1, w_1], + [u_2, v_1, w_1], + [u_1, v_2, w_1], + [u_2, v_2, w_1], + [u_1, v_1, w_2], + [u_2, v_1, w_2], + [u_1, v_2, w_2], + [u_2, v_2, w_2], + ] + ) + + return self._rotate(vertices) + + def _rotate(self, vertices: np.ndarray) -> np.ndarray: + """Rotate vertices and adjust for reference point.""" + theta = -1 * self.params.direction + phi = -1 * self.params.dip + rotated_vertices = rotate_xyz(vertices, list(self.params.origin), theta, phi) + + return rotated_vertices + def inside_plate( points: np.ndarray, diff --git a/tests/plates_test.py b/tests/plates_test.py index 9e5bcb5..52cb217 100644 --- a/tests/plates_test.py +++ b/tests/plates_test.py @@ -37,7 +37,9 @@ def test_inside_plate(tmp_path): strike_length=strike_length, dip_length=dip_length, width=width, - origin=(1.0, 0.0, 0.0), + easting=1.0, + northing=0.0, + elevation=0.0, ), ) model = np.zeros(len(grid.centroids)) @@ -77,6 +79,9 @@ def test_make_plate(tmp_path): strike_length=strike_length, dip_length=dip_length, width=width, + easting=0.0, + northing=0.0, + elevation=0.0, direction=direction, dip=dip, ), @@ -102,6 +107,9 @@ def test_make_plate(tmp_path): strike_length=strike_length, dip_length=dip_length, width=width, + easting=0.0, + northing=0.0, + elevation=0.0, direction=direction, dip=dip, ), @@ -140,16 +148,18 @@ def test_make_plate_multiple(tmp_path): strike_length=strike_length, dip_length=dip_length, width=width, + easting=-1.0, + northing=0.0, + elevation=0.0, direction=direction, dip=dip, - origin=(-1, 0, 0), ) model = make_plate( grid.centroids, plate=plate, background=0.0, ) - plate.origin = (1, 0, 0) + plate.easting = 1.0 model = make_plate(grid.centroids, plate=plate, background=model) grid.add_data({"plate model": {"values": model}}) @@ -162,3 +172,31 @@ def test_make_plate_multiple(tmp_path): & (grid.centroids[:, 2] <= (dip_length / 2)) ) assert np.all(model[mask] == 1.0) + + +def test_plate_alias(): + plate = PlateModel( + strike_length=15, + dip_length=7, + width=2, + easting=0.0, + northing=0.0, + elevation=0.0, + direction=90, + dip=0, + ) + assert plate.direction == 90 + assert "direction" in plate.model_dump() + + plate = PlateModel( + strike_length=15, + dip_length=7, + width=2, + easting=0.0, + northing=0.0, + elevation=0.0, + dip_direction=90, + dip=0, + ) + assert plate.direction == 90 + assert "dip_direction" in plate.model_dump(by_alias=True)