Source code for torx.basic_functions_m

"""Simple functions which are commonly used."""
import numpy as np
import xarray as xr
from torx.autodoc_decorators_m import autodoc_function



@autodoc_function
def smoothstep(test_point: float, step_center: float, step_width: float, order=3):
    """
    Return a hermite smoothstep of a specific order.

    The step will be exactly 0 for test_point values less than step_center - step_width / 2
    The step will be exactly 1 for test_point values greater than step_center + step_width / 2

    For a step of order A, the derivative of order (A-1) will also be zero
    at the endpoints.
    """
    xn = (np.atleast_1d(test_point) - step_center) / step_width + 0.5

    # Ignore the runtime warning that comes from comparing NaN
    # NaN will return False in any comparison
    with np.errstate(invalid="ignore"):
        xn[xn <= 0] = 0
        xn[xn >= 1] = 1

    if order == 1:
        step_values = xn
    elif order == 2:
        step_values = -2.0 * xn**3 + 3.0 * xn**2
    elif order == 3:
        step_values = 6.0 * xn**5 - 15.0 * xn**4 + 10.0 * xn**3
    elif order == 4:
        step_values = -20.0 * xn**7 + 70.0 * xn**6 - 84.0 * xn**5 + 35.0 * xn**4
    elif order == 5:
        step_values = (
            70.0 * xn**9
            - 315.0 * xn**8
            + 540.0 * xn**7
            - 420.0 * xn**6
            + 126.0 * xn**5
        )
    elif order == 6:
        step_values = (
            -252.0 * xn**11
            + 1386.0 * xn**10
            - 3080.0 * xn**9
            + 3465.0 * xn**8
            - 1980.0 * xn**7
            + 462.0 * xn**6
        )
    elif order == 7:
        step_values = (
            924.0 * xn**13
            - 6006.0 * xn**12
            + 16380.0 * xn**11
            - 24024.0 * xn**10
            + 20020.0 * xn**9
            - 9009.0 * xn**8
            + 1716.0 * xn**7
        )
    else:
        return NotImplemented

    if isinstance(test_point, xr.DataArray):
        return xr.DataArray(
            step_values,
            coords=test_point.coords,
            dims=test_point.dims,
            attrs=test_point.attrs,
        )
    else:
        return step_values




@autodoc_function
def find_nearest_index(array, value):
    """Return the index in the array that is nearest to the value."""
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return idx




@autodoc_function
def second_smallest(array):
    """Find the second smallest, unique entry in an array."""
    m1 = m2 = float("inf")
    for x in array.flat:
        if x < m1:
            m1, m2 = x, m1
        elif x < m2 and x != m1:
            m2 = x
    return m2