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