1import os 2import wx 3import wx.lib.agw.flatnotebook as fnb 4import abipy.gui.awx as awx 5 6from monty.string import marquee 7from monty.collections import AttrDict 8from wx.py.shell import Shell 9from abipy.abilab import abiopen 10from abipy.electrons import SigresPlotter 11from abipy.gui.scissors import ScissorsBuilderFrame 12from abipy.gui.baseviewer import MultiViewerFrame 13from abipy.gui import mixins as mix 14from abipy.gui.kpoints import SpinKpointBandPanel 15 16 17class SigresViewerFrame(MultiViewerFrame, mix.Has_Structure, mix.Has_MultipleEbands, mix.Has_Tools, mix.Has_NetcdfFiles): 18 VERSION = "0.1" 19 20 HELP_MSG = """Quick help: 21 22 Kpoint list: 23 24 Left-Click: to visualize table with QP results for the selected spin, k-point 25 Right-Click: display popup menu with choices. 26 27Also, these key bindings can be used 28(For Mac OSX, replace 'Ctrl' with 'Apple'): 29 30 Ctrl-Q: quit 31""" 32 33 @property 34 def codename(self): 35 return "SigmaResultsViewer" 36 37 @property 38 def active_sigres(self): 39 """The active SIGRES file i.e. the SIGRES associated to the active tab.""" 40 return self.active_tab.sigres 41 42 @property 43 def structure(self): 44 """`Structure` associated to the active tab.""" 45 return self.active_sigres.structure 46 47 @property 48 def ebands(self): 49 """The electron bands associated to the active SIGRES file.""" 50 return self.active_sigres.ebands 51 52 @property 53 def ebands_list(self): 54 """List of `ElectronBands`.""" 55 ebands_list = [] 56 for page in range(self.notebook.GetPageCount()): 57 tab = self.notebook.GetPage(page) 58 ebands_list.append(tab.sigres.ebands) 59 return ebands_list 60 61 @property 62 def ebands_filepaths(self): 63 """ 64 Return a list with the absolute paths of the files 65 from which the `ElectronBands` have been read. 66 """ 67 paths = [] 68 for page in range(self.notebook.GetPageCount()): 69 tab = self.notebook.GetPage(page) 70 paths.append(tab.sigres.filepath) 71 return paths 72 73 @property 74 def nc_filepaths(self): 75 """String with the absolute path of the netcdf file.""" 76 paths = [] 77 for page in range(self.notebook.GetPageCount()): 78 tab = self.notebook.GetPage(page) 79 paths.append(tab.sigres.filepath) 80 return paths 81 82 @property 83 def sigres_filepaths(self): 84 """ 85 Return a list with the absolute paths of the files 86 from which the `SigresFile` have been read. 87 """ 88 paths = [] 89 for page in range(self.notebook.GetPageCount()): 90 tab = self.notebook.GetPage(page) 91 paths.append(tab.sigres.filepath) 92 return paths 93 94 def makeMenu(self): 95 """Creates the main menu.""" 96 # Base menu 97 menu_bar = super(SigresViewerFrame, self).makeMenu() 98 99 # Add mixin menus. 100 menu_bar.Append(self.CreateStructureMenu(), "Structure") 101 menu_bar.Append(self.CreateEbandsMenu(), "Ebands") 102 menu_bar.Append(self.CreateQpDataMenu(), "QP Data") 103 menu_bar.Append(self.CreateToolsMenu(), "Tools") 104 menu_bar.Append(self.CreateNetcdfMenu(), "Netcdf") 105 106 # Help menu 107 help_menu = self.makeHelpMenu() 108 menu_bar.Append(help_menu, "Help") 109 110 self.SetMenuBar(menu_bar) 111 112 def CreateQpDataMenu(self): 113 """Build and return the QP menu.""" 114 # Structure Menu ID's 115 self.ID_QPGAPS_COMPARE = wx.NewId() 116 self.ID_QPENES_COMPARE = wx.NewId() 117 118 menu = wx.Menu() 119 menu.Append(self.ID_QPGAPS_COMPARE, "Compare QP-gaps", "Compare QP gaps") 120 self.Bind(wx.EVT_MENU, self.OnQpGapsCompare, id=self.ID_QPGAPS_COMPARE) 121 122 menu.Append(self.ID_QPENES_COMPARE, "Compare QP-enes", "Compare QP energies") 123 self.Bind(wx.EVT_MENU, self.OnQpEnesCompare, id=self.ID_QPENES_COMPARE) 124 125 return menu 126 127 def OnQpGapsCompare(self, event): 128 """Compare the QP gaps reported in the SIGRES files.""" 129 # Instanciate the plotter and add the filepaths to the plotter. 130 plotter = SigresPlotter() 131 plotter.add_files(self.sigres_filepaths) 132 133 # Plot the convergence of the QP gaps. 134 plotter.plot_qpgaps(title="Convergence of QP gaps", hspan=0.05) 135 136 def OnQpEnesCompare(self, event): 137 """Compare the QP energies reported in the SIGRES files.""" 138 # Instanciate the plotter and add the filepaths to the plotter. 139 plotter = SigresPlotter() 140 plotter.add_files(self.sigres_filepaths) 141 142 # Plot the convergence of the QP energies. 143 plotter.plot_qpenes(title="Convergence of QP energies", hspan=0.05) 144 145 def makeToolBar(self): 146 """Creates the toolbar.""" 147 self.toolbar = toolbar = self.CreateToolBar() 148 toolbar.SetToolBitmapSize(wx.Size(48, 48)) 149 150 def bitmap(path): 151 return wx.Bitmap(awx.path_img(path)) 152 153 self.ID_SCISSORS = wx.NewId() 154 self.ID_PLOTQPSE0 = wx.NewId() 155 self.ID_PLOTKSWITHMARKS = wx.NewId() 156 157 artBmp = wx.ArtProvider.GetBitmap 158 #toolbar.AddSimpleTool(wx.ID_OPEN, artBmp(wx.ART_FILE_OPEN, wx.ART_TOOLBAR), "Open") 159 toolbar.AddSimpleTool(self.ID_PLOTQPSE0, bitmap("qpresults.png"), "Plot QPState Results.") 160 toolbar.AddSimpleTool(self.ID_PLOTKSWITHMARKS, bitmap("qpmarkers.png"), "Plot KS energies with QPState markers.") 161 toolbar.AddSimpleTool(self.ID_SCISSORS, bitmap("qpscissor.png"), "Build energy-dependent scissors from GW correction.") 162 163 # Associate menu/toolbar items with their handlers. 164 menu_handlers = [ 165 (self.ID_PLOTQPSE0, self.OnPlotQpsE0), 166 (self.ID_PLOTKSWITHMARKS, self.OnPlotKSwithQPmarkers), 167 (self.ID_SCISSORS, self.OnScissors), 168 ] 169 170 for combo in menu_handlers: 171 mid, handler = combo[:2] 172 self.Bind(wx.EVT_MENU, handler, id=mid) 173 174 self.toolbar.Realize() 175 176 def addFileTab(self, parent, filepath): 177 """Open a SIGRES file and create a new tab in the notebook.""" 178 gsr = abiopen(filepath) 179 tab = SigresFileTab(self.notebook, gsr) 180 self.notebook.AddPage(tab, os.path.basename(filepath)) 181 182 def OnPlotQpsE0(self, event): 183 """Plot QPState results as function of the KS energy.""" 184 if self.active_sigres is None: return 185 self.active_sigres.plot_qps_vs_e0() 186 187 def OnPlotKSwithQPmarkers(self, event): 188 """Plot KS energies with QPState markers.""" 189 if self.active_sigres is None: return 190 sigres = self.active_sigres 191 192 band_range = (sigres.min_gwbstart, sigres.max_gwbstop) 193 try: 194 EbandsWithMarkersPlotFrame(self, sigres.ebands, band_range=band_range).Show() 195 except: 196 awx.showErrorMessage(self) 197 198 def OnScissors(self, event): 199 """Build the scissors operator.""" 200 if self.active_sigres is None: return 201 ScissorsBuilderFrame(self, self.active_sigres.filepath).Show() 202 203 204 205class SigresFileTab(wx.Panel): 206 """Tab showing information on a single SIGRES file.""" 207 def __init__(self, parent, sigres, **kwargs): 208 """ 209 Args: 210 parent: 211 parent window. 212 sigres: 213 `SigresFile` object 214 """ 215 super(SigresFileTab, self).__init__(parent, -1, **kwargs) 216 self.sigres = sigres 217 218 splitter = wx.SplitterWindow(self, style=wx.SP_LIVE_UPDATE) 219 splitter.SetSashGravity(0.95) 220 221 self.skb_panel = SpinKpointBandPanel(splitter, sigres.structure, sigres.nsppol, sigres.gwkpoints, sigres.max_gwbstop, 222 bstart=sigres.min_gwbstart) 223 224 # Set the callback for double click on k-point row.. 225 self.Bind(self.skb_panel.MYEVT_SKB_ACTIVATED, self.onShowQPTable) 226 227 # Add Python shell 228 msg = "SIGRES_File object is accessible via the sigres variable. Use sigres.<TAB> to access the list of methods." 229 msg = marquee(msg, width=len(msg) + 8, mark="#") 230 msg = "#"*len(msg) + "\n" + msg + "\n" + "#"*len(msg) + "\n" 231 232 pyshell = Shell(splitter, introText=msg, locals={"sigres": sigres}) 233 splitter.SplitHorizontally(self.skb_panel, pyshell) 234 235 sizer = wx.BoxSizer(wx.VERTICAL) 236 sizer.Add(splitter, 1, wx.EXPAND, 5) 237 self.SetSizerAndFit(sizer) 238 239 @property 240 def statusbar(self): 241 return self.viewer_frame.statusbar 242 243 @property 244 def viewer_frame(self): 245 """The parent frame `WfkViewerFrame`.""" 246 try: 247 return self._viewer_frame 248 249 except AttributeError: 250 self._viewer_frame = self.getParentWithType(SigresViewerFrame) 251 return self._viewer_frame 252 253 def onShowQPTable(self, event): 254 """Show a table with the QP results for the selected spin, kpoint, band.""" 255 spin, kpoint, band = event.skb 256 qplist = self.sigres.get_qplist(spin, kpoint) 257 table = qplist.to_table() 258 title = "spin: %d, kpoint: %s" % (spin, kpoint) 259 260 awx.SimpleGridFrame(self, table, labels_from_table=True, title=title).Show() 261 262 263class EbandsWithMarkersPlotFrame(awx.Frame): 264 265 def __init__(self, parent, ebands, **kwargs): 266 """ 267 Args: 268 parent: 269 Parent window 270 ebands: 271 `ElectronBands` object. 272 """ 273 self.band_range = kwargs.pop("band_range", None) 274 275 super(EbandsWithMarkersPlotFrame, self).__init__(parent, -1, **kwargs) 276 self.SetTitle("Select parameters") 277 278 self.ebands = ebands 279 280 if not ebands.markers: 281 raise awx.Error("Found empty markers dictionary in ebands %s" % repr(ebands)) 282 283 self.BuildUi() 284 285 def BuildUi(self): 286 287 main_sizer = wx.BoxSizer(wx.VERTICAL) 288 289 hsizer1 = wx.BoxSizer( wx.HORIZONTAL ) 290 291 marker_label = wx.StaticText(self, -1, "Available Markers:") 292 marker_label.Wrap(-1) 293 hsizer1.Add(marker_label, 0, wx.ALIGN_CENTER_VERTICAL|wx.ALL, 5 ) 294 295 marker_choices = list(self.ebands.markers.keys()) 296 self.marker_choice = wx.Choice( self, -1, wx.DefaultPosition, wx.DefaultSize, marker_choices, 0 ) 297 self.marker_choice.SetSelection(0) 298 hsizer1.Add(self.marker_choice, 0, wx.ALIGN_CENTER_VERTICAL|wx.ALL, 5 ) 299 300 scale_label = wx.StaticText(self, -1, "Scale Factor:") 301 scale_label.Wrap(-1) 302 hsizer1.Add(scale_label, 0, wx.ALIGN_CENTER_VERTICAL|wx.TOP|wx.BOTTOM|wx.LEFT, 5 ) 303 304 self.scale_ctrl = wx.SpinCtrlDouble(self, id=-1, value=str(50), min=0.0, max=1.e+10, inc=50) 305 self.scale_ctrl.SetToolTipString("Multiplicative factor used to render the markers more visible.") 306 hsizer1.Add( self.scale_ctrl, 0, wx.ALL, 5 ) 307 308 main_sizer.Add( hsizer1, 0, wx.ALIGN_CENTER_HORIZONTAL, 5 ) 309 310 hsizer2 = wx.BoxSizer( wx.HORIZONTAL ) 311 312 ok_button = wx.Button(self, wx.ID_OK, label='Ok') 313 cancel_button = wx.Button(self, wx.ID_CANCEL, label='Cancel') 314 315 ok_button.Bind(wx.EVT_BUTTON, self.OnOkButton) 316 cancel_button.Bind(wx.EVT_BUTTON, self.OnCloseButton) 317 318 hsizer2.Add(ok_button, 0, wx.ALL, 5 ) 319 hsizer2.Add(cancel_button, 0, wx.ALL, 5 ) 320 321 main_sizer.Add( hsizer2, 0, wx.ALIGN_CENTER_HORIZONTAL, 5 ) 322 323 self.SetSizerAndFit(main_sizer) 324 325 def OnCloseButton(self, event): 326 self.Destroy() 327 328 def GetParams(self): 329 return AttrDict( 330 qpattr=self.marker_choice.GetStringSelection(), 331 fact=float(self.scale_ctrl.GetValue())) 332 333 def OnOkButton(self, event): 334 p = self.GetParams() 335 with_marker = p.qpattr + ":" + str(p.fact) 336 337 self.ebands.plot(marker=with_marker, band_range=self.band_range) 338 339 340class SigresViewerApp(awx.App): 341 pass 342 343 344def wxapp_sigresviewer(sigres_filepaths): 345 app = SigresViewerApp() 346 frame = SigresViewerFrame(None, filepaths=sigres_filepaths) 347 app.SetTopWindow(frame) 348 frame.Show() 349 return app 350