diff --git a/grid_apps/grid_model_merger/driver.py b/grid_apps/grid_model_merger/driver.py index 1dbeca4..92eb3a8 100644 --- a/grid_apps/grid_model_merger/driver.py +++ b/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 diff --git a/grid_apps/grid_model_merger/options.py b/grid_apps/grid_model_merger/options.py index 3949a4d..c2a04de 100644 --- a/grid_apps/grid_model_merger/options.py +++ b/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.") diff --git a/grid_apps/utils.py b/grid_apps/utils.py index ffa9f64..05e9343 100644 --- a/grid_apps/utils.py +++ b/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 diff --git a/tests/run_tests/grid_model_merger_test.py b/tests/run_tests/grid_model_merger_test.py index dcd6ec9..fc6b0d5 100644 --- a/tests/run_tests/grid_model_merger_test.py +++ b/tests/run_tests/grid_model_merger_test.py @@ -12,8 +12,9 @@ from pathlib import Path import numpy as np -from geoh5py.objects import BlockModel, Octree, Points +from geoh5py.objects import BlockModel, Grid2D, Octree, Points from geoh5py.workspace import Workspace +from pytest import mark from grid_apps.block_models.driver import Driver as BlockModelDriver from grid_apps.grid_model_merger.driver import Driver @@ -25,7 +26,7 @@ def setup_block_model(workspace) -> BlockModel: # padding in the W/E/N/S directions should make create locs at least as # far as the core hull plus the padding distances - top = 500 + top = 200 depth_core = 300.0 height = 300 width = 1000 @@ -41,6 +42,18 @@ def setup_block_model(workspace) -> BlockModel: return mesh +def setup_grid2d_model(workspace) -> Grid2D: + mesh = Grid2D.create( + workspace, + origin=[0, 0, 0], + u_cell_size=50.0, + v_cell_size=50.0, + u_count=10, + v_count=15, + ) + return mesh + + def setup_octree(workspace, locations, refinement, params_dict) -> Octree: points = Points.create(workspace, vertices=locations) @@ -69,7 +82,7 @@ def test_merge_block_model(tmp_path: Path): # pylint: disable=too-many-locals with Workspace.create(tmp_path / f"{__name__}.geoh5") as ws: mesh = setup_block_model(ws) - other = mesh.copy(origin=(515, 10, 500)) + other = mesh.copy(origin=(515, 10, 300)) model_b = other.add_data({"values": {"values": np.full(mesh.n_cells, 2.0)}}) model_a = mesh.add_data({"values": {"values": np.full(mesh.n_cells, 1.0)}}) @@ -87,7 +100,7 @@ def test_merge_block_model(tmp_path: Path): # pylint: disable=too-many-locals out_grid = driver.run() merged_model = out_grid.children[0] - np.testing.assert_almost_equal(merged_model.values[2649], 1.5, decimal=1) + np.testing.assert_almost_equal(merged_model.values[3406], 1.5, decimal=1) # Repeat with a hole in the first model values = model_a.values @@ -96,7 +109,33 @@ def test_merge_block_model(tmp_path: Path): # pylint: disable=too-many-locals out_grid = driver.run() merged_model = out_grid.children[0] - np.testing.assert_almost_equal(merged_model.values[2649], 2.0, decimal=1) + np.testing.assert_almost_equal(merged_model.values[3406], 2.0, decimal=1) + + +def test_merge_grid2d_model(tmp_path: Path): # pylint: disable=too-many-locals + + with Workspace.create(tmp_path / f"{__name__}.geoh5") as ws: + mesh = setup_grid2d_model(ws) + other = mesh.copy(origin=(150, 100, 60)) + model_b = other.add_data({"values": {"values": np.full(mesh.n_cells, 2.0)}}) + model_a = mesh.add_data({"values": {"values": np.full(mesh.n_cells, 1.0)}}) + + options = GridModelMergerOptions.build( + { + "geoh5": ws, + "input_a_grid": mesh, + "input_b_grid": other, + "input_a_model": model_a, + "input_b_model": model_b, + } + ) + + driver = Driver(options) + out_grid = driver.run() + + assert out_grid.origin["z"] == 30.0 + merged_model = out_grid.children[0] + np.testing.assert_almost_equal(merged_model.values[110], 1.5, decimal=1) def test_merge_octree_model(tmp_path: Path, setup_test_octree): # pylint: disable=too-many-locals @@ -135,23 +174,45 @@ def test_merge_octree_model(tmp_path: Path, setup_test_octree): # pylint: disab np.testing.assert_almost_equal(merged_model.values[1568], 2.0, decimal=2) -def test_merge_mixed_model(tmp_path: Path, setup_test_octree): # pylint: disable=too-many-locals - (locations, refinement, _, params_dict) = setup_test_octree - - with Workspace.create(tmp_path / f"{__name__}.geoh5") as ws: - octree = setup_octree(ws, locations, refinement, params_dict) - block_model = setup_block_model(ws) - - model_a = octree.add_data({"values": {"values": np.full(octree.n_cells, 1.0)}}) - model_b = block_model.add_data( - {"values": {"values": np.full(block_model.n_cells, 2.0)}} - ) +@mark.parametrize( + "type_a, type_b", + [ + (BlockModel, Octree), + (BlockModel, Grid2D), + (Grid2D, BlockModel), + (Grid2D, Octree), + (Octree, BlockModel), + (Octree, Grid2D), + ], +) +def test_merge_mixed_model( + tmp_path: Path, setup_test_octree, type_a: type, type_b: type +): # pylint: disable=too-many-locals + + def make_octree_model(workspace): + (locations, refinement, _, params_dict) = setup_test_octree + return setup_octree(workspace, locations, refinement, params_dict) + + dict_builder = { + BlockModel: setup_block_model, + Octree: make_octree_model, + Grid2D: setup_grid2d_model, + } + + with Workspace.create( + tmp_path / f"{type_a.__name__}_and_{type_b.__name__}.geoh5" + ) as ws: + mesh_a = dict_builder[type_a](ws) + mesh_b = dict_builder[type_b](ws) + + model_a = mesh_a.add_data({"values": {"values": np.full(mesh_a.n_cells, 1.0)}}) + model_b = mesh_b.add_data({"values": {"values": np.full(mesh_b.n_cells, 2.0)}}) options = GridModelMergerOptions.build( { "geoh5": ws, - "input_a_grid": octree, - "input_b_grid": block_model, + "input_a_grid": mesh_a, + "input_b_grid": mesh_b, "input_a_model": model_a, "input_b_model": model_b, } @@ -160,19 +221,4 @@ def test_merge_mixed_model(tmp_path: Path, setup_test_octree): # pylint: disabl driver = Driver(options) out_grid = driver.run() - assert isinstance(out_grid, Octree) - - options = GridModelMergerOptions.build( - { - "geoh5": ws, - "input_b_grid": octree, - "input_a_grid": block_model, - "input_b_model": model_a, - "input_a_model": model_b, - } - ) - - driver = Driver(options) - out_grid = driver.run() - - assert isinstance(out_grid, BlockModel) + assert type(out_grid) is type_a