GEOPY-1302: Organize geoapps.utils.testing.setup_inversion_workspace

geoapps_utils/utils/locations.py

@@ -15,12 +15,46 @@
 import numpy as np
 from geoh5py import Workspace
 from geoh5py.data import Data
-from geoh5py.objects import Grid2D, Points
+from geoh5py.objects import CellObject, Grid2D, Points
 from geoh5py.objects.grid_object import GridObject
 from scipy.interpolate import LinearNDInterpolator
 from scipy.spatial import Delaunay, cKDTree
 
 
+def gaussian(
+    x: np.ndarray, y: np.ndarray, amplitude: float, width: float
+) -> np.ndarray:
+    """
+    Gaussian function for 2D data.
+
+    :param x: X-coordinates.
+    :param y: Y-coordinates.
+    :param amplitude: Amplitude of the Gaussian.
+    :param width: Width parameter of the Gaussian.
+    """
+
+    return amplitude * np.exp(-0.5 * ((x / width) ** 2.0 + (y / width) ** 2.0))
+
+
+def mask_large_connections(cell_object: CellObject, distance_threshold: float):
+    """
+    Trim connections in cell based objects.
+
+    :param cell_object: Cell object containing segments with small vertex spacing
+        along-line, but large spacing between segments.
+
+    :return: Cleaned object without cells exceeding the distance threshold.
+    """
+
+    dist = np.linalg.norm(
+        cell_object.vertices[cell_object.cells[:, 0], :]
+        - cell_object.vertices[cell_object.cells[:, 1], :],
+        axis=1,
+    )
+
+    return np.where(dist > distance_threshold)[0]
+
+
 def mask_under_horizon(locations: np.ndarray, horizon: np.ndarray) -> np.ndarray:
     """
     Mask locations under a horizon.

tests/locations_test.py

@@ -12,17 +12,39 @@
 
 import numpy as np
 from geoh5py import Workspace
-from geoh5py.objects import Grid2D, Points
+from geoh5py.objects import Curve, Grid2D, Points
 
 from geoapps_utils.utils.locations import (
+    gaussian,
     get_locations,
     get_overlapping_limits,
     map_indices_to_coordinates,
+    mask_large_connections,
     mask_under_horizon,
 )
 from geoapps_utils.utils.transformations import rotate_points, z_rotation_matrix
 
 
+def test_gaussian():
+    x = np.linspace(-10, 10, 100)
+    y = np.linspace(-10, 10, 100)
+    x_grid, y_grid = np.meshgrid(x, y)
+    z_grid = gaussian(x_grid, y_grid, 10, 5)
+    assert np.isclose(z_grid.max(), 10, rtol=1e-3)
+
+
+def test_mask_large_connections(tmp_path):
+    with Workspace(tmp_path / "test.geoh5") as ws:
+        x = np.linspace(0, 100, 11)
+        y = np.linspace(0, 300, 4)
+        x_grid, y_grid = np.meshgrid(x, y)
+        z_grid = np.zeros_like(x_grid)
+        vertices = np.column_stack([x_grid.ravel(), y_grid.ravel(), z_grid.ravel()])
+        crv = Curve.create(ws, name="test_curve", vertices=vertices)
+        mask = mask_large_connections(crv, distance_threshold=50.0)
+        assert len(mask) == 3
+
+
 def test_rotate_points():
     points = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1], [2, 3, 4]])
     validation = rotate_points(