# -*-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))
