From c3da038e879369b28501a0eb3df8f4bf3b8fb0d1 Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Tue, 25 Mar 2025 14:41:36 -0700
Subject: [PATCH 01/29] Translate notebook into utilities

---
 simpeg_drivers/utils/regularization.py | 384 +++++++++++++++++++++++++
 1 file changed, 384 insertions(+)
 create mode 100644 simpeg_drivers/utils/regularization.py

diff --git a/simpeg_drivers/utils/regularization.py b/simpeg_drivers/utils/regularization.py
new file mode 100644
index 000000000..e61927f2f
--- /dev/null
+++ b/simpeg_drivers/utils/regularization.py
@@ -0,0 +1,384 @@
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+#  Copyright (c) 2025 Mira Geoscience Ltd.                                          '
+#                                                                                   '
+#  This file is part of simpeg-drivers package.                                     '
+#                                                                                   '
+#  simpeg-drivers is distributed under the terms and conditions of the MIT License  '
+#  (see LICENSE file at the root of this source code package).                      '
+#                                                                                   '
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+
+import numpy as np
+import scipy.sparse as ssp
+from discretize import TreeMesh
+from geoh5py.data import FloatData
+from simpeg.utils import mkvc, sdiag
+
+
+def cell_neighbors_along_axis(mesh: TreeMesh, axis: str) -> np.ndarray:
+    """
+    Get adjacent cells along provided axis for all cells in the mesh
+
+    :param mesh: Input TreeMesh.
+    :param axis: Cartesian axis along which to find neighbors.  Must be
+        'x', 'y', or 'z'.
+    """
+
+    if axis not in "xyz":
+        raise ValueError("Argument 'axis' must be one of 'x', 'y', or 'z'.")
+
+    stencil = getattr(mesh, f"stencil_cell_gradient_{axis}")
+    ith_neighbor, jth_neighbor, _ = ssp.find(stencil)
+    n_stencils = int(ith_neighbor.shape[0] / 2)
+    stencil_indices = jth_neighbor[np.argsort(ith_neighbor)].reshape((n_stencils, 2))
+
+    return np.sort(stencil_indices, axis=1)
+
+
+def cell_neighbors(mesh: TreeMesh) -> np.ndarray:
+    """Find all cell neighbors in a TreeMesh."""
+
+    x_neighbors = cell_neighbors_along_axis(mesh, "x")
+    x_neighbors_backward = np.fliplr(x_neighbors)
+    y_neighbors = cell_neighbors_along_axis(mesh, "y")
+    y_neighbors_backward = np.fliplr(y_neighbors)
+    max_index = np.max([x_neighbors.max(), y_neighbors.max()])
+    if mesh.dim == 3:
+        z_neighbors = cell_neighbors_along_axis(mesh, "z")
+        z_neighbors_backward = np.fliplr(z_neighbors)
+        max_index = np.max([max_index, z_neighbors.max()])
+
+    all_neighbors = []  # Store
+    x_adjacent = np.ones(max_index + 1, dtype="int") * -1
+    y_adjacent = np.ones(max_index + 1, dtype="int") * -1
+    x_adjacent_backward = np.ones(max_index + 1, dtype="int") * -1
+    y_adjacent_backward = np.ones(max_index + 1, dtype="int") * -1
+
+    x_adjacent[y_neighbors[:, 0]] = y_neighbors[:, 1]
+    y_adjacent[x_neighbors[:, 1]] = x_neighbors[:, 0]
+
+    x_adjacent_backward[y_neighbors_backward[:, 0]] = y_neighbors_backward[:, 1]
+    y_adjacent_backward[x_neighbors_backward[:, 1]] = x_neighbors_backward[:, 0]
+
+    all_neighbors += [x_neighbors]
+    all_neighbors += [y_neighbors]
+
+    all_neighbors += [np.c_[x_neighbors[:, 0], x_adjacent[x_neighbors[:, 1]]]]
+    all_neighbors += [np.c_[x_neighbors[:, 1], x_adjacent[x_neighbors[:, 0]]]]
+
+    all_neighbors += [np.c_[y_adjacent[y_neighbors[:, 0]], y_neighbors[:, 1]]]
+    all_neighbors += [np.c_[y_adjacent[y_neighbors[:, 1]], y_neighbors[:, 0]]]
+
+    # Repeat backward for Treemesh
+    all_neighbors += [x_neighbors_backward]
+    all_neighbors += [y_neighbors_backward]
+
+    all_neighbors += [
+        np.c_[
+            x_neighbors_backward[:, 0], x_adjacent_backward[x_neighbors_backward[:, 1]]
+        ]
+    ]
+    all_neighbors += [
+        np.c_[
+            x_neighbors_backward[:, 1], x_adjacent_backward[x_neighbors_backward[:, 0]]
+        ]
+    ]
+
+    # Stack all and keep only unique pairs
+    all_neighbors = np.vstack(all_neighbors)
+    all_neighbors = np.unique(all_neighbors, axis=0)
+
+    # Remove all the -1 for TreeMesh
+    all_neighbors = all_neighbors[
+        (all_neighbors[:, 0] != -1) & (all_neighbors[:, 1] != -1), :
+    ]
+
+    # Use all the neighbours on the xy plane to find neighbours in z
+    if mesh.dim == 3:
+        all_neighbors_z = []
+        z_adjacent = np.ones(max_index + 1, dtype="int") * -1
+        z_adjacent_backward = np.ones(max_index + 1, dtype="int") * -1
+
+        z_adjacent[z_neighbors[:, 0]] = z_neighbors[:, 1]
+        z_adjacent_backward[z_neighbors_backward[:, 0]] = z_neighbors_backward[:, 1]
+
+        all_neighbors_z += [z_neighbors]
+        all_neighbors_z += [z_neighbors_backward]
+
+        all_neighbors_z += [np.c_[all_neighbors[:, 0], z_adjacent[all_neighbors[:, 1]]]]
+        all_neighbors_z += [np.c_[all_neighbors[:, 1], z_adjacent[all_neighbors[:, 0]]]]
+
+        all_neighbors_z += [
+            np.c_[all_neighbors[:, 0], z_adjacent_backward[all_neighbors[:, 1]]]
+        ]
+        all_neighbors_z += [
+            np.c_[all_neighbors[:, 1], z_adjacent_backward[all_neighbors[:, 0]]]
+        ]
+
+        # Stack all and keep only unique pairs
+        all_neighbors = np.vstack([all_neighbors, np.vstack(all_neighbors_z)])
+        all_neighbors = np.unique(all_neighbors, axis=0)
+
+        # Remove all the -1 for TreeMesh
+        all_neighbors = all_neighbors[
+            (all_neighbors[:, 0] != -1) & (all_neighbors[:, 1] != -1), :
+        ]
+
+    return all_neighbors
+
+
+def rotation_matrix_2d(phi: float | FloatData, n_cells: int) -> ssp.csr_matrix:
+    """
+    Create a 2d rotation matrix for the xz plane.
+
+    :param phi: Angle in radians for rotation about the y-axis (xz plan).
+    :param n_cells: Number of cells in the mesh.
+    """
+
+    rza = mkvc(np.c_[np.cos(phi), np.cos(phi)].T)
+    rzb = mkvc(np.c_[np.sin(phi), np.zeros(n_cells)].T)
+    rzc = mkvc(np.c_[-np.sin(phi), np.zeros(n_cells)].T)
+    Rz = ssp.diags([rzb[:-1], rza, rzc[:-1]], [-1, 0, 1])
+
+    return Rz
+
+
+def rotation_matrix_3d(
+    psi: float | FloatData,
+    theta: float | FloatData,
+    phi: float | FloatData,
+    n_cells: int,
+) -> ssp.csr_matrix:
+    """
+    Create a rotation matrix for rotation about the z-axis, y-axis, and x-axis.
+
+    :param psi: Angle in radians for rotation about the z-axis.
+    :param theta: Angle in radians for rotation about the x-axis.
+    :param phi: Angle in radians for rotation about the y-axis.
+    :param n_cells: number of cells in the mesh.
+    """
+    rxa = mkvc(np.c_[np.ones(n_cells), np.cos(psi), np.cos(psi)].T)
+    rxb = mkvc(np.c_[np.zeros(n_cells), np.sin(psi), np.zeros(n_cells)].T)
+    rxc = mkvc(np.c_[np.zeros(n_cells), -np.sin(psi), np.zeros(n_cells)].T)
+    Rx = ssp.diags([rxb[:-1], rxa, rxc[:-1]], [-1, 0, 1])
+
+    rya = mkvc(np.c_[np.cos(theta), np.ones(n_cells), np.cos(theta)].T)
+    ryb = mkvc(np.c_[np.sin(theta), np.zeros(n_cells), np.zeros(n_cells)].T)
+    ryc = mkvc(np.c_[-np.sin(theta), np.zeros(n_cells), np.zeros(n_cells)].T)
+    Ry = ssp.diags([ryb[:-2], rya, ryc[:-2]], [-2, 0, 2])
+
+    rza = mkvc(np.c_[np.cos(phi), np.cos(phi), np.ones(n_cells)].T)
+    rzb = mkvc(np.c_[np.sin(phi), np.zeros(n_cells), np.zeros(n_cells)].T)
+    rzc = mkvc(np.c_[-np.sin(phi), np.zeros(n_cells), np.zeros(n_cells)].T)
+    Rz = ssp.diags([rzb[:-1], rza, rzc[:-1]], [-1, 0, 1])
+
+    return Rz * (Ry * Rx)
+
+
+def rotation_matrix(
+    mesh: TreeMesh,
+    phi: float | FloatData,
+    theta: float | FloatData | None = None,
+    psi: float | FloatData | None = None,
+) -> ssp.csr_matrix:
+    """
+    Create a rotation matrix for rotation about the z-axis, x-axis, and y-axis.
+
+    :param mesh: Input TreeMesh.
+    :param phi: Angle in radians for rotation about the y-axis.
+    :param theta: Angle in radians for rotation about the x-axis.
+    :param psi: Angle in radians for rotation about the z-axis.
+    """
+
+    def to_vector(angle):
+        if isinstance(angle, float):
+            return np.ones(mesh.n_cells) * angle
+        return angle
+
+    hx = mesh.h_gridded[:, 0]
+    hy = mesh.h_gridded[:, 1]
+    phi = to_vector(phi)
+    phi = np.arctan2((np.sin(phi) / hy), (np.cos(phi) / hx))
+
+    if mesh.dim == 2:
+        rotation = rotation_matrix_2d(phi, mesh.n_cells)
+
+    elif mesh.dim == 3:
+        if any(k is None for k in [psi, theta, phi]):
+            raise ValueError(
+                "Input must include 'psi', 'theta', and 'phi' for 3D mesh."
+            )
+
+        hz = mesh.h_gridded[:, 2]
+        psi = to_vector(psi)
+        psi = np.arctan2((np.sin(psi) / hz), (np.cos(psi) / hy))
+        theta = to_vector(theta)
+        theta = np.arctan2((np.sin(theta) / hz), (np.cos(theta) / hx))
+
+        rotation = rotation_matrix_3d(psi, theta, phi, mesh.n_cells)
+
+    else:
+        raise ValueError("Rotation matrix only implemented for 2D and 3D meshes")
+
+    return rotation
+
+
+def get_cell_normals(n_cells: int, axis: str, outward: bool) -> np.ndarray:
+    """
+    Returns cell normals for given axis and all cells.
+
+    :param n_cells: Number of cells in the mesh.
+    :param axis: Cartesian axis (one of 'x', 'y', or 'z'
+    :param outward: Direction of the normal. True for outward facing,
+        False for inward facing normals.
+    """
+
+    ind = 1 if outward else -1
+
+    if axis == "x":
+        normals = np.kron(np.ones(n_cells), np.c_[ind, 0, 0])
+    elif axis == "y":
+        normals = np.kron(np.ones(n_cells), np.c_[0, ind, 0])
+    elif axis == "z":
+        normals = np.kron(np.ones(n_cells), np.c_[0, 0, ind])
+    else:
+        raise ValueError("Axis must be one of 'x', 'y', or 'z'.")
+
+    return normals
+
+
+def get_cell_corners(
+    mesh: TreeMesh,
+    neighbors: np.ndarray,
+    normals: np.ndarray,
+) -> list[np.ndarray]:
+    """
+    Return the bottom southwest and top northeast nodes of all cells.
+
+    :param mesh: Input TreeMesh.
+    :param neighbors: Cell neighbors array.
+    :param normals: Cell normals array.
+    """
+
+    bottom_southwest = (
+        mesh.gridCC[neighbors[:, 0], :]
+        - mesh.h_gridded[neighbors[:, 0], :] / 2
+        + normals[neighbors[:, 0], :] * mesh.h_gridded[neighbors[:, 0], :]
+    )
+    top_northeast = (
+        mesh.gridCC[neighbors[:, 0], :]
+        + mesh.h_gridded[neighbors[:, 0], :] / 2
+        + normals[neighbors[:, 0], :] * mesh.h_gridded[neighbors[:, 0], :]
+    )
+
+    return [bottom_southwest, top_northeast]
+
+
+def get_neighbor_corners(mesh: TreeMesh, neighbors: np.ndarray):
+    """
+    Return the bottom southwest and top northeast corners.
+
+    :param mesh: Input TreeMesh.
+    :param neighbors: Cell neighbors array.
+    """
+
+    bottom_southwest = (
+        mesh.gridCC[neighbors[:, 1], :] - mesh.h_gridded[neighbors[:, 1], :] / 2
+    )
+    top_northeast = (
+        mesh.gridCC[neighbors[:, 1], :] + mesh.h_gridded[neighbors[:, 1], :] / 2
+    )
+
+    corners = [bottom_southwest, top_northeast]
+
+    return corners
+
+
+def partial_volumes(
+    mesh: TreeMesh, neighbors: np.ndarray, normals: np.ndarray
+) -> np.ndarray:
+    """
+    Compute partial volumes created by intersecting rotated and unrotated cells.
+
+    :param mesh: Input TreeMesh.
+    :param neighbors: Cell neighbors array.
+    :param normals: Cell normals array.
+    """
+    cell_corners = get_cell_corners(mesh, neighbors, normals)
+    neighbor_corners = get_neighbor_corners(mesh, neighbors)
+
+    volumes = np.ones(neighbors.shape[0])
+    for i in range(mesh.dim):
+        volumes *= np.max(
+            [
+                np.min([neighbor_corners[1][:, i], cell_corners[1][:, i]], axis=0)
+                - np.max([neighbor_corners[0][:, i], cell_corners[0][:, i]], axis=0),
+                np.zeros(neighbors.shape[0]),
+            ],
+            axis=0,
+        )
+
+    # Remove all rows of zero
+    ind = (volumes > 0) * (neighbors[:, 0] != neighbors[:, 1])
+    neighbors = neighbors[ind, :]
+    volumes = volumes[ind]
+
+    return volumes, neighbors
+
+
+def gradient_operator(
+    neighbors: np.ndarray, volumes: np.ndarray, n_cells: int
+) -> ssp.csr_matrix:
+    """
+    Assemble the sparse gradient operator.
+
+    :param neighbors: Cell neighbor array.
+    :param volumes: Partial volume array.
+    :param n_cells: Number of cells in mesh.
+    """
+    Grad = ssp.csr_matrix(
+        (volumes, (neighbors[:, 0], neighbors[:, 1])), shape=(n_cells, n_cells)
+    )
+
+    # Normalize rows
+    Vol = mkvc(Grad.sum(axis=1))
+    Vol[Vol > 0] = 1.0 / Vol[Vol > 0]
+    Grad = -sdiag(Vol) * Grad
+
+    diag = np.ones(n_cells)
+    diag[Vol == 0] = 0
+    Grad = sdiag(diag) + Grad
+
+    return Grad
+
+
+def rotated_gradient(
+    mesh: TreeMesh,
+    neighbors: np.ndarray,
+    axis: str,
+    phi: float | FloatData,
+    theta: float | FloatData | None = None,
+    psi: float | FloatData | None = None,
+    forward: bool = True,
+) -> ssp.csr_matrix:
+    """
+    Calculated rotated gradient operator using partial volumes.
+
+    :param mesh: Input TreeMesh.
+    :param neighbors: Cell neighbors array.
+    :param axis: Regularization axis.
+    :param phi: Angle in radians for rotation about the y-axis.
+    :param theta: Angle in radians for rotation about the x-axis.
+    :param psi: Angle in radians for rotation about the z-axis
+    :param forward: Whether to use forward or backward difference for
+        derivative approximations.
+    """
+
+    n_cells = mesh.n_cells
+
+    rot = rotation_matrix(mesh, phi, theta, psi)
+    normals = get_cell_normals(n_cells, axis, forward)
+    rotated_normals = (rot * normals.T).reshape(n_cells, mesh.dim)
+    volumes, neighbors = partial_volumes(mesh, neighbors, rotated_normals)
+
+    return gradient_operator(neighbors, volumes, n_cells)

From 17efc2dd05e8f4f5a80d8439300f3963e2cc952e Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Tue, 25 Mar 2025 15:33:44 -0700
Subject: [PATCH 02/29] support phi/theta/psi input

---
 .../direct_current_2d_inversion.ui.json       | 36 +++++++++++++++++++
 .../direct_current_3d_inversion.ui.json       | 36 +++++++++++++++++++
 .../direct_current_batch2d_inversion.ui.json  | 36 +++++++++++++++++++
 .../uijson/fem_inversion.ui.json              | 36 +++++++++++++++++++
 .../uijson/gravity_inversion.ui.json          | 36 +++++++++++++++++++
 .../induced_polarization_2d_inversion.ui.json | 36 +++++++++++++++++++
 .../induced_polarization_3d_inversion.ui.json | 36 +++++++++++++++++++
 ...ced_polarization_batch2d_inversion.ui.json | 36 +++++++++++++++++++
 .../uijson/joint_surveys_inversion.ui.json    | 36 +++++++++++++++++++
 .../uijson/magnetic_scalar_inversion.ui.json  | 36 +++++++++++++++++++
 .../uijson/magnetic_vector_inversion.ui.json  | 36 +++++++++++++++++++
 .../uijson/magnetotellurics_inversion.ui.json | 36 +++++++++++++++++++
 .../uijson/tdem_inversion.ui.json             | 36 +++++++++++++++++++
 .../uijson/tipper_inversion.ui.json           | 36 +++++++++++++++++++
 simpeg_drivers/driver.py                      | 12 ++++++-
 simpeg_drivers/params.py                      |  6 ++++
 .../potential_fields/gravity/uijson.py        |  3 ++
 17 files changed, 524 insertions(+), 1 deletion(-)

diff --git a/simpeg_drivers-assets/uijson/direct_current_2d_inversion.ui.json b/simpeg_drivers-assets/uijson/direct_current_2d_inversion.ui.json
index 6a1f5e6b0..4437bb9bc 100644
--- a/simpeg_drivers-assets/uijson/direct_current_2d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/direct_current_2d_inversion.ui.json
@@ -277,6 +277,42 @@
         "property": "",
         "enabled": true
     },
+    "gradient_phi": {
+        "group": "Regularization",
+        "label": "Rotated gradient phi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_theta": {
+        "group": "Regularization",
+        "label": "Rotated gradient theta",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_psi": {
+        "group": "Regularization",
+        "label": "Rotated gradient psi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
     "s_norm": {
         "association": "Cell",
         "dataType": "Float",
diff --git a/simpeg_drivers-assets/uijson/direct_current_3d_inversion.ui.json b/simpeg_drivers-assets/uijson/direct_current_3d_inversion.ui.json
index 6498ea8f8..8fce461ce 100644
--- a/simpeg_drivers-assets/uijson/direct_current_3d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/direct_current_3d_inversion.ui.json
@@ -232,6 +232,42 @@
         "property": "",
         "enabled": true
     },
+    "gradient_phi": {
+        "group": "Regularization",
+        "label": "Rotated gradient phi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_theta": {
+        "group": "Regularization",
+        "label": "Rotated gradient theta",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_psi": {
+        "group": "Regularization",
+        "label": "Rotated gradient psi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
     "s_norm": {
         "association": "Cell",
         "dataType": "Float",
diff --git a/simpeg_drivers-assets/uijson/direct_current_batch2d_inversion.ui.json b/simpeg_drivers-assets/uijson/direct_current_batch2d_inversion.ui.json
index a0d7a9f25..2c85caf41 100644
--- a/simpeg_drivers-assets/uijson/direct_current_batch2d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/direct_current_batch2d_inversion.ui.json
@@ -259,6 +259,42 @@
         "property": "",
         "enabled": true
     },
+    "gradient_phi": {
+        "group": "Regularization",
+        "label": "Rotated gradient phi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_theta": {
+        "group": "Regularization",
+        "label": "Rotated gradient theta",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_psi": {
+        "group": "Regularization",
+        "label": "Rotated gradient psi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
     "s_norm": {
         "association": "Cell",
         "dataType": "Float",
diff --git a/simpeg_drivers-assets/uijson/fem_inversion.ui.json b/simpeg_drivers-assets/uijson/fem_inversion.ui.json
index f85c96277..ca8ff5395 100644
--- a/simpeg_drivers-assets/uijson/fem_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/fem_inversion.ui.json
@@ -268,6 +268,42 @@
         "property": "",
         "enabled": true
     },
+    "gradient_phi": {
+        "group": "Regularization",
+        "label": "Rotated gradient phi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_theta": {
+        "group": "Regularization",
+        "label": "Rotated gradient theta",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_psi": {
+        "group": "Regularization",
+        "label": "Rotated gradient psi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
     "s_norm": {
         "association": "Cell",
         "dataType": "Float",
diff --git a/simpeg_drivers-assets/uijson/gravity_inversion.ui.json b/simpeg_drivers-assets/uijson/gravity_inversion.ui.json
index 9f61329f7..31fe3c467 100644
--- a/simpeg_drivers-assets/uijson/gravity_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/gravity_inversion.ui.json
@@ -501,6 +501,42 @@
         "property": "",
         "enabled": true
     },
+    "gradient_phi": {
+        "group": "Regularization",
+        "label": "Rotated gradient phi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_theta": {
+        "group": "Regularization",
+        "label": "Rotated gradient theta",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_psi": {
+        "group": "Regularization",
+        "label": "Rotated gradient psi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
     "s_norm": {
         "association": "Cell",
         "dataType": "Float",
diff --git a/simpeg_drivers-assets/uijson/induced_polarization_2d_inversion.ui.json b/simpeg_drivers-assets/uijson/induced_polarization_2d_inversion.ui.json
index f7cdec993..6dd2414d8 100644
--- a/simpeg_drivers-assets/uijson/induced_polarization_2d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/induced_polarization_2d_inversion.ui.json
@@ -286,6 +286,42 @@
         "property": "",
         "enabled": true
     },
+    "gradient_phi": {
+        "group": "Regularization",
+        "label": "Rotated gradient phi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_theta": {
+        "group": "Regularization",
+        "label": "Rotated gradient theta",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_psi": {
+        "group": "Regularization",
+        "label": "Rotated gradient psi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
     "s_norm": {
         "association": "Cell",
         "dataType": "Float",
diff --git a/simpeg_drivers-assets/uijson/induced_polarization_3d_inversion.ui.json b/simpeg_drivers-assets/uijson/induced_polarization_3d_inversion.ui.json
index d32f6e303..23694922e 100644
--- a/simpeg_drivers-assets/uijson/induced_polarization_3d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/induced_polarization_3d_inversion.ui.json
@@ -248,6 +248,42 @@
         "property": "",
         "enabled": true
     },
+    "gradient_phi": {
+        "group": "Regularization",
+        "label": "Rotated gradient phi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_theta": {
+        "group": "Regularization",
+        "label": "Rotated gradient theta",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_psi": {
+        "group": "Regularization",
+        "label": "Rotated gradient psi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
     "s_norm": {
         "association": "Cell",
         "dataType": "Float",
diff --git a/simpeg_drivers-assets/uijson/induced_polarization_batch2d_inversion.ui.json b/simpeg_drivers-assets/uijson/induced_polarization_batch2d_inversion.ui.json
index 5b74c35c5..cd68f23cd 100644
--- a/simpeg_drivers-assets/uijson/induced_polarization_batch2d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/induced_polarization_batch2d_inversion.ui.json
@@ -270,6 +270,42 @@
         "property": "",
         "enabled": true
     },
+    "gradient_phi": {
+        "group": "Regularization",
+        "label": "Rotated gradient phi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_theta": {
+        "group": "Regularization",
+        "label": "Rotated gradient theta",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_psi": {
+        "group": "Regularization",
+        "label": "Rotated gradient psi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
     "s_norm": {
         "association": "Cell",
         "dataType": "Float",
diff --git a/simpeg_drivers-assets/uijson/joint_surveys_inversion.ui.json b/simpeg_drivers-assets/uijson/joint_surveys_inversion.ui.json
index 5b8b832c0..b26411e29 100644
--- a/simpeg_drivers-assets/uijson/joint_surveys_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/joint_surveys_inversion.ui.json
@@ -237,6 +237,42 @@
         "property": "",
         "enabled": true
     },
+    "gradient_phi": {
+        "group": "Regularization",
+        "label": "Rotated gradient phi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_theta": {
+        "group": "Regularization",
+        "label": "Rotated gradient theta",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_psi": {
+        "group": "Regularization",
+        "label": "Rotated gradient psi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
     "s_norm": {
         "association": "Cell",
         "dataType": "Float",
diff --git a/simpeg_drivers-assets/uijson/magnetic_scalar_inversion.ui.json b/simpeg_drivers-assets/uijson/magnetic_scalar_inversion.ui.json
index 3e2630f63..033864756 100644
--- a/simpeg_drivers-assets/uijson/magnetic_scalar_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/magnetic_scalar_inversion.ui.json
@@ -533,6 +533,42 @@
         "property": "",
         "enabled": true
     },
+    "gradient_phi": {
+        "group": "Regularization",
+        "label": "Rotated gradient phi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_theta": {
+        "group": "Regularization",
+        "label": "Rotated gradient theta",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_psi": {
+        "group": "Regularization",
+        "label": "Rotated gradient psi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
     "s_norm": {
         "association": "Cell",
         "dataType": "Float",
diff --git a/simpeg_drivers-assets/uijson/magnetic_vector_inversion.ui.json b/simpeg_drivers-assets/uijson/magnetic_vector_inversion.ui.json
index 05c44729e..55d0d293e 100644
--- a/simpeg_drivers-assets/uijson/magnetic_vector_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/magnetic_vector_inversion.ui.json
@@ -597,6 +597,42 @@
         "property": "",
         "enabled": true
     },
+    "gradient_phi": {
+        "group": "Regularization",
+        "label": "Rotated gradient phi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_theta": {
+        "group": "Regularization",
+        "label": "Rotated gradient theta",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_psi": {
+        "group": "Regularization",
+        "label": "Rotated gradient psi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
     "s_norm": {
         "association": "Cell",
         "dataType": "Float",
diff --git a/simpeg_drivers-assets/uijson/magnetotellurics_inversion.ui.json b/simpeg_drivers-assets/uijson/magnetotellurics_inversion.ui.json
index 220c8152a..47e7712b8 100644
--- a/simpeg_drivers-assets/uijson/magnetotellurics_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/magnetotellurics_inversion.ui.json
@@ -453,6 +453,42 @@
         "property": "",
         "enabled": true
     },
+    "gradient_phi": {
+        "group": "Regularization",
+        "label": "Rotated gradient phi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_theta": {
+        "group": "Regularization",
+        "label": "Rotated gradient theta",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_psi": {
+        "group": "Regularization",
+        "label": "Rotated gradient psi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
     "s_norm": {
         "association": "Cell",
         "dataType": "Float",
diff --git a/simpeg_drivers-assets/uijson/tdem_inversion.ui.json b/simpeg_drivers-assets/uijson/tdem_inversion.ui.json
index 67a75b245..aadabeaae 100644
--- a/simpeg_drivers-assets/uijson/tdem_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/tdem_inversion.ui.json
@@ -305,6 +305,42 @@
         "property": "",
         "enabled": true
     },
+    "gradient_phi": {
+        "group": "Regularization",
+        "label": "Rotated gradient phi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_theta": {
+        "group": "Regularization",
+        "label": "Rotated gradient theta",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_psi": {
+        "group": "Regularization",
+        "label": "Rotated gradient psi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
     "s_norm": {
         "association": "Cell",
         "dataType": "Float",
diff --git a/simpeg_drivers-assets/uijson/tipper_inversion.ui.json b/simpeg_drivers-assets/uijson/tipper_inversion.ui.json
index 335f75289..428a80c68 100644
--- a/simpeg_drivers-assets/uijson/tipper_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/tipper_inversion.ui.json
@@ -333,6 +333,42 @@
         "property": "",
         "enabled": true
     },
+    "gradient_phi": {
+        "group": "Regularization",
+        "label": "Rotated gradient phi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_theta": {
+        "group": "Regularization",
+        "label": "Rotated gradient theta",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
+    "gradient_psi": {
+        "group": "Regularization",
+        "label": "Rotated gradient psi",
+        "value": 0,
+        "isValue": true,
+        "parent": "mesh",
+        "association": "Cell",
+        "dataType": "Float",
+        "property": "",
+        "optional": true,
+        "enabled": false
+    },
     "s_norm": {
         "association": "Cell",
         "dataType": "Float",
diff --git a/simpeg_drivers/driver.py b/simpeg_drivers/driver.py
index aff86b949..2d9fbe7f0 100644
--- a/simpeg_drivers/driver.py
+++ b/simpeg_drivers/driver.py
@@ -59,6 +59,7 @@
 )
 from simpeg_drivers.joint.params import BaseJointOptions
 from simpeg_drivers.utils.utils import tile_locations
+from simpeg_drivers.utils.regularization import cell_neighbors, rotated_gradient
 
 mlogger = logging.getLogger("distributed")
 mlogger.setLevel(logging.WARNING)
@@ -452,6 +453,7 @@ def get_regularization(self):
             comps = "sxz" if "2d" in self.params.inversion_type else "sxyz"
             avg_comps = "sxy" if "2d" in self.params.inversion_type else "sxyz"
             weights = ["alpha_s"] + [f"length_scale_{k}" for k in comps[1:]]
+            neighbors = cell_neighbors(reg.regularization_mesh)
             for comp, avg_comp, objfct, weight in zip(
                 comps, avg_comps, reg.objfcts, weights
             ):
@@ -466,7 +468,15 @@ def get_regularization(self):
                         getattr(reg.regularization_mesh, f"aveCC2F{avg_comp}") * weight
                     )
                     norm = getattr(reg.regularization_mesh, f"aveCC2F{avg_comp}") * norm
-
+                    Grad = rotated_gradient(
+                        mesh=reg.regularization_mesh,
+                        neighbors=neighbors,
+                        axis=comp,
+                        phi=self.params.gradient_phi,
+                        theta=self.params.gradient_theta,
+                        psi=self.params.gradient_psi,
+                    )
+                    setattr(objfct.regularization_mesh, f"cell_gradient_{comp}", Grad)
                 objfct.set_weights(**{comp: weight})
                 objfct.norm = norm
 
diff --git a/simpeg_drivers/params.py b/simpeg_drivers/params.py
index 2060c9aa1..34aeef282 100644
--- a/simpeg_drivers/params.py
+++ b/simpeg_drivers/params.py
@@ -262,6 +262,9 @@ class BaseInversionOptions(CoreOptions):
     :param length_scale_x: Length scale x.
     :param length_scale_y: Length scale y.
     :param length_scale_z: Length scale z.
+    :param gradient_phi: Gradient rotation about the y-axis
+    :param gradient_theta: Gradient rotation about the x-axis
+    :param gradient_psi: Gradient rotation about the z-axis
 
     :param s_norm: S norm.
     :param x_norm: X norm.
@@ -324,6 +327,9 @@ class BaseInversionOptions(CoreOptions):
     length_scale_x: float | FloatData = 1.0
     length_scale_y: float | FloatData | None = 1.0
     length_scale_z: float | FloatData = 1.0
+    gradient_phi: float | FloatData | None = None
+    gradient_theta: float | FloatData | None = None
+    gradient_psi: float | FloatData | None = None
 
     s_norm: float | FloatData | None = 0.0
     x_norm: float | FloatData = 2.0
diff --git a/simpeg_drivers/potential_fields/gravity/uijson.py b/simpeg_drivers/potential_fields/gravity/uijson.py
index 9957f4196..6738920a3 100644
--- a/simpeg_drivers/potential_fields/gravity/uijson.py
+++ b/simpeg_drivers/potential_fields/gravity/uijson.py
@@ -108,6 +108,9 @@ class GravityInversionUIJson(SimPEGDriversUIJson):
     length_scale_x: DataForm
     length_scale_y: DataForm
     length_scale_z: DataForm
+    gradient_phi: DataForm
+    gradient_theta: DataForm
+    gradient_psi: DataForm
     s_norm: DataForm
     x_norm: DataForm
     y_norm: DataForm

From 1cd25d9b749029751cc9e550ebab4cc406d1f2e4 Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Thu, 27 Mar 2025 10:55:57 -0700
Subject: [PATCH 03/29] Running with property group for gradient dip/direction

---
 .../direct_current_2d_inversion.ui.json       | 35 ++---------
 .../direct_current_3d_inversion.ui.json       | 35 ++---------
 .../direct_current_batch2d_inversion.ui.json  | 35 ++---------
 .../uijson/fem_inversion.ui.json              | 35 ++---------
 .../uijson/gravity_inversion.ui.json          | 35 ++---------
 .../induced_polarization_2d_inversion.ui.json | 35 ++---------
 .../induced_polarization_3d_inversion.ui.json | 35 ++---------
 ...ced_polarization_batch2d_inversion.ui.json | 35 ++---------
 .../uijson/joint_surveys_inversion.ui.json    | 35 ++---------
 .../uijson/magnetic_scalar_inversion.ui.json  | 35 ++---------
 .../uijson/magnetic_vector_inversion.ui.json  | 35 ++---------
 .../uijson/magnetotellurics_inversion.ui.json | 35 ++---------
 .../uijson/tdem_inversion.ui.json             | 35 ++---------
 .../uijson/tipper_inversion.ui.json           | 35 +++--------
 simpeg_drivers/components/models.py           | 18 ++++++
 simpeg_drivers/driver.py                      | 30 ++++++----
 simpeg_drivers/params.py                      | 20 ++++++-
 .../potential_fields/gravity/uijson.py        |  4 +-
 simpeg_drivers/utils/regularization.py        | 58 +++++++++----------
 19 files changed, 154 insertions(+), 466 deletions(-)

diff --git a/simpeg_drivers-assets/uijson/direct_current_2d_inversion.ui.json b/simpeg_drivers-assets/uijson/direct_current_2d_inversion.ui.json
index 4437bb9bc..b08f2469f 100644
--- a/simpeg_drivers-assets/uijson/direct_current_2d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/direct_current_2d_inversion.ui.json
@@ -277,41 +277,16 @@
         "property": "",
         "enabled": true
     },
-    "gradient_phi": {
+    "gradient_rotation": {
         "group": "Regularization",
-        "label": "Rotated gradient phi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
-    },
-    "gradient_theta": {
-        "group": "Regularization",
-        "label": "Rotated gradient theta",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
         "association": "Cell",
         "dataType": "Float",
-        "property": "",
+        "dataGroupType": "Dip direction & dip",
+        "label": "Gradient rotation",
         "optional": true,
-        "enabled": false
-    },
-    "gradient_psi": {
-        "group": "Regularization",
-        "label": "Rotated gradient psi",
-        "value": 0,
-        "isValue": true,
+        "enabled": false,
         "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "value": ""
     },
     "s_norm": {
         "association": "Cell",
diff --git a/simpeg_drivers-assets/uijson/direct_current_3d_inversion.ui.json b/simpeg_drivers-assets/uijson/direct_current_3d_inversion.ui.json
index 8fce461ce..563028b3a 100644
--- a/simpeg_drivers-assets/uijson/direct_current_3d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/direct_current_3d_inversion.ui.json
@@ -232,41 +232,16 @@
         "property": "",
         "enabled": true
     },
-    "gradient_phi": {
+    "gradient_rotation": {
         "group": "Regularization",
-        "label": "Rotated gradient phi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
-    },
-    "gradient_theta": {
-        "group": "Regularization",
-        "label": "Rotated gradient theta",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
         "association": "Cell",
         "dataType": "Float",
-        "property": "",
+        "dataGroupType": "Dip direction & dip",
+        "label": "Gradient rotation",
         "optional": true,
-        "enabled": false
-    },
-    "gradient_psi": {
-        "group": "Regularization",
-        "label": "Rotated gradient psi",
-        "value": 0,
-        "isValue": true,
+        "enabled": false,
         "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "value": ""
     },
     "s_norm": {
         "association": "Cell",
diff --git a/simpeg_drivers-assets/uijson/direct_current_batch2d_inversion.ui.json b/simpeg_drivers-assets/uijson/direct_current_batch2d_inversion.ui.json
index 2c85caf41..d834dfc7c 100644
--- a/simpeg_drivers-assets/uijson/direct_current_batch2d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/direct_current_batch2d_inversion.ui.json
@@ -259,41 +259,16 @@
         "property": "",
         "enabled": true
     },
-    "gradient_phi": {
+    "gradient_rotation": {
         "group": "Regularization",
-        "label": "Rotated gradient phi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
-    },
-    "gradient_theta": {
-        "group": "Regularization",
-        "label": "Rotated gradient theta",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
         "association": "Cell",
         "dataType": "Float",
-        "property": "",
+        "dataGroupType": "Dip direction & dip",
+        "label": "Gradient rotation",
         "optional": true,
-        "enabled": false
-    },
-    "gradient_psi": {
-        "group": "Regularization",
-        "label": "Rotated gradient psi",
-        "value": 0,
-        "isValue": true,
+        "enabled": false,
         "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "value": ""
     },
     "s_norm": {
         "association": "Cell",
diff --git a/simpeg_drivers-assets/uijson/fem_inversion.ui.json b/simpeg_drivers-assets/uijson/fem_inversion.ui.json
index ca8ff5395..fc3c00cca 100644
--- a/simpeg_drivers-assets/uijson/fem_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/fem_inversion.ui.json
@@ -268,41 +268,16 @@
         "property": "",
         "enabled": true
     },
-    "gradient_phi": {
+    "gradient_rotation": {
         "group": "Regularization",
-        "label": "Rotated gradient phi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
-    },
-    "gradient_theta": {
-        "group": "Regularization",
-        "label": "Rotated gradient theta",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
         "association": "Cell",
         "dataType": "Float",
-        "property": "",
+        "dataGroupType": "Dip direction & dip",
+        "label": "Gradient rotation",
         "optional": true,
-        "enabled": false
-    },
-    "gradient_psi": {
-        "group": "Regularization",
-        "label": "Rotated gradient psi",
-        "value": 0,
-        "isValue": true,
+        "enabled": false,
         "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "value": ""
     },
     "s_norm": {
         "association": "Cell",
diff --git a/simpeg_drivers-assets/uijson/gravity_inversion.ui.json b/simpeg_drivers-assets/uijson/gravity_inversion.ui.json
index 31fe3c467..389ba5405 100644
--- a/simpeg_drivers-assets/uijson/gravity_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/gravity_inversion.ui.json
@@ -501,41 +501,16 @@
         "property": "",
         "enabled": true
     },
-    "gradient_phi": {
+    "gradient_rotation": {
         "group": "Regularization",
-        "label": "Rotated gradient phi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
-    },
-    "gradient_theta": {
-        "group": "Regularization",
-        "label": "Rotated gradient theta",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
         "association": "Cell",
         "dataType": "Float",
-        "property": "",
+        "dataGroupType": "Dip direction & dip",
+        "label": "Gradient rotation",
         "optional": true,
-        "enabled": false
-    },
-    "gradient_psi": {
-        "group": "Regularization",
-        "label": "Rotated gradient psi",
-        "value": 0,
-        "isValue": true,
+        "enabled": false,
         "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "value": ""
     },
     "s_norm": {
         "association": "Cell",
diff --git a/simpeg_drivers-assets/uijson/induced_polarization_2d_inversion.ui.json b/simpeg_drivers-assets/uijson/induced_polarization_2d_inversion.ui.json
index 6dd2414d8..04fab5563 100644
--- a/simpeg_drivers-assets/uijson/induced_polarization_2d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/induced_polarization_2d_inversion.ui.json
@@ -286,41 +286,16 @@
         "property": "",
         "enabled": true
     },
-    "gradient_phi": {
+    "gradient_rotation": {
         "group": "Regularization",
-        "label": "Rotated gradient phi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
-    },
-    "gradient_theta": {
-        "group": "Regularization",
-        "label": "Rotated gradient theta",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
         "association": "Cell",
         "dataType": "Float",
-        "property": "",
+        "dataGroupType": "Dip direction & dip",
+        "label": "Gradient rotation",
         "optional": true,
-        "enabled": false
-    },
-    "gradient_psi": {
-        "group": "Regularization",
-        "label": "Rotated gradient psi",
-        "value": 0,
-        "isValue": true,
+        "enabled": false,
         "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "value": ""
     },
     "s_norm": {
         "association": "Cell",
diff --git a/simpeg_drivers-assets/uijson/induced_polarization_3d_inversion.ui.json b/simpeg_drivers-assets/uijson/induced_polarization_3d_inversion.ui.json
index 23694922e..829f197ed 100644
--- a/simpeg_drivers-assets/uijson/induced_polarization_3d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/induced_polarization_3d_inversion.ui.json
@@ -248,41 +248,16 @@
         "property": "",
         "enabled": true
     },
-    "gradient_phi": {
+    "gradient_rotation": {
         "group": "Regularization",
-        "label": "Rotated gradient phi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
-    },
-    "gradient_theta": {
-        "group": "Regularization",
-        "label": "Rotated gradient theta",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
         "association": "Cell",
         "dataType": "Float",
-        "property": "",
+        "dataGroupType": "Dip direction & dip",
+        "label": "Gradient rotation",
         "optional": true,
-        "enabled": false
-    },
-    "gradient_psi": {
-        "group": "Regularization",
-        "label": "Rotated gradient psi",
-        "value": 0,
-        "isValue": true,
+        "enabled": false,
         "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "value": ""
     },
     "s_norm": {
         "association": "Cell",
diff --git a/simpeg_drivers-assets/uijson/induced_polarization_batch2d_inversion.ui.json b/simpeg_drivers-assets/uijson/induced_polarization_batch2d_inversion.ui.json
index cd68f23cd..dfa48ffd5 100644
--- a/simpeg_drivers-assets/uijson/induced_polarization_batch2d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/induced_polarization_batch2d_inversion.ui.json
@@ -270,41 +270,16 @@
         "property": "",
         "enabled": true
     },
-    "gradient_phi": {
+    "gradient_rotation": {
         "group": "Regularization",
-        "label": "Rotated gradient phi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
-    },
-    "gradient_theta": {
-        "group": "Regularization",
-        "label": "Rotated gradient theta",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
         "association": "Cell",
         "dataType": "Float",
-        "property": "",
+        "dataGroupType": "Dip direction & dip",
+        "label": "Gradient rotation",
         "optional": true,
-        "enabled": false
-    },
-    "gradient_psi": {
-        "group": "Regularization",
-        "label": "Rotated gradient psi",
-        "value": 0,
-        "isValue": true,
+        "enabled": false,
         "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "value": ""
     },
     "s_norm": {
         "association": "Cell",
diff --git a/simpeg_drivers-assets/uijson/joint_surveys_inversion.ui.json b/simpeg_drivers-assets/uijson/joint_surveys_inversion.ui.json
index b26411e29..3cf0c1452 100644
--- a/simpeg_drivers-assets/uijson/joint_surveys_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/joint_surveys_inversion.ui.json
@@ -237,41 +237,16 @@
         "property": "",
         "enabled": true
     },
-    "gradient_phi": {
+    "gradient_rotation": {
         "group": "Regularization",
-        "label": "Rotated gradient phi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
-    },
-    "gradient_theta": {
-        "group": "Regularization",
-        "label": "Rotated gradient theta",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
         "association": "Cell",
         "dataType": "Float",
-        "property": "",
+        "dataGroupType": "Dip direction & dip",
+        "label": "Gradient rotation",
         "optional": true,
-        "enabled": false
-    },
-    "gradient_psi": {
-        "group": "Regularization",
-        "label": "Rotated gradient psi",
-        "value": 0,
-        "isValue": true,
+        "enabled": false,
         "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "value": ""
     },
     "s_norm": {
         "association": "Cell",
diff --git a/simpeg_drivers-assets/uijson/magnetic_scalar_inversion.ui.json b/simpeg_drivers-assets/uijson/magnetic_scalar_inversion.ui.json
index 033864756..20defc673 100644
--- a/simpeg_drivers-assets/uijson/magnetic_scalar_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/magnetic_scalar_inversion.ui.json
@@ -533,41 +533,16 @@
         "property": "",
         "enabled": true
     },
-    "gradient_phi": {
+    "gradient_rotation": {
         "group": "Regularization",
-        "label": "Rotated gradient phi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
-    },
-    "gradient_theta": {
-        "group": "Regularization",
-        "label": "Rotated gradient theta",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
         "association": "Cell",
         "dataType": "Float",
-        "property": "",
+        "dataGroupType": "Dip direction & dip",
+        "label": "Gradient rotation",
         "optional": true,
-        "enabled": false
-    },
-    "gradient_psi": {
-        "group": "Regularization",
-        "label": "Rotated gradient psi",
-        "value": 0,
-        "isValue": true,
+        "enabled": false,
         "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "value": ""
     },
     "s_norm": {
         "association": "Cell",
diff --git a/simpeg_drivers-assets/uijson/magnetic_vector_inversion.ui.json b/simpeg_drivers-assets/uijson/magnetic_vector_inversion.ui.json
index 55d0d293e..d08a75f02 100644
--- a/simpeg_drivers-assets/uijson/magnetic_vector_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/magnetic_vector_inversion.ui.json
@@ -597,41 +597,16 @@
         "property": "",
         "enabled": true
     },
-    "gradient_phi": {
+    "gradient_rotation": {
         "group": "Regularization",
-        "label": "Rotated gradient phi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
-    },
-    "gradient_theta": {
-        "group": "Regularization",
-        "label": "Rotated gradient theta",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
         "association": "Cell",
         "dataType": "Float",
-        "property": "",
+        "dataGroupType": "Dip direction & dip",
+        "label": "Gradient rotation",
         "optional": true,
-        "enabled": false
-    },
-    "gradient_psi": {
-        "group": "Regularization",
-        "label": "Rotated gradient psi",
-        "value": 0,
-        "isValue": true,
+        "enabled": false,
         "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "value": ""
     },
     "s_norm": {
         "association": "Cell",
diff --git a/simpeg_drivers-assets/uijson/magnetotellurics_inversion.ui.json b/simpeg_drivers-assets/uijson/magnetotellurics_inversion.ui.json
index 47e7712b8..0f65c8b40 100644
--- a/simpeg_drivers-assets/uijson/magnetotellurics_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/magnetotellurics_inversion.ui.json
@@ -453,41 +453,16 @@
         "property": "",
         "enabled": true
     },
-    "gradient_phi": {
+    "gradient_rotation": {
         "group": "Regularization",
-        "label": "Rotated gradient phi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
-    },
-    "gradient_theta": {
-        "group": "Regularization",
-        "label": "Rotated gradient theta",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
         "association": "Cell",
         "dataType": "Float",
-        "property": "",
+        "dataGroupType": "Dip direction & dip",
+        "label": "Gradient rotation",
         "optional": true,
-        "enabled": false
-    },
-    "gradient_psi": {
-        "group": "Regularization",
-        "label": "Rotated gradient psi",
-        "value": 0,
-        "isValue": true,
+        "enabled": false,
         "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "value": ""
     },
     "s_norm": {
         "association": "Cell",
diff --git a/simpeg_drivers-assets/uijson/tdem_inversion.ui.json b/simpeg_drivers-assets/uijson/tdem_inversion.ui.json
index aadabeaae..b6e49561f 100644
--- a/simpeg_drivers-assets/uijson/tdem_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/tdem_inversion.ui.json
@@ -305,41 +305,14 @@
         "property": "",
         "enabled": true
     },
-    "gradient_phi": {
+    "gradient_rotation": {
         "group": "Regularization",
-        "label": "Rotated gradient phi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
         "association": "Cell",
         "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
-    },
-    "gradient_theta": {
-        "group": "Regularization",
-        "label": "Rotated gradient theta",
-        "value": 0,
-        "isValue": true,
+        "dataGroupType": "Dip direction & dip",
+        "label": "Gradient rotation",
         "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
-    },
-    "gradient_psi": {
-        "group": "Regularization",
-        "label": "Rotated gradient psi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "value": ""
     },
     "s_norm": {
         "association": "Cell",
diff --git a/simpeg_drivers-assets/uijson/tipper_inversion.ui.json b/simpeg_drivers-assets/uijson/tipper_inversion.ui.json
index 428a80c68..a865739cb 100644
--- a/simpeg_drivers-assets/uijson/tipper_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/tipper_inversion.ui.json
@@ -335,39 +335,20 @@
     },
     "gradient_phi": {
         "group": "Regularization",
-        "label": "Rotated gradient phi",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "main": true,
+        "label": "test",
+        "value": 9
     },
-    "gradient_theta": {
+    "gradient_rotation": {
         "group": "Regularization",
-        "label": "Rotated gradient theta",
-        "value": 0,
-        "isValue": true,
-        "parent": "mesh",
         "association": "Cell",
         "dataType": "Float",
-        "property": "",
+        "dataGroupType": "Dip direction & dip",
+        "label": "Gradient rotation",
         "optional": true,
-        "enabled": false
-    },
-    "gradient_psi": {
-        "group": "Regularization",
-        "label": "Rotated gradient psi",
-        "value": 0,
-        "isValue": true,
+        "enabled": false,
         "parent": "mesh",
-        "association": "Cell",
-        "dataType": "Float",
-        "property": "",
-        "optional": true,
-        "enabled": false
+        "value": ""
     },
     "s_norm": {
         "association": "Cell",
diff --git a/simpeg_drivers/components/models.py b/simpeg_drivers/components/models.py
index bba96ba92..745a0480f 100644
--- a/simpeg_drivers/components/models.py
+++ b/simpeg_drivers/components/models.py
@@ -77,6 +77,8 @@ def __init__(self, driver: InversionDriver):
         self._length_scale_x = InversionModel(driver, "length_scale_x")
         self._length_scale_y = InversionModel(driver, "length_scale_y")
         self._length_scale_z = InversionModel(driver, "length_scale_z")
+        self._gradient_dip = InversionModel(driver, "gradient_dip")
+        self._gradient_direction = InversionModel(driver, "gradient_direction")
         self._s_norm = InversionModel(driver, "s_norm")
         self._x_norm = InversionModel(driver, "x_norm")
         self._y_norm = InversionModel(driver, "y_norm")
@@ -256,6 +258,20 @@ def length_scale_z(self) -> np.ndarray | None:
 
         return self._length_scale_z.model.copy()
 
+    @property
+    def gradient_dip(self) -> np.ndarray | None:
+        if self._gradient_dip.model is None:
+            return None
+
+        return self._gradient_dip.model.copy()
+
+    @property
+    def gradient_direction(self) -> np.ndarray | None:
+        if self._gradient_direction.model is None:
+            return None
+
+        return self._gradient_direction.model.copy()
+
     @property
     def s_norm(self) -> np.ndarray | None:
         if self._s_norm.model is None:
@@ -358,6 +374,8 @@ class InversionModel:
         "length_scale_x",
         "length_scale_y",
         "length_scale_z",
+        "gradient_dip",
+        "gradient_direction",
         "s_norm",
         "x_norm",
         "y_norm",
diff --git a/simpeg_drivers/driver.py b/simpeg_drivers/driver.py
index 2d9fbe7f0..3ccea7391 100644
--- a/simpeg_drivers/driver.py
+++ b/simpeg_drivers/driver.py
@@ -441,6 +441,7 @@ def get_regularization(self):
             return BaseRegularization(mesh=self.inversion_mesh.mesh)
 
         reg_funcs = []
+        is_rotated = self.params.gradient_rotation is not None
         for mapping in self.mapping:
             reg = Sparse(
                 self.inversion_mesh.mesh,
@@ -449,11 +450,14 @@ def get_regularization(self):
                 reference_model=self.models.reference,
             )
 
+            if is_rotated:
+                neighbors = cell_neighbors(reg.regularization_mesh)
+
             # Adjustment for 2D versus 3D problems
             comps = "sxz" if "2d" in self.params.inversion_type else "sxyz"
             avg_comps = "sxy" if "2d" in self.params.inversion_type else "sxyz"
             weights = ["alpha_s"] + [f"length_scale_{k}" for k in comps[1:]]
-            neighbors = cell_neighbors(reg.regularization_mesh)
+
             for comp, avg_comp, objfct, weight in zip(
                 comps, avg_comps, reg.objfcts, weights
             ):
@@ -468,15 +472,21 @@ def get_regularization(self):
                         getattr(reg.regularization_mesh, f"aveCC2F{avg_comp}") * weight
                     )
                     norm = getattr(reg.regularization_mesh, f"aveCC2F{avg_comp}") * norm
-                    Grad = rotated_gradient(
-                        mesh=reg.regularization_mesh,
-                        neighbors=neighbors,
-                        axis=comp,
-                        phi=self.params.gradient_phi,
-                        theta=self.params.gradient_theta,
-                        psi=self.params.gradient_psi,
-                    )
-                    setattr(objfct.regularization_mesh, f"cell_gradient_{comp}", Grad)
+
+                    if is_rotated:
+                        Grad = rotated_gradient(
+                            mesh=reg.regularization_mesh.mesh,
+                            neighbors=neighbors,
+                            axis=comp,
+                            dip=self.models.gradient_dip,
+                            direction=self.models.gradient_direction,
+                        )
+                        setattr(
+                            reg.regularization_mesh.mesh,
+                            f"_stencil_cell_gradient_{comp}",
+                            Grad,
+                        )
+
                 objfct.set_weights(**{comp: weight})
                 objfct.norm = norm
 
diff --git a/simpeg_drivers/params.py b/simpeg_drivers/params.py
index 34aeef282..d274c5b46 100644
--- a/simpeg_drivers/params.py
+++ b/simpeg_drivers/params.py
@@ -327,9 +327,7 @@ class BaseInversionOptions(CoreOptions):
     length_scale_x: float | FloatData = 1.0
     length_scale_y: float | FloatData | None = 1.0
     length_scale_z: float | FloatData = 1.0
-    gradient_phi: float | FloatData | None = None
-    gradient_theta: float | FloatData | None = None
-    gradient_psi: float | FloatData | None = None
+    gradient_rotation: PropertyGroup | None = None
 
     s_norm: float | FloatData | None = 0.0
     x_norm: float | FloatData = 2.0
@@ -362,6 +360,22 @@ class BaseInversionOptions(CoreOptions):
     percentile: float = 95.0
     epsilon_cooling_factor: float = 1.2
 
+    @property
+    def gradient_dip(self):
+        """Returns the dip of the gradient rotation."""
+        if self.gradient_rotation is not None:
+            dip_uid = self.gradient_rotation.properties[1]
+            return self.geoh5.get_entity(dip_uid)[0].values
+        return None
+
+    @property
+    def gradient_direction(self):
+        """Returns the direction of the dip of the gradient rotation"""
+        if self.gradient_rotation is not None:
+            direction_uid = self.gradient_rotation.properties[0]
+            return self.geoh5.get_entity(direction_uid)[0].values
+        return None
+
 
 class EMDataMixin:
     """
diff --git a/simpeg_drivers/potential_fields/gravity/uijson.py b/simpeg_drivers/potential_fields/gravity/uijson.py
index 6738920a3..408b4a922 100644
--- a/simpeg_drivers/potential_fields/gravity/uijson.py
+++ b/simpeg_drivers/potential_fields/gravity/uijson.py
@@ -108,9 +108,7 @@ class GravityInversionUIJson(SimPEGDriversUIJson):
     length_scale_x: DataForm
     length_scale_y: DataForm
     length_scale_z: DataForm
-    gradient_phi: DataForm
-    gradient_theta: DataForm
-    gradient_psi: DataForm
+    gradient_rotation: GroupForm
     s_norm: DataForm
     x_norm: DataForm
     y_norm: DataForm
diff --git a/simpeg_drivers/utils/regularization.py b/simpeg_drivers/utils/regularization.py
index e61927f2f..5aaa95c0e 100644
--- a/simpeg_drivers/utils/regularization.py
+++ b/simpeg_drivers/utils/regularization.py
@@ -11,7 +11,6 @@
 import numpy as np
 import scipy.sparse as ssp
 from discretize import TreeMesh
-from geoh5py.data import FloatData
 from simpeg.utils import mkvc, sdiag
 
 
@@ -27,7 +26,7 @@ def cell_neighbors_along_axis(mesh: TreeMesh, axis: str) -> np.ndarray:
     if axis not in "xyz":
         raise ValueError("Argument 'axis' must be one of 'x', 'y', or 'z'.")
 
-    stencil = getattr(mesh, f"stencil_cell_gradient_{axis}")
+    stencil = getattr(mesh, f"cell_gradient_{axis}")
     ith_neighbor, jth_neighbor, _ = ssp.find(stencil)
     n_stencils = int(ith_neighbor.shape[0] / 2)
     stencil_indices = jth_neighbor[np.argsort(ith_neighbor)].reshape((n_stencils, 2))
@@ -127,27 +126,25 @@ def cell_neighbors(mesh: TreeMesh) -> np.ndarray:
     return all_neighbors
 
 
-def rotation_matrix_2d(phi: float | FloatData, n_cells: int) -> ssp.csr_matrix:
+def rotation_matrix_2d(phi: np.ndarray) -> ssp.csr_matrix:
     """
     Create a 2d rotation matrix for the xz plane.
 
     :param phi: Angle in radians for rotation about the y-axis (xz plan).
-    :param n_cells: Number of cells in the mesh.
     """
-
+    n_cells = len(phi)
     rza = mkvc(np.c_[np.cos(phi), np.cos(phi)].T)
     rzb = mkvc(np.c_[np.sin(phi), np.zeros(n_cells)].T)
     rzc = mkvc(np.c_[-np.sin(phi), np.zeros(n_cells)].T)
-    Rz = ssp.diags([rzb[:-1], rza, rzc[:-1]], [-1, 0, 1])
+    Ry = ssp.diags([rzb[:-1], rza, rzc[:-1]], [-1, 0, 1])
 
-    return Rz
+    return Ry
 
 
 def rotation_matrix_3d(
-    psi: float | FloatData,
-    theta: float | FloatData,
-    phi: float | FloatData,
-    n_cells: int,
+    psi: np.ndarray,
+    theta: np.ndarray,
+    phi: np.ndarray,
 ) -> ssp.csr_matrix:
     """
     Create a rotation matrix for rotation about the z-axis, y-axis, and x-axis.
@@ -157,6 +154,7 @@ def rotation_matrix_3d(
     :param phi: Angle in radians for rotation about the y-axis.
     :param n_cells: number of cells in the mesh.
     """
+    n_cells = len(phi)
     rxa = mkvc(np.c_[np.ones(n_cells), np.cos(psi), np.cos(psi)].T)
     rxb = mkvc(np.c_[np.zeros(n_cells), np.sin(psi), np.zeros(n_cells)].T)
     rxc = mkvc(np.c_[np.zeros(n_cells), -np.sin(psi), np.zeros(n_cells)].T)
@@ -177,9 +175,9 @@ def rotation_matrix_3d(
 
 def rotation_matrix(
     mesh: TreeMesh,
-    phi: float | FloatData,
-    theta: float | FloatData | None = None,
-    psi: float | FloatData | None = None,
+    phi: np.ndarray,
+    theta: np.ndarray | None = None,
+    psi: np.ndarray | None = None,
 ) -> ssp.csr_matrix:
     """
     Create a rotation matrix for rotation about the z-axis, x-axis, and y-axis.
@@ -190,18 +188,12 @@ def rotation_matrix(
     :param psi: Angle in radians for rotation about the z-axis.
     """
 
-    def to_vector(angle):
-        if isinstance(angle, float):
-            return np.ones(mesh.n_cells) * angle
-        return angle
-
     hx = mesh.h_gridded[:, 0]
     hy = mesh.h_gridded[:, 1]
-    phi = to_vector(phi)
     phi = np.arctan2((np.sin(phi) / hy), (np.cos(phi) / hx))
 
     if mesh.dim == 2:
-        rotation = rotation_matrix_2d(phi, mesh.n_cells)
+        rotation = rotation_matrix_2d(phi)
 
     elif mesh.dim == 3:
         if any(k is None for k in [psi, theta, phi]):
@@ -210,12 +202,10 @@ def to_vector(angle):
             )
 
         hz = mesh.h_gridded[:, 2]
-        psi = to_vector(psi)
         psi = np.arctan2((np.sin(psi) / hz), (np.cos(psi) / hy))
-        theta = to_vector(theta)
         theta = np.arctan2((np.sin(theta) / hz), (np.cos(theta) / hx))
 
-        rotation = rotation_matrix_3d(psi, theta, phi, mesh.n_cells)
+        rotation = rotation_matrix_3d(psi, theta, phi)
 
     else:
         raise ValueError("Rotation matrix only implemented for 2D and 3D meshes")
@@ -356,9 +346,8 @@ def rotated_gradient(
     mesh: TreeMesh,
     neighbors: np.ndarray,
     axis: str,
-    phi: float | FloatData,
-    theta: float | FloatData | None = None,
-    psi: float | FloatData | None = None,
+    dip: np.ndarray,
+    direction: np.ndarray,
     forward: bool = True,
 ) -> ssp.csr_matrix:
     """
@@ -367,18 +356,23 @@ def rotated_gradient(
     :param mesh: Input TreeMesh.
     :param neighbors: Cell neighbors array.
     :param axis: Regularization axis.
-    :param phi: Angle in radians for rotation about the y-axis.
-    :param theta: Angle in radians for rotation about the x-axis.
-    :param psi: Angle in radians for rotation about the z-axis
+    :param dip: Angle in radians for rotation from the horizon.
+    :param direction: Angle in radians for rotation about the z-axis.
     :param forward: Whether to use forward or backward difference for
         derivative approximations.
     """
 
     n_cells = mesh.n_cells
+    if any(len(k) != n_cells for k in [dip, direction]):
+        raise ValueError(
+            "Input angle arrays are not the same size as the number of "
+            "cells in the mesh."
+        )
 
-    rot = rotation_matrix(mesh, phi, theta, psi)
+    Rx = rotation_matrix(mesh, 0, dip, 0)
+    Rz = rotation_matrix(mesh, 0, 0, direction)
     normals = get_cell_normals(n_cells, axis, forward)
-    rotated_normals = (rot * normals.T).reshape(n_cells, mesh.dim)
+    rotated_normals = (Rz * (Rx * normals.T)).reshape(n_cells, mesh.dim)
     volumes, neighbors = partial_volumes(mesh, neighbors, rotated_normals)
 
     return gradient_operator(neighbors, volumes, n_cells)

From 8a7249df8a1b35c7a6196fa225f38fc61ee45082 Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Thu, 27 Mar 2025 11:00:27 -0700
Subject: [PATCH 04/29] convert input dip/direction to radians

---
 simpeg_drivers/params.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/simpeg_drivers/params.py b/simpeg_drivers/params.py
index d274c5b46..7c889944d 100644
--- a/simpeg_drivers/params.py
+++ b/simpeg_drivers/params.py
@@ -365,7 +365,8 @@ def gradient_dip(self):
         """Returns the dip of the gradient rotation."""
         if self.gradient_rotation is not None:
             dip_uid = self.gradient_rotation.properties[1]
-            return self.geoh5.get_entity(dip_uid)[0].values
+            dips = self.geoh5.get_entity(dip_uid)[0].values
+            return np.deg2rad(dips)
         return None
 
     @property
@@ -373,7 +374,8 @@ def gradient_direction(self):
         """Returns the direction of the dip of the gradient rotation"""
         if self.gradient_rotation is not None:
             direction_uid = self.gradient_rotation.properties[0]
-            return self.geoh5.get_entity(direction_uid)[0].values
+            directions = self.geoh5.get_entity(direction_uid)[0].values
+            return np.deg2rad(directions)
         return None
 
 

From 4fe068b1582cf87ffcba8795ffa048240d0a70b1 Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Thu, 27 Mar 2025 14:08:23 -0700
Subject: [PATCH 05/29] Simplify rotations

---
 simpeg_drivers/params.py               |   8 +-
 simpeg_drivers/utils/regularization.py | 104 +++++++++++--------------
 2 files changed, 48 insertions(+), 64 deletions(-)

diff --git a/simpeg_drivers/params.py b/simpeg_drivers/params.py
index 7c889944d..80579e832 100644
--- a/simpeg_drivers/params.py
+++ b/simpeg_drivers/params.py
@@ -361,8 +361,8 @@ class BaseInversionOptions(CoreOptions):
     epsilon_cooling_factor: float = 1.2
 
     @property
-    def gradient_dip(self):
-        """Returns the dip of the gradient rotation."""
+    def gradient_dip(self) -> np.ndarray | None:
+        """Gradient dip angle in clockwise radians from horizontal."""
         if self.gradient_rotation is not None:
             dip_uid = self.gradient_rotation.properties[1]
             dips = self.geoh5.get_entity(dip_uid)[0].values
@@ -370,8 +370,8 @@ def gradient_dip(self):
         return None
 
     @property
-    def gradient_direction(self):
-        """Returns the direction of the dip of the gradient rotation"""
+    def gradient_direction(self) -> np.ndarray | None:
+        """Gradient direction angle in clockwise radians from north"""
         if self.gradient_rotation is not None:
             direction_uid = self.gradient_rotation.properties[0]
             directions = self.geoh5.get_entity(direction_uid)[0].values
diff --git a/simpeg_drivers/utils/regularization.py b/simpeg_drivers/utils/regularization.py
index 5aaa95c0e..7c393085e 100644
--- a/simpeg_drivers/utils/regularization.py
+++ b/simpeg_drivers/utils/regularization.py
@@ -126,13 +126,24 @@ def cell_neighbors(mesh: TreeMesh) -> np.ndarray:
     return all_neighbors
 
 
-def rotation_matrix_2d(phi: np.ndarray) -> ssp.csr_matrix:
+def rotate_xz_2d(mesh: TreeMesh, phi: np.ndarray) -> ssp.csr_matrix:
     """
-    Create a 2d rotation matrix for the xz plane.
+    Create a 2d ellipsoidal rotation matrix for the xz plane.
 
-    :param phi: Angle in radians for rotation about the y-axis (xz plan).
+    :param mesh: TreeMesh used to adjust angle of rotation to
+        compensate for cell aspect ratio.
+    :param phi: Angle in radians for clockwise rotation about the
+        y-axis (xz plane).
     """
+
+    if mesh.dim != 2:
+        raise ValueError("Must pass a 2 dimensional mesh.")
+
     n_cells = len(phi)
+    hx = mesh.h_gridded[:, 0]
+    hz = mesh.h_gridded[:, 1]
+    phi = -np.arctan2((np.sin(phi) / hz), (np.cos(phi) / hx))
+
     rza = mkvc(np.c_[np.cos(phi), np.cos(phi)].T)
     rzb = mkvc(np.c_[np.sin(phi), np.zeros(n_cells)].T)
     rzc = mkvc(np.c_[-np.sin(phi), np.zeros(n_cells)].T)
@@ -141,76 +152,49 @@ def rotation_matrix_2d(phi: np.ndarray) -> ssp.csr_matrix:
     return Ry
 
 
-def rotation_matrix_3d(
-    psi: np.ndarray,
-    theta: np.ndarray,
-    phi: np.ndarray,
-) -> ssp.csr_matrix:
+def rotate_yz_3d(mesh: TreeMesh, theta: np.ndarray) -> ssp.csr_matrix:
     """
-    Create a rotation matrix for rotation about the z-axis, y-axis, and x-axis.
+    Create a 3D ellipsoidal rotation matrix for the yz plane.
 
-    :param psi: Angle in radians for rotation about the z-axis.
-    :param theta: Angle in radians for rotation about the x-axis.
-    :param phi: Angle in radians for rotation about the y-axis.
-    :param n_cells: number of cells in the mesh.
+    :param mesh: TreeMesh used to adjust angle of rotation to
+        compensate for cell aspect ratio.
+    :param theta: Angle in radians for clockwise rotation about the
+        x-axis (yz plane).
     """
-    n_cells = len(phi)
-    rxa = mkvc(np.c_[np.ones(n_cells), np.cos(psi), np.cos(psi)].T)
-    rxb = mkvc(np.c_[np.zeros(n_cells), np.sin(psi), np.zeros(n_cells)].T)
-    rxc = mkvc(np.c_[np.zeros(n_cells), -np.sin(psi), np.zeros(n_cells)].T)
-    Rx = ssp.diags([rxb[:-1], rxa, rxc[:-1]], [-1, 0, 1])
 
-    rya = mkvc(np.c_[np.cos(theta), np.ones(n_cells), np.cos(theta)].T)
-    ryb = mkvc(np.c_[np.sin(theta), np.zeros(n_cells), np.zeros(n_cells)].T)
-    ryc = mkvc(np.c_[-np.sin(theta), np.zeros(n_cells), np.zeros(n_cells)].T)
-    Ry = ssp.diags([ryb[:-2], rya, ryc[:-2]], [-2, 0, 2])
+    n_cells = len(theta)
+    hy = mesh.h_gridded[:, 1]
+    hz = mesh.h_gridded[:, 2]
+    theta = -np.arctan2((np.sin(theta) / hz), (np.cos(theta) / hy))
 
-    rza = mkvc(np.c_[np.cos(phi), np.cos(phi), np.ones(n_cells)].T)
-    rzb = mkvc(np.c_[np.sin(phi), np.zeros(n_cells), np.zeros(n_cells)].T)
-    rzc = mkvc(np.c_[-np.sin(phi), np.zeros(n_cells), np.zeros(n_cells)].T)
-    Rz = ssp.diags([rzb[:-1], rza, rzc[:-1]], [-1, 0, 1])
+    rxa = mkvc(np.c_[np.ones(n_cells), np.cos(theta), np.cos(theta)].T)
+    rxb = mkvc(np.c_[np.zeros(n_cells), np.sin(theta), np.zeros(n_cells)].T)
+    rxc = mkvc(np.c_[np.zeros(n_cells), -np.sin(theta), np.zeros(n_cells)].T)
+    Rx = ssp.diags([rxb[:-1], rxa, rxc[:-1]], [-1, 0, 1])
 
-    return Rz * (Ry * Rx)
+    return Rx
 
 
-def rotation_matrix(
-    mesh: TreeMesh,
-    phi: np.ndarray,
-    theta: np.ndarray | None = None,
-    psi: np.ndarray | None = None,
-) -> ssp.csr_matrix:
+def rotate_xy_3d(mesh: TreeMesh, phi: np.ndarray) -> ssp.csr_matrix:
     """
-    Create a rotation matrix for rotation about the z-axis, x-axis, and y-axis.
+    Create a 3D ellipsoidal rotation matrix for the xy plane.
 
-    :param mesh: Input TreeMesh.
-    :param phi: Angle in radians for rotation about the y-axis.
-    :param theta: Angle in radians for rotation about the x-axis.
-    :param psi: Angle in radians for rotation about the z-axis.
+    :param mesh: TreeMesh used to adjust angle of rotation to
+        compensate for cell aspect ratio.
+    :param phi: Angle in radians for clockwise rotation about the
+        z-axis (xy plane).
     """
-
+    n_cells = len(phi)
     hx = mesh.h_gridded[:, 0]
     hy = mesh.h_gridded[:, 1]
-    phi = np.arctan2((np.sin(phi) / hy), (np.cos(phi) / hx))
-
-    if mesh.dim == 2:
-        rotation = rotation_matrix_2d(phi)
-
-    elif mesh.dim == 3:
-        if any(k is None for k in [psi, theta, phi]):
-            raise ValueError(
-                "Input must include 'psi', 'theta', and 'phi' for 3D mesh."
-            )
-
-        hz = mesh.h_gridded[:, 2]
-        psi = np.arctan2((np.sin(psi) / hz), (np.cos(psi) / hy))
-        theta = np.arctan2((np.sin(theta) / hz), (np.cos(theta) / hx))
+    phi = -np.arctan2((np.sin(phi) / hy), (np.cos(phi) / hx))
 
-        rotation = rotation_matrix_3d(psi, theta, phi)
-
-    else:
-        raise ValueError("Rotation matrix only implemented for 2D and 3D meshes")
+    rza = mkvc(np.c_[np.cos(phi), np.cos(phi), np.ones(n_cells)].T)
+    rzb = mkvc(np.c_[np.sin(phi), np.zeros(n_cells), np.zeros(n_cells)].T)
+    rzc = mkvc(np.c_[-np.sin(phi), np.zeros(n_cells), np.zeros(n_cells)].T)
+    Rz = ssp.diags([rzb[:-1], rza, rzc[:-1]], [-1, 0, 1])
 
-    return rotation
+    return Rz
 
 
 def get_cell_normals(n_cells: int, axis: str, outward: bool) -> np.ndarray:
@@ -369,8 +353,8 @@ def rotated_gradient(
             "cells in the mesh."
         )
 
-    Rx = rotation_matrix(mesh, 0, dip, 0)
-    Rz = rotation_matrix(mesh, 0, 0, direction)
+    Rx = rotate_yz_3d(mesh, dip)
+    Rz = rotate_xy_3d(mesh, direction)
     normals = get_cell_normals(n_cells, axis, forward)
     rotated_normals = (Rz * (Rx * normals.T)).reshape(n_cells, mesh.dim)
     volumes, neighbors = partial_volumes(mesh, neighbors, rotated_normals)

From b422d9a81286da2dd84118df0fcfb8336d279d61 Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Thu, 27 Mar 2025 14:45:23 -0700
Subject: [PATCH 06/29] Fix type form for gradient_rotations in gravity/ui.json
 and add gradient_rotations to base joint params

---
 simpeg_drivers/joint/params.py                    | 3 ++-
 simpeg_drivers/potential_fields/gravity/uijson.py | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/simpeg_drivers/joint/params.py b/simpeg_drivers/joint/params.py
index b93f69530..139e304f0 100644
--- a/simpeg_drivers/joint/params.py
+++ b/simpeg_drivers/joint/params.py
@@ -15,7 +15,7 @@
 
 from geoapps_utils.driver.data import BaseData
 from geoh5py.data import FloatData
-from geoh5py.groups import SimPEGGroup, UIJsonGroup
+from geoh5py.groups import PropertyGroup, SimPEGGroup, UIJsonGroup
 from geoh5py.objects import DrapeModel, Octree
 from geoh5py.shared.utils import fetch_active_workspace
 from pydantic import ConfigDict, field_validator, model_validator
@@ -69,6 +69,7 @@ class BaseJointOptions(BaseData):
     length_scale_x: float | FloatData = 1.0
     length_scale_y: float | FloatData | None = 1.0
     length_scale_z: float | FloatData = 1.0
+    gradient_rotation: PropertyGroup | None = None
     s_norm: float | FloatData | None = 0.0
     x_norm: float | FloatData = 2.0
     y_norm: float | FloatData | None = 2.0
diff --git a/simpeg_drivers/potential_fields/gravity/uijson.py b/simpeg_drivers/potential_fields/gravity/uijson.py
index 408b4a922..0d0196b9d 100644
--- a/simpeg_drivers/potential_fields/gravity/uijson.py
+++ b/simpeg_drivers/potential_fields/gravity/uijson.py
@@ -108,7 +108,7 @@ class GravityInversionUIJson(SimPEGDriversUIJson):
     length_scale_x: DataForm
     length_scale_y: DataForm
     length_scale_z: DataForm
-    gradient_rotation: GroupForm
+    gradient_rotation: DataForm
     s_norm: DataForm
     x_norm: DataForm
     y_norm: DataForm

From 58694c69731d2078b14dd0d7e0ab8a3d1de1bebe Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Fri, 28 Mar 2025 08:42:50 -0700
Subject: [PATCH 07/29] Dom's comments

---
 simpeg_drivers-assets/uijson/tipper_inversion.ui.json | 6 ------
 simpeg_drivers/params.py                              | 8 +++++---
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/simpeg_drivers-assets/uijson/tipper_inversion.ui.json b/simpeg_drivers-assets/uijson/tipper_inversion.ui.json
index a865739cb..50d783891 100644
--- a/simpeg_drivers-assets/uijson/tipper_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/tipper_inversion.ui.json
@@ -333,12 +333,6 @@
         "property": "",
         "enabled": true
     },
-    "gradient_phi": {
-        "group": "Regularization",
-        "main": true,
-        "label": "test",
-        "value": 9
-    },
     "gradient_rotation": {
         "group": "Regularization",
         "association": "Cell",
diff --git a/simpeg_drivers/params.py b/simpeg_drivers/params.py
index 80579e832..63eebc144 100644
--- a/simpeg_drivers/params.py
+++ b/simpeg_drivers/params.py
@@ -262,9 +262,7 @@ class BaseInversionOptions(CoreOptions):
     :param length_scale_x: Length scale x.
     :param length_scale_y: Length scale y.
     :param length_scale_z: Length scale z.
-    :param gradient_phi: Gradient rotation about the y-axis
-    :param gradient_theta: Gradient rotation about the x-axis
-    :param gradient_psi: Gradient rotation about the z-axis
+    :param gradient_rotation: Property group for gradient rotation angles.
 
     :param s_norm: S norm.
     :param x_norm: X norm.
@@ -373,8 +371,12 @@ def gradient_dip(self) -> np.ndarray | None:
     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
+
             direction_uid = self.gradient_rotation.properties[0]
             directions = self.geoh5.get_entity(direction_uid)[0].values
+            if self.gradient_rotation.property_group_type == GroupTypeEnum.DIPDIR:
+                directions -= 90.0
             return np.deg2rad(directions)
         return None
 

From d62ef56f2cb0200875ebaa70af5029450346a7c0 Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Fri, 28 Mar 2025 10:38:36 -0700
Subject: [PATCH 08/29] fix strike and dip handling

---
 simpeg_drivers/params.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/simpeg_drivers/params.py b/simpeg_drivers/params.py
index 63eebc144..ed25cbfcc 100644
--- a/simpeg_drivers/params.py
+++ b/simpeg_drivers/params.py
@@ -375,8 +375,8 @@ def gradient_direction(self) -> np.ndarray | None:
 
             direction_uid = self.gradient_rotation.properties[0]
             directions = self.geoh5.get_entity(direction_uid)[0].values
-            if self.gradient_rotation.property_group_type == GroupTypeEnum.DIPDIR:
-                directions -= 90.0
+            if self.gradient_rotation.property_group_type == GroupTypeEnum.STRIKEDIP:
+                directions += 90.0
             return np.deg2rad(directions)
         return None
 

From a9ae058e94deaa03761e2aff0b37dfc5b4b74020 Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Fri, 28 Mar 2025 11:54:07 -0700
Subject: [PATCH 09/29] Add rotated_gradient subtest in gravity runtest

---
 tests/run_tests/driver_grav_test.py | 70 +++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/tests/run_tests/driver_grav_test.py b/tests/run_tests/driver_grav_test.py
index 4475dda68..8f55b1c08 100644
--- a/tests/run_tests/driver_grav_test.py
+++ b/tests/run_tests/driver_grav_test.py
@@ -15,6 +15,7 @@
 
 import numpy as np
 from geoapps_utils.utils.importing import GeoAppsError
+from geoh5py.groups import PropertyGroup
 from geoh5py.workspace import Workspace
 from pytest import raises
 
@@ -188,6 +189,75 @@ def test_gravity_run(
             assert np.all(nan_ind == inactive_ind)
 
 
+def test_rotated_gradient(
+    tmp_path: Path,
+    max_iterations=1,
+    pytest=True,
+):
+    workpath = tmp_path / "inversion_test.ui.geoh5"
+    if pytest:
+        workpath = tmp_path.parent / "test_gravity_fwr_run0" / "inversion_test.ui.geoh5"
+
+    with Workspace(workpath) as geoh5:
+        gz = geoh5.get_entity("Iteration_0_gz")[0]
+        mesh = geoh5.get_entity("mesh")[0]
+
+        inds = (mesh.centroids[:, 0] > -35) & (mesh.centroids[:, 0] < 35)
+        norms = np.ones(mesh.n_cells) * 2
+        norms[inds] = 0
+        gradient_norms = mesh.add_data({"norms": {"values": norms}})
+
+        topography = geoh5.get_entity("topography")[0]
+
+        # Turn some values to nan
+        values = gz.values.copy()
+        values[0] = np.nan
+        gz.values = values
+
+        # Rotate the gradients
+
+        direction, dip = mesh.add_data(
+            {
+                "gradient_direction": {"values": 90 * np.ones(mesh.n_cells)},
+                "gradient_dip": {"values": 45 * np.ones(mesh.n_cells)},
+            }
+        )
+        gradient_rotation = PropertyGroup(
+            parent=mesh,
+            name="gradient_rotation",
+            property_group_type="Dip direction & dip",
+            properties=[direction, dip],
+        )
+
+        # Run the inverse
+        active_cells = ActiveCellsOptions(topography_object=topography)
+        params = GravityInversionOptions(
+            geoh5=geoh5,
+            mesh=mesh,
+            active_cells=active_cells,
+            data_object=gz.parent,
+            starting_model=1e-4,
+            reference_model=0.0,
+            gradient_rotation=gradient_rotation,
+            s_norm=0.0,
+            x_norm=gradient_norms,
+            y_norm=gradient_norms,
+            z_norm=gradient_norms,
+            gradient_type="components",
+            gz_channel=gz,
+            gz_uncertainty=2e-3,
+            lower_bound=0.0,
+            max_global_iterations=max_iterations,
+            initial_beta_ratio=1e-2,
+            percentile=100,
+            store_sensitivities="ram",
+            save_sensitivities=True,
+        )
+        params.write_ui_json(path=tmp_path / "Inv_run.ui.json")
+
+    _ = GravityInversionDriver.start(str(tmp_path / "Inv_run.ui.json"))
+
+
 if __name__ == "__main__":
     # Full run
     test_gravity_fwr_run(

From 79bfb77fe3badd9cf27ab16d0267c52d99c94d72 Mon Sep 17 00:00:00 2001
From: dominiquef <dominiquef@mirageoscience.com>
Date: Fri, 28 Mar 2025 12:41:41 -0700
Subject: [PATCH 10/29] Fix for array assignement. Bring back length scales in
 rotated operator

---
 simpeg_drivers/driver.py | 32 ++++++++++++++++++++++----------
 1 file changed, 22 insertions(+), 10 deletions(-)

diff --git a/simpeg_drivers/driver.py b/simpeg_drivers/driver.py
index 3ccea7391..7b3e226c0 100644
--- a/simpeg_drivers/driver.py
+++ b/simpeg_drivers/driver.py
@@ -42,6 +42,7 @@
     objective_function,
     optimization,
 )
+from simpeg.utils import sdiag
 from simpeg.regularization import BaseRegularization, Sparse
 
 from simpeg_drivers import DRIVER_MAP
@@ -451,7 +452,7 @@ def get_regularization(self):
             )
 
             if is_rotated:
-                neighbors = cell_neighbors(reg.regularization_mesh)
+                neighbors = cell_neighbors(reg.regularization_mesh.mesh)
 
             # Adjustment for 2D versus 3D problems
             comps = "sxz" if "2d" in self.params.inversion_type else "sxyz"
@@ -468,13 +469,8 @@ def get_regularization(self):
                 weight = mapping * getattr(self.models, weight)
                 norm = mapping * getattr(self.models, f"{comp}_norm")
                 if comp in "xyz":
-                    weight = (
-                        getattr(reg.regularization_mesh, f"aveCC2F{avg_comp}") * weight
-                    )
-                    norm = getattr(reg.regularization_mesh, f"aveCC2F{avg_comp}") * norm
-
                     if is_rotated:
-                        Grad = rotated_gradient(
+                        grad_op = rotated_gradient(
                             mesh=reg.regularization_mesh.mesh,
                             neighbors=neighbors,
                             axis=comp,
@@ -482,9 +478,25 @@ def get_regularization(self):
                             direction=self.models.gradient_direction,
                         )
                         setattr(
-                            reg.regularization_mesh.mesh,
-                            f"_stencil_cell_gradient_{comp}",
-                            Grad,
+                            reg.regularization_mesh,
+                            f"_cell_gradient_{comp}",
+                            reg.regularization_mesh.Pac.T
+                            @ (grad_op @ reg.regularization_mesh.Pac),
+                        )
+                        setattr(
+                            reg.regularization_mesh,
+                            f"_aveCC2F{avg_comp}",
+                            sdiag(np.ones(reg.regularization_mesh.n_cells)),
+                        )
+
+                    else:
+                        weight = (
+                            getattr(reg.regularization_mesh, f"aveCC2F{avg_comp}")
+                            * weight
+                        )
+                        norm = (
+                            getattr(reg.regularization_mesh, f"aveCC2F{avg_comp}")
+                            * norm
                         )
 
                 objfct.set_weights(**{comp: weight})

From b3bbaa830a637cfa186efd06c65988dca2fadbda Mon Sep 17 00:00:00 2001
From: dominiquef <dominiquef@mirageoscience.com>
Date: Fri, 28 Mar 2025 12:43:32 -0700
Subject: [PATCH 11/29] Bring back length_scales

---
 simpeg_drivers/utils/regularization.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/simpeg_drivers/utils/regularization.py b/simpeg_drivers/utils/regularization.py
index 7c393085e..c7be3d844 100644
--- a/simpeg_drivers/utils/regularization.py
+++ b/simpeg_drivers/utils/regularization.py
@@ -26,7 +26,11 @@ def cell_neighbors_along_axis(mesh: TreeMesh, axis: str) -> np.ndarray:
     if axis not in "xyz":
         raise ValueError("Argument 'axis' must be one of 'x', 'y', or 'z'.")
 
-    stencil = getattr(mesh, f"cell_gradient_{axis}")
+    if isinstance(mesh, TreeMesh):
+        stencil = getattr(mesh, f"stencil_cell_gradient_{axis}")
+    else:
+        stencil = getattr(mesh, f"cell_gradient_{axis}")
+
     ith_neighbor, jth_neighbor, _ = ssp.find(stencil)
     n_stencils = int(ith_neighbor.shape[0] / 2)
     stencil_indices = jth_neighbor[np.argsort(ith_neighbor)].reshape((n_stencils, 2))
@@ -359,4 +363,5 @@ def rotated_gradient(
     rotated_normals = (Rz * (Rx * normals.T)).reshape(n_cells, mesh.dim)
     volumes, neighbors = partial_volumes(mesh, neighbors, rotated_normals)
 
-    return gradient_operator(neighbors, volumes, n_cells)
+    unit_grad = gradient_operator(neighbors, volumes, n_cells)
+    return sdiag(1 / mesh.h_gridded[:, "xyz".find(axis)]) @ unit_grad

From dacfb8829e667c7ae0bf489e8bf38922a8c1ab66 Mon Sep 17 00:00:00 2001
From: dominiquef <dominiquef@mirageoscience.com>
Date: Fri, 28 Mar 2025 12:47:55 -0700
Subject: [PATCH 12/29] Add rotated gradient run test

---
 .../driver_rotated_gradients_test.py          | 154 ++++++++++++++++++
 1 file changed, 154 insertions(+)
 create mode 100644 tests/run_tests/driver_rotated_gradients_test.py

diff --git a/tests/run_tests/driver_rotated_gradients_test.py b/tests/run_tests/driver_rotated_gradients_test.py
new file mode 100644
index 000000000..bb1735c73
--- /dev/null
+++ b/tests/run_tests/driver_rotated_gradients_test.py
@@ -0,0 +1,154 @@
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+#  Copyright (c) 2025 Mira Geoscience Ltd.                                          '
+#                                                                                   '
+#  This file is part of simpeg-drivers package.                                     '
+#                                                                                   '
+#  simpeg-drivers is distributed under the terms and conditions of the MIT License  '
+#  (see LICENSE file at the root of this source code package).                      '
+#                                                                                   '
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+
+from __future__ import annotations
+
+from pathlib import Path
+from unittest.mock import patch
+
+import numpy as np
+from geoapps_utils.utils.importing import GeoAppsError
+from geoh5py.groups.property_group import PropertyGroup
+from geoh5py.workspace import Workspace
+from pytest import raises
+
+from simpeg_drivers.params import ActiveCellsOptions
+from simpeg_drivers.potential_fields import (
+    GravityForwardOptions,
+    GravityInversionOptions,
+)
+from simpeg_drivers.potential_fields.gravity.driver import (
+    GravityForwardDriver,
+    GravityInversionDriver,
+)
+from simpeg_drivers.utils.testing import check_target, setup_inversion_workspace
+from simpeg_drivers.utils.utils import get_inversion_output
+
+
+# To test the full run and validate the inversion.
+# Move this file out of the test directory and run.
+
+target_run = {"data_norm": 0.0028055269276044915, "phi_d": 8.32e-05, "phi_m": 0.00333}
+
+
+def test_gravity_fwr_run(
+    tmp_path: Path,
+    n_grid_points=2,
+    refinement=(2,),
+):
+    # Run the forward
+    geoh5, _, model, survey, topography = setup_inversion_workspace(
+        tmp_path,
+        background=0.0,
+        anomaly=0.75,
+        n_electrodes=n_grid_points,
+        n_lines=n_grid_points,
+        refinement=refinement,
+        center=(0.0, 0.0, 15.0),
+        flatten=False,
+    )
+
+    active_cells = ActiveCellsOptions(topography_object=topography)
+    params = GravityForwardOptions(
+        geoh5=geoh5,
+        mesh=model.parent,
+        active_cells=active_cells,
+        topography_object=topography,
+        data_object=survey,
+        starting_model=model,
+        gz_channel_bool=True,
+    )
+    fwr_driver = GravityForwardDriver(params)
+    fwr_driver.run()
+
+
+def test_rotated_grad_run(
+    tmp_path: Path,
+    max_iterations=1,
+    pytest=True,
+):
+    workpath = tmp_path / "inversion_test.ui.geoh5"
+    if pytest:
+        workpath = tmp_path.parent / "test_gravity_fwr_run0" / "inversion_test.ui.geoh5"
+
+    with Workspace(workpath) as geoh5:
+        gz = geoh5.get_entity("Iteration_0_gz")[0]
+        orig_gz = gz.values.copy()
+        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
+
+        data_list = mesh.add_data(
+            {
+                "azimuth": {"values": azimuth},
+                "dip": {"values": dip},
+            }
+        )
+        pg = PropertyGroup(
+            mesh, properties=data_list, property_group_type="Dip direction & dip"
+        )
+        topography = geoh5.get_entity("topography")[0]
+
+        # Run the inverse
+        active_cells = ActiveCellsOptions(topography_object=topography)
+        params = GravityInversionOptions(
+            geoh5=geoh5,
+            mesh=mesh,
+            active_cells=active_cells,
+            data_object=gz.parent,
+            gradient_rotation=pg,
+            starting_model=1e-4,
+            reference_model=0.0,
+            s_norm=0.0,
+            x_norm=0.0,
+            y_norm=0.0,
+            z_norm=0.0,
+            gradient_type="components",
+            gz_channel=gz,
+            gz_uncertainty=2e-3,
+            lower_bound=0.0,
+            max_global_iterations=max_iterations,
+            initial_beta_ratio=1e-1,
+            percentile=95,
+            store_sensitivities="ram",
+            save_sensitivities=True,
+        )
+        params.write_ui_json(path=tmp_path / "Inv_run.ui.json")
+
+    driver = GravityInversionDriver.start(str(tmp_path / "Inv_run.ui.json"))
+
+    with Workspace(driver.params.geoh5.h5file) as run_ws:
+        output = get_inversion_output(
+            driver.params.geoh5.h5file, driver.params.out_group.uid
+        )
+        output["data"] = orig_gz
+
+        if pytest:
+            check_target(output, target_run)
+            nan_ind = np.isnan(run_ws.get_entity("Iteration_0_model")[0].values)
+            inactive_ind = run_ws.get_entity("active_cells")[0].values == 0
+            assert np.all(nan_ind == inactive_ind)
+
+
+if __name__ == "__main__":
+    # Full run
+    test_gravity_fwr_run(
+        Path("./"),
+        n_grid_points=10,
+        refinement=(4, 8),
+    )
+
+    test_rotated_grad_run(
+        Path("./"),
+        max_iterations=40,
+        pytest=False,
+    )

From b47a375ecff816e4b4a18d4ffb4de19a98b321cb Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Mon, 31 Mar 2025 09:14:51 -0700
Subject: [PATCH 13/29] Add backward diff gradient regularization if rotated

---
 simpeg_drivers/driver.py | 89 ++++++++++++++++++++++++----------------
 1 file changed, 53 insertions(+), 36 deletions(-)

diff --git a/simpeg_drivers/driver.py b/simpeg_drivers/driver.py
index 7b3e226c0..a324e07e7 100644
--- a/simpeg_drivers/driver.py
+++ b/simpeg_drivers/driver.py
@@ -442,17 +442,28 @@ def get_regularization(self):
             return BaseRegularization(mesh=self.inversion_mesh.mesh)
 
         reg_funcs = []
+        total_reg_funcs = []
         is_rotated = self.params.gradient_rotation is not None
         for mapping in self.mapping:
-            reg = Sparse(
-                self.inversion_mesh.mesh,
-                active_cells=self.models.active_cells,
-                mapping=mapping,
-                reference_model=self.models.reference,
+            reg_funcs.append(
+                Sparse(
+                    self.inversion_mesh.mesh,
+                    active_cells=self.models.active_cells,
+                    mapping=mapping,
+                    reference_model=self.models.reference,
+                )
             )
 
             if is_rotated:
-                neighbors = cell_neighbors(reg.regularization_mesh.mesh)
+                reg_funcs.append(
+                    Sparse(
+                        self.inversion_mesh.mesh,
+                        active_cells=self.models.active_cells,
+                        mapping=mapping,
+                        reference_model=self.models.reference,
+                    )
+                )
+                neighbors = cell_neighbors(reg_funcs[0].regularization_mesh.mesh)
 
             # Adjustment for 2D versus 3D problems
             comps = "sxz" if "2d" in self.params.inversion_type else "sxyz"
@@ -460,42 +471,47 @@ def get_regularization(self):
             weights = ["alpha_s"] + [f"length_scale_{k}" for k in comps[1:]]
 
             for comp, avg_comp, objfct, weight in zip(
-                comps, avg_comps, reg.objfcts, weights
+                comps, avg_comps, reg_funcs[0].objfcts, weights
             ):
                 if getattr(self.models, weight) is None:
-                    setattr(reg, weight, 0.0)
+                    setattr(reg_funcs[0], weight, 0.0)
                     continue
 
                 weight = mapping * getattr(self.models, weight)
                 norm = mapping * getattr(self.models, f"{comp}_norm")
                 if comp in "xyz":
                     if is_rotated:
-                        grad_op = rotated_gradient(
-                            mesh=reg.regularization_mesh.mesh,
-                            neighbors=neighbors,
-                            axis=comp,
-                            dip=self.models.gradient_dip,
-                            direction=self.models.gradient_direction,
-                        )
-                        setattr(
-                            reg.regularization_mesh,
-                            f"_cell_gradient_{comp}",
-                            reg.regularization_mesh.Pac.T
-                            @ (grad_op @ reg.regularization_mesh.Pac),
-                        )
-                        setattr(
-                            reg.regularization_mesh,
-                            f"_aveCC2F{avg_comp}",
-                            sdiag(np.ones(reg.regularization_mesh.n_cells)),
-                        )
-
+                        for reg, forward in zip(reg_funcs, [True, False]):
+                            grad_op = rotated_gradient(
+                                mesh=reg.regularization_mesh.mesh,
+                                neighbors=neighbors,
+                                axis=comp,
+                                dip=self.models.gradient_dip,
+                                direction=self.models.gradient_direction,
+                                forward=forward,
+                            )
+                            setattr(
+                                reg.regularization_mesh,
+                                f"_cell_gradient_{comp}",
+                                reg.regularization_mesh.Pac.T
+                                @ (grad_op @ reg.regularization_mesh.Pac),
+                            )
+                            setattr(
+                                reg.regularization_mesh,
+                                f"_aveCC2F{avg_comp}",
+                                sdiag(np.ones(reg.regularization_mesh.n_cells)),
+                            )
                     else:
                         weight = (
-                            getattr(reg.regularization_mesh, f"aveCC2F{avg_comp}")
+                            getattr(
+                                reg_funcs[0].regularization_mesh, f"aveCC2F{avg_comp}"
+                            )
                             * weight
                         )
                         norm = (
-                            getattr(reg.regularization_mesh, f"aveCC2F{avg_comp}")
+                            getattr(
+                                reg_funcs[0].regularization_mesh, f"aveCC2F{avg_comp}"
+                            )
                             * norm
                         )
 
@@ -503,13 +519,14 @@ def get_regularization(self):
                 objfct.norm = norm
 
             if getattr(self.params, "gradient_type") is not None:
-                setattr(
-                    reg,
-                    "gradient_type",
-                    getattr(self.params, "gradient_type"),
-                )
-
-            reg_funcs.append(reg)
+                for reg in reg_funcs:
+                    setattr(
+                        reg,
+                        "gradient_type",
+                        getattr(self.params, "gradient_type"),
+                    )
+
+            total_reg_funcs += reg_funcs
 
         return objective_function.ComboObjectiveFunction(objfcts=reg_funcs)
 

From 9895c255df093b50ec3a6885d6b89b1ffdd7e87c Mon Sep 17 00:00:00 2001
From: dominiquef <dominiquef@mirageoscience.com>
Date: Tue, 1 Apr 2025 11:06:01 -0700
Subject: [PATCH 14/29] Simplify mechanics for backward reg

---
 simpeg_drivers/driver.py               | 135 +++++++++++++------------
 simpeg_drivers/utils/regularization.py |  38 +++++++
 2 files changed, 108 insertions(+), 65 deletions(-)

diff --git a/simpeg_drivers/driver.py b/simpeg_drivers/driver.py
index a324e07e7..8a1118e6a 100644
--- a/simpeg_drivers/driver.py
+++ b/simpeg_drivers/driver.py
@@ -14,7 +14,7 @@
 from __future__ import annotations
 
 import multiprocessing
-
+from copy import copy
 import sys
 from datetime import datetime, timedelta
 import logging
@@ -42,8 +42,13 @@
     objective_function,
     optimization,
 )
-from simpeg.utils import sdiag
-from simpeg.regularization import BaseRegularization, Sparse
+
+from simpeg.regularization import (
+    BaseRegularization,
+    RegularizationMesh,
+    Sparse,
+    SparseSmoothness,
+)
 
 from simpeg_drivers import DRIVER_MAP
 from simpeg_drivers.components import (
@@ -60,7 +65,7 @@
 )
 from simpeg_drivers.joint.params import BaseJointOptions
 from simpeg_drivers.utils.utils import tile_locations
-from simpeg_drivers.utils.regularization import cell_neighbors, rotated_gradient
+from simpeg_drivers.utils.regularization import cell_neighbors, set_rotated_operators
 
 mlogger = logging.getLogger("distributed")
 mlogger.setLevel(logging.WARNING)
@@ -442,91 +447,91 @@ def get_regularization(self):
             return BaseRegularization(mesh=self.inversion_mesh.mesh)
 
         reg_funcs = []
-        total_reg_funcs = []
         is_rotated = self.params.gradient_rotation is not None
+        neighbors = None
+        backward_mesh = None
+        forward_mesh = None
         for mapping in self.mapping:
-            reg_funcs.append(
-                Sparse(
-                    self.inversion_mesh.mesh,
-                    active_cells=self.models.active_cells,
-                    mapping=mapping,
-                    reference_model=self.models.reference,
-                )
+            reg_func = Sparse(
+                forward_mesh or self.inversion_mesh.mesh,
+                active_cells=self.models.active_cells,
+                mapping=mapping,
+                reference_model=self.models.reference,
             )
 
-            if is_rotated:
-                reg_funcs.append(
-                    Sparse(
-                        self.inversion_mesh.mesh,
-                        active_cells=self.models.active_cells,
-                        mapping=mapping,
-                        reference_model=self.models.reference,
-                    )
+            if is_rotated and not neighbors:
+                backward_mesh = RegularizationMesh(
+                    self.inversion_mesh.mesh, active_cells=self.models.active_cells
                 )
-                neighbors = cell_neighbors(reg_funcs[0].regularization_mesh.mesh)
+                neighbors = cell_neighbors(reg_func.regularization_mesh.mesh)
 
             # Adjustment for 2D versus 3D problems
-            comps = "sxz" if "2d" in self.params.inversion_type else "sxyz"
-            avg_comps = "sxy" if "2d" in self.params.inversion_type else "sxyz"
-            weights = ["alpha_s"] + [f"length_scale_{k}" for k in comps[1:]]
-
-            for comp, avg_comp, objfct, weight in zip(
-                comps, avg_comps, reg_funcs[0].objfcts, weights
+            components = "sxz" if "2d" in self.params.inversion_type else "sxyz"
+            weight_names = ["alpha_s"] + [f"length_scale_{k}" for k in components[1:]]
+            functions = []
+            for comp, weight_name, fun in zip(
+                components, weight_names, reg_func.objfcts
             ):
-                if getattr(self.models, weight) is None:
-                    setattr(reg_funcs[0], weight, 0.0)
+                if getattr(self.models, weight_name) is None:
+                    setattr(reg_funcs, weight_name, 0.0)
+                    functions.append(fun)
                     continue
 
-                weight = mapping * getattr(self.models, weight)
+                weight = mapping * getattr(self.models, weight_name)
                 norm = mapping * getattr(self.models, f"{comp}_norm")
-                if comp in "xyz":
-                    if is_rotated:
-                        for reg, forward in zip(reg_funcs, [True, False]):
-                            grad_op = rotated_gradient(
-                                mesh=reg.regularization_mesh.mesh,
-                                neighbors=neighbors,
-                                axis=comp,
-                                dip=self.models.gradient_dip,
-                                direction=self.models.gradient_direction,
-                                forward=forward,
-                            )
-                            setattr(
-                                reg.regularization_mesh,
-                                f"_cell_gradient_{comp}",
-                                reg.regularization_mesh.Pac.T
-                                @ (grad_op @ reg.regularization_mesh.Pac),
-                            )
-                            setattr(
-                                reg.regularization_mesh,
-                                f"_aveCC2F{avg_comp}",
-                                sdiag(np.ones(reg.regularization_mesh.n_cells)),
-                            )
+
+                if isinstance(fun, SparseSmoothness):
+                    if is_rotated and not forward_mesh:
+                        fun = set_rotated_operators(
+                            fun,
+                            neighbors,
+                            comp,
+                            self.models.gradient_dip,
+                            self.models.gradient_direction,
+                        )
+
                     else:
                         weight = (
                             getattr(
-                                reg_funcs[0].regularization_mesh, f"aveCC2F{avg_comp}"
+                                reg_func.regularization_mesh,
+                                f"aveCC2F{fun.orientation}",
                             )
                             * weight
                         )
                         norm = (
                             getattr(
-                                reg_funcs[0].regularization_mesh, f"aveCC2F{avg_comp}"
+                                reg_func.regularization_mesh,
+                                f"aveCC2F{fun.orientation}",
                             )
                             * norm
                         )
 
-                objfct.set_weights(**{comp: weight})
-                objfct.norm = norm
-
-            if getattr(self.params, "gradient_type") is not None:
-                for reg in reg_funcs:
-                    setattr(
-                        reg,
-                        "gradient_type",
-                        getattr(self.params, "gradient_type"),
-                    )
+                fun.set_weights(**{comp: weight})
+                fun.norm = norm
+                functions.append(fun)
+
+                if isinstance(fun, SparseSmoothness) and is_rotated:
+                    fun.gradient_type = "components"
+                    backward_fun = copy(fun)
+                    setattr(backward_fun, "_regularization_mesh", backward_mesh)
+
+                    # Only do it once for MVI
+                    if not forward_mesh:
+                        backward_fun = set_rotated_operators(
+                            backward_fun,
+                            neighbors,
+                            comp,
+                            self.models.gradient_dip,
+                            self.models.gradient_direction,
+                            forward=False,
+                        )
+                    functions.append(backward_fun)
 
-            total_reg_funcs += reg_funcs
+            # Will avoid recomputing operators if the regularization mesh is the same
+            forward_mesh = reg_func.regularization_mesh
+            reg_func.objfcts = functions
+            reg_func.norms = [fun.norm for fun in functions]
+            reg_funcs.append(reg_func)
 
         return objective_function.ComboObjectiveFunction(objfcts=reg_funcs)
 
diff --git a/simpeg_drivers/utils/regularization.py b/simpeg_drivers/utils/regularization.py
index c7be3d844..d41c7e919 100644
--- a/simpeg_drivers/utils/regularization.py
+++ b/simpeg_drivers/utils/regularization.py
@@ -11,6 +11,7 @@
 import numpy as np
 import scipy.sparse as ssp
 from discretize import TreeMesh
+from simpeg.regularization import SparseSmoothness
 from simpeg.utils import mkvc, sdiag
 
 
@@ -365,3 +366,40 @@ def rotated_gradient(
 
     unit_grad = gradient_operator(neighbors, volumes, n_cells)
     return sdiag(1 / mesh.h_gridded[:, "xyz".find(axis)]) @ unit_grad
+
+
+def set_rotated_operators(
+    function: SparseSmoothness,
+    neighbors: np.ndarray,
+    axis: str,
+    dip: np.ndarray,
+    direction: np.ndarray,
+    forward: bool = True,
+) -> SparseSmoothness:
+    """
+    Calculated rotated gradient operator using partial volumes.
+
+    :param function: Smoothness regularization to change operator for.
+    :param neighbors: Cell neighbors array.
+    :param axis: Regularization axis.
+    :param dip: Angle in radians for rotation from the horizon.
+    :param direction: Angle in radians for rotation about the z-axis.
+    :param forward: Whether to use forward or backward difference for
+        derivative approximations.
+    """
+    grad_op = rotated_gradient(
+        function.regularization_mesh.mesh, neighbors, axis, dip, direction, forward
+    )
+    setattr(
+        function.regularization_mesh,
+        f"_cell_gradient_{function.orientation}",
+        function.regularization_mesh.Pac.T
+        @ (grad_op @ function.regularization_mesh.Pac),
+    )
+    setattr(
+        function.regularization_mesh,
+        f"_aveCC2F{function.orientation}",
+        sdiag(np.ones(function.regularization_mesh.n_cells)),
+    )
+
+    return function

From 373217512d7033a79891ad89991ab68c2dc14271 Mon Sep 17 00:00:00 2001
From: dominiquef <dominiquef@mirageoscience.com>
Date: Tue, 1 Apr 2025 11:45:59 -0700
Subject: [PATCH 15/29] Use deepcopy

---
 simpeg_drivers/driver.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/simpeg_drivers/driver.py b/simpeg_drivers/driver.py
index 8a1118e6a..8c2822a2d 100644
--- a/simpeg_drivers/driver.py
+++ b/simpeg_drivers/driver.py
@@ -14,7 +14,7 @@
 from __future__ import annotations
 
 import multiprocessing
-from copy import copy
+from copy import deepcopy
 import sys
 from datetime import datetime, timedelta
 import logging
@@ -512,7 +512,7 @@ def get_regularization(self):
 
                 if isinstance(fun, SparseSmoothness) and is_rotated:
                     fun.gradient_type = "components"
-                    backward_fun = copy(fun)
+                    backward_fun = deepcopy(fun)
                     setattr(backward_fun, "_regularization_mesh", backward_mesh)
 
                     # Only do it once for MVI

From 0dfa19bc394df01a1406eab1e42df3dfefc33ceb Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Tue, 1 Apr 2025 12:02:03 -0700
Subject: [PATCH 16/29] refactor neighbor/corner finding and add unit test

---
 simpeg_drivers/utils/regularization.py | 98 ++++++++++++++------------
 tests/utils_regularization_test.py     | 92 ++++++++++++++++++++++++
 2 files changed, 145 insertions(+), 45 deletions(-)
 create mode 100644 tests/utils_regularization_test.py

diff --git a/simpeg_drivers/utils/regularization.py b/simpeg_drivers/utils/regularization.py
index c7be3d844..6ea98f84f 100644
--- a/simpeg_drivers/utils/regularization.py
+++ b/simpeg_drivers/utils/regularization.py
@@ -38,52 +38,32 @@ def cell_neighbors_along_axis(mesh: TreeMesh, axis: str) -> np.ndarray:
     return np.sort(stencil_indices, axis=1)
 
 
-def cell_neighbors(mesh: TreeMesh) -> np.ndarray:
-    """Find all cell neighbors in a TreeMesh."""
-
-    x_neighbors = cell_neighbors_along_axis(mesh, "x")
-    x_neighbors_backward = np.fliplr(x_neighbors)
-    y_neighbors = cell_neighbors_along_axis(mesh, "y")
-    y_neighbors_backward = np.fliplr(y_neighbors)
-    max_index = np.max([x_neighbors.max(), y_neighbors.max()])
-    if mesh.dim == 3:
-        z_neighbors = cell_neighbors_along_axis(mesh, "z")
-        z_neighbors_backward = np.fliplr(z_neighbors)
-        max_index = np.max([max_index, z_neighbors.max()])
-
+def collect_all_neighbors(neighbors, neighbors_backwards, corners, corners_backwards):
     all_neighbors = []  # Store
-    x_adjacent = np.ones(max_index + 1, dtype="int") * -1
-    y_adjacent = np.ones(max_index + 1, dtype="int") * -1
-    x_adjacent_backward = np.ones(max_index + 1, dtype="int") * -1
-    y_adjacent_backward = np.ones(max_index + 1, dtype="int") * -1
-
-    x_adjacent[y_neighbors[:, 0]] = y_neighbors[:, 1]
-    y_adjacent[x_neighbors[:, 1]] = x_neighbors[:, 0]
 
-    x_adjacent_backward[y_neighbors_backward[:, 0]] = y_neighbors_backward[:, 1]
-    y_adjacent_backward[x_neighbors_backward[:, 1]] = x_neighbors_backward[:, 0]
+    all_neighbors += [neighbors[0]]
+    all_neighbors += [neighbors[1]]
 
-    all_neighbors += [x_neighbors]
-    all_neighbors += [y_neighbors]
+    all_neighbors += [np.c_[neighbors[0][:, 0], corners[0][neighbors[0][:, 1]]]]
+    all_neighbors += [np.c_[neighbors[0][:, 1], corners[0][neighbors[0][:, 0]]]]
 
-    all_neighbors += [np.c_[x_neighbors[:, 0], x_adjacent[x_neighbors[:, 1]]]]
-    all_neighbors += [np.c_[x_neighbors[:, 1], x_adjacent[x_neighbors[:, 0]]]]
-
-    all_neighbors += [np.c_[y_adjacent[y_neighbors[:, 0]], y_neighbors[:, 1]]]
-    all_neighbors += [np.c_[y_adjacent[y_neighbors[:, 1]], y_neighbors[:, 0]]]
+    all_neighbors += [np.c_[corners[1][neighbors[1][:, 0]], neighbors[1][:, 1]]]
+    all_neighbors += [np.c_[corners[1][neighbors[1][:, 1]], neighbors[1][:, 0]]]
 
     # Repeat backward for Treemesh
-    all_neighbors += [x_neighbors_backward]
-    all_neighbors += [y_neighbors_backward]
+    all_neighbors += [neighbors_backwards[0]]
+    all_neighbors += [neighbors_backwards[1]]
 
     all_neighbors += [
         np.c_[
-            x_neighbors_backward[:, 0], x_adjacent_backward[x_neighbors_backward[:, 1]]
+            neighbors_backwards[0][:, 0],
+            corners_backwards[0][neighbors_backwards[0][:, 1]],
         ]
     ]
     all_neighbors += [
         np.c_[
-            x_neighbors_backward[:, 1], x_adjacent_backward[x_neighbors_backward[:, 0]]
+            neighbors_backwards[0][:, 1],
+            corners_backwards[0][neighbors_backwards[0][:, 0]],
         ]
     ]
 
@@ -97,25 +77,20 @@ def cell_neighbors(mesh: TreeMesh) -> np.ndarray:
     ]
 
     # Use all the neighbours on the xy plane to find neighbours in z
-    if mesh.dim == 3:
+    if len(neighbors) == 3:
         all_neighbors_z = []
-        z_adjacent = np.ones(max_index + 1, dtype="int") * -1
-        z_adjacent_backward = np.ones(max_index + 1, dtype="int") * -1
-
-        z_adjacent[z_neighbors[:, 0]] = z_neighbors[:, 1]
-        z_adjacent_backward[z_neighbors_backward[:, 0]] = z_neighbors_backward[:, 1]
 
-        all_neighbors_z += [z_neighbors]
-        all_neighbors_z += [z_neighbors_backward]
+        all_neighbors_z += [neighbors[2]]
+        all_neighbors_z += [neighbors_backwards[2]]
 
-        all_neighbors_z += [np.c_[all_neighbors[:, 0], z_adjacent[all_neighbors[:, 1]]]]
-        all_neighbors_z += [np.c_[all_neighbors[:, 1], z_adjacent[all_neighbors[:, 0]]]]
+        all_neighbors_z += [np.c_[all_neighbors[:, 0], corners[2][all_neighbors[:, 1]]]]
+        all_neighbors_z += [np.c_[all_neighbors[:, 1], corners[2][all_neighbors[:, 0]]]]
 
         all_neighbors_z += [
-            np.c_[all_neighbors[:, 0], z_adjacent_backward[all_neighbors[:, 1]]]
+            np.c_[all_neighbors[:, 0], corners_backwards[2][all_neighbors[:, 1]]]
         ]
         all_neighbors_z += [
-            np.c_[all_neighbors[:, 1], z_adjacent_backward[all_neighbors[:, 0]]]
+            np.c_[all_neighbors[:, 1], corners_backwards[2][all_neighbors[:, 0]]]
         ]
 
         # Stack all and keep only unique pairs
@@ -130,6 +105,39 @@ def cell_neighbors(mesh: TreeMesh) -> np.ndarray:
     return all_neighbors
 
 
+def cell_adjacent(neighbors: list[np.ndarray]) -> list[np.ndarray]:
+    """Find all cell corners from cell neighbor array."""
+
+    dim = len(neighbors)
+    max_index = np.max(neighbors)
+    corners = -1 * np.ones((dim, max_index + 1), dtype="int")
+
+    corners[0, neighbors[1][:, 0]] = neighbors[1][:, 1]
+    corners[1, neighbors[0][:, 1]] = neighbors[0][:, 0]
+    if dim == 3:
+        corners[2, neighbors[2][:, 0]] = neighbors[2][:, 1]
+
+    return [np.array(k) for k in corners.tolist()]
+
+
+def cell_neighbors(mesh: TreeMesh) -> np.ndarray:
+    """Find all cell neighbors in a TreeMesh."""
+
+    neighbors = []
+    neighbors.append(cell_neighbors_along_axis(mesh, "x"))
+    neighbors.append(cell_neighbors_along_axis(mesh, "y"))
+    if mesh.dim == 3:
+        neighbors.append(cell_neighbors_along_axis(mesh, "z"))
+
+    neighbors_backwards = [np.fliplr(k) for k in neighbors]
+    corners = cell_adjacent(neighbors)
+    corners_backwards = cell_adjacent(neighbors_backwards)
+
+    return collect_all_neighbors(
+        neighbors, neighbors_backwards, corners, corners_backwards
+    )
+
+
 def rotate_xz_2d(mesh: TreeMesh, phi: np.ndarray) -> ssp.csr_matrix:
     """
     Create a 2d ellipsoidal rotation matrix for the xz plane.
diff --git a/tests/utils_regularization_test.py b/tests/utils_regularization_test.py
new file mode 100644
index 000000000..d0f0fefe1
--- /dev/null
+++ b/tests/utils_regularization_test.py
@@ -0,0 +1,92 @@
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+#  Copyright (c) 2025 Mira Geoscience Ltd.                                          '
+#                                                                                   '
+#  This file is part of simpeg-drivers package.                                     '
+#                                                                                   '
+#  simpeg-drivers is distributed under the terms and conditions of the MIT License  '
+#  (see LICENSE file at the root of this source code package).                      '
+#                                                                                   '
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+
+import numpy as np
+from discretize import TreeMesh
+
+from simpeg_drivers.utils.regularization import (
+    cell_adjacent,
+    cell_neighbors_along_axis,
+    collect_all_neighbors,
+)
+
+
+def get_mesh():
+    mesh = TreeMesh(h=[[10.0] * 4, [10.0] * 4, [10.0] * 4], diagonal_balance=False)
+    mesh.refine(2)
+    return mesh
+
+
+def test_cell_neighbors_along_axis():
+    mesh = get_mesh()
+    centers = mesh.cell_centers
+    neighbors = cell_neighbors_along_axis(mesh, "x")
+    assert np.allclose(centers[7], [15.0, 15.0, 15.0])
+    assert np.allclose(
+        centers[neighbors[neighbors[:, 0] == 7][0][1]], [25.0, 15.0, 15.0]
+    )
+    assert np.allclose(
+        centers[neighbors[neighbors[:, 1] == 7][0][0]], [5.0, 15.0, 15.0]
+    )
+    neighbors = cell_neighbors_along_axis(mesh, "y")
+    assert np.allclose(
+        centers[neighbors[neighbors[:, 0] == 7][0][1]], [15.0, 25.0, 15.0]
+    )
+    assert np.allclose(
+        centers[neighbors[neighbors[:, 1] == 7][0][0]], [15.0, 5.0, 15.0]
+    )
+    neighbors = cell_neighbors_along_axis(mesh, "z")
+    assert np.allclose(
+        centers[neighbors[neighbors[:, 0] == 7][0][1]], [15.0, 15.0, 25.0]
+    )
+    assert np.allclose(
+        centers[neighbors[neighbors[:, 1] == 7][0][0]], [15.0, 15.0, 5.0]
+    )
+
+
+def test_collect_all_neighbors():
+    mesh = get_mesh()
+    centers = mesh.cell_centers
+    neighbors = [cell_neighbors_along_axis(mesh, k) for k in "xyz"]
+    neighbors_bck = [np.fliplr(k) for k in neighbors]
+    corners = cell_adjacent(neighbors)
+    corners_bck = cell_adjacent(neighbors_bck)
+    all_neighbors = collect_all_neighbors(
+        neighbors, neighbors_bck, corners, corners_bck
+    )
+    assert np.allclose(centers[7], [15.0, 15.0, 15.0])
+    neighbor_centers = centers[all_neighbors[all_neighbors[:, 0] == 7][:, 1]].tolist()
+    assert [5, 5, 5] in neighbor_centers
+    assert [15, 5, 5] in neighbor_centers
+    assert [25, 5, 5] in neighbor_centers
+    assert [5, 15, 5] in neighbor_centers
+    assert [15, 15, 5] in neighbor_centers
+    assert [25, 15, 5] in neighbor_centers
+    assert [5, 25, 5] in neighbor_centers
+    assert [15, 25, 5] in neighbor_centers
+    assert [25, 25, 5] in neighbor_centers
+    assert [5, 5, 15] in neighbor_centers
+    assert [15, 5, 15] in neighbor_centers
+    assert [25, 5, 15] in neighbor_centers
+    assert [5, 15, 15] in neighbor_centers
+    assert [25, 15, 15] in neighbor_centers
+    assert [5, 25, 15] in neighbor_centers
+    assert [15, 25, 15] in neighbor_centers
+    assert [25, 25, 15] in neighbor_centers
+    assert [5, 5, 25] in neighbor_centers
+    assert [15, 5, 25] in neighbor_centers
+    assert [25, 5, 25] in neighbor_centers
+    assert [5, 15, 25] in neighbor_centers
+    assert [15, 15, 25] in neighbor_centers
+    assert [25, 15, 25] in neighbor_centers
+    assert [5, 25, 25] in neighbor_centers
+    assert [15, 25, 25] in neighbor_centers
+    assert [25, 25, 25] in neighbor_centers
+    assert [15, 15, 15] not in neighbor_centers

From 6271153046a142853138680acdf61a7ba4f88f94 Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Tue, 1 Apr 2025 12:09:16 -0700
Subject: [PATCH 17/29] Improve docstring/typing for new methods

---
 simpeg_drivers/utils/regularization.py | 41 ++++++++++++++++++--------
 1 file changed, 29 insertions(+), 12 deletions(-)

diff --git a/simpeg_drivers/utils/regularization.py b/simpeg_drivers/utils/regularization.py
index 6ea98f84f..9e359ed76 100644
--- a/simpeg_drivers/utils/regularization.py
+++ b/simpeg_drivers/utils/regularization.py
@@ -38,17 +38,30 @@ def cell_neighbors_along_axis(mesh: TreeMesh, axis: str) -> np.ndarray:
     return np.sort(stencil_indices, axis=1)
 
 
-def collect_all_neighbors(neighbors, neighbors_backwards, corners, corners_backwards):
+def collect_all_neighbors(
+    neighbors: list[np.ndarray],
+    neighbors_backwards: list[np.ndarray],
+    adjacent: list[np.ndarray],
+    adjacent_backwards: list[np.ndarray],
+) -> np.ndarray:
+    """
+    Collect all neighbors for cells in the mesh.
+
+    :param neighbors: Direct neighbors in each principle axes.
+    :param neighbors_backwards: Direct neighbors in reverse order.
+    :param adjacent: Adjacent neighbors (corners).
+    :param adjacent_backwards: Adjacent neighbors in reverse order.
+    """
     all_neighbors = []  # Store
 
     all_neighbors += [neighbors[0]]
     all_neighbors += [neighbors[1]]
 
-    all_neighbors += [np.c_[neighbors[0][:, 0], corners[0][neighbors[0][:, 1]]]]
-    all_neighbors += [np.c_[neighbors[0][:, 1], corners[0][neighbors[0][:, 0]]]]
+    all_neighbors += [np.c_[neighbors[0][:, 0], adjacent[0][neighbors[0][:, 1]]]]
+    all_neighbors += [np.c_[neighbors[0][:, 1], adjacent[0][neighbors[0][:, 0]]]]
 
-    all_neighbors += [np.c_[corners[1][neighbors[1][:, 0]], neighbors[1][:, 1]]]
-    all_neighbors += [np.c_[corners[1][neighbors[1][:, 1]], neighbors[1][:, 0]]]
+    all_neighbors += [np.c_[adjacent[1][neighbors[1][:, 0]], neighbors[1][:, 1]]]
+    all_neighbors += [np.c_[adjacent[1][neighbors[1][:, 1]], neighbors[1][:, 0]]]
 
     # Repeat backward for Treemesh
     all_neighbors += [neighbors_backwards[0]]
@@ -57,13 +70,13 @@ def collect_all_neighbors(neighbors, neighbors_backwards, corners, corners_backw
     all_neighbors += [
         np.c_[
             neighbors_backwards[0][:, 0],
-            corners_backwards[0][neighbors_backwards[0][:, 1]],
+            adjacent_backwards[0][neighbors_backwards[0][:, 1]],
         ]
     ]
     all_neighbors += [
         np.c_[
             neighbors_backwards[0][:, 1],
-            corners_backwards[0][neighbors_backwards[0][:, 0]],
+            adjacent_backwards[0][neighbors_backwards[0][:, 0]],
         ]
     ]
 
@@ -83,14 +96,18 @@ def collect_all_neighbors(neighbors, neighbors_backwards, corners, corners_backw
         all_neighbors_z += [neighbors[2]]
         all_neighbors_z += [neighbors_backwards[2]]
 
-        all_neighbors_z += [np.c_[all_neighbors[:, 0], corners[2][all_neighbors[:, 1]]]]
-        all_neighbors_z += [np.c_[all_neighbors[:, 1], corners[2][all_neighbors[:, 0]]]]
+        all_neighbors_z += [
+            np.c_[all_neighbors[:, 0], adjacent[2][all_neighbors[:, 1]]]
+        ]
+        all_neighbors_z += [
+            np.c_[all_neighbors[:, 1], adjacent[2][all_neighbors[:, 0]]]
+        ]
 
         all_neighbors_z += [
-            np.c_[all_neighbors[:, 0], corners_backwards[2][all_neighbors[:, 1]]]
+            np.c_[all_neighbors[:, 0], adjacent_backwards[2][all_neighbors[:, 1]]]
         ]
         all_neighbors_z += [
-            np.c_[all_neighbors[:, 1], corners_backwards[2][all_neighbors[:, 0]]]
+            np.c_[all_neighbors[:, 1], adjacent_backwards[2][all_neighbors[:, 0]]]
         ]
 
         # Stack all and keep only unique pairs
@@ -106,7 +123,7 @@ def collect_all_neighbors(neighbors, neighbors_backwards, corners, corners_backw
 
 
 def cell_adjacent(neighbors: list[np.ndarray]) -> list[np.ndarray]:
-    """Find all cell corners from cell neighbor array."""
+    """Find all adjacent cells (corners) from cell neighbor array."""
 
     dim = len(neighbors)
     max_index = np.max(neighbors)

From e25f99d8cf95730a8ca92f923cef5213ed966937 Mon Sep 17 00:00:00 2001
From: dominiquef <dominiquef@mirageoscience.com>
Date: Tue, 1 Apr 2025 13:34:35 -0700
Subject: [PATCH 18/29] Bring back gradient_type temporarly

---
 simpeg_drivers/driver.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/simpeg_drivers/driver.py b/simpeg_drivers/driver.py
index 8c2822a2d..d6d06adc4 100644
--- a/simpeg_drivers/driver.py
+++ b/simpeg_drivers/driver.py
@@ -533,6 +533,15 @@ def get_regularization(self):
             reg_func.norms = [fun.norm for fun in functions]
             reg_funcs.append(reg_func)
 
+        # TODO - To be deprcated on GEOPY-2109
+        if getattr(self.params, "gradient_type") is not None:
+            for reg in reg_funcs:
+                setattr(
+                    reg,
+                    "gradient_type",
+                    getattr(self.params, "gradient_type"),
+                )
+
         return objective_function.ComboObjectiveFunction(objfcts=reg_funcs)
 
     def get_tiles(self):

From 27934f8260f29591e2eec48ef705452d1f41ff29 Mon Sep 17 00:00:00 2001
From: dominiquef <dominiquef@mirageoscience.com>
Date: Tue, 1 Apr 2025 13:57:37 -0700
Subject: [PATCH 19/29] Clean up of inversion model module. Allow to not trim
 the dip azimut models

---
 simpeg_drivers/components/models.py | 123 +++++++++++++---------------
 simpeg_drivers/driver.py            |   2 +-
 2 files changed, 58 insertions(+), 67 deletions(-)

diff --git a/simpeg_drivers/components/models.py b/simpeg_drivers/components/models.py
index 745a0480f..29efc1fd7 100644
--- a/simpeg_drivers/components/models.py
+++ b/simpeg_drivers/components/models.py
@@ -29,6 +29,25 @@
     from simpeg_drivers.driver import InversionDriver
 
 
+MODEL_TYPES = [
+    "starting",
+    "reference",
+    "lower_bound",
+    "upper_bound",
+    "conductivity",
+    "alpha_s",
+    "length_scale_x",
+    "length_scale_y",
+    "length_scale_z",
+    "gradient_dip",
+    "gradient_direction",
+    "s_norm",
+    "x_norm",
+    "y_norm",
+    "z_norm",
+]
+
+
 class InversionModelCollection:
     """
     Collection of inversion models.
@@ -39,22 +58,6 @@ class InversionModelCollection:
 
     """
 
-    model_types = [
-        "starting",
-        "reference",
-        "lower_bound",
-        "upper_bound",
-        "conductivity",
-        "alpha_s",
-        "length_scale_x",
-        "length_scale_y",
-        "length_scale_z",
-        "s_norm",
-        "x_norm",
-        "y_norm",
-        "z_norm",
-    ]
-
     def __init__(self, driver: InversionDriver):
         """
         :param driver: Parental InversionDriver class.
@@ -62,27 +65,34 @@ def __init__(self, driver: InversionDriver):
         self._active_cells: np.ndarray | None = None
         self._driver = driver
         self.is_sigma = self.driver.params.physical_property == "conductivity"
-        self.is_vector = (
+        is_vector = (
             True if self.driver.params.inversion_type == "magnetic vector" else False
         )
-        self.n_blocks = (
-            3 if self.driver.params.inversion_type == "magnetic vector" else 1
+        self._starting = InversionModel(driver, "starting", is_vector=is_vector)
+        self._reference = InversionModel(driver, "reference", is_vector=is_vector)
+        self._lower_bound = InversionModel(driver, "lower_bound", is_vector=is_vector)
+        self._upper_bound = InversionModel(driver, "upper_bound", is_vector=is_vector)
+        self._conductivity = InversionModel(driver, "conductivity", is_vector=is_vector)
+        self._alpha_s = InversionModel(driver, "alpha_s", is_vector=is_vector)
+        self._length_scale_x = InversionModel(
+            driver, "length_scale_x", is_vector=is_vector
+        )
+        self._length_scale_y = InversionModel(
+            driver, "length_scale_y", is_vector=is_vector
+        )
+        self._length_scale_z = InversionModel(
+            driver, "length_scale_z", is_vector=is_vector
+        )
+        self._gradient_dip = InversionModel(
+            driver, "gradient_dip", trim_active_cells=False
         )
-        self._starting = InversionModel(driver, "starting")
-        self._reference = InversionModel(driver, "reference")
-        self._lower_bound = InversionModel(driver, "lower_bound")
-        self._upper_bound = InversionModel(driver, "upper_bound")
-        self._conductivity = InversionModel(driver, "conductivity")
-        self._alpha_s = InversionModel(driver, "alpha_s")
-        self._length_scale_x = InversionModel(driver, "length_scale_x")
-        self._length_scale_y = InversionModel(driver, "length_scale_y")
-        self._length_scale_z = InversionModel(driver, "length_scale_z")
-        self._gradient_dip = InversionModel(driver, "gradient_dip")
-        self._gradient_direction = InversionModel(driver, "gradient_direction")
-        self._s_norm = InversionModel(driver, "s_norm")
-        self._x_norm = InversionModel(driver, "x_norm")
-        self._y_norm = InversionModel(driver, "y_norm")
-        self._z_norm = InversionModel(driver, "z_norm")
+        self._gradient_direction = InversionModel(
+            driver, "gradient_direction", trim_active_cells=False
+        )
+        self._s_norm = InversionModel(driver, "s_norm", is_vector=is_vector)
+        self._x_norm = InversionModel(driver, "x_norm", is_vector=is_vector)
+        self._y_norm = InversionModel(driver, "y_norm", is_vector=is_vector)
+        self._z_norm = InversionModel(driver, "z_norm", is_vector=is_vector)
 
     @property
     def n_active(self) -> int:
@@ -307,7 +317,7 @@ def z_norm(self) -> np.ndarray | None:
     def _model_method_wrapper(self, method, name=None, **kwargs):
         """wraps individual model's specific method and applies in loop over model types."""
         returned_items = {}
-        for mtype in self.model_types:
+        for mtype in MODEL_TYPES:
             model = getattr(self, f"_{mtype}")
             if model.model is not None:
                 f = getattr(model, method)
@@ -364,43 +374,24 @@ class InversionModel:
     remove_air: Use active cells vector to remove air cells from model.
     """
 
-    model_types = [
-        "starting",
-        "reference",
-        "lower_bound",
-        "upper_bound",
-        "conductivity",
-        "alpha_s",
-        "length_scale_x",
-        "length_scale_y",
-        "length_scale_z",
-        "gradient_dip",
-        "gradient_direction",
-        "s_norm",
-        "x_norm",
-        "y_norm",
-        "z_norm",
-    ]
-
     def __init__(
         self,
         driver: InversionDriver,
         model_type: str,
+        is_vector: bool = False,
+        trim_active_cells: bool = True,
     ):
         """
         :param driver: InversionDriver object.
-        :param model_type: Type of inversion model, can be any of "starting", "reference",
-            "lower_bound", "upper_bound".
+        :param model_type: Type of inversion model, can be any of MODEL_TYPES.
+        :param is_vector: If True, model is a vector.
+        :param trim_active_cells: If True, remove air cells from model.
         """
         self.driver = driver
         self.model_type = model_type
         self.model: np.ndarray | None = None
-        self.is_vector = (
-            True if self.driver.params.inversion_type == "magnetic vector" else False
-        )
-        self.n_blocks = (
-            3 if self.driver.params.inversion_type == "magnetic vector" else 1
-        )
+        self.is_vector = is_vector
+        self.trim_active_cells = trim_active_cells
         self._initialize()
 
     def _initialize(self):
@@ -452,7 +443,7 @@ def _initialize(self):
                 and self.is_vector
                 and model.shape[0] == self.driver.inversion_mesh.n_cells
             ):
-                model = np.tile(model, self.n_blocks)
+                model = np.tile(model, 3 if self.is_vector else 1)
 
         if model is not None:
             self.model = mkvc(model)
@@ -461,8 +452,8 @@ def _initialize(self):
     def remove_air(self, active_cells):
         """Use active cells vector to remove air cells from model"""
 
-        if self.model is not None:
-            self.model = self.model[np.tile(active_cells, self.n_blocks)]
+        if self.model is not None and self.trim_active_cells:
+            self.model = self.model[np.tile(active_cells, 3 if self.is_vector else 1)]
 
     def permute_2_octree(self) -> np.ndarray | None:
         """
@@ -604,7 +595,7 @@ def model_type(self):
 
     @model_type.setter
     def model_type(self, v):
-        if v not in self.model_types:
-            msg = f"Invalid model_type: {v}. Must be one of {(*self.model_types,)}."
+        if v not in MODEL_TYPES:
+            msg = f"Invalid model_type: {v}. Must be one of {(*MODEL_TYPES,)}."
             raise ValueError(msg)
         self._model_type = v
diff --git a/simpeg_drivers/driver.py b/simpeg_drivers/driver.py
index d6d06adc4..35b93a094 100644
--- a/simpeg_drivers/driver.py
+++ b/simpeg_drivers/driver.py
@@ -459,7 +459,7 @@ def get_regularization(self):
                 reference_model=self.models.reference,
             )
 
-            if is_rotated and not neighbors:
+            if is_rotated and neighbors is None:
                 backward_mesh = RegularizationMesh(
                     self.inversion_mesh.mesh, active_cells=self.models.active_cells
                 )

From f63c829546f4e58645b258b657ccd0ec845d975d Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Tue, 1 Apr 2025 14:12:45 -0700
Subject: [PATCH 20/29] vstack neighbors before asking for max

---
 simpeg_drivers/utils/regularization.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/simpeg_drivers/utils/regularization.py b/simpeg_drivers/utils/regularization.py
index 9e359ed76..d19751bdf 100644
--- a/simpeg_drivers/utils/regularization.py
+++ b/simpeg_drivers/utils/regularization.py
@@ -126,7 +126,7 @@ def cell_adjacent(neighbors: list[np.ndarray]) -> list[np.ndarray]:
     """Find all adjacent cells (corners) from cell neighbor array."""
 
     dim = len(neighbors)
-    max_index = np.max(neighbors)
+    max_index = np.max(np.vstack(neighbors))
     corners = -1 * np.ones((dim, max_index + 1), dtype="int")
 
     corners[0, neighbors[1][:, 0]] = neighbors[1][:, 1]

From 0557fa28d68261d9686fbc8b542325888cd003af Mon Sep 17 00:00:00 2001
From: dominiquef <dominiquef@mirageoscience.com>
Date: Tue, 1 Apr 2025 14:29:34 -0700
Subject: [PATCH 21/29] Fix for MVI

---
 simpeg_drivers/driver.py | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/simpeg_drivers/driver.py b/simpeg_drivers/driver.py
index 35b93a094..25be67d3f 100644
--- a/simpeg_drivers/driver.py
+++ b/simpeg_drivers/driver.py
@@ -454,7 +454,7 @@ def get_regularization(self):
         for mapping in self.mapping:
             reg_func = Sparse(
                 forward_mesh or self.inversion_mesh.mesh,
-                active_cells=self.models.active_cells,
+                active_cells=self.models.active_cells if forward_mesh is None else None,
                 mapping=mapping,
                 reference_model=self.models.reference,
             )
@@ -481,14 +481,15 @@ def get_regularization(self):
                 norm = mapping * getattr(self.models, f"{comp}_norm")
 
                 if isinstance(fun, SparseSmoothness):
-                    if is_rotated and not forward_mesh:
-                        fun = set_rotated_operators(
-                            fun,
-                            neighbors,
-                            comp,
-                            self.models.gradient_dip,
-                            self.models.gradient_direction,
-                        )
+                    if is_rotated:
+                        if forward_mesh is None:
+                            fun = set_rotated_operators(
+                                fun,
+                                neighbors,
+                                comp,
+                                self.models.gradient_dip,
+                                self.models.gradient_direction,
+                            )
 
                     else:
                         weight = (

From f5d71737ffdb32932c11052439a295ffa3a0f007 Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Wed, 2 Apr 2025 08:58:53 -0700
Subject: [PATCH 22/29] remove gradient_rotations from 2d/batch2d ui.json files

---
 simpeg_drivers-assets/uijson/direct_current_2d_inversion.ui.json | 1 +
 .../uijson/direct_current_batch2d_inversion.ui.json              | 1 +
 .../uijson/induced_polarization_2d_inversion.ui.json             | 1 +
 .../uijson/induced_polarization_batch2d_inversion.ui.json        | 1 +
 4 files changed, 4 insertions(+)

diff --git a/simpeg_drivers-assets/uijson/direct_current_2d_inversion.ui.json b/simpeg_drivers-assets/uijson/direct_current_2d_inversion.ui.json
index b08f2469f..2d6357f44 100644
--- a/simpeg_drivers-assets/uijson/direct_current_2d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/direct_current_2d_inversion.ui.json
@@ -285,6 +285,7 @@
         "label": "Gradient rotation",
         "optional": true,
         "enabled": false,
+        "visible": false,
         "parent": "mesh",
         "value": ""
     },
diff --git a/simpeg_drivers-assets/uijson/direct_current_batch2d_inversion.ui.json b/simpeg_drivers-assets/uijson/direct_current_batch2d_inversion.ui.json
index d834dfc7c..6337065a4 100644
--- a/simpeg_drivers-assets/uijson/direct_current_batch2d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/direct_current_batch2d_inversion.ui.json
@@ -267,6 +267,7 @@
         "label": "Gradient rotation",
         "optional": true,
         "enabled": false,
+        "visible": false,
         "parent": "mesh",
         "value": ""
     },
diff --git a/simpeg_drivers-assets/uijson/induced_polarization_2d_inversion.ui.json b/simpeg_drivers-assets/uijson/induced_polarization_2d_inversion.ui.json
index 04fab5563..940c72e0e 100644
--- a/simpeg_drivers-assets/uijson/induced_polarization_2d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/induced_polarization_2d_inversion.ui.json
@@ -294,6 +294,7 @@
         "label": "Gradient rotation",
         "optional": true,
         "enabled": false,
+        "visible": false,
         "parent": "mesh",
         "value": ""
     },
diff --git a/simpeg_drivers-assets/uijson/induced_polarization_batch2d_inversion.ui.json b/simpeg_drivers-assets/uijson/induced_polarization_batch2d_inversion.ui.json
index dfa48ffd5..bec2dade8 100644
--- a/simpeg_drivers-assets/uijson/induced_polarization_batch2d_inversion.ui.json
+++ b/simpeg_drivers-assets/uijson/induced_polarization_batch2d_inversion.ui.json
@@ -278,6 +278,7 @@
         "label": "Gradient rotation",
         "optional": true,
         "enabled": false,
+        "visible": false,
         "parent": "mesh",
         "value": ""
     },

From 812aa95c72d63f833125cf0de84b2823e21cbb2b Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Wed, 2 Apr 2025 09:04:26 -0700
Subject: [PATCH 23/29] Set target for gradient rotation runtest after
 validating solution in fine resolution mesh

---
 tests/run_tests/driver_rotated_gradients_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/run_tests/driver_rotated_gradients_test.py b/tests/run_tests/driver_rotated_gradients_test.py
index bb1735c73..c326843db 100644
--- a/tests/run_tests/driver_rotated_gradients_test.py
+++ b/tests/run_tests/driver_rotated_gradients_test.py
@@ -35,7 +35,7 @@
 # To test the full run and validate the inversion.
 # Move this file out of the test directory and run.
 
-target_run = {"data_norm": 0.0028055269276044915, "phi_d": 8.32e-05, "phi_m": 0.00333}
+target_run = {"data_norm": 0.006830937520353864, "phi_d": 0.0276, "phi_m": 0.0288}
 
 
 def test_gravity_fwr_run(

From 553f63774ef6121d31d953cbffb5254fc8d18e2d Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Wed, 2 Apr 2025 09:42:12 -0700
Subject: [PATCH 24/29] remove rotated gradients from gravity runtest

---
 tests/run_tests/driver_grav_test.py | 69 -----------------------------
 1 file changed, 69 deletions(-)

diff --git a/tests/run_tests/driver_grav_test.py b/tests/run_tests/driver_grav_test.py
index 8f55b1c08..ca4d00d2b 100644
--- a/tests/run_tests/driver_grav_test.py
+++ b/tests/run_tests/driver_grav_test.py
@@ -189,75 +189,6 @@ def test_gravity_run(
             assert np.all(nan_ind == inactive_ind)
 
 
-def test_rotated_gradient(
-    tmp_path: Path,
-    max_iterations=1,
-    pytest=True,
-):
-    workpath = tmp_path / "inversion_test.ui.geoh5"
-    if pytest:
-        workpath = tmp_path.parent / "test_gravity_fwr_run0" / "inversion_test.ui.geoh5"
-
-    with Workspace(workpath) as geoh5:
-        gz = geoh5.get_entity("Iteration_0_gz")[0]
-        mesh = geoh5.get_entity("mesh")[0]
-
-        inds = (mesh.centroids[:, 0] > -35) & (mesh.centroids[:, 0] < 35)
-        norms = np.ones(mesh.n_cells) * 2
-        norms[inds] = 0
-        gradient_norms = mesh.add_data({"norms": {"values": norms}})
-
-        topography = geoh5.get_entity("topography")[0]
-
-        # Turn some values to nan
-        values = gz.values.copy()
-        values[0] = np.nan
-        gz.values = values
-
-        # Rotate the gradients
-
-        direction, dip = mesh.add_data(
-            {
-                "gradient_direction": {"values": 90 * np.ones(mesh.n_cells)},
-                "gradient_dip": {"values": 45 * np.ones(mesh.n_cells)},
-            }
-        )
-        gradient_rotation = PropertyGroup(
-            parent=mesh,
-            name="gradient_rotation",
-            property_group_type="Dip direction & dip",
-            properties=[direction, dip],
-        )
-
-        # Run the inverse
-        active_cells = ActiveCellsOptions(topography_object=topography)
-        params = GravityInversionOptions(
-            geoh5=geoh5,
-            mesh=mesh,
-            active_cells=active_cells,
-            data_object=gz.parent,
-            starting_model=1e-4,
-            reference_model=0.0,
-            gradient_rotation=gradient_rotation,
-            s_norm=0.0,
-            x_norm=gradient_norms,
-            y_norm=gradient_norms,
-            z_norm=gradient_norms,
-            gradient_type="components",
-            gz_channel=gz,
-            gz_uncertainty=2e-3,
-            lower_bound=0.0,
-            max_global_iterations=max_iterations,
-            initial_beta_ratio=1e-2,
-            percentile=100,
-            store_sensitivities="ram",
-            save_sensitivities=True,
-        )
-        params.write_ui_json(path=tmp_path / "Inv_run.ui.json")
-
-    _ = GravityInversionDriver.start(str(tmp_path / "Inv_run.ui.json"))
-
-
 if __name__ == "__main__":
     # Full run
     test_gravity_fwr_run(

From a7bb33a9b99c65b20f880a323f3db8eecb0107c4 Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Wed, 2 Apr 2025 09:56:45 -0700
Subject: [PATCH 25/29] Fixing tests

---
 simpeg_drivers/driver.py            | 2 +-
 tests/models_test.py                | 8 ++++----
 tests/run_tests/driver_grav_test.py | 1 -
 3 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/simpeg_drivers/driver.py b/simpeg_drivers/driver.py
index 25be67d3f..da360be33 100644
--- a/simpeg_drivers/driver.py
+++ b/simpeg_drivers/driver.py
@@ -473,7 +473,7 @@ def get_regularization(self):
                 components, weight_names, reg_func.objfcts
             ):
                 if getattr(self.models, weight_name) is None:
-                    setattr(reg_funcs, weight_name, 0.0)
+                    setattr(reg_func, weight_name, 0.0)
                     functions.append(fun)
                     continue
 
diff --git a/tests/models_test.py b/tests/models_test.py
index de1969b3c..b7855427f 100644
--- a/tests/models_test.py
+++ b/tests/models_test.py
@@ -72,7 +72,7 @@ def test_zero_reference_model(tmp_path: Path):
     geoh5 = params.geoh5
     with geoh5.open():
         driver = MVIInversionDriver(params)
-        _ = InversionModel(driver, "reference")
+        _ = InversionModel(driver, "reference", is_vector=True)
         incl = np.unique(geoh5.get_entity("reference_inclination")[0].values)
         decl = np.unique(geoh5.get_entity("reference_declination")[0].values)
     assert len(incl) == 1
@@ -87,7 +87,7 @@ def test_collection(tmp_path: Path):
         driver = MVIInversionDriver(params)
         models = InversionModelCollection(driver)
         models.remove_air(driver.models.active_cells)
-        starting = InversionModel(driver, "starting")
+        starting = InversionModel(driver, "starting", is_vector=True)
         starting.remove_air(driver.models.active_cells)
         np.testing.assert_allclose(models.starting, starting.model, atol=1e-7)
 
@@ -96,7 +96,7 @@ def test_initialize(tmp_path: Path):
     params = get_mvi_params(tmp_path)
     with params.geoh5.open():
         driver = MVIInversionDriver(params)
-        starting_model = InversionModel(driver, "starting")
+        starting_model = InversionModel(driver, "starting", is_vector=True)
         assert len(starting_model.model) == 3 * driver.inversion_mesh.n_cells
         assert len(np.unique(starting_model.model)) == 3
 
@@ -117,7 +117,7 @@ def test_model_from_object(tmp_path: Path):
         point_object.add_data({"test_data": {"values": vals}})
         data_object = geoh5.get_entity("test_data")[0]
         params.lower_bound = data_object
-        lower_bound = InversionModel(driver, "lower_bound")
+        lower_bound = InversionModel(driver, "lower_bound", is_vector=True)
         nc = int(len(lower_bound.model) / 3)
         A = driver.inversion_mesh.mesh.cell_centers
         b = lower_bound.model[:nc]
diff --git a/tests/run_tests/driver_grav_test.py b/tests/run_tests/driver_grav_test.py
index ca4d00d2b..4475dda68 100644
--- a/tests/run_tests/driver_grav_test.py
+++ b/tests/run_tests/driver_grav_test.py
@@ -15,7 +15,6 @@
 
 import numpy as np
 from geoapps_utils.utils.importing import GeoAppsError
-from geoh5py.groups import PropertyGroup
 from geoh5py.workspace import Workspace
 from pytest import raises
 

From 8dd9307af1d29726cf00549a953d9dbc2453cc50 Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Wed, 2 Apr 2025 09:59:34 -0700
Subject: [PATCH 26/29] more fixes

---
 simpeg_drivers/components/models.py              |  2 ++
 tests/run_tests/driver_rotated_gradients_test.py | 10 +++++++---
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/simpeg_drivers/components/models.py b/simpeg_drivers/components/models.py
index 29efc1fd7..10ef0f8d9 100644
--- a/simpeg_drivers/components/models.py
+++ b/simpeg_drivers/components/models.py
@@ -58,6 +58,8 @@ class InversionModelCollection:
 
     """
 
+    model_types = MODEL_TYPES
+
     def __init__(self, driver: InversionDriver):
         """
         :param driver: Parental InversionDriver class.
diff --git a/tests/run_tests/driver_rotated_gradients_test.py b/tests/run_tests/driver_rotated_gradients_test.py
index c326843db..1858258a0 100644
--- a/tests/run_tests/driver_rotated_gradients_test.py
+++ b/tests/run_tests/driver_rotated_gradients_test.py
@@ -38,7 +38,7 @@
 target_run = {"data_norm": 0.006830937520353864, "phi_d": 0.0276, "phi_m": 0.0288}
 
 
-def test_gravity_fwr_run(
+def test_gravity_rotated_grad_fwr_run(
     tmp_path: Path,
     n_grid_points=2,
     refinement=(2,),
@@ -76,7 +76,11 @@ def test_rotated_grad_run(
 ):
     workpath = tmp_path / "inversion_test.ui.geoh5"
     if pytest:
-        workpath = tmp_path.parent / "test_gravity_fwr_run0" / "inversion_test.ui.geoh5"
+        workpath = (
+            tmp_path.parent
+            / "test_gravity_rotated_grad_fwr_0"
+            / "inversion_test.ui.geoh5"
+        )
 
     with Workspace(workpath) as geoh5:
         gz = geoh5.get_entity("Iteration_0_gz")[0]
@@ -141,7 +145,7 @@ def test_rotated_grad_run(
 
 if __name__ == "__main__":
     # Full run
-    test_gravity_fwr_run(
+    test_gravity_rotated_grad_fwr_run(
         Path("./"),
         n_grid_points=10,
         refinement=(4, 8),

From e068ab381daa139c255d2911256b8f95566798e8 Mon Sep 17 00:00:00 2001
From: benjamink <benjamink@mirageoscience.com>
Date: Wed, 2 Apr 2025 10:04:41 -0700
Subject: [PATCH 27/29] update phi_m target on gravity runtest

---
 tests/run_tests/driver_grav_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/run_tests/driver_grav_test.py b/tests/run_tests/driver_grav_test.py
index 4475dda68..38d1980d0 100644
--- a/tests/run_tests/driver_grav_test.py
+++ b/tests/run_tests/driver_grav_test.py
@@ -34,7 +34,7 @@
 # To test the full run and validate the inversion.
 # Move this file out of the test directory and run.
 
-target_run = {"data_norm": 0.0028055269276044915, "phi_d": 8.32e-05, "phi_m": 0.00333}
+target_run = {"data_norm": 0.0028055269276044915, "phi_d": 8.32e-05, "phi_m": 0.0038}
 
 
 def test_gravity_fwr_run(

From 2f31e142640ac2bd9f159399c023ffd32f6dbc02 Mon Sep 17 00:00:00 2001
From: dominiquef <dominiquef@mirageoscience.com>
Date: Wed, 2 Apr 2025 10:33:06 -0700
Subject: [PATCH 28/29] Clip out rows of zero gradients. Simplify
 get_regularization

---
 simpeg_drivers/driver.py               | 55 +++++++++++++-------------
 simpeg_drivers/utils/regularization.py | 10 +++--
 2 files changed, 34 insertions(+), 31 deletions(-)

diff --git a/simpeg_drivers/driver.py b/simpeg_drivers/driver.py
index 25be67d3f..d8b5da4f9 100644
--- a/simpeg_drivers/driver.py
+++ b/simpeg_drivers/driver.py
@@ -480,38 +480,31 @@ def get_regularization(self):
                 weight = mapping * getattr(self.models, weight_name)
                 norm = mapping * getattr(self.models, f"{comp}_norm")
 
-                if isinstance(fun, SparseSmoothness):
-                    if is_rotated:
-                        if forward_mesh is None:
-                            fun = set_rotated_operators(
-                                fun,
-                                neighbors,
-                                comp,
-                                self.models.gradient_dip,
-                                self.models.gradient_direction,
-                            )
-
-                    else:
-                        weight = (
-                            getattr(
-                                reg_func.regularization_mesh,
-                                f"aveCC2F{fun.orientation}",
-                            )
-                            * weight
-                        )
-                        norm = (
-                            getattr(
-                                reg_func.regularization_mesh,
-                                f"aveCC2F{fun.orientation}",
-                            )
-                            * norm
+                if not isinstance(fun, SparseSmoothness):
+                    fun.set_weights(**{comp: weight})
+                    fun.norm = norm
+                    functions.append(fun)
+                    continue
+
+                if is_rotated:
+                    if forward_mesh is None:
+                        fun = set_rotated_operators(
+                            fun,
+                            neighbors,
+                            comp,
+                            self.models.gradient_dip,
+                            self.models.gradient_direction,
                         )
 
-                fun.set_weights(**{comp: weight})
-                fun.norm = norm
+                average_op = getattr(
+                    reg_func.regularization_mesh,
+                    f"aveCC2F{fun.orientation}",
+                )
+                fun.set_weights(**{comp: average_op @ weight})
+                fun.norm = average_op @ norm
                 functions.append(fun)
 
-                if isinstance(fun, SparseSmoothness) and is_rotated:
+                if is_rotated:
                     fun.gradient_type = "components"
                     backward_fun = deepcopy(fun)
                     setattr(backward_fun, "_regularization_mesh", backward_mesh)
@@ -526,6 +519,12 @@ def get_regularization(self):
                             self.models.gradient_direction,
                             forward=False,
                         )
+                    average_op = getattr(
+                        backward_fun.regularization_mesh,
+                        f"aveCC2F{fun.orientation}",
+                    )
+                    backward_fun.set_weights(**{comp: average_op @ weight})
+                    backward_fun.norm = average_op @ norm
                     functions.append(backward_fun)
 
             # Will avoid recomputing operators if the regularization mesh is the same
diff --git a/simpeg_drivers/utils/regularization.py b/simpeg_drivers/utils/regularization.py
index 843ab5372..943760cb1 100644
--- a/simpeg_drivers/utils/regularization.py
+++ b/simpeg_drivers/utils/regularization.py
@@ -415,16 +415,20 @@ def set_rotated_operators(
     grad_op = rotated_gradient(
         function.regularization_mesh.mesh, neighbors, axis, dip, direction, forward
     )
+    grad_op_active = function.regularization_mesh.Pac.T @ (
+        grad_op @ function.regularization_mesh.Pac
+    )
+    active_faces = grad_op_active.max(axis=1).toarray().ravel() > 0
+
     setattr(
         function.regularization_mesh,
         f"_cell_gradient_{function.orientation}",
-        function.regularization_mesh.Pac.T
-        @ (grad_op @ function.regularization_mesh.Pac),
+        grad_op_active[active_faces, :],
     )
     setattr(
         function.regularization_mesh,
         f"_aveCC2F{function.orientation}",
-        sdiag(np.ones(function.regularization_mesh.n_cells)),
+        sdiag(np.ones(function.regularization_mesh.n_cells))[active_faces, :],
     )
 
     return function

From 87cfa93ba130830432d8006c6fe181e90c665972 Mon Sep 17 00:00:00 2001
From: dominiquef <dominiquef@mirageoscience.com>
Date: Wed, 2 Apr 2025 13:52:04 -0700
Subject: [PATCH 29/29] Fix driver for 1d inversions

---
 simpeg_drivers/driver.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/simpeg_drivers/driver.py b/simpeg_drivers/driver.py
index fe7049b03..1b1cb2a97 100644
--- a/simpeg_drivers/driver.py
+++ b/simpeg_drivers/driver.py
@@ -466,7 +466,14 @@ def get_regularization(self):
                 neighbors = cell_neighbors(reg_func.regularization_mesh.mesh)
 
             # Adjustment for 2D versus 3D problems
-            components = "sxz" if "2d" in self.params.inversion_type else "sxyz"
+            components = (
+                "sxz"
+                if (
+                    "2d" in self.params.inversion_type
+                    or "1d" in self.params.inversion_type
+                )
+                else "sxyz"
+            )
             weight_names = ["alpha_s"] + [f"length_scale_{k}" for k in components[1:]]
             functions = []
             for comp, weight_name, fun in zip(