# -*-Python-*-
# Created by bgrierson at 27 Jul 2017 13:48
"""
This script demonstrates the use of an MDS+ elm detector to filter data
"""
ser = 'DIII-D'
sh = 155196
thresh = 1e16
madthresh = 3.5 * 3
tmin = 2e3
tmax = 5e3
phase_keep = [0.8, 0.99]
period_keep = [20.0, 200.0]
fs06f = OMFITmdsValue(server=ser, shot=sh, TDI='fs06f')
peped = OMFITmdsValue(server=ser, shot=sh, TDI='prmtan_peped')
elmstart = OMFITmdsValue(server=ser, shot=sh, TDI='elmstart')
elmend = OMFITmdsValue(server=ser, shot=sh, TDI='elmend')
elmsmarker = OMFITmdsValue(server=ser, shot=sh, TDI='elmsmarker')
elmpeak = OMFITmdsValue(server=ser, shot=sh, TDI='elmpeak')
# Find where ELM start is above threshold
wh = where((elmstart.data() > thresh) * (elmstart.dim_of(0) >= tmin) * (elmstart.dim_of(0) <= tmax))
# Number of ELMs detected
nelm = len(wh[0])
print('Detected {} ELMs'.format(nelm))
# Time for ELM start
elm_start = elmstart.dim_of(0)[wh[0]]
# Time for ELM end
elm_end = elmend.dim_of(0)[wh[0]]
# ELM phase needs four ponts per ELM (start, start+eps, end, end+eps)
# We ignore first and last ELM-free periods
elm_phase_time = np.zeros(nelm * 4)
elm_phase = np.zeros(nelm * 4)
elm_period = np.zeros(nelm * 4)
for count, i in enumerate(range(0, nelm * 4 - 4, 4)):
elm_phase_time[i] = elm_start[count]
elm_phase_time[i + 1] = elm_start[count] + 1e-6
elm_phase_time[i + 2] = elm_end[count]
elm_phase_time[i + 3] = elm_end[count] + 1e-6
elm_phase[i] = 1
elm_phase[i + 1] = -1
elm_phase[i + 2] = 0
elm_phase[i + 3] = 0
elm_period[i + 1] = elm_start[count] - elm_end[count]
elm_period[i + 2] = elm_start[count] - elm_end[count]
elm_period[i + 3] = elm_start[count + 1] - elm_end[count]
elm_period[i + 4] = elm_start[count + 1] - elm_end[count]
# Ignore last ELM
elm_phase_time = elm_phase_time[0:-4]
elm_phase = elm_phase[0:-4]
elm_period = elm_period[0:-4]
# Put ELM phase and period on pePed timebase
elm_phase_on_peped = interpolate.interp1d(elm_phase_time, elm_phase, bounds_error=False)(peped.dim_of(0))
elm_period_on_peped = interpolate.interp1d(elm_phase_time, elm_period, bounds_error=False)(peped.dim_of(0))
# Select the ELM phase and period we want
wh = where(
(elm_phase_on_peped >= phase_keep[0])
* (elm_phase_on_peped <= phase_keep[1])
* (elm_period_on_peped > period_keep[0])
* (elm_period_on_peped < period_keep[1])
)
peped_use_time = peped.dim_of(0)[wh[0]]
peped_use = peped.data()[wh[0]]
# Plot it
fn = FigureNotebook(0)
fig, ax = fn.subplots(nrows=3, sharex=True, label='Filterscope')
ax[0].plot(fs06f.dim_of(0), fs06f.data(), label='fs06f')
ax[0].plot(elmstart.dim_of(0), elmstart.data(), marker='o', ms=4, label='start')
ax[0].plot(elmend.dim_of(0), elmend.data(), marker='o', ms=4, label='end')
ax[0].plot(elmsmarker.dim_of(0), elmsmarker.data(), marker='o', ms=4, label='marker')
ax[0].plot(elmpeak.dim_of(0), elmpeak.data(), marker='o', ms=4, label='peak')
ax[0].axhline(thresh, color='black', ls='dashed')
ax[0].set_ylim([0, 5e16])
ax[0].legend()
ax[1].plot(elm_phase_time, elm_phase, marker='o', label='ELM Phase')
ax[1].axhline(0, color='black')
ax[1].legend()
ax[2].plot(elm_phase_time, elm_period, marker='o', label='ELM Period')
ax[2].axhline(0, color='black')
ax[2].legend()
fig, ax = fn.subplots(nrows=3, sharex=True, label='pePed')
ax[0].plot(peped.dim_of(0), peped.data(), color='black', label='pePed')
ax[0].scatter(peped.dim_of(0), peped.data(), c=elm_phase_on_peped, label='Phase')
ax[0].plot(peped_use_time, peped_use, marker='o', color='y', lw=4, alpha=0.4, label='Keep')
ax[0].plot(fs06f.dim_of(0), fs06f.data() * 5e-17)
ax[0].set_ylim([0, 6])
ax[0].legend()
ax[1].plot(elm_phase_time, elm_phase, color='black', label='ELM Phase')
ax[1].scatter(peped.dim_of(0), elm_phase_on_peped, c=elm_phase_on_peped)
ax[1].axhline(phase_keep[0], color='black', ls='dashed')
ax[1].axhline(phase_keep[1], color='black', ls='dashed')
ax[1].axhline(0, color='black')
ax[1].legend()
ax[2].plot(elm_phase_time, elm_period, marker='o', label='ELM Period')
ax[2].axhline(period_keep[0], color='black', ls='dashed')
ax[2].axhline(period_keep[1], color='black', ls='dashed')
ax[2].legend()
# Try to automatically detect ELMs without relying on a threhold using MAD median-absolute-deviation
fig, ax = plt.subplots(nrows=3, sharex=True)
ax[0].plot(fs06f.dim_of(0), fs06f.data(), label='fs06f')
ax[0].plot(elmstart.dim_of(0), elmstart.data(), marker='o', ms=4, label='start')
ax[0].plot(elmend.dim_of(0), elmend.data(), marker='o', ms=4, label='end')
ax[0].plot(elmsmarker.dim_of(0), elmsmarker.data(), marker='o', ms=4, label='marker')
ax[0].plot(elmpeak.dim_of(0), elmpeak.data(), marker='o', ms=4, label='peak')
ax[0].axhline(thresh, color='black', ls='dashed')
ax[0].set_ylim([0, 5e16])
ax[0].legend()
med = np.median(elmstart.data())
adiff = np.abs(elmstart.data() - med)
mad = np.median(adiff)
mod_z_score = 0.6745 * adiff / mad
ax[1].plot(fs06f.dim_of(0), fs06f.data(), label='fs06f')
ax[1].legend()
ax[2].plot(elmstart.dim_of(0), mod_z_score, marker='o', ms=4, label='Mod Z Score')
ax[2].axhline(3.5, color='black', ls='dashed')
ax[2].axhline(3.5 * 2, color='black', ls='dashed')
ax[2].set_yscale('log')
