Source code for torx.analysis.lineouts.on_grid_m

"""Checks whether a point lies on the numerical grid."""
import numpy as np
from numba import njit
from torx.autodoc_decorators_m import autodoc_function



@autodoc_function
@njit
def on_grid(
    r_grid: np.ndarray,
    z_grid: np.ndarray,
    r_test: np.ndarray,
    z_test: np.ndarray,
    spacing,
    expansion: float = 1.01,
):
    """
    Return the distance from a sample point to its four nearest grid neighbors.

    May be very expensive for large grids.
    """
    # Check which points are within a grid-spacing of a grid point in each dimension
    r_match = np.abs(r_grid - np.expand_dims(r_test, axis=-1)) < spacing * expansion
    z_match = np.abs(z_grid - np.expand_dims(z_test, axis=-1)) < spacing * expansion

    # Check which points are within a grid-spacing of a grid point in both dimensions
    match_both = np.logical_and(r_match, z_match)

    # Return True if a sample point has at least 4 grid-points within a grid-spacing
    return np.count_nonzero(match_both, axis=match_both.ndim - 1) >= 4