Source code for torx.analysis.smoothing.smoothing_1d_m

"""Provides smoothing functionality for 1D arrays."""
import warnings
import numpy as np
from scipy.interpolate import interp1d
from torx.decorators import autodoc_function



@autodoc_function
def smooth_1d_block(array: np.ndarray, axis: np.ndarray,
                    block_length: int):
    """
    Smooth a 1D array by averaging over non-overlapping blocks.

    The array is split into consecutive blocks of block_length, each block
    is collapsed to its mean, and the result is linearly interpolated back
    onto the original axis. Neighboring blocks do not overlap, so the
    output is piecewise linear on the block scale and the endpoints are
    extrapolated rather than preserved. See smooth_1d_window for an
    overlapping full-resolution alternative.

    Block length should be picked sensibly since too small blocks have little
    effect and too big blocks will smooth out large scale structures.
    Can be used for easier interpretation of noisy signals.

    Parameters
    ----------
    array : numpy.ndarray
        Input signal.
    axis : numpy.ndarray
        Monotonic independent variable, same shape as array.
    block_length : int
        Number of samples averaged per block. A value of 1 is a no-op.

    Returns
    -------
    numpy.ndarray
        Smoothed signal, same shape as array.
    """
    assert len(array.shape) == 1, "Only 1D arrays can be smoothed"
    assert array.shape == axis.shape, "Axis and array shapes do not match"

    # Check for common denominator
    if len(array) % block_length:
        blocked_array = np.zeros((len(array) // block_length + 1))
        blocked_axis = np.copy(blocked_array)
    else:
        blocked_array = np.zeros((len(array) // block_length))
        blocked_axis = np.copy(blocked_array)

    # Average array values for each block
    for i in range(len(blocked_array)):
        # The last block is possibly shorter, so take everything left.
        if i == len(blocked_array) - 1:
            blocked_array[i] = np.mean(array[i * block_length:])
            blocked_axis[i] = np.mean(axis[i * block_length:])
        else:
            block = slice(i * block_length, (i + 1) * block_length)
            blocked_array[i] = np.mean(array[block])
            blocked_axis[i] = np.mean(axis[block])

    return interp1d(
        blocked_axis, blocked_array, kind="linear", fill_value="extrapolate"
    )(axis)




@autodoc_function
def smooth_1d_window(array: np.ndarray,
                     smooth_length: int) -> np.ndarray:
    """
    Smooth a 1D array with a sliding, overlapping averaging window.

    Every sample is replaced by the mean of the window centered on it, so
    consecutive windows overlap and the output keeps the full resolution
    of the input. The window length is forced odd and the window shrinks
    symmetrically toward both ends so the first and last points stay
    unchanged. Useful for de-spiking the ragged outer field-line-traced
    flux surfaces in the scrape-off layer. See smooth_1d_block for the
    non-overlapping variant.

    No axis is required: the window is defined in sample index space, so
    the input is assumed to be evenly sampled.

    Parameters
    ----------
    array : numpy.ndarray
        Input signal.
    smooth_length : int
        Length of the averaging window. Even values are rounded down to
        the next odd value so the window stays symmetric about each
        point. A length below 3 (or an array shorter than 3) returns a
        copy unchanged.

    Returns
    -------
    numpy.ndarray
        Smoothed signal, same shape as ``array``.

    Warns
    -----
    UserWarning
        If the effective window is below 3, or the array has fewer than
        3 points. Either case is a no-op, so the warning distinguishes
        it from smoothing that ran but had little effect. Note that an
        even smooth_length is rounded down first, so a value of 2 warns
        while 4 does not.
    """
    y = np.asarray(array, float)
    n = y.size
    requested_length = smooth_length
    if smooth_length % 2 == 0:
        smooth_length -= 1
    if smooth_length < 3:
        warnings.warn(
            f"smooth_length={requested_length} gives an effective window of "
            f"{smooth_length}, which is below the minimum of 3. No smoothing "
            "is applied and the input is returned unchanged.",
            UserWarning, stacklevel=2)
        return y.copy()
    if n < 3:
        warnings.warn(
            f"An array of {n} point(s) is shorter than the minimum window of "
            "3. No smoothing is applied and the input is returned unchanged.",
            UserWarning, stacklevel=2)
        return y.copy()
    half = smooth_length // 2
    # Prefix sums so a window mean is a single subtraction.
    cumulative = np.concatenate([[0.0], np.cumsum(y)])
    smoothed = y.copy()
    for i in range(n):
        # Symmetric half-window, shrinking toward both ends.
        m = min(i, n - 1 - i, half)
        if m > 0:
            smoothed[i] = (cumulative[i + m + 1] - cumulative[i - m]) \
                / (2 * m + 1)
    return smoothed