class GEO:
def __init__(self, geoframe=None):
if geoframe is None:
if __name__ == '__main__':
geoframe = tk.Tk()
else:
geoframe = tk.Toplevel()
geoframe.geometry('1600x950')
pol_frame = ttk.Frame(geoframe)
butframe = ttk.Frame(geoframe)
toolframe = ttk.Frame(geoframe)
for frame in (pol_frame, butframe, toolframe):
frame.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
# Canvas frame
self.pol_fig = Figure(figsize=(8.3, 8.3), dpi=100)
self.pol_fig.set_facecolor(figc)
self.pol_fig.subplots_adjust(left=0.1, bottom=0.1, right=0.9 , top=0.94)
self.pol_canvas = FigureCanvasTkAgg(self.pol_fig, master=pol_frame)
self.pol_canvas._tkcanvas.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
axsep = self.pol_fig.add_subplot(111, aspect='equal')
axsep.set_xlabel('R [m]')
axsep.set_ylabel('z [m]')
self.sepscat, = axsep.plot([], [], 'k+')
self.sepfit, = axsep.plot([], [], 'g-')
gc_r, gc_z = map_equ_20180130.get_gc()
for key in gc_r.keys():
axsep.plot(gc_r[key], gc_z[key], 'b-')
# Player buttons
locdir = os.path.dirname(os.path.realpath(__file__))
self.playfig = tk.PhotoImage(file='%s/play.gif' %locdir)
self.pausefig = tk.PhotoImage(file='%s/pause.gif' %locdir)
self.forwardfig = tk.PhotoImage(file='%s/forward.gif' %locdir)
self.backwardfig = tk.PhotoImage(file='%s/backward.gif' %locdir)
self.play_button = ttk.Button(butframe, command=self.Play, image=self.playfig)
backward_button = ttk.Button(butframe, command=self.Backward,image=self.backwardfig)
forward_button = ttk.Button(butframe, command=self.Forward, image=self.forwardfig)
createToolTip(forward_button , 'Go to next timestep')
createToolTip(backward_button , 'Go backward by one timestep')
createToolTip(self.play_button, 'Forward animation/pause')
for but in backward_button, self.play_button, forward_button:
but.pack(side=tk.LEFT)
# Entry frame
geoinit = {'Nmom':'6', 'exp':'AUGD', 'dia':'EQH', 'ed':'0'}
self.geodict = {}
n_col=0
for key in ('ed', 'dia', 'exp', 'Nmom'):
var = ttk.Entry(butframe, width=6)
var.insert(0, geoinit[key])
var.pack(side=tk.RIGHT, padx='0 10')
ttk.Label(butframe, text='%s:' %key, width=6).pack(side=tk.RIGHT, padx=5)
self.geodict[key] = var
# Toolbar
toolbar = NavigationToolbar2TkAgg(self.pol_canvas, toolframe)
toolbar.config(background=tbc)
toolbar.update()
self.pol_canvas.draw()
if __name__ == '__main__':
geoframe.mainloop()
def sepx(self, nshot, tbeg=0., tend=10., tok='AUGD'):
self.nshot = nshot
self.tbeg = tbeg
self.tend = tend
self.diag = self.geodict['dia'].get().strip()
self.exp = self.geodict['exp'].get().strip()
ed = int(self.geodict['ed'].get())
mom_order = int(self.geodict['Nmom'].get())
# Execute fit
self.pol_fig.clf()
sliderax = self.pol_fig.add_axes([0.1, 0.01, 0.6, 0.03], axisbg='yellow')
self.sl_time = Slider(sliderax,'Time:', 0., 1., valinit=0.)
self.sl_time.valtext.set_text('0.0000')
def update_time(val):
self.update_plot(val=val)
self.sl_time.on_changed(update_time)
self.pol_fig.canvas.mpl_connect('button_press_event', lambda event:self.pol_fig.canvas._tkcanvas.focus_set())
axsep = self.pol_fig.add_subplot(111, aspect='equal')
axsep.set_xlabel('R [m]')
axsep.set_ylabel('z [m]')
self.sepscat, = axsep.plot([], [], 'k+')
self.sepfit, = axsep.plot([], [], 'g-')
self.mag, = axsep.plot([], [], 'ro', label='Mag. axis')
if tok == 'JET':
self.septxt = axsep.text(3.6, -1.8, 't =%6.3f s' %tbeg, ha='center')
self.fit = fit_sep.FIT_SEP_JET(self.nshot, self.tbeg, self.tend, ed=ed, mom_order=mom_order)
axsep.set_xlim((1.5, 4.0))
axsep.set_ylim((-2, 2))
else:
self.septxt = axsep.text(2.6, -1.2, 't =%6.3f s' %tbeg, ha='center')
gc_r, gc_z = map_equ_20180130.get_gc()
for key in gc_r.keys():
axsep.plot(gc_r[key], gc_z[key], 'b-')
axsep.set_xlim((0.8, 3.0))
axsep.set_ylim((-1.4, 1.4))
self.u, = axsep.plot([], [], 'ko')
self.o, = axsep.plot([], [], 'go')
self.xu, = axsep.plot([], [], 'bo', label='Lower X')
self.xo, = axsep.plot([], [], 'mo', label='Upper X')
self.fit = fit_sep.FIT_SEP(self.nshot, self.tbeg, self.tend, diag=self.diag, exp=self.exp, ed=ed, mom_order=mom_order)
axsep.set_title('%s:%s(%d) #%d' \
%(self.exp, self.diag, self.fit.ed, self.nshot) )
axsep.legend(numpoints=1, loc=1)
if len(self.fit.ind_ok) <= 0:
print('Problems with the fit')
return
self.nt = len(self.fit.tarr)
self.t_u = self.fit.tarr[self.fit.ind_ok]
self.moments_u = 100.*self.fit.moments[self.fit.ind_ok, :, :]
self.store_mry()
self.astra_ubnd()
self.begin()
#--------------------------
# Animation separatrix plot
#--------------------------
def begin(self):
self.stop = True
self.jt = 0
self.update_plot()
def Pause(self):
self.stop = True
self.play_button['image'] = self.playfig
self.play_button['command'] = self.Play
def Forward(self):
self.stop = True
self.play_button['image'] = self.playfig
self.play_button['command'] = self.Play
self.jt += 1
self.jt = self.jt%self.nt
self.update_plot()
def Backward(self):
if not hasattr(self, 'jt'):
print('Run\n File-> Read and fit separatrix\nfirst')
return
self.stop = True
self.play_button['image'] = self.playfig
self.play_button['command'] = self.Play
self.jt -= 1
self.jt = self.jt%self.nt
self.update_plot()
def Play(self):
self.stop = False
self.play_button['image'] = self.pausefig
self.play_button['command'] = self.Pause
timeout = 1e-10
while (self.jt < self.nt-1) and (not self.stop):
self.jt += 1
self.update_plot()
self.pol_canvas.start_event_loop(timeout)
def updateTimeSlider(self, val=None):
if val == None:
progress = (self.fit.tarr[self.jt] - self.fit.tarr[0])/(self.fit.tarr[-1] - self.fit.tarr[0])
else:
progress = val
self.sl_time.val = progress
poly = self.sl_time.poly.get_xy()
poly[2:4, 0] = progress
self.sl_time.poly.set_xy(poly)
self.sl_time.valtext.set_text('%8.4f' %self.fit.tarr[self.jt])
def update_plot(self, val=None):
if val != None:
tloc = self.fit.tarr[0] + val*(self.fit.tarr[-1] - self.fit.tarr[0])
self.jt = np.argmin(np.abs(self.fit.tarr - tloc))
if not hasattr(self, 'fit'):
return
self.sepscat.set_data(self.fit.rscat[self.jt], self.fit.zscat[self.jt])
if self.fit.ind_ok[self.jt]:
self.sepfit.set_visible(True)
self.sepfit.set_data(self.fit.r_bnd[self.jt], self.fit.z_bnd[self.jt])
else:
self.sepfit.set_visible(False)
self.mag.set_data(self.fit.Rmag[self.jt], self.fit.Zmag[self.jt])
if hasattr(self.fit, 'Rzunt'):
self.u.set_data( self.fit.Rzunt[self.jt], self.fit.Zunt[self.jt])
self.o.set_data(self.fit.Rzoben[self.jt], self.fit.Zoben[self.jt])
self.xu.set_data( self.fit.Rxpu[self.jt], self.fit.Zxpu[self.jt])
self.xo.set_data( self.fit.Rxpo[self.jt], self.fit.Zxpo[self.jt])
self.septxt.set_text('t =%6.3f s' %self.fit.tarr[self.jt])
self.updateTimeSlider(val=val)
self.pol_canvas.draw()
# ASCII output
def astra_ubnd(self):
if not hasattr(self, 'fit'):
return None
# ASTRA BND u-file
nt, n_the = self.fit.r_bnd[self.fit.ind_ok, :].shape
n_the -= 1
posarr = 1 + np.arange(2*n_the)
tlbl = 'Time'.ljust(20) + 'Seconds'
poslbl = 'POSITION'
bndlbl = 'BOUNDARY [M]'
self.drzarr = np.append(self.fit.r_bnd[self.fit.ind_ok, :-1], self.fit.z_bnd[self.fit.ind_ok, :-1], axis=1)
uf_d = { 'pre': 'M', 'ext': 'BND', 'shot': self.nshot, \
'grid': { 'X':{'lbl': tlbl , 'arr': self.t_u}, \
'Y':{'lbl': poslbl, 'arr': posarr} }, \
'data': {'lbl':bndlbl, 'arr': self.drzarr} }
ufiles.WU(uf_d)
def astra_expbnd(self):
nt, n_the2 = self.drzarr.shape
n_the = int(n_the2/2)
sbnd = 'NAMEXP BND NTIMES %3d POINTS %3d\n ' %(nt, n_the)
for jt in range(nt):
sbnd += '%13.5E' %self.t_u[jt]
if (jt%6 == 5) or (jt == nt-1):
sbnd += '\n '
for jt in range(nt):
for jthe in range(n_the):
sbnd += '%13.5E' %self.drzarr[jt, jthe]
if (jthe%6 == 5) or (jthe == n_the-1):
sbnd += '\n '
for jthe in range(n_the, n_the2):
sbnd += '%13.5E' %self.drzarr[jt, jthe]
if ((jthe-n_the)%6 == 5) or (jthe == n_the2-1):
sbnd += '\n '
return sbnd[:-1] # Remove trailing space
def store_mry(self):
nt, mom_order, mom_type = self.moments_u.shape
marr = 1 + np.arange(mom_order)
karr = 1 + np.arange(mom_type)
tlbl = 'Time'.ljust(20) + 'Seconds'
mlbl = 'MOMENT INDEX'
klbl = 'MOMENT TYPE'
dlbl = 'FOURIER MOMENTS'.ljust(20) + 'CM'
uf_d = { 'pre': 'M', 'ext': 'MRY', 'shot': self.nshot, \
'grid': { 'X':{'lbl': tlbl, 'arr': self.t_u}, \
'Y':{'lbl': mlbl, 'arr': marr}, \
'Z':{'lbl': klbl, 'arr': karr} }, \
'data': {'lbl':dlbl, 'arr':self.moments_u} }
uf_d['comm'] = '!exp=%s\n !diag=%s\n !ed=%d' %(self.exp, self.diag, self.fit.ed)
ufiles.WU(uf_d)
class ContTweaker:
def __init__(self,iwvlngth,ispectrum,espectrum,oldxspline,oldyspline,outfile):
#Initialize input variables for the class
self.outfile=outfile
self.iwvlngth=iwvlngth
self.ispectrum=ispectrum
self.espectrum=espectrum
#intialize CT.XSPLINE/YSPLINE to the old values at first (will be replaced)
self.xspline=oldxspline
self.yspline=oldyspline
self.oldxspline=oldxspline
self.oldyspline=oldyspline
#Interpolate cubic spline of the continuum for both "NEW" and old spline
self.oldspline=scipy.interpolate.interp1d(oldxspline,oldyspline,kind='cubic')
self.spline=scipy.interpolate.interp1d(oldxspline,oldyspline,kind='cubic')
#Make the MPL figures appear
self.MakePlot()
#Add buttons to the bottom of the main TB_CONTFIT widget menu
self.popup=Tkinter.Frame()
self.popup.grid()
But1=Tkinter.Button(self.popup,text='Refresh Plot',command=self.Refresh)
But1.grid(column=0,row=0)
But3=Tkinter.Button(self.popup,text='Remove Point',command=self.RemovePoint)
But3.grid(column=0,row=1)
But4=Tkinter.Button(self.popup,text='Add Point',command=self.AddPoint)
But4.grid(column=0,row=2)
But5=Tkinter.Button(self.popup,text='Reopen Window',command=self.MakePlot)
But5.grid(column=0,row=4)
But6=Tkinter.Button(self.popup,text='Save Spline',command=self.Save)
But6.grid(column=0,row=5)
But7=Tkinter.Button(self.popup,text='Exit',command=self.Exit)
But7.grid(column=0,row=6)
#Toggle between a Continous-editing and single click modes
MODES=[("Single",False),("Continuous",True)]
self.useContinuous=Tkinter.BooleanVar()
self.useContinuous.set(False)
ii=1
labelMode=Tkinter.StringVar()
labelmode=Tkinter.Label(self.popup,textvariable=labelMode,anchor="w",fg="black")
labelmode.grid(column=1, row=0, columnspan=2, sticky='S')
labelMode.set(u"Point-editing mode")
for text,mode in MODES:
b=Tkinter.Radiobutton(self.popup,text=text,variable=self.useContinuous,value=mode)
b.grid(column=ii,row=1)
ii+=1
if ii==1: b.select()
#Wait until the window is closed (i.e. CT.EXIT() is run by clicking EXIT button)
self.popup.wait_window()
def MakePlot(self):
#This function generates the plot window that will be the vidual interface for how
#the continuum tweaking is going
#Make the Spline from the CURRENT spline values
xcont=self.iwvlngth
ycont=self.spline(xcont)
#Make the spline from the OLD spline values
ospectrum=self.oldspline(xcont)
#Create a MPL figure
self.tweakfig=plt.figure(figsize=(8,8))
#First subplot is for displaying the original fit to the spectrum
self.ax1=self.tweakfig.add_subplot(3,1,1)
#Second subplot is for displaying the current fit to the spectrum
#To help with looking at all regions of the same time, share the X and Y
#axes to CT.AX1.
self.ax2=self.tweakfig.add_subplot(3,1,2,sharex=self.ax1,sharey=self.ax1)
#Third subplot for displaying the continuum-fitted spectrum based on the
#current spline. Because this is now 0 and 1ish, only share the X axis
#to CT.AX1
self.ax3=self.tweakfig.add_subplot(3,1,3,sharex=self.ax1)
#Create the TK window for embedding the figure, set it up, and draw
self.TFroot=Tkinter.Tk()
self.TFroot.wm_title("Continuum Tweaker")
self.TweakPlot=FigureCanvasTkAgg(self.tweakfig,master=self.TFroot)
#This toolbar is mainly for zooming. By sharing the x (and y) axes, when zooming in
#on one axes, all three will zoom appropriately
NavTweakPlot=NavigationToolbar2TkAgg(self.TweakPlot, self.TFroot)
self.TweakPlot.get_tk_widget().pack(side=Tkinter.TOP, fill=Tkinter.BOTH, expand=1)
#For CT.AX1, plot the original specturm, the error spectrum, and the original continuum fitted
self.ax1.plot(self.iwvlngth,self.ispectrum,'k',drawstyle='steps',label='Spectrum')
self.ax1.plot(self.iwvlngth,self.espectrum,'g',drawstyle='steps',label='Error')
self.ax1.plot(xcont,ycont,'b',label='Continuum')
self.ax1.plot(self.oldxspline,self.oldyspline,'or',label='Spline Pts.')
#Make a legend
self.ax1.legend(ncol=4,frameon=False, loc=9, bbox_to_anchor=(0.5, 1.3))
self.ymin,self.ymax=self.ax1.get_ylim()
self.ax1.set_ylabel('Flux\n(original)')
#For CT.AX2, plot the spectrum, and the current spline points and the resulting continuum
self.ax2.plot(self.iwvlngth,self.ispectrum,'k',drawstyle='steps')
self.ax2.plot(xcont,ycont,'b')
self.ax2.plot(self.xspline,self.yspline,'or',picker=5)#Picker needed to pick which point
self.ax2.set_ylabel('Flux\n(tweaked fit)')
#For CT.AX3, plot the spectrum upon dividing by the continuum. Plot the
#error spectrum as 1-err and 1+err to reflect how much noise one MIGHT expect
#As a rule of thumb, one should aim to have the continuum fluctuations within the error
self.ax3.plot(xcont,self.ispectrum/ospectrum,'k',drawstyle='steps')
self.ax3.plot(xcont,1.0-self.espectrum/ospectrum,'--g',drawstyle='steps')
self.ax3.plot(xcont,1.0+self.espectrum/ospectrum,'--g',drawstyle='steps')
self.xmin,self.xmax=self.ax3.get_xlim()
self.ax3.plot([self.xmin,self.xmax],[1,1],'--r')
self.ax3.plot([self.xmin,self.xmax],[0,0],'--r')
#restrict CT.AX3 to only show the normalized range of values
self.ax3.set_ylim(-1,2)
self.ax3.set_ylabel('Relative flux')
self.ax3.set_xlabel('Wavlength')
#Draw all the new plotted stuff
self.TweakPlot.draw()
def Refresh(self,yrefresh=True):
#This function takes any modifications that have taken place and updates
#the plots accordingly. If YREFRESH=TRUE, this sets the y-axis to the
#original scaling. Otherwise keep the current values.
#Get the current x values (incase the user has zoomed in using toolbar)
xmin,xmax=self.ax2.get_xlim()
ymin,ymax=self.ax2.get_ylim()
#if YREFRESH=TRUE, set to original values (SELF.YMIN/YMAX)
if yrefresh:
ymin=self.ymin
ymax=self.ymax
#Clear CT.AX2/AX3 of all previous information
self.ax2.clear()
self.ax3.clear()
#Remake the spline based on the new XSPLINE/YSPLINE points, and save
self.spline=scipy.interpolate.interp1d(self.xspline,self.yspline,kind='cubic')
#Generate the continuum based on the new spline
xcont=self.iwvlngth
ycont=self.spline(xcont)
#Update CT.AX1, perserve the x-axis bounds, and return to original y-axis
self.ax1.set_ylim(ymin,ymax)
self.ax1.set_xlim(xmin,xmax)
#Update CT.AX2 by plotting new spline, also perserve the x-axis bounds
self.ax2.plot(self.iwvlngth,self.ispectrum,'k',drawstyle='steps')
self.ax2.plot(xcont,ycont,'b')
self.ax2.plot(self.xspline,self.yspline,'or',picker=5)
self.ax2.set_ylabel('Flux\n(tweaked fit)')
self.ax2.set_ylim(ymin,ymax)
self.ax2.set_xlim(xmin,xmax)
#in CT.AX2, Divide out the spectrum&errorspectrum by the continuum, and plot
self.ax3.plot(xcont,self.ispectrum/ycont,'k',drawstyle='steps')
self.ax3.plot(xcont,1.0-self.espectrum/ycont,'--g',drawstyle='steps')
self.ax3.plot(xcont,1.0+self.espectrum/ycont,'--g',drawstyle='steps')
self.ax3.set_ylabel('Relative flux')
self.ax3.set_xlabel('Wavlength')
self.ax3.set_ylim(-1,2)
self.ax3.plot([self.xmin,self.xmax],[1,1],'--r')
self.ax3.plot([self.xmin,self.xmax],[0,0],'--r')
self.ax3.set_xlim(xmin,xmax)
#Update plotting window
self.TweakPlot.draw()
#Function for closing the current MPL event by it's ID, and stop the event loop
#It might seem redundant that I have both things, but I intend to have a continous
#editing method, which would need the looper.
def QuitEdit(self,cid):
#Close event ID
self.TweakPlot.mpl_disconnect(cid)
#Stop event loop
self.TweakPlot.stop_event_loop()
#Function when "Add Point" button is clicked.
def AddPoint(self):
#Show Tutorial message for what to do.
if usetutorial: tkMessageBox.showinfo("Help Message", "Click where to add point.")
#Start mouse click event, and run CT.ClickAdd
self.cidbut=self.TweakPlot.mpl_connect('button_press_event',self.ClickAdd)
self.TweakPlot.start_event_loop(0)
#If use continuous button is on, repeat adding points
while self.useContinuous.get():
if usetutorial: tkMessageBox.showinfo("Help Message", "Continuous point addition on. Keep adding points.")
try:
self.cidbut=self.TweakPlot.mpl_connect('button_press_event',self.ClickAdd)
self.TweakPlot.start_event_loop(0)
except: self.useContinuous.set(False)
#Given a mouse event for adding a point...
def ClickAdd(self,event):
#Grab the x/y coordiantes of the click, and add to spline
self.xspline.append(event.xdata)
self.yspline.append(event.ydata)
#Sort the spline data to be in order by wavelength
self.xspline,self.yspline=SortList(self.xspline,self.yspline)
#Refresh the plot with new data, but keep y-axis
self.Refresh(yrefresh=False)
#Close the MPL event stuff
self.QuitEdit(self.cidbut)
#Function ro remove a point when "Remove Point" button pressed
def RemovePoint(self):
#Show tutorial message on what to do
if usetutorial: tkMessageBox.showinfo("Help Message", "Click point to remove.")
#Start MPL event for picking an MPL artist, and start the loop. Run CT.ClickRemove
self.cidpick=self.TweakPlot.mpl_connect('pick_event',self.ClickRemove)
self.TweakPlot.start_event_loop(0)
#If Use continuous button is on, repeat removing points
while self.useContinuous.get():
if usetutorial: tkMessageBox.showinfo("Help Message", "Continuous point removal on. Keep removing points.")
try:
self.cidpick=self.TweakPlot.mpl_connect('pick_event',self.ClickRemove)
self.TweakPlot.start_event_loop(0)
except:
self.useContinuous.set(False)
#Given a picker event for removing a point...
def ClickRemove(self,event):
#Get the spline point that you picked, it's x and y coordinates
splinepoint = event.artist
xsplineval=splinepoint.get_xdata()
ysplineval=splinepoint.get_ydata()
#Index of the artist
ind = event.ind
#Make sure the point is in the spline point lists
if xsplineval[ind] in self.xspline:
if ysplineval[ind] in self.yspline:
#Remove that point from the spline, I think this is where sorting is important...
self.xspline.pop(ind)
self.yspline.pop(ind)
#Refresh the plot with new spline, but keep y-axis
self.Refresh(yrefresh=False)
#Close the event and stop the event loop
self.QuitEdit(self.cidpick)
#Function saves the spline using SaveSpline function
def Save(self):
print "Saving Spline"
SaveSpline(self.xspline,self.yspline,self.iwvlngth,self.outfile)
#Destroy CT.POPUP and CT.TFROOT (the masters for the CT buttons and plots) and leave CT.
def Exit(self):
self.popup.destroy()
self.TFroot.destroy()
return
class MPLWID:
timeout = 1e-10
def __init__(self, opt_d, rdd, myframe=None, method='rec_spl', n_spl=51):
self.method = method
self.ext_str = ''
for diag_sig in sorted(opt_d.keys()):
diag = diag_sig.split(':')[0]
self.ext_str += diag
# Choosing last diag's info for plotting
self.exp = opt_d[diag_sig]['exp' ]
self.diag = opt_d[diag_sig]['diag']
self.sig = opt_d[diag_sig]['sig']
self.eq_exp = opt_d[diag_sig]['Map_exp' ]
self.eq_diag = opt_d[diag_sig]['Map_diag']
self.nshot = int(opt_d[diag_sig]['shot'])
self.ed = int( opt_d[diag_sig]['ed'] )
self.eq_ed = int( opt_d[diag_sig]['Map_ed'] )
self.tbeg = float(opt_d[diag_sig]['tbeg' ] )
self.tend = float(opt_d[diag_sig]['tend' ] )
self.ylab = '%s %s' %(self.diag, self.sig)
self.xlab = r'$\rho_{pol}$'
# Combine diagnostics
for diag_sig in rdd.keys():
if not hasattr(self, 'xexp'):
self.xexp = np.array(rdd[diag_sig]['rhop'])
self.yexp = np.array(rdd[diag_sig]['data'])
self.yerr = np.array(rdd[diag_sig]['data_err'])
else:
self.xexp = np.append(self.xexp, rdd[diag_sig]['rhop'] , axis=1)
self.yexp = np.append(self.yexp, rdd[diag_sig]['data'] , axis=1)
self.yerr = np.append(self.yerr, rdd[diag_sig]['data_err'], axis=1)
ind_sort = np.argsort(self.xexp, axis=1)
for jt in range(self.xexp.shape[0]):
self.xexp[jt,:] = self.xexp[jt, ind_sort[jt]]
self.yexp[jt,:] = self.yexp[jt, ind_sort[jt]]
self.yerr[jt,:] = self.yerr[jt, ind_sort[jt]]
self.tgrid = rdd[diag_sig]['tgrid']
self.jt = 0
self.nt = len(self.tgrid)
self.nR = n_spl
self.rho_out = np.linspace(0, 1, n_spl)
self.R_out = np.linspace(1.05, 2.25, self.nR)
self.rho12_out = np.linspace(0, 1.2, n_spl)
# Frame structure
if myframe is None:
myframe = tk.Toplevel()
myframe.geometry('1200x740')
myframe.title('View and fit profiles, method %s' %self.method)
figframe = ttk.Frame(myframe)
entframe = ttk.Frame(myframe)
rbframe = ttk.Frame(myframe)
btframe = ttk.Frame(myframe)
lblframe = ttk.Frame(myframe)
toolframe = ttk.Frame(myframe)
for box in (figframe, toolframe, entframe, \
rbframe, btframe, lblframe):
box.pack(side=tk.TOP, fill=tk.X)
figframe3 = ttk.Frame(figframe)
figframe2 = ttk.Frame(figframe)
for box in (figframe3, figframe2):
box.pack(side=tk.LEFT, fill=tk.X)
figframe1 = ttk.Frame(figframe3)
entfframe = ttk.Frame(figframe3)
for box in (figframe1, entfframe):
box.pack(side=tk.TOP, fill=tk.X)
self.nb2d = ttk.Notebook(figframe2)
self.nb2d.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
frame2d = ttk.Frame(self.nb2d)
ctframe = ttk.Frame(self.nb2d)
self.nb2d.add(frame2d, text='Wired plot')
self.nb2d.add(ctframe, text='Contour plot')
# Plot settings
fig1d = Figure(figsize=(5.6, 4.9), dpi=100)
self.fig2d = Figure(figsize=(6.2, 5.1), dpi=100)
self.figct = Figure(figsize=(6.2, 5.1), dpi=100)
self.can_1d = FigureCanvasTkAgg(fig1d, master=figframe1)
self.can_2d = FigureCanvasTkAgg(self.fig2d, master=frame2d)
self.can_ct = FigureCanvasTkAgg(self.figct, master=ctframe)
self.can_1d._tkcanvas.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
self.can_2d._tkcanvas.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
self.can_ct._tkcanvas.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
toolbar = NavigationToolbar2TkAgg( self.can_1d, toolframe)
toolbar.update()
self.ax1d = fig1d.add_subplot(1, 1, 1)
fig1d.subplots_adjust(left=0.15, bottom=0.16, right=0.95, top=0.9)
self.ylab = '%s %s' %(self.diag, self.sig)
self.rho_lab = r'$\rho_{pol}$'
self.ax1d.set_xlim([0, 1.1])
self.ax1d.set_xlabel(self.rho_lab, fontsize=fsize)
self.ax1d.set_ylabel(self.ylab , fontsize=fsize)
self.ok, self.caplines, self.barlinecols = \
self.ax1d.errorbar(0, 0, 0, fmt='g.')
self.rej, = self.ax1d.plot([], [], 'ro')
self.del_pt, = self.ax1d.plot([], [], 'bo')
self.pfit, = self.ax1d.plot([], [], 'm-')
self.sep_line, = self.ax1d.plot([1, 1], [0, 1], 'k--')
# Player buttons
locdir = os.path.dirname(os.path.realpath(__file__))
self.playfig = tk.PhotoImage(file='%s/play.gif' %locdir)
self.pausefig = tk.PhotoImage(file='%s/pause.gif' %locdir)
self.forwardfig = tk.PhotoImage(file='%s/forward.gif' %locdir)
self.backwardfig = tk.PhotoImage(file='%s/backward.gif' %locdir)
self.backward_button = ttk.Button(entfframe, command=self.Backward, image=self.backwardfig)
self.play_button = ttk.Button(entfframe, command=self.Play, image=self.playfig)
self.forward_button = ttk.Button(entfframe, command=self.Forward, image=self.forwardfig)
createToolTip(self.forward_button , 'Go to next timestep')
createToolTip(self.backward_button , 'Go backward by one timestep')
createToolTip(self.play_button , 'Forward animation/pause')
sliderax = fig1d.add_axes([0.2, 0.02, 0.6, 0.03], axisbg='yellow')
self.sl_time = Slider(sliderax,'Time:', 0., 1., valinit=0.)
self.sl_time.valtext.set_text('0.0000')
def update_time(val):
self.update_plot(val=val)
self.sl_time.on_changed(update_time)
fig1d.canvas.mpl_connect('button_press_event', self.MouseInteraction)
fig1d.canvas.mpl_connect('key_press_event' , self.on_key)
fig1d.canvas.mpl_connect('key_release_event' , self.off_key)
fig1d.canvas.mpl_connect('button_press_event', lambda event:fig1d.canvas._tkcanvas.focus_set())
for but in self.backward_button, self.play_button, self.forward_button:
but.pack(side=tk.LEFT)
# Fit entries
sigma_init = 1
if self.diag in ['IDA']:
erbar_init = 0.
else:
erbar_init = .3
ttk.Label(entframe, text='Fit tolerance', width=11).pack(side=tk.LEFT, padx=2)
self.cfit_val = ttk.Entry(entframe, width=8)
self.cfit_val.insert(0, erbar_init)
self.cfit_val.pack(side=tk.LEFT)
if self.method == 'rec_spl':
ttk.Label(entframe, text='# sigmas').pack(side=tk.LEFT, padx='10 2')
self.sigma_val = ttk.Entry(entframe, width=8)
self.sigma_val.insert(0, sigma_init)
self.sigma_val.pack(side=tk.LEFT)
ttk.Label(rbframe, text='U-file coord', width=11).pack(side=tk.LEFT, padx=2)
self.rho_event = tk.StringVar()
self.rho_event.set('rho_tor')
for val in ('rho_pol', 'rho_pol12', 'rho_tor', 'R_hor'):
ttk.Radiobutton(rbframe, variable=self.rho_event, value=val, text=val).pack(side=tk.LEFT)
btfit = ttk.Button(btframe, text='Fit', command=self.fit)
btfit.pack(side=tk.LEFT, padx=2)
lblframe1 = ttk.Frame(lblframe)
lblframe2 = ttk.Frame(lblframe)
lblframe1.pack(side=tk.LEFT)
lblframe2.pack(side=tk.RIGHT)
scat_help = tk.Label(lblframe1, text='Left plot, scatter points:')
scat_blue = tk.Label(lblframe1, text='Blue: hand-deleted', fg='blue')
scat_red = tk.Label(lblframe1, text='Red: rejected by spline', fg='red')
scat_green = tk.Label(lblframe1, text='Green: valid points', fg='#00a000')
for w in (scat_help, scat_blue, scat_red, scat_green):
w.pack(side=tk.TOP, anchor=tk.W)
del_lb = ttk.Label(lblframe2, text='Mouse-L: deletes points')
del_all = ttk.Label(lblframe2, text='Ctrl+Mouse-L: deletes frame')
undolb = ttk.Label(lblframe2, text='Mouse-Mid: un-delete points for this frame')
righlb = ttk.Label(lblframe2, text='Mouse-R: move to next frame plot')
for w in (del_lb, del_all, undolb, righlb):
w.pack(side=tk.TOP, anchor=tk.W)
# Create GUI
self.update_plot1d()
self.stop = True
self.ctrl = False
def rhomap(self, t_ind):
t_ind = np.atleast_1d(t_ind)
if self.rho_lbl == 'rho_tor':
self.xlab = r'$\rho_{tor}$'
self.x_out = self.rho_out
# Map rho_pol-> rho_tor
if not hasattr(self, 'rho_t'):
status = eqm.Open(self.nshot, diag=self.eq_diag, exp=self.eq_exp, ed=self.eq_ed)
if not status:
print('Problems opening shotfile')
return
self.eq_ed = eqm.ed
self.rho_t = eqm.rho2rho(self.x_out, t_in=self.tgrid, coord_in='rho_pol', coord_out='rho_tor')
print('Edition for mapping %d' %self.eq_ed)
# Interpolate on fixed equispaced rho_tor grid
for jt in t_ind:
self.d_out[jt, :] = np.interp(self.x_out, self.rho_t[jt, :], self.y_fit[jt, :])
elif self.rho_lbl == 'R_hor':
self.xlab = r'$R_{hor}$'
self.x_out = self.R_out
status = eqm.Open(self.nshot, diag=self.eq_diag, exp=self.eq_exp, ed=self.eq_ed)
if not status:
print('Problems opening shotfile')
return
self.eq_ed = eqm.ed
# Map rho_pol-> R
print('Rhor')
R1, z1 = eqm.rhoTheta2rz(self.rho_out, [0, np.pi], t_in=self.tgrid[t_ind])
for jt in t_ind:
R_tmp = np.append(R1[jt, 0, :], R1[jt, 1, :])
d_tmp = np.append(self.y_fit[jt], self.y_fit[jt])
index = np.argsort(R_tmp)
self.d_out[jt, :] = np.interp(self.R_out, R_tmp[index], d_tmp[index])
print('Edition for mapping %d' %self.eq_ed)
else:
if self.rho_lbl == 'rho_pol12':
self.x_out = self.rho12_out
self.ax1d.set_xlim([0, 1.2])
else:
self.x_out = self.rho_out
self.ax1d.set_xlim([0, 1])
n_spl = len(self.rho_out)
self.xlab = r'$\rho_{pol}$'
self.d_out = np.array(self.y_fit).reshape((self.nt, n_spl))
def fit(self):
self.rho_lbl = self.rho_event.get().strip()
# Input ready, start execution
self.ind_del = []
self.tnot = []
zero = 1e-8
n_spl = len(self.rho_out)
self.d_out = np.zeros((self.nt, n_spl))
for jt in range(self.nt):
self.ind_del.append([])
print('Fitting experimental profiles')
self.fit_tol = float(self.cfit_val.get())
if self.method == 'rec_spl':
self.sigmas = float(self.sigma_val.get())
else:
self.sigmas = 1e12
if self.fit_tol == 0:
self.sigmas = 1e12
if self.rho_lbl == 'rho_pol12':
self.rho_spl = self.rho12_out
else:
self.rho_spl = self.rho_out
self.del2d = {}
log = logging.getLogger()
log.info('Starting 2d spline')
# Fit 2D
nt, nx_in = self.xexp.shape
nx_fit = len(self.rho_spl)
if self.method in ('rec_spl', 'gauss'):
pool = Pool(cpu_count())
out = pool.map(fit1d, [(self.xexp[jt], self.yexp[jt], self.yerr[jt], self.rho_spl, self.fit_tol, self.sigmas, self.method) for jt in range(nt)])
pool.close()
pool.join()
out2 = np.array(out)
self.y_fit = out2[:, :nx_fit]
self.dat_ok = out2[:, nx_fit: nx_fit+nx_in]
self.confid = out2[:, nx_fit+nx_in:]
self.rhomap(range(self.nt))
log.info('Done 2d spline')
# Contour plot
xg, self.tg = np.meshgrid(self.x_out, self.tgrid)
axct = self.figct.add_subplot(1, 1, 1)
axct.set_xlabel('Time [s]', fontsize=fsize3d)
axct.set_ylabel(self.xlab , fontsize=fsize3d)
ctr = axct.contourf(self.tg, xg, self.d_out)
self.fig2d.colorbar(ctr, aspect=10, shrink=0.9)
self.line, = axct.plot([], [], 'k-')
# 3D plot
self.ax2d = p3.Axes3D(self.fig2d)
self.ax2d.set_xlabel(self.xlab , fontsize=fsize3d)
self.ax2d.set_ylabel('Time [s]', fontsize=fsize3d)
self.ax2d.plot_wireframe(xg, self.tg, self.d_out, linewidth=.3, rstride=1, cstride=3)
self.sing2d, = self.ax2d.plot([], [], [], 'r-')
self.can_2d.draw()
self.update_plot()
#------------
# Interaction
#------------
def MouseInteraction(self,event):
if event.button == 1 and self.ctrl:
self.delete_frame(event)
elif event.button == 1:
self.delete_point(event)
elif event.button == 2: # press mouse-middle
self.undo(event)
elif event.button == 3 and self.ctrl:
self.stop = True
self.Backward()
elif event.button == 3:
self.stop = True
self.Forward()
def undo(self, event):
self.ind_del[self.jt] = []
out = fit1d((self.xexp[self.jt], self.yexp[self.jt], self.yerr[self.jt], \
self.rho_spl, self.fit_tol, self.sigmas, self.method ))
nx_fit = len(self.rho_spl)
nx_in = self.xexp.shape[1]
self.y_fit[self.jt] = out[: nx_fit]
self.dat_ok[self.jt] = out[nx_fit: nx_fit+nx_in]
self.confid[self.jt] = out[nx_fit+nx_in: ]
self.rhomap(self.jt)
if self.jt in self.tnot:
self.tnot.remove(self.jt)
self.del2d[self.jt].set_data([], [])
self.del2d[self.jt].set_3d_properties([])
self.update_plot()
def delete_frame(self, event):
if self.ax1d != event.inaxes:
return
if not self.stop:
return
if self.jt not in self.tnot:
self.tnot.append(self.jt)
self.ind_del[self.jt] = range(len(self.xexp[self.jt,:]))
self.del2d[self.jt], = self.ax2d.plot(self.x_out, self.tg[self.jt], self.d_out[self.jt, :], 'b-')
self.update_plot()
def delete_point(self, event):
if self.ax1d != event.inaxes:
return
if not self.stop:
return
sx = self.ax1d.get_xlim()
dx = sx[-1] - sx[0]
sy = self.ax1d.get_ylim()
dy = sy[-1] - sy[0]
x = np.zeros(self.yexp.shape[1])
y = np.zeros(self.yexp.shape[1])
yerr = np.zeros(self.yexp.shape[1])
x[:] = self.xexp[self.jt, :]
y[:] = self.yexp[self.jt, :]
yerr[:] = self.yerr[self.jt, :]
y[self.ind_del[self.jt]] = np.nan
yerr[self.ind_del[self.jt]] = np.infty
idel = np.nanargmin(np.hypot((x - event.xdata)/dx, (y - event.ydata)/dy))
self.ind_del[self.jt].append(idel)
y[idel] = np.nan
if np.sum(np.isfinite(y)) > 1:
out = fit1d( (x, y, yerr, self.rho_spl, self.fit_tol, self.sigmas, self.method) )
nx_fit = len(self.rho_spl)
nx_in = self.xexp.shape[1]
self.y_fit[self.jt] = out[: nx_fit]
self.dat_ok[self.jt] = out[nx_fit: nx_fit+nx_in]
self.confid[self.jt] = out[nx_fit+nx_in: ]
self.rhomap(self.jt)
else: # Treat as delete_frame
self.tnot.append(self.jt)
self.ind_del[self.jt] = range(len(self.xexp[self.jt,:]))
self.del2d[self.jt], = self.ax2d.plot(self.x_out, self.tg[self.jt], self.d_out[self.jt, :], 'b-')
self.update_plot()
def on_key(self, event):
if 'control' == event.key:
self.ctrl = True
if ' ' == event.key:
if self.stop:
self.Play()
else:
self.Pause()
def off_key(self, event):
if event.key in ('ctrl+control', 'control') :
self.ctrl = False
#------------
# Animation
#------------
def Pause(self):
self.stop = True
self.play_button['image'] = self.playfig
self.play_button['command'] = self.Play
def Forward(self):
self.stop = True
self.play_button['image'] = self.playfig
self.play_button['command'] = self.Play
self.jt += 1
self.jt = self.jt%self.nt
self.update_plot()
def Backward(self):
self.stop = True
self.play_button['image'] = self.playfig
self.play_button['command'] = self.Play
self.jt -= 1
self.jt = self.jt%self.nt
self.update_plot()
def Play(self):
self.stop = False
self.play_button['image'] = self.pausefig
self.play_button['command'] = self.Pause
while self.jt < self.nt-1 and not self.stop:
self.jt += 1
self.update_plot()
self.can_1d.start_event_loop(self.timeout)
def updateTimeSlider(self, val=None):
if val is None:
progress = (self.tgrid[self.jt] - self.tgrid[0])/(self.tgrid[-1] - self.tgrid[0])
else:
progress = val
self.sl_time.val = progress
poly = self.sl_time.poly.get_xy()
poly[2:4, 0] = progress
self.sl_time.poly.set_xy(poly)
self.sl_time.valtext.set_text('%8.4f' %self.tgrid[self.jt])
def update_plot(self, val=None):
self.update_plot1d(val=val)
if hasattr(self, 'y_fit'):
self.update_plot2d(val=val)
def update_plot1d(self, val=None):
if val is not None:
t_arr = self.tgrid[0] + val*(self.tgrid[-1] - self.tgrid[0])
self.jt = np.argmin(np.abs(self.tgrid - t_arr))
rho = self.xexp[self.jt]
dat = self.yexp[self.jt]
self.tlab = '#%d Time = %7.4f [s]' %(self.nshot, self.tgrid[self.jt])
self.ax1d.set_title(self.tlab, fontsize=fsize)
ymin = min(np.min(dat), 0)
if hasattr(self, 'y_fit'):
yfit = self.y_fit[self.jt]
sigma = self.confid[self.jt]
ymax = 1.01*max(np.max(yfit), np.max(dat))
if self.rho_lbl == 'rho_pol12':
xfit = self.rho12_out
else:
xfit =self.rho_out
self.pfit.set_data(xfit, yfit)
if np.max(sigma) > 0:
if hasattr(self, 'fill'):
self.fill.remove()
self.fill = self.ax1d.fill_between(xfit, yfit - sigma, yfit + sigma, facecolor='r', edgecolor='None', alpha=0.1)
rho_del = rho[self.ind_del[self.jt]]
dat_del = dat[self.ind_del[self.jt]]
self.del_pt.set_data(rho_del , dat_del )
ind_ok = (np.isfinite(self.dat_ok[self.jt]))
else:
ind_ok = np.ones(len(rho), dtype=bool)
ymax = 1.01*np.max(dat)
x = rho[ind_ok]
y = dat[ind_ok]
err = self.yerr[self.jt, ind_ok]
self.ok.set_data( x, y)
self.rej.set_data(rho[~ind_ok], dat[~ind_ok])
# Find the ending points of the errorbars
err_pos = (x, y - err), (x, y + err)
# Update the caplines
try:
for j, pos in enumerate(err_pos):
self.caplines[j].set_data(pos)
except:
print('Problem setting calpines')
#Update the error bars
self.barlinecols[0].set_segments(zip(zip(x, y - err), zip(x, y + err)))
self.ax1d.set_ylim([ymin, ymax])
self.sep_line.set_ydata([ymin, ymax])
self.updateTimeSlider(val=val)
self.can_1d.draw()
def update_plot2d(self, val=None):
if val is not None:
t_arr = self.tgrid[0] + val*(self.tgrid[-1] - self.tgrid[0])
self.jt = np.argmin(np.abs(self.tgrid - t_arr))
if self.nb2d.index('current') == 0:
self.sing2d.set_data(self.x_out, self.tg[self.jt])
self.sing2d.set_3d_properties(self.d_out[self.jt])
n_updates = 1000
dn = int(self.nt/n_updates) + 1
if self.stop or self.jt%dn == 0:
self.can_2d.draw()
else:
tim = self.tgrid[self.jt]
self.line.set_data([tim, tim], [self.x_out[0], self.x_out[-1]])
self.can_ct.draw()
class GEO_NEMEC:
timeout = 1e-10
def __init__(self, nshot, fmframe=None):
if fmframe is None:
if __name__ == '__main__':
fmframe = tk.Tk()
import trgui_style
else:
fmframe = tk.Toplevel()
fmframe.geometry('1300x950')
# Widgets style
# Menu bar
menubar = tk.Menu(fmframe)
filemenu = tk.Menu(menubar, tearoff=0)
filemenu.add_command(label="Run NEMEC", command=self.surf)
filemenu.add_command(label="Store u-files", command=self.store_u)
filemenu.add_separator()
filemenu.add_command(label="Close", command=fmframe.destroy)
menubar.add_cascade(label = "File", menu=filemenu)
fmframe.config(menu = menubar)
canvframe = ttk.Frame(fmframe, height=850)
entframe = ttk.Frame(fmframe)
self.toolframe = ttk.Frame(fmframe)
for frame in canvframe, entframe, self.toolframe:
frame.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
self.nbplot = ttk.Notebook(canvframe)
self.nbplot.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
self.nbplot.bind('<Button-1>', self.on_click)
pol_frame = ttk.Frame(self.nbplot)
mom_frame = ttk.Frame(self.nbplot)
profframe = ttk.Frame(self.nbplot)
self.nbplot.add(pol_frame, text='Poloidal plot')
self.nbplot.add(mom_frame, text='Fourier moments')
self.nbplot.add(profframe, text='EQuilibrium profiles')
self.proffig = Figure(figsize=(4., 8.), dpi=100)
self.pol_fig = Figure(figsize=(6., 8.), dpi=100)
self.mom_fig = Figure(figsize=(6., 8.), dpi=100)
self.proffig.subplots_adjust(left=0.2, bottom=0.1, right=0.8 , top=0.9)
self.pol_fig.subplots_adjust(left=0.06, bottom=0.1, right=0.98, top=0.9)
self.mom_fig.subplots_adjust(left=0.08, bottom=0.1, right=0.98, top=0.9)
self.profcanvas = FigureCanvasTkAgg(self.proffig, master=profframe)
self.pol_canvas = FigureCanvasTkAgg(self.pol_fig, master=pol_frame)
self.mom_canvas = FigureCanvasTkAgg(self.mom_fig, master=mom_frame)
self.profcanvas._tkcanvas.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
self.pol_canvas._tkcanvas.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
self.mom_canvas._tkcanvas.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
#-----------------
# Initialise plots
#-----------------
self.pol_fig.text(.5, .99, r'NEMEC equilibrium, constant $\rho$ and constant $\theta$ contours', \
ha='center', va='top')
self.mom_fig.text(.5, .99, 'NEMEC Fourier moments', ha='center', va='top')
# Profiles
fsize = 16
self.qsub = self.proffig.add_subplot(311)
self.qsub.set_xlim(0,1)
self.qsub.set_xlabel(r'$\rho_{tor}$',fontsize=fsize)
self.qsub.set_ylabel('q profile')
self.psub = self.proffig.add_subplot(312)
self.psub.set_xlim((0, 1))
self.psub.set_xlabel(r'$\rho_{tor}$',fontsize=fsize)
self.psub.set_ylabel('Pressure [Pa]')
self.rbsub = self.proffig.add_subplot(313)
self.rbsub.set_xlim(0,1)
self.rbsub.set_xlabel(r'$\rho_{tor}$',fontsize=fsize)
self.rbsub.set_ylabel('R*B$_\phi$')
#---------------
# Player buttons
locdir = os.path.dirname(os.path.realpath(__file__))
self.playfig = tk.PhotoImage(file='%s/play.gif' %locdir)
self.pausefig = tk.PhotoImage(file='%s/pause.gif' %locdir)
self.forwardfig = tk.PhotoImage(file='%s/forward.gif' %locdir)
self.backwardfig = tk.PhotoImage(file='%s/backward.gif' %locdir)
backward_button = ttk.Button(self.toolframe, command=self.Backward,image=self.backwardfig)
self.play_button = ttk.Button(self.toolframe, command=self.Play, image=self.playfig)
forward_button = ttk.Button(self.toolframe, command=self.Forward, image=self.forwardfig)
createToolTip(forward_button , 'Go to next timestep')
createToolTip(backward_button , 'Go backward by one timestep')
createToolTip(self.play_button, 'Forward animation/pause')
for but in backward_button, self.play_button, forward_button:
but.pack(side=tk.LEFT)
# Entry frame
geoinit = {'shot': nshot, 'Nmom': 7, 'ntheta': 101, 'nrho': 81, \
'exp': 'AUGD', 'dia': 'EQH', 'ed': 0, \
'tbeg': 4, 'tend': 6., 'DeltaT': 0.1}
self.geodict = {}
for key in ('shot', 'tbeg', 'tend', 'DeltaT', \
'exp', 'dia', 'ed', \
'Nmom', 'ntheta', 'nrho'):
raw = ttk.Frame(entframe)
raw.pack(side=tk.TOP)
lbl = ttk.Label(entframe, text=key, width=5).pack(side=tk.LEFT)
var = ttk.Entry(entframe, width=5)
var.insert(0, geoinit[key])
var.pack(side=tk.LEFT, padx='2 7')
self.geodict[key] = var
self.psep = self.pol_fig.add_subplot(111, aspect='equal')
self.psep.set_xlim((0.8, 3.0))
self.psep.set_ylim((-1.4, 1.4))
self.psep.set_xlabel('R [m]')
self.psep.set_ylabel('z [m]')
nshot_in = int(geoinit['shot'])
gc_r, gc_z = map_equ_20180130.get_gc(nshot_in)
for key in gc_r.keys():
self.psep.plot(gc_r[key], gc_z[key], 'b-')
toolbar = NavigationToolbar2TkAgg(self.pol_canvas, pol_frame)
toolbar.update()
toolbar = NavigationToolbar2TkAgg(self.mom_canvas, mom_frame)
toolbar.update()
toolbar = NavigationToolbar2TkAgg(self.profcanvas, profframe)
toolbar.update()
if __name__ == '__main__':
fmframe.mainloop()
def surf(self):
self.nshot = int(self.geodict['shot'].get())
dianam = self.geodict['dia'].get().strip()
expnam = self.geodict['exp'].get().strip()
ed = int(self.geodict['ed'].get().strip())
delta_t = float(self.geodict['DeltaT'].get().strip())
# Parameters for kk
npfm = 1001
mpfm = 1001
nrzmax = 1000
mom_type = len(momtyp)
self.mom_order = int(self.geodict['Nmom'].get().strip())
n_the_u = int(self.geodict['ntheta'].get().strip())
n_rho_u = int(self.geodict['nrho'].get().strip())
tbeg = float(self.geodict['tbeg'].get().strip())
tend = float(self.geodict['tend'].get().strip())
#==============
# kk-time grid
#==============
if eqm.Open(self.nshot, diag=dianam, exp=expnam, ed=ed):
eqm._read_profiles()
ind_ok = []
tim_old = -np.infty
# Allow delta_t
for jt, tim in enumerate(eqm.t_eq):
if (tim >= tbeg) and (tim <= tend) and (tim > tim_old + delta_t):
ind_ok.append(jt)
tim_old = tim
self.tarr = eqm.t_eq[ind_ok]
self.nt = len(self.tarr)
self.rho_u = np.linspace(0, 1, n_rho_u)
self.q_u = np.zeros((n_rho_u, self.nt))
self.pres_u = np.zeros((n_rho_u, self.nt))
self.rbphi_u = np.zeros((n_rho_u, self.nt))
vmec_dic = { 'nshot':self.nshot, 'mdescur':self.mom_order, \
'exp':expnam, 'dia':dianam, 'ed':ed, \
'nrho_nem': 101, 'nthe_nem': 60}
self.rsurf = np.zeros((self.nt, n_rho_u, n_the_u))
self.zsurf = np.zeros((self.nt, n_rho_u, n_the_u))
#===========
# Time loop
#===========
tf = eqm.tf[ ::-1, ind_ok]
q = - eqm.q[ ::-1, ind_ok]
pres = eqm.pres[ ::-1, ind_ok]
ffs = - 2e-7*eqm.jpol[::-1, ind_ok]
q[-1, :] = 2.*q[-2, :] - q[-3, :]
nmod = self.mom_order
n_mom = nmod*mom_type
self.mom_u = np.zeros((self.nt, n_rho_u, n_mom))
self.rc = np.zeros((self.nt, n_rho_u, nmod))
self.rs = np.zeros((self.nt, n_rho_u, nmod))
self.zc = np.zeros((self.nt, n_rho_u, nmod))
self.zs = np.zeros((self.nt, n_rho_u, nmod))
for jt in range(self.nt):
rhot = np.sqrt((tf[:, jt] - tf[0, jt])/(tf[-1, jt] - tf[0, jt]))
# Profiles: q, pressure, R*Bphi
self.q_u[:, jt] = np.interp(self.rho_u, rhot, q[:, jt])
self.pres_u[:, jt] = np.interp(self.rho_u, rhot, pres[:, jt])
self.rbphi_u[:, jt] = np.interp(self.rho_u, rhot, ffs[:, jt])
# NEMEC for magnetic moments
pool = Pool(cpu_count())
out = pool.map(nemec.nemec, [(vmec_dic, self.tarr[jt], '%d' %jt) for jt in range(self.nt)])
pool.close()
pool.join()
for jt in range(self.nt):
rho_nem = out[jt][4]
for jord in range(nmod):
# rho_nem is a grid equispaced in tor. flux, not in rho_tor
self.rc[jt, :, jord] = np.interp(self.rho_u, rho_nem, out[jt][0][:, jord])
self.rs[jt, :, jord] = np.interp(self.rho_u, rho_nem, out[jt][1][:, jord])
self.zc[jt, :, jord] = np.interp(self.rho_u, rho_nem, out[jt][2][:, jord])
self.zs[jt, :, jord] = np.interp(self.rho_u, rho_nem, out[jt][3][:, jord])
for jord in range(nmod):
fac = 1./(self.rho_u[1:]**jord)
self.mom_u[..., 1:, jord] = fac[:]*self.rc[..., 1:, jord]
self.mom_u[..., 1:, jord + nmod] = fac[:]*self.rs[..., 1:, jord]
self.mom_u[..., 1:, jord + 2*nmod] = fac[:]*self.zc[..., 1:, jord]
self.mom_u[..., 1:, jord + 3*nmod] = fac[:]*self.zs[..., 1:, jord]
self.mom_u[..., 0, jord] = self.rc[..., 0, jord]
self.mom_u[..., 0, jord + nmod] = 0.
self.mom_u[..., 0, jord + 2*nmod] = self.zc[..., 0, jord]
self.mom_u[..., 0, jord + 3*nmod] = 0.
self.mom_u *= 100.
self.rsurf, self.zsurf = mom2rz.mom2rz( self.rc, self.rs, \
self.zc, self.zs, nthe=n_the_u)
# End of time loop
self.begin()
def store_u(self):
tlbl = 'Time'.ljust(20) + 'Seconds'
mlbl = 'MOMENT INDEX'
dlbl = 'FOURIER MOMENTS'.ljust(20) + 'CM'
rholbl = 'RHO_TOR'
thlbl = 'THETA'.ljust(20) + 'RAD'
rlbl = 'MAJOR RADIUS'.ljust(20) + 'M'
zlbl = 'VERTICAL POSITION'.ljust(20) + 'M'
n_the_u = self.rsurf.shape[2]
theta = np.linspace(0, 2*np.pi, n_the_u)
# NUBEAM R,z=f(t,rho,theta)
uf_d = { 'pre': 'R', 'ext': 'SURF', 'shot': self.nshot, \
'grid': {'X': {'lbl': tlbl , 'arr': self.tarr}, \
'Y': {'lbl': rholbl, 'arr': self.rho_u}, \
'Z': {'lbl': thlbl , 'arr': theta} }, \
'data': {'lbl': rlbl, 'arr': self.rsurf} }
ufiles.WU(uf_d)
uf_d = { 'pre': 'Z', 'ext': 'SURF', 'shot': self.nshot, \
'grid': {'X': {'lbl': tlbl , 'arr': self.tarr}, \
'Y': {'lbl': rholbl, 'arr': self.rho_u}, \
'Z': {'lbl': thlbl , 'arr': theta} }, \
'data': {'lbl': zlbl, 'arr': self.zsurf} }
ufiles.WU(uf_d)
# q profile
qlbl = 'Q'
uf_d = { 'pre': 'Q', 'ext': 'EQ', 'shot': self.nshot, \
'grid': {'X': {'lbl': tlbl , 'arr': self.tarr}, \
'Y': {'lbl': rholbl, 'arr': self.rho_u} }, \
'data': {'lbl': qlbl, 'arr': self.q_u} }
ufiles.WU(uf_d)
# Pressure profile
plbl = 'P'.ljust(20) + '[Pa]'
uf_d={ 'pre': 'P', 'ext': 'EQ', 'shot': self.nshot, \
'grid': {'X': {'lbl': tlbl , 'arr': self.tarr}, \
'Y': {'lbl': rholbl, 'arr': self.rho_u} }, \
'data': {'lbl': plbl, 'arr': self.pres_u} }
ufiles.WU(uf_d)
# R*Bphi
rblbl = 'R*Bphi'.ljust(20) + '[T*m]'
uf_d={ 'pre': 'F', 'ext': 'EQ', 'shot': self.nshot, \
'grid': {'X': {'lbl': tlbl , 'arr': self.tarr}, \
'Y': {'lbl': rholbl, 'arr': self.rho_u} },
'data': {'lbl': rblbl, 'arr': self.rbphi_u} }
ufiles.WU(uf_d)
# Fourier moments
n_mom = self.mom_u.shape[-1]
marr = 1 + np.arange(n_mom)
uf_d={ 'pre': 'M', 'ext': 'MMX', 'shot': self.nshot, \
'grid': {'X': {'lbl': tlbl , 'arr': self.tarr}, \
'Y': {'lbl': rholbl, 'arr': self.rho_u}, \
'Z': {'lbl': mlbl , 'arr': marr} }, \
'data': {'lbl': dlbl, 'arr': self.mom_u} }
ufiles.WU(uf_d)
#------------
# Animation
#------------
def on_click(self, event):
if event.widget.identify(event.x, event.y) == 'label':
self.jtab = event.widget.index('@%d,%d' % (event.x, event.y))
self.update_plot()
def begin(self):
self.stop = True
self.jtab = 0
self.jt = 0
self.set_plots()
def Pause(self):
self.stop = True
self.play_button['image'] = self.playfig
self.play_button['command'] = self.Play
def Forward(self):
self.stop = True
self.play_button['image'] = self.playfig
self.play_button['command'] = self.Play
self.jt += 1
self.jt = self.jt%self.nt
self.update_plot()
def Backward(self):
self.stop = True
self.play_button['image'] = self.playfig
self.play_button['command'] = self.Play
self.jt -= 1
self.jt = self.jt%self.nt
self.update_plot()
def Play(self):
self.stop = False
self.play_button['image'] = self.pausefig
self.play_button['command'] = self.Pause
while self.jt < self.nt-1 and not self.stop:
self.jt += 1
self.update_plot()
self.pol_canvas.start_event_loop(self.timeout)
def set_plots(self):
# Scale bar
self.crntsc = tk.Scale(self.toolframe, command=self.jump, \
orient=tk.HORIZONTAL, length=300)
self.crntsc.pack(side=tk.LEFT)
# Poloidal plot
self.tim0 = self.pol_fig.text(.5, .96, '', ha='center', va='top')
self.scat = self.pol_fig.add_subplot(111, aspect='equal')
self.scat.set_xlim((0.8, 3.0))
self.scat.set_ylim((-1.4, 1.4))
self.scat.set_xlabel('R [m]')
self.scat.set_ylabel('z [m]')
gc_r, gc_z = map_equ_20180130.get_gc()
for key in gc_r.keys():
self.scat.plot(gc_r[key], gc_z[key], 'b-')
self.rhoplot = {}
self.theplot = {}
self.theplot2 = {}
nt, n_rho_u, n_the_u = self.rsurf.shape
for jrho in range(n_rho_u):
self.theplot[jrho], = self.scat.plot([], [], 'g-')
self.theplot2[jrho], = self.scat.plot([], [], 'g-')
for jthe in range(n_the_u):
self.rhoplot[jthe], = self.scat.plot([], [], 'r-')
# Moments
self.tim1 = self.mom_fig.text(.5, .96, '', ha='center', va='top')
mom_type = len(momtyp)
ncols = self.mom_order
nrows = mom_type
self.rcp = {}
self.rsp = {}
self.zcp = {}
self.zsp = {}
for jm in range(self.mom_order):
jplot = jm + 1
xtext = 0.1 + (float(jm)*0.94)/float(self.mom_order)
self.mom_fig.text(xtext, 0.93, '%dth moment' %jm, ha='center', va='top')
momp = self.mom_fig.add_subplot(nrows, ncols, jplot)
momp.set_xlabel(r'$\rho_{tor}$')
self.rcp[jm], = momp.plot(self.rho_u, self.rc[0, :, jm])
if jm == 0:
momp.set_ylabel('R cos')
momp = self.mom_fig.add_subplot(nrows, ncols, jplot + self.mom_order)
self.rsp[jm], = momp.plot(self.rho_u, self.rs[0, :, jm])
if jm == 0:
momp.set_ylabel('R sin')
momp = self.mom_fig.add_subplot(nrows, ncols, jplot + 2*self.mom_order)
self.zcp[jm], = momp.plot(self.rho_u, self.zc[0, :, jm])
if jm == 0:
momp.set_ylabel('z cos')
momp = self.mom_fig.add_subplot(nrows, ncols, jplot + 3*self.mom_order)
self.zsp[jm], = momp.plot(self.rho_u, self.zs[0, :, jm])
if jm == 0:
momp.set_ylabel('z sin')
# Profiles
self.tim2 = self.proffig.text(.5, .97, '', ha='center', va='top')
self.qplot, = self.qsub.plot( self.rho_u, self.q_u[:, 0])
self.pplot, = self.psub.plot( self.rho_u, self.pres_u[:, 0])
self.rbplot, = self.rbsub.plot(self.rho_u, self.rbphi_u[:, 0])
# Update
self.update_plot()
def jump(self, arg):
self.jt = int(float(arg)/100. * (self.nt-1))
self.update_plot(reset=False)
def update_plot(self, reset=True):
if reset:
self.crntsc.set(float(self.jt)/(self.nt - 1)*100.)
# return
tstr = '# %d, Time:%6.3f s' %(self.nshot, self.tarr[self.jt])
if self.jtab == 0: # Poloidal
nt, n_rho_u, n_the_u = self.rsurf.shape
for jrho in range(n_rho_u):
self.theplot[jrho].set_data(self.rsurf[self.jt, jrho, :], self.zsurf[self.jt, jrho, :])
self.theplot2[jrho].set_data(self.rsurf[self.jt, jrho, [-1, 0]], self.zsurf[self.jt, jrho, [-1, 0]])
for jthe in range(n_the_u):
self.rhoplot[jthe].set_data(self.rsurf[self.jt, :, jthe], self.zsurf[self.jt, :, jthe])
self.tim0.set_text(tstr)
self.pol_canvas.draw()
elif self.jtab == 1: # Moments
for jm in range(self.mom_order):
self.rcp[jm].set_ydata(self.rc[self.jt][:, jm])
self.rsp[jm].set_ydata(self.rs[self.jt][:, jm])
self.zcp[jm].set_ydata(self.zc[self.jt][:, jm])
self.zsp[jm].set_ydata(self.zs[self.jt][:, jm])
self.tim1.set_text(tstr)
self.mom_canvas.draw()
elif self.jtab == 2: # Profiles
self.qsub.set_ylim((0, 1.1*np.max(self.q_u[:, self.jt])))
self.qplot.set_ydata(self.q_u[:, self.jt])
self.psub.set_ylim((0, 1.1*np.max(self.pres_u[:, self.jt])))
self.pplot.set_ydata(self.pres_u[:, self.jt])
self.rbsub.set_ylim((0, 1.1*np.max(self.rbphi_u[:, self.jt])))
self.rbplot.set_ydata(self.rbphi_u[:, self.jt])
self.tim2.set_text(tstr)
self.profcanvas.draw()
class GEO_FULL:
timeout = 1e-10
def __init__(self, sshot):
if __name__ == '__main__':
fmframe = tk.Tk()
else:
fmframe = tk.Toplevel()
fmframe.geometry('1600x950')
# Menu bar
menubar = tk.Menu(fmframe)
filemenu = tk.Menu(menubar, tearoff=0)
filemenu.add_command(label="Run & fit", command=self.surf)
filemenu.add_command(label="Store u-files", command=self.store_u)
filemenu.add_command(label="Plot surfaces", command=self.plot_sur)
filemenu.add_command(label="Plot moments" , command=self.plot_mom)
filemenu.add_separator()
filemenu.add_command(label="Close", command=fmframe.destroy)
menubar.add_cascade(label = "File", menu=filemenu)
fmframe.config(menu = menubar)
canvframe = tk.Frame(fmframe, width=1600, height=800)
butframe = tk.Frame(fmframe)
entframe = tk.Frame(fmframe)
for frame in canvframe, butframe, entframe:
frame.pack(side=tk.TOP)
profframe = tk.Frame(canvframe, width=590)
pol_frame = tk.Frame(canvframe, width=1000)
for frame in profframe, pol_frame:
frame.pack(side=tk.TOP)
canvframe.pack_propagate(0)
entframe.pack_propagate(0)
profframe.pack_propagate(0)
pol_frame.pack_propagate(0)
self.proffig = Figure()
self.pol_fig = Figure((4, 7))
self.proffig.subplots_adjust(left=0.15, bottom=0.05, right=0.9 , top=0.9)
self.pol_fig.subplots_adjust(left=0.06, bottom=0.05, right=0.98, top=0.95)
self.profcanvas = FigureCanvasTkAgg(self.proffig, master=profframe)
self.pol_canvas = FigureCanvasTkAgg(self.pol_fig, master=pol_frame)
pol_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=1)
profframe.pack(side=tk.LEFT, fill=tk.BOTH, expand=1)
canvframe.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
butframe.pack (side=tk.TOP, fill=tk.BOTH, expand=1)
entframe.pack (side=tk.TOP, fill=tk.BOTH, expand=1)
self.profcanvas._tkcanvas.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
self.pol_canvas._tkcanvas.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
#-----------------
# Initialise plots
#-----------------
# Profiles
self.proffig.clf()
fsize = 16
self.qsub = self.proffig.add_subplot(311)
self.qsub.set_xlim(0,1)
self.qsub.set_xlabel(r'$\rho_{tor}$',fontsize=fsize)
self.qsub.set_ylabel('q profile')
self.qplot, = self.qsub.plot([], [], 'g-')
self.psub = self.proffig.add_subplot(312)
self.psub.set_xlim((0, 1))
self.psub.set_xlabel(r'$\rho_{tor}$',fontsize=fsize)
self.psub.set_ylabel('Pressure [Pa]')
self.pplot, = self.psub.plot([], [], 'm-')
self.rbsub = self.proffig.add_subplot(313)
self.rbsub.set_xlim(0,1)
self.rbsub.set_xlabel(r'$\rho_{tor}$',fontsize=fsize)
self.rbsub.set_ylabel('R*B$_\phi$')
self.rbplot, = self.rbsub.plot([], [], 'm-')
#---------------
# Player buttons
locdir = os.path.dirname(os.path.realpath(__file__))
self.playfig = tk.PhotoImage(file='%s/play.gif' %locdir)
self.pausefig = tk.PhotoImage(file='%s/pause.gif' %locdir)
self.forwardfig = tk.PhotoImage(file='%s/forward.gif' %locdir)
self.backwardfig = tk.PhotoImage(file='%s/backward.gif' %locdir)
backward_button = ttk.Button(self.toolframe, command=self.Backward,image=self.backwardfig)
self.play_button = ttk.Button(self.toolframe, command=self.Play, image=self.playfig)
forward_button = ttk.Button(self.toolframe, command=self.Forward, image=self.forwardfig)
createToolTip(forward_button , 'Go to next timestep')
createToolTip(backward_button , 'Go backward by one timestep')
createToolTip(self.play_button, 'Forward animation/pause')
for but in backward_button, self.play_button, forward_button:
but.pack(side=tk.LEFT)
plot_sur_btn = tk.Button(butframe, command=self.plot_sur, text='Surfaces')
plot_mom_btn = tk.Button(butframe, command=self.plot_mom, text='Moments')
createToolTip(forward_button , 'Go to next timestep')
createToolTip(backward_button , 'Go backward by one timestep')
createToolTip(self.play_button, 'Forward animation/pause')
sliderax = self.pol_fig.add_axes([0.1, 0.01, 0.6, 0.03], axisbg='yellow')
self.sl_time = Slider(sliderax,'Time:', 0., 1., valinit=0.)
self.sl_time.valtext.set_text('0.0000')
def update_time(val):
self.update_plot(val=val)
self.sl_time.on_changed(update_time)
self.cid2 = self.pol_fig.canvas.mpl_connect('scroll_event' , self.WheelInteraction)
self.cid5 = self.pol_fig.canvas.mpl_connect('button_press_event', lambda event:self.pol_fig.canvas._tkcanvas.focus_set())
for but in backward_button, self.play_button, forward_button, plot_sur_btn, plot_mom_btn:
but.pack(side=tk.LEFT)
# Entry frame
geoinit = {'shot':sshot, 'Nmom':7, 'ntheta':101, 'nrho':41, \
'exp':'AUGD', 'dia':'EQH', 'ed':0, \
'tbeg':4, 'tend':4.1, 'DeltaT':0.1}
self.geodict = {}
for key in ('shot', 'tbeg', 'tend', 'DeltaT', \
'exp', 'dia', 'ed', \
'Nmom', 'ntheta', 'nrho'):
raw = ttk.Frame(entframe)
raw.pack(side=tk.TOP)
lbl = ttk.Label(entframe, text=key, width=5).pack(side=tk.LEFT)
var = ttk.Entry(entframe, width=5)
var.insert(0, geoinit[key])
var.pack(side=tk.LEFT, padx='2 7')
self.geodict[key] = var
self.psep = self.pol_fig.add_subplot(111, aspect='equal')
self.psep.set_xlim((0.8, 3.0))
self.psep.set_ylim((-1.4, 1.4))
self.psep.set_xlabel('R [m]')
self.psep.set_ylabel('z [m]')
gc_r, gc_z = map_equ_20180130.get_gc()
for key in gc_r.iterjeys():
self.psep.plot(gc_r[key], gc_z[key], 'b-')
toolbar1 = NavigationToolbar2TkAgg(self.pol_canvas, pol_frame)
toolbar1.update()
toolbar2 = NavigationToolbar2TkAgg(self.profcanvas, profframe)
toolbar2.update()
self.pol_canvas.draw()
fmframe.mainloop()
def surf(self):
sshot = self.geodict['shot'].get().strip()
self.nshot = int(sshot)
dianam = self.geodict['dia'].get().strip()
expnam = self.geodict['exp'].get().strip()
ed = int(self.geodict['ed'].get().strip())
delta_t = float(self.geodict['DeltaT'].get().strip())
# Parameters
npfm = 1001
mpfm = 1001
nrzmax = 1000
mom_type = len(momtyp)
self.mom_order = int(self.geodict['Nmom'].get().strip())
n_the_u = int(self.geodict['ntheta'].get().strip())
n_rho_u = int(self.geodict['nrho'].get().strip())
tbeg = float(self.geodict['tbeg'].get().strip())
tend = float(self.geodict['tend'].get().strip())
n_mom = mom_type*self.mom_order
#----------------
# Initialise plot
#----------------
# Magnetic surfaces
#==============
# equ-time grid
#==============
if eqm.Open(self.nshot, diag=dianam, exp=expnam, ed=ed):
eqm._read_profiles()
ind_ok = []
tim_old = -np.infty
# Allow delta_t
for jt, tim in enumerate(eqm.t_eq):
if (tim >= tbeg) and (tim <= tend) and (tim > tim_old + delta_t):
ind_ok.append(jt)
tim_old = tim
self.tarr = eqm.t_eq[ind_ok]
self.nt = len(self.tarr)
self.rho_eq_grid = (1 + np.arange(n_rho_u))/float(n_rho_u)
self.rho1d = np.linspace(0, 1, n_rho_u)
self.rho1d[-1] = -0.999
self.q_u = np.zeros((n_rho_u, self.nt))
self.pres_u = np.zeros((n_rho_u, self.nt))
self.rbphi_u = np.zeros((n_rho_u, self.nt))
self.rsurf = np.zeros((self.nt, n_rho_u, n_the_u))
self.zsurf = np.zeros((self.nt, n_rho_u, n_the_u))
#===========
# Time loop
#===========
tf = eqm.tf[ ::-1, ind_ok]
q = - eqm.q[ ::-1, ind_ok]
pres = eqm.pres[ ::-1, ind_ok]
ffs = - 2e-7*eqm.jpol[::-1, ind_ok]
self.mom_u = np.zeros((self.nt, n_rho_u, n_mom))
self.rc = np.zeros((self.nt, n_rho_u, self.mom_order))
self.rs = np.zeros((self.nt, n_rho_u, self.mom_order))
self.zc = np.zeros((self.nt, n_rho_u, self.mom_order))
self.zs = np.zeros((self.nt, n_rho_u, self.mom_order))
self.rscat, self.zscat = eqm.rhoTheta2rz(self.rho1d, t_in=self.tarr, coord_in='rho_tor')
for jt, tshot in enumerate(self.tarr):
print('\n\nTIME = %8.4f\n\n' %tshot)
rhot = np.sqrt((tf[:, jt] - tf[0, jt])/(tf[-1, jt] - tf[0, jt]))
# Profiles: q, pressure, R*Bphi
self.q_u[:, jt] = np.interp(self.rho1d, rhot, q[:, jt])
self.q_u[-1, jt] = 2.*self.q_u[-2, jt] - self.q_u[-3, jt]
self.pres_u[:, jt] = np.interp(self.rho1d, rhot, pres[:, jt])
self.rbphi_u[:, jt] = np.interp(self.rho1d, rhot, ffs[:, jt])
for jrho in range(2, n_rho_u):
arr = descu(self.rscat[jt][jrho], self.zscat[jt][jrho], self.mom_order)
self.rc[jt, jrho, :] = arr[:, 0]
self.rs[jt, jrho, :] = arr[:, 1]
self.zc[jt, jrho, :] = arr[:, 2]
self.zs[jt, jrho, :] = arr[:, 3]
print self.rho1d[jrho], self.tarr[jt]
for jord in range(self.mom_order):
fac = 1./(self.rho_eq_grid**jord)
self.mom_u[..., :, jord] = fac[:]*self.rc[..., :, jord]
self.mom_u[..., :, jord + self.mom_order] = fac[:]*self.rs[..., :, jord]
self.mom_u[..., :, jord + 2*self.mom_order] = fac[:]*self.zc[..., :, jord]
self.mom_u[..., :, jord + 3*self.mom_order] = fac[:]*self.zs[..., :, jord]
self.rsurf, self.zsurf = mom2rz.mom2rz( self.rc, self.rs, \
self.zc, self.zs, nthe=n_the_u)
self.begin()
def store_u(self):
tlbl = 'Time'.ljust(20) + 'Seconds'
mlbl = 'MOMENT INDEX'
dlbl = 'FOURIER MOMENTS'.ljust(20) + 'CM'
rholbl = 'RHO_TOR'
thlbl = 'THETA'.ljust(20) + 'RAD'
rlbl = 'MAJOR RADIUS'.ljust(20) + 'M'
zlbl = 'VERTICAL POSITION'.ljust(20) + 'M'
sshot = str(self.nshot)
n_the_u = self.rsurf.shape[2]
theta = np.linspace(0, 2*np.pi, n_the_u)
# NUBEAM R,z=f(t,rho,theta)
uf_d = { 'pre':'R', 'ext':'SURF', 'shot':sshot, \
'grid': {'X': {'lbl':tlbl , 'arr':self.tarr}, \
'Y': {'lbl':rholbl, 'arr':self.rho_eq_grid}, \
'Z': {'lbl':thlbl , 'arr':theta} }, \
'data': {'lbl':rlbl, 'arr':self.rsurf} }
ufiles.WU(uf_d)
uf_d = { 'pre':'Z', 'ext':'SURF', 'shot':sshot, \
'grid': {'X': {'lbl':tlbl , 'arr':self.tarr}, \
'Y': {'lbl':rholbl, 'arr':self.rho_eq_grid}, \
'Z': {'lbl':thlbl , 'arr':theta} }, \
'data': {'lbl':zlbl, 'arr':self.zsurf} }
ufiles.WU(uf_d)
# q profile
qlbl = 'Q'
uf_d = { 'pre':'Q', 'ext':'EQ', 'shot':sshot, \
'grid': {'X':{'lbl':tlbl , 'arr':self.tarr}, \
'Y':{'lbl':rholbl, 'arr':self.rho1d} }, \
'data': {'lbl':qlbl, 'arr':self.q_u} }
ufiles.WU(uf_d)
# Pressure profile
plbl = 'P'.ljust(20) + '[Pa]'
uf_d={ 'pre':'P', 'ext':'EQ', 'shot':sshot, \
'grid': {'X':{'lbl':tlbl , 'arr':self.tarr}, \
'Y':{'lbl':rholbl, 'arr':self.rho1d} }, \
'data': {'lbl':plbl, 'arr':self.pres_u} }
ufiles.WU(uf_d)
# R*Bphi
rblbl = 'R*Bphi'.ljust(20) + '[T*m]'
uf_d={ 'pre':'F', 'ext':'EQ', 'shot':sshot, \
'grid': {'X':{'lbl':tlbl , 'arr':self.tarr}, \
'Y':{'lbl':rholbl, 'arr':self.rho1d} },
'data': {'lbl':rblbl, 'arr':self.rbphi_u} }
ufiles.WU(uf_d)
# Fourier moments
n_mom = self.mom_u.shape[-1]
marr = 1 + np.arange(n_mom)
self.mom_u *= 100.
uf_d={ 'pre':'M', 'ext':'MMX', 'shot':sshot, \
'grid': {'X':{'lbl':tlbl , 'arr':self.tarr}, \
'Y':{'lbl':rholbl, 'arr':self.rho_eq_grid}, \
'Z':{'lbl':mlbl , 'arr':marr} }, \
'data': {'lbl':dlbl, 'arr':self.mom_u} }
ufiles.WU(uf_d)
#------------
# Animation
#------------
def WheelInteraction(self, event):
self.jt += event.step
self.jt = min(self.jt, self.nt-1)
self.update_plot()
def begin(self):
self.stop = True
self.mom_flag = False
self.set_sur()
self.jt = 0
self.update_plot()
def Pause(self):
self.stop = True
self.play_button['image'] = self.playfig
self.play_button['command'] = self.Play
def Forward(self):
self.stop = True
self.play_button['image'] = self.playfig
self.play_button['command'] = self.Play
self.jt += 1
self.jt = self.jt%self.nt
self.update_plot()
def Backward(self):
self.stop = True
self.play_button['image'] = self.playfig
self.play_button['command'] = self.Play
self.jt -= 1
self.jt = self.jt%self.nt
self.update_plot()
def Play(self):
self.stop = False
self.play_button['image'] = self.pausefig
self.play_button['command'] = self.Pause
while self.jt < self.nt-1 and not self.stop:
self.jt += 1
self.update_plot()
self.pol_canvas.start_event_loop(self.timeout)
def set_sur(self):
# Magnetic surfaces from NEMEC, interpolated on regular rho grid
self.pol_fig.clf()
self.pol_fig.text(.5, .95, r'NEMEC equilibrium, constant $\rho$ and constant $\theta$ contours', \
ha='center', va='top')
self.scat = self.pol_fig.add_subplot(111, aspect='equal')
self.scat.set_xlim((0.8, 3.0))
self.scat.set_ylim((-1.4, 1.4))
self.scat.set_xlabel('R [m]')
self.scat.set_ylabel('z [m]')
gc_r, gc_z = map_equ_20180130.get_gc()
for key in gc_r.iterjeys():
self.scat.plot(gc_r[key], gc_z[key], 'b-')
self.rhoplot = {}
self.theplot = {}
self.theplot2 = {}
nt, n_rho_u, n_the_u = self.rsurf.shape
for jrho in range(n_rho_u):
self.theplot[jrho], = self.scat.plot([], [], 'g-')
self.theplot2[jrho], = self.scat.plot([], [], 'g-')
for jthe in range(n_the_u):
self.rhoplot[jthe], = self.scat.plot([], [], 'r-')
sliderax = self.pol_fig.add_axes([0.1, 0.01, 0.6, 0.03], axisbg='yellow')
self.sl_time = Slider(sliderax,'Time:', 0., 1., valinit=0.)
self.sl_time.valtext.set_text('0.0000')
def update_time(val):
self.update_plot(val=val)
self.sl_time.on_changed(update_time)
self.cid2 = self.pol_fig.canvas.mpl_connect('scroll_event' , self.WheelInteraction)
self.cid5 = self.pol_fig.canvas.mpl_connect('button_press_event', lambda event:self.pol_fig.canvas._tkcanvas.focus_set())
def set_mom(self):
self.pol_fig.clf()
self.pol_fig.text(.5, .95, r'NEMEC Fourier moments', \
ha='center', va='top')
mom_type = len(momtyp)
ncols = self.mom_order
nrows = mom_type
self.rcp = {}
self.rsp = {}
self.zcp = {}
self.zsp = {}
for jm in range(self.mom_order):
jplot = jm + 1
xtext = 0.1 + (float(jm)*0.94)/float(self.mom_order)
self.pol_fig.text(xtext, 0.95, '%dth moment' %jm)
momp = self.pol_fig.add_subplot(nrows, ncols, jplot)
momp.set_xlabel(r'$\rho_{tor}$')
# self.rcp[jm], = momp.plot([], [], 'b-')
self.rcp[jm], = momp.plot(self.rho_eq_grid, self.rc[0, :, jm])
if jm == 0:
momp.set_ylabel('R cos')
momp = self.pol_fig.add_subplot(nrows, ncols, jplot + self.mom_order)
# self.rsp[jm], = momp.plot([], [], 'b-')
self.rsp[jm], = momp.plot(self.rho_eq_grid, self.rs[0, :, jm])
if jm == 0:
momp.set_ylabel('R sin')
momp = self.pol_fig.add_subplot(nrows, ncols, jplot + 2*self.mom_order)
# self.zcp[jm], = momp.plot([], [], 'b-')
self.zcp[jm], = momp.plot(self.rho_eq_grid, self.zc[0, :, jm])
if jm == 0:
momp.set_ylabel('z cos')
momp = self.pol_fig.add_subplot(nrows, ncols, jplot + 3*self.mom_order)
# self.zsp[jm], = momp.plot([], [], 'b-')
self.zsp[jm], = momp.plot(self.rho_eq_grid, self.zs[0, :, jm])
if jm == 0:
momp.set_ylabel('z sin')
sliderax = self.pol_fig.add_axes([0.1, 0.01, 0.6, 0.03], axisbg='yellow')
self.sl_time = Slider(sliderax,'Time:', 0., 1., valinit=0.)
self.sl_time.valtext.set_text('0.0000')
def update_time(val):
self.update_plot(val=val)
self.sl_time.on_changed(update_time)
self.cid2 = self.pol_fig.canvas.mpl_connect('scroll_event' , self.WheelInteraction)
self.cid5 = self.pol_fig.canvas.mpl_connect('button_press_event', lambda event:self.pol_fig.canvas._tkcanvas.focus_set())
def plot_sur(self):
self.mom_flag = False
self.set_sur()
self.update_plot()
def plot_mom(self):
self.mom_flag = True
self.set_mom()
self.update_plot()
def updateTimeSlider(self, val=None):
if val is None:
progress = (self.tarr[self.jt] - self.tarr[0])/(self.tarr[-1] - self.tarr[0])
else:
progress = val
self.sl_time.val = progress
poly = self.sl_time.poly.get_xy()
poly[2:4, 0] = progress
self.sl_time.poly.set_xy(poly)
self.sl_time.valtext.set_text('%8.4f' %self.tarr[self.jt])
def update_plot(self, val=None):
if val is not None:
t_arr = self.tarr[0] + val*(self.tarr[-1] - self.tarr[0])
self.jt = np.argmin(np.abs(self.tarr - t_arr))
self.qsub.set_ylim((0, 1.1*np.max(self.q_u[:, self.jt])))
self.qplot.set_data(self.rho1d, self.q_u[:, self.jt])
self.psub.set_ylim((0, 1.1*np.max(self.pres_u[:, self.jt])))
self.pplot.set_data(self.rho1d, self.pres_u[:, self.jt])
self.rbsub.set_ylim((0, 1.1*np.max(self.rbphi_u[:, self.jt])))
self.rbplot.set_data(self.rho1d, self.rbphi_u[:, self.jt])
self.profcanvas.draw()
if self.mom_flag:
for jm in range(self.mom_order):
self.rcp[jm].set_data(self.rho_eq_grid, self.rc[self.jt][:, jm])
self.rsp[jm].set_data(self.rho_eq_grid, self.rs[self.jt][:, jm])
self.zcp[jm].set_data(self.rho_eq_grid, self.zc[self.jt][:, jm])
self.zsp[jm].set_data(self.rho_eq_grid, self.zs[self.jt][:, jm])
else:
nt, n_rho_u, n_the_u = self.rsurf.shape
for jrho in range(n_rho_u):
self.theplot[jrho].set_data(self.rsurf[self.jt, jrho, :], self.zsurf[self.jt, jrho, :])
self.theplot2[jrho].set_data(self.rsurf[self.jt, jrho, [-1, 0]], self.zsurf[self.jt, jrho, [-1, 0]])
for jthe in range(n_the_u):
self.rhoplot[jthe].set_data(self.rsurf[self.jt, :, jthe], self.zsurf[self.jt, :, jthe])
self.updateTimeSlider(val=val)
self.pol_canvas.draw()
