SCRIPTS FUNCTIONS wrap_uncertaintiesΒΆ

# -*-Python-*-
# Created by smithsp at 01 Sep 2017  12:25

"""
This script demonstrates how to wrap a function that takes an array
as an argument and returns an array so that the function will take
arrays with uncertainties

defaultVars parameters
----------------------
:param i: The i'th sample profile (i<10)
"""
defaultVars(i=0)

# Get a sample data profile to fit
pkl = OMFITpickle(OMFITsrc + '/../samples/data_pickled.pkl')  # 10 55-pt profiles
order = argsort(pkl['x'][i, :])
x = pkl['x'][i, order]
y = pkl['y'][i, order]
yerr = pkl['e'][i, order]
f = fitSL(x, y, yerr)
fits = f.fits
params = f.fits['default'].params
# print f.get_best_fit()
xfit = linspace(0, 1.1, 110)
yfit = f(xfit)
yd = []
for yk in yfit:
    yd.extend(list(yk.derivatives.keys()))
print(len(set(yd)))
# print yfit[1].derivatives.keys(),yfit[0].derivatives.keys()
# OMFITx.End()
uerrorbar(x, uarray(y, yerr))
uband(xfit, yfit)
unc_tree = root['OUTPUTS'].setdefault('UNCERTAINTIES', dict())
close('all')
# def fit_spline(data,unc_tag='nominal',xfit=linspace(0,1,100)):
#    if unc_tag in unc_tree:
#        return unc_tree[unc_tag](xfit)
#    x = data[0,:]
#    y = data[1,:]
#    yerr = data[2,:]
#    uerrorbar(x,uarray(y,yerr))
#    unc_tree[unc_tag] = interpolate.LSQUnivariateSpline(x,y,[0,.2,.4,.9,max(x)],k=10,w=1/yerr)
#    return unc_tree[unc_tag](xfit)
# wrapped_fit_spline = wrap_array_func(fit_spline)
# uband(xfit,wrapped_fit_spline(array([x,uarray(y,yerr),yerr]),xfit=xfit))
q = xfit ** 2 + 1 + xfit


def transit_e(te, q=q, R0=1.6955, unc_tag='nominal'):
    nom_tag = 'transit_e;q=%r,R0=%r' % (q, R0)
    if unc_tag == 'nominal':
        scratch[nom_tag] = {}
    tmp = scratch[nom_tag]
    if unc_tag not in tmp:
        tmp[unc_tag] = utils_fusion.omega_transit_e(q=q, te=te, R0=R0)
    return tmp[unc_tag]


utrans = wrap_array_func(transit_e)(yfit, q=q)
uband(xfit, utrans)
figure()
uband(xfit, deriv(xfit, utrans))