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

simpeg_drivers-assets/uijson/direct_current_2d_inversion.ui.json

@@ -283,7 +283,8 @@
         "dataType": "Float",
         "dataGroupType": [
             "Strike & dip",
-            "Dip direction & dip"
+            "Dip direction & dip",
+            "3D vector"
         ],
         "label": "Gradient rotation",
         "optional": true,

simpeg_drivers-assets/uijson/direct_current_3d_inversion.ui.json

@@ -238,7 +238,8 @@
         "dataType": "Float",
         "dataGroupType": [
             "Strike & dip",
-            "Dip direction & dip"
+            "Dip direction & dip",
+            "3D vector"
         ],
         "label": "Gradient rotation",
         "optional": true,

simpeg_drivers-assets/uijson/direct_current_batch2d_inversion.ui.json

@@ -265,7 +265,8 @@
         "dataType": "Float",
         "dataGroupType": [
             "Strike & dip",
-            "Dip direction & dip"
+            "Dip direction & dip",
+            "3D vector"
         ],
         "label": "Gradient rotation",
         "optional": true,

simpeg_drivers-assets/uijson/fem_inversion.ui.json

@@ -274,7 +274,8 @@
         "dataType": "Float",
         "dataGroupType": [
             "Strike & dip",
-            "Dip direction & dip"
+            "Dip direction & dip",
+            "3D vector"
         ],
         "label": "Gradient rotation",
         "optional": true,

simpeg_drivers-assets/uijson/gravity_inversion.ui.json

@@ -507,7 +507,8 @@
         "dataType": "Float",
         "dataGroupType": [
             "Strike & dip",
-            "Dip direction & dip"
+            "Dip direction & dip",
+            "3D vector"
         ],
         "label": "Gradient rotation",
         "optional": true,

simpeg_drivers-assets/uijson/induced_polarization_2d_inversion.ui.json

@@ -293,7 +293,8 @@
         "dataType": "Float",
         "dataGroupType": [
             "Strike & dip",
-            "Dip direction & dip"
+            "Dip direction & dip",
+            "3D vector"
         ],
         "label": "Gradient rotation",
         "optional": true,

simpeg_drivers-assets/uijson/induced_polarization_3d_inversion.ui.json

@@ -254,7 +254,8 @@
         "dataType": "Float",
         "dataGroupType": [
             "Strike & dip",
-            "Dip direction & dip"
+            "Dip direction & dip",
+            "3D vector"
         ],
         "label": "Gradient rotation",
         "optional": true,

simpeg_drivers-assets/uijson/induced_polarization_batch2d_inversion.ui.json

@@ -276,7 +276,8 @@
         "dataType": "Float",
         "dataGroupType": [
             "Strike & dip",
-            "Dip direction & dip"
+            "Dip direction & dip",
+            "3D vector"
         ],
         "label": "Gradient rotation",
         "optional": true,

simpeg_drivers-assets/uijson/joint_cross_gradient_inversion.ui.json

@@ -190,7 +190,11 @@
         "group": "Regularization",
         "association": "Cell",
         "dataType": "Float",
-        "dataGroupType": "Dip direction & dip",
+        "dataGroupType": [
+            "Strike & dip",
+            "Dip direction & dip",
+            "3D vector"
+        ],
         "label": "Gradient rotation",
         "optional": true,
         "enabled": false,

simpeg_drivers-assets/uijson/joint_surveys_inversion.ui.json

@@ -243,7 +243,8 @@
         "dataType": "Float",
         "dataGroupType": [
             "Strike & dip",
-            "Dip direction & dip"
+            "Dip direction & dip",
+            "3D vector"
         ],
         "label": "Gradient rotation",
         "optional": true,

simpeg_drivers-assets/uijson/magnetic_scalar_inversion.ui.json

@@ -539,7 +539,8 @@
         "dataType": "Float",
         "dataGroupType": [
             "Strike & dip",
-            "Dip direction & dip"
+            "Dip direction & dip",
+            "3D vector"
         ],
         "label": "Gradient rotation",
         "optional": true,

simpeg_drivers-assets/uijson/magnetic_vector_inversion.ui.json

@@ -604,7 +604,8 @@
         "dataType": "Float",
         "dataGroupType": [
             "Strike & dip",
-            "Dip direction & dip"
+            "Dip direction & dip",
+            "3D vector"
         ],
         "label": "Gradient rotation",
         "optional": true,

simpeg_drivers-assets/uijson/magnetotellurics_inversion.ui.json

@@ -459,7 +459,8 @@
         "dataType": "Float",
         "dataGroupType": [
             "Strike & dip",
-            "Dip direction & dip"
+            "Dip direction & dip",
+            "3D vector"
         ],
         "label": "Gradient rotation",
         "optional": true,

simpeg_drivers-assets/uijson/tdem_inversion.ui.json

@@ -311,7 +311,8 @@
         "dataType": "Float",
         "dataGroupType": [
             "Strike & dip",
-            "Dip direction & dip"
+            "Dip direction & dip",
+            "3D vector"
         ],
         "label": "Gradient rotation",
         "parent": "mesh",

simpeg_drivers-assets/uijson/tipper_inversion.ui.json

@@ -339,7 +339,8 @@
         "dataType": "Float",
         "dataGroupType": [
             "Strike & dip",
-            "Dip direction & dip"
+            "Dip direction & dip",
+            "3D vector"
         ],
         "label": "Gradient rotation",
         "optional": true,

simpeg_drivers/params.py

@@ -26,12 +26,15 @@
     ReferencedData,
 )
 from geoh5py.groups import PropertyGroup, SimPEGGroup, UIJsonGroup
+from geoh5py.groups.property_group_type import GroupTypeEnum
 from geoh5py.objects import DrapeModel, Octree, Points
 from geoh5py.shared.utils import fetch_active_workspace
 from geoh5py.ui_json import InputFile
 from pydantic import BaseModel, ConfigDict, field_validator, model_validator
+from simpeg.utils.mat_utils import cartesian2amplitude_dip_azimuth
 
 import simpeg_drivers
+from simpeg_drivers.utils.regularization import direction_and_dip
 
 
 logger = getLogger(__name__)
@@ -371,26 +374,32 @@ class BaseInversionOptions(CoreOptions):
     epsilon_cooling_factor: float = 1.2
 
     @property
-    def gradient_dip(self) -> np.ndarray | None:
-        """Gradient dip angle in clockwise radians from horizontal."""
+    def gradient_orientations(self) -> tuple(float, float):
+        """
+        Direction and dip angles for rotated gradient regularization.
+
+        Angles are in radians and are clockwise from North for direction
+        and clockwise from horizontal for dip.
+        """
+
         if self.gradient_rotation is not None:
-            dip_uid = self.gradient_rotation.properties[1]
-            dips = self.geoh5.get_entity(dip_uid)[0].values
-            return np.deg2rad(dips)
+            orientations = direction_and_dip(self.gradient_rotation)
+
+            return np.deg2rad(orientations)
+
         return None
 
     @property
-    def gradient_direction(self) -> np.ndarray | None:
-        """Gradient direction angle in clockwise radians from north"""
-        if self.gradient_rotation is not None:
-            from geoh5py.groups.property_group_type import GroupTypeEnum
+    def gradient_direction(self) -> np.ndarray:
+        if self.gradient_orientations is None:
+            return None
+        return self.gradient_orientations[:, 0]
 
-            direction_uid = self.gradient_rotation.properties[0]
-            directions = self.geoh5.get_entity(direction_uid)[0].values
-            if self.gradient_rotation.property_group_type == GroupTypeEnum.STRIKEDIP:
-                directions += 90.0
-            return np.deg2rad(directions)
-        return None
+    @property
+    def gradient_dip(self) -> np.ndarray:
+        if self.gradient_orientations is None:
+            return None
+        return self.gradient_orientations[:, 1]
 
 
 class EMDataMixin:

simpeg_drivers/utils/regularization.py

@@ -11,8 +11,11 @@
 import numpy as np
 import scipy.sparse as ssp
 from discretize import TreeMesh
+from geoh5py.groups import PropertyGroup
+from geoh5py.groups.property_group_type import GroupTypeEnum
 from simpeg.regularization import SparseSmoothness
 from simpeg.utils import mkvc, sdiag
+from simpeg.utils.mat_utils import cartesian2amplitude_dip_azimuth
 
 
 def cell_neighbors_along_axis(mesh: TreeMesh, axis: str) -> np.ndarray:
@@ -393,6 +396,49 @@
     return sdiag(1 / mesh.h_gridded[:, "xyz".find(axis)]) @ unit_grad
 
 
+def ensure_dip_direction_convention(orientations: np.ndarray, group_type: str):
+    """
+    Ensure orientations array has dip and direction convention.
+
+    :param orientations: Array of orientations.  Either n * 2 if Strike & dip
+        or Dip direction & dip group_type, or n * 3 if 3D Vector group_type.
+    :param group_type as specified in geoh5py.GroupTypeEnum.
+    """
+
+    if group_type == GroupTypeEnum.VECTOR:
+        orientations = orientations / np.c_[np.linalg.norm(orientations, axis=1)]
+        dip, direction = cartesian2amplitude_dip_azimuth(orientations)[:, 1:].T
+        dip += 90.0
+        orientations = np.c_[direction, dip]
+
+    if group_type in [GroupTypeEnum.STRIKEDIP]:
+        orientations[:, 0] = 90.0 + orientations[:, 0]
+
+    return orientations
+
+
+def direction_and_dip(property_group: PropertyGroup) -> list[np.ndarray]:
+    """Conversion of orientation group to direction and dip."""
+
+    group_type = property_group.property_group_type
+    if group_type not in [
+        GroupTypeEnum.VECTOR,
+        GroupTypeEnum.STRIKEDIP,
+        GroupTypeEnum.DIPDIR,
+    ]:
+        raise ValueError(
+            "Property group does not contain orientation data. "
+            "Type must be one of '3D vector', 'Strike & dip', or "
+            "'Dip direction & dip'."
+        )
+
+    orientations = np.vstack(
+        [property_group.parent.get_data(k)[0].values for k in property_group.properties]
+    ).T
+
+    return ensure_dip_direction_convention(orientations, group_type)
+
+
 def set_rotated_operators(
     function: SparseSmoothness,
     neighbors: np.ndarray,

tests/run_tests/driver_joint_cross_gradient_test.py

@@ -260,7 +260,6 @@ def test_joint_cross_gradient_inv_run(
             x_norm=0.0,
             y_norm=0.0,
             z_norm=0.0,
-            gradient_type="components",
             percentile=100,
             store_sensitivities="ram",
         )

tests/run_tests/driver_joint_surveys_test.py

@@ -154,7 +154,6 @@ def test_joint_surveys_inv_run(
             x_norm=0.0,
             y_norm=0.0,
             z_norm=0.0,
-            gradient_type="components",
             lower_bound=0.0,
             max_global_iterations=max_iterations,
             initial_beta_ratio=1e-2,

tests/run_tests/driver_rotated_gradients_test.py

@@ -88,18 +88,18 @@ def test_rotated_grad_run(
         mesh = geoh5.get_entity("mesh")[0]
 
         # Create property group with orientation
-        dip = np.ones(mesh.n_cells) * 45
-        azimuth = np.ones(mesh.n_cells) * 90
+        i = np.ones(mesh.n_cells)
+        j = np.zeros(mesh.n_cells)
+        k = np.ones(mesh.n_cells)
 
         data_list = mesh.add_data(
             {
-                "azimuth": {"values": azimuth},
-                "dip": {"values": dip},
+                "i": {"values": i},
+                "j": {"values": j},
+                "k": {"values": k},
             }
         )
-        pg = PropertyGroup(
-            mesh, properties=data_list, property_group_type="Dip direction & dip"
-        )
+        pg = PropertyGroup(mesh, properties=data_list, property_group_type="3D vector")
         topography = geoh5.get_entity("topography")[0]
 
         # Run the inverse

tests/utils_regularization_test.py

@@ -15,6 +15,8 @@
     cell_adjacent,
     cell_neighbors_along_axis,
     collect_all_neighbors,
+    direction_and_dip,
+    ensure_dip_direction_convention,
 )
 
 
@@ -90,3 +92,45 @@ def test_collect_all_neighbors():
     assert [15, 25, 25] in neighbor_centers
     assert [25, 25, 25] in neighbor_centers
     assert [15, 15, 15] not in neighbor_centers
+
+
+def test_ensure_dip_direction_convention():
+    # Rotate the vertical unit vector 37 degrees to the west and then 16
+    # degrees counter-clockwise about the z-axis.  Should result in a dip
+    # direction of 270 - 16 = 254 and a dip of 37.
+    Ry = np.array(
+        [
+            [
+                np.cos(np.deg2rad(-37)),
+                0,
+                np.sin(np.deg2rad(-37)),
+            ],
+            [0, 1, 0],
+            [-np.sin(np.deg2rad(-37)), 0, np.cos(np.deg2rad(-37))],
+        ]
+    )
+    Rz = np.array(
+        [
+            [np.cos(np.deg2rad(16)), -np.sin(np.deg2rad(16)), 0],
+            [np.sin(np.deg2rad(16)), np.cos(np.deg2rad(16)), 0],
+            [0, 0, 1],
+        ]
+    )
+    arbitrary_vector = Rz.dot(Ry.dot([0, 0, 1]))
+
+    orientations = np.array(
+        [
+            [1, 0, 1],
+            [0, 1, 1],
+            [-1, 0, 1],
+            [0, -1, 1],
+            [1, 0, np.sqrt(3)],
+            [0, 1, np.sqrt(3)],
+            [-1, 0, np.sqrt(3)],
+            [0, -1, np.sqrt(3)],
+            arbitrary_vector.tolist(),
+        ]
+    )
+    dir_dip = ensure_dip_direction_convention(orientations, group_type="3D vector")
+    assert np.allclose(dir_dip[:, 0], [90, 0, 270, 180] * 2 + [254])
+    assert np.allclose(dir_dip[:, 1], [45] * 4 + [30] * 4 + [37])