1#!/usr/local/bin/python3.8
2"""Gui for the oncvpsp norm-conserving pseudopotential generator."""
3import os
4import copy
5import time
6import shutil
7import abc
8import sys
9import wx
10import wx.lib.mixins.listctrl as listmix
11import numpy as np
12
13from collections import OrderedDict
14from monty.dev import get_ncpus
15from monty.collections import AttrDict
16from pymatgen.core.periodic_table import Element
17from abipy.gui.editor import TextNotebookFrame, SimpleTextViewer
18from abipy.gui.oncvtooltips import oncv_tip
19from abipy.gui import mixins as mix
20from abipy.gui import awx
21from abipy.gui.awx.elements_gui import WxPeriodicTable, PeriodicPanel, ElementButton
22from abipy.flowtk import Pseudo
23
24try:
25    from pseudo_dojo.core.dojoreport import DojoReport
26    from pseudo_dojo.refdata.nist import database as nist
27    from pseudo_dojo.ppcodes.ppgen import OncvGenerator
28    from pseudo_dojo.ppcodes.oncvpsp import OncvOutputParser, MultiPseudoGenDataPlotter
29except ImportError as exc:
30    print("Error while trying to import pseudo_dojo modules:\n%s" % str(exc))
31    #raise
32
33# TODO
34# Change oncvpsp so that
35#   1) we always write the logarithmic derivative
36#   2) better error handling
37
38
39_char2l = {
40    "s": 0,
41    "p": 1,
42    "d": 2,
43    "f": 3,
44    "g": 4,
45    "h": 5,
46    "i": 6,
47}
48
49def char2l(char):
50    return _char2l[char]
51
52
53def all_symbols():
54    return [e.symbol for e in Element]
55
56
57def add_size(kwargs, size=(800, 600)):
58    """Add size to kwargs if not present."""
59    if "size" not in kwargs:
60        kwargs["size"] = size
61
62    return kwargs
63
64
65
66def my_periodic_table(parent):
67    """
68    A periodic table that allows the user to select the element
69    before starting the pseudopotential generation.
70    """
71    class MyElementButton(ElementButton):
72
73        def makePopupMenu(self):
74            # Get the menu of the super class.
75            menu = super(MyElementButton, self).makePopupMenu()
76
77            self.ID_POPUP_ONCVPSP = wx.NewId()
78            menu.Append(self.ID_POPUP_ONCVPSP, "Generate NC pseudo with oncvpsp")
79
80            # Associate menu/toolbar items with their handlers.
81            menu_handlers = [
82                (self.ID_POPUP_ONCVPSP, self.onOncvpsp),
83            ]
84
85            for combo in menu_handlers:
86                mid, handler = combo[:2]
87                self.Bind(wx.EVT_MENU, handler, id=mid)
88
89            return menu
90
91        def onOncvpsp(self, event):
92            """Open a frame for the initialization of oncvpsp."""
93            frame = OncvParamsFrame(self, self.Z)
94            frame.Show()
95
96    class MyPeriodicPanel(PeriodicPanel):
97        element_button_class = MyElementButton
98
99        def OnSelect(self, event):
100            # Get the atomic number Z, open a dialog to get basic configuration parameters from the user.
101            # The dialog will then generate the main Frame for the pseudo generation.
102            super(MyPeriodicPanel, self).OnSelect(event)
103            z = event.GetId() - 100
104            print("select", z)
105
106    class MyPeriodicTable(WxPeriodicTable):
107        periodic_panel_class = MyPeriodicPanel
108
109    return MyPeriodicTable(parent)
110
111
112class OncvParamsFrame(awx.Frame):
113    """
114    This frame allows the user to specify the most important parameters
115    to generate the pseudopotential once the chemical element has been selected.
116    """
117
118    HELP_MSG = """\
119Quick help:
120    Use this window to select the AE reference configuration and how
121    to separate states into core and valence.
122"""
123
124    def __init__(self, parent, z, **kwargs):
125        super(OncvParamsFrame, self).__init__(parent, **kwargs)
126        self.element = Element.from_Z(z)
127        self.buildUI()
128
129    def buildUI(self):
130        # Build controller with the dimensions.
131        panel = wx.Panel(self, -1)
132        self.wxdims = awx.RowMultiCtrl(self, ctrl_params=OrderedDict([
133                    ("nc", dict(dtype="i", tooltip="Number of core states")),
134                    ("nv", dict(dtype="i", tooltip="Number of valence states")),
135                    ("lmax", dict(dtype="i", tooltip="Maximum angular momentum for pseudo"))
136               ]))
137
138        # Initialize the quantum numbers of the AE atom with the ground-state configuration.
139        # E.g., The electronic structure for Fe is represented as:
140        # [(1, "s", 2), (2, "s", 2), (2, "p", 6), (3, "s", 2), (3, "p", 6), (3, "d", 6), (4, "s", 2)]
141        ele_struct = self.element.full_electronic_structure
142
143        ctrls = OrderedDict([
144            ("n", dict(dtype="i")),
145            ("l", dict(dtype="i")),
146            ("f", dict(dtype="f"))])
147
148        self.wxaeconf = awx.TableMultiCtrl(self, nrows=len(ele_struct), ctrls=ctrls)
149
150        for wxrow, (n, lchar, f) in zip(self.wxaeconf, ele_struct):
151            row = OrderedDict()
152            row["n"], row["l"], row["f"] = n, char2l(lchar), f
153            wxrow.SetParams(row)
154
155        add_button = wx.Button(self, -1, "Add row")
156        add_button.Bind(wx.EVT_BUTTON, self.onAddButton)
157        del_button = wx.Button(self, -1, "Delete row")
158        del_button.Bind(wx.EVT_BUTTON, self.onDelButton)
159        hsz = wx.BoxSizer(wx.HORIZONTAL)
160        hsz.Add(add_button)
161        hsz.Add(del_button)
162
163        main_sizer = wx.BoxSizer(wx.VERTICAL)
164        main_sizer.Add(hsz, 0,flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
165        main_sizer.Add(self.wxdims, 0, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
166        main_sizer.Add(self.wxaeconf, 0, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
167
168        help_button = wx.Button(self, wx.ID_HELP)
169        help_button.Bind(wx.EVT_BUTTON, self.onHelp)
170        main_sizer.Add(help_button, 0, flag=wx.ALL | wx.ALIGN_RIGHT)
171
172        self.SetSizerAndFit(main_sizer)
173
174    def get_oncv_params(self):
175        """Return the basic dimensions used in oncvpsp."""
176        print(self.wxaeconf.GetParams())
177        return AttrDict(
178            dims=self.wxdims.GetParams(),
179        )
180
181    def show_nist_lda_levels(self):
182        # Get the LDA levels of the neutral atom.
183        # (useful to decide if semicore states should be included in the valence).
184        entry = nist.get_neutral_entry(symbol=self.element.symbol)
185        frame = awx.Frame(self)
186        awx.ListCtrlFromTable(frame, table=entry.to_table())
187        frame.Show()
188
189    def onAddButton(self, event):
190        """Add a new row."""
191        self.get_oncv_params()
192
193    def onDelButton(self, event):
194        """Delete last row."""
195        self.show_nist_lda_levels()
196
197
198class WxOncvFrame(awx.Frame, mix.Has_Tools):
199    """The main frame of the GUI"""
200    VERSION = "0.1"
201
202    HELP_MSG = """\
203This window shows a template input file with the variables
204used to generated the pseudopotential. The `optimize` buttons
205allows you to scan a set of possible values for the generation of the pseudopotential.
206"""
207
208    def __init__(self, parent, filepath=None):
209        super(WxOncvFrame, self).__init__(parent, id=-1, title=self.codename)
210
211        # This combination of options for config seems to work on my Mac.
212        self.config = wx.FileConfig(appName=self.codename, localFilename=self.codename + ".ini",
213                                    style=wx.CONFIG_USE_LOCAL_FILE)
214
215        # Build menu, toolbar and status bar.
216        self.SetMenuBar(self.makeMenu())
217        self.statusbar = self.CreateStatusBar()
218        self.Centre()
219        self.makeToolBar()
220        #self.toolbar.Enable(False)
221
222        self.input_file = None
223        if filepath is not None:
224            if os.path.exists(filepath):
225                self.input_file = os.path.abspath(filepath)
226                self.BuildUI(notebook=OncvNotebook.from_file(self, filepath))
227            else:
228                # Assume symbol
229                self.BuildUI(notebook=OncvNotebook.from_symbol(self, filepath))
230        else:
231            self.BuildUI()
232
233    @property
234    def codename(self):
235        """Name of the application."""
236        return "WxOncvGui"
237
238    def BuildUI(self, notebook=None):
239        """Build user-interface."""
240        old_selection = None
241        if hasattr(self, "main_sizer"):
242            # Remove old notebook
243            main_sizer = self.main_sizer
244            main_sizer.Hide(0)
245            main_sizer.Remove(0)
246
247            # Save the active tab so that we can set it afterwards.
248            old_selection = self.notebook.GetSelection()
249            del self.notebook
250        else:
251            self.main_sizer = main_sizer = wx.BoxSizer(wx.VERTICAL)
252
253        if notebook is not None:
254            self.notebook = notebook
255        else:
256            # Default is, of course, silicon.
257            self.notebook = OncvNotebook.from_symbol(self, "Si")
258
259        main_sizer.Add(self.notebook, flag=wx.EXPAND)
260        self.SetSizerAndFit(main_sizer)
261        main_sizer.Layout()
262
263        # Reinstate the old selection
264        if old_selection is not None:
265            self.notebook.SetSelection(old_selection)
266
267    def AddFileToHistory(self, filepath):
268        """Add the absolute filepath to the file history."""
269        self.file_history.AddFileToHistory(filepath)
270        self.file_history.Save(self.config)
271        self.config.Flush()
272
273    def onFileHistory(self, event):
274        fileNum = event.GetId() - wx.ID_FILE1
275        filepath = self.file_history.GetHistoryFile(fileNum)
276        self.file_history.AddFileToHistory(filepath)
277
278        self.BuildUI(notebook=OncvNotebook.from_file(self, filepath))
279
280    def makeMenu(self):
281        """Creates the main menu."""
282        menu_bar = wx.MenuBar()
283
284        file_menu = wx.Menu()
285        file_menu.Append(wx.ID_OPEN, "&Open", help="Open an input file")
286        file_menu.Append(wx.ID_SAVE, "&Save", help="Save the input file")
287        file_menu.Append(wx.ID_CLOSE, "&Close", help="Close the Gui")
288        #file_menu.Append(wx.ID_EXIT, "&Quit", help="Exit the application")
289
290        # Make sub-menu with the list of supported visualizers.
291        symbol_menu = wx.Menu()
292        self._id2symbol = {}
293
294        for symbol in all_symbols():
295            _id = wx.NewId()
296            symbol_menu.Append(_id, symbol)
297            self._id2symbol[_id] = symbol
298            self.Bind(wx.EVT_MENU, self.onNewNotebookFromSymbol, id=_id)
299
300        file_menu.AppendMenu(-1, 'Template from element', symbol_menu)
301
302        file_history = self.file_history = wx.FileHistory(8)
303        file_history.Load(self.config)
304        recent = wx.Menu()
305        file_history.UseMenu(recent)
306        file_history.AddFilesToMenu()
307        file_menu.AppendMenu(-1, "&Recent Files", recent)
308        self.Bind(wx.EVT_MENU_RANGE, self.onFileHistory, id=wx.ID_FILE1, id2=wx.ID_FILE9)
309        menu_bar.Append(file_menu, "File")
310
311        # Add Mixin menus.
312        menu_bar.Append(self.CreateToolsMenu(), "Tools")
313
314        help_menu = wx.Menu()
315        help_menu.Append(wx.ID_HELP, "Help ", help="Quick help")
316        help_menu.Append(wx.ID_ABOUT, "About " + self.codename, help="Info on the application")
317        menu_bar.Append(help_menu, "Help")
318
319        # Associate menu/toolbar items with their handlers.
320        menu_handlers = [
321            (wx.ID_OPEN, self.onOpen),
322            (wx.ID_CLOSE, self.onClose),
323            #(wx.ID_EXIT, self.onExit),
324            (wx.ID_SAVE, self.onSave),
325            (wx.ID_HELP, self.onHelp),
326            (wx.ID_ABOUT, self.onAbout),
327        ]
328
329        for combo in menu_handlers:
330            mid, handler = combo[:2]
331            self.Bind(wx.EVT_MENU, handler, id=mid)
332
333        return menu_bar
334
335    def makeToolBar(self):
336        """Creates the toolbar."""
337        self.toolbar = toolbar = self.CreateToolBar()
338        self.toolbar.SetToolBitmapSize(wx.Size(48, 48))
339
340        def bitmap(path):
341            return wx.Bitmap(awx.path_img(path))
342
343        self.ID_SHOW_INPUT = wx.NewId()
344        self.ID_RUN_INPUT = wx.NewId()
345
346        toolbar.AddSimpleTool(self.ID_SHOW_INPUT, bitmap("in.png"), "Visualize the input file")
347        toolbar.AddSimpleTool(self.ID_RUN_INPUT, bitmap("run.png"), "Run the input file.")
348
349        toolbar.Realize()
350
351        # Associate menu/toolbar items with their handlers.
352        menu_handlers = [
353            (self.ID_SHOW_INPUT, self.onShowInput),
354            (self.ID_RUN_INPUT, self.onRunInput),
355        ]
356
357        for combo in menu_handlers:
358            mid, handler = combo[:2]
359            self.Bind(wx.EVT_MENU, handler, id=mid)
360
361    def onNewNotebookFromSymbol(self, event):
362        symbol = self._id2symbol[event.GetId()]
363        self.BuildUI(notebook=OncvNotebook.from_symbol(self, symbol))
364
365    def onShowInput(self, event):
366        """Show the input file in a new frame."""
367        text = self.notebook.makeInputString()
368        SimpleTextViewer(self, text=text, title="Oncvpsp Input").Show()
369
370    def onRunInput(self, event):
371        """Build a new generator from the input file, and add it to the queue."""
372        text = self.notebook.makeInputString()
373        try:
374            psgen = OncvGenerator(text, calc_type=self.notebook.calc_type)
375        except:
376            return awx.showErrorMessage(self)
377
378        frame = PseudoGeneratorsFrame(self, [psgen], title="Run Input")
379        frame.launch_psgens()
380        frame.Show()
381
382    def onOpen(self, event):
383        """Open a file"""
384        dialog = wx.FileDialog(self, message="Choose an input file", style=wx.OPEN)
385        if dialog.ShowModal() == wx.ID_CANCEL: return
386
387        filepath = dialog.GetPath()
388        dialog.Destroy()
389
390        # Add to the history.
391        self.file_history.AddFileToHistory(filepath)
392        self.file_history.Save(self.config)
393        self.config.Flush()
394
395        self.BuildUI(notebook=OncvNotebook.from_file(self, filepath))
396
397    def onSave(self, event):
398        """Save a file"""
399        dialog = wx.FileDialog(self, message="Save file as...", style=wx.SAVE | wx.OVERWRITE_PROMPT,
400                               wildcard="Dat files (*.dat)|*.dat")
401        if dialog.ShowModal() == wx.ID_CANCEL: return
402
403        filepath = dialog.GetPath()
404        dialog.Destroy()
405
406        # Add to the history.
407        self.file_history.AddFileToHistory(filepath)
408        self.file_history.Save(self.config)
409        self.config.Flush()
410
411        with open(filepath, "w") as fh:
412            fh.write(self.notebook.makeInputString())
413
414    def onClose(self, event):
415        """ Respond to the "Close" menu command."""
416        self.Destroy()
417
418    def onAbout(self, event):
419        return awx.makeAboutBox(
420            codename=self.codename,
421            version=self.VERSION,
422            description="oncvgui is a front-end for the pseudopotential generator oncvpsp",
423            developers=["Matteo Giantomassi"],
424            website="http://www.mat-simresearch.com/")
425
426
427class OptimizationFrame(awx.Frame, metaclass=abc.ABCMeta):
428    """Base class for optimization frames."""
429    def __init__(self, parent, **kwargs):
430        super(OptimizationFrame, self).__init__(parent, **kwargs)
431
432        # All optimization buttons are disabled when we start an optimization.
433        #self.main_frame.notebook.enable_all_optimize_buttons(False)
434        #self.main_frame.Enable(False)
435        #self.Bind(wx.EVT_WINDOW_DESTROY, self.onDestroy)
436
437    @property
438    def main_frame(self):
439        return self.getParentWithType(WxOncvFrame)
440
441    def onDestroy(self, event):
442        """Enable all optimize_buttons before destroying the Frame."""
443        #self.main_frame.notebook.enable_all_optimize_buttons(True)
444        return super(OptimizationFrame, self).Destroy()
445
446    def onCloseButton(self, event):
447        self.onDestroy(event)
448
449    def onOkButton(self, event):
450        """
451        Get input from user, generate new input files by changing some parameters
452        and open a new frame for running the calculations.
453        """
454        # Build the PseudoGenerators and open a new frame to run them.
455        psgens = []
456        for inp in self.build_new_inps():
457            try:
458                psgen = OncvGenerator(str(inp), calc_type=self.notebook.calc_type)
459                psgens.append(psgen)
460            except:
461                return awx.showErrorMessage(self)
462
463        frame = PseudoGeneratorsFrame(self, psgens, title=self.opt_type)
464        frame.launch_psgens()
465        frame.Show()
466
467    def make_buttons(self, parent=None):
468        """
469        Build the three buttons (Generate, Cancel, Help), binds them and return the sizer.
470        """
471        parent = self if parent is None else parent
472
473        add_opts = dict(flag=wx.ALIGN_CENTER_VERTICAL | wx.ALL, border=5)
474        gen_button = wx.Button(parent, wx.ID_OK, label='Generate')
475        gen_button.Bind(wx.EVT_BUTTON, self.onOkButton)
476
477        close_button = wx.Button(parent, wx.ID_CANCEL, label='Cancel')
478        close_button.Bind(wx.EVT_BUTTON, self.onCloseButton)
479
480        help_button = wx.Button(parent, wx.ID_HELP)
481        help_button.Bind(wx.EVT_BUTTON, self.onHelp)
482
483        hbox = wx.BoxSizer(wx.HORIZONTAL)
484        hbox.Add(gen_button, **add_opts)
485        hbox.Add(close_button, **add_opts)
486        hbox.Add(help_button, **add_opts)
487
488        return hbox
489
490    @abc.abstractproperty
491    def opt_type(self):
492        """Human-readable string describing the optimization type."""
493
494    @abc.abstractmethod
495    def build_new_inps(self):
496        """Returns a list of new inputs."""
497
498
499class LlocOptimizationFrame(OptimizationFrame):
500    """
501    This frame allows the user to optimize the parameters for the local part of the pseudopotential
502    """
503    HELP_MSG = """\
504This window allows you to change/optimize the parameters governing the local part"""
505
506    def __init__(self, parent, notebook, **kwargs):
507        """
508        Args:
509            notebook:
510                `OncvNotebool` containing the parameters of the template.
511        """
512        super(LlocOptimizationFrame, self).__init__(parent, **kwargs)
513
514        # Save reference to the input panel.
515        self.notebook = notebook
516
517        panel = wx.Panel(self, -1)
518        main_sizer = wx.BoxSizer(wx.VERTICAL)
519        add_opts = dict(flag=wx.ALIGN_CENTER_VERTICAL | wx.ALL, border=5)
520
521        #self.fcfact_ctl = awx.IntervalControl(self, start=0.25, num=6, step=0.05, choices=[">", "centered", "<"])
522        #check_l.SetToolTipString("Enable/Disable optimization for this l-channel")
523        #main_sizer.Add(self.fcfact_ctl, 1, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
524
525        buttons_sizer = self.make_buttons()
526        main_sizer.Add(buttons_sizer, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
527
528        self.SetSizerAndFit(main_sizer)
529
530    @property
531    def opt_type(self):
532        return "lloc optimization"
533
534    def build_new_inps(self):
535        # Generate new list of inputs.
536        base_inp = self.notebook.makeInput()
537        return base_inp.optimize_vloc()
538
539
540class Rc5OptimizationFrame(OptimizationFrame):
541    """
542    This frame allows the user to optimize the parameters for the local part of the pseudopotential
543    """
544    HELP_MSG = """\
545This window allows you to change/optimize the rc5 parameter"""
546
547    def __init__(self, parent, notebook, **kwargs):
548        """
549        Args:
550            notebook:
551                `OncvNotebool` containing the parameters of the template.
552        """
553        super(Rc5OptimizationFrame, self).__init__(parent, **kwargs)
554
555        # Save reference to the input panel.
556        self.notebook = notebook
557
558        panel = wx.Panel(self, -1)
559        main_sizer = wx.BoxSizer(wx.VERTICAL)
560        add_opts = dict(flag=wx.ALIGN_CENTER_VERTICAL | wx.ALL, border=5)
561
562        self.rc5_ctlr = awx.IntervalControl(self, start=0.25, num=6, step=0.05, choices=[">", "centered", "<"])
563        main_sizer.Add(self.rc5_ctlr, 1, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
564
565        buttons_sizer = self.make_buttons()
566        main_sizer.Add(buttons_sizer, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
567
568        self.SetSizerAndFit(main_sizer)
569
570    @property
571    def opt_type(self):
572        return "rc5 optimization"
573
574    def build_new_inps(self):
575        # Generate new list of inputs.
576        base_inp = self.notebook.makeInput()
577        # TODO
578        return base_inp.optimize_rc5()
579
580
581class DeblOptimizationFrame(OptimizationFrame):
582    """
583    This frame allows the user to optimize the VKB projectors
584    """
585
586    HELP_MSG = """\
587This window allows you to optimize the parameters used to construct the VKB projectors."""
588
589    def __init__(self, parent, notebook, **kwargs):
590        """
591        Args:
592            notebook:
593                Notebook containing the parameters of the template.
594        """
595        super(DeblOptimizationFrame, self).__init__(parent, **kwargs)
596
597        # Save reference to the input panel.
598        self.notebook = notebook
599        lmax = notebook.lmax
600
601        panel = wx.Panel(self, -1)
602        main_sizer = wx.BoxSizer(wx.VERTICAL)
603        add_opts = dict(flag=wx.ALIGN_CENTER_VERTICAL | wx.ALL, border=5)
604
605        # Build list of controls to allow the user to select the list of
606        # debl for the different l-channels. One can disable l-channels via checkboxes.
607        self.checkbox_l = [None] * (lmax + 1)
608        self.debl_range_l = [None] * (lmax + 1)
609
610        debl_l = notebook.makeInput().debl_l
611
612        for l in range(lmax + 1):
613            sbox_sizer = wx.StaticBoxSizer(wx.StaticBox(self, -1, "Angular Channel L=%d" % l), wx.VERTICAL)
614            vsz = wx.BoxSizer(wx.HORIZONTAL)
615
616            self.checkbox_l[l] = check_l = wx.CheckBox(self, -1)
617            check_l.SetToolTipString("Enable/Disable optimization for this l-channel")
618            check_l.SetValue(True)
619            self.debl_range_l[l] = debl_range_l = awx.IntervalControl(self, start=debl_l[l], num=3, step=0.5)
620
621            vsz.Add(check_l, **add_opts)
622            vsz.Add(debl_range_l, **add_opts)
623
624            sbox_sizer.Add(vsz, 1, wx.ALL, 5)
625            main_sizer.Add(sbox_sizer, 1, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
626
627        buttons_sizer = self.make_buttons()
628        main_sizer.Add(buttons_sizer, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
629
630        self.SetSizerAndFit(main_sizer)
631
632    @property
633    def opt_type(self):
634        return "Vkb optimization"
635
636    def build_new_inps(self):
637        # Get the list of angular channels activated and the corresponding arrays.
638        l_list, deblvals_list = [], []
639        for l, (checkbox, wxrange) in enumerate(zip(self.checkbox_l, self.debl_range_l)):
640            if checkbox.IsChecked():
641                l_list.append(l)
642                deblvals_list.append(wxrange.getValues())
643
644        # Generate new list of inputs.
645        base_inp = self.notebook.makeInput()
646        new_inps = []
647        for l, new_debls in zip(l_list, deblvals_list):
648            new_inps.extend(base_inp.optimize_debls_for_l(l=l, new_debls=new_debls))
649
650        return new_inps
651
652
653class FcfactOptimizationFrame(OptimizationFrame):
654    """
655    This frame allows the user to optimize the model core charge
656    """
657
658    HELP_MSG = """\
659This window allows you to change/optimize the parameters governing the model core charge"""
660
661    def __init__(self, parent, notebook, **kwargs):
662        """
663        Args:
664            notebook:
665                Notebook containing the parameters of the template.
666        """
667        super(FcfactOptimizationFrame, self).__init__(parent, **kwargs)
668
669        # Save reference to the input panel.
670        self.notebook = notebook
671
672        panel = wx.Panel(self, -1)
673        main_sizer = wx.BoxSizer(wx.VERTICAL)
674        add_opts = dict(flag=wx.ALIGN_CENTER_VERTICAL | wx.ALL, border=5)
675
676        self.fcfact_ctl = awx.IntervalControl(self, start=0.25, num=6, step=0.05, choices=[">", "centered", "<"])
677        #check_l.SetToolTipString("Enable/Disable optimization for this l-channel")
678        main_sizer.Add(self.fcfact_ctl, 1, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
679
680        buttons_sizer = self.make_buttons()
681        main_sizer.Add(buttons_sizer, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
682
683        self.SetSizerAndFit(main_sizer)
684
685    @property
686    def opt_type(self):
687        return "fcfact optimization"
688
689    def build_new_inps(self):
690        fcfact_list = self.fcfact_ctl.getValues()
691        print(fcfact_list)
692
693        # Generate new list of inputs.
694        base_inp = self.notebook.makeInput()
695        return base_inp.optimize_modelcore(fcfact_list, add_icmod0=True)
696
697
698class QcutOptimizationFrame(OptimizationFrame):
699    """
700    This frame allows the user to select the l-channels and
701    the list of values of qcut_l to be analyzed.
702    """
703
704    HELP_MSG = """\
705This window allows you to change/optimize the value of the qcut parameters for
706the different angular channel. Use the checkboxes to select the l-channel(s) to be
707analyzed, and the other controls to specify the list of qc values to test.
708"""
709
710    def __init__(self, parent, notebook, **kwargs):
711        """
712        Args:
713            notebook:
714                Notebook containing the parameters of the template.
715        """
716        super(QcutOptimizationFrame, self).__init__(parent, **kwargs)
717
718        # Save reference to the input panel.
719        self.notebook = notebook
720        lmax = notebook.lmax
721
722        panel = wx.Panel(self, -1)
723        main_sizer = wx.BoxSizer(wx.VERTICAL)
724        add_opts = dict(flag=wx.ALIGN_CENTER_VERTICAL | wx.ALL, border=5)
725
726        # Build list of controls to allow the user to select the list of
727        # qcuts for the different l-channels. One can disable l-channels via checkboxes.
728        self.checkbox_l = [None] * (lmax + 1)
729        self.wxqcut_range_l = [None] * (lmax + 1)
730
731        qcut_l = notebook.makeInput().qcut_l
732
733        for l in range(lmax + 1):
734            sbox_sizer = wx.StaticBoxSizer(wx.StaticBox(self, -1, "Angular Channel L=%d" % l), wx.VERTICAL)
735            vsz = wx.BoxSizer(wx.HORIZONTAL)
736
737            self.checkbox_l[l] = check_l = wx.CheckBox(self, -1)
738            check_l.SetToolTipString("Enable/Disable optimization for this l-channel")
739            check_l.SetValue(True)
740
741            self.wxqcut_range_l[l] = qcrange_l = awx.IntervalControl(self, start=qcut_l[l], num=4, step=0.5)
742
743            vsz.Add(check_l, **add_opts)
744            vsz.Add(qcrange_l, **add_opts)
745
746            sbox_sizer.Add(vsz, 1, wx.ALL, 5)
747            main_sizer.Add(sbox_sizer, 1, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
748
749        buttons_sizer = self.make_buttons()
750        main_sizer.Add(buttons_sizer, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
751
752        self.SetSizerAndFit(main_sizer)
753
754    @property
755    def opt_type(self):
756        return "Qcut optimization"
757
758    def build_new_inps(self):
759        # Get the list of angular channels activated and the corresponding arrays.
760        l_list, qcvals_list = [], []
761        for l, (checkbox, wxrange) in enumerate(zip(self.checkbox_l, self.wxqcut_range_l)):
762            if checkbox.IsChecked():
763                l_list.append(l)
764                qcvals_list.append(wxrange.getValues())
765        #print("\nl_list:", l_list, "\nqcvals_list", qcvals_list)
766
767        # Generate new list of inputs.
768        base_inp = self.notebook.makeInput()
769        new_inps = []
770        for l, new_qcuts in zip(l_list, qcvals_list):
771            new_inps.extend(base_inp.optimize_qcuts_for_l(l=l, new_qcuts=new_qcuts))
772
773        return new_inps
774
775
776class RcOptimizationFrame(OptimizationFrame):
777    """
778    This frame allows the user to select the l-channels and
779    the list of values of rc_l to be analyzed.
780    """
781
782    HELP_MSG = """\
783This window allows you to change/optimize the value of the rc parameters (core radius)
784for  the different angular channel."""
785
786    def __init__(self, parent, notebook, **kwargs):
787        """
788        Args:
789            notebook:
790                Notebook containing the parameters of the template.
791        """
792        super(RcOptimizationFrame, self).__init__(parent, **kwargs)
793
794        panel = wx.Panel(self, -1)
795        main_sizer = wx.BoxSizer(wx.VERTICAL)
796
797        # Save reference to the input panel.
798        self.notebook = notebook
799        lmax = notebook.lmax
800
801        # Build list of controls to allow the user to select the list of
802        # qcut values for the different l-channels.
803        # One can disable particular l-channels via checkboxes.
804        self.checkbox_l = [None] * (lmax + 1)
805        self.wxrc_range_l = [None] * (lmax + 1)
806
807        rc_l = notebook.makeInput().rc_l
808
809        add_opts = dict(flag=wx.ALIGN_CENTER_VERTICAL | wx.ALL, border=5)
810        for l in range(lmax + 1):
811            sbox_sizer = wx.StaticBoxSizer(wx.StaticBox(self, -1, "Angular Channel L=%d" % l), wx.VERTICAL)
812            vsz = wx.BoxSizer(wx.HORIZONTAL)
813
814            self.checkbox_l[l] = check_l = wx.CheckBox(self, -1)
815            check_l.SetToolTipString("Enable/Disable optimization for this l-channel")
816            check_l.SetValue(True)
817
818            self.wxrc_range_l[l] = qcrange_l = awx.IntervalControl(self, start=rc_l[l], num=4, step=0.05)
819
820            vsz.Add(check_l, **add_opts)
821            vsz.Add(qcrange_l, **add_opts)
822
823            sbox_sizer.Add(vsz, 1, wx.ALL, 5)
824            main_sizer.Add(sbox_sizer, 1, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
825
826        buttons_sizer = self.make_buttons()
827        main_sizer.Add(buttons_sizer, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
828
829        self.SetSizerAndFit(main_sizer)
830
831    @property
832    def opt_type(self):
833        return "Rc optimization"
834
835    def build_new_inps(self):
836        l_list, rc_list = [], []
837        for l, (checkbox, wxrange) in enumerate(zip(self.checkbox_l, self.wxrc_range_l)):
838            if checkbox.IsChecked():
839                l_list.append(l)
840                rc_list.append(wxrange.getValues())
841
842        # Generate new list of inputs.
843        base_inp = self.notebook.makeInput()
844        new_inps = []
845        for l, new_rcs in zip(l_list, rc_list):
846            new_inps.extend(base_inp.optimize_rcs_for_l(l, new_rcs))
847
848        return new_inps
849
850
851def empty_field(tag, oncv_dims=None):
852    """Returns an empty field."""
853    # Get the subclass from tag, initialize data with None and call __init__
854    cls = _FIELD_LIST[tag]
855    data = cls.make_empty_data(oncv_dims=oncv_dims)
856
857    return cls(tag, data, oncv_dims)
858
859
860class Field(object):
861    # Flags used to define the type of field
862    # Subclasses should define the class attribute type
863    FTYPE_ROW = 1
864    FTYPE_TABLE = 2
865    FTYPE_RAGGED = 3
866
867    # TODO: Change convention cbox --> slist
868    parser_for_dtype = dict(
869        i=int,
870        f=float,
871        cbox=str,
872    )
873
874    def __init__(self, tag, data, oncv_dims):
875        """
876        Args:
877
878            tag:
879                Tag used to identify the field.
880            data:
881                Accepts Ordered dict or List of ordered dicts with varname=value.
882
883            oncv_dims:
884                Dictionary with all the dimensions of the calculation.
885        """
886        #print("tag", tag, "data", data, "oncv_dims", oncv_dims)
887        self.oncv_dims = AttrDict(**oncv_dims)
888        self.tag = tag
889        self.data = data
890
891    @classmethod
892    def make_empty_data(cls, oncv_dims=None):
893        """Initialize data and fill it with None."""
894        if cls.ftype == cls.FTYPE_ROW:
895            data = OrderedDict([(k, v.get("value")) for k, v in cls.WXCTRL_PARAMS.items()])
896
897        elif cls.ftype == cls.FTYPE_TABLE:
898            nrows = cls.nrows_from_dims(oncv_dims)
899            data = nrows * [None]
900            for i in range(nrows):
901                data[i] = OrderedDict([(k, v.get("value")) for k, v in cls.WXCTRL_PARAMS.items()])
902
903        else:
904            raise NotImplementedError(str(cls.ftype))
905
906        return data
907
908    #@property
909    #def filepos(self):
910    #    for pos, field in enumerate(_FIELD_LIST):
911    #        if isinstance(self, field):
912    #            return pos
913    #    else:
914    #        raise ValueError("Cannot find position of class:" % self.__class__.__name__)
915
916    @property
917    def nrows(self):
918        """
919        Number of rows i.e. the number of lines in the section
920        as specified in the input file.
921        """
922        if self.ftype == self.FTYPE_ROW:
923            return 1
924        elif self.ftype == self.FTYPE_TABLE:
925            return len(self.data)
926        else:
927            raise NotImplementedError()
928
929    #@property
930    #def cols_per_row(self):
931    #    if self.type == self.FTYPE_RAGGED:
932    #    raise NotImplementedError()
933
934    def __str__(self):
935        """Returns a string with the input variables."""
936        lines = []
937        app = lines.append
938
939        # Some fields have a single row but we need a list in the loop below.
940        entries = self.data
941        if self.ftype == self.FTYPE_ROW:
942            entries = [self.data]
943
944        for i, entry in enumerate(entries):
945            if i == 0:
946                # Put comment with the name of the variables only once.
947                app("# " + " ".join((str(k) for k in entry.keys())))
948            app(" ".join(str(v) for v in entry.values()))
949
950        return "\n".join(lines)
951
952    def has_var(self, key):
953        """Return True if variable belongs to self."""
954        if self.ftype == self.FTYPE_ROW:
955            return key in self.data
956        else:
957            return key in self.data[0]
958
959    def set_var(self, key, value):
960        """Set the value of a variable."""
961        assert self.has_var(key)
962        if self.ftype == self.FTYPE_ROW:
963            self.data[key] = value
964
965        elif self.ftype == self.FTYPE_TABLE:
966            for r in range(self.nrows):
967                self.data[r][key] = value
968
969        else:
970            raise NotImplementedError()
971
972    def set_vars(self, ord_vars):
973        """
974        Set the value of the variables inside a field
975
976        Args:
977            ord_vars:
978                OrderedDict or list of OrderedDict (depending on field_idx).
979        """
980        assert len(ord_vars) == len(self.data)
981
982        if self.ftype == self.FTYPE_ROW:
983            #assert isinstance(ord_vars, OrderedDict)
984            for k, v in ord_vars.items():
985                self.data[k] = v
986
987        elif self.ftype == self.FTYPE_TABLE:
988            # List of ordered dicts.
989            for i, od in enumerate(ord_vars):
990                for k, v in od.items():
991                    self.data[i][k] = v
992
993        else:
994            raise NotImplementedError()
995
996    def set_vars_from_lines(self, lines):
997        """The the value of the variables from a list of strings."""
998        #print("About to read: ", type(self), "\nlines=\n", "\n".join(lines))
999        okeys = self.WXCTRL_PARAMS.keys()
1000        odtypes = [v["dtype"] for v in self.WXCTRL_PARAMS.values()]
1001        parsers = [self.parser_for_dtype[ot] for ot in odtypes]
1002        #print("okeys", okeys, "odtypes", odtypes)
1003
1004        if self.ftype == self.FTYPE_ROW:
1005            assert len(lines) == 1
1006            tokens = lines[0].split()
1007            #print("row tokens", tokens)
1008            #if self.__class__ == VlocalField: tokens[-1] = int(tokens[-1])
1009
1010            for key, p, tok in zip(okeys, parsers, tokens):
1011                #print(key)
1012                try:
1013                    self.data[key] = p(tok)
1014                except Exception:
1015                    print("Exception while trying to convert: key= %s, tok= %s" % (key, tok))
1016                    raise
1017
1018        elif self.ftype == self.FTYPE_TABLE:
1019            assert len(lines) == self.nrows
1020            for i in range(self.nrows):
1021                tokens = lines[i].split()
1022                #print("table tokens: ", tokens)
1023                for key, p, tok in zip(okeys, parsers, tokens):
1024                    self.data[i][key] = p(tok)
1025
1026        else:
1027            raise NotImplementedError()
1028
1029    def get_vars(self):
1030        return self.data
1031
1032    @classmethod
1033    def from_wxctrl(cls, wxctrl, tag, oncv_dims):
1034        """Build the object from a wxpython widget."""
1035        # Get the variables from the controller
1036        ord_vars = wxctrl.GetParams()
1037        # Build empty field and set its variables.
1038        new = empty_field(tag, oncv_dims)
1039        new.set_vars(ord_vars)
1040
1041        return new
1042
1043    def make_wxctrl(self, parent, **kwargs):
1044        """"Build the wx controller associated to this field."""
1045        if self.ftype == self.FTYPE_ROW:
1046            return awx.RowMultiCtrl(parent, self._customize_wxctrl(**kwargs))
1047
1048        elif self.ftype == self.FTYPE_TABLE:
1049            return awx.TableMultiCtrl(parent, self.nrows, self._customize_wxctrl(**kwargs))
1050
1051        else:
1052            # Ragged case e.g. test configurations:
1053            #dims
1054            raise NotImplementedError()
1055
1056    def _customize_wxctrl(self, **kwargs):
1057        """
1058        Start with the default parameters for the wx controller
1059        and override them with those given in kwargs
1060        """
1061        # Make a deep copy since WXTRL_PARAMS is mutable.
1062        ctrl_params = copy.deepcopy(self.WXCTRL_PARAMS)
1063
1064        for label, params in ctrl_params.items():
1065            value = kwargs.pop("label", None)
1066            if value is not None:
1067                params["value"] = str(value)
1068
1069        return ctrl_params
1070
1071
1072class RowField(Field):
1073    """A field made of a single row."""
1074    ftype = Field.FTYPE_ROW
1075
1076
1077class TableField(Field):
1078    """A field made of multiple rows, all with the same number of columns."""
1079    ftype = Field.FTYPE_TABLE
1080
1081    @classmethod
1082    def nrows_from_dims(cls, oncv_dims):
1083        """Return the number of rows from a dictionary with the dimensions."""
1084        raise NotImplementedError("Subclasses should define nrows_from_dims")
1085
1086    def get_col(self, colname):
1087        """Returns an array with the values of column colname."""
1088        col = [None] * len(self.data)
1089        for i, row in enumerate(self.data):
1090            col[i] = row[colname]
1091
1092        return col
1093
1094
1095class RaggedField(Field):
1096    """
1097    A field made of ragged rows, i.e. multiple rows with different number of columns.
1098    """
1099    ftype = Field.FTYPE_RAGGED
1100
1101    @classmethod
1102    def nrows_from_dims(cls, oncv_dims):
1103        """Return the number of rows from a dictionary with the dimensions."""
1104        raise NotImplementedError("Subclasses should define nrows_from_dims")
1105
1106    @classmethod
1107    def ncols_of_rows(cls, oncv_dims):
1108        """Return the number of columns in each row from a dictionary with the dimensions."""
1109        raise NotImplementedError("Subclasses should define nrows_from_dims")
1110
1111
1112def add_tooltips(cls):
1113    """Class decorator that add tooltips to WXCTRL_PARAMS."""
1114    d = cls.WXCTRL_PARAMS
1115    for key, params in d.items():
1116        params["tooltip"] = oncv_tip(key)
1117
1118    return cls
1119
1120
1121@add_tooltips
1122class AtomConfField(RowField):
1123    name = "ATOMIC CONFIGURATION"
1124
1125    WXCTRL_PARAMS = OrderedDict([
1126        ("atsym", dict(dtype="cbox", choices=all_symbols())),
1127        ("z", dict(dtype="i")),
1128        ("nc", dict(dtype="i", value=0, tooltip="number of core states"),),
1129        ("nv", dict(dtype="i", value=0, tooltip="number of valence states")),
1130        #("iexc", dict(dtype="i", value=4, tooltip="xc functional")),
1131        # GGA-PBE
1132        #("iexc", dict(dtype="f", value=4, tooltip="xc functional")),
1133        #("iexc", dict(dtype="cbox", value="4", choices=["-001013", "4"], tooltip="xc functional")),
1134        # LDA
1135        ("iexc", dict(dtype="cbox", value="4", choices=["-001012", "4"], tooltip="xc functional")),
1136        # PBEsol
1137        #("iexc", dict(dtype="cbox", value="4", choices=["-116133", "4"], tooltip="xc functional")),
1138        ("psfile", dict(dtype="cbox", value="psp8", choices=["psp8", "upf", "both"]))])
1139
1140
1141@add_tooltips
1142class RefConfField(TableField):
1143    name = "REFERENCE CONFIGURATION"
1144
1145    WXCTRL_PARAMS = OrderedDict([
1146        ("n", dict(dtype="i")),
1147        ("l", dict(dtype="i")),
1148        ("f", dict(dtype="f"))])
1149
1150    #@classmethod
1151    #def neutral_from_symbol(cls, symbol):
1152    #    # TODO
1153    #    element = periodic_table.Element(symbol)
1154    #    # E.g., The electronic structure for Fe is represented as:
1155    #    # [(1, "s", 2), (2, "s", 2), (2, "p", 6), (3, "s", 2), (3, "p", 6), (3, "d", 6), (4, "s", 2)]
1156    #    #new = empty_field(cls.tag, oncv_dims=dict(nc))
1157    #    for row, (n, lchar, f) in zip(new, element.full_electronic_structure):
1158    #        row["n"], row["l"], row["f"] = n, periodic_table.char2l(lchar), f
1159    #    return new
1160
1161    @classmethod
1162    def nrows_from_dims(cls, oncv_dims):
1163        return oncv_dims["nv"] + oncv_dims["nc"]
1164
1165    @property
1166    def nrows(self):
1167        return self.oncv_dims["nv"] + self.oncv_dims["nc"]
1168
1169
1170@add_tooltips
1171class PseudoConfField(TableField):
1172    name = "PSEUDOPOTENTIAL AND OPTIMIZATION"
1173
1174    WXCTRL_PARAMS = OrderedDict([
1175        ("l", dict(dtype="i", value=0)),
1176        ("rc", dict(dtype="f", value=3.0)),
1177        ("ep", dict(dtype="f", value=0.0)),
1178        ("ncon", dict(dtype="i", value=4)),
1179        ("nbas", dict(dtype="i", value=7)),
1180        ("qcut", dict(dtype="f", value=6.0))])
1181
1182    @classmethod
1183    def nrows_from_dims(cls, oncv_dims):
1184        return oncv_dims["lmax"] + 1
1185
1186    @property
1187    def nrows(self):
1188        return self.oncv_dims["lmax"] + 1
1189
1190
1191@add_tooltips
1192class LmaxField(RowField):
1193    name = "LMAX"
1194
1195    WXCTRL_PARAMS = OrderedDict([
1196        ("lmax", dict(dtype="i", value=2))])
1197
1198
1199@add_tooltips
1200class VlocalField(RowField):
1201    name = "LOCAL POTENTIAL"
1202
1203    WXCTRL_PARAMS = OrderedDict([
1204        ("lloc", dict(dtype="i", value=4)),
1205        ("lpopt", dict(dtype="i", value=5)),
1206        ("rc5", dict(dtype="f", value=3.0)),
1207        ("dvloc0", dict(dtype="f", value=0))])
1208
1209
1210@add_tooltips
1211class VkbConfsField(TableField):
1212    name = "VANDERBILT-KLEINMAN-BYLANDER PROJECTORs"
1213
1214    WXCTRL_PARAMS = OrderedDict([
1215        ("l", dict(dtype="i", value=0)),
1216        ("nproj", dict(dtype="i", value=2)),
1217        ("debl", dict(dtype="f", value=1.0))])
1218
1219    @classmethod
1220    def nrows_from_dims(cls, oncv_dims):
1221        return oncv_dims["lmax"] + 1
1222
1223    @property
1224    def nrows(self):
1225        return self.oncv_dims["lmax"] + 1
1226
1227
1228@add_tooltips
1229class ModelCoreField(RowField):
1230    name = "MODEL CORE CHARGE"
1231
1232    WXCTRL_PARAMS = OrderedDict([
1233        ("icmod", dict(dtype="i", value=0)),
1234        ("fcfact", dict(dtype="f", value=0.25)),
1235        ("rcfact", dict(dtype="f", value=0.0)),
1236        ])
1237
1238
1239@add_tooltips
1240class LogDerField(RowField):
1241    name = "LOG DERIVATIVE ANALYSIS"
1242
1243    WXCTRL_PARAMS = OrderedDict([
1244        ("epsh1", dict(dtype="f", value=-12.0)),
1245        ("epsh2", dict(dtype="f", value=+12.0)),
1246        ("depsh", dict(dtype="f", value=0.02))])
1247
1248
1249@add_tooltips
1250class RadGridField(RowField):
1251    name = "OUTPUT GRID"
1252
1253    WXCTRL_PARAMS = OrderedDict([
1254        ("rlmax", dict(dtype="f", value=6.0, step=1.0)),
1255        ("drl", dict(dtype="f", value=0.01))])
1256
1257
1258#@add_tooltips
1259#class TestConfigsField(RaggedField):
1260    #    name = "TEST CONFIGURATIONS"
1261    #    WXCTRL_PARAMS = OrderedDict([
1262    #        ("ncnf", dict(dtype="i", value="0")),
1263    #        ("nvcnf", dict(dtype="i", value="0")),
1264    #        ("n", dict(dtype="i")),
1265    #        ("l", dict(dtype="i")),
1266    #        ("f", dict(dtype="f"))])
1267
1268    #@classmethod
1269    #def from_oxidation_states(cls, symbol, only_common=True):
1270    #    """
1271    #    Initialize the test configurations with the most common oxidation states.
1272
1273    #    Args:
1274    #        symbol:
1275    #            Chemical symbol.:w
1276    #        only_common:
1277    #            If False all the known oxidations states are considered, else only
1278    #            the most common ones.
1279    #    """
1280    #    element = periodic_table.Element(symbol)
1281
1282    #    if only_common:
1283    #        oxi_states = element.common_oxidation_states
1284    #    else:
1285    #        oxi_states = element.oxidation_states
1286
1287    #    for oxi in oxi_states:
1288    #        # Get the electronic configuration of atom with Z = Z + oxi
1289    #        if oxi == 0:
1290    #            continue
1291    #        oxiele = periodic_table.Element.from_Z(element.Z + oxi)
1292
1293    #        # Here we found the valence configuration by comparing
1294    #        # the full configuration of oxiele and the one of the initial element.
1295
1296
1297    #    return new
1298
1299    #@property
1300    #def nrows(self):
1301    #    return self.oncv_dims["ncnf"]
1302
1303    #@classmethod
1304    #def nrows_from_dims(cls, oncv_dims):
1305    #    return oncv_dims["ncnf"]
1306
1307    #@property
1308    #def nlines_for_row(self, row):
1309
1310
1311# List with the field in the same order as the one used in the input file.
1312_FIELD_LIST = [
1313    AtomConfField,
1314    RefConfField,
1315    LmaxField,
1316    PseudoConfField,
1317    VlocalField,
1318    VkbConfsField,
1319    ModelCoreField,
1320    LogDerField,
1321    RadGridField,
1322    #TestConfigsField,
1323]
1324
1325_NFIELDS = len(_FIELD_LIST)
1326
1327
1328class OncvInput(object):
1329    """
1330    This object stores the variables needed for generating a pseudo with oncvsps.
1331    One can initialize this object either from a prexisting file
1332    or programmatically from the input provided by the user in a GUI.
1333
1334    An input consistst of _NFIELDS fields. Each field is either a OrderedDict
1335    or a list of ordered dicts with the input variables.
1336    """
1337    @classmethod
1338    def from_file(cls, filepath):
1339        """Initialize the object from an external input file."""
1340        # Read input lines: ignore empty lines or line starting with #
1341        lines = []
1342        with open(filepath) as fh:
1343            for line in fh:
1344                line = line.strip()
1345                if line and not line.startswith("#"):
1346                    lines.append(line)
1347
1348        # Read dimensions
1349        # nc and nv from the first line.
1350        tokens = lines[0].split()
1351        atsym, z, nc, nv = tokens[0:4]
1352        z, nc, nv = map(int, (z, nc, nv))
1353
1354        # Read lmax and ncfn
1355        lmax = int(lines[nc + nv + 1])
1356        #print("lmax = ",lmax)
1357
1358        # TODO
1359        # number of tests and number of rows for the different configurations.
1360        ncnf = 0
1361
1362        # Initialize the object
1363        new = OncvInput(oncv_dims=dict(atsym=atsym, nc=nc, nv=nv, lmax=lmax, ncnf=ncnf))
1364
1365        # TODO
1366        # Fill it
1367        start = 0
1368        for field in new:
1369            stop = start + field.nrows
1370            #print(type(field))
1371            field.set_vars_from_lines(lines[start:stop])
1372            start = stop
1373
1374        return new
1375
1376    @classmethod
1377    def from_symbol(cls, symbol):
1378        """
1379        Return a tentative input file for generating a pseudo for the given chemical symbol
1380
1381        .. note:
1382              Assume default values that might not be optimal.
1383        """
1384        nc, nv, lmax = 0, 0, 0
1385        #atom = nist.get_neutral_entry(symbol=symbol)
1386        #for state in atom.states:
1387        #    lmax = max(lmax, state.l)
1388
1389        # E.g., The electronic structure for Fe is represented as:
1390        # [(1, "s", 2), (2, "s", 2), (2, "p", 6), (3, "s", 2), (3, "p", 6), (3, "d", 6), (4, "s", 2)]
1391        element = Element[symbol]
1392        for (n, lchar, f) in element.full_electronic_structure:
1393            nc += 1
1394            lmax = max(lmax, char2l(lchar))
1395
1396        # FIXME
1397        lmax = 1
1398        lmax = 2
1399        #lmax = 3
1400
1401        ncnf = 0
1402        oncv_dims = dict(atsym=symbol, nc=nc, nv=nv, lmax=lmax, ncnf=ncnf)
1403
1404        new = cls(oncv_dims)
1405
1406        field = new.fields[_FIELD_LIST.index(RefConfField)]
1407        for row, (n, lchar, f) in zip(field.data, element.full_electronic_structure):
1408            row["n"], row["l"], row["f"] = n, char2l(lchar), f
1409
1410        return new
1411
1412    def __init__(self, oncv_dims, fields=None):
1413        """
1414        Initialize the object from a dict with the fundamental dimensions.
1415        If fields is None, we create an empty dict, else we use fields.
1416        """
1417        self.dims = AttrDict(**oncv_dims)
1418        #print("oncv_dims", self.dims)
1419
1420        if fields is None:
1421            # Default fields.
1422            self.fields = _NFIELDS * [None]
1423            for i in range(_NFIELDS):
1424                self.fields[i] = empty_field(i, self.dims)
1425
1426            header = self.fields[_FIELD_LIST.index(AtomConfField)]
1427            header.set_var("atsym", self.dims.atsym)
1428            header.set_var("z", Element[self.dims.atsym].Z)
1429            header.set_var("nc", self.dims.nc)
1430            header.set_var("nv", self.dims.nv)
1431
1432        else:
1433            self.fields = fields
1434
1435        # 1) ATOM CONFIGURATION
1436        # atsym, z, nc, nv, iexc   psfile
1437        # O    8     1   2   3   psp8
1438        #
1439        # 2) REFERENCE CONFIGURATION
1440        # n, l, f  (nc+nv lines)
1441        # 1    0    2.0
1442        # 2    0    2.0
1443        # 2    1    4.0
1444        #
1445        # 3) LMAX FIELD
1446        # lmax
1447        # 1
1448        # 4) PSEUDOPOTENTIAL AND OPTIMIZATION
1449        # l, rc, ep, ncon, nbas, qcut  (lmax+1 lines, l's must be in order)
1450        # 0    1.60    0.00    4    7    8.00
1451        # 1    1.60    0.00    4    7    8.00
1452        #
1453        # 5) LOCAL POTENTIAL
1454        # lloc, lpopt, rc(5), dvloc0
1455        # 4    5    1.4    0.0
1456        #
1457        # 6) VANDERBILT-KLEINMAN-BYLANDER PROJECTORs
1458        # l, nproj, debl  (lmax+1 lines, l's in order)
1459        # 0    2    1.50
1460        # 1    2    1.00
1461        #
1462        # 7) MODEL CORE CHARGE
1463        # icmod, fcfact
1464        # 0    0.0
1465        #
1466        # 8) LOG DERIVATIVE ANALYSIS
1467        # epsh1, epsh2, depsh
1468        # -2.0  2.0  0.02
1469        #
1470        # 9) OUTPUT GRID
1471        # rlmax, drl
1472        # 4.0  0.01
1473
1474        # TODO
1475        # 10) TEST CONFIGURATIONS
1476        # ncnf
1477        # 2
1478        # nvcnf    (repeated ncnf times)
1479        # n, l, f  (nvcnf lines, repeated follwing nvcnf's ncnf times)
1480        # 2
1481        # 2    0    2.0
1482        # 2    1    3.0
1483        #
1484        # 2
1485        # 2    0    1.0
1486        # 2    1    4.0
1487        #ncnf = 2
1488        #nvcnf = 2
1489
1490    @property
1491    def lmax(self):
1492        return self.dims.lmax
1493
1494    def __iter__(self):
1495        return self.fields.__iter__()
1496
1497    def __str__(self):
1498        """Returns a string with the input variables."""
1499        lines = []
1500        app = lines.append
1501
1502        for i, field in enumerate(self):
1503            # FIXME This breaks the output parser!!!!!
1504            #pre = "\n#" if i > 0 else ""
1505            #app(pre + field.name)
1506            app(str(field))
1507
1508        s = "\n".join(lines)
1509        # FIXME needed to bypass problems with tests
1510        return s + "\n 0\n"
1511
1512    def __setitem__(self, key, value):
1513        ncount = 0
1514        for f in self.fields:
1515            if f.has_var(key):
1516                ncount += 1
1517                f.set_var(key, value)
1518
1519        assert ncount == 1
1520
1521    def deepcopy(self):
1522        """Deep copy of the input."""
1523        return copy.deepcopy(self)
1524
1525    def optimize_vloc(self):
1526        """Produce a list of new input files by changing the lloc option for vloc."""
1527        # Test all possible vloc up to lmax
1528        inps, new = [], self.deepcopy()
1529        for il in range(self.lmax+1):
1530            new["lloc"] = il
1531            inps.append(new.deepcopy())
1532
1533        # Add option for smooth polynomial
1534        new["lloc"] = 4
1535        inps.append(new)
1536
1537        return inps
1538
1539    def optimize_modelcore(self, fcfact_list, add_icmod0=True):
1540        """Produce a list of new input files by changing the icmod option for model core."""
1541        inps, new = [], self.deepcopy()
1542
1543        if add_icmod0:
1544            new["icmod"] = 0
1545            inps.append(new.deepcopy())
1546
1547        for fcfact in fcfact_list:
1548            new["icmod"] = 1
1549            new["fcfact"] = fcfact
1550            inps.append(new.deepcopy())
1551
1552        return inps
1553
1554    @property
1555    def qcut_l(self):
1556        """List with the values of qcuts as function of l."""
1557        i = _FIELD_LIST.index(PseudoConfField)
1558        return self.fields[i].get_col("qcut")
1559
1560    @property
1561    def debl_l(self):
1562        """List with the values of debl as function of l."""
1563        i = _FIELD_LIST.index(VkbConfsField)
1564        return self.fields[i].get_col("debl")
1565
1566    @property
1567    def rc_l(self):
1568        """List with the values of rc as function of l."""
1569        i = _FIELD_LIST.index(PseudoConfField)
1570        return self.fields[i].get_col("rc")
1571
1572    def optimize_qcuts_for_l(self, l, new_qcuts):
1573        """
1574        Returns a list of new input objects in which the qcut parameter for
1575        the given l has been replaced by the values listed in new_qcuts.
1576
1577        Args:
1578            l:
1579                Angular momentum
1580            new_qcuts:
1581                Iterable with the new values of qcut.
1582                The returned list will have len(new_qcuts) input objects.
1583        """
1584        # Find the field with the configuration parameters.
1585        i = _FIELD_LIST.index(PseudoConfField)
1586
1587        # Find the row with the given l.
1588        for irow, row in enumerate(self.fields[i].data):
1589            if row["l"] == l:
1590                break
1591        else:
1592            raise ValueError("Cannot find l %s in the PseudoConfField" % l)
1593
1594        # This is the dict we want to change
1595        inps = []
1596        for qc in new_qcuts:
1597            new_inp = self.deepcopy()
1598            new_inp.fields[i].data[irow]["qcut"] = qc
1599            inps.append(new_inp)
1600
1601        return inps
1602
1603    def optimize_rcs_for_l(self, l, new_rcs):
1604        """
1605        Returns a list of new input objects in which the rc parameter for
1606        the given l has been replaced by the values listed in new_rcs.
1607
1608        Args:
1609            l:
1610                Angular momentum
1611            new_rcs:
1612                Iterable with the new values of rcs.
1613                The returned list will have len(new_rcs) input objects.
1614        """
1615        # Find the field with the configuration parameters.
1616        i = _FIELD_LIST.index(PseudoConfField)
1617
1618        # Find the row with the given l.
1619        for irow, row in enumerate(self.fields[i].data):
1620            if row["l"] == l:
1621                break
1622        else:
1623            raise ValueError("Cannot find l %s in the PseudoConfField" % l)
1624
1625        # This is the dict we want to change
1626        inps = []
1627        for rc in new_rcs:
1628            new_inp = self.deepcopy()
1629            new_inp.fields[i].data[irow]["rc"] = rc
1630            inps.append(new_inp)
1631
1632        return inps
1633
1634    def optimize_debls_for_l(self, l, new_debls):
1635        """
1636        Returns a list of new input objects in which the debls parameter for
1637        the given l has been replaced by the values listed in new_debls.
1638
1639        Args:
1640            l:
1641                Angular momentum
1642            new_debls:
1643                Iterable with the new values of debls.
1644                The returned list will have len(new_debls) input objects.
1645        """
1646        # Find the field with the configuration parameters.
1647        i = _FIELD_LIST.index(VkbConfsField)
1648
1649        # Find the row with the given l.
1650        for irow, row in enumerate(self.fields[i].data):
1651            if row["l"] == l:
1652                break
1653        else:
1654            raise ValueError("Cannot find l %s in the VkbConfsField" % l)
1655
1656        # This is the dict we want to change
1657        inps = []
1658        for debl in new_debls:
1659            new_inp = self.deepcopy()
1660            new_inp.fields[i].data[irow]["debl"] = debl
1661            inps.append(new_inp)
1662
1663        return inps
1664
1665
1666class OncvNotebook(wx.Notebook):
1667
1668    @classmethod
1669    def from_file(cls, parent, filename):
1670        inp = OncvInput.from_file(filename)
1671        new = cls(parent, inp.dims)
1672
1673        for field in inp:
1674            wxctrl = new.wxctrls[field.__class__]
1675            wxctrl.SetParams(field.data)
1676
1677        return new
1678
1679    @classmethod
1680    def from_symbol(cls, parent, symbol):
1681        inp = OncvInput.from_symbol(symbol)
1682        new = cls(parent, inp.dims)
1683
1684        for field in inp:
1685            wxctrl = new.wxctrls[field.__class__]
1686            wxctrl.SetParams(field.data)
1687
1688        return new
1689
1690    #def fromInput(cls)
1691
1692    def __init__(self, parent, oncv_dims):
1693        super(OncvNotebook, self).__init__(parent)
1694
1695        # Build tabs
1696        self.oncv_dims = oncv_dims
1697        self.ae_tab = AeConfTab(self, oncv_dims)
1698        self.ps_tab = PsConfTab(self, oncv_dims)
1699        self.pstests_tab = PsTestsTab(self, oncv_dims)
1700
1701        # Add tabs
1702        self.AddPage(self.ae_tab, "AE config")
1703        self.AddPage(self.ps_tab, "PP config")
1704        self.AddPage(self.pstests_tab, "Tests")
1705
1706    @property
1707    def tabs(self):
1708        return (self.ae_tab, self.ps_tab, self.pstests_tab)
1709
1710    @property
1711    def wxctrls(self):
1712        d = {}
1713        for tab in self.tabs:
1714            d.update(tab.wxctrls)
1715        return d
1716
1717    @property
1718    def calc_type(self):
1719        return self.ae_tab.calctype_cbox.GetValue()
1720
1721    @property
1722    def lmax(self):
1723        # TODO: property or method?
1724        return self.ps_tab.wxctrls[LmaxField].GetParams()["lmax"]
1725
1726    def getElement(self):
1727        symbol = self.ae_tab.wxctrls[AtomConfField].GetParams()["atsym"]
1728        return Element[symbol]
1729
1730    def makeInput(self):
1731        """Build an OncvInput instance from the values specified in the controllers."""
1732        inp = OncvInput(self.oncv_dims)
1733
1734        for cls, wxctrl in self.wxctrls.items():
1735            i = _FIELD_LIST.index(cls)
1736            inp.fields[i].set_vars(wxctrl.GetParams())
1737
1738        return inp
1739
1740    def makeInputString(self):
1741        """Return a string with the input passed to the pp generator."""
1742        return str(self.makeInput())
1743
1744
1745class AeConfTab(awx.Panel):
1746    def __init__(self, parent, oncv_dims):
1747        super(AeConfTab, self).__init__(parent)
1748
1749        # Set the dimensions and build the widgets.
1750        self.oncv_dims = oncv_dims
1751        self.buildUI()
1752
1753    def buildUI(self):
1754        self.main_sizer = main_sizer = wx.BoxSizer(wx.VERTICAL)
1755
1756        stext = wx.StaticText(self, -1, "Calculation type:")
1757        choices = ["scalar-relativistic", "fully-relativistic", "non-relativistic"]
1758        self.calctype_cbox = wx.ComboBox(
1759            self, id=-1, name='Calculation type', choices=choices, value=choices[0], style=wx.CB_READONLY)
1760
1761        add_opts = dict(proportion=0, flag=wx.ALIGN_CENTER_VERTICAL | wx.ALL, border=5)
1762        sizer_addopts = dict(proportion=0, flag=wx.ALL, border=5)
1763
1764        hbox0 = wx.BoxSizer(wx.HORIZONTAL)
1765        hbox0.Add(stext, **add_opts)
1766        hbox0.Add(self.calctype_cbox)
1767
1768        main_sizer.Add(hbox0, **add_opts)
1769
1770        #(list_of_classes_on_row, show)
1771        layout = [
1772            AtomConfField,
1773            RefConfField,
1774            RadGridField,
1775            ]
1776
1777        # Each field has a widget that returns the variables in a dictionary
1778        self.wxctrls = {}
1779        self.sbox_sizers = {}
1780        #self.sboxes = {}
1781
1782        for cls in layout:
1783            i = _FIELD_LIST.index(cls)
1784            f = empty_field(i, self.oncv_dims)
1785            wxctrl = f.make_wxctrl(self)
1786            sbox = wx.StaticBox(self, -1, f.name + ":")
1787            sbox_sizer = wx.StaticBoxSizer(sbox, wx.VERTICAL)
1788            #sbox_sizer = wx.BoxSizer(wx.VERTICAL)
1789            sbox_sizer.Add(wxctrl, **sizer_addopts)
1790            #sz = wx.FlexGridSizer(wx.VERTICAL)
1791            #sz.Add(sbox_sizer)
1792
1793            self.wxctrls[cls] = wxctrl
1794            self.sbox_sizers[cls] = sbox_sizer
1795            #self.sboxes[cls] = sz
1796
1797            main_sizer.Add(sbox_sizer, **add_opts)
1798
1799        lda_levels_button = wx.Button(self, -1, "LDA levels (NIST)")
1800        lda_levels_button.Bind(wx.EVT_BUTTON, self.onShowLdaLevels)
1801        main_sizer.Add(lda_levels_button, **add_opts)
1802
1803        add_button = wx.Button(self, -1, "Add level")
1804        add_button.action = "add"
1805        add_button.Bind(wx.EVT_BUTTON, self.onAddRemoveLevel)
1806        main_sizer.Add(add_button, **add_opts)
1807
1808        remove_button = wx.Button(self, -1, "Remove level")
1809        remove_button.action = "remove"
1810        remove_button.Bind(wx.EVT_BUTTON, self.onAddRemoveLevel)
1811        main_sizer.Add(remove_button, **add_opts)
1812
1813        self.SetSizerAndFit(main_sizer)
1814
1815    @property
1816    def atomconf(self):
1817        return self.wxctrls[AtomConfField].GetParams()
1818
1819    @property
1820    def nc(self):
1821        return self.atomconf["nc"]
1822
1823    @property
1824    def nv(self):
1825        return self.atomconf["nv"]
1826
1827    @property
1828    def symbol(self):
1829        return self.atomconf["atsym"]
1830
1831    def get_core(self):
1832        # E.g., The electronic structure for Fe is represented as:
1833        # [(1, "s", 2), (2, "s", 2), (2, "p", 6), (3, "s", 2), (3, "p", 6), (3, "d", 6), (4, "s", 2)]
1834        core = []
1835        refconf = self.wxctrls[RefConfField]
1836        for ic in range(self.nc):
1837            d = refconf[ic].GetParams()
1838            t = [d[k] for k in ("n", "l", "f")]
1839            core.append(tuple(t))
1840
1841        return core
1842
1843    def get_valence(self):
1844        valence = []
1845        refconf = self.wxctrls[RefConfField]
1846        for iv in range(self.nc, self.nc + self.nv):
1847            d = refconf[iv].GetParams()
1848            t = [d[k] for k in ("n", "l", "f")]
1849            valence.append(tuple(t))
1850
1851        return valence
1852
1853    def onAddRemoveLevel(self, event):
1854        button = event.GetEventObject()
1855        sbox_sizer = self.sbox_sizers[RefConfField]
1856        old = self.wxctrls[RefConfField]
1857        sbox_sizer.Hide(0)
1858        sbox_sizer.Remove(0)
1859
1860        if button.action == "add":
1861            old.appendRow()
1862        elif button.action == "remove":
1863            old.removeRow()
1864
1865        sizer_addopts = dict(proportion=0, flag=wx.ALL, border=5)
1866        sbox_sizer.Insert(0, old, **sizer_addopts)
1867
1868        #self.sboxes[RefConfField].Layout()
1869        sbox_sizer.Show(0)
1870        sbox_sizer.Layout()
1871        #self.main_sizer.Layout()
1872        self.Layout()
1873        self.Fit()
1874
1875        #frame = self.GetParent().GetParent()
1876        #frame.fSizer.Layout()
1877        #frame.Fit()
1878
1879    def onShowLdaLevels(self, event):
1880        # Get the LDA levels of the neutral atom.
1881        # (useful to decide if semicore states should be included in the valence).
1882        entry = nist.get_neutral_entry(self.symbol)
1883        frame = awx.Frame(self, title="LDA levels for neutral %s (NIST database)" % self.symbol)
1884        awx.ListCtrlFromTable(frame, table=entry.to_table())
1885        frame.Show()
1886
1887
1888class PsConfTab(awx.Panel):
1889    def __init__(self, parent, oncv_dims):
1890        super(PsConfTab, self).__init__(parent)
1891
1892        self.notebook = parent
1893
1894        # Set the dimensions and build the widgets.
1895        self.oncv_dims = oncv_dims
1896        self.buildUI()
1897
1898    def buildUI(self):
1899        sizer_addopts = dict(proportion=0, flag=wx.ALL, border=5)
1900        add_opts = dict(proportion=0, flag=wx.ALIGN_CENTER_VERTICAL | wx.ALL, border=5)
1901
1902        self.main_sizer = main_sizer = wx.BoxSizer(wx.VERTICAL)
1903
1904        # list_of_classes
1905        layout = [
1906            LmaxField,
1907            PseudoConfField,
1908            VkbConfsField,
1909            VlocalField,
1910            ModelCoreField,
1911        ]
1912
1913        # FieldClass: [(label, OptimizationFrame), ....]
1914        fields_with_optimization = {
1915            PseudoConfField: [
1916                ("Change rc", RcOptimizationFrame),
1917                ("Change qcut", QcutOptimizationFrame),
1918                ],
1919            VlocalField: [
1920                ("Change lloc", LlocOptimizationFrame),
1921                ("Change rc5", Rc5OptimizationFrame),
1922                ],
1923            ModelCoreField: [("Change fcfact", FcfactOptimizationFrame)],
1924            VkbConfsField: [("Change debl", DeblOptimizationFrame)],
1925        }
1926
1927        # Each field has a widget that returns the variables in a dictionary
1928        self.wxctrls = {}
1929
1930        # Keep an internal list of buttons so that we can disable them easily.
1931        self.all_optimize_buttons = []
1932        self.sboxes = {}
1933
1934        for cls in layout:
1935            i = _FIELD_LIST.index(cls)
1936            f = empty_field(i, self.oncv_dims)
1937            wxctrl = f.make_wxctrl(self)
1938
1939            sbox_sizer = wx.StaticBoxSizer(wx.StaticBox(self, -1, f.name + ":"), wx.VERTICAL)
1940            #sbox_sizer = wx.BoxSizer(wx.VERTICAL)
1941            sbox_sizer.Add(wxctrl, **sizer_addopts)
1942
1943            self.sboxes[cls] = sbox_sizer
1944            self.wxctrls[cls] = wxctrl
1945
1946            # add optimization button if the field supports it.
1947            if f.__class__ in fields_with_optimization:
1948                hsz = wx.BoxSizer(wx.HORIZONTAL)
1949                for label, opt_frame in fields_with_optimization[f.__class__]:
1950                    optimize_button = wx.Button(self, -1, label)
1951                    optimize_button.Bind(wx.EVT_BUTTON, self.onOptimize)
1952                    optimize_button.field_class = f.__class__
1953                    optimize_button.opt_frame = opt_frame
1954
1955                    self.all_optimize_buttons.append(optimize_button)
1956                    hsz.Add(optimize_button, 0, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL, border=5)
1957
1958                sbox_sizer.Add(hsz, 0, flag=wx.ALL | wx.ALIGN_CENTER_HORIZONTAL)
1959
1960            main_sizer.Add(sbox_sizer, **add_opts)
1961
1962        #self.old_lmax = self.lmax_ctrl.GetParams()["lmax"]
1963        #self.wxctrls[LmaxField].Bind(
1964
1965        add_button = wx.Button(self, -1, "Add row")
1966        add_button.Bind(wx.EVT_BUTTON, self.onLmaxChanged)
1967        main_sizer.Add(add_button)
1968
1969        self.SetSizerAndFit(main_sizer)
1970
1971    @property
1972    def lmax_ctrl(self):
1973        return self.wxctrls[LmaxField]
1974
1975    def onLmaxChanged(self, event):
1976        #self.wxctrls[PseudoConfField].appendRow()
1977        #self.wxctrls[VkbConfsField].appendRow()
1978
1979        #main_sizer = self.main_sizer
1980        #main_sizer.Hide(1)
1981        #main_sizer.Remove(1)
1982        #main_sizer.Insert(1, self.wxctrls[PseudoConfField], 0, wx.EXPAND | wx.ALL, 5)
1983
1984        #for sbox in self.sboxes.values():
1985        #    sbox.Layout()
1986        self.main_sizer.Layout()
1987
1988    def onOptimize(self, event):
1989        button = event.GetEventObject()
1990        #self.enable_all_optimize_buttons(False)
1991        opt_frame = button.opt_frame
1992
1993        try:
1994            opt_frame(self, notebook=self.notebook).Show()
1995        finally:
1996            self.enable_all_optimize_buttons(True)
1997
1998    def enable_all_optimize_buttons(self, enable=True):
1999        """Enable/disable the optimization buttons."""
2000        for button in self.all_optimize_buttons:
2001            button.Enable(enable)
2002
2003
2004class PsTestsTab(awx.Panel):
2005    def __init__(self, parent, oncv_dims):
2006        super(PsTestsTab, self).__init__(parent)
2007        self.notebook = parent
2008
2009        self.main_sizer = main_sizer = wx.BoxSizer(wx.VERTICAL)
2010        add_opts = dict(proportion=0, flag=wx.ALIGN_CENTER_VERTICAL | wx.ALL | wx.EXPAND, border=5)
2011
2012        layout = [
2013            LogDerField,
2014        ]
2015
2016        self.wxctrls = {}
2017
2018        for cls in layout:
2019            i = _FIELD_LIST.index(cls)
2020            f = empty_field(i, oncv_dims)
2021            wxctrl = f.make_wxctrl(self)
2022            self.wxctrls[cls] = wxctrl
2023            sbox_sizer = wx.StaticBoxSizer(wx.StaticBox(self, -1, f.name + ":"), wx.VERTICAL)
2024            sbox_sizer.Add(wxctrl, **add_opts)
2025            main_sizer.Add(sbox_sizer, **add_opts)
2026
2027        self.conf_txtctrl = wx.TextCtrl(self, -1, "#TEST CONFIGURATIONS\n  0", style=wx.TE_MULTILINE | wx.TE_LEFT)
2028        main_sizer.Add(self.conf_txtctrl, **add_opts)
2029
2030        common_button = wx.Button(self, -1, label='Common Oxidation States')
2031        common_button.Bind(wx.EVT_BUTTON, self.addOxiStates)
2032        common_button.test_mode = "common_oxis"
2033
2034        all_button = wx.Button(self, -1, label='All Oxidation States')
2035        all_button.Bind(wx.EVT_BUTTON, self.addOxiStates)
2036        all_button.test_mode = "all_oxis"
2037
2038        hsz = wx.BoxSizer(wx.HORIZONTAL)
2039        hsz.Add(common_button, **add_opts)
2040        hsz.Add(all_button, **add_opts)
2041        main_sizer.Add(hsz, **add_opts)
2042
2043        self.SetSizerAndFit(main_sizer)
2044
2045    def conftests_str(self):
2046        return self.conf_txtctrl.GetValue()
2047
2048    def addOxiStates(self, event):
2049        button = event.GetEventObject()
2050        test_mode = button.test_mode
2051
2052        element = self.notebook.getElement()
2053
2054        if test_mode == "common_oxis":
2055            oxi_states = element.common_oxidation_states
2056            title = "(common oxidation states)"
2057        elif test_mode == "all_oxis":
2058            oxi_states = element.oxidation_states
2059            title = "(all oxidation states)"
2060        else:
2061            raise ValueError("Wrong test_mode %s" % test_mode)
2062
2063        self.conf_txtctrl.Clear()
2064        self.conf_txtctrl.WriteText("# TEST CONFIGURATIONS" + title + "\n")
2065
2066        core_aeatom = self.notebook.ae_tab.get_core()
2067        val_aeatom = self.notebook.ae_tab.get_valence()
2068        print("core", core_aeatom)
2069        print("val", val_aeatom)
2070
2071        """
2072        # TEST CONFIGURATIONS
2073        # ncnf
2074            3
2075        #
2076        #   nvcnf (repeated ncnf times)
2077        #   n, l, f  (nvcnf lines, repeated follwing nvcnf's ncnf times)
2078            2
2079            2    0    2.0
2080            2    1    3.0
2081        #
2082            2
2083            2    0    1.0
2084            2    1    4.0
2085        #
2086            2
2087            2    0    1.0
2088            2    1    3.0
2089        """
2090        test_confs = []
2091
2092        for oxi in oxi_states:
2093            #self.conf_txtctrl.WriteText(str(oxi) + "\n")
2094
2095            # Get the electronic configuration of atom with Z = Z + oxi
2096            if oxi == 0: continue
2097
2098            # Here we found the valence configuration by comparing
2099            # the full configuration of oxiele and the one of the initial element.
2100            oxi_element = Element.from_Z(element.Z - oxi)
2101            oxi_estruct = oxi_element.full_electronic_structure
2102            oxi_estruct = [(t[0], char2l(t[1]), t[2]) for t in oxi_estruct]
2103
2104            if oxi < 0:
2105                test_conf = [t for t in oxi_estruct if t not in core_aeatom]
2106            else:
2107                test_conf = [t for t in oxi_estruct if t not in core_aeatom]
2108
2109            self.conf_txtctrl.WriteText(str(oxi) + "\n")
2110            self.conf_txtctrl.WriteText(str(oxi_element) + "\n")
2111            self.conf_txtctrl.WriteText(str(test_conf) + "\n")
2112
2113        self.conf_txtctrl.WriteText("# ncnf\n" + str(len(test_confs)) + "\n")
2114        self.conf_txtctrl.WriteText("""\
2115#
2116#   nvcnf (repeated ncnf times)
2117#   n, l, f  (nvcnf lines, repeated follwing nvcnf's ncnf times)\n""")
2118
2119        #for test in test_confs:
2120        #    self.conf_txtctrl.WriteText(len(test) + "\n")
2121        #    for row in test:
2122        #        self.conf_txtctrl.WriteText(str(row) + "\n")
2123
2124        self.main_sizer.Layout()
2125
2126
2127# Event posted when we start an optimization.
2128#EVT_OPTIMIZATION_TYPE = wx.NewEventType()
2129#EVT_OPTIMIZATION = wx.PyEventBinder(EVT_CONSOLE_TYPE, 1)
2130#
2131#class OptimizationEvent(wx.PyEvent):
2132#    """
2133#    This event is triggered when we start/end the optimization process
2134#    """
2135#    def __init__(self, kind, msg):
2136#        wx.PyEvent.__init__(self)
2137#        self.SetEventType(EVT_OPTIMIZATION_TYPE)
2138#        self.kind, self.msg = kind, msg
2139#
2140#    #@classmethod
2141#    #def start_optimization(cls, msg)
2142#    #@classmethod
2143#    #def end_optimization(cls, msg)
2144#    #wx.PostEvent(self.console, event)
2145
2146
2147class PseudoGeneratorListCtrl(wx.ListCtrl, listmix.ColumnSorterMixin, listmix.ListCtrlAutoWidthMixin):
2148    """
2149    ListCtrl that allows the user to interact with a list of pseudogenerators. Supports column sorting
2150    """
2151    # List of columns
2152    _COLUMNS = ["#", 'status', "max_ecut", "atan_logder_err", "max_psexc_abserr", "herm_err", "warnings"]
2153
2154    def __init__(self, parent, psgens=(), **kwargs):
2155        """
2156        Args:
2157            parent:
2158                Parent window.
2159            psgens:
2160                List of `PseudoGenerator` instances.
2161        """
2162        super(PseudoGeneratorListCtrl, self).__init__(
2163            parent, id=-1, style=wx.LC_REPORT | wx.BORDER_SUNKEN, **kwargs)
2164
2165        self.psgens = psgens if psgens else []
2166
2167        for index, col in enumerate(self._COLUMNS):
2168            self.InsertColumn(index, col)
2169
2170        # Used to store the Max width in pixels for the data in the column.
2171        column_widths = [awx.get_width_height(self, s)[0] for s in self._COLUMNS]
2172
2173        # Used by the ColumnSorterMixin, see wx/lib/mixins/listctrl.py
2174        self.itemDataMap = {}
2175
2176        for index, psgen in enumerate(self.psgens):
2177            entry = self.make_entry(index, psgen)
2178            self.Append(entry)
2179            self.SetItemData(index, index)
2180            self.itemDataMap[index] = entry
2181
2182            w = [awx.get_width_height(self, s)[0] for s in entry]
2183            column_widths = map(max, zip(w, column_widths))
2184
2185        for index, col in enumerate(self._COLUMNS):
2186            self.SetColumnWidth(index, column_widths[index])
2187
2188        # Now that the list exists we can init the other base class, see wx/lib/mixins/listctrl.py
2189        listmix.ColumnSorterMixin.__init__(self, len(self._COLUMNS))
2190        listmix.ListCtrlAutoWidthMixin.__init__(self)
2191
2192        self.Bind(wx.EVT_LIST_ITEM_ACTIVATED, self.onItemActivated)
2193        self.Bind(wx.EVT_LIST_ITEM_RIGHT_CLICK, self.onRightClick)
2194
2195    @property
2196    def notebook(self):
2197        """Reference to the Wx window with the input file."""
2198        try:
2199            return self._notebook
2200        except AttributeError:
2201            parent = self.GetParent()
2202            cls = WxOncvFrame
2203
2204            while True:
2205                if parent is None:
2206                    raise RuntimeError("Cannot find parent with class %s, reached None parent!" % cls)
2207
2208                if isinstance(parent, cls):
2209                    break
2210                else:
2211                    parent = parent.GetParent()
2212
2213            # The input we need is an attribute of WxOncvFrame.
2214            self._notebook = parent.notebook
2215            return self._notebook
2216
2217    @staticmethod
2218    def make_entry(index, psgen):
2219        """Returns the entry associated to the generator psgen with the given index."""
2220        max_ecut = None
2221        max_atan_logder_l1err = None
2222        max_psexc_abserr = None
2223        herm_err = None
2224
2225        if psgen.results is None and psgen.status == psgen.S_OK:
2226            psgen.check_status()
2227
2228        if psgen.results is not None:
2229            max_ecut = psgen.results.max_ecut
2230            max_atan_logder_l1err = psgen.results.max_atan_logder_l1err
2231            max_psexc_abserr = psgen.results.max_psexc_abserr
2232            herm_err = psgen.results.herm_err
2233
2234        return [
2235            "%d\t\t" % index,
2236            "%s" % psgen.status,
2237            "%s" % max_ecut,
2238            "%s" % max_atan_logder_l1err,
2239            "%s" % max_psexc_abserr,
2240            "%s" % herm_err,
2241            "%s" % len(psgen.warnings),
2242        ]
2243
2244    def doRefresh(self):
2245        """Refresh the panel and redraw it."""
2246        column_widths = [awx.get_width_height(self, s)[0] for s in self._COLUMNS]
2247
2248        for index, psgen in enumerate(self.psgens):
2249            entry = self.make_entry(index, psgen)
2250            self.SetItemData(index, index)
2251            self.itemDataMap[index] = entry
2252
2253            w = [awx.get_width_height(self, s)[0] for s in entry]
2254            column_widths = map(max, zip(w, column_widths))
2255
2256        for index, col in enumerate(self._COLUMNS):
2257            self.SetColumnWidth(index, column_widths[index])
2258
2259    def add_psgen(self, psgen):
2260        """Add a PseudoGenerator to the list."""
2261        index = len(self.psgens)
2262        entry = self.make_entry(index, psgen)
2263        self.Append(entry)
2264        self.SetItemData(index, index)
2265        self.itemDataMap[index] = entry
2266
2267        # Add it to the list and update column widths.
2268        self.psgens.append(psgen)
2269        self.doRefresh()
2270
2271    def GetListCtrl(self):
2272        """Used by the ColumnSorterMixin, see wx/lib/mixins/listctrl.py"""
2273        return self
2274
2275    def getSelectedPseudoGen(self):
2276        """
2277        Returns the PseudoGenerators selected by the user.
2278        None if no selection has been done.
2279        """
2280        # Get selected index, map to index in psgens and return the object.
2281        item = self.GetFirstSelected()
2282        if item == -1: return None
2283        index = self.GetItemData(item)
2284        return self.psgens[index]
2285
2286    def onItemActivated(self, event):
2287        """Call psgen.plot_results."""
2288        psgen = self.getSelectedPseudoGen()
2289        if psgen is None: return
2290        psgen.plot_results()
2291
2292    def onPlotSubMenu(self, event):
2293        """Called by plot submenu."""
2294        psgen = self.getSelectedPseudoGen()
2295        if psgen is None or psgen.plotter is None:
2296            return
2297
2298        key = self._id2plotdata[event.GetId()]
2299        psgen.plotter.plot_key(key)
2300
2301    def onRightClick(self, event):
2302        """Generate the popup menu."""
2303        popup_menu = self.makePopupMenu()
2304        self.PopupMenu(popup_menu, event.GetPoint())
2305        popup_menu.Destroy()
2306
2307    def makePopupMenu(self):
2308        """
2309        Build and return a popup menu. Subclasses can extend or replace this base method.
2310        """
2311        self.ID_POPUP_STDIN = wx.NewId()
2312        self.ID_POPUP_STDOUT = wx.NewId()
2313        self.ID_POPUP_STDERR = wx.NewId()
2314        self.ID_POPUP_CHANGE_INPUT = wx.NewId()
2315        self.ID_POPUP_COMPUTE_HINTS = wx.NewId()
2316        self.ID_POPUP_COMPUTE_GBRV = wx.NewId()
2317        self.ID_POPUP_COMPUTE_PSPS = wx.NewId()
2318        self.ID_POPUP_SAVE_PSGEN = wx.NewId()
2319
2320        menu = wx.Menu()
2321
2322        # Make sub-menu with the list of supported quantities
2323        plot_submenu = wx.Menu()
2324
2325        # TODO: this list could be taken from the class or from the plotter instance.
2326        all_keys = [
2327            "radial_wfs",
2328            "projectors",
2329            "densities",
2330            "potentials",
2331            "der_potentials",
2332            "atan_logders",
2333            "ene_vs_ecut",
2334        ]
2335
2336        self._id2plotdata = {}
2337
2338        for aqt in all_keys:
2339            _id = wx.NewId()
2340            plot_submenu.Append(_id, aqt)
2341            self._id2plotdata[_id] = aqt
2342            self.Bind(wx.EVT_MENU, self.onPlotSubMenu, id=_id)
2343        menu.AppendMenu(-1, 'Plot', plot_submenu)
2344
2345        menu.Append(self.ID_POPUP_STDIN, "Show standard input")
2346        menu.Append(self.ID_POPUP_STDOUT, "Show standard output")
2347        menu.Append(self.ID_POPUP_STDERR, "Show standard error")
2348        menu.Append(self.ID_POPUP_CHANGE_INPUT, "Use these variables as new template")
2349        menu.Append(self.ID_POPUP_COMPUTE_HINTS, "Compute hints for ecut")
2350        #menu.Append(self.ID_POPUP_COMPUTE_GBRV, "Perform GBRV tests")
2351        menu.Append(self.ID_POPUP_COMPUTE_PSPS, "Get PSPS.nc file and plot data")
2352        menu.Append(self.ID_POPUP_SAVE_PSGEN, "Save PS generation")
2353
2354        # Associate menu/toolbar items with their handlers.
2355        menu_handlers = [
2356            (self.ID_POPUP_STDIN, self.onShowStdin),
2357            (self.ID_POPUP_STDOUT, self.onShowStdout),
2358            (self.ID_POPUP_STDERR, self.onShowStderr),
2359            (self.ID_POPUP_CHANGE_INPUT, self.onChangeInput),
2360            (self.ID_POPUP_COMPUTE_HINTS, self.onComputeHints),
2361            #(self.ID_POPUP_COMPUTE_GBRV, self.onGBRV),
2362            (self.ID_POPUP_COMPUTE_PSPS, self.onPsps),
2363            (self.ID_POPUP_SAVE_PSGEN, self.onSavePsgen),
2364        ]
2365
2366        for combo in menu_handlers:
2367            mid, handler = combo[:2]
2368            self.Bind(wx.EVT_MENU, handler, id=mid)
2369
2370        return menu
2371
2372    def onChangeInput(self, event):
2373        """Change the template input file."""
2374        psgen = self.getSelectedPseudoGen()
2375        if psgen is None: return
2376
2377        # Change the parameters in the notebook using those in psgen.stdin_path
2378        if not os.path.exists(psgen.stdin_path):
2379            return awx.showErrorMessage(self, "Input file %s does not exist" % psgen.stdin_path)
2380
2381        main_frame = self.notebook.GetParent()
2382        new_notebook = OncvNotebook.from_file(main_frame, psgen.stdin_path)
2383        main_frame.BuildUI(notebook=new_notebook)
2384
2385    def onComputeHints(self, event):
2386        psgen = self.getSelectedPseudoGen()
2387        if psgen is None: return
2388        print("workdir", psgen.workdir)
2389
2390        # Change the parameters in the notebook using those in psgen.stdin_path
2391        #if not os.path.exists(psgen.stdin_path):
2392        #    return awx.showErrorMessage(self, "Input file %s does not exist" % psgen.stdin_path)
2393        from abipy import abilab
2394        from pseudo_dojo.dojo.dojo_workflows import PPConvergenceFactory
2395        factory = PPConvergenceFactory()
2396
2397        print("pseudo", psgen.pseudo)
2398        workdir = os.path.join("HINTS_", psgen.pseudo.name)
2399        flow = abilab.AbinitFlow(workdir, pickle_protocol=0)
2400
2401        work = factory.work_for_pseudo(psgen.pseudo, ecut_slice=slice(4, None, 1), nlaunch=2)
2402        flow.register_work(work)
2403        flow.allocate()
2404        flow.build_and_pickle_dump()
2405        scheduler = abilab.PyFlowScheduler.from_user_config()
2406        scheduler.add_flow(flow)
2407        scheduler.start()
2408
2409    def onSavePsgen(self, event):
2410        # Update the notebook
2411        self.onChangeInput(event)
2412
2413        psgen = self.getSelectedPseudoGen()
2414        if psgen is None: return
2415
2416        main_frame = self.notebook.GetParent()
2417        input_file = main_frame.input_file
2418
2419        if input_file is None:
2420            raise NotImplementedError()
2421
2422        dirpath = os.path.dirname(input_file)
2423        basename = os.path.basename(input_file).replace(".in", "")
2424
2425        ps_dest = os.path.join(dirpath, basename + ".psp8")
2426        upf_dest = os.path.join(dirpath, basename + ".upf")
2427        out_dest = os.path.join(dirpath, basename + ".out")
2428        djrepo_dest = os.path.join(dirpath, basename + ".djrepo")
2429
2430        exists = []
2431        for f in [input_file, ps_dest, upf_dest, out_dest, djrepo_dest]:
2432            if os.path.exists(f): exists.append(os.path.basename(f))
2433
2434        if exists:
2435            msg = "File(s):\n%s already exist.\nDo you want to owerwrite them?" % "\n".join(exists)
2436            answer = awx.askUser(self, msg)
2437            if not answer: return
2438
2439        # Update the input file, then copy the pseudo file and the output file.
2440        with open(input_file, "wt") as fh:
2441            fh.write(self.notebook.makeInputString())
2442
2443        print(psgen.pseudo.path)
2444        shutil.copy(psgen.pseudo.path, ps_dest)
2445        upf_src = psgen.pseudo.path.replace(".psp8", ".upf")
2446        if os.path.exists(upf_src):
2447            shutil.copy(upf_src, upf_dest)
2448        shutil.copy(psgen.stdout_path, out_dest)
2449
2450        # Parse the output file
2451        onc_parser = OncvOutputParser(out_dest)
2452        onc_parser.scan()
2453        if not onc_parser.run_completed:
2454            raise RuntimeError("oncvpsp output is not complete. Exiting")
2455
2456        # Build dojoreport
2457        pseudo = Pseudo.from_file(ps_dest)
2458        report = DojoReport.empty_from_pseudo(pseudo, onc_parser.hints, devel=False)
2459        report.json_write()
2460
2461    def onGBRV(self, event):
2462        psgen = self.getSelectedPseudoGen()
2463        if psgen is None: return
2464
2465        # Change the parameters in the notebook using those in psgen.stdin_path
2466        #if not os.path.exists(psgen.stdin_path):
2467        #    return awx.showErrorMessage(self, "Input file %s does not exist" % psgen.stdin_path)
2468        from abipy import abilab
2469        from pseudo_dojo.dojo.dojo_workflows import GbrvFactory
2470        factory = GbrvFactory()
2471
2472        flow = abilab.AbinitFlow(workdir="GBRV", pickle_protocol=0)
2473        print("pseudo", psgen.pseudo)
2474        for struct_type in ["fcc", "bcc"]:
2475            work = factory.relax_and_eos_work(psgen.pseudo, struct_type)
2476            flow.register_work(work)
2477
2478        flow.allocate()
2479        flow.build_and_pickle_dump()
2480        scheduler = abilab.PyFlowScheduler.from_user_config()
2481        scheduler.add_flow(flow)
2482        scheduler.start()
2483
2484    def onPsps(self, event):
2485        psgen = self.getSelectedPseudoGen()
2486        if psgen is None: return
2487
2488        with psgen.pseudo.open_pspsfile(ecut=30) as psps:
2489            print("Printing data from:", psps.filepath)
2490            psps.plot(ecut_ffnl=60)
2491
2492    #def onAddToHistory(self, event):
2493    #    psgen = self.getSelectedPseudoGen()
2494    #    if psgen is None: return
2495    #    # Add it to history.
2496    #    self.history.append(psgen)
2497
2498    def _showStdfile(self, event, stdfile):
2499        """
2500        Helper function used to access the std files
2501        of the PseudoGenerator selected by the user.
2502        """
2503        psgen = self.getSelectedPseudoGen()
2504        if psgen is None: return
2505        call = dict(
2506            stdin=psgen.get_stdin,
2507            stdout=psgen.get_stdout,
2508            stderr=psgen.get_stderr,
2509        )[stdfile]
2510
2511        SimpleTextViewer(self, text=call()).Show()
2512
2513    def onShowStdin(self, event):
2514        """Open a frame with the input file."""
2515        self._showStdfile(event, "stdin")
2516
2517    def onShowStdout(self, event):
2518        """Open a frame with the output file."""
2519        self._showStdfile(event, "stdout")
2520
2521    def onShowStderr(self, event):
2522        """Open a frame with the stderr file."""
2523        self._showStdfile(event, "stderr")
2524
2525    def plot_columns(self, **kwargs):
2526        """Use matplotlib to plot the values reported in the columns."""
2527        keys = self._COLUMNS[2:]
2528        table = OrderedDict([(k, []) for k in keys])
2529
2530        for index, psgen in enumerate(self.psgens):
2531            entry = self.make_entry(index, psgen)
2532            for k, v in zip(keys, entry[2:]):
2533                #print("k", k, "v", v)
2534                try:
2535                    v = float(v)
2536                except ValueError:
2537                    # This happens if the run is not completed.
2538                    v = None
2539
2540                table[k].append(v)
2541
2542        # Return immediately if all entries are None i.e.
2543        # if all calculations are still running.
2544        count = 0
2545        for items in table.values():
2546            if all(item is None for item in items): count += 1
2547        if count == len(table): return
2548
2549        import matplotlib.pyplot as plt
2550
2551        # Build grid of plots.
2552        fig, ax_list = plt.subplots(nrows=len(table), ncols=1, sharex=False, squeeze=True)
2553
2554        for key, ax in zip(table.keys(), ax_list):
2555            ax.grid(True)
2556            ax.set_title(key)
2557            row = table[key]
2558            xs = np.arange(len(row))
2559            ys = np.array(table[key]).astype(np.double)
2560            # Use mask to exclude None values from the plot.
2561            mask = np.isfinite(ys)
2562            line, = ax.plot(xs[mask], ys[mask], linewidth=2, markersize=10, marker="o")
2563
2564        plt.show()
2565
2566
2567class PseudoGeneratorsPanel(awx.Panel):
2568    """
2569    A panel with a list control providing info on the status pseudogenerators
2570    Provides popup menus for interacting with the generators.
2571    """
2572    def __init__(self, parent, psgens=(), **kwargs):
2573        """
2574        Args:
2575            parent:
2576                Parent window.
2577            psgens:
2578                List of `PseudoGenerator` objects.
2579        """
2580        super(PseudoGeneratorsPanel, self).__init__(parent, **kwargs)
2581
2582        main_sizer = wx.BoxSizer(wx.VERTICAL)
2583        self.psgen_list_ctrl = PseudoGeneratorListCtrl(self, psgens)
2584
2585        main_sizer.Add(self.psgen_list_ctrl, 1, wx.ALL | wx.EXPAND | wx.ALIGN_CENTER_VERTICAL, 5)
2586        self.SetSizerAndFit(main_sizer)
2587
2588    @property
2589    def psgens(self):
2590        return self.psgen_list_ctrl.psgens
2591
2592    def plot_columns(self, **kwargs):
2593        self.psgen_list_ctrl.plot_columns(**kwargs)
2594
2595
2596class PseudoGeneratorsFrame(awx.Frame):
2597    """
2598    This frame contains a list of pseudopotential generators,
2599    It provides controls to run the calculation, and inspect/plot the results.
2600    """
2601    REFRESH_INTERVAL = 120
2602
2603    HELP_MSG = """\
2604This window allows you to generate and analyze multiple pseudopotentials.
2605"""
2606
2607    def __init__(self, parent, psgens=(), **kwargs):
2608        """
2609        Args:
2610            psgens:
2611                List of `PseudoGenerators`.
2612        """
2613        super(PseudoGeneratorsFrame, self).__init__(parent, -1, **add_size(kwargs))
2614
2615        # Build menu, toolbar and status bar.
2616        self.makeToolBar()
2617        self.statusbar = self.CreateStatusBar()
2618        self.Centre()
2619
2620        self.panel = panel = wx.Panel(self, -1)
2621        self.main_sizer = main_sizer = wx.BoxSizer(wx.VERTICAL)
2622
2623        self.psgens_wxlist = PseudoGeneratorsPanel(panel, psgens)
2624        main_sizer.Add(self.psgens_wxlist, 1, wx.ALL | wx.EXPAND, 5)
2625
2626        submit_button = wx.Button(panel, -1, label='Submit')
2627        submit_button.Bind(wx.EVT_BUTTON, self.OnSubmitButton)
2628
2629        text = wx.StaticText(panel, -1, "Max nlaunch:")
2630        text.Wrap(-1)
2631        text.SetToolTipString("Maximum number of tasks that can be submitted. Use -1 for unlimited launches.")
2632        self.max_nlaunch = wx.SpinCtrl(panel, -1, value=str(get_ncpus()), min=-1)
2633
2634        help_button = wx.Button(panel, wx.ID_HELP)
2635        help_button.Bind(wx.EVT_BUTTON, self.onHelp)
2636        main_sizer.Add(help_button, 0, flag=wx.ALL | wx.ALIGN_RIGHT)
2637
2638        hsizer = wx.BoxSizer(wx.HORIZONTAL)
2639        hsizer.Add(submit_button, 0, wx.ALL | wx.ALIGN_CENTER_VERTICAL, 5)
2640        hsizer.Add(text, 0, wx.ALL | wx.ALIGN_CENTER_VERTICAL, 5)
2641        hsizer.Add(self.max_nlaunch, 0, wx.ALL | wx.ALIGN_CENTER_VERTICAL, 5)
2642        hsizer.Add(help_button, 0, wx.ALL | wx.ALIGN_CENTER_VERTICAL, 5)
2643
2644        main_sizer.Add(hsizer, 0, wx.ALL | wx.ALIGN_CENTER_HORIZONTAL, 5)
2645        panel.SetSizerAndFit(main_sizer)
2646
2647        # Register this event when the GUI is IDLE
2648        self.last_refresh = time.time()
2649        self.Bind(wx.EVT_IDLE, self.OnIdle)
2650
2651    @property
2652    def psgens(self):
2653        """List of PseudoGenerators."""
2654        return self.psgens_wxlist.psgens
2655
2656    def launch_psgens(self):
2657        return self.psgens_wxlist.psgen_list_ctrl.launch_psges()
2658
2659    def makeToolBar(self):
2660        """Create toolbar."""
2661        self.toolbar = toolbar = self.CreateToolBar()
2662        toolbar.SetToolBitmapSize(wx.Size(48, 48))
2663
2664        def bitmap(path):
2665            return wx.Bitmap(awx.path_img(path))
2666
2667        self.ID_SHOW_INPUTS = wx.NewId()
2668        self.ID_SHOW_OUTPUTS = wx.NewId()
2669        self.ID_SHOW_ERRS = wx.NewId()
2670        self.ID_CHECK_STATUS = wx.NewId()
2671        self.ID_MULTI_PLOTTER = wx.NewId()
2672        self.ID_PLOT_COLUMNS = wx.NewId()
2673
2674        toolbar.AddSimpleTool(self.ID_SHOW_INPUTS, bitmap("in.png"), "Visualize the input file(s) of the generators.")
2675        toolbar.AddSimpleTool(self.ID_SHOW_OUTPUTS, bitmap("out.png"), "Visualize the output file(s) of the generators.")
2676        toolbar.AddSimpleTool(self.ID_SHOW_ERRS, bitmap("log.png"), "Visualize the errors file(s) of the generators.")
2677        toolbar.AddSimpleTool(self.ID_MULTI_PLOTTER, bitmap("log.png"), "Multi plotter.")
2678        toolbar.AddSimpleTool(self.ID_PLOT_COLUMNS, bitmap("log.png"), "Plot columns.")
2679        toolbar.AddSeparator()
2680        toolbar.AddSimpleTool(self.ID_CHECK_STATUS, bitmap("refresh.png"), "Check the status of the workflow(s).")
2681
2682        toolbar.Realize()
2683
2684        # Associate menu/toolbar items with their handlers.
2685        menu_handlers = [
2686            (self.ID_SHOW_INPUTS, self.onShowInputs),
2687            (self.ID_SHOW_OUTPUTS, self.onShowOutputs),
2688            (self.ID_SHOW_ERRS, self.onShowErrors),
2689            (self.ID_CHECK_STATUS, self.onCheckStatus),
2690            (self.ID_MULTI_PLOTTER, self.onMultiPlotter),
2691            (self.ID_PLOT_COLUMNS, self.onPlotColumns),
2692        ]
2693
2694        for combo in menu_handlers:
2695            mid, handler = combo[:2]
2696            self.Bind(wx.EVT_MENU, handler, id=mid)
2697
2698    def OnSubmitButton(self, event):
2699        """
2700        Called when Run button is pressed.
2701        Run the calculation in a subprocess in non-blocking mode and add it to
2702        the list containing the generators in executions
2703        Submit up to max_nlauch tasks (-1 to run'em all)
2704        """
2705        self.launch_psgens()
2706
2707    def launch_psgens(self):
2708        max_nlaunch = int(self.max_nlaunch.GetValue())
2709
2710        nlaunch = 0
2711        for psgen in self.psgens:
2712            nlaunch += psgen.start()
2713            if nlaunch == max_nlaunch:
2714                break
2715
2716        self.statusbar.PushStatusText("Submitted %d tasks" % nlaunch)
2717
2718    def onCheckStatus(self, event):
2719        """
2720        Callback triggered by the checkstatus button.
2721        Check the status of the `PseudoGenerators` and refresh the panel.
2722        """
2723        self.checkStatusAndRedraw()
2724
2725    def checkStatusAndRedraw(self):
2726        """Check the status of all the workflows and redraw the panel."""
2727        self.statusbar.PushStatusText("Checking status...")
2728        start = time.time()
2729
2730        for psgen in self.psgens:
2731            psgen.check_status()
2732        self.statusbar.PushStatusText("Check completed in %.1f [s]" % (time.time() - start))
2733
2734        # Redraw the panel
2735        main_sizer = self.main_sizer
2736        main_sizer.Hide(0)
2737        main_sizer.Remove(0)
2738        new_psgen_wxlist = PseudoGeneratorsPanel(self.panel, self.psgens)
2739        main_sizer.Insert(0, new_psgen_wxlist, 1, wx.ALL | wx.EXPAND, 5)
2740        self.psgen_wxlist = new_psgen_wxlist
2741
2742        self.panel.Layout()
2743
2744        # Write number of jobs with given status.
2745        #message = ", ".join("%s: %s" % (k, v) for (k, v) in counter.items())
2746        #self.statusbar.PushStatusText(message)
2747
2748    def OnIdle(self, event):
2749        """Function executed when the GUI is idle."""
2750        now = time.time()
2751        if (now - self.last_refresh) > self.REFRESH_INTERVAL:
2752            self.checkStatusAndRedraw()
2753            self.last_refresh = time.time()
2754
2755    def onShowInputs(self, event):
2756        """Show all input files."""
2757        TextNotebookFrame(self, text_list=[psgen.get_stdin() for psgen in self.psgens],
2758                          page_names=["PSGEN # %d" % i for i in range(len(self.psgens))]).Show()
2759
2760    def onShowOutputs(self, event):
2761        """Show all output files."""
2762        TextNotebookFrame(self, text_list=[psgen.get_stdout() for psgen in self.psgens],
2763                          page_names=["PSGEN # %d" % i for i in range(len(self.psgens))]).Show()
2764
2765    def onShowErrors(self, event):
2766        """Show all error files."""
2767        TextNotebookFrame(self, text_list=[psgen.get_stderr() for psgen in self.psgens],
2768                          page_names=["PSGEN # %d" % i for i in range(len(self.psgens))]).Show()
2769
2770    def onPlotColumns(self, event):
2771        self.psgens_wxlist.plot_columns()
2772
2773    def onMultiPlotter(self, event):
2774        """Open a dialog that allows the user to plot the results of multiple generators."""
2775        multi_plotter = MultiPseudoGenDataPlotter()
2776
2777        # Add psgen only if run is OK.
2778        for i, psgen in enumerate(self.psgens):
2779            if psgen.status == psgen.S_OK:
2780                multi_plotter.add_psgen(label="%d" % i, psgen=psgen)
2781
2782        # Return immediately if no calculation is OK.
2783        if not len(multi_plotter):
2784            return
2785
2786        keys = list(multi_plotter.keys())
2787
2788        class MyFrame(awx.FrameWithChoice):
2789            """Get a string with the quantity to plot and call multi_plotter.plot_key"""
2790            def onOkButton(self, event):
2791                multi_plotter.plot_key(key=self.getChoice())
2792
2793        MyFrame(self, choices=keys, title="MultiPlotter").Show()
2794
2795
2796def wxapp_oncvpsp(filepath=None):
2797    """Standalone application."""
2798    class OncvApp(awx.App):
2799        def OnInit(self):
2800            # Enforce WXAgg as matplotlib backend to avoid nasty SIGSEGV in the C++ layer
2801            # occurring when WX Guis produce plots with other backends.
2802            import matplotlib
2803            matplotlib.use('WXAgg')
2804
2805            frame = WxOncvFrame(None, filepath)
2806            #frame = my_periodic_table(None)
2807            #frame = OncvParamsFrame(None, z=12)
2808            frame.Show(True)
2809            self.SetTopWindow(frame)
2810            return True
2811
2812    app = OncvApp()
2813    return app
2814
2815if __name__ == "__main__":
2816    import sys
2817    #onc_inp = OncvInput.from_file("08_O.dat")
2818    #print(onc_inp)
2819
2820    #print("optimizing qcut")
2821    #for inp in onc_inp.optimize_qcuts_for_l(l=0, new_qcuts=[1, 9]):
2822    #    a = 1
2823    #    #print("new model\n", inp)
2824
2825    #for inp in onc_inp.optimize_vloc():
2826    #    print("new\n", inp)
2827    #for inp in onc_inp.optimize_modelcore():
2828    #    print("new model\n", inp)
2829    #sys.exit(0)
2830    try:
2831        filepaths = sys.argv[1:]
2832    except IndexError:
2833        filepaths = None
2834
2835    if filepaths is not None:
2836        if filepaths[0] == "table":
2837            for symbol in all_symbols():
2838                path = symbol + ".dat"
2839                if os.path.exists(path):
2840                    print("Will open file %s" % path)
2841                    wxapp_oncvpsp(path).MainLoop()
2842        else:
2843            for filepath in filepaths:
2844                #print(filepath)
2845                wxapp_oncvpsp(filepath).MainLoop()
2846    else:
2847        wxapp_oncvpsp().MainLoop()
2848
2849    #app = awx.App()
2850    #frame = QcutOptimizationFrame(None, lmax=1)
2851    #frame.Show()
2852    #app.MainLoop()
2853
2854