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()