GEOPY-2994: Support Grid2D in grid_model_merger application

grid_apps/grid_model_merger/driver.py

@@ -17,7 +17,7 @@
 from discretize.utils import mesh_utils
 from geoapps_utils.base import Driver as BaseDriver
 from geoapps_utils.utils.plotting import inv_symlog, symlog
-from geoh5py.objects import Octree
+from geoh5py.objects import Grid2D, Octree
 from geoh5py.shared.utils import fetch_active_workspace
 from scipy.spatial import cKDTree
 
@@ -27,6 +27,7 @@
     get_boundary_active_cells,
     refine_tree_by_mesh,
     tensor_to_block_model,
+    tensor_to_grid2d,
     treemesh_2_octree,
 )
 
@@ -82,6 +83,14 @@ def get_global_mesh_specs(self) -> tuple[np.ndarray, np.ndarray]:
                     ],
                     axis=0,
                 )
+            elif isinstance(grid, Grid2D):
+                cell_size = np.min(
+                    [
+                        cell_size,
+                        np.r_[grid.u_cell_size, grid.v_cell_size, np.inf],
+                    ],
+                    axis=0,
+                )
             else:
                 cell_size = np.min(
                     [
@@ -115,6 +124,12 @@ def get_output_grid(self):
             # Use type of the first entry
             mesh_type = type(self.params.selections[0].grid)
 
+            elevation = 0.0
+            if mesh_type is Grid2D:
+                elevation = np.mean(extent[:, 2])
+                extent = extent[:, :2]
+                cell_size = cell_size[:2]
+
             logger.info("Merging selected grids to '%s' . . .", mesh_type.__name__)
 
             mesh = mesh_utils.mesh_builder_xyz(
@@ -132,6 +147,13 @@ def get_output_grid(self):
                 output_grid = treemesh_2_octree(
                     self.params.geoh5, mesh, parent=self.params.out_group
                 )
+            elif mesh_type is Grid2D:
+                output_grid = tensor_to_grid2d(
+                    self.params.geoh5,
+                    mesh,
+                    elevation=elevation,
+                    parent=self.params.out_group,
+                )
             else:
                 output_grid = tensor_to_block_model(
                     self.params.geoh5, mesh, parent=self.params.out_group
@@ -155,7 +177,6 @@ def interpolate_models_to_output_grid(self):
                     continue
 
                 active = ~np.isnan(model)
-
                 logger.info(
                     "Interpolating model '%s' from grid '%s' to output grid '%s' . . .",
                     selection.model.name,
@@ -214,7 +235,9 @@ def cosine_taper_weights(edge_locations: np.ndarray, target_locations: np.ndarra
         rad, _ = tree.query(target_locations, workers=-1)
 
         rad_max = rad.max() + 1e-8  # Avoid zero division
-        cosine_taper = -0.5 * np.cos(-rad / rad_max * np.pi) + 0.5
+        cosine_taper = (
+            -0.5 * np.cos(-rad / rad_max * np.pi) + 0.501
+        )  # Weights from [0.001 to 1.001]
 
         # Find nearest neighbors and apply weighted model
         del tree

grid_apps/grid_model_merger/options.py

@@ -18,13 +18,14 @@
 from discretize import TensorMesh, TreeMesh
 from geoapps_utils.base import Options
 from geoh5py.data import NumericData
-from geoh5py.objects import BlockModel, Octree
+from geoh5py.objects import BlockModel, Grid2D, Octree
 from geoh5py.objects.grid_object import GridObject
 from pydantic import BaseModel, ConfigDict, model_serializer, model_validator
 
 from grid_apps import assets_path
 from grid_apps.utils import (
     block_model_to_tensor,
+    grid2d_to_tensor,
     octree_2_treemesh,
     tensor_mesh_ordering,
 )
@@ -129,6 +130,9 @@ def to_discretize(self) -> tuple[TreeMesh | TensorMesh, np.ndarray | None]:
         elif isinstance(self.grid, Octree):
             mesh = octree_2_treemesh(self.grid)
 
+        elif isinstance(self.grid, Grid2D):
+            mesh = grid2d_to_tensor(self.grid)
+
         if mesh is None:
             raise TypeError(f"Mesh type {type(self.grid)} currently not supported.")
 

grid_apps/utils.py

@@ -15,7 +15,7 @@
 from discretize import TensorMesh, TreeMesh
 from geoh5py import Workspace
 from geoh5py.data import FloatData, ReferencedData
-from geoh5py.objects import BlockModel, Curve, ObjectBase, Octree, Points
+from geoh5py.objects import BlockModel, Curve, Grid2D, ObjectBase, Octree, Points
 from geoh5py.ui_json.utils import fetch_active_workspace
 from pydantic import ConfigDict, validate_call
 from scipy.interpolate import interp1d
@@ -77,6 +77,31 @@ def tensor_to_block_model(
     return block_model
 
 
+@typed_call
+def tensor_to_grid2d(
+    workspace: Workspace, mesh: TensorMesh, elevation: float = 0.0, **kwargs
+) -> Grid2D:
+    """
+    Convert a tensor mesh to a 2D grid object.
+
+    :param workspace: Workspace to create the Grid2D object.
+    :param mesh: Tensor mesh object from discretize
+    :param kwargs: Extra parameters to pass to the Grid2D object.
+
+    :return: Grid2D entity.
+    """
+    grid = Grid2D.create(
+        workspace,
+        origin=[mesh.x0[0], mesh.x0[1], elevation],
+        u_cell_size=np.mean(mesh.h[0]),
+        v_cell_size=np.mean(mesh.h[1]),
+        u_count=len(mesh.h[0]),
+        v_count=len(mesh.h[1]),
+        **kwargs,
+    )
+    return grid
+
+
 @typed_call
 def block_model_to_treemesh(
     entity: BlockModel, diagonal_balance=True, finalize=True
@@ -188,14 +213,13 @@ def containing_cell_indices(mesh: TreeMesh | TensorMesh, locations: np.ndarray):
     else:
         in_x = np.searchsorted(mesh.nodes_x, locations[:, 0]) - 1
         in_y = np.searchsorted(mesh.nodes_y, locations[:, 1]) - 1
-        in_z = np.searchsorted(mesh.nodes_z, locations[:, 2]) - 1
-        indices = (
-            in_x
-            + mesh.shape_cells[0] * in_y
-            + in_z * mesh.shape_cells[0] * mesh.shape_cells[1]
-        ).astype(int)
+        indices = (in_x + mesh.shape_cells[0] * in_y).astype(int)
+
+        if mesh.dim > 2:
+            in_z = np.searchsorted(mesh.nodes_z, locations[:, 2]) - 1
+            indices += in_z * mesh.shape_cells[0] * mesh.shape_cells[1]
 
-    indices[~mesh.is_inside(locations)] = -1
+    indices[~mesh.is_inside(locations[:, : mesh.dim])] = -1
 
     return indices
 
@@ -249,7 +273,9 @@ def create_octree_from_octrees(meshes: list[Octree | TreeMesh]) -> TreeMesh:
 
 @typed_call
 def refine_tree_by_mesh(
-    tree: TreeMesh, mesh: TreeMesh | Octree | BlockModel, finalize: bool = False
+    tree: TreeMesh,
+    mesh: TreeMesh | Octree | BlockModel | Grid2D,
+    finalize: bool = False,
 ) -> TreeMesh:
     """
     Given a TreeMesh, insert cells at the corresponding octree level.
@@ -267,6 +293,13 @@ def refine_tree_by_mesh(
         centers = mesh.centroids
         levels = tree.max_level - np.log2(mesh.octree_cells["NCells"])
 
+    elif isinstance(mesh, Grid2D):
+        centers = mesh.centroids
+        octree_level = np.max(
+            [0, int(np.min([mesh.u_cell_size, mesh.v_cell_size]) // np.min(tree.h)) - 1]
+        )
+        levels = np.full(mesh.n_cells, tree.max_level - octree_level, dtype=int)
+
     else:
         centers = mesh.cell_centers
         levels = (
@@ -445,10 +478,10 @@ def get_boundary_active_cells(
 
     # Find actives horizontal mesh boundary cells
     if horizontal_edges:
-        for face in mesh.cell_boundary_indices[:-2]:
+        for face in mesh.cell_boundary_indices[0:4]:
             is_face[face] = True
 
-    if vertical_edges:
+    if vertical_edges and mesh.dim > 2:
         for face in mesh.cell_boundary_indices[-2:]:
             is_face[face] = True
 
@@ -717,3 +750,34 @@ def treemesh_2_octree(workspace: Workspace, treemesh: TreeMesh, **kwargs) -> Oct
     )
 
     return mesh_object
+
+
+@typed_call
+def grid2d_to_tensor(
+    entity: Grid2D,
+) -> TensorMesh:
+    """
+    Convert a Grid2D object to a discretize.TensorMesh.
+
+    :param entity: The Grid2D object to convert.
+
+    :return: An equivalent TensorMesh object.
+    """
+
+    if entity.rotation != 0.0 or entity.dip != 0.0:
+        raise NotImplementedError(
+            "Conversion of rotated or dipping 2D grid not supported."
+        )
+
+    origin = [
+        entity.origin["x"],
+        entity.origin["y"],
+    ]
+    mesh = TensorMesh(
+        [
+            np.full(entity.u_count, entity.u_cell_size),
+            np.full(entity.v_count, entity.v_cell_size),
+        ],
+        x0=origin,
+    )
+    return mesh