diff --git a/geoapps_utils/utils/locations.py b/geoapps_utils/utils/locations.py index a8a9d958..97bf3e6b 100644 --- a/geoapps_utils/utils/locations.py +++ b/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. diff --git a/tests/locations_test.py b/tests/locations_test.py index 1309bb25..9277e7a1 100644 --- a/tests/locations_test.py +++ b/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(