SCRIPTS FITTING spl_mod

# -*-Python-*-
# Created by bgrierson at 13 Dec 2016  07:57

# Demonstration of the GAprofiles spline fitter "spl_mod"


def spl_mod_idl(x, y, ye, knotloc, autoknot, sol, bcs):

    # write the data to text file
    scratch['spl_mod_input'] = OMFITascii('spl_mod_input.dat')
    with open(scratch['spl_mod_input'].filename, 'w') as index:
        for i in range(len(x)):
            index.write('{} {} {}\n'.format(x[i], y[i], ye[i]))

    nk = len(knotloc)

    if autoknot:
        ak = 'yes'
    else:
        ak = 'no'

    if sol:
        edge = 'fit'
    else:
        edge = 'ignore'

    if bcs['y0'] is not None:
        y0 = bcs['y0']
    else:
        y0 = 'free'

    if bcs['y1'] is not None:
        y1 = bcs['y1']
    else:
        y1 = 'free'

    if bcs['y1p'] is not None:
        y1p = bcs['y1p']
    else:
        y1p = 'free'

    model = '''
model = create_struct('autoknot','{}')
model = create_struct(model,'numknot',{})
model = create_struct(model,'knotloc',knotloc)
model = create_struct(model,'type','point')
model = create_struct(model,'axis','{}')
model = create_struct(model,'boundary','{}')
model = create_struct(model,'edgeslope','{}')
model = create_struct(model,'edge','{}')
model = create_struct(model,'x_axis','rho')
model = create_struct(model,'normalize','no')
model = create_struct(model,'ftol',0.01)
model = create_struct(model,'maxit',200)
'''.format(
        ak, nk, y0, y1, y1p, edge
    ).strip()

    # Define IDL text
    runName = 'idl_input.pro'
    idl_cmd = []
    idl_cmd.append("!PATH = !PATH + ':' + getenv('VERSION4D')+'FITTING/'")
    idl_cmd.append(".compile spl_mod")
    idl_cmd.append('data = FLTARR(3,{})'.format(len(x)))
    idl_cmd.append("OPENR,lun,'spl_mod_input.dat',/GET_LUN")
    idl_cmd.append('READF,lun,data')
    idl_cmd.append('x = data[0,*]')
    idl_cmd.append('y = data[1,*]')
    idl_cmd.append('ye = data[2,*]')
    idl_cmd.append('knotloc={}'.format(knotloc.tolist()))
    idl_cmd.append("model_interp = 'spline'")
    idl_cmd.append(model)
    idl_cmd.append('f = spl_mod(x,y,ye,MODEL=model)')
    idl_cmd.append("SAVE,f,FILENAME='spl_mod_f.sav'")
    idl_cmd.append('exit')

    # Define executable variables for remote execution
    runScript = OMFITascii(runName, fromString='\n'.join(idl_cmd))
    server = IDE['idl']['server']
    tunnel = IDE['idl']['tunnel']
    remotedir = OMFITworkDir(root, IDE['idl']['server']) + '/spl_mod/'
    workdir = MainSettings['SETUP']['workDir'] + '/spl_mod/'

    # Execute the fit
    OMFITx.executable(
        root,
        inputs=[runScript, scratch['spl_mod_input']],
        outputs=['spl_mod_f.sav'],
        executable=IDE['idl']['idl'] + ' ' + runName + ';ls *',
        workdir=workdir,
        server=server,
        tunnel=tunnel,
        remotedir=remotedir,
    )

    # Save result to the OMFIT tree
    if 'spl_mod_f' in scratch:
        del scratch['spl_mod_f']
    scratch['spl_mod_f'] = OMFITidlSav('spl_mod_f.sav')


nxdata = 51
xdata = linspace(0, 1, nxdata)
ydata = np.sin(2.0 * np.pi * xdata) + 2.0
yerr = np.repeat(0.1, nxdata)
wgt = 1.0 / yerr
# Some fixed knots
knotloc = np.array([0.0, 0.33, 0.66, 1.0])
# Boundary conditions
bcs = {'y0': None, 'y0p': 0.0, 'y1': None, 'y1p': None}
# Fit SOL?
sol = False
# Autoknotter?
autoknot = False

spl_mod_idl(xdata, ydata, yerr, knotloc, autoknot, sol, bcs)

# Compute the residual (data-fit)/err
yfit = interpolate.interp1d(scratch['spl_mod_f']['f']['XFIT'], scratch['spl_mod_f']['f']['YFIT'])(xdata)
resid = (ydata - yfit) * wgt

fig, ax = plt.subplots(nrows=2)
uerrorbar(xdata, uarray(ydata, yerr), ax=ax[0], markersize=3.0, label='Data')
ax[0].plot(scratch['spl_mod_f']['f']['XFIT'], scratch['spl_mod_f']['f']['YFIT'], label='IDL')
ax[0].plot(
    scratch['spl_mod_f']['f']['KNOTLOC'], np.repeat(np.mean(ydata), len(knotloc)), marker='o', ls='None', color='r', label='knot locations'
)
ax[0].legend(loc='best')
ax[0].text(0.05, 3.2, '$\\chi_r^2={0:0.2f}$'.format(scratch['spl_mod_f']['f']['REDCHISQ']))
ax[1].plot(xdata, resid, marker='o', ls='None')
ax[1].axhline(0.0, ls='dashed', color='black')
ax[1].set_ylim([-np.max(np.abs(ax[1].get_ylim())), np.max(np.abs(ax[1].get_ylim()))])