GUIS shot_designer_prof

# -*-Python-*-
# Created by bgrierson at 20 Dec 2016  14:23

# Location of variables and namelist
loc = "root['INPUTS']['GUIS_shot_designer']"
nml = eval(loc)

# Profile varables
ext = ['NEL', 'TEL']

# Initialize
nml.setdefault('prof_times', nml['tinit'])
for e in ext:
    nml.setdefault('{}_sep'.format(e), np.zeros(1))
    nml.setdefault('{}_ped'.format(e), np.zeros(1))
    nml.setdefault('{}_axis'.format(e), np.zeros(1))

# Check to guarantee times are floats or npumpy arrays, and the initial time is in the time list.
def check_time(value):
    return nml['tinit'] in np.atleast_1d(value)


# Check that values are of the right length
def check_profs(value):
    return len(value) == len(np.atleast_1d(nml['prof_times']))


# Preview profiles in time
def preview():
    x = np.array([0.0, 0.95, 1.0])
    fn = FigureNotebook(0, 'Profiles')
    for e in ext:
        fig, ax = fn.subplots(label=e)
        for i, t in enumerate(np.atleast_1d(nml['prof_times'])):
            y = np.array([nml['{}_axis'.format(e)][i], nml['{}_ped'.format(e)][i], nml['{}_sep'.format(e)][i]])
            ax.plot(x, y, label='Time={}'.format(t))
        ax.legend()


# Time vertices
OMFITx.Entry(loc + "['prof_times']", 'Times (s)', default=nml['tinit'], check=check_time, updateGUI=True)

# Profile parameters
for e in ext:
    if check_profs(nml['{}_sep'.format(e)]) and check_profs(nml['{}_ped'.format(e)]) and check_profs(nml['{}_axis'.format(e)]):
        OMFITx.Tab(e)
    else:
        OMFITx.Tab('*' + e + '*')
    OMFITx.Entry(
        loc + "['{}_sep']".format(e), e + ' sep', default=np.zeros(len(np.atleast_1d(nml['prof_times']))), check=check_profs, updateGUI=True
    )
    OMFITx.Entry(
        loc + "['{}_ped']".format(e), e + ' ped', default=np.zeros(len(np.atleast_1d(nml['prof_times']))), check=check_profs, updateGUI=True
    )
    OMFITx.Entry(
        loc + "['{}_axis']".format(e),
        e + ' axis',
        default=np.zeros(len(np.atleast_1d(nml['prof_times']))),
        check=check_profs,
        updateGUI=True,
    )

# Preview
OMFITx.Tab()
OMFITx.Button('Preview', preview)