GEOPY-2108: Support 3D vector property group for orientation of rotated gradients

geoapps_utils/utils/transformations.py

@@ -9,17 +9,85 @@
 # '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
 
 import numpy as np
+import scipy.sparse as ssp
+
+
+def z_rotation_matrix(angle: np.ndarray) -> ssp.csr_matrix:
+    """
+    Sparse matrix for heterogeneous vector rotation about the z axis.
+
+    To be used in a matrix-vector product with an array of shape (n, 3)
+    where n is the number of 3-component vectors.
+
+    :param angle: Array of angles in radians for counterclockwise rotation
+        about the z-axis.
+    """
+    n = len(angle)
+    rza = np.c_[np.cos(angle), np.cos(angle), np.ones(n)].T
+    rza = rza.flatten(order="F")
+    rzb = np.c_[np.sin(angle), np.zeros(n), np.zeros(n)].T
+    rzb = rzb.flatten(order="F")
+    rzc = np.c_[-np.sin(angle), np.zeros(n), np.zeros(n)].T
+    rzc = rzc.flatten(order="F")
+    rot_z = ssp.diags([rzb[:-1], rza, rzc[:-1]], [-1, 0, 1])
+
+    return rot_z
+
+
+def x_rotation_matrix(angle: np.ndarray) -> ssp.csr_matrix:
+    """
+    Sparse matrix for heterogeneous vector rotation about the x axis.
+
+    To be used in a matrix-vector product with an array of shape (n, 3)
+    where n is the number of 3-component vectors.
+
+    :param angle: Array of angles in radians for counterclockwise rotation
+        about the x-axis.
+    """
+    n = len(angle)
+    rxa = np.c_[np.ones(n), np.cos(angle), np.cos(angle)].T
+    rxa = rxa.flatten(order="F")
+    rxb = np.c_[np.zeros(n), np.sin(angle), np.zeros(n)].T
+    rxb = rxb.flatten(order="F")
+    rxc = np.c_[np.zeros(n), -np.sin(angle), np.zeros(n)].T
+    rxc = rxc.flatten(order="F")
+    rot_x = ssp.diags([rxb[:-1], rxa, rxc[:-1]], [-1, 0, 1])
+
+    return rot_x
+
+
+def y_rotation_matrix(angle: np.ndarray) -> ssp.csr_matrix:
+    """
+    Sparse matrix for heterogeneous vector rotation about the y axis.
+
+    To be used in a matrix-vector product with an array of shape (n, 3)
+    where n is the number of 3-component vectors.
+
+    :param angle: Array of angles in radians for counterclockwise rotation
+        about the y-axis.
+    """
+    n = len(angle)
+    rxa = np.c_[np.cos(angle), np.ones(n), np.cos(angle)].T
+    rxa = rxa.flatten(order="F")
+    rxb = np.c_[-np.sin(angle), np.zeros(n), np.zeros(n)].T
+    rxb = rxb.flatten(order="F")
+    rxc = np.c_[np.sin(angle), np.zeros(n), np.zeros(n)].T
+    rxc = rxc.flatten(order="F")
+    rot_y = ssp.diags([rxb[:-2], rxa, rxc[:-2]], [-2, 0, 2])
+
+    return rot_y
 
 
 def rotate_xyz(xyz: np.ndarray, center: list, theta: float, phi: float = 0.0):
     """
-    Perform a counterclockwise rotation of scatter points around the x-axis,
-        then z-axis, about a center point.
+    Rotate points counterclockwise around the x then z axes, about a center point.
 
     :param xyz: shape(*, 2) or shape(*, 3) Input coordinates.
     :param center: len(2) or len(3) Coordinates for the center of rotation.
-    :param theta: Angle of rotation around z-axis in degrees.
-    :param phi: Angle of rotation around x-axis in degrees.
+    :param theta: Angle of rotation in counterclockwise degree about the z-axis
+        as viewed from above.
+    :param phi: Angle of rotation in couterclockwise degrees around x-axis
+        as viewed from the east.
     """
     return2d = False
     locs = xyz.copy()
@@ -55,3 +123,75 @@ def rotate_xyz(xyz: np.ndarray, center: list, theta: float, phi: float = 0.0):
         # Return 2-dimensional data if the input xyz was 2-dimensional.
         return xyz_out[:, :2]
     return xyz_out
+
+
+def ccw_east_to_cw_north(azimuth: np.ndarray) -> np.ndarray:
+    """Convert counterclockwise azimuth from east to clockwise from north."""
+    return (((5 * np.pi) / 2) - azimuth) % (2 * np.pi)
+
+
+def inclination_to_dip(inclination: np.ndarray) -> np.ndarray:
+    """Convert inclination from positive z-axis to dip from horizon."""
+    return inclination - (np.pi / 2)
+
+
+def cartesian_to_spherical(points: np.ndarray) -> np.ndarray:
+    """
+    Convert cartesian to spherical coordinate system.
+
+    :param points: Array of shape (n, 3) representing x, y, z coordinates of a point
+        in 3D space.
+
+    :returns: Array of shape (n, 3) representing the magnitude, azimuth and inclination
+        in spherical coordinates. The azimuth angle is measured in radians
+        counterclockwise from east in the range of 0 to 2pi as viewed from above, and
+        inclination angle is measured in radians from the positive z-axis.
+    """
+
+    magnitude = np.linalg.norm(points, axis=1)
+    inclination = np.arccos(points[:, 2] / magnitude)
+    azimuth = np.arctan2(points[:, 1], points[:, 0])
+    return np.column_stack((magnitude, azimuth, inclination))
+
+
+def spherical_normal_to_direction_and_dip(angles: np.ndarray) -> np.ndarray:
+    """
+    Convert normals in spherical coordinates to dip and direction of the tangent plane.
+
+    :param angles: Array of shape (n, 2) representing the azimuth and inclination angles
+        of a normal vector in spherical coordinates. The azimuth angle is measured in
+        radians counterclockwise from east in the range of 0 to 2pi as viewed from above,
+        and inclination angle is measured in radians from the positive z-axis.
+
+    :returns: Array of shape (n, 2) representing direction from 0 to 2pi radians
+        clockwise from north as viewed from above and dip from -pi to pi in positive
+        radians below the horizon and negative above.
+    """
+
+    rot_z = z_rotation_matrix(angles[:, 0])
+    rot_y = y_rotation_matrix(angles[:, 1])
+    tangents = np.tile([1, 0, 0], len(angles))
+    tangents = np.reshape(rot_z * rot_y * tangents, (-1, 3))
+    angles = cartesian_to_spherical(tangents)
+
+    return np.column_stack(
+        (ccw_east_to_cw_north(angles[:, 1]), inclination_to_dip(angles[:, 2]))
+    )
+
+
+def cartesian_normal_to_direction_and_dip(normals: np.ndarray) -> np.ndarray:
+    """
+    Convert 3D normal vectors to dip and direction.
+
+    :param normals: Array of shape (n, 3) representing the x, y, z components of a
+        normal vector in 3D space.
+
+    :returns: Array of shape (n, 2) representing azimuth from 0 to 2pi radians
+        clockwise from north as viewed from above and dip from -pi to pi in positive
+        radians below the horizon and negative above.
+    """
+
+    spherical_normals = cartesian_to_spherical(normals)
+    direction_and_dip = spherical_normal_to_direction_and_dip(spherical_normals[:, 1:])
+
+    return direction_and_dip

tests/transformations_test.py

@@ -11,8 +11,19 @@
 from __future__ import annotations
 
 import numpy as np
+import pytest
+from geoh5py import Workspace
+from geoh5py.groups.property_group import GroupTypeEnum
+from geoh5py.objects import Points
 
-from geoapps_utils.utils.transformations import rotate_xyz
+from geoapps_utils.utils.transformations import (
+    cartesian_normal_to_direction_and_dip,
+    cartesian_to_spherical,
+    rotate_xyz,
+    spherical_normal_to_direction_and_dip,
+    x_rotation_matrix,
+    z_rotation_matrix,
+)
 
 
 def test_positive_rotation_xyz():
@@ -33,3 +44,102 @@ def test_negative_rotation_xyz():
 
 def test_2d_input():
     assert (rotate_xyz(np.c_[1, 0], [0, 0], 45)).shape[1] == 2, "Error on 2D input."
+
+
+testdata = [
+    (-60, -30, 1, [30, 30]),
+    (-150, -30, 1, [-60, 30]),
+    (-240, -30, 1, [-150, 30]),
+    (-330, -30, 1, [120, 30]),
+    (-60, 30, -1, [30, 150]),
+    (-150, 30, -1, [-60, 150]),
+    (-240, 30, -1, [-150, 150]),
+    (-330, 30, -1, [120, 150]),
+]
+
+
+@pytest.mark.parametrize("theta,phi,polarity,expected", testdata)
+def test_cartesian_to_spherical(theta, phi, polarity, expected):
+    pole = rotate_xyz(xyz=np.c_[0, 0, polarity], center=[0, 0, 0], theta=theta, phi=phi)
+    angles = np.rad2deg(cartesian_to_spherical(pole)[:, 1:])
+    assert np.allclose(angles, [expected])
+
+
+testdata = [
+    (-60, -30, 1, [60, 30]),
+    (-150, -30, 1, [150, 30]),
+    (-240, -30, 1, [240, 30]),
+    (-330, -30, 1, [330, 30]),
+    (-60, 30, -1, [240, 30]),
+    (-150, 30, -1, [330, 30]),
+    (-240, 30, -1, [60, 30]),
+    (-330, 30, -1, [150, 30]),
+]
+
+
+@pytest.mark.parametrize("theta,phi,polarity,expected", testdata)
+def test_spherical_to_direction_and_dip_upwards(theta, phi, polarity, expected):
+    pole = rotate_xyz(xyz=np.c_[0, 0, polarity], center=[0, 0, 0], theta=theta, phi=phi)
+    spherical_coords = cartesian_to_spherical(pole)[:, 1:]
+    angles = np.rad2deg(spherical_normal_to_direction_and_dip(spherical_coords))
+    assert np.allclose(angles, [expected])
+
+
+def test_spherical_values(tmp_path):  # pylint: disable=too-many-locals
+    theta, phi = np.meshgrid(np.arange(0, 360, 10), np.arange(-90, 90, 10))
+    theta = theta.flatten()
+    phi = phi.flatten()
+
+    rad = 100.0
+    x = rad * np.cos(np.radians(theta)) * np.cos(np.radians(phi))
+    y = rad * np.sin(np.radians(theta)) * np.cos(np.radians(phi))
+    z = rad * np.sin(np.radians(phi))
+
+    xyz = np.c_[x, y, z]
+    rad_azim_incl = cartesian_to_spherical(xyz)
+    dirdip = cartesian_normal_to_direction_and_dip(xyz)
+    assert np.allclose(rad_azim_incl[:, 0], rad)
+    rot_x = x_rotation_matrix(-1 * dirdip[:, 1])
+    rot_z = z_rotation_matrix(-1 * dirdip[:, 0])
+    tangents = np.reshape(rot_z * rot_x * np.tile([0, 1, 0], len(dirdip)), (-1, 3))
+    assert np.allclose(np.linalg.norm(np.cross(xyz, tangents), axis=1), 100)
+
+    # Create a workspace and add the data for visual inspection
+    with Workspace.create(tmp_path / f"{__name__}.geoh5") as workspace:
+        points = Points.create(
+            workspace,
+            name="points_on_sphere",
+            vertices=xyz,
+        )
+
+        unit_normals = points.add_data(
+            {
+                "x": {"values": x / rad},
+                "y": {"values": y / rad},
+                "z": {"values": z / rad},
+            }
+        )
+        _ = points.create_property_group(
+            name="sphere_normals_vector",
+            properties=unit_normals,
+            property_group_type=GroupTypeEnum.VECTOR,
+        )
+
+        _ = points.add_data(
+            {
+                "azimuth": {"values": np.rad2deg(rad_azim_incl[:, 1])},
+                "inclination": {"values": np.rad2deg(rad_azim_incl[:, 2])},
+            }
+        )
+
+        direction, dip = points.add_data(
+            {
+                "direction": {"values": np.rad2deg(dirdip[:, 0])},
+                "dip": {"values": np.rad2deg(dirdip[:, 1])},
+            }
+        )
+        _ = points.create_property_group(
+            name="direction_dip",
+            properties=[direction, dip],
+            property_group_type=GroupTypeEnum.DIPDIR,
+        )