Source code for torx.analysis.fourier.fourier_base_m

"""Basic Fourier analysis routines."""
import xarray as xr
import numpy as np
from scipy import fft
from scipy import interpolate
from torx.autodoc_decorators_m import autodoc_function



@autodoc_function
def equally_sampled_interval(x_vals, n_samples):
    """Return the equally sampled interval within the given x values."""
    if(type(x_vals) == xr.DataArray):
        x_vals = x_vals.values

    x_sampled, x_spacing = np.linspace(
        np.min(x_vals), np.max(x_vals), num=n_samples, retstep=True
    )
    return x_sampled, x_spacing




@autodoc_function
def equally_sampled_data(x_vals, y_vals, n_samples):
    """
    Return samples of equally-spaced points within the given values.

    Fast fourier transform algorithm requires equally spaced data.

    Supplying n_samples as a power of 2 improves the computational efficiency
    of the FFT.
    """
    # Sort the data
    order = np.argsort(x_vals)
    x_sorted = x_vals[order]
    y_sorted = y_vals[order]

    # Remove repeat values
    zero_positions = np.where(np.diff(x_sorted) == 0.0)
    x_sorted = np.delete(x_sorted, zero_positions)
    y_sorted = np.delete(y_sorted, zero_positions)

    # Resample the data at evenly spaced points
    x_interpolated, x_spacing = equally_sampled_interval(x_vals, n_samples)
    # Interpolate the sorted data on the new interval
    y_interpolated = interpolate.interp1d(x_sorted, y_sorted, kind="cubic")(
        x_interpolated
    )

    return x_interpolated, y_interpolated, x_spacing




@autodoc_function
def fourier_transform_1d(
    y_vals: np.ndarray,
    x_vals: np.ndarray,
    n_samples: int,
    apply_window: bool=False,
    remove_mean: bool=True,
    window_beta: int=14.0,
    normalize: bool=True
):
    """
    Perform a one dimensional fourier transform on real, unequally-spaced data.

    The function will interpolate the given data onto an equally-spaced x-grid,
    and then perform the fourier transform. The returned complex_amplitudes and
    frequencies are given only for positive frequencies, since real input data
    is assumed.

    If 'apply_window" is false (e.g. for flux surfaces in the scrape-of layer),
    a window function is applied to make the signal 'apply_window'.

    Note
    ----
    The order of x_vals and y_vals is switched compared to standard ordering,
    in order to pass x_vals as a keyword argument to xr.apply_ufunc.
    """
    # Perform interpolation (cubic by default)
    _, y_interp, _ = equally_sampled_data(x_vals, y_vals, n_samples)

    if remove_mean:
        y_interp = y_interp - np.mean(y_interp)

    if normalize:
        norm = "forward"
        norm_fac = 2.0 * np.pi
    else:
        norm = "backward"
        norm_fac = 1.0

    # Make signal apply_window if in open field line region
    if apply_window:
        window = np.kaiser(n_samples, beta=window_beta)
        amplitude_full = fft.rfft(y_interp * window, norm=norm)
    else:
        amplitude_full = fft.rfft(y_interp, norm=norm)

    # RFFT works on real data and will only return the positive part of the
    # spectrum, thus we do not need to remove the negative part (as would be
    # the case when using fft).
    # However we crop the data to floor(n_samples / 2) because the length may
    # vary depending on even and odd n_samples.
    amplitude_pos = 2.0 * amplitude_full[: (n_samples // 2)]

    return amplitude_pos * norm_fac




@autodoc_function
def fourier_frequencies(x_vals, n_samples):
    """Compute the FFT-frequencies from the sampling points x_vals."""
    # Get the spacing for the rfftfreq command
    _, spacing = equally_sampled_interval(x_vals, n_samples)
    frequencies = fft.rfftfreq(n_samples, d=spacing)
    # Do the same cropping as in fourier transform to get the correct length
    return frequencies[: (n_samples // 2)]