tests passing

Author: benk-mira

simpeg_drivers/components/topography.py

@@ -38,6 +38,8 @@
     floating_active,
     get_containing_cells,
     get_neighbouring_cells,
+    mask_vertices_and_cells,
+    octree_extents,
 )
 
 
@@ -100,25 +102,26 @@ def active_cells(self, mesh: InversionMesh, data: InversionData) -> np.ndarray:
             active_cells = InversionModel.obj_2_mesh(
                 self.params.active_cells.active_model, mesh.entity
             )
+
         else:
-            topography = self.params.active_cells.topography_object
-            locations = getattr(topography, "centroids", None) or self.locations
-            xmin, xmax, ymin, ymax, zmin, zmax = mesh.entity.extent.ravel(order="F")
-            mask = (
-                (locations[:, 0] > xmin)
-                & (locations[:, 0] < xmax)
-                & (locations[:, 1] > ymin)
-                & (locations[:, 1] < ymax)
-                & (locations[:, 2] > zmin)
-                & (locations[:, 2] < zmax)
-            )
+            if any(k in self.params.inversion_type for k in ["2d", "p3d"]):
+                vertices = self.locations
+                cells = getattr(
+                    self.params.active_cells.topography_object, "cells", None
+                )
+            else:
+                extent = octree_extents(mesh.entity)
+                vertices, cells = mask_vertices_and_cells(
+                    extent.ravel(order="F"),
+                    self.locations,
+                    getattr(self.params.active_cells.topography_object, "cells", None),
+                )
 
-            cells = getattr(topography, "cells", None)
             active_cells = active_from_xyz(
                 mesh.entity,
-                locations,
+                vertices,
                 grid_reference="bottom" if forced_to_surface else "center",
-                triangulation=cells[mask, :] or None,
+                triangulation=cells,
             )
 
         active_cells = (mesh.permutation @ active_cells).astype(bool)

simpeg_drivers/utils/utils.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from copy import deepcopy
 from typing import TYPE_CHECKING
 
@@ -43,6 +44,62 @@
     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].
+    """
+
+    def half_cell(axis):
+        return (
+            getattr(octree, f"{axis}_cell_size") * octree.octree_cells["NCells"]
+        ) / 2
+
+    xmin = (octree.centroids[:, 0] - half_cell("u")).min()
+    xmax = (octree.centroids[:, 0] + half_cell("u")).max()
+    ymin = (octree.centroids[:, 1] - half_cell("v")).min()
+    ymax = (octree.centroids[:, 1] + half_cell("v")).max()
+
+    return np.array([xmin, xmax, ymin, ymax])
+
+
+def mask_vertices_and_cells(
+    extent: Sequence, 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 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])
+    )
+    if cells is None:
+        return vertices[vertex_mask], None
+
+    cell_mask = np.any(vertex_mask[cells], axis=1)
+    vertex_mask = np.zeros_like(vertex_mask, dtype=bool)
+    vertex_mask[cells[cell_mask].flatten()] = True
+
+    new_cells = cells.copy()[cell_mask]
+    cell_map = np.arange(len(vertices))[vertex_mask]
+    new_cells = np.searchsorted(cell_map, new_cells)
+
+    return vertices[vertex_mask], new_cells
+
+
 def calculate_2D_trend(
     points: np.ndarray, values: np.ndarray, order: int = 0, method: str = "all"
 ):
@@ -495,7 +552,7 @@ def active_from_xyz(
         raise ValueError("'grid_reference' must be one of 'center', 'top', or 'bottom'")
 
     # Return the active cell array
-    return mask_under_horizon(locations, topo, triangulation=triangulation)
+    return mask_under_horizon(locations, horizon=topo, triangulation=triangulation)
 
 
 def truncate_locs_depths(locs: np.ndarray, depth_core: float) -> np.ndarray: