1"""
2Objects used to extract and plot results from output files in text format.
3"""
4import os
5import numpy as np
6import pandas as pd
7
8from collections import OrderedDict
9from io import StringIO
10from monty.string import is_string, marquee
11from monty.functools import lazy_property
12from monty.termcolor import cprint
13from pymatgen.core.units import bohr_to_ang
14from abipy.core.symmetries import AbinitSpaceGroup
15from abipy.core.structure import Structure, dataframes_from_structures
16from abipy.core.kpoints import has_timrev_from_kptopt
17from abipy.core.mixins import TextFile, AbinitNcFile, NotebookWriter
18from abipy.abio.inputs import GEOVARS
19from abipy.abio.timer import AbinitTimerParser
20from abipy.abio.robots import Robot
21from abipy.flowtk import EventsParser, NetcdfReader, GroundStateScfCycle, D2DEScfCycle
22
23
24class AbinitTextFile(TextFile):
25    """
26    Base class for the ABINIT main output files and log files.
27    """
28    @property
29    def events(self):
30        """
31        List of ABINIT events reported in the file.
32        """
33        # Parse the file the first time the property is accessed or when mtime is changed.
34        stat = os.stat(self.filepath)
35        if stat.st_mtime != self._last_mtime or not hasattr(self, "_events"):
36            self._events = EventsParser().parse(self.filepath)
37        return self._events
38
39    def get_timer(self):
40        """
41        Timer data.
42        """
43        timer = AbinitTimerParser()
44        timer.parse(self.filepath)
45        return timer
46
47
48class AbinitLogFile(AbinitTextFile, NotebookWriter):
49    """
50    Class representing the Abinit log file.
51
52    .. rubric:: Inheritance Diagram
53    .. inheritance-diagram:: AbinitLogFile
54    """
55
56    def to_string(self, verbose=0):
57        return str(self.events)
58
59    def plot(self, **kwargs):
60        """Empty placeholder."""
61        return None
62
63    def yield_figs(self, **kwargs):  # pragma: no cover
64        """
65        This function *generates* a predefined list of matplotlib figures with minimal input from the user.
66        """
67        yield None
68
69    def write_notebook(self, nbpath=None):
70        """
71        Write a jupyter_ notebook to ``nbpath``. If nbpath is None, a temporay file in the current
72        working directory is created. Return path to the notebook.
73        """
74        nbformat, nbv, nb = self.get_nbformat_nbv_nb(title=None)
75
76        nb.cells.extend([
77            nbv.new_code_cell("abilog = abilab.abiopen('%s')" % self.filepath),
78            nbv.new_code_cell("print(abilog.events)"),
79        ])
80
81        return self._write_nb_nbpath(nb, nbpath)
82
83
84class AbinitOutputFile(AbinitTextFile, NotebookWriter):
85    """
86    Class representing the main Abinit output file.
87
88    .. rubric:: Inheritance Diagram
89    .. inheritance-diagram:: AbinitOutputFile
90    """
91    # TODO: Extract number of errors and warnings.
92
93    def __init__(self, filepath):
94        super().__init__(filepath)
95        self.debug_level = 0
96        self._parse()
97
98    def _parse(self):
99        """
100        header: String with the input variables
101        footer: String with the output variables
102        datasets: Dictionary mapping dataset index to list of strings.
103        """
104        # Get code version and find magic line signaling that the output file is completed.
105        self.version, self.run_completed = None, False
106        self.overall_cputime, self.overall_walltime = 0.0, 0.0
107        self.proc0_cputime, self.proc0_walltime = 0.0, 0.0
108
109        with open(self.filepath) as fh:
110            for line in fh:
111                if self.version is None and line.startswith(".Version"):
112                    self.version = line.split()[1]
113
114                if line.startswith("- Proc."):
115                    #- Proc.   0 individual time (sec): cpu=         25.5  wall=         26.1
116                    tokens = line.split()
117                    self.proc0_walltime = float(tokens[-1])
118                    self.proc0_cputime = float(tokens[-3])
119
120                if line.startswith("+Overall time"):
121                    #+Overall time at end (sec) : cpu=         25.5  wall=         26.1
122                    tokens = line.split()
123                    self.overall_cputime = float(tokens[-3])
124                    self.overall_walltime = float(tokens[-1])
125
126                if " Calculation completed." in line:
127                    self.run_completed = True
128
129        # Parse header to get important dimensions and variables
130        self.header, self.footer, self.datasets = [], [], OrderedDict()
131        where = "in_header"
132
133        with open(self.filepath, "rt") as fh:
134            for line in fh:
135                if "== DATASET" in line:
136                    # Save dataset number
137                    # == DATASET  1 ==================================================================
138                    where = int(line.replace("=", "").split()[-1])
139                    assert where not in self.datasets
140                    self.datasets[where] = []
141                elif "== END DATASET(S) " in line:
142                    where = "in_footer"
143
144                if where == "in_header":
145                    self.header.append(line)
146                elif where == "in_footer":
147                    self.footer.append(line)
148                else:
149                    # dataset number --> lines
150                    self.datasets[where].append(line)
151
152        self.header = "".join(self.header)
153        if self.debug_level: print("header:\n", self.header)
154        # Output files produced in dryrun_mode contain the following line:
155        # abinit : before driver, prtvol=0, debugging mode => will skip driver
156        self.dryrun_mode = "debugging mode => will skip driver" in self.header
157        #print("dryrun_mode:", self.dryrun_mode)
158
159        #if " jdtset " in self.header: raise NotImplementedError("jdtset is not supported")
160        #if " udtset " in self.header: raise NotImplementedError("udtset is not supported")
161
162        self.ndtset = len(self.datasets)
163        if not self.datasets:
164            #raise NotImplementedError("Empty dataset sections.")
165            self.ndtset = 1
166            self.datasets[1] = "Empty dataset"
167
168        for key, data in self.datasets.items():
169            if self.debug_level: print("data")
170            self.datasets[key] = "".join(data)
171            if self.debug_level: print(self.datasets[key])
172
173        self.footer = "".join(self.footer)
174        if self.debug_level: print("footer:\n", self.footer)
175
176        self.initial_vars_global, self.initial_vars_dataset = self._parse_variables("header")
177        self.final_vars_global, self.final_vars_dataset = None, None
178        if self.run_completed:
179            if self.dryrun_mode:
180                # footer is not present. Copy values from header.
181                self.final_vars_global, self.final_vars_dataset = self.initial_vars_global, self.initial_vars_dataset
182            else:
183                self.final_vars_global, self.final_vars_dataset = self._parse_variables("footer")
184
185    def _parse_variables(self, what):
186        vars_global = OrderedDict()
187        vars_dataset = OrderedDict([(k, OrderedDict()) for k in self.datasets.keys()])
188        #print("keys", vars_dataset.keys())
189
190        lines = getattr(self, what).splitlines()
191        if what == "header":
192            magic_start = " -outvars: echo values of preprocessed input variables --------"
193        elif what == "footer":
194            magic_start = " -outvars: echo values of variables after computation  --------"
195        else:
196            raise ValueError("Invalid value for what: `%s`" % str(what))
197
198        magic_stop = "================================================================================"
199
200        # Select relevant portion with variables.
201        for i, line in enumerate(lines):
202            if magic_start in line:
203                break
204        else:
205            raise ValueError("Cannot find magic_start line: `%s`\nPerhaps this is not an Abinit output file!" % magic_start)
206        lines = lines[i+1:]
207
208        for i, line in enumerate(lines):
209            if magic_stop in line:
210                break
211        else:
212            raise ValueError("Cannot find magic_stop line: `%s`\nPerhaps this is not an Abinit output file!" % magic_stop)
213        lines = lines[:i]
214
215        # Parse data. Assume format:
216        #   timopt          -1
217        #    tnons      0.0000000  0.0000000  0.0000000     0.2500000  0.2500000  0.2500000
218        #               0.0000000  0.0000000  0.0000000     0.2500000  0.2500000  0.2500000
219        def get_dtindex_key_value(line):
220            tokens = line.split()
221            s, value = tokens[0], " ".join(tokens[1:])
222            l = []
223            for i, c in enumerate(s[::-1]):
224                if c.isalpha():
225                    key = s[:len(s)-i]
226                    break
227                l.append(c)
228            else:
229                raise ValueError("Cannot find dataset index in token: %s\n" % s)
230
231            #print(line, "\n", l)
232            dtindex = None
233            if l:
234                l.reverse()
235                dtindex = int("".join(l))
236            return dtindex, key, value
237
238        # (varname, dtindex), [line1, line2 ...]
239        stack_var, stack_lines = None, []
240
241        def pop_stack():
242            if stack_lines:
243                key, dtidx = stack_var
244                value = " ".join(stack_lines)
245                if dtidx is None:
246                    vars_global[key] = value
247                else:
248                    vars_dataset[dtidx][key] = value
249
250        for line in lines:
251            if not line: continue
252            # Ignore first char
253            line = line[1:].lstrip().rstrip()
254            if not line: continue
255            #print("line", line)
256            if line[0].isalpha():
257                pop_stack()
258                stack_lines = []
259                dtidx, key, value = get_dtindex_key_value(line)
260                stack_var = (key, dtidx)
261                stack_lines.append(value)
262            else:
263                stack_lines.append(line)
264
265        pop_stack()
266
267        return vars_global, vars_dataset
268
269    def _get_structures(self, what):
270        if what == "header":
271            vars_global, vars_dataset = self.initial_vars_global, self.initial_vars_dataset
272        elif what == "footer":
273            vars_global, vars_dataset = self.final_vars_global, self.final_vars_dataset
274        else:
275            raise ValueError("Invalid value for what: `%s`" % str(what))
276
277        #print("global", vars_global["acell"])
278        from abipy.abio.abivars import is_abiunit
279        inigeo = {k: vars_global[k] for k in GEOVARS if k in vars_global}
280
281        spgvars = ("spgroup", "symrel", "tnons", "symafm")
282        spgd_global = {k: vars_global[k] for k in spgvars if k in vars_global}
283        global_kptopt = vars_global.get("kptopt", 1)
284
285        structures = []
286        for i in self.datasets:
287            # This code breaks down if there are conflicting GEOVARS in globals and dataset.
288            d = inigeo.copy()
289            d.update({k: vars_dataset[i][k] for k in GEOVARS if k in vars_dataset[i]})
290
291            for key, value in d.items():
292                # Must handle possible unit.
293                fact = 1.0
294                tokens = [t.lower() for t in value.split()]
295                if is_abiunit(tokens[-1]):
296                    tokens, unit = tokens[:-1], tokens[-1]
297                    if unit in ("angstr", "angstrom", "angstroms"):
298                        fact = 1.0 / bohr_to_ang
299                    elif unit in ("bohr", "bohrs", "au"):
300                        fact = 1.0
301                    else:
302                        raise ValueError("Don't know how to handle unit: %s" % unit)
303
304                s = " ".join(tokens)
305                dtype = float if key not in ("ntypat", "typat", "natom") else int
306                try:
307                    #print(key, s)
308                    value = np.fromstring(s, sep=" ", dtype=dtype)
309                    #print(key, value)
310                    if fact != 1.0: value *= fact # Do not change integer arrays e.g typat!
311                    d[key] = value
312                except ValueError as exc:
313                    print(key, s)
314                    raise exc
315
316            if "rprim" not in d and "angdeg" not in d: d["rprim"] = np.eye(3)
317            if "natom" in d and d["natom"] == 1 and all(k not in d for k in ("xred", "xcart", "xangst")):
318                d["xred"] = np.zeros(3)
319            #print(d)
320            abistr = Structure.from_abivars(d)
321
322            # Extract Abinit spacegroup.
323            spgd = spgd_global.copy()
324            spgd.update({k: vars_dataset[i][k] for k in spgvars if k in vars_dataset[i]})
325
326            spgid = int(spgd.get("spgroup", 0))
327            if "symrel" not in spgd:
328                symrel = np.reshape(np.eye(3, 3, dtype=int), (1, 3, 3))
329                spgd["symrel"] = " ".join((str(i) for i in symrel.flatten()))
330            else:
331                symrel = np.reshape(np.array([int(n) for n in spgd["symrel"].split()], dtype=int), (-1, 3, 3))
332            nsym = len(symrel)
333            assert nsym == spgd.get("nsym", nsym) #; print(symrel.shape)
334
335            if "tnons" in spgd:
336                tnons = np.reshape(np.array([float(t) for t in spgd["tnons"].split()], dtype=float), (nsym, 3))
337            else:
338                tnons = np.zeros((nsym, 3))
339
340            if "symafm" in spgd:
341                symafm = np.array([int(n) for n in spgd["symafm"].split()], dtype=int)
342                symafm.shape = (nsym,)
343            else:
344                symafm = np.ones(nsym, dtype=int)
345
346            try:
347                has_timerev = has_timrev_from_kptopt(vars_dataset[i].get("kptopt", global_kptopt))
348                abi_spacegroup = AbinitSpaceGroup(spgid, symrel, tnons, symafm, has_timerev, inord="C")
349                abistr.set_abi_spacegroup(abi_spacegroup)
350            except Exception as exc:
351                print("Cannot build AbinitSpaceGroup from the variables reported in file!\n", str(exc))
352
353            structures.append(abistr)
354
355        return structures
356
357    @lazy_property
358    def initial_structures(self):
359        """List of initial |Structure|."""
360        return self._get_structures("header")
361
362    @property
363    def has_same_initial_structures(self):
364        """True if all initial structures are equal."""
365        return all(self.initial_structures[0] == s for s in self.initial_structures)
366
367    @lazy_property
368    def final_structures(self):
369        """List of final |Structure|."""
370        if self.run_completed:
371            return self._get_structures("footer")
372        else:
373            cprint("Cannot extract final structures from file.\n %s" % self.filepath, "red")
374            return []
375
376    @lazy_property
377    def initial_structure(self):
378        """
379        The |Structure| defined in the output file.
380
381        If the input file contains multiple datasets **AND** the datasets
382        have different structures, this property returns None.
383        In this case, one has to access the structure of the individual datasets.
384        For example:
385
386            self.initial_structures[0]
387
388        gives the structure of the first dataset.
389        """
390        if not self.has_same_initial_structures:
391            print("Datasets have different structures. Returning None. Use initial_structures[0]")
392            return None
393        return self.initial_structures[0]
394
395    @property
396    def has_same_final_structures(self):
397        """True if all initial structures are equal."""
398        return all(self.final_structures[0] == s for s in self.final_structures)
399
400    @lazy_property
401    def final_structure(self):
402        """
403        The |Structure| defined in the output file.
404
405        If the input file contains multiple datasets **AND** the datasets
406        have different structures, this property returns None.
407        In this case, one has to access the structure of the individual datasets.
408        For example:
409
410            self.final_structures[0]
411
412        gives the structure of the first dataset.
413        """
414        if not self.has_same_final_structures:
415            print("Datasets have different structures. Returning None. Use final_structures[0]")
416            return None
417        return self.final_structures[0]
418
419    def diff_datasets(self, dt_list1, dt_list2, with_params=True, differ="html", dryrun=False):
420        """
421        Compare datasets
422        """
423        if not isinstance(dt_list1, (list, tuple)): dt_list1 = [dt_list1]
424        if not isinstance(dt_list2, (list, tuple)): dt_list2 = [dt_list2]
425
426        dt_lists = [dt_list1, dt_list2]
427        import tempfile
428        tmp_names = []
429        for i in range(2):
430            _, tmpname = tempfile.mkstemp(text=True)
431            tmp_names.append(tmpname)
432            with open(tmpname, "wt") as fh:
433                if with_params: fh.write(self.header)
434                for idt in dt_lists[i]:
435                    fh.write(self.datasets[idt])
436                if with_params: fh.write(self.footer)
437
438        if differ == "html":
439            from abipy.tools.devtools import HtmlDiff
440            diff = HtmlDiff(tmp_names)
441            if dryrun:
442                return diff
443            else:
444                return diff.open_browser()
445        else:
446            cmd = "%s %s %s" % (differ, tmp_names[0], tmp_names[1])
447            if dryrun:
448                return cmd
449            else:
450                return os.system(cmd)
451
452    def __str__(self):
453        return self.to_string()
454
455    def to_string(self, verbose=0):
456        """String representation."""
457        lines = ["ndtset: %d, completed: %s" % (self.ndtset, self.run_completed)]
458        app = lines.append
459
460        # Different cases depending whether final structures are available
461        # and whether structures are equivalent.
462        if self.run_completed:
463            if self.has_same_final_structures:
464                if self.initial_structure != self.final_structure:
465                    # Structural relaxation.
466                    df = dataframes_from_structures([self.initial_structure, self.final_structure],
467                                                    index=["initial", "final"])
468                    app("Lattice parameters:")
469                    app(str(df.lattice))
470                    app("Atomic coordinates:")
471                    app(str(df.coords))
472                else:
473                    # initial == final. Print final structure.
474                    app(self.final_structure.to_string(verbose=verbose))
475        else:
476            # Final structures are not available.
477            if self.has_same_initial_structures:
478                app(self.initial_structure.to_string(verbose=verbose))
479            else:
480                df = dataframes_from_structures(self.initial_structures,
481                                                index=[i+1 for i in range(self.ndtset)])
482                app("Lattice parameters:")
483                app(str(df.lattice))
484                app("Atomic coordinates:")
485                app(str(df.coords))
486
487        # Print dataframe with dimensions.
488        df = self.get_dims_spginfo_dataframe(verbose=verbose)
489        from abipy.tools.printing import print_dataframe
490        strio = StringIO()
491        print_dataframe(df, file=strio)
492        strio.seek(0)
493        app("")
494        app(marquee("Dimensions of calculation", mark="="))
495        app("".join(strio))
496
497        return "\n".join(lines)
498
499    def get_dims_spginfo_dataframe(self, verbose=0):
500        """
501        Parse the section with the dimensions of the calculation. Return Dataframe.
502        """
503        dims_dataset, spginfo_dataset = self.get_dims_spginfo_dataset(verbose=verbose)
504        rows = []
505        for dtind, dims in dims_dataset.items():
506            d = OrderedDict()
507            d["dataset"] = dtind
508            d.update(dims)
509            d.update(spginfo_dataset[dtind])
510            rows.append(d)
511
512        df = pd.DataFrame(rows, columns=list(rows[0].keys()) if rows else None)
513        df = df.set_index('dataset')
514        return df
515
516    def get_dims_spginfo_dataset(self, verbose=0):
517        """
518        Parse the section with the dimensions of the calculation. Return dictionaries
519
520        Args:
521            verbose: Verbosity level.
522
523        Return: (dims_dataset, spginfo_dataset)
524            where dims_dataset[i] is an OrderedDict with the dimensions of dataset `i`
525            spginfo_dataset[i] is a dictionary with space group information.
526        """
527        # If single dataset, we have to parse
528        #
529        #  Symmetries : space group Fd -3 m (#227); Bravais cF (face-center cubic)
530        # ================================================================================
531        #  Values of the parameters that define the memory need of the present run
532        #      intxc =       0    ionmov =       0      iscf =       7    lmnmax =       6
533        #      lnmax =       6     mgfft =      18  mpssoang =       3    mqgrid =    3001
534        #      natom =       2  nloc_mem =       1    nspden =       1   nspinor =       1
535        #     nsppol =       1      nsym =      48    n1xccc =    2501    ntypat =       1
536        #     occopt =       1   xclevel =       2
537        # -    mband =           8        mffmem =           1         mkmem =          29
538        #        mpw =         202          nfft =        5832          nkpt =          29
539        # ================================================================================
540        # P This job should need less than                       3.389 Mbytes of memory.
541        #   Rough estimation (10% accuracy) of disk space for files :
542        # _ WF disk file :      0.717 Mbytes ; DEN or POT disk file :      0.046 Mbytes.
543        # ================================================================================
544
545        # If multi datasets we have to parse:
546
547        #  DATASET    2 : space group F-4 3 m (#216); Bravais cF (face-center cubic)
548        # ================================================================================
549        #  Values of the parameters that define the memory need for DATASET  2.
550        #      intxc =       0    ionmov =       0      iscf =       7    lmnmax =       2
551        #      lnmax =       2     mgfft =      12  mpssoang =       3    mqgrid =    3001
552        #      natom =       2  nloc_mem =       1    nspden =       1   nspinor =       1
553        #     nsppol =       1      nsym =      24    n1xccc =    2501    ntypat =       2
554        #     occopt =       1   xclevel =       1
555        # -    mband =          10        mffmem =           1         mkmem =           2
556        #        mpw =          69          nfft =        1728          nkpt =           2
557        # ================================================================================
558        # P This job should need less than                       1.331 Mbytes of memory.
559        #   Rough estimation (10% accuracy) of disk space for files :
560        # _ WF disk file :      0.023 Mbytes ; DEN or POT disk file :      0.015 Mbytes.
561        # ================================================================================
562
563        magic = "Values of the parameters that define the memory need"
564        memory_pre = "P This job should need less than"
565        magic_exit = "------------- Echo of variables that govern the present computation"
566        filesizes_pre = "_ WF disk file :"
567        #verbose = 1
568
569        def parse_spgline(line):
570            """Parse the line with space group info, return dict."""
571            # Could use regular expressions ...
572            i = line.find("space group")
573
574            if i == -1:
575                # the unit cell is not primitive
576                return {}
577
578            spg_str, brav_str = line[i:].replace("space group", "").split(";")
579            toks = spg_str.split()
580            return {
581                "spg_symbol": "".join(toks[:-1]),
582                "spg_number": int(toks[-1].replace("(", "").replace(")", "").replace("#", "")),
583                "bravais": brav_str.strip(),
584            }
585
586        from abipy.tools.numtools import grouper
587        dims_dataset, spginfo_dataset = OrderedDict(), OrderedDict()
588        inblock = 0
589        with open(self.filepath, "rt") as fh:
590            for line in fh:
591                line = line.strip()
592                if verbose > 1: print("inblock:", inblock, " at line:", line)
593
594                if line.startswith(magic_exit): break
595
596                if (not line or line.startswith("===") or line.startswith("---")
597                    #or line.startswith("P")
598                    or line.startswith("Rough estimation") or line.startswith("PAW method is used")):
599                    continue
600
601                if line.startswith("DATASET") or line.startswith("Symmetries :"):
602                    # Get dataset index, parse space group and lattice info, init new dims dict.
603                    inblock = 1
604                    if line.startswith("Symmetries :"):
605                        # No multidataset
606                        dtindex = 1
607                    else:
608                        tokens = line.split()
609                        dtindex = int(tokens[1])
610
611                    dims_dataset[dtindex] = dims = OrderedDict()
612                    spginfo_dataset[dtindex] = parse_spgline(line)
613                    continue
614
615                if inblock == 1 and line.startswith(magic):
616                    inblock = 2
617                    continue
618
619                if inblock == 2:
620                    # Lines with data.
621                    if line.startswith("For the susceptibility"): continue
622
623                    if line.startswith(memory_pre):
624                        dims["mem_per_proc_mb"] = float(line.replace(memory_pre, "").split()[0])
625                    elif line.startswith(filesizes_pre):
626                        tokens = line.split()
627                        mbpos = [i - 1 for i, t in enumerate(tokens) if t.startswith("Mbytes")]
628                        assert len(mbpos) == 2
629                        dims["wfk_size_mb"] = float(tokens[mbpos[0]])
630                        dims["denpot_size_mb"] = float(tokens[mbpos[1]])
631                    elif line.startswith("Pmy_natom="):
632                        dims.update(my_natom=int(line.replace("Pmy_natom=", "").strip()))
633                        #print("my_natom", dims["my_natom"])
634                    else:
635                        if line and line[0] == "-": line = line[1:]
636                        tokens = grouper(2, line.replace("=", "").split())
637                        if verbose > 1: print("tokens:", tokens)
638                        dims.update([(t[0], int(t[1])) for t in tokens])
639
640            return dims_dataset, spginfo_dataset
641
642    def next_gs_scf_cycle(self):
643        """
644        Return the next :class:`GroundStateScfCycle` in the file. None if not found.
645        """
646        return GroundStateScfCycle.from_stream(self)
647
648    def get_all_gs_scf_cycles(self):
649        """Return list of :class:`GroundStateScfCycle` objects. Empty list if no entry is found."""
650        # NOTE: get_all should not used with next because of the call to self.seek(0)
651        # The API should be refactored
652        cycles = []
653        self.seek(0)
654        while True:
655            cycle = self.next_gs_scf_cycle()
656            if cycle is None: break
657            cycles.append(cycle)
658        self.seek(0)
659        return cycles
660
661    def next_d2de_scf_cycle(self):
662        """
663        Return :class:`D2DEScfCycle` with information on the DFPT iterations. None if not found.
664        """
665        return D2DEScfCycle.from_stream(self)
666
667    def get_all_d2de_scf_cycles(self):
668        """Return list of :class:`D2DEScfCycle` objects. Empty list if no entry is found."""
669        cycles = []
670        self.seek(0)
671        while True:
672            cycle = self.next_d2de_scf_cycle()
673            if cycle is None: break
674            cycles.append(cycle)
675        return cycles
676
677    def plot(self, tight_layout=True, with_timer=False, show=True):
678        """
679        Plot GS/DFPT SCF cycles and timer data found in the output file.
680
681        Args:
682            with_timer: True if timer section should be plotted
683        """
684        from abipy.tools.plotting import MplExpose
685        with MplExpose(slide_mode=False, slide_timeout=5.0) as e:
686            e(self.yield_figs(tight_layout=tight_layout, with_timer=with_timer))
687
688    # TODO: Use header and vars to understand if we have SCF/DFPT/Relaxation
689    def yield_figs(self, **kwargs):  # pragma: no cover
690        """
691        This function *generates* a predefined list of matplotlib figures with minimal input from the user.
692        """
693        tight_layout = kwargs.pop("tight_layout", False)
694        with_timer = kwargs.pop("with_timer", True)
695
696        for icycle, cycle in enumerate(self.get_all_gs_scf_cycles()):
697            yield cycle.plot(title="SCF cycle #%d" % icycle, tight_layout=tight_layout, show=False)
698
699        for icycle, cycle in enumerate(self.get_all_d2de_scf_cycles()):
700            yield cycle.plot(title="DFPT cycle #%d" % icycle, tight_layout=tight_layout, show=False)
701
702        if with_timer:
703            self.seek(0)
704            try:
705                yield self.get_timer().plot_all(tight_layout=tight_layout, show=False)
706            except Exception:
707                print("Abinit output files does not contain timopt data")
708
709    def compare_gs_scf_cycles(self, others, show=True):
710        """
711        Produce and returns a list of matplotlib_ figure comparing the GS self-consistent
712        cycle in self with the ones in others.
713
714        Args:
715            others: list of :class:`AbinitOutputFile` objects or strings with paths to output files.
716            show: True to diplay plots.
717        """
718        # Open file here if we receive a string. Files will be closed before returning
719        close_files = []
720        for i, other in enumerate(others):
721            if is_string(other):
722                others[i] = self.__class__.from_file(other)
723                close_files.append(i)
724
725        fig, figures = None, []
726        while True:
727            cycle = self.next_gs_scf_cycle()
728            if cycle is None: break
729
730            fig = cycle.plot(show=False)
731            for i, other in enumerate(others):
732                other_cycle = other.next_gs_scf_cycle()
733                if other_cycle is None: break
734                last = (i == len(others) - 1)
735                fig = other_cycle.plot(ax_list=fig.axes, show=show and last)
736                if last:
737                    fig.tight_layout()
738                    figures.append(fig)
739
740        self.seek(0)
741        for other in others: other.seek(0)
742
743        if close_files:
744            for i in close_files: others[i].close()
745
746        return figures
747
748    def compare_d2de_scf_cycles(self, others, show=True):
749        """
750        Produce and returns a matplotlib_ figure comparing the DFPT self-consistent
751        cycle in self with the ones in others.
752
753        Args:
754            others: list of :class:`AbinitOutputFile` objects or strings with paths to output files.
755            show: True to diplay plots.
756        """
757        # Open file here if we receive a string. Files will be closed before returning
758        close_files = []
759        for i, other in enumerate(others):
760            if is_string(other):
761                others[i] = self.__class__.from_file(other)
762                close_files.append(i)
763
764        fig, figures = None, []
765        while True:
766            cycle = self.next_d2de_scf_cycle()
767            if cycle is None: break
768
769            fig = cycle.plot(show=False)
770            for i, other in enumerate(others):
771                other_cycle = other.next_d2de_scf_cycle()
772                if other_cycle is None: break
773                last = (i == len(others) - 1)
774                fig = other_cycle.plot(ax_list=fig.axes, show=show and last)
775                if last:
776                    fig.tight_layout()
777                    figures.append(fig)
778
779        self.seek(0)
780        for other in others: other.seek(0)
781
782        if close_files:
783            for i in close_files: others[i].close()
784
785        return figures
786
787    def get_panel(self):
788        """
789        Build panel with widgets to interact with the Abinit output file either in a notebook or in panel app.
790        """
791        from abipy.panels.outputs import AbinitOutputFilePanel
792        return AbinitOutputFilePanel(self).get_panel()
793
794    def write_notebook(self, nbpath=None):
795        """
796        Write a jupyter_ notebook to nbpath. If ``nbpath`` is None, a temporay file in the current
797        working directory is created. Return path to the notebook.
798        """
799        nbformat, nbv, nb = self.get_nbformat_nbv_nb(title=None)
800
801        nb.cells.extend([
802            nbv.new_code_cell("abo = abilab.abiopen('%s')" % self.filepath),
803            nbv.new_code_cell("print(abo.events)"),
804            nbv.new_code_cell("abo.plot()"),
805        ])
806
807        return self._write_nb_nbpath(nb, nbpath)
808
809
810def validate_output_parser(abitests_dir=None, output_files=None):  # pragma: no cover
811    """
812    Validate/test Abinit output parser.
813
814    Args:
815        dirpath: Abinit tests directory.
816        output_files: List of Abinit output files.
817
818    Return: Exit code.
819    """
820    def is_abinit_output(path):
821        """
822        True if path is one of the output files used in the Abinit Test suite.
823        """
824        if not path.endswith(".abo"): return False
825        if not path.endswith(".out"): return False
826
827        with open(path, "rt") as fh:
828            for i, line in enumerate(fh):
829                if i == 1:
830                    return line.rstrip().lower().endswith("abinit")
831            return False
832
833    # Files are collected in paths.
834    paths = []
835
836    if abitests_dir is not None:
837        print("Analyzing directory %s for input files" % abitests_dir)
838
839        for dirpath, dirnames, filenames in os.walk(abitests_dir):
840            for fname in filenames:
841                path = os.path.join(dirpath, fname)
842                if is_abinit_output(path): paths.append(path)
843
844    if output_files is not None:
845        print("Analyzing files:", str(output_files))
846        for arg in output_files:
847            if is_abinit_output(arg): paths.append(arg)
848
849    nfiles = len(paths)
850    if nfiles == 0:
851        cprint("Empty list of input files.", "red")
852        return 0
853
854    print("Found %d Abinit output files" % len(paths))
855    errpaths = []
856    for path in paths:
857        print(path + ": ", end="")
858        try:
859            out = AbinitOutputFile.from_file(path)
860            s = out.to_string(verbose=2)
861            assert out.run_completed
862            cprint("OK", "green")
863        except Exception as exc:
864            if not isinstance(exc, NotImplementedError):
865                cprint("FAILED", "red")
866                errpaths.append(path)
867                import traceback
868                print(traceback.format_exc())
869                #print("[%s]: Exception:\n%s" % (path, str(exc)))
870                #with open(path, "rt") as fh:
871                #    print(10*"=" + "Input File" + 10*"=")
872                #    print(fh.read())
873                #    print()
874            else:
875                cprint("NOTIMPLEMENTED", "magenta")
876
877    if errpaths:
878        cprint("failed: %d/%d [%.1f%%]" % (len(errpaths), nfiles, 100 * len(errpaths)/nfiles), "red")
879        for i, epath in enumerate(errpaths):
880            cprint("[%d] %s" % (i, epath), "red")
881    else:
882        cprint("All input files successfully parsed!", "green")
883
884    return len(errpaths)
885
886
887class AboRobot(Robot):
888    """
889    This robot analyzes the results contained in multiple Abinit output files.
890    Can compare dimensions, SCF cycles, analyze timers.
891
892    .. rubric:: Inheritance Diagram
893    .. inheritance-diagram:: AboRobot
894    """
895    EXT = "abo"
896
897    def get_dims_dataframe(self, with_time=True, index=None):
898        """
899        Build and return |pandas-DataFrame| with the dimensions of the calculation.
900
901        Args:
902            with_time: True if walltime and cputime should be added
903            index: Index of the dataframe. Use relative paths of files if None.
904        """
905        rows, my_index = [], []
906        for i, abo in enumerate(self.abifiles):
907            try:
908                dims_dataset, spg_dataset = abo.get_dims_spginfo_dataset()
909            except Exception as exc:
910                cprint("Exception while trying to get dimensions from %s\n%s" % (abo.relpath, str(exc)), "yellow")
911                continue
912
913            for dtindex, dims in dims_dataset.items():
914                dims = dims.copy()
915                dims.update({"dtset": dtindex})
916                # Add walltime and cputime in seconds
917                if with_time:
918                    dims.update(OrderedDict([(k, getattr(abo, k)) for k in
919                        ("overall_cputime", "proc0_cputime", "overall_walltime", "proc0_walltime")]))
920                rows.append(dims)
921                my_index.append(abo.relpath if index is None else index[i])
922
923        return pd.DataFrame(rows, index=my_index, columns=list(rows[0].keys()))
924
925    def get_dataframe(self, with_geo=True, with_dims=True, abspath=False, funcs=None):
926        """
927        Return a |pandas-DataFrame| with the most important results and the filenames as index.
928
929        Args:
930            with_geo: True if structure info should be added to the dataframe
931            with_dims: True if dimensions should be added
932            abspath: True if paths in index should be absolute. Default: Relative to getcwd().
933            funcs: Function or list of functions to execute to add more data to the DataFrame.
934                Each function receives a |GsrFile| object and returns a tuple (key, value)
935                where key is a string with the name of column and value is the value to be inserted.
936        """
937        rows, row_names = [], []
938        for label, abo in self.items():
939            row_names.append(label)
940            d = OrderedDict()
941
942            if with_dims:
943                dims_dataset, spg_dataset = abo.get_dims_spginfo_dataset()
944                if len(dims_dataset) > 1:
945                    cprint("Multiple datasets are not supported. ARGH!", "yellow")
946                d.update(dims_dataset[1])
947
948            # Add info on structure.
949            if with_geo and abo.run_completed:
950                d.update(abo.final_structure.get_dict4pandas(with_spglib=True))
951
952            # Execute functions
953            if funcs is not None: d.update(self._exec_funcs(funcs, abo))
954            rows.append(d)
955
956        row_names = row_names if not abspath else self._to_relpaths(row_names)
957        return pd.DataFrame(rows, index=row_names, columns=list(rows[0].keys()))
958
959    def get_time_dataframe(self):
960        """
961        Return a |pandas-DataFrame| with the wall-time, cpu time in seconds and the filenames as index.
962        """
963        rows, row_names = [], []
964
965        for label, abo in self.items():
966            row_names.append(label)
967            d = OrderedDict([(k, getattr(abo, k)) for k in
968                ("overall_cputime", "proc0_cputime", "overall_walltime", "proc0_walltime")])
969            rows.append(d)
970
971        return pd.DataFrame(rows, index=row_names, columns=list(rows[0].keys()))
972
973    # TODO
974    #def gridplot_timer(self)
975
976    def yield_figs(self, **kwargs):  # pragma: no cover
977        """
978        This function *generates* a predefined list of matplotlib figures with minimal input from the user.
979        """
980        yield None
981
982    def write_notebook(self, nbpath=None):
983        """
984        Write a jupyter_ notebook to nbpath. If nbpath is None, a temporay file in the current
985        working directory is created. Return path to the notebook.
986        """
987        nbformat, nbv, nb = self.get_nbformat_nbv_nb(title=None)
988
989        args = [(l, f.filepath) for l, f in self.items()]
990        nb.cells.extend([
991            #nbv.new_markdown_cell("# This is a markdown cell"),
992            nbv.new_code_cell("robot = abilab.AboRobot(*%s)\nrobot.trim_paths()\nrobot" % str(args)),
993            nbv.new_code_cell("# robot.get_dims_dataframe()"),
994            nbv.new_code_cell("robot.get_dataframe()"),
995        ])
996
997        # Mixins
998        nb.cells.extend(self.get_baserobot_code_cells())
999
1000        return self._write_nb_nbpath(nb, nbpath)
1001
1002
1003class OutNcFile(AbinitNcFile):
1004    """
1005    Class representing the _OUT.nc file containing the dataset results
1006    produced at the end of the run. The netcdf variables can be accessed
1007    via instance attribute e.g. ``outfile.ecut``. Provides integration with ipython_.
1008    """
1009    # TODO: This object is deprecated
1010    def __init__(self, filepath):
1011        super().__init__(filepath)
1012        self.reader = NetcdfReader(filepath)
1013        self._varscache = {k: None for k in self.reader.rootgrp.variables}
1014
1015    def __dir__(self):
1016        """Ipython integration."""
1017        return sorted(list(self._varscache.keys()))
1018
1019    def __getattribute__(self, name):
1020        try:
1021            return super().__getattribute__(name)
1022        except AttributeError:
1023            # Look in self._varscache
1024            varscache = super().__getattribute__("_varscache")
1025            if name not in varscache:
1026                raise AttributeError("Cannot find attribute %s" % name)
1027            reader = super().__getattribute__("reader")
1028            if varscache[name] is None:
1029                varscache[name] = reader.read_value(name)
1030            return varscache[name]
1031
1032    @lazy_property
1033    def params(self):
1034        """:class:`OrderedDict` with parameters that might be subject to convergence studies."""
1035        return {}
1036
1037    def close(self):
1038        """Close the file."""
1039        self.reader.close()
1040
1041    def get_allvars(self):
1042        """
1043        Read all netcdf_ variables present in the file.
1044        Return dictionary varname --> value
1045        """
1046        for k, v in self._varscache.items():
1047            if v is not None: continue
1048            self._varscache[k] = self.reader.read_value(k)
1049        return self._varscache
1050