refactor neighbor/corner finding and add unit test

Author: benk-mira

simpeg_drivers/utils/regularization.py

@@ -38,52 +38,32 @@ def cell_neighbors_along_axis(mesh: TreeMesh, axis: str) -> np.ndarray:
     return np.sort(stencil_indices, axis=1)
 
 
-def cell_neighbors(mesh: TreeMesh) -> np.ndarray:
-    """Find all cell neighbors in a TreeMesh."""
-
-    x_neighbors = cell_neighbors_along_axis(mesh, "x")
-    x_neighbors_backward = np.fliplr(x_neighbors)
-    y_neighbors = cell_neighbors_along_axis(mesh, "y")
-    y_neighbors_backward = np.fliplr(y_neighbors)
-    max_index = np.max([x_neighbors.max(), y_neighbors.max()])
-    if mesh.dim == 3:
-        z_neighbors = cell_neighbors_along_axis(mesh, "z")
-        z_neighbors_backward = np.fliplr(z_neighbors)
-        max_index = np.max([max_index, z_neighbors.max()])
-
+def collect_all_neighbors(neighbors, neighbors_backwards, corners, corners_backwards):
     all_neighbors = []  # Store
-    x_adjacent = np.ones(max_index + 1, dtype="int") * -1
-    y_adjacent = np.ones(max_index + 1, dtype="int") * -1
-    x_adjacent_backward = np.ones(max_index + 1, dtype="int") * -1
-    y_adjacent_backward = np.ones(max_index + 1, dtype="int") * -1
-
-    x_adjacent[y_neighbors[:, 0]] = y_neighbors[:, 1]
-    y_adjacent[x_neighbors[:, 1]] = x_neighbors[:, 0]
 
-    x_adjacent_backward[y_neighbors_backward[:, 0]] = y_neighbors_backward[:, 1]
-    y_adjacent_backward[x_neighbors_backward[:, 1]] = x_neighbors_backward[:, 0]
+    all_neighbors += [neighbors[0]]
+    all_neighbors += [neighbors[1]]
 
-    all_neighbors += [x_neighbors]
-    all_neighbors += [y_neighbors]
+    all_neighbors += [np.c_[neighbors[0][:, 0], corners[0][neighbors[0][:, 1]]]]
+    all_neighbors += [np.c_[neighbors[0][:, 1], corners[0][neighbors[0][:, 0]]]]
 
-    all_neighbors += [np.c_[x_neighbors[:, 0], x_adjacent[x_neighbors[:, 1]]]]
-    all_neighbors += [np.c_[x_neighbors[:, 1], x_adjacent[x_neighbors[:, 0]]]]
-
-    all_neighbors += [np.c_[y_adjacent[y_neighbors[:, 0]], y_neighbors[:, 1]]]
-    all_neighbors += [np.c_[y_adjacent[y_neighbors[:, 1]], y_neighbors[:, 0]]]
+    all_neighbors += [np.c_[corners[1][neighbors[1][:, 0]], neighbors[1][:, 1]]]
+    all_neighbors += [np.c_[corners[1][neighbors[1][:, 1]], neighbors[1][:, 0]]]
 
     # Repeat backward for Treemesh
-    all_neighbors += [x_neighbors_backward]
-    all_neighbors += [y_neighbors_backward]
+    all_neighbors += [neighbors_backwards[0]]
+    all_neighbors += [neighbors_backwards[1]]
 
     all_neighbors += [
         np.c_[
-            x_neighbors_backward[:, 0], x_adjacent_backward[x_neighbors_backward[:, 1]]
+            neighbors_backwards[0][:, 0],
+            corners_backwards[0][neighbors_backwards[0][:, 1]],
         ]
     ]
     all_neighbors += [
         np.c_[
-            x_neighbors_backward[:, 1], x_adjacent_backward[x_neighbors_backward[:, 0]]
+            neighbors_backwards[0][:, 1],
+            corners_backwards[0][neighbors_backwards[0][:, 0]],
         ]
     ]
 
@@ -97,25 +77,20 @@ def cell_neighbors(mesh: TreeMesh) -> np.ndarray:
     ]
 
     # Use all the neighbours on the xy plane to find neighbours in z
-    if mesh.dim == 3:
+    if len(neighbors) == 3:
         all_neighbors_z = []
-        z_adjacent = np.ones(max_index + 1, dtype="int") * -1
-        z_adjacent_backward = np.ones(max_index + 1, dtype="int") * -1
-
-        z_adjacent[z_neighbors[:, 0]] = z_neighbors[:, 1]
-        z_adjacent_backward[z_neighbors_backward[:, 0]] = z_neighbors_backward[:, 1]
 
-        all_neighbors_z += [z_neighbors]
-        all_neighbors_z += [z_neighbors_backward]
+        all_neighbors_z += [neighbors[2]]
+        all_neighbors_z += [neighbors_backwards[2]]
 
-        all_neighbors_z += [np.c_[all_neighbors[:, 0], z_adjacent[all_neighbors[:, 1]]]]
-        all_neighbors_z += [np.c_[all_neighbors[:, 1], z_adjacent[all_neighbors[:, 0]]]]
+        all_neighbors_z += [np.c_[all_neighbors[:, 0], corners[2][all_neighbors[:, 1]]]]
+        all_neighbors_z += [np.c_[all_neighbors[:, 1], corners[2][all_neighbors[:, 0]]]]
 
         all_neighbors_z += [
-            np.c_[all_neighbors[:, 0], z_adjacent_backward[all_neighbors[:, 1]]]
+            np.c_[all_neighbors[:, 0], corners_backwards[2][all_neighbors[:, 1]]]
         ]
         all_neighbors_z += [
-            np.c_[all_neighbors[:, 1], z_adjacent_backward[all_neighbors[:, 0]]]
+            np.c_[all_neighbors[:, 1], corners_backwards[2][all_neighbors[:, 0]]]
         ]
 
         # Stack all and keep only unique pairs
@@ -130,6 +105,39 @@ def cell_neighbors(mesh: TreeMesh) -> np.ndarray:
     return all_neighbors
 
 
+def cell_adjacent(neighbors: list[np.ndarray]) -> list[np.ndarray]:
+    """Find all cell corners from cell neighbor array."""
+
+    dim = len(neighbors)
+    max_index = np.max(neighbors)
+    corners = -1 * np.ones((dim, max_index + 1), dtype="int")
+
+    corners[0, neighbors[1][:, 0]] = neighbors[1][:, 1]
+    corners[1, neighbors[0][:, 1]] = neighbors[0][:, 0]
+    if dim == 3:
+        corners[2, neighbors[2][:, 0]] = neighbors[2][:, 1]
+
+    return [np.array(k) for k in corners.tolist()]
+
+
+def cell_neighbors(mesh: TreeMesh) -> np.ndarray:
+    """Find all cell neighbors in a TreeMesh."""
+
+    neighbors = []
+    neighbors.append(cell_neighbors_along_axis(mesh, "x"))
+    neighbors.append(cell_neighbors_along_axis(mesh, "y"))
+    if mesh.dim == 3:
+        neighbors.append(cell_neighbors_along_axis(mesh, "z"))
+
+    neighbors_backwards = [np.fliplr(k) for k in neighbors]
+    corners = cell_adjacent(neighbors)
+    corners_backwards = cell_adjacent(neighbors_backwards)
+
+    return collect_all_neighbors(
+        neighbors, neighbors_backwards, corners, corners_backwards
+    )
+
+
 def rotate_xz_2d(mesh: TreeMesh, phi: np.ndarray) -> ssp.csr_matrix:
     """
     Create a 2d ellipsoidal rotation matrix for the xz plane.

tests/utils_regularization_test.py

@@ -0,0 +1,92 @@
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+#  Copyright (c) 2025 Mira Geoscience Ltd.                                          '
+#                                                                                   '
+#  This file is part of simpeg-drivers package.                                     '
+#                                                                                   '
+#  simpeg-drivers is distributed under the terms and conditions of the MIT License  '
+#  (see LICENSE file at the root of this source code package).                      '
+#                                                                                   '
+# '''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+
+import numpy as np
+from discretize import TreeMesh
+
+from simpeg_drivers.utils.regularization import (
+    cell_adjacent,
+    cell_neighbors_along_axis,
+    collect_all_neighbors,
+)
+
+
+def get_mesh():
+    mesh = TreeMesh(h=[[10.0] * 4, [10.0] * 4, [10.0] * 4], diagonal_balance=False)
+    mesh.refine(2)
+    return mesh
+
+
+def test_cell_neighbors_along_axis():
+    mesh = get_mesh()
+    centers = mesh.cell_centers
+    neighbors = cell_neighbors_along_axis(mesh, "x")
+    assert np.allclose(centers[7], [15.0, 15.0, 15.0])
+    assert np.allclose(
+        centers[neighbors[neighbors[:, 0] == 7][0][1]], [25.0, 15.0, 15.0]
+    )
+    assert np.allclose(
+        centers[neighbors[neighbors[:, 1] == 7][0][0]], [5.0, 15.0, 15.0]
+    )
+    neighbors = cell_neighbors_along_axis(mesh, "y")
+    assert np.allclose(
+        centers[neighbors[neighbors[:, 0] == 7][0][1]], [15.0, 25.0, 15.0]
+    )
+    assert np.allclose(
+        centers[neighbors[neighbors[:, 1] == 7][0][0]], [15.0, 5.0, 15.0]
+    )
+    neighbors = cell_neighbors_along_axis(mesh, "z")
+    assert np.allclose(
+        centers[neighbors[neighbors[:, 0] == 7][0][1]], [15.0, 15.0, 25.0]
+    )
+    assert np.allclose(
+        centers[neighbors[neighbors[:, 1] == 7][0][0]], [15.0, 15.0, 5.0]
+    )
+
+
+def test_collect_all_neighbors():
+    mesh = get_mesh()
+    centers = mesh.cell_centers
+    neighbors = [cell_neighbors_along_axis(mesh, k) for k in "xyz"]
+    neighbors_bck = [np.fliplr(k) for k in neighbors]
+    corners = cell_adjacent(neighbors)
+    corners_bck = cell_adjacent(neighbors_bck)
+    all_neighbors = collect_all_neighbors(
+        neighbors, neighbors_bck, corners, corners_bck
+    )
+    assert np.allclose(centers[7], [15.0, 15.0, 15.0])
+    neighbor_centers = centers[all_neighbors[all_neighbors[:, 0] == 7][:, 1]].tolist()
+    assert [5, 5, 5] in neighbor_centers
+    assert [15, 5, 5] in neighbor_centers
+    assert [25, 5, 5] in neighbor_centers
+    assert [5, 15, 5] in neighbor_centers
+    assert [15, 15, 5] in neighbor_centers
+    assert [25, 15, 5] in neighbor_centers
+    assert [5, 25, 5] in neighbor_centers
+    assert [15, 25, 5] in neighbor_centers
+    assert [25, 25, 5] in neighbor_centers
+    assert [5, 5, 15] in neighbor_centers
+    assert [15, 5, 15] in neighbor_centers
+    assert [25, 5, 15] in neighbor_centers
+    assert [5, 15, 15] in neighbor_centers
+    assert [25, 15, 15] in neighbor_centers
+    assert [5, 25, 15] in neighbor_centers
+    assert [15, 25, 15] in neighbor_centers
+    assert [25, 25, 15] in neighbor_centers
+    assert [5, 5, 25] in neighbor_centers
+    assert [15, 5, 25] in neighbor_centers
+    assert [25, 5, 25] in neighbor_centers
+    assert [5, 15, 25] in neighbor_centers
+    assert [15, 15, 25] in neighbor_centers
+    assert [25, 15, 25] in neighbor_centers
+    assert [5, 25, 25] in neighbor_centers
+    assert [15, 25, 25] in neighbor_centers
+    assert [25, 25, 25] in neighbor_centers
+    assert [15, 15, 15] not in neighbor_centers