diff --git a/grid_apps-assets/uijson/grid_model_merger.ui.json b/grid_apps-assets/uijson/grid_model_merger.ui.json new file mode 100644 index 0000000..54ffac3 --- /dev/null +++ b/grid_apps-assets/uijson/grid_model_merger.ui.json @@ -0,0 +1,220 @@ +{ + "version": "0.1.0", + "title": "Grid Model Merger", + "conda_environment": "grid_apps", + "run_command": "grid_apps.grid_model_merger.driver", + "geoh5": "", + "monitoring_directory": "", + "output_grid": { + "optional": true, + "enabled": false, + "label": "Output grid", + "meshType": [ + "{B020A277-90E2-4CD7-84D6-612EE3F25051}", + "{48f5054a-1c5c-4ca4-9048-80f36dc60a06}", + "{4ea87376-3ece-438b-bf12-3479733ded46}" + ], + "value": "", + "tooltip": "Grid object to transfer models onto.", + "main": true + }, + "scaling_type": { + "choiceList": [ + "linear", + "log" + ], + "main": true, + "label": "Scaling type", + "tooltip": "Select a scaling to apply to the models before merging.", + "value": "log" + }, + "input_a_grid": { + "groupOptional": true, + "enabled": false, + "group": "Input A", + "label": "Object", + "meshType": [ + "{B020A277-90E2-4CD7-84D6-612EE3F25051}", + "{48f5054a-1c5c-4ca4-9048-80f36dc60a06}", + "{4ea87376-3ece-438b-bf12-3479733ded46}" + ], + "value": "", + "tooltip": "Grid object to be merged. The first grid type decides the type of the output grid.", + "main": true + }, + "input_a_model": { + "association": [ + "Cell" + ], + "dataType": [ + "Float", + "Integer" + ], + "group": "Input A", + "main": true, + "label": "Model", + "parent": "input_a_grid", + "enabled": false, + "optional": true, + "value": "" + }, + "input_b_grid": { + "groupOptional": true, + "enabled": false, + "group": "Input B", + "label": "Object", + "meshType": [ + "{B020A277-90E2-4CD7-84D6-612EE3F25051}", + "{48f5054a-1c5c-4ca4-9048-80f36dc60a06}", + "{4ea87376-3ece-438b-bf12-3479733ded46}" + ], + "value": "", + "tooltip": "Grid object to be merged.", + "main": true + }, + "input_b_model": { + "association": [ + "Cell" + ], + "dataType": [ + "Float", + "Integer" + ], + "group": "Input B", + "main": true, + "label": "Model", + "parent": "input_b_grid", + "enabled": false, + "optional": true, + "value": "" + }, + "input_c_grid": { + "groupOptional": true, + "enabled": false, + "group": "Input C", + "label": "Object", + "meshType": [ + "{B020A277-90E2-4CD7-84D6-612EE3F25051}", + "{48f5054a-1c5c-4ca4-9048-80f36dc60a06}", + "{4ea87376-3ece-438b-bf12-3479733ded46}" + ], + "value": "", + "tooltip": "Grid object to be merged.", + "main": true + }, + "input_c_model": { + "association": [ + "Cell" + ], + "dataType": [ + "Float", + "Integer" + ], + "group": "Input C", + "main": true, + "label": "Model", + "parent": "input_c_grid", + "enabled": false, + "optional": true, + "value": "" + }, + "input_d_grid": { + "groupOptional": true, + "enabled": false, + "group": "Input D", + "label": "Object", + "meshType": [ + "{B020A277-90E2-4CD7-84D6-612EE3F25051}", + "{48f5054a-1c5c-4ca4-9048-80f36dc60a06}", + "{4ea87376-3ece-438b-bf12-3479733ded46}" + ], + "value": "", + "tooltip": "Grid object to be merged.", + "main": true + }, + "input_d_model": { + "association": [ + "Cell" + ], + "dataType": [ + "Float", + "Integer" + ], + "group": "Input D", + "main": true, + "label": "Model", + "parent": "input_d_grid", + "enabled": false, + "optional": true, + "value": "" + }, + "input_e_grid": { + "groupOptional": true, + "enabled": false, + "group": "Input E", + "label": "Object", + "meshType": [ + "{B020A277-90E2-4CD7-84D6-612EE3F25051}", + "{48f5054a-1c5c-4ca4-9048-80f36dc60a06}", + "{4ea87376-3ece-438b-bf12-3479733ded46}" + ], + "value": "", + "tooltip": "Grid object to be merged.", + "main": true + }, + "input_e_model": { + "association": [ + "Cell" + ], + "dataType": [ + "Float", + "Integer" + ], + "group": "Input E", + "main": true, + "label": "Model", + "parent": "input_e_grid", + "enabled": false, + "optional": true, + "value": "" + }, + "input_f_grid": { + "groupOptional": true, + "enabled": false, + "group": "Input F", + "label": "Object", + "meshType": [ + "{B020A277-90E2-4CD7-84D6-612EE3F25051}", + "{48f5054a-1c5c-4ca4-9048-80f36dc60a06}", + "{4ea87376-3ece-438b-bf12-3479733ded46}" + ], + "value": "", + "tooltip": "Grid object to be merged.", + "main": true + }, + "input_f_model": { + "association": [ + "Cell" + ], + "dataType": [ + "Float", + "Integer" + ], + "group": "Input F", + "main": true, + "label": "Model", + "parent": "input_f_grid", + "enabled": false, + "optional": true, + "value": "" + }, + "out_group": { + "group": "Output preferences", + "label": "UIJson group", + "value": "", + "groupType": "{BB50AC61-A657-4926-9C82-067658E246A0}", + "visible": true, + "optional": true, + "enabled": false + } +} diff --git a/grid_apps/block_model_to_octree/driver.py b/grid_apps/block_model_to_octree/driver.py index ad5a788..c8d88b0 100644 --- a/grid_apps/block_model_to_octree/driver.py +++ b/grid_apps/block_model_to_octree/driver.py @@ -14,18 +14,16 @@ from pathlib import Path import numpy as np -from discretize import TreeMesh from geoapps_utils.base import Driver as BaseDriver -from geoh5py.data import FloatData, ReferencedData -from geoh5py.objects import BlockModel, Octree +from geoh5py.objects import Octree from geoh5py.ui_json.utils import fetch_active_workspace from scipy.spatial import cKDTree from grid_apps.block_model_to_octree.options import BlockModel2OctreeOptions from grid_apps.utils import ( - block_model_to_discretize, - boundary_value_indices, - tensor_mesh_ordering, + block_model_to_treemesh, + refine_by_cell_volumes, + refine_by_values, treemesh_2_octree, ) @@ -51,50 +49,6 @@ def run(self): return octree - @staticmethod - def block_model_to_treemesh( - entity: BlockModel, diagonal_balance=True, finalize=True - ) -> TreeMesh: - """ - Convert a block model to an octree mesh with the same base cell size and - centered. - - :param entity: BlockModel object to be converted - :param diagonal_balance: Whether to balance the mesh diagonally. - :param finalize: Whether to finalize the treemesh after creation. - - :return: TreeMesh object. - """ - origin = [] - octree_cells = [] - for ii, ax in zip("xyz", "uvz", strict=True): - cell_sizes = np.abs(getattr(entity, f"{ax}_cells")) - h_core = cell_sizes.min() - - # Compute number of octree cells to span the extent - n_c = np.ceil(np.log2(np.sum(cell_sizes) / h_core)) - cell_sizes_octree = np.ones(int(2**n_c)) * h_core - octree_cells.append(cell_sizes_octree) - - # Colocate the center of the octree with the center of the block model - ind_core = np.where(cell_sizes == h_core)[0] - center = ( - entity.origin[ii] - + entity.local_axis_centers(ax)[ind_core[len(ind_core) // 2]] - ) - - axis_center = len(cell_sizes_octree) // 2 - origin.append(center - np.sum(cell_sizes_octree[:axis_center]) - h_core / 2) - - treemesh = TreeMesh( - octree_cells, - x0=origin, - finalize=finalize, - diagonal_balance=diagonal_balance, - ) - - return treemesh - def make_grid(self) -> Octree: """ Convert the block model and output the octree mesh. @@ -104,16 +58,12 @@ def make_grid(self) -> Octree: with fetch_active_workspace(self.params.geoh5, mode="r+"): entity = self.params.entity - treemesh = Driver.block_model_to_treemesh(entity, finalize=False) + treemesh = block_model_to_treemesh(entity, finalize=False) model = None if self.params.data is None: - treemesh = Driver.refine_by_cell_volumes( - treemesh, entity, finalize=True - ) + treemesh = refine_by_cell_volumes(treemesh, entity, finalize=True) else: - treemesh = Driver.refine_by_values( - treemesh, self.params.data, finalize=True - ) + treemesh = refine_by_values(treemesh, self.params.data, finalize=True) # Transfer the model ind = treemesh.get_containing_cells(entity.centroids) model = ( @@ -147,98 +97,6 @@ def make_grid(self) -> Octree: return octree - @staticmethod - def refine_by_cell_volumes( - mesh: TreeMesh, - entity: BlockModel, - finalize: bool = True, - mask: np.ndarray | None = None, - ) -> TreeMesh: - """ - Refine the octree mesh by the cell volumes of the block model. - - :param mesh: TreeMesh object to be refined. - :param entity: BlockModel object to be used for refinement. - :param finalize: Whether to finalize the treemesh after refinement. - :param mask: Optional mask on the block model centroids to apply the refinement over. - - :return: TreeMesh object with refined levels. - """ - if not isinstance(entity, BlockModel): - raise TypeError("entity must be an instance of BlockModel.") - - tensor_oct_level = [] - for ax in "uvz": - cell_sizes = np.abs(getattr(entity, f"{ax}_cells")) - h_core = cell_sizes.min() - # Find the core region - tensor_oct_level.append(np.log2(cell_sizes / h_core).astype(int)) - - e_x, e_y, e_z = np.meshgrid(*tensor_oct_level) - max_level = np.c_[np.ravel(e_x), np.ravel(e_y), np.ravel(e_z)].max(axis=1) - - locations = entity.centroids - if mask is not None: - locations = locations[mask] - max_level = max_level[mask] - - mesh.insert_cells(locations, mesh.max_level - max_level, finalize=finalize) - - return mesh - - @staticmethod - def refine_by_values( - mesh: TreeMesh, data: FloatData | ReferencedData, finalize=True - ) -> TreeMesh: - """ - Increase the mesh resolution based on the gradient of data values. - - :param mesh: Input TreeMesh object. - :param data: FloatData or ReferencedData object containing the values to - be used for refinement. - :param finalize: Whether to finalize the treemesh after refinement. - - :return: TreeMesh object with refined levels. - """ - if not isinstance(data, FloatData | ReferencedData): - raise TypeError( - "Argument 'data' must be an instance of FloatData or ReferencedData." - ) - - entity = data.parent - if not isinstance(entity, BlockModel): - raise TypeError("The parent of 'data' must be an instance of BlockModel.") - - tensor = block_model_to_discretize(entity) - indices = tensor_mesh_ordering(entity) - - gradients = np.abs(tensor.cell_gradient @ data.values[indices]) - levels = np.zeros(gradients.shape, dtype=int) - isnan = np.isnan(gradients) - - if isinstance(data, FloatData): - actives = gradients[~isnan] - bins = np.percentile( - actives[actives > 0], np.linspace(5, 95, mesh.max_level) - ) - levels[~isnan] = np.searchsorted(bins, actives) - else: - levels[gradients > 0] = mesh.max_level - - # Refine on the value/nan interface, without boundary cells - if any(isnan): - horizon = boundary_value_indices( - tensor, data.values[indices], data.nan_value - ) - mesh = Driver.refine_by_cell_volumes( - mesh, entity, finalize=False, mask=horizon[np.argsort(indices)] - ) - - locs = tensor.average_cell_to_face @ tensor.cell_centers - mesh.insert_cells(locs[~isnan], levels[~isnan].astype(int), finalize=finalize) - - return mesh - if __name__ == "__main__": file = Path(sys.argv[1]).resolve() diff --git a/grid_apps/block_models/driver.py b/grid_apps/block_models/driver.py index d8064c6..c29f4d9 100644 --- a/grid_apps/block_models/driver.py +++ b/grid_apps/block_models/driver.py @@ -22,6 +22,7 @@ from scipy.spatial import cKDTree from grid_apps.block_models.options import BlockModelOptions +from grid_apps.utils import tensor_to_block_model logger = logging.getLogger(__name__) @@ -188,14 +189,7 @@ def get_block_model( # pylint: disable=too-many-arguments, too-many-positional- expansion_factor=expansion_factor, ) - object_out = BlockModel.create( - workspace, - origin=[mesh.x0[0], mesh.x0[1], mesh.x0[2] + mesh.h[2].sum()], - u_cell_delimiters=mesh.nodes_x - mesh.x0[0], - v_cell_delimiters=mesh.nodes_y - mesh.x0[1], - z_cell_delimiters=-(mesh.x0[2] + mesh.h[2].sum() - mesh.nodes_z[::-1]), - name=name, - ) + object_out = tensor_to_block_model(workspace, mesh, name=name) return object_out diff --git a/grid_apps/grid_model_merger/__init__.py b/grid_apps/grid_model_merger/__init__.py new file mode 100644 index 0000000..7e91710 --- /dev/null +++ b/grid_apps/grid_model_merger/__init__.py @@ -0,0 +1,8 @@ +# ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' +# Copyright (c) 2024-2026 Mira Geoscience Ltd. ' +# ' +# This file is part of grid-apps package. ' +# ' +# grid-apps is distributed under the terms and conditions of the MIT License ' +# (see LICENSE file at the root of this source code package). ' +# ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' diff --git a/grid_apps/grid_model_merger/driver.py b/grid_apps/grid_model_merger/driver.py new file mode 100644 index 0000000..1dbeca4 --- /dev/null +++ b/grid_apps/grid_model_merger/driver.py @@ -0,0 +1,227 @@ +# ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' +# Copyright (c) 2024-2026 Mira Geoscience Ltd. ' +# ' +# This file is part of grid-apps package. ' +# ' +# grid-apps 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 logging +import sys +from pathlib import Path + +import numpy as np +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.shared.utils import fetch_active_workspace +from scipy.spatial import cKDTree + +from grid_apps.grid_model_merger.options import GridModelMergerOptions, ScalingTypeEnum +from grid_apps.utils import ( + containing_cell_indices, + get_boundary_active_cells, + refine_tree_by_mesh, + tensor_to_block_model, + treemesh_2_octree, +) + + +logger = logging.getLogger(__name__) + + +class Driver(BaseDriver): + """ + Merge multiple grids and models from selection. + + :param params: Options for merging multiple grids and models. + """ + + _params_class = GridModelMergerOptions + + def __init__(self, params: GridModelMergerOptions): + + super().__init__(params) + self.output_grid = None + + def run(self): + """Create an octree mesh from input values.""" + with fetch_active_workspace(self.params.geoh5, mode="r+"): + self.output_grid = self.get_output_grid() + self.interpolate_models_to_output_grid() + output = self.params.out_group or self.output_grid + self.update_monitoring_directory(output) + logger.info("Done.") + + return self.output_grid + + def get_global_mesh_specs(self) -> tuple[np.ndarray, np.ndarray]: + """ + Loop through selections to get global mesh specifications for output grid. + + :return: Array of outer extent and core cell dimensions. + """ + extent = np.vstack([[np.inf] * 3, [-np.inf] * 3]) + cell_size = np.hstack([np.inf] * 3) + + for selection in self.params.selections: + grid = selection.grid + + extent[0, :] = np.min([extent[0, :], grid.extent[0, :]], axis=0) + extent[1, :] = np.max([extent[1, :], grid.extent[1, :]], axis=0) + + if isinstance(grid, Octree): + cell_size = np.min( + [ + cell_size, + np.r_[grid.u_cell_size, grid.v_cell_size, grid.w_cell_size], + ], + axis=0, + ) + else: + cell_size = np.min( + [ + cell_size, + np.r_[ + grid.u_cells.min(), + grid.v_cells.min(), + grid.z_cells.min(), + ], + ], + axis=0, + ) + + return extent, np.abs(cell_size) + + def get_output_grid(self): + """ + Make block model object from input data. + """ + with fetch_active_workspace(self.params.geoh5, mode="r+"): + if self.params.output_grid is not None: + if self.out_group is not None: + return self.params.output_grid.copy( + parent=self.out_group, copy_children=False + ) + + return self.params.output_grid + + extent, cell_size = self.get_global_mesh_specs() + + # Use type of the first entry + mesh_type = type(self.params.selections[0].grid) + + logger.info("Merging selected grids to '%s' . . .", mesh_type.__name__) + + mesh = mesh_utils.mesh_builder_xyz( + extent, + cell_size, + mesh_type="tree" if mesh_type is Octree else "tensor", + tree_diagonal_balance=True, + ) + + if mesh_type is Octree: + for selection in self.params.selections: + mesh = refine_tree_by_mesh(mesh, selection.grid, finalize=False) + + mesh.finalize() + output_grid = treemesh_2_octree( + self.params.geoh5, mesh, parent=self.params.out_group + ) + else: + output_grid = tensor_to_block_model( + self.params.geoh5, mesh, parent=self.params.out_group + ) + + return output_grid + + def interpolate_models_to_output_grid(self): + """ + Interpolate models from selections to output grid. + """ + with fetch_active_workspace(self.params.geoh5, mode="r+"): + out_model = np.full(self.output_grid.n_cells, np.nan, dtype=float) + weights = np.full(self.output_grid.n_cells, np.nan, dtype=float) + threshold = None + + for selection in self.params.selections: + mesh, model = selection.to_discretize() + + if model is None or not np.any(~np.isnan(model)): + continue + + active = ~np.isnan(model) + + logger.info( + "Interpolating model '%s' from grid '%s' to output grid '%s' . . .", + selection.model.name, + selection.grid.name, + self.output_grid.name, + ) + + if self.params.scaling_type == ScalingTypeEnum.LOG: + if threshold is None: + threshold = np.percentile(np.abs(model[active]), 10) + + model = symlog(model, threshold=threshold) + + active_boundary = get_boundary_active_cells( + mesh, active, horizontal_edges=True + ) + weight_model = np.full(active.shape[0], np.nan, dtype=float) + weight_model[active] = self.cosine_taper_weights( + mesh.cell_centers[active_boundary], mesh.cell_centers[active] + ) + + # Find nearest neighbors and apply weighted model + cell_id = containing_cell_indices(mesh, self.output_grid.centroids) + + # Trim weights for cells outside the extent of the input mesh + cell_weights = np.full(self.output_grid.n_cells, np.nan, dtype=float) + cell_weights[cell_id != -1] = weight_model[cell_id][cell_id != -1] + + out_model = np.nansum( + [out_model, cell_weights * model[cell_id]], axis=0 + ) + weights = np.nansum([weights, cell_weights], axis=0) + + # Normalizes weighted sum + non_zero = weights > 0 + out_model[non_zero] /= weights[non_zero] + out_model[~non_zero] = np.nan + + if self.params.scaling_type == ScalingTypeEnum.LOG: + out_model = inv_symlog(out_model, threshold=threshold) + + if np.any(~np.isnan(out_model)): + self.output_grid.add_data( + { + "merged_model": { + "values": out_model, + "entity_type": self.params.selections[0].model.entity_type, + } + } + ) + + @staticmethod + def cosine_taper_weights(edge_locations: np.ndarray, target_locations: np.ndarray): + # Compute weights based on distance to boundary cells + tree = cKDTree(edge_locations) + 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 + + # Find nearest neighbors and apply weighted model + del tree + + return cosine_taper + + +if __name__ == "__main__": + file = Path(sys.argv[1]).resolve() + Driver.start(file) diff --git a/grid_apps/grid_model_merger/options.py b/grid_apps/grid_model_merger/options.py new file mode 100644 index 0000000..3949a4d --- /dev/null +++ b/grid_apps/grid_model_merger/options.py @@ -0,0 +1,159 @@ +# ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' +# Copyright (c) 2024-2026 Mira Geoscience Ltd. ' +# ' +# This file is part of grid-apps package. ' +# ' +# grid-apps 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 string +from enum import StrEnum +from pathlib import Path +from typing import Any, ClassVar + +import numpy as np +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.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, + octree_2_treemesh, + tensor_mesh_ordering, +) + + +class ScalingTypeEnum(StrEnum): + LINEAR = "linear" + LOG = "log" + + +class GridModelMergerOptions(Options): + """ + Block model parameters for use with `block_models.driver`. + + :param output_grid: Grid object to hold the merged model. + :param scaling_type: Scaling type for the merged model, either linear or log. + Interpolation performed in log scale attenuates extreme values. + :param selections: List of grid and model selections. + """ + + name: ClassVar[str] = "grid_model_merger" + default_ui_json: ClassVar[Path] = assets_path() / "uijson/grid_model_merger.ui.json" + title: str = "Grid Model Merger" + run_command: str = "grid_apps.grid_model_merger.driver" + + conda_environment: str = "grid_apps" + output_grid: GridObject | None = None + scaling_type: ScalingTypeEnum = ScalingTypeEnum.LOG + selections: list[MeshModelSelection] + + @model_validator(mode="before") + @classmethod + def collect_selections(cls, values: dict): + """Collect selections from the input dictionary.""" + if "selections" not in values: + selections = collect_selections_from_dict(values) + values["selections"] = selections + return values + + @model_serializer(mode="wrap") + def distribute_selections(self, handler, info): + """Convert selections to a individual parameters.""" + dump = handler(self, info) + selections = dump.pop("selections") + refinement_params: dict[str, Any] = {} + for i, group in enumerate(selections): + group_id = string.ascii_lowercase[i] + if group is None: + refinement_params[f"input_{group_id}_grid"] = None + refinement_params[f"input_{group_id}_model"] = None + else: + for param, value in group.items(): + param_name = f"input_{group_id}_{param}" + refinement_params[param_name] = value + + return dict(dump, **refinement_params) + + @classmethod + def collect_input_from_dict(cls, model: type[BaseModel], data: dict[str, Any]): + """ + Recursively replace BaseModel objects with dictionary of 'data' values. + + :param model: BaseModel object holding data and possibly other nested + BaseModel objects. + :param data: Dictionary of parameters and values without nesting structure. + """ + + update = super().collect_input_from_dict(model, data) + update["selections"] = collect_selections_from_dict(data) + + return update + + +class MeshModelSelection(BaseModel): + """ + Input parameters. + + :param grid: Grid object with data. + :param model: NumericalData (model) entity. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + grid: GridObject + model: NumericData | None = None + + def to_discretize(self) -> tuple[TreeMesh | TensorMesh, np.ndarray | None]: + """ + Convert the geoh5 entity to discretized mesh object. + + :return: Discretize mesh object and model if provided. + """ + model = None + mesh = None + if self.model is not None: + model = self.model.values + + if isinstance(self.grid, BlockModel): + if model is not None: + model = model[tensor_mesh_ordering(self.grid)] + mesh = block_model_to_tensor(self.grid) + + elif isinstance(self.grid, Octree): + mesh = octree_2_treemesh(self.grid) + + if mesh is None: + raise TypeError(f"Mesh type {type(self.grid)} currently not supported.") + + return mesh, model + + +def collect_selections_from_dict(data: dict) -> list[dict | None]: + """Collect active input dictionaries from input dictionary.""" + inputs: list[dict | None] = [] + for identifier in input_identifiers(data): + input_params = {} + for param in ["grid", "model"]: + input_name = f"input_{identifier}_{param}" + value = data.get(input_name, None) + + if value is not None: + input_params[param] = value + + if input_params.get("grid", None) is not None: + inputs.append(input_params) + + return inputs + + +def input_identifiers(data: dict) -> list[str]: + """Return identifiers for active inputs (object not none).""" + active = [k for k in data if "_grid" in k] + return np.unique([k.split("_")[1] for k in active]) diff --git a/grid_apps/utils.py b/grid_apps/utils.py index 5445eca..ffa9f64 100644 --- a/grid_apps/utils.py +++ b/grid_apps/utils.py @@ -14,23 +14,29 @@ import numpy as np 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.ui_json.utils import fetch_active_workspace +from pydantic import ConfigDict, validate_call from scipy.interpolate import interp1d +from scipy.sparse import find from scipy.spatial import cKDTree -def block_model_to_discretize( +typed_call = validate_call(config=ConfigDict(arbitrary_types_allowed=True)) + + +@typed_call +def block_model_to_tensor( entity: BlockModel, -) -> TensorMesh | tuple[TensorMesh, np.ndarray]: +) -> TensorMesh: """ Convert a block model to a discretize.TensorMesh. :param entity: The block model to convert. - """ - if not isinstance(entity, BlockModel): - raise TypeError("entity must be an instance of BlockModel.") + :return: An equivalent TensorMesh object. + """ origin = [ entity.origin["x"] + entity.u_cells[entity.u_cells < 0].sum(), entity.origin["y"] + entity.v_cells[entity.v_cells < 0].sum(), @@ -47,35 +53,78 @@ def block_model_to_discretize( return mesh -def boundary_value_indices( - mesh: TensorMesh | TreeMesh, values: np.ndarray, target: float | int -) -> np.ndarray: +@typed_call +def tensor_to_block_model( + workspace: Workspace, mesh: TensorMesh, **kwargs +) -> BlockModel: """ - Get a mask of the boundary cells in a mesh based on a target value. + Convert a tensor mesh to a block model. - :param mesh: The discretize mesh. - :param values: The values associated with the cells. - :param target: The target value to identify boundary cells. + :param workspace: Workspace to create the block model. + :param mesh: Tensor mesh object from discretize + :param kwargs: Extra parameters to pass to the block model. - :return: Mask of boundary cells. + :return: BlockModel entity. """ - if not isinstance(mesh, TensorMesh | TreeMesh): - raise TypeError("Mesh must be an instance of TensorMesh or TreeMesh.") + block_model = BlockModel.create( + workspace, + origin=[mesh.x0[0], mesh.x0[1], mesh.x0[2] + mesh.h[2].sum()], + u_cell_delimiters=mesh.nodes_x - mesh.x0[0], + v_cell_delimiters=mesh.nodes_y - mesh.x0[1], + z_cell_delimiters=-(mesh.x0[2] + mesh.h[2].sum() - mesh.nodes_z[::-1]), + **kwargs, + ) + return block_model - if not isinstance(values, np.ndarray): - raise TypeError("Values must be a numpy array.") - if target is np.nan: - is_target = np.isnan(values) - else: - is_target = values == target +@typed_call +def block_model_to_treemesh( + entity: BlockModel, diagonal_balance=True, finalize=True +) -> TreeMesh: + """ + Convert a block model to an octree mesh with the same base cell size and + centered. - on_face = (mesh.cell_gradient @ is_target).astype(bool) - boundary_cells = (mesh.average_face_to_cell @ on_face).astype(bool) + :param entity: BlockModel object to be converted + :param diagonal_balance: Whether to balance the mesh diagonally. + :param finalize: Whether to finalize the treemesh after creation. - return boundary_cells + :return: TreeMesh object. + """ + origin = [] + octree_cells = [] + for ii, ax in zip("xyz", "uvz", strict=True): + cell_sizes = np.abs(getattr(entity, f"{ax}_cells")) + h_core = cell_sizes.min() + + # Compute number of octree cells to span the extent + n_c = np.ceil(np.log2(np.sum(cell_sizes) / h_core)) + cell_sizes_octree = np.ones(int(2**n_c)) * h_core + octree_cells.append(cell_sizes_octree) + + # Colocate the center of the octree with the center of the block model + ind_core = np.where(np.isclose(cell_sizes, h_core, atol=1e-1))[0] + center = ( + entity.origin[ii] + + entity.local_axis_centers(ax)[ind_core[len(ind_core) // 2]] + ) + axis_center = len(cell_sizes_octree) // 2 + origin.append(center - np.sum(cell_sizes_octree[:axis_center]) - h_core / 2) + treemesh = TreeMesh( + octree_cells, + x0=origin, + diagonal_balance=diagonal_balance, + ) + + if finalize: + treemesh.finalize() + + return treemesh + + +@typed_call def collocate_octrees(global_mesh: Octree, local_meshes: list[Octree]): """ Collocate a list of octree meshes into a global octree mesh. @@ -121,13 +170,44 @@ def collocate_octrees(global_mesh: Octree, local_meshes: list[Octree]): workspace.update_attribute(local_mesh, "attributes") +@typed_call +def containing_cell_indices(mesh: TreeMesh | TensorMesh, locations: np.ndarray): + """ + Return indices of cells containing the list of locations. + + Points that do not intersect return -1. + + :param mesh: Input discretize mesh object. + :param locations: Input locations array. + + :return: Array of indices + """ + if isinstance(mesh, TreeMesh): + indices = mesh.get_containing_cells(locations) + + 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[~mesh.is_inside(locations)] = -1 + + return indices + + +@typed_call def create_octree_from_octrees(meshes: list[Octree | TreeMesh]) -> TreeMesh: """ Create an all encompassing octree mesh from a list of meshes. :param meshes: List of Octree or TreeMesh meshes. - :return octree: A global Octree. + :return: An all-encompassing TreeMesh object """ cell_size = [] dimensions = None @@ -160,29 +240,132 @@ def create_octree_from_octrees(meshes: list[Octree | TreeMesh]) -> TreeMesh: treemesh = TreeMesh(cells, origin=origin, diagonal_balance=False) for mesh in meshes: - if isinstance(mesh, Octree) and mesh.octree_cells is not None: - centers = mesh.centroids - levels = treemesh.max_level - np.log2(mesh.octree_cells["NCells"]) - elif isinstance(mesh, TreeMesh) and mesh.cell_centers is not None: - centers = mesh.cell_centers - levels = ( - treemesh.max_level - - mesh.max_level - + mesh.cell_levels_by_index(np.arange(mesh.nC)) - ) - else: - raise TypeError( - f"All meshes must be Octree or TreeMesh, not {type(mesh)} " - "and must have octree cells defined." - ) - - treemesh.insert_cells(centers, levels, finalize=False) + treemesh = refine_tree_by_mesh(treemesh, mesh) treemesh.finalize() return treemesh +@typed_call +def refine_tree_by_mesh( + tree: TreeMesh, mesh: TreeMesh | Octree | BlockModel, finalize: bool = False +) -> TreeMesh: + """ + Given a TreeMesh, insert cells at the corresponding octree level. + + :param tree: TreeMesh to be refined. + :param mesh: Input TreeMesh or Octree mesh to refine with. + :param finalize: Whether to finalize the refined mesh. + :return: Refined mesh. + """ + if isinstance(mesh, BlockModel): + treemesh = block_model_to_treemesh(mesh, finalize=False) + mesh = refine_by_cell_volumes(treemesh, mesh) + + if isinstance(mesh, Octree) and mesh.octree_cells is not None: + centers = mesh.centroids + levels = tree.max_level - np.log2(mesh.octree_cells["NCells"]) + + else: + centers = mesh.cell_centers + levels = ( + tree.max_level + - mesh.max_level + + mesh.cell_levels_by_index(np.arange(mesh.nC)) + ) + + tree.insert_cells(centers, levels, finalize=finalize) + + return tree + + +@typed_call +def refine_by_cell_volumes( + mesh: TreeMesh, + entity: BlockModel, + finalize: bool = True, + mask: np.ndarray | None = None, +) -> TreeMesh: + """ + Refine the octree mesh by the cell volumes of the block model. + + :param mesh: TreeMesh object to be refined. + :param entity: BlockModel object to be used for refinement. + :param finalize: Whether to finalize the treemesh after refinement. + :param mask: Optional mask on the block model centroids to apply the refinement over. + + :return: TreeMesh object with refined levels. + """ + tensor_oct_level = [] + for ax in "uvz": + cell_sizes = np.abs(getattr(entity, f"{ax}_cells")) + h_core = cell_sizes.min() + # Find the core region + tensor_oct_level.append(np.log2(cell_sizes / h_core).astype(int)) + + e_x, e_y, e_z = np.meshgrid(*tensor_oct_level) + max_level = np.c_[np.ravel(e_x), np.ravel(e_y), np.ravel(e_z)].max(axis=1) + + locations = entity.centroids + if mask is not None: + locations = locations[mask] + max_level = max_level[mask] + + mesh.insert_cells(locations, mesh.max_level - max_level, finalize=finalize) + + return mesh + + +@typed_call +def refine_by_values( + mesh: TreeMesh, data: FloatData | ReferencedData, finalize=True +) -> TreeMesh: + """ + Increase the mesh resolution based on the gradient of data values. + + :param mesh: Input TreeMesh object. + :param data: FloatData or ReferencedData object containing the values to + be used for refinement. + :param finalize: Whether to finalize the treemesh after refinement. + + :return: TreeMesh object with refined levels. + """ + entity = data.parent + + if not isinstance(entity, BlockModel): + raise TypeError("The parent of 'data' must be an instance of BlockModel.") + + tensor = block_model_to_tensor(entity) + indices = tensor_mesh_ordering(entity) + + gradients = np.abs(tensor.cell_gradient @ data.values[indices]) + levels = np.zeros(gradients.shape, dtype=int) + isnan = np.isnan(gradients) + + if isinstance(data, FloatData): + actives = gradients[~isnan] + bins = np.percentile(actives[actives > 0], np.linspace(5, 95, mesh.max_level)) + levels[~isnan] = np.searchsorted(bins, actives) + else: + levels[gradients > 0] = mesh.max_level + + # Refine on the value/nan interface, without boundary cells + if any(isnan): + horizon = get_boundary_active_cells( + tensor, data.values[indices] == data.nan_value + ) + mesh = refine_by_cell_volumes( + mesh, entity, finalize=False, mask=horizon[np.argsort(indices)] + ) + + locs = tensor.average_cell_to_face @ tensor.cell_centers + mesh.insert_cells(locs[~isnan], levels[~isnan].astype(int), finalize=finalize) + + return mesh + + +@typed_call def densify_curve(curve: Curve, increment: float) -> np.ndarray: """ Refine a curve by adding points along the curve at a given increment. @@ -218,6 +401,8 @@ def find_endpoints(points: np.ndarray) -> np.ndarray: Find the endpoints of a co-linear array of points. :param points: locations array of shape (n, 3). + + :return: Array of shape (n, 2) containing the endpoints. """ xmin = points[:, 0].min() @@ -235,6 +420,47 @@ def find_endpoints(points: np.ndarray) -> np.ndarray: return np.array(endpoints) +@typed_call +def get_boundary_active_cells( + mesh: TreeMesh | TensorMesh, + actives: np.ndarray, + horizontal_edges: bool = False, + vertical_edges: bool = False, +) -> np.ndarray: + """ + Given a mesh and a set of active cells, return the active cells + that are on the boundary of the active domain. + + :param mesh: Tree or TensorMesh object. + :param actives: Bool array of active cells. + :param horizontal_edges: Include the cells on the horizontal edges of the mesh. + :param vertical_edges: Include the cells on the top and bottom edges of the mesh. + + :return: Bool array of boundary cells of the active domain. + """ + if actives.ndim != 1 or actives.shape[0] != mesh.n_cells: + raise ValueError("Input array 'actives' must have length mesh.n_cells.") + + is_face = np.zeros_like(actives, dtype=bool) + + # Find actives horizontal mesh boundary cells + if horizontal_edges: + for face in mesh.cell_boundary_indices[:-2]: + is_face[face] = True + + if vertical_edges: + for face in mesh.cell_boundary_indices[-2:]: + is_face[face] = True + + # Find boundary active cells + face_diff = ~np.isclose(mesh.stencil_cell_gradient @ actives, 0, atol=0.1) + _, cols, _ = find(mesh.stencil_cell_gradient[face_diff, :]) + is_face[cols] = True + + return is_face & actives + + +@typed_call def get_neighbouring_cells(mesh: TreeMesh, indices: list | np.ndarray) -> tuple: """ Get the indices of neighbouring cells along a given axis for a given list of @@ -248,12 +474,6 @@ def get_neighbouring_cells(mesh: TreeMesh, indices: list | np.ndarray) -> tuple: axis[1] = (south, north) axis[2] = (down, up) """ - if not isinstance(indices, list | np.ndarray): - raise TypeError("Input 'indices' must be a list or numpy.ndarray of indices.") - - if not isinstance(mesh, TreeMesh): - raise TypeError("Input 'mesh' must be a discretize.TreeMesh object.") - neighbors: dict[int, list] = {ax: [[], []] for ax in range(mesh.dim)} for ind in indices: @@ -267,16 +487,15 @@ def get_neighbouring_cells(mesh: TreeMesh, indices: list | np.ndarray) -> tuple: ) +@typed_call def get_octree_attributes(mesh: Octree | TreeMesh) -> dict[str, list]: """ Get mesh attributes. :param mesh: Input Octree or TreeMesh object. + :return mesh_attributes: Dictionary of mesh attributes. """ - if not isinstance(mesh, Octree | TreeMesh): - raise TypeError(f"All meshes must be Octree or TreeMesh, not {type(mesh)}") - cell_size = [] cell_count = [] dimensions = [] @@ -308,6 +527,7 @@ def get_octree_attributes(mesh: Octree | TreeMesh) -> dict[str, list]: } +@typed_call def octree_2_treemesh( # pylint: disable=too-many-locals mesh: Octree, ) -> TreeMesh | None: @@ -397,6 +617,8 @@ def surface_strip( strip. The surrounding strip will be 2*width wider and longer than the input points. :param name: Name of the new Points objects. + + :return: New points object """ assert points.locations is not None @@ -430,6 +652,7 @@ def surface_strip( return Points.create(points.workspace, vertices=vertices, name=name) +@typed_call def tensor_mesh_ordering( entity: BlockModel, ) -> np.ndarray: @@ -440,9 +663,6 @@ def tensor_mesh_ordering( :return indices: Array of indices to reorder cell-based values. """ - if not isinstance(entity, BlockModel): - raise TypeError("mesh must be an instance of BlockModel.") - indices = np.arange(entity.n_cells) indices = indices.reshape( ( @@ -461,6 +681,7 @@ def tensor_mesh_ordering( return indices +@typed_call def treemesh_2_octree(workspace: Workspace, treemesh: TreeMesh, **kwargs) -> Octree: """ Converts a :obj:`discretize.TreeMesh` to :obj:`geoh5py.objects.Octree` entity. diff --git a/tests/run_tests/block_2_octree_test.py b/tests/run_tests/block_2_octree_test.py index 9487380..4a6494d 100644 --- a/tests/run_tests/block_2_octree_test.py +++ b/tests/run_tests/block_2_octree_test.py @@ -109,7 +109,7 @@ def test_float_refine_octree(tmp_path): driver = BlockModelToOctreeDriver(params) octree = driver.make_grid() - assert octree.n_cells == 7491 + assert octree.n_cells == 6140 def test_integer_refine_octree(tmp_path): diff --git a/tests/run_tests/grid_model_merger_test.py b/tests/run_tests/grid_model_merger_test.py new file mode 100644 index 0000000..dcd6ec9 --- /dev/null +++ b/tests/run_tests/grid_model_merger_test.py @@ -0,0 +1,178 @@ +# ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' +# Copyright (c) 2024-2026 Mira Geoscience Ltd. ' +# ' +# This file is part of grid-apps package. ' +# ' +# grid-apps 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 + +import numpy as np +from geoh5py.objects import BlockModel, Octree, Points +from geoh5py.workspace import Workspace + +from grid_apps.block_models.driver import Driver as BlockModelDriver +from grid_apps.grid_model_merger.driver import Driver +from grid_apps.grid_model_merger.options import GridModelMergerOptions +from grid_apps.octree_creation.driver import OctreeDriver +from grid_apps.octree_creation.options import OctreeOptions + + +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 + depth_core = 300.0 + height = 300 + width = 1000 + n = 100 + + x_grid, y_grid = np.meshgrid(np.arange(0, width, n), np.arange(0, height, n)) + z_grid = np.around((top / 2) * np.sin(x_grid) + (top / 2), -1) + locs = np.c_[x_grid.ravel(), y_grid.ravel(), z_grid.ravel()] + pads = [100, 150, 200, 300, 0, 0] + mesh = BlockModelDriver.get_block_model( + workspace, locs, [50, 50, 50], depth_core, pads, 1.1, name="test" + ) + return mesh + + +def setup_octree(workspace, locations, refinement, params_dict) -> Octree: + points = Points.create(workspace, vertices=locations) + + params_dict.update( + { + "geoh5": workspace, + "objects": points, + "u_cell_size": 25.0, + "v_cell_size": 25.0, + "w_cell_size": 25.0, + "refinements": [ + { + "refinement_object": points, + "levels": refinement, + "horizon": False, + } + ], + } + ) + params = OctreeOptions(**params_dict) + driver = OctreeDriver(params) + return driver.run() + + +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)) + 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() + + merged_model = out_grid.children[0] + np.testing.assert_almost_equal(merged_model.values[2649], 1.5, decimal=1) + + # Repeat with a hole in the first model + values = model_a.values + values[(mesh.centroids[:, 0] > 750) & (mesh.centroids[:, 1] > 250)] = np.nan + model_a.values = values + + out_grid = driver.run() + merged_model = out_grid.children[0] + np.testing.assert_almost_equal(merged_model.values[2649], 2.0, decimal=1) + + +def test_merge_octree_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: + mesh = setup_octree(ws, locations, refinement, params_dict) + other = mesh.copy(origin=(-600, -600, mesh.origin["z"])) + + 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() + + merged_model = out_grid.children[0] + np.testing.assert_almost_equal(merged_model.values[1558], 1.5, decimal=2) + + # Repeat with a hole in the first model + values = model_a.values + values[(mesh.centroids[:, 0] > 200) & (mesh.centroids[:, 1] > 200)] = np.nan + model_a.values = values + + out_grid = driver.run() + merged_model = out_grid.children[0] + 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)}} + ) + + options = GridModelMergerOptions.build( + { + "geoh5": ws, + "input_a_grid": octree, + "input_b_grid": block_model, + "input_a_model": model_a, + "input_b_model": model_b, + } + ) + + 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) diff --git a/tests/utils_test.py b/tests/utils_test.py index 04b1efd..8db5f73 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -14,16 +14,17 @@ from discretize import TreeMesh from geoh5py import Workspace from geoh5py.objects import BlockModel, Curve, Octree, Points +from pydantic import ValidationError from pytest import raises from grid_apps.block_models.driver import Driver as BlockModelDriver from grid_apps.utils import ( - block_model_to_discretize, - boundary_value_indices, + block_model_to_tensor, collocate_octrees, create_octree_from_octrees, densify_curve, find_endpoints, + get_boundary_active_cells, get_neighbouring_cells, get_octree_attributes, octree_2_treemesh, @@ -134,17 +135,17 @@ def test_block_model_to_discretize(tmp_path): name="TestBlockModel", ) - with raises(TypeError): - block_model_to_discretize("abc") + with raises(ValidationError): + block_model_to_tensor("abc") - tensor = block_model_to_discretize(block_model) + tensor = block_model_to_tensor(block_model) indices = tensor_mesh_ordering(block_model) # Check the shape of the discretized points np.testing.assert_allclose(block_model.centroids[indices], tensor.cell_centers) -def test_tensor_boundary_value_indices(tmp_path): +def test_tensor_get_boundary_active_cells(tmp_path): # Create a test block model h5file_path = tmp_path / f"{__name__}.geoh5" with Workspace.create(h5file_path) as workspace: @@ -155,51 +156,75 @@ def test_tensor_boundary_value_indices(tmp_path): z_cell_delimiters=np.cumsum(np.ones(16)), ) - tensor = block_model_to_discretize(block) + tensor = block_model_to_tensor(block) values = np.ones(tensor.n_cells) - values[int(tensor.n_cells / 2)] = 2 - indices = boundary_value_indices(tensor, values, 2) + first_quadrant = np.all(tensor.cell_centers > 8, axis=1) + values[first_quadrant] = 2 - assert indices.sum() == 7 + indices = get_boundary_active_cells(tensor, values == 2) # Just for visual validation block.add_data( { + "model": { + "values": values[np.argsort(tensor_mesh_ordering(block))], + }, "boundary_indices": { "values": indices[np.argsort(tensor_mesh_ordering(block))], - } + }, } ) + assert indices.sum() == 127 -def test_octree_boundary_value_indices(tmp_path): +def test_octree_get_boundary_active_cells(tmp_path): treemesh = TreeMesh([16, 16, 16]) treemesh.refine(4, finalize=True) values = np.ones(treemesh.n_cells) - values[7] = 2 - with pytest.raises(TypeError, match="Mesh must be an instance"): - indices = boundary_value_indices("abc", values, 2) + first_quadrant = np.all(treemesh.cell_centers > 0.5, axis=1) + values[first_quadrant] = 2 - with pytest.raises(TypeError, match="Values must be a numpy array"): - indices = boundary_value_indices(treemesh, 123, 2) + with pytest.raises(ValidationError): + get_boundary_active_cells("abc", values == 2) + + with pytest.raises(ValidationError): + indices = get_boundary_active_cells(treemesh, True) + + indices = get_boundary_active_cells(treemesh, values == 2) + + assert indices.sum() == 169 + + # Repeat with horizontal boundary cells + indices_he = get_boundary_active_cells(treemesh, values == 2, horizontal_edges=True) + assert indices_he.sum() == 260 + + # Repeat with vertical boundary cells + indices_ve = get_boundary_active_cells(treemesh, values == 2, vertical_edges=True) + assert indices_ve.sum() == 218 - indices = boundary_value_indices(treemesh, values, 2) h5file_path = tmp_path / f"{__name__}.geoh5" with Workspace.create(h5file_path) as workspace: octree = treemesh_2_octree(workspace, treemesh, name="TestOctree") octree.add_data( { + "model": { + "values": values, + }, "boundary_indices": { "values": indices, - } + }, + "horizontal_boundary": { + "values": indices_he, + }, + "vertical_boundary": { + "values": indices_ve, + }, } ) - assert indices.sum() == 7 - def test_find_endpoints(): locs = np.array([[-1, -1, 0], [0, 0, 0], [1, 1, 0]]) @@ -387,14 +412,10 @@ def test_get_neighbouring_cells(): mesh.insert_cells([100, 100, 100], mesh.max_level, finalize=True) ind = mesh.get_containing_cells([95.0, 95.0, 95.0]) - with pytest.raises( - TypeError, match=r"Input 'indices' must be a list or numpy.ndarray of indices\." - ): + with pytest.raises(ValidationError): get_neighbouring_cells(mesh, ind) - with pytest.raises( - TypeError, match=r"Input 'mesh' must be a discretize.TreeMesh object\." - ): + with pytest.raises(ValidationError): get_neighbouring_cells(1, [ind]) neighbours = get_neighbouring_cells(mesh, [ind])