SCRIPTS FITTING polyexp

# -*-Python-*-
# Created by bgrierson at 14 Dec 2016  08:00

# Demonstration of the polyexp fitting in OMFITprofiles

from pylab import random


def mtanh(x, A, B, xsym, hwid, core, edge):
    """
    Groebner APS 2004
    y = A * MTANH(alpha,z) + B
    MTANH(alpha,z) = [(1+alpha*z)*exp(z) - exp(-z)] / [exp(z) + exp(-z)]
    generalized to have a arbitrary order poly on numerator exp(z) and exp(-z)
    to allow for more variation in the core and edge.  Re-formulated as
    MTANH(alpha,z) = [nc*exp(z) - ne*exp(-z)] / [exp(z) + exp(-z)]
    where nc and nc are the core and edge polynomials
    Pedestal height is A+B
    Pedestal offset is A-B
    Location of the knee is xsym-hwid
    Location of the foot is xsym+hwid
    :param x: x-axis values
    :param A: Amplitude
    :param B: Offset
    :param xsym: Symmetry location
    :param hwid: half-width
    :param core: core polynomial coefficients [1.0, ...]
    :param edge: edge polynomial coefficients [1.0, ...]
    :return: Function y(x)
    """
    # Z shifted function
    z = (xsym - x) / hwid
    # Core polynomial
    nc = core[0]
    for i in range(1, len(core)):
        nc += core[i] * z ** (i)
    # Edge polynomial
    ne = edge[0]
    for i in range(1, len(edge)):
        ne += edge[i] * z ** (i)
    # mtanh in all its glory
    mt = (nc * np.exp(z) - ne * np.exp(-z)) / (np.exp(z) + np.exp(-z))
    # Final function
    y = A * mt + B
    return y


# ----
# Set default profile parameters
# ----
xsym = 1.0
hwid = 0.02
core = np.array([1.0, 1e-2])
edge = np.array([1.0, -1e-2, -3e-4])
pedestal = 5.0
offset = 0.1
B = (pedestal + offset) / 2.0
A = B - offset

# ----
#  Form "data" with noise and uncertainty
# ----
x = linspace(0.0, 1.2, 101)
ydata = mtanh(x, A, B, xsym, hwid, core, edge) + 0.5 * random(len(x))
yerr = np.repeat(0.1, len(x))

# ----
# Now fit the data
# ----
pow_core = 2
onAxis_gradzero = True
onAxis_value = None
fitEdge = True
maxiter = 2000
blend_width_min = 0.01
edge_width_min = 0.01
fitout = MtanhPolyExpFit(
    x,
    ydata,
    yerr,
    pow_core=pow_core,
    onAxis_gradzero=onAxis_gradzero,
    onAxis_value=onAxis_value,
    fitEdge=fitEdge,
    maxiter=maxiter,
    blend_width_min=blend_width_min,
    edge_width_min=edge_width_min,
)
y_unc = fitout(x)
print('Fit reduced chi2: {}'.format(fitout.redchi))

fig, ax = plt.subplots(nrows=2)
uerrorbar(x, uarray(ydata, yerr), ax=ax[0], markersize=2, ls='None', label='Data')
ax[0].plot(x, nominal_values(y_unc), label='Fit')
ax[0].legend()
ax[1].plot(x, (ydata - nominal_values(y_unc)) / yerr, marker='o', ls='None', label='Residual')
ax[1].axhline(0.0, ls='dashed')
ax[1].legend()