# JonKragskow/OrbiPlot

Fetching contributors…
Cannot retrieve contributors at this time
478 lines (373 sloc) 15.9 KB
 ######################################## # # # # # OrbiPlot by J G C Kragskow # # Alpha Build # # # # # ######################################## from numpy import arange, exp, sqrt import wx import matplotlib matplotlib.use('WXAgg') from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg as FigureCanvas from matplotlib.backends.backend_wx import NavigationToolbar2Wx from matplotlib.figure import Figure r = arange(0,200,0.01) orbitals = ["1s", "2s", "3s", "4s","5s","6s","2p","3p","4p","5p","6p","3d","4d","5d","6d","4f","5f"] def s_1(r,mode): rho = 2.*r/1. O1s = 2.*exp(-rho/2) if mode==1: return r**2.*O1s**2. if mode==2: return O1s def s_2(r,mode): rho = 2.*r/2. O2s = 1/(2*sqrt(2))*(2-rho)*exp(-rho/2) if mode==1: return r**2.*O2s**2. if mode==2: return O2s def s_3(r,mode): rho = 2.*r/3. O3s = 1/(9*sqrt(3))*(6-6*rho+rho**2)*exp(-rho/2) if mode==1: return r**2.*O3s**2. if mode==2: return O3s def s_4(r,mode): rho = 2.*r/4. O4s = (1./96)*(24-36*rho+12*rho**2-rho**3)*exp(-rho/2) if mode==1: return r**2.*O4s**2. if mode==2: return O4s def s_5(r,mode): rho = 2.*r/5. O5s = (1/(300*sqrt(5)))*(120-240*rho+120*rho**2-20*rho**3+rho**4)*exp(-rho/2) if mode==1: return r**2.*O5s**2. if mode==2: return O5s def s_6(r,mode): rho = 2.*r/6. O6s = (1/(2160*sqrt(6)))*(720-1800*rho+1200*rho**2-300*rho**3+30*rho**4-rho**5)*exp(-rho/2) if mode==1: return r**2.*O6s**2. if mode==2: return O6s def p_2(r,mode): rho = 2.*r/2. O2p = 1/(2*sqrt(6))*rho*exp(-rho/2) if mode==1: return r**2.*O2p**2. if mode==2: return O2p def p_3(r,mode): rho = 2.*r/3. O3p = 1/(9*sqrt(6))*rho*(4-rho)*exp(-rho/2) if mode==1: return r**2.*O3p**2. if mode==2: return O3p def p_4(r,mode): rho = 2.*r/4. O4p = 1/(32*sqrt(15))*rho*(20-10*rho+rho**2)*exp(-rho/2) if mode==1: return r**2.*O4p**2. if mode==2: return O4p def p_5(r,mode): rho = 2.*r/5. O5p = 1/(150*sqrt(30))*rho*(120-90*rho+18*rho**2-rho**3)*exp(-rho/2) if mode==1: return r**2.*O5p**2. if mode==2: return O5p def p_6(r,mode): rho = 2.*r/6. O6p = 1/(432*sqrt(210))*rho*(840-840*rho+252*rho**2-28*rho**3+rho**4)*exp(-rho/2) if mode==1: return r**2.*O6p**2. if mode==2: return O6p def d_3(r,mode): rho = 2.*r/3. O3d = 1/(9*sqrt(30))*rho**2*exp(-rho/2) if mode==1: return r**2.*O3d**2. if mode==2: return O3d def d_4(r,mode): rho = 2.*r/4. O4d = 1/(96*sqrt(5))*(6-rho)*rho**2*exp(-rho/2) if mode==1: return r**2.*O4d**2. if mode==2: return O4d def d_5(r,mode): rho = 2.*r/5. O5d = 1/(150*sqrt(70))*(42-14*rho+rho**2)*rho**2*exp(-rho/2) if mode==1: return r**2.*O5d**2. if mode==2: return O5d def d_6(r,mode): rho = 2.*r/6. O6d = 1/(864*sqrt(105))*(336-168*rho+24*rho**2-rho**3)*rho**2*exp(-rho/2) if mode==1: return r**2.*O6d**2. if mode==2: return O6d def f_4(r,mode): rho = 2.*r/4. O4f = 1/(96*sqrt(35))*rho**3*exp(-rho/2); if mode==1: return r**2.*O4f**2. if mode==2: return O4f def f_5(r,mode): rho = 2.*r/5. O5f = 1/(300*sqrt(70))*(8-rho)*rho**3*exp(-rho/2); if mode==1: return r**2.*O5f**2. if mode==2: return O5f def f_6(r,mode): rho = 2.*r/5. O6f = 1/(2592*sqrt(35))*(rho**2-18*rho+72)*rho**3*exp(-rho/2); if mode==1: return r**2.*O6f**2. if mode==2: return O6f #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% GUI CODE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% app = wx.App() frame = wx.Frame(None,-1,'OrbiPlot Alpha Build',style = wx.MAXIMIZE_BOX | wx.SYSTEM_MENU | wx.RESIZE_BORDER | wx.CLOSE_BOX | wx.CAPTION | wx.MINIMIZE_BOX |wx.DEFAULT_FRAME_STYLE, size = (1000,600)) checkbox_1s = wx.CheckBox(frame, wx.ID_ANY, "1s") checkbox_2s = wx.CheckBox(frame, wx.ID_ANY, "2s") checkbox_2p = wx.CheckBox(frame, wx.ID_ANY, "2p") checkbox_3s = wx.CheckBox(frame, wx.ID_ANY, "3s") checkbox_3p = wx.CheckBox(frame, wx.ID_ANY, "3p") checkbox_3d = wx.CheckBox(frame, wx.ID_ANY, "3d") checkbox_4s = wx.CheckBox(frame, wx.ID_ANY, "4s") checkbox_4p = wx.CheckBox(frame, wx.ID_ANY, "4p") checkbox_4d = wx.CheckBox(frame, wx.ID_ANY, "4d") checkbox_4f = wx.CheckBox(frame, wx.ID_ANY, "4f") checkbox_5s = wx.CheckBox(frame, wx.ID_ANY, "5s") checkbox_5p = wx.CheckBox(frame, wx.ID_ANY, "5p") checkbox_5d = wx.CheckBox(frame, wx.ID_ANY, "5d") checkbox_5f = wx.CheckBox(frame, wx.ID_ANY, "5f") checkbox_6s = wx.CheckBox(frame, wx.ID_ANY, "6s") checkbox_6p = wx.CheckBox(frame, wx.ID_ANY, "6p") checkbox_6d = wx.CheckBox(frame, wx.ID_ANY, "6d") checkbox_6f = wx.CheckBox(frame, wx.ID_ANY, "6f") checkbox_SaveFileType = wx.CheckBox(frame, wx.ID_ANY, "Png (Untick) or Svg (Tick)") checkbox_RadWav = wx.CheckBox(frame, wx.ID_ANY, "Radial Wavefunction") checkbox_RadDist = wx.CheckBox(frame, wx.ID_ANY, "Radial Distribution Function") checkbox_Gridlines = wx.CheckBox(frame, wx.ID_ANY, "Gridlines") UpperXText = wx.TextCtrl(frame, wx.ID_ANY, "0", style=wx.TE_CENTRE) LowerXText = wx.TextCtrl(frame, wx.ID_ANY, "100", style=wx.TE_CENTRE) LinewidthSlider = wx.Slider(frame, wx.ID_ANY, 3, 0, 5, style=wx.SL_HORIZONTAL | wx.SL_LABELS) PlotButton = wx.Button(frame, wx.ID_ANY, "Plot") SaveButton = wx.Button(frame, wx.ID_ANY, "Save") UntickButton = wx.Button(frame, wx.ID_ANY, "Untick All") UpperXLabel = wx.StaticText(frame, wx.ID_ANY, "Upper x limit") LowerXLabel = wx.StaticText(frame, wx.ID_ANY, "Lower x limit") LinewidthLabel = wx.StaticText(frame, wx.ID_ANY, "Linewidth") plotzone = Figure() a = plotzone.add_subplot(111) canvas = FigureCanvas(frame, -1, plotzone) #Create sizer for entire window, split down the middle LHS=plot, RHS=options WholeWindowSizer = wx.GridSizer(rows=1, cols=2, hgap=10, vgap=0) #Add plot to LHS of window inside an inset box PlotSizer = wx.StaticBoxSizer(wx.StaticBox(frame, -1),wx.VERTICAL) PlotSizer.Add(canvas,wx.EXPAND) WholeWindowSizer.Add(PlotSizer, 1, wx.ALIGN_RIGHT | wx.EXPAND) #Create column for RHS elements RHSSizer = wx.GridSizer(rows=7, cols=1, hgap=1, vgap=10) RHSSizer.Add((5,5), proportion=1) #Create grid for orbital checkboxes and add boxes CheckBoxGridSizer = wx.GridSizer(rows=7, cols=4, hgap=5, vgap=150) CheckBoxGridSizer.Add(checkbox_1s, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(0, 0, wx.ALIGN_CENTER,0) CheckBoxGridSizer.Add(0, 0, wx.ALIGN_CENTER,0) CheckBoxGridSizer.Add(0, 0, wx.ALIGN_CENTER,0) CheckBoxGridSizer.Add(checkbox_2s, 1, wx.ALIGN_CENTER,0) CheckBoxGridSizer.Add(checkbox_2p, 1, wx.ALIGN_CENTER,0) CheckBoxGridSizer.Add(0, 0, wx.ALIGN_CENTER,0) CheckBoxGridSizer.Add(0, 0, wx.ALIGN_CENTER,0) CheckBoxGridSizer.Add(checkbox_3s, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(checkbox_3p, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(checkbox_3d, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(0, 0, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(checkbox_4s, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(checkbox_4p, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(checkbox_4d, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(checkbox_4f, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(checkbox_5s, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(checkbox_5p, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(checkbox_5d, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(checkbox_5f, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(checkbox_6s, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(checkbox_6p, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(checkbox_6d, 1, wx.ALIGN_CENTER, 0) CheckBoxGridSizer.Add(checkbox_6f, 1, wx.ALIGN_CENTER, 0) RHSSizer.Add(CheckBoxGridSizer,1, wx.EXPAND|wx.ALL, 0) #Add check boxes for gridlines and plot type PlotOptionsSizer = wx.GridSizer(rows=4, cols=3, hgap=5, vgap=100) PlotOptionsSizer.Add(checkbox_RadWav, 1, wx.ALIGN_CENTER, 0) PlotOptionsSizer.Add(checkbox_RadDist, 1, wx.ALIGN_CENTER, 0) PlotOptionsSizer.Add(checkbox_Gridlines, 1, wx.ALIGN_CENTER, 0) PlotOptionsSizer.Add(PlotButton, 0, wx.ALIGN_CENTER, 0) PlotOptionsSizer.Add(SaveButton, 0, wx.ALIGN_CENTER,0) PlotOptionsSizer.Add(UntickButton, 0, wx.ALIGN_CENTER, 0) PlotOptionsSizer.Add(0, 0, wx.ALIGN_CENTER, 0) PlotOptionsSizer.Add(checkbox_SaveFileType, 1, wx.ALIGN_CENTER, 0) RHSSizer.Add((5,5), proportion=1) RHSSizer.Add(PlotOptionsSizer, 0, wx.ALIGN_CENTER|wx.EXPAND , 5) #Add Limits entry box LimitsOptionsSizer = wx.GridSizer(rows=1, cols=4, hgap=5, vgap=10) LimitsOptionsSizer.Add(UpperXLabel, 0, wx.ALIGN_CENTER, 0) LimitsOptionsSizer.Add(UpperXText, 0, wx.ALIGN_CENTER, 0) LimitsOptionsSizer.Add(LowerXLabel, 0, wx.ALIGN_CENTER, 0) LimitsOptionsSizer.Add(LowerXText, 0, wx.ALIGN_CENTER, 0) RHSSizer.Add(LimitsOptionsSizer, 0, wx.ALIGN_CENTER , 5) #Add Thickness Slider SliderGridSizer = wx.GridSizer(rows=2, cols=1, hgap=5, vgap=1) SliderGridSizer.Add(LinewidthLabel, 0, wx.ALIGN_CENTER, 0) SliderGridSizer.Add(LinewidthSlider, 10,0,0,wx.ALIGN_CENTER) RHSSizer.Add(SliderGridSizer, 0, wx.ALIGN_CENTER) #Self Promotion SelfPromoSizer = wx.GridSizer(rows=3, cols=1, hgap=5, vgap=1) SelfPromoSizer.Add(wx.StaticText(frame, wx.ID_ANY, "OrbiPlot by J G C Kragskow"), 1, wx.ALIGN_CENTER, 0) SelfPromoSizer.Add(wx.StaticText(frame, wx.ID_ANY, "@kragskow"), 1, wx.ALIGN_CENTER, 0) SelfPromoSizer.Add(wx.StaticText(frame, wx.ID_ANY, "www.kragskow.com"), 1, wx.ALIGN_CENTER, 0) RHSSizer.Add(SelfPromoSizer, 0, wx.ALIGN_CENTER) WholeWindowSizer.Add(RHSSizer,1, wx.EXPAND|wx.ALIGN_CENTER) wx.Panel.SetSizer(frame,WholeWindowSizer) frame.Layout() def SaveButtonAction(event): SaveFileDialog = wx.FileDialog(frame, message="Save file as ...", style=wx.FD_SAVE | wx.FD_FILE_MUST_EXIST) SaveFileDialog.ShowModal() SaveLocation = SaveFileDialog.GetPath() fig = a.get_figure() #standard savefig command does not work on subplot if checkbox_SaveFileType.GetValue()==0: fig.savefig(SaveLocation,Transparent=True,format='png') elif checkbox_SaveFileType.GetValue()==1: fig.savefig(SaveLocation,Transparent=True,format='svg') SaveFileDialog.Destroy() def PlotButtonAction(event): x_list_1 = [] #clear list of x data to be plotted y_list_1 = [] #clear list of x data to be plotted a.clear() #clear axes nplots = 0 #create counter for number of lines in plot a.set_xlim((float(UpperXText.GetValue()), float(LowerXText.GetValue()))) #set x axis limits again based on input boxes ########Read radial wavefunction and probabilty checkboxes and assign plot flag to decide which gets plotted if checkbox_RadWav.GetValue()==checkbox_RadDist.GetValue():#If both are selected then set modeflag to plot nothing modeflag=0 elif checkbox_RadDist.GetValue()==1: #If probability is selected then set modeflag to plot radial distribution function modeflag=1 a.set_xlabel("$r / a_{0}$") #set xlabel a.set_ylabel("$4\pi r^{2}R(r)^{2}$") #set ylabel elif checkbox_RadWav.GetValue()==1:#If Wave Function is selected then set modeflag to plot radial wave function modeflag=2 a.set_xlabel("$r / a_{0}$") #set xlabel a.set_ylabel("$R(r)$") #set ylabel ##########Plot the actual lines using orbital functions if modeflag !=0:#If modeflag=1 or 2 if checkbox_1s.GetValue()==1: a.plot(r,s_1(r,modeflag),label='1s',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_2s.GetValue()==1: a.plot(r,s_2(r,modeflag),label='2s',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_3s.GetValue()==1: a.plot(r,s_3(r,modeflag),label='3s',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_4s.GetValue()==1: a.plot(r,s_4(r,modeflag),label='4s',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_5s.GetValue()==1: a.plot(r,s_5(r,modeflag),label='5s',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_6s.GetValue()==1: a.plot(r,s_6(r,modeflag),label='6s',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_2p.GetValue()==1: a.plot(r,p_2(r,modeflag),label='2p',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_3p.GetValue()==1: a.plot(r,p_3(r,modeflag),label='3p',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_4p.GetValue()==1: a.plot(r,p_4(r,modeflag),label='4p',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_5p.GetValue()==1: a.plot(r,p_5(r,modeflag),label='5p',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_6p.GetValue()==1: a.plot(r,p_6(r,modeflag),label='6p',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_3d.GetValue()==1: a.plot(r,d_3(r,modeflag),label='3d',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_4d.GetValue()==1: a.plot(r,d_4(r,modeflag),label='4d',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_5d.GetValue()==1: a.plot(r,d_5(r,modeflag),label='5d',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_6d.GetValue()==1: a.plot(r,d_6(r,modeflag),label='6d',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_4f.GetValue()==1: a.plot(r,f_4(r,modeflag),label='4f',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_5f.GetValue()==1: a.plot(r,f_5(r,modeflag),label='5f',linewidth=LinewidthSlider.GetValue()) nplots += 1 if checkbox_6f.GetValue()==1: a.plot(r,f_6(r,modeflag),label='6f',linewidth=LinewidthSlider.GetValue()) nplots += 1 plotzone.tight_layout() #Stops axis labels from clipping outside of figure area ######## Options for gridlines and plot legend if checkbox_Gridlines.GetValue()==1: a.grid(True) if checkbox_Gridlines.GetValue()==0: a.grid(False) if nplots > 0: a.legend(loc=0,fontsize=(10),numpoints=1,ncol=nplots/4+1,frameon=False) #only draw legend if >0 lines drawn ######## Draw the plot canvas.draw() #wxpython equivalent of plt.show(), displays the freshly assembled canvas def UntickButtonAction(event): checkbox_1s.SetValue(0) checkbox_2s.SetValue(0) checkbox_3s.SetValue(0) checkbox_4s.SetValue(0) checkbox_5s.SetValue(0) checkbox_6s.SetValue(0) checkbox_2p.SetValue(0) checkbox_3p.SetValue(0) checkbox_4p.SetValue(0) checkbox_5p.SetValue(0) checkbox_6p.SetValue(0) checkbox_3d.SetValue(0) checkbox_4d.SetValue(0) checkbox_5d.SetValue(0) checkbox_6d.SetValue(0) checkbox_4f.SetValue(0) checkbox_5f.SetValue(0) checkbox_6f.SetValue(0) PlotButton.Bind(wx.EVT_BUTTON, PlotButtonAction) SaveButton.Bind(wx.EVT_BUTTON, SaveButtonAction) UntickButton.Bind(wx.EVT_BUTTON, UntickButtonAction) frame.Show() frame.SetMinSize((1000,600)) frame.Centre() app.MainLoop()