Source code for torx.analysis.fourier.fourier_analysis_m

"""Routines for performing Fourier analysis on data."""
from .fourier_base_m import fourier_transform_1d
import xarray as xr
import numpy as np
from torx.autodoc_decorators_m import autodoc_function



@autodoc_function
def fourier_analysis(
    snaps: xr.Dataset,
    lineout,
    variables: str,
    n_samples: int,
    rho: float,
    allow_rechunk: bool = True,
    remove_mean: bool = True,
    no_interp: bool = False
):
    """Perform a basic fourier analysis of the given datasets on the lineout."""
    # Drop variables that you don't want to perform fourier analysis on
    snaps = snaps[variables]

    # Interpolate the snaps onto the lineout (the resulting snaps have
    # dimension [phi, tau, interp_points])
    if(not no_interp):
        snaps = lineout.interpolate(snaps)

    # NOTE: Field-alignment in the confined region is given by the symmetry
    #       angle, while in the open field line region we do the analysis on
    #       the normalized arc length.
    if rho < 1:
        x_vals = lineout.symmetry_angle / (2.0 * np.pi)
        apply_window = False
    elif rho > 1:
        arclen = lineout.poloidal_arc_length()
        x_vals = arclen / np.max(arclen)
        apply_window = True
    else:
        raise NotImplementedError(
            "Fourier analysis not implemented on the separatrix rho=1!"
        )

    # Dictionaries for apply_ufunc
    kwargs_dict = dict(
        x_vals=x_vals,
        n_samples=n_samples,
        apply_window=apply_window,
        remove_mean=remove_mean,
        normalize=True
    )
    dask_gufunc_kwargs_dict = dict(
        output_sizes=dict(mode_number=n_samples // 2),
        allow_rechunk=allow_rechunk,
    )
    attrs_dict = dict(
        rho=rho,
    )

    # Parallel ufunc over 1D-Fourier transform
    fourier_ds = xr.apply_ufunc(
        fourier_transform_1d,
        snaps.chunk({"interp_points": -1}),
        input_core_dims=[["interp_points"],],
        output_core_dims=[["mode_number"]],
        kwargs=kwargs_dict,
        vectorize=True,
        dask="parallelized",
        dask_gufunc_kwargs=dask_gufunc_kwargs_dict,
        output_dtypes=[np.cdouble]
    ).assign_attrs(attrs_dict)

    return fourier_ds




@autodoc_function
def cross_spectrum_analysis(
    fourier_ds: xr.Dataset,
    variables: str,
    reference: str
):
    """
    Perform a cross spectrum analysis between variables and a reference.

    The xr.Dataset is assumed to contain already Fourier analyzed variables.
    """
    cross_spectrum_ds = fourier_ds[variables] * fourier_ds[reference].conj()

    return cross_spectrum_ds.assign_attrs(fourier_ds.attrs)