Source code for torx.analysis.fit_eich_profile_m

"""Routines to fit Eich-type heat flux profiles."""
import warnings
import numpy as np
import matplotlib.pyplot as plt
import scipy.special
from scipy.optimize import curve_fit
from torx import Quantity
from torx.autodoc_decorators_m import autodoc_function

def erfc(x):
    """
    Complimentary error function.

    Handles quantity objects.

    Only dimensionless quantities can be handled.

    erfc(x) = 1 - erf(x)
    """
    if isinstance(x, Quantity):
        return scipy.special.erfc(x.to("").magnitude)
    else:
        return scipy.special.erfc(x)



@autodoc_function
def eich_profile(s_bar, q0, lambda_q, S, q_BG):
    """
    Fit an Eich profile to heat flux data.

    Using upstream-mapped distance to remove the explicit flux-expansion factor.

    See Eich Nuclear Fusion 2013. Using s_bar as the OMP-mapped distance
    means that you should set 'f_x' = 1.

    You can pass the parameters to the profile either all as dimensional
    Quantity objects, or all as dimensionless raw arrays/floats.

    s_bar: OMP-mapped radial distance [mm]
    q0: peak heat flux [MW/m^2]
    lambda_q: heat flux decay width [mm]
    S: spreading factor [mm]
    q_BG: background heat-flux [MW/m^2]
    """
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)
        return (
            q0
            / 2.0
            * np.exp((S / (2.0 * lambda_q)) ** 2 - s_bar / lambda_q)
            * erfc(S / (2.0 * lambda_q) - s_bar / S)
            + q_BG
        )




@autodoc_function
def fit_eich_profile(position, heat_flux, p0, plot=False):
    """
    Use non-linear least-squares to fit an Eich profile to heat flux data.

    Should pass
    position: OMP-mapped radial distance [mm]
    heat_flux: total heat flux [MW/m^2]
    p0: initial guess off the parameters for the Eich profile

    We use the Levenberg-Marquardt algorithm (default) for curve fitting.

    Note that the algorithm takes into account the error if provided.
    """
    heat_flux_loc = heat_flux * heat_flux.attrs["norm"].to("MW/m^2").magnitude
    position_loc = position * position.attrs["norm"].to("mm").magnitude

    # Transform p0 to the correct units
    p0_loc = [
        p0[0].to("MW/m^2").magnitude,
        p0[1].to("mm").magnitude,
        p0[2].to("mm").magnitude,
        p0[3].to("MW/m^2").magnitude,
    ]

    popt, pcov = curve_fit(
        eich_profile,
        xdata=position_loc,
        ydata=heat_flux_loc,
        p0=p0_loc,
        method="lm",
        xtol=1e-6,
    )

    q0, lambda_q, S, q_BG = popt
    perr = np.sqrt(np.diag(pcov))
    q0_err, lambda_q_err, S_err, q_BG_err = perr

    if plot:
        plt.plot(position_loc, heat_flux_loc, label="signal")
        plt.plot(position, eich_profile(position_loc, *popt), label="eich")
        plt.legend()
        plt.title("Eich profile fitting")

    return (
        Quantity((q0, q0_err), "MW/m^2"),
        Quantity((lambda_q, lambda_q_err), "mm"),
        Quantity((S, S_err), "mm"),
        Quantity((q_BG, q_BG_err), "MW/m^2"),
    )