Source code for torx.analysis.lineouts.lineout_operations_m

"""Operations that can be performed on lineouts or collections of lineouts."""
import numpy as np
import xarray as xr
from scipy.interpolate import interp1d
from torx.autodoc_decorators_m import autodoc_function



@autodoc_function
def combine_quants_on_radial_lineouts(
    quants_on_lineouts: dict,
    theta: np.array,
    npoints: int=100
):
    """
    Combine quantities from rho values on different lineouts to common rho.

    Applies a resampling to the grid spanning from the min to the max value of
    rho present in all datasets and combines the result into a single dataset
    along the theta dimension provided.
    """
    # Check if all quantities have dimension rho
    if any([not "rho" in r.dims for r in quants_on_lineouts.values()]):
        raise RuntimeError(
                    f"The input quantities must contain the dimension rho."
                )

    # TODO: Add support for other datatypes
    if(isinstance(quants_on_lineouts, dict)):
        rho_min = np.min([np.min(ql.rho) for ql in quants_on_lineouts.values()])
        rho_max = np.max([np.max(ql.rho) for ql in quants_on_lineouts.values()])
    else:
        raise RuntimeError(
                    f"The input quantities must be given as a dictionary."
                )
    # Apply interpolation to the same rho for all quantities
    rho = np.linspace(rho_min, rho_max, npoints)
    ql_new = {}
    for i, ql in enumerate(quants_on_lineouts.values()):

        interp_fun = lambda y, x: interp1d(x, y, fill_value="extrapolate")(rho)

        ql_new[i] = xr.apply_ufunc(
            interp_fun,
            ql,
            input_core_dims=[["rho"]],
            output_core_dims=[["rho"]],
            exclude_dims={"rho"},
            vectorize=True,
            dask="parallelized",
            dask_gufunc_kwargs={"output_sizes":{"rho":npoints}},
            output_dtypes=[np.float64],
            kwargs={"x":ql.rho}
        )
        ql_new[i].coords["rho"] = rho
        ql_new[i] = ql_new[i].assign_coords({"theta":theta[i]})

    return xr.concat(ql_new.values(), dim="theta")