Remove duplicate octree_extent function. Update extent for ndarray

Author: domfournier

simpeg_drivers/components/topography.py

@@ -39,7 +39,6 @@
     get_containing_cells,
     get_neighbouring_cells,
     mask_vertices_and_cells,
-    octree_extents,
 )
 
 
@@ -113,9 +112,8 @@ def active_cells(self, mesh: InversionMesh, data: InversionData) -> np.ndarray:
                     self.params.active_cells.topography_object, "cells", None
                 )
             else:
-                extent = octree_extents(mesh.entity)[:4]
                 vertices, cells = mask_vertices_and_cells(
-                    extent.ravel(order="F"),
+                    mesh.entity.extent,
                     self.locations,
                     getattr(self.params.active_cells.topography_object, "cells", None),
                 )

simpeg_drivers/utils/utils.py

@@ -12,7 +12,6 @@
 from __future__ import annotations
 
 import multiprocessing
-from collections.abc import Sequence
 from copy import deepcopy
 from pathlib import Path
 from typing import TYPE_CHECKING
@@ -34,7 +33,7 @@
 )
 from geoh5py.objects.surveys.electromagnetics.base import LargeLoopGroundEMSurvey
 from geoh5py.shared import INTEGER_NDV
-from geoh5py.shared.utils import fetch_active_workspace, stringify
+from geoh5py.shared.utils import fetch_active_workspace, mask_by_extent, stringify
 from geoh5py.ui_json import InputFile
 from grid_apps.utils import octree_2_treemesh
 from scipy.interpolate import interp1d
@@ -48,46 +47,20 @@
     from simpeg_drivers.driver import InversionDriver
 
 
-def octree_extents(octree: Octree) -> np.ndarray:
-    """
-    Get the true extents of an octree (min/max of the perimeter).
-
-    The octree.extents property returns min/max of the centroids
-
-    :param octree: Octree mesh object.
-
-    :returns: Array of [xmin, xmax, ymin, ymax].
-    """
-
-    origin = np.array(list(octree.origin.tolist()))
-    span = np.array(
-        [
-            getattr(octree, f"{axis}_cell_size") * getattr(octree, f"{axis}_count")
-            for axis in "uvw"
-        ]
-    )
-
-    return np.stack([origin, origin + span]).flatten(order="F")
-
-
 def mask_vertices_and_cells(
-    extent: Sequence, vertices: np.ndarray, cells: np.ndarray | None
+    extent: np.ndarray, vertices: np.ndarray, cells: np.ndarray | None
 ) -> tuple[np.ndarray, np.ndarray]:
     """
     Mask vertices and remove cells whose vertices are all outside the extent.
 
-    :param extent: Array-like object of [xmin, xmax, ymin, ymax].
+    :param extent: Array-like object of [[xmin, ymin], [xmax, ymax]].
     :param vertices: Array of shape (n_vertices, 3) containing the x, y, z coordinates.
     :param cells: Array of shape (n_cells, 3) containing the indices of the vertices
         that make up each cell.
     """
 
-    vertex_mask = (
-        (vertices[:, 0] >= extent[0])
-        & (vertices[:, 0] <= extent[1])
-        & (vertices[:, 1] >= extent[2])
-        & (vertices[:, 1] <= extent[3])
-    )
+    vertex_mask = mask_by_extent(vertices, extent=extent)
+
     if cells is None:
         return vertices[vertex_mask], None
 

tests/utils_test.py

@@ -9,33 +9,8 @@
 # '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
 
 import numpy as np
-from geoh5py import Workspace
-from geoh5py.objects import Octree, Points
-from grid_apps.octree_creation.driver import OctreeDriver
-from grid_apps.octree_creation.options import OctreeOptions, RefinementOptions
 
-from simpeg_drivers.utils.utils import mask_vertices_and_cells, octree_extents
-
-
-def test_octree_extents(tmp_path):
-    with Workspace(tmp_path / "test.geoh5") as ws:
-        X, Y = np.meshgrid(np.linspace(0, 1000, 51), np.linspace(0, 1000, 51))
-        Z = np.zeros_like(X)
-        vertices = np.column_stack([X.ravel(), Y.ravel(), Z.ravel()])
-        pts = Points.create(ws, name="points", vertices=vertices)
-        options = OctreeOptions(
-            geoh5=ws,
-            objects=pts,
-            refinements=[
-                RefinementOptions(
-                    refinement_object=pts, levels=[4, 2], horizon=False, distance=100
-                ),
-            ],
-        )
-        octree = OctreeDriver.octree_from_params(options)
-
-    extents = octree_extents(octree)
-    assert np.allclose(extents, [-1112.5, 2087.5, -1112.5, 2087.5, -1062.5, 537.5])
+from simpeg_drivers.utils.utils import mask_vertices_and_cells
 
 
 def test_mask_vertices_and_cells():
@@ -54,11 +29,11 @@ def test_mask_vertices_and_cells():
             [7, 5, 8],
         ]
     )
-    extent = [0.5, 2, 0, 2, 0, 1]
+    extent = np.vstack([[0.5, 0, 0], [2, 2, 1]])
     masked_vertices, masked_cells = mask_vertices_and_cells(extent, vertices, cells)
     assert len(masked_vertices) == len(vertices)
     assert len(masked_cells) == len(cells)
-    extent = [1.5, 2, 0, 2, 0, 1]
+    extent = np.vstack([[1.5, 0, 0], [2, 2, 1]])
     masked_vertices, masked_cells = mask_vertices_and_cells(extent, vertices, cells)
     assert len(masked_vertices) == 6
     assert len(masked_cells) == 4