GEOPY-2173: Hook up 2D rotated gradients to batch 2D inversions

simpeg_drivers-assets/uijson/direct_current_batch2d_inversion.ui.json

@@ -270,7 +270,6 @@
         "label": "Gradient rotation",
         "optional": true,
         "enabled": false,
-        "visible": false,
         "parent": "mesh",
         "value": ""
     },

simpeg_drivers-assets/uijson/induced_polarization_batch2d_inversion.ui.json

@@ -281,7 +281,6 @@
         "label": "Gradient rotation",
         "optional": true,
         "enabled": false,
-        "visible": false,
         "parent": "mesh",
         "value": ""
     },

simpeg_drivers/electricals/driver.py

@@ -20,6 +20,7 @@
 from geoapps_utils.utils.locations import get_locations
 from geoapps_utils.utils.numerical import weighted_average
 from geoh5py.data import Data
+from geoh5py.groups import PropertyGroup
 from geoh5py.objects import DrapeModel
 from geoh5py.ui_json.ui_json import fetch_active_workspace
 from geoh5py.workspace import Workspace
@@ -105,6 +106,14 @@ def transfer_models(self, mesh: DrapeModel) -> dict[str, uuid.UUID | float]:
             for model in ["reference_model", "lower_bound", "upper_bound"]:
                 models[model] = getattr(self.batch2d_params.models, model)
 
+            if self.batch2d_params.models.gradient_rotation is not None:
+                group_properties = {}
+                for prop in self.batch2d_params.models.gradient_rotation.properties:
+                    model = self.batch2d_params.mesh.get_data(prop)[0]
+                    group_properties[model.name] = model
+
+                models.update(group_properties)
+
         if self.batch2d_params.mesh is not None:
             xyz_in = get_locations(self.workspace, self.batch2d_params.mesh)
             xyz_out = mesh.centroids
@@ -122,6 +131,19 @@ def transfer_models(self, mesh: DrapeModel) -> dict[str, uuid.UUID | float]:
                 model_object = mesh.add_data({name: {"values": model_values}})
                 models[name] = model_object
 
+            if (
+                not self.batch2d_params.forward_only
+                and self.batch2d_params.models.gradient_rotation is not None
+            ):
+                pg = PropertyGroup(
+                    mesh,
+                    properties=[models[prop] for prop in group_properties],
+                    property_group_type=self.batch2d_params.models.gradient_rotation.property_group_type,
+                )
+                models["gradient_rotation"] = pg
+                del models["azimuth"]
+                del models["dip"]
+
         return models
 
     def write_files(self, lookup):

tests/run_tests/driver_2d_rotated_gradients_test.py

@@ -14,6 +14,7 @@
 from unittest.mock import patch
 
 import numpy as np
+from geoapps_utils.modelling.plates import PlateModel
 from geoapps_utils.utils.importing import GeoAppsError
 from geoh5py.groups.property_group import PropertyGroup
 from geoh5py.workspace import Workspace
@@ -39,7 +40,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.5963828772690819, "phi_d": 2870, "phi_m": 18.7}
+target_run = {"data_norm": 0.6345768240307744, "phi_d": 1090, "phi_m": 3.99}
 
 
 def test_dc2d_rotated_grad_fwr_run(
@@ -48,11 +49,21 @@ def test_dc2d_rotated_grad_fwr_run(
     n_lines=3,
     refinement=(4, 6),
 ):
+    plate_model = PlateModel(
+        strike_length=1000.0,
+        dip_length=150.0,
+        width=20.0,
+        origin=(50.0, 0.0, -30),
+        direction=90,
+        dip=45,
+    )
+
     # Run the forward
     geoh5, _, model, survey, topography = setup_inversion_workspace(
         tmp_path,
+        plate_model=plate_model,
         background=0.01,
-        anomaly=10,
+        anomaly=10.0,
         n_electrodes=n_electrodes,
         n_lines=n_lines,
         refinement=refinement,
@@ -104,22 +115,23 @@ def test_dc2d_rotated_grad_run(
         mesh = geoh5.get_entity("Models")[0]
 
         # Create property group with orientation
-        dip = np.ones(mesh.n_cells) * 45
-        azimuth = np.ones(mesh.n_cells) * 90
+        i = np.ones(mesh.n_cells)
+        j = np.zeros(mesh.n_cells)
+        k = np.ones(mesh.n_cells)
 
         data_list = mesh.add_data(
             {
-                "azimuth": {"values": azimuth},
-                "dip": {"values": dip},
+                "i": {"values": i},
+                "j": {"values": j},
+                "k": {"values": k},
             }
         )
-        pg = PropertyGroup(
-            mesh, properties=data_list, property_group_type="Dip direction & dip"
-        )
+        pg = PropertyGroup(mesh, properties=data_list, property_group_type="3D vector")
 
         # Run the inverse
         params = DC2DInversionOptions.build(
             geoh5=geoh5,
+            mesh=mesh,
             drape_model=DrapeModelOptions(
                 u_cell_size=5.0,
                 v_cell_size=5.0,
@@ -140,7 +152,7 @@ def test_dc2d_rotated_grad_run(
             model_type="Resistivity (Ohm-m)",
             starting_model=100.0,
             reference_model=100.0,
-            s_norm=0.0,
+            s_norm=1.0,
             x_norm=0.0,
             z_norm=0.0,
             max_global_iterations=max_iterations,

tests/run_tests/driver_airborne_tem_test.py

@@ -13,6 +13,7 @@
 from pathlib import Path
 
 import numpy as np
+from geoapps_utils.modelling.plates import PlateModel
 from geoh5py.groups import SimPEGGroup
 from geoh5py.workspace import Workspace
 from pytest import raises
@@ -38,6 +39,7 @@
 def test_bad_waveform(tmp_path: Path):
     n_grid_points = 3
     refinement = (2,)
+
     geoh5, _, model, survey, topography = setup_inversion_workspace(
         tmp_path,
         background=0.001,

tests/run_tests/driver_dc_b2d_rotated_gradients_test.py

@@ -0,0 +1,195 @@
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+#  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
+
+import json
+from pathlib import Path
+
+import numpy as np
+from geoapps_utils.modelling.plates import PlateModel
+from geoh5py.groups import PropertyGroup, SimPEGGroup
+from geoh5py.workspace import Workspace
+
+from simpeg_drivers.electricals.direct_current.pseudo_three_dimensions.driver import (
+    DCBatch2DForwardDriver,
+    DCBatch2DInversionDriver,
+)
+from simpeg_drivers.electricals.direct_current.pseudo_three_dimensions.options import (
+    DCBatch2DForwardOptions,
+    DCBatch2DInversionOptions,
+)
+from simpeg_drivers.electricals.options import (
+    FileControlOptions,
+)
+from simpeg_drivers.options import (
+    ActiveCellsOptions,
+    DrapeModelOptions,
+    LineSelectionOptions,
+)
+from simpeg_drivers.utils.utils import get_inversion_output
+from tests.testing_utils import check_target, setup_inversion_workspace
+
+
+# To test the full run and validate the inversion.
+# Move this file out of the test directory and run.
+
+target_run = {"data_norm": 1.1038993594022803, "phi_d": 308, "phi_m": 0.0237}
+
+
+def test_dc_rotated_p3d_fwr_run(
+    tmp_path: Path, n_electrodes=10, n_lines=3, refinement=(4, 6)
+):
+    plate_model = PlateModel(
+        strike_length=1000.0,
+        dip_length=150.0,
+        width=20.0,
+        origin=(0.0, 0.0, -50),
+        direction=90,
+        dip=45,
+    )
+
+    geoh5, _, model, survey, topography = setup_inversion_workspace(
+        tmp_path,
+        plate_model=plate_model,
+        background=0.01,
+        anomaly=10.0,
+        n_electrodes=n_electrodes,
+        n_lines=n_lines,
+        refinement=refinement,
+        inversion_type="direct current pseudo 3d",
+        drape_height=0.0,
+        flatten=False,
+    )
+    params = DCBatch2DForwardOptions.build(
+        geoh5=geoh5,
+        mesh=model.parent,
+        drape_model=DrapeModelOptions(
+            u_cell_size=5.0,
+            v_cell_size=5.0,
+            depth_core=100.0,
+            expansion_factor=1.1,
+            horizontal_padding=1000.0,
+            vertical_padding=1000.0,
+        ),
+        topography_object=topography,
+        data_object=survey,
+        starting_model=model,
+        line_selection=LineSelectionOptions(
+            line_object=geoh5.get_entity("line_ids")[0]
+        ),
+    )
+    fwr_driver = DCBatch2DForwardDriver(params)
+    fwr_driver.run()
+
+
+def test_dc_rotated_gradient_p3d_run(
+    tmp_path: Path,
+    max_iterations=1,
+    pytest=True,
+):
+    workpath = tmp_path / "inversion_test.ui.geoh5"
+    if pytest:
+        workpath = (
+            tmp_path.parent / "test_dc_rotated_p3d_fwr_run0" / "inversion_test.ui.geoh5"
+        )
+
+    with Workspace(workpath) as geoh5:
+        potential = geoh5.get_entity("Iteration_0_dc")[0]
+        out_group = geoh5.get_entity("Line 1")[0].parent
+        mesh = out_group.get_entity("mesh")[0]  # Finds the octree mesh
+        topography = geoh5.get_entity("topography")[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"
+        )
+
+        # Run the inverse
+        params = DCBatch2DInversionOptions.build(
+            geoh5=geoh5,
+            mesh=mesh,
+            drape_model=DrapeModelOptions(
+                u_cell_size=5.0,
+                v_cell_size=5.0,
+                depth_core=100.0,
+                expansion_factor=1.1,
+                horizontal_padding=1000.0,
+                vertical_padding=1000.0,
+            ),
+            topography_object=topography,
+            data_object=potential.parent,
+            gradient_rotation=pg,
+            potential_channel=potential,
+            potential_uncertainty=1e-3,
+            line_selection=LineSelectionOptions(
+                line_object=geoh5.get_entity("line_ids")[0]
+            ),
+            starting_model=1e-2,
+            reference_model=1e-2,
+            s_norm=0.0,
+            x_norm=0.0,
+            z_norm=0.0,
+            gradient_type="components",
+            max_global_iterations=max_iterations,
+            initial_beta=None,
+            initial_beta_ratio=10.0,
+            percentile=100,
+            upper_bound=10,
+            cooling_rate=1,
+            file_control=FileControlOptions(cleanup=False),
+        )
+        params.write_ui_json(path=tmp_path / "Inv_run.ui.json")
+
+    driver = DCBatch2DInversionDriver.start(str(tmp_path / "Inv_run.ui.json"))
+
+    basepath = workpath.parent
+    with open(basepath / "lookup.json", encoding="utf8") as f:
+        lookup = json.load(f)
+        middle_line_id = next(k for k, v in lookup.items() if v["line_id"] == 101)
+
+    with Workspace(basepath / f"{middle_line_id}.ui.geoh5", mode="r") as workspace:
+        middle_inversion_group = next(
+            k for k in workspace.groups if isinstance(k, SimPEGGroup)
+        )
+        filedata = middle_inversion_group.get_entity("SimPEG.out")[0]
+
+        with driver.batch2d_params.out_group.workspace.open(mode="r+"):
+            filedata.copy(parent=driver.batch2d_params.out_group)
+
+    output = get_inversion_output(
+        driver.batch2d_params.geoh5.h5file, driver.batch2d_params.out_group.uid
+    )
+    if geoh5.open():
+        output["data"] = potential.values
+    if pytest:
+        check_target(output, target_run)
+
+
+if __name__ == "__main__":
+    # Full run
+    test_dc_rotated_p3d_fwr_run(
+        Path("./"), n_electrodes=20, n_lines=3, refinement=(4, 8)
+    )
+    test_dc_rotated_gradient_p3d_run(
+        Path("./"),
+        max_iterations=20,
+        pytest=False,
+    )

tests/run_tests/driver_dc_b2d_test.py

@@ -53,7 +53,7 @@ def test_dc_p3d_fwr_run(
     geoh5, _, model, survey, topography = setup_inversion_workspace(
         tmp_path,
         background=0.01,
-        anomaly=10,
+        anomaly=10.0,
         n_electrodes=n_electrodes,
         n_lines=n_lines,
         refinement=refinement,

tests/run_tests/driver_fem_test.py

@@ -16,6 +16,7 @@
 from pathlib import Path
 
 import numpy as np
+from geoapps_utils.modelling.plates import PlateModel
 from geoh5py import Workspace
 from geoh5py.groups import SimPEGGroup
 
@@ -72,8 +73,15 @@ def test_fem_fwr_run(
     cell_size=(20.0, 20.0, 20.0),
 ):
     # Run the forward
+    plate_model = PlateModel(
+        strike_length=40.0,
+        dip_length=40.0,
+        width=40.0,
+        origin=(0.0, 0.0, -50.0),
+    )
     geoh5, _, model, survey, topography = setup_inversion_workspace(
         tmp_path,
+        plate_model=plate_model,
         background=1e-3,
         anomaly=1.0,
         n_electrodes=n_grid_points,