GEOPY-1302: Organize geoapps.utils.testing.setup_inversion_workspace

.pre-commit-config.yaml

@@ -91,7 +91,7 @@ repos:
     -   id: mixed-line-ending
         exclude: ^\.idea/.*\.xml$
     -   id: name-tests-test
-        exclude: testing_utils.py
+        exclude: targets.py
     -   id: pretty-format-json
         args:
             - --autofix

simpeg_drivers/utils/synthetics/__init__.py

@@ -0,0 +1,9 @@
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+#  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).                      '
+#                                                                                   '
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''

simpeg_drivers/utils/synthetics/driver.py

@@ -0,0 +1,109 @@
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+#  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 geoh5py import Workspace
+from geoh5py.data import FloatData
+from geoh5py.objects import DrapeModel, ObjectBase, Octree, Surface
+
+from simpeg_drivers.utils.synthetics.meshes.factory import get_mesh
+from simpeg_drivers.utils.synthetics.models import get_model
+from simpeg_drivers.utils.synthetics.options import SyntheticsComponentsOptions
+from simpeg_drivers.utils.synthetics.surveys.factory import get_survey
+from simpeg_drivers.utils.synthetics.topography import (
+    get_active,
+    get_topography_surface,
+)
+
+
+class SyntheticsComponents:
+    """Creates a workspace populated with objects for simulation and subsequent inversion."""
+
+    def __init__(
+        self,
+        geoh5: Workspace,
+        options: SyntheticsComponentsOptions | None = None,
+    ):
+        if options is None:
+            options = SyntheticsComponentsOptions()
+
+        self.geoh5 = geoh5
+        self.options = options
+        self._topography: Surface | None = None
+        self._survey: ObjectBase | None = None
+        self._mesh: Octree | DrapeModel | None = None
+        self._active: FloatData | None = None
+        self._model: FloatData | None = None
+
+    @property
+    def topography(self):
+        entity = self.geoh5.get_entity("topography")[0]
+        assert isinstance(entity, Surface | type(None))
+        if entity is None:
+            assert self.options is not None
+            entity = get_topography_surface(
+                geoh5=self.geoh5,
+                options=self.options.survey,
+            )
+        self._topography = entity
+        return self._topography
+
+    @property
+    def survey(self):
+        entity = self.geoh5.get_entity(self.options.survey.name)[0]
+        assert isinstance(entity, ObjectBase | type(None))
+        if entity is None:
+            assert self.options is not None
+            entity = get_survey(
+                geoh5=self.geoh5,
+                method=self.options.method,
+                options=self.options.survey,
+            )
+        self._survey = entity
+        return self._survey
+
+    @property
+    def mesh(self):
+        entity = self.geoh5.get_entity(self.options.mesh.name)[0]
+        assert isinstance(entity, Octree | DrapeModel | type(None))
+        if entity is None:
+            assert self.options is not None
+            entity = get_mesh(
+                self.options.method,
+                survey=self.survey,
+                topography=self.topography,
+                options=self.options.mesh,
+            )
+        self._mesh = entity
+        return self._mesh
+
+    @property
+    def active(self):
+        entity = self.geoh5.get_entity(self.options.active.name)[0]
+        assert isinstance(entity, FloatData | type(None))
+        if entity is None:
+            entity = get_active(self.mesh, self.topography)
+        self._active = entity
+        return self._active
+
+    @property
+    def model(self):
+        entity = self.geoh5.get_entity(self.options.model.name)[0]
+        assert isinstance(entity, FloatData | type(None))
+        if entity is None:
+            assert self.options is not None
+            entity = get_model(
+                method=self.options.method,
+                mesh=self.mesh,
+                active=self.active.values,
+                options=self.options.model,
+            )
+        self._model = entity
+        return self._model

simpeg_drivers/utils/synthetics/meshes/__init__.py

@@ -0,0 +1,9 @@
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+#  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).                      '
+#                                                                                   '
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''

simpeg_drivers/utils/synthetics/meshes/factory.py

@@ -0,0 +1,45 @@
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+#  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 discretize import TensorMesh, TreeMesh
+from geoh5py.objects import CellObject, DrapeModel, Octree, Points, Surface
+
+from simpeg_drivers.utils.synthetics.options import MeshOptions
+
+from .octrees import get_octree_mesh
+from .tensors import get_tensor_mesh
+
+
+def get_mesh(
+    method: str,
+    survey: Points,
+    topography: Surface,
+    options: MeshOptions,
+) -> DrapeModel | Octree:
+    """Factory for mesh creation with behaviour modified by the provided method."""
+
+    if "2d" in method:
+        assert isinstance(survey, CellObject)
+        return get_tensor_mesh(
+            survey=survey,
+            cell_size=options.cell_size,
+            padding_distance=options.padding_distance,
+            name=options.name,
+        )
+
+    return get_octree_mesh(
+        survey=survey,
+        topography=topography,
+        cell_size=options.cell_size,
+        refinement=options.refinement,
+        padding_distance=options.padding_distance,
+        refine_on_receivers=method in ["fdem", "airborne tdem"],
+        name=options.name,
+    )

simpeg_drivers/utils/synthetics/meshes/octrees.py

@@ -0,0 +1,90 @@
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+#  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 discretize.utils import mesh_builder_xyz
+from geoh5py.objects import Octree, Points, Surface
+from octree_creation_app.driver import OctreeDriver
+from octree_creation_app.utils import treemesh_2_octree
+
+
+def get_base_octree(
+    survey: Points,
+    topography: Surface,
+    cell_size: tuple[float, float, float],
+    refinement: tuple,
+    padding: float,
+) -> TreeMesh:
+    """
+    Generate a survey centered TreeMesh object with topography refinement.
+
+    :param survey: Survey object with vertices that define the core of the
+        tensor mesh.
+    :param topography: Surface used to refine the topography.
+    :param cell_size: Tuple defining the cell size in all directions.
+    :param refinement: Tuple containing the number of cells to refine at each
+        level around the topography.
+    :param padding: Distance to pad the mesh in all directions.
+
+    :return mesh: The discretize TreeMesh object for computations.
+    """
+    padding_distance = np.ones((3, 2)) * padding
+    mesh = mesh_builder_xyz(
+        survey.vertices - np.r_[cell_size] / 2.0,
+        cell_size,
+        depth_core=100.0,
+        padding_distance=padding_distance,
+        mesh_type="TREE",
+        tree_diagonal_balance=False,
+    )
+    mesh = OctreeDriver.refine_tree_from_surface(
+        mesh, topography, levels=refinement, finalize=False
+    )
+
+    return mesh
+
+
+def get_octree_mesh(
+    survey: Points,
+    topography: Surface,
+    cell_size: tuple[float, float, float],
+    refinement: tuple,
+    padding_distance: float,
+    refine_on_receivers: bool,
+    name: str = "mesh",
+) -> Octree:
+    """Generate a survey centered mesh with topography and survey refinement.
+
+    :param survey: Survey object with vertices that define the core of the
+        tensor mesh and the source refinement for EM methods.
+    :param topography: Surface used to refine the topography.
+    :param cell_size: Tuple defining the cell size in all directions.
+    :param refinement: Tuple containing the number of cells to refine at each
+        level around the topography.
+    :param padding: Distance to pad the mesh in all directions.
+    :param refine_on_receivers: Refine on the survey locations or not.
+
+    :return entity: The geoh5py Octree object to store the results of
+        computation in the shared cells of the computational mesh.
+    :return mesh: The discretize TreeMesh object for computations.
+    """
+
+    mesh = get_base_octree(survey, topography, cell_size, refinement, padding_distance)
+
+    if refine_on_receivers:
+        mesh = OctreeDriver.refine_tree_from_points(
+            mesh, survey.vertices, levels=[2], finalize=False
+        )
+
+    mesh.finalize()
+    entity = treemesh_2_octree(survey.workspace, mesh, name=name)
+
+    return entity

simpeg_drivers/utils/synthetics/meshes/tensors.py

@@ -0,0 +1,55 @@
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+#  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 geoh5py.data import IntegerData
+from geoh5py.objects import CellObject, DrapeModel
+
+from simpeg_drivers.utils.utils import get_drape_model
+
+
+def get_tensor_mesh(
+    survey: CellObject,
+    cell_size: tuple[float, float, float],
+    padding_distance: float,
+    line_id: int = 101,
+    name: str = "mesh",
+) -> DrapeModel:
+    """
+    Generate a tensor mesh and the colocated DrapeModel.
+
+    :param survey: Survey object with vertices that define the core of the
+        tensor mesh.
+    :param cell_size: Tuple defining the cell size in all directions.
+    :param padding_distance: Distance to pad the mesh in all directions.
+    :param line_id: Chooses line from the survey to define the drape model.
+
+    :return entity: The DrapeModel object that shares cells with the discretize
+        tensor mesh and which stores the results of computations.
+    :return mesh: The discretize tensor mesh object for computations.
+    """
+
+    line_data = survey.get_entity("line_ids")[0]
+    assert isinstance(line_data, IntegerData)
+    lines = line_data.values
+    entity, mesh, _ = get_drape_model(  # pylint: disable=unbalanced-tuple-unpacking
+        survey.workspace,
+        name,
+        survey.vertices[np.unique(survey.cells[lines == line_id, :]), :],
+        [cell_size[0], cell_size[2]],
+        100.0,
+        [padding_distance] * 2 + [padding_distance] * 2,
+        1.1,
+        parent=None,
+        return_colocated_mesh=True,
+        return_sorting=True,
+    )
+
+    return entity

simpeg_drivers/utils/synthetics/models.py

@@ -0,0 +1,53 @@
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+#  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 geoapps_utils.modelling.plates import make_plate
+from geoh5py.data import FloatData
+from geoh5py.objects import DrapeModel, Octree
+
+from simpeg_drivers.utils.synthetics.options import ModelOptions
+
+
+def get_model(
+    method: str,
+    mesh: Octree | DrapeModel,
+    active: np.ndarray,
+    options: ModelOptions,
+) -> FloatData:
+    """
+    Create a halfspace and plate model in the cell centers of the provided mesh.
+
+    :param method: The geophysical method controlling the factory behaviour
+    :param mesh: The mesh whose cell centers the model will be defined on.
+    :param plate: The plate object defining the location and orientation of the
+        plate anomaly.
+    :param background: Value given to the halfspace.
+    :param anomaly: Value given to the plate.
+    """
+
+    cell_centers = mesh.centroids.copy()
+
+    model = make_plate(
+        points=cell_centers,
+        plate=options.plate,
+        background=options.background,
+        anomaly=options.anomaly,
+    )
+
+    if "1d" in method:
+        model = options.background * np.ones(mesh.n_cells)
+        inside_anomaly = (mesh.centroids[:, 2] < 0) & (mesh.centroids[:, 2] > -20)
+        model[inside_anomaly] = options.anomaly
+
+    model[~active] = np.nan
+    model = mesh.add_data({options.name: {"values": model}})
+    assert isinstance(model, FloatData)
+    return model