1#!/usr/local/bin/python3.8
2"""
3Script to analyze/compare results stored in multiple netcdf/output files.
4By default the script displays the results/plots in the shell.
5Use --ipython to start an ipython terminal or -nb to generate an ipython notebook.
6"""
7
8import sys
9import os
10import argparse
11import numpy as np
12
13from pprint import pprint
14from monty.functools import prof_main
15from monty.termcolor import cprint
16from abipy import abilab
17from abipy.tools.plotting import get_ax_fig_plt, GenericDataFilesPlotter
18
19
20def remove_disordered(structures, paths):
21    """Remove disordered structures and print warning message."""
22    slist = []
23    for s, p in zip(structures, paths):
24        if not s.is_ordered:
25            cprint("Removing disordered structure: %s found in %s" % (s.formula, p), "magenta")
26        else:
27            slist.append(s)
28    return slist
29
30
31def df_to_clipboard(options, df):
32    """Copy dataframe to clipboard if options.clipboard."""
33    if getattr(options, "clipboard", False):
34        cprint("Copying dataframe to the system clipboard.", "green")
35        cprint("This can be pasted into Excel, for example", "green")
36        df.to_clipboard()
37
38
39def abiview_fields(options):
40    """Animate fields with Mayavi. Accept any file with density or potential ..."""
41    from abipy.display.mvtk import MayaviFieldAnimator
42    a = MayaviFieldAnimator(options.filepath)
43    a.volume_animate()
44    return 0
45
46
47def abicomp_structure(options):
48    """
49    Compare crystalline structures. Use `--group` to compare for similarity.
50    """
51    if options.group:
52        return compare_structures(options)
53
54    paths = options.paths
55    index = [os.path.relpath(p) for p in paths]
56
57    if options.notebook:
58        import nbformat
59        nbv = nbformat.v4
60        nb = nbv.new_notebook()
61
62        nb.cells.extend([
63            nbv.new_code_cell("""\
64import sys
65import os
66
67%matplotlib notebook
68from IPython.display import display
69
70from abipy import abilab"""),
71            nbv.new_code_cell("dfs = abilab.dataframes_from_structures(%s, index=%s)" % (paths, index)),
72            # Analyze dataframes.
73            nbv.new_code_cell("dfs.lattice"),
74            nbv.new_code_cell("dfs.coords"),
75            nbv.new_code_cell("# for structure in dfs.structures: display(structure)"),
76        ])
77
78        import io, tempfile # os,
79        _, nbpath = tempfile.mkstemp(prefix="abinb_", suffix='.ipynb', dir=os.getcwd(), text=True)
80
81        # Write notebook
82        import nbformat
83        with io.open(nbpath, 'wt', encoding="utf8") as fh:
84            nbformat.write(nb, fh)
85
86        cmd = "jupyter notebook %s" % nbpath
87        return os.system(cmd)
88
89    dfs = abilab.dataframes_from_structures(paths, index=index,
90            symprec=options.symprec, angle_tolerance=options.angle_tolerance)
91
92    if options.ipython:
93        import IPython
94        IPython.embed(header="Type `dfs` in the terminal and use <TAB> to list its methods", dfs=dfs)
95    else:
96        print("Spglib options: symprec=", options.symprec, "angle_tolerance=", options.angle_tolerance)
97        abilab.print_dataframe(dfs.lattice, title="Lattice parameters:")
98        df_to_clipboard(options, dfs.lattice)
99
100        if options.verbose:
101            abilab.print_dataframe(dfs.coords, title="Atomic positions (columns give the site index):")
102
103    return 0
104
105
106def compare_structures(options):
107    """Inspired to a similar function in pmg_structure."""
108    paths = options.paths
109    if len(paths) < 2:
110        print("You need more than one structure to compare!")
111        return 1
112
113    try:
114        structures = [abilab.Structure.from_file(p) for p in paths]
115    except Exception as ex:
116        print("Error reading structures from files. Are they in the right format?")
117        print(str(ex))
118        return 1
119
120    from pymatgen.analysis.structure_matcher import StructureMatcher, ElementComparator
121    compareby = "species" if options.anonymous else "element"
122    m = StructureMatcher() if compareby == "species" else StructureMatcher(comparator=ElementComparator())
123    print("Grouping %s structures by `%s` with `anonymous: %s`" % (len(structures), compareby, options.anonymous))
124
125    for i, grp in enumerate(m.group_structures(structures, anonymous=options.anonymous)):
126        print("Group {}: ".format(i))
127        for s in grp:
128            spg_symbol, international_number = s.get_space_group_info()
129            print("\t- {} ({}), vol: {:.2f} A^3, {} ({})".format(
130                  paths[structures.index(s)], s.formula, s.volume, spg_symbol, international_number))
131        print()
132
133    if options.verbose:
134        pprint(m.as_dict())
135
136
137def abicomp_spg(options):
138    """
139    Compare the space group found by Abinit with the spglib results
140    for a set of crystalline structure(s) read from FILE(s).
141    """
142    try:
143        structures = [abilab.Structure.from_file(p) for p in options.paths]
144    except Exception as ex:
145        print("Error reading structures from files. Are they in the right format?")
146        print(str(ex))
147        return 1
148
149    # Remove disordered structures.
150    structures = remove_disordered(structures, options.paths)
151
152    rows, index = [], []
153    symprec, angle_tolerance = options.symprec, options.angle_tolerance
154    for structure in structures:
155        index.append(structure.formula)
156        # Call Abinit
157        row = structure.abiget_spginfo(tolsym=options.tolsym, pre="abi_")
158        # Call spglib.
159        spglib_symbol, spglib_number = structure.get_space_group_info(symprec=symprec, angle_tolerance=angle_tolerance)
160        spglib_lattice_type = structure.spget_lattice_type(symprec=symprec, angle_tolerance=angle_tolerance)
161        row.update(spglib_symbol=spglib_symbol, spglib_number=spglib_number, spglib_lattice=spglib_lattice_type)
162        rows.append(row)
163
164    import pandas as pd
165    df = pd.DataFrame(rows, index=index, columns=list(rows[0].keys()) if rows else None)
166
167    print("Spglib options: symprec=", options.symprec, "angle_tolerance=", options.angle_tolerance)
168    print("Abinit options: tolsym=", options.tolsym)
169    print("")
170    abilab.print_dataframe(df, title="Spacegroup found by Abinit and Spglib:")
171    df_to_clipboard(options, df)
172
173    return 0
174
175
176def abicomp_mp_structure(options):
177    """
178    Compare the crystalline structure(s) read from FILE with the one(s)
179    reported in the materials project database.
180    """
181    return _compare_with_database(options)
182
183
184def abicomp_cod_structure(options):
185    """
186    Compare the crystalline structure(s) read from FILE with the one(s)
187    given in the COD database (http://www.crystallography.net/cod).
188    """
189    return _compare_with_database(options)
190
191
192def _compare_with_database(options):
193    structures = [abilab.Structure.from_file(p) for p in options.paths]
194    dbname = {"mp_structure": "materials project", "cod_structure": "COD"}[options.command]
195    if options.command == "mp_structure":
196        mpres = [abilab.mp_search(struct.composition.formula) for struct in structures]
197    elif options.command == "cod_structure":
198        mpres = [abilab.cod_search(struct.composition.formula) for struct in structures]
199    else:
200        raise NotImplementedError(str(options.command))
201
202    # Filter by spglib space group number.
203    if getattr(options, "same_spgnum", False):
204        spgnums = [struct.get_space_group_info()[1] for struct in structures]
205        mpres = [r.filter_by_spgnum(spgnum) for spgnum, r in zip(spgnums, mpres)]
206
207    retcode = 0
208    for this_structure, r in zip(structures, mpres):
209        if r.structures:
210            if options.notebook:
211                new = r.add_entry(this_structure, "this")
212                retcode += new.make_and_open_notebook(foreground=options.foreground,
213                                                      classic_notebook=options.classic_notebook,
214                                                      no_browser=options.no_browser)
215
216            else:
217                print()
218                dfs = abilab.dataframes_from_structures(r.structures + [this_structure], index=r.ids + ["this"])
219                abilab.print_dataframe(dfs.lattice, title="Lattice parameters:", sortby="spglib_num")
220                if options.verbose:
221                    abilab.print_dataframe(dfs.coords, title="Atomic positions (columns give the site index):")
222                else:
223                    print("Use --verbose to print atomic positions.")
224                print()
225
226        else:
227            print("Couldn't find %s database entries with formula `%s`" % (dbname, this_structure.composition.formula))
228            retcode += 1
229
230    return retcode
231
232
233def abicomp_xrd(options):
234    """
235    Compare X-ray diffraction plots (requires FILES with structure).
236    """
237    if len(options.paths) < 2:
238        print("You need more than one structure to compare!")
239        return 1
240
241    structures = [abilab.Structure.from_file(p) for p in options.paths]
242
243    dfs = abilab.dataframes_from_structures(structures, index=[os.path.relpath(p) for p in options.paths])
244    abilab.print_dataframe(dfs.lattice, title="Lattice parameters:")
245    if options.verbose:
246        abilab.print_dataframe(dfs.coords, title="Atomic positions (columns give the site index):")
247
248    from pymatgen.analysis.diffraction.xrd import XRDCalculator
249    two_theta_range = tuple(float(t) for t in options.two_theta_range)
250    xrd = XRDCalculator(wavelength=options.wavelength, symprec=options.symprec)
251    xrd.plot_structures(structures, two_theta_range=two_theta_range, fontsize=6,
252                        annotate_peaks=not options.no_annotate_peaks, tight_layout=True)
253    return 0
254
255
256def abicomp_data(options):
257    """
258    Compare results stored in multiple files with data in tabular format.
259    """
260    plotter = GenericDataFilesPlotter.from_files(options.paths)
261    print(plotter.to_string(verbose=options.verbose))
262    plotter.plot(use_index=options.use_index)
263    return 0
264
265
266def abicomp_ebands(options):
267    """
268    Plot electron bands on a grid.
269    """
270    paths, e0 = options.paths, options.e0
271    plotter = abilab.ElectronBandsPlotter(key_ebands=[(os.path.relpath(p), p) for p in paths])
272
273    if options.ipython:
274        import IPython
275        IPython.embed(header=str(plotter) + "\n\nType `plotter` in the terminal and use <TAB> to list its methods",
276                      plotter=plotter)
277
278    elif options.notebook:
279        plotter.make_and_open_notebook(foreground=options.foreground,
280                                       classic_notebook=options.classic_notebook,
281                                       no_browser=options.no_browser)
282    else:
283        # Print pandas Dataframe.
284        df = plotter.get_ebands_frame()
285        abilab.print_dataframe(df)
286        df_to_clipboard(options, df)
287
288        # Optionally, print info on gaps and their location
289        if not options.verbose:
290            print("\nUse --verbose for more information")
291        else:
292            for ebands in plotter.ebands_list:
293                print(ebands)
294
295        # Here I select the plot method to call.
296        if options.plot_mode != "None":
297            plotfunc = getattr(plotter, options.plot_mode, None)
298            if plotfunc is None:
299                raise ValueError("Don't know how to handle plot_mode: %s" % options.plot_mode)
300            plotfunc(e0=e0)
301
302    return 0
303
304
305def abicomp_edos(options):
306    """
307    Plot electron DOSes on a grid.
308    """
309    paths, e0 = options.paths, options.e0
310    plotter = abilab.ElectronDosPlotter(key_edos=[(os.path.relpath(p), p) for p in paths])
311
312    if options.ipython:
313        import IPython
314        IPython.embed(header=str(plotter) + "\n\nType `plotter` in the terminal and use <TAB> to list its methods",
315                      plotter=plotter)
316
317    elif options.notebook:
318        plotter.make_and_open_notebook(foreground=options.foreground,
319                                       classic_notebook=options.classic_notebook,
320                                       no_browser=options.no_browser)
321
322    elif options.expose:
323        plotter.expose(slide_mode=options.slide_mode, slide_timeout=options.slide_timeout,
324                       verbose=options.verbose)
325
326    else:
327        # Optionally, print info on gaps and their location
328        if not options.verbose:
329            print("\nUse --verbose for more information")
330        else:
331            for edos in plotter.edos_list:
332                print(edos)
333
334        # Here I select the plot method to call.
335        if options.plot_mode != "None":
336            plotfunc = getattr(plotter, options.plot_mode, None)
337            if plotfunc is None:
338                raise ValueError("Don't know how to handle plot_mode: %s" % options.plot_mode)
339            plotfunc(e0=e0)
340
341    return 0
342
343
344def abicomp_phbands(options):
345    """
346    Plot phonon bands on a grid.
347    """
348    paths = options.paths
349    plotter = abilab.PhononBandsPlotter(key_phbands=[(os.path.relpath(p), p) for p in paths])
350
351    if options.ipython:
352        import IPython
353        IPython.embed(header=str(plotter) + "\n\nType `plotter` in the terminal and use <TAB> to list its methods",
354                      plotter=plotter)
355
356    elif options.notebook:
357        plotter.make_and_open_notebook(foreground=options.foreground,
358                                       classic_notebook=options.classic_notebook,
359                                       no_browser=options.no_browser)
360
361    elif options.expose:
362        plotter.expose(slide_mode=options.slide_mode, slide_timeout=options.slide_timeout,
363                       verbose=options.verbose)
364    else:
365        # Print pandas Dataframe.
366        abilab.print_dataframe(plotter.get_phbands_frame())
367
368        # Optionally, print info on gaps and their location
369        if not options.verbose:
370            print("\nUse --verbose for more information")
371        else:
372            for phbands in plotter.phbands_list:
373                print(phbands)
374
375        # Select the plot method to call.
376        if options.plot_mode == "panel":
377            plotter.get_panel().show()
378
379        elif options.plot_mode != "None":
380            plotfunc = getattr(plotter, options.plot_mode, None)
381            if plotfunc is None:
382                raise ValueError("Don't know how to handle plot_mode: %s" % options.plot_mode)
383            plotfunc()
384
385    return 0
386
387
388def abicomp_phdos(options):
389    """
390    Compare multiple PHDOS files.
391    """
392    paths = options.paths
393    plotter = abilab.PhononDosPlotter(key_phdos=[(os.path.relpath(p), p) for p in paths])
394
395    if options.ipython:
396        import IPython
397        IPython.embed(header=str(plotter) + "\n\nType `plotter` in the terminal and use <TAB> to list its methods",
398                      plotter=plotter)
399
400    elif options.expose:
401        plotter.expose(slide_mode=options.slide_mode, slide_timeout=options.slide_timeout,
402                       verbose=options.verbose)
403
404    elif options.notebook:
405        plotter.make_and_open_notebook(foreground=options.foreground,
406                                       classic_notebook=options.classic_notebook,
407                                       no_browser=options.no_browser)
408
409    else:
410        # Optionally, print info on gaps and their location
411        if not options.verbose:
412            print("\nUse --verbose for more information")
413        else:
414            for phdos in plotter.phdos_list:
415                print(phdos)
416
417        # Here I select the plot method to call.
418        if options.plot_mode != "None":
419            plotfunc = getattr(plotter, options.plot_mode, None)
420            if plotfunc is None:
421                raise ValueError("Don't know how to handle plot_mode: %s" % options.plot_mode)
422            plotfunc()
423
424    return 0
425
426
427def abicomp_getattr(options):
428    """
429    Extract attribute from Abipy object for all files listed on the command line and print them.
430    Use `--list` option to list the attributes available (in the first file).
431    """
432    files = []
433    attr_name = options.paths[0]
434    values = []
435    for p in options.paths[1:]:
436        with abilab.abiopen(p) as abifile:
437            if options.list:
438                print("List of attributes available in %s" % p)
439                pprint(dir(abifile))
440                return 0
441
442            v = getattr(abifile, attr_name)
443            print(v, "   # File: ", p)
444            if options.plot:
445                try:
446                    values.append(float(v))
447                except TypeError as exc:
448                    print("Cannot plot data. Exception:\n", str(exc))
449
450    if options.plot and len(values) == len(options.paths[1:]):
451        # Plot values.
452        ax, fig, plt = get_ax_fig_plt()
453        xs = np.arange(len(options.paths[1:]))
454        ax.plot(xs, values)
455        ax.set_ylabel(attr_name)
456        ax.set_xticks(xs)
457        xlabels = options.paths[1:]
458        s = set((os.path.basename(s) for s in xlabels))
459        if len(s) == len(xlabels): xlabels = s
460        ax.set_xticklabels(xlabels) #, rotation='vertical')
461        plt.show()
462
463    return 0
464
465
466##################
467# Robot commands #
468##################
469
470def abicomp_gsr(options):
471    """
472    Compare multiple GSR files.
473    """
474    return _invoke_robot(options)
475
476
477def abicomp_hist(options):
478    """
479    Compare multiple HIST files.
480    """
481    return _invoke_robot(options)
482
483
484def abicomp_ddb(options):
485    """
486    Compare multiple DDB files. Assume DDB files with a list of q-points in the IBZ
487    corresponding to homogeneous sampling i.e. files that have been merged with mrgddb.
488    """
489    return _invoke_robot(options)
490
491
492def abicomp_anaddb(options):
493    """
494    Compare multiple anaddb.nc files.
495    """
496    return _invoke_robot(options)
497
498
499def abicomp_phbst(options):
500    """
501    Compare multiple PHBST.nc files.
502    """
503    return _invoke_robot(options)
504
505
506def abicomp_sigres(options):
507    """
508    Compare multiple SIGRES files.
509    """
510    return _invoke_robot(options)
511
512
513def abicomp_mdf(options):
514    """
515    Compare macroscopic dielectric functions stored in multiple MDF files.
516    """
517    return _invoke_robot(options)
518
519
520def abicomp_optic(options):
521    """
522    Compare results stored in OPTIC.nc files.
523    """
524    return _invoke_robot(options)
525
526
527def abicomp_a2f(options):
528    """
529    Compare results stored in A2f.nc files.
530    """
531    return _invoke_robot(options)
532
533
534def abicomp_gkq(options):
535    """
536    Compare multiple GKQ files with EPH matrix elements for a given q-point.
537    """
538    if options.diff:
539        robot = _build_robot(options, trim_paths=True)
540        robot.plot_gkq2_diff()
541    else:
542        return _invoke_robot(options)
543
544
545#def abicomp_wrmax(options):
546#    """
547#    Compare multiple WRmax files with first order potential in real-space.
548#    """
549#    return _invoke_robot(options)
550
551
552def abicomp_v1qavg(options):
553    """
554    Compare multiple V1QAVG files with the average of the DFPT V1 potentials as function of q-point.
555    """
556    return _invoke_robot(options)
557
558
559def abicomp_sigeph(options):
560    """
561    Compare multiple SIGEPH files storing the e-ph self-energy.
562    """
563    return _invoke_robot(options)
564
565
566def abicomp_rta(options):
567    """
568    Compare multiple RTA files.
569    """
570    return _invoke_robot(options)
571
572
573def abicomp_abiwan(options):
574    """
575    Compare multiple ABIWAN files.
576    """
577    return _invoke_robot(options)
578
579
580def abicomp_pseudos(options):
581    """"Compare multiple pseudos Print table to terminal."""
582    # Make sure entries in index are unique.
583    index = [os.path.basename(p) for p in options.paths]
584    if len(index) != len(set(index)): index = [os.path.relpath(p) for p in options.paths]
585    from abipy.electrons.psps import dataframe_from_pseudos
586    df = dataframe_from_pseudos(options.paths, index=index)
587    abilab.print_dataframe(df, sortby="Z_val")
588    return 0
589
590
591def _build_robot(options, trim_paths=False):
592    """Build robot instance from CLI options."""
593    robot_cls = abilab.Robot.class_for_ext(options.command.upper())
594
595    # To define an Help action
596    # http://stackoverflow.com/questions/20094215/argparse-subparser-monolithic-help-output?rq=1
597    paths = options.paths
598    if options.verbose > 1:
599        print("In _invoke_robot with paths", paths)
600
601    if os.path.isdir(paths[0]):
602        # Assume directory.
603        robot = robot_cls.from_dir(top=paths[0], walk=not options.no_walk)
604    else:
605        # Assume file.
606        robot = robot_cls.from_files([paths[0]])
607
608    if len(paths) > 1:
609        # Handle multiple arguments. Could be either other directories or files.
610        for p in paths[1:]:
611            if os.path.isdir(p):
612                robot.scan_dir(top=p, walk=not options.no_walk)
613            elif os.path.isfile(p):
614                robot.add_file(os.path.abspath(p), p)
615            else:
616                cprint("Ignoring %s. Neither file or directory." % str(p), "red")
617
618    if len(robot) == 0:
619        raise RuntimeError("Empty robot --> No file associated to this robot has been found")
620
621    if trim_paths: robot.trim_paths()
622    return robot
623
624
625def _invoke_robot(options):
626    """
627    Analyze multiple files with a robot. Support list of files and/or list of directories passed on the CLI.
628
629    By default, the script with call `robot.to_string(options.verbose)` to print info to terminal.
630    For finer control, use --ipy to start an ipython console to interact with the robot directly
631    or --nb to generate a jupyter notebook.
632    """
633    robot = _build_robot(options)
634
635    if options.notebook:
636        robot.make_and_open_notebook(foreground=options.foreground,
637                                     classic_notebook=options.classic_notebook,
638                                     no_browser=options.no_browser)
639
640    elif options.panel:
641        try:
642            import panel  # noqa: F401
643        except ImportError as exc:
644            cprint("Use `conda install panel` or `pip install panel` to install the python package.", "red")
645            raise exc
646
647        robot.get_panel().show()
648        return 0
649
650    elif options.print or options.expose:
651        robot.trim_paths()
652        # Print dataframe if robot provides get_dataframe method.
653        if hasattr(robot, "get_dataframe"):
654            try:
655                df = robot.get_dataframe()
656                abilab.print_dataframe(df, title="Output of robot.get_dataframe():")
657                df_to_clipboard(options, df)
658
659            except Exception as exc:
660                cprint("Exception:\n%s\n\nwhile invoking get_dataframe. Falling back to to_string" % str(exc), "red")
661                print(robot.to_string(verbose=options.verbose))
662
663        else:
664            cprint("%s does not provide `get_dataframe` method. Using `to_string`" % (
665                    robot.__class__.__name__), "yellow")
666            print(robot.to_string(verbose=options.verbose))
667
668        if not options.verbose:
669            print("\nUse --verbose for more information")
670
671        if options.expose and hasattr(robot, "expose"):
672            robot.expose(slide_mode=options.slide_mode, slide_timeout=options.slide_timeout,
673                         verbose=options.verbose)
674
675    #elif options.ipython:
676    else:
677        import IPython
678        robot.trim_paths()
679        IPython.embed(header=repr(robot) + "\n\nType `robot` in the terminal and use <TAB> to list its methods",
680                      robot=robot)
681
682    return 0
683
684
685def abicomp_gs_scf(options):
686    """
687    Compare ground-state SCF cycles.
688    """
689    paths = options.paths
690    f0 = abilab.AbinitOutputFile(paths[0])
691    figures = f0.compare_gs_scf_cycles(paths[1:])
692    if not figures:
693        cprint("Cannot find GS-SCF sections in output files.", "yellow")
694    return 0
695
696
697def abicomp_dfpt2_scf(options):
698    """
699    Compare DFPT SCF cycles.
700    """
701    paths = options.paths
702    f0 = abilab.AbinitOutputFile(paths[0])
703    figures = f0.compare_d2de_scf_cycles(paths[1:])
704    if not figures:
705        cprint("Cannot find DFPT-SCF sections in output files.", "yellow")
706    return 0
707
708
709def abicomp_text(options):
710    """
711    Compare 2+ text files in the browser
712    """
713    from abipy.tools.devtools import HtmlDiff
714    return HtmlDiff(options.paths).open_browser(diffmode=options.diffmode)
715
716
717def abicomp_time(options):
718    """
719    Analyze/plot the timing data of single or multiple runs.
720    """
721    paths = options.paths
722    from abipy.abio.timer import AbinitTimerParser
723
724    if len(options.paths) == 1 and os.path.isdir(paths[0]):
725        # Scan directory tree
726        top = options.paths[0]
727        print("Walking directory tree from top:", top, "Looking for file extension:", options.ext)
728        parser, paths_found, okfiles = AbinitTimerParser.walk(top=top, ext=options.ext)
729
730        if not paths_found:
731            cprint("Empty file list!", color="magenta")
732            return 1
733
734        print("Found %d files\n" % len(paths_found))
735        if okfiles != paths_found:
736            badfiles = [f for f in paths_found if f not in okfiles]
737            cprint("Cannot parse timing data from the following files:", color="magenta")
738            for bad in badfiles: print(bad)
739
740    else:
741        # Parse list of files.
742        parser = AbinitTimerParser()
743        okfiles = parser.parse(options.paths)
744
745        if okfiles != options.paths:
746            badfiles = [f for f in options.paths if f not in okfiles]
747            cprint("Cannot parse timing data from the following files:", color="magenta")
748            for bad in badfiles: print(bad)
749
750    if parser is None:
751        cprint("Cannot analyze timing data. parser is None", color="magenta")
752        return 1
753
754    print(parser.summarize())
755
756    if options.verbose:
757        for timer in parser.timers():
758            print(timer.get_dataframe())
759
760    if options.ipython:
761        cprint("Invoking ipython shell. Use parser to access the object inside ipython", color="blue")
762        import IPython
763        IPython.start_ipython(argv=[], user_ns={"parser": parser})
764    elif options.notebook:
765        parser.make_and_open_notebook(foreground=options.foreground,
766                                      classic_notebook=options.classic_notebook,
767                                      no_browser=options.no_browser)
768    else:
769        parser.plot_all()
770
771    return 0
772
773
774def get_epilog():
775    return """\
776Usage example:
777
778############
779# Structures
780############
781
782  abicomp.py structure */*/outdata/out_GSR.nc   => Compare structures in multiple files.
783                                                   Use `--group` to compare for similarity.
784  abicomp.py spg si.cif out_GSR.nc              => Compare spacegroup(s) found by Abinit and spglib.
785  abicomp.py hist FILE(s)                       => Compare final structures read from HIST.nc files.
786  abicomp.py mp_structure FILE(s)               => Compare structure(s) read from FILE(s) with the one(s)
787                                                   given in the materials project database.
788  abicomp.py cod_structure FILE(s)              => Compare structure(s) read from FILE(s) with the one(s)
789                                                   given in the COD database (http://www.crystallography.net/cod).
790  abicomp.py xrd *.cif *.GSR.nc                 => Compare X-ray diffraction plots (requires FILES with structure).
791
792###########
793# Electrons
794###########
795
796  abicomp.py ebands out1_GSR.nc out2_WFK.nc     => Plot electron bands on a grid (Use `-p` to change plot mode)
797  abicomp.py ebands *_GSR.nc -ipy               => Build plotter object and start ipython console.
798  abicomp.py ebands *_GSR.nc -nb                => Interact with the plotter in the jupyter notebook.
799  abicomp.py ebands `find . -name "*_GSR.nc"` -c = Find all GSR.nc files startign from current working directory
800                                                   Copy dataframe to the system clipboard.
801                                                   This can be pasted into Excel, for example
802  abicomp.py edos *_WFK.nc -nb                  => Compare electron DOS in the jupyter notebook.
803  abicomp.py optic DIR -nb                      => Compare optic results in the jupyter notebook.
804  abicomp.py abiwan *_ABIWAN.nc --expose        => Compare ABIWAN results, produce matplotlib figures.
805
806#########
807# Phonons
808#########
809
810  abicomp.py phbands *_PHBST.nc -p gridplot     => Compare phonon bands with matplotlib grid.
811  abicomp.py phbands *_PHBST.nc -p panel        => Compare phonon bands in the panel dashboard (GUI)
812  abicomp.py phbst *_PHBST.nc -ipy              => Compare phonon bands with robot in ipython terminal.
813  abicomp.py phdos *_PHDOS.nc -nb               => Compare phonon DOSes in the jupyter notebook.
814  abicomp.py ddb outdir1 outdir2 out_DDB -nb    => Analyze all DDB files in directories outdir1, outdir2 and out_DDB file.
815
816###############
817# Anaddb netcdf
818###############
819
820  abicomp anaddb tutorespfn_telast_2-telast_3/anaddb.nc
821
822#########
823# E-PH
824#########
825
826  abicomp.py a2f *_A2F.nc -nb                   => Compare A2f results in the jupyter notebook.
827  abicomp.py sigeph *_SIGEPH.nc -nb             => Compare Fan-Migdal self-energy in the jupyter notebook.
828  abicomp.py gkq out1_GKQ.nc out1_GKQ.nc -d     => Plot difference between matrix elements (supports 2+ files).
829  abicomp.py v1qavg out_V1QAVG.nc               => Compare V1QAVG files.
830
831########
832# GW/BSE
833########
834
835  abicomp.py sigres *_SIGRES.nc                 => Compare multiple SIGRES files.
836  abicomp.py mdf *_MDF.nc --seaborn             => Compare macroscopic dielectric functions.
837                                                   Use seaborn settings.
838
839###############
840# Miscelleanous
841###############
842
843  abicomp.py data FILE1 FILE2 ...                 => Read data from files with results in tabular format and
844                                                     compare results. Mainly used for text files without any schema.
845  abicomp.py getattr energy *_GSR.nc              => Extract the `energy` attribute from a list of GSR files
846                                                     and print results. Use `--list` to get list of possible names.
847  abicomp.py pseudos PSEUDO_FILES                 => Compare pseudopotential files.
848
849############
850# Text files
851############
852
853  abicomp.py gs_scf run1.abo run2.abo             => Compare the SCF cycles in two output files.
854  abicomp.py dfpt2_scf run1.abo run2.abo          => Compare the DFPT SCF cycles in two output files.
855  abicomp.py.py time [OUT_FILES]                  => Parse timing data in files and plot results
856  abicomp.py.py time . --ext=abo                  => Scan directory tree from `.`, look for files with extension `abo`
857                                                     parse timing data and plot results.
858  abicomp.py text run1.abo run2.abo               => Produce diff of 2+ text files in the browser.
859
860TIP:
861
862The python code operates on a list of files/directories passed via the command line interface.
863The arguments are interpreted by the shell before invoking the script.
864This means that one can use the bash syntax and the unix tools to precompute the list of files/directories.
865
866For example, one can use Unix find to select all files with the a given extension and pass them to abicomp.py.
867For command:
868
869    abicomp.py structure `find . -name "*_GSR.nc"`
870
871will compare the structures extracted from all GSR.nc files found within the current working directory (note backticks).
872
873Also, remember that in bash {#..#} generates a sequence of numbers or chars, similarly to range() in Python
874For instance:
875
876    {1..5} --> 1 2 3 4 5
877
878and this trick can be used to select files/directories as in:
879
880    abicomp.py structure w1/t{2..4}/outdata/*_GSR.nc
881
882The [!name] syntax can be used to exclude patterns, so
883
884    abicomp.py structure w1/t[!2]*/outdata/*_GSR.nc
885
886excludes all the GSR.nc files in the t2/outdata directory.
887
888See also http://wiki.bash-hackers.org/syntax/pattern
889
890NOTE: The `gsr`, `ddb`, `sigres`, `mdf` commands use robots to analyze files.
891In this case, one can provide a list of files and/or list of directories on the command-line interface e.g.:
892
893    $ abicomp.py ddb dir1 out_DDB dir2
894
895Directories will be scanned recursively to find files with the extension associated to the robot, e.g.
896`abicompy.py mdf .` will read all *_MDF.nc files inside the current directory including sub-directories (if any).
897Use --no-walk to ignore sub-directories when robots are used.
898
899Use `abicomp.py --help` for help and `abicomp.py COMMAND --help` to get the documentation for `COMMAND`.
900Use `-v` to increase verbosity level (can be supplied multiple times e.g -vv).
901"""
902
903
904def get_parser(with_epilog=False):
905
906    # Parent parser for common options.
907    copts_parser = argparse.ArgumentParser(add_help=False)
908    copts_parser.add_argument('paths', nargs="+", help="List of files to compare.")
909    copts_parser.add_argument('-v', '--verbose', default=0, action='count', # -vv --> verbose=2
910        help='Verbose, can be supplied multiple times to increase verbosity.')
911    copts_parser.add_argument('-sns', "--seaborn", const="paper", default=None, action='store', nargs='?', type=str,
912        help='Use seaborn settings. Accept value defining context in ("paper", "notebook", "talk", "poster"). Default: paper')
913    copts_parser.add_argument('-mpl', "--mpl-backend", default=None,
914        help=("Set matplotlib interactive backend. "
915              "Possible values: GTKAgg, GTK3Agg, GTK, GTKCairo, GTK3Cairo, WXAgg, WX, TkAgg, Qt4Agg, Qt5Agg, macosx."
916              "See also: https://matplotlib.org/faq/usage_faq.html#what-is-a-backend."))
917    copts_parser.add_argument('--loglevel', default="ERROR", type=str,
918        help="Set the loglevel. Possible values: CRITICAL, ERROR (default), WARNING, INFO, DEBUG.")
919
920    # Parent parser for commands calling spglib.
921    spgopt_parser = argparse.ArgumentParser(add_help=False)
922    spgopt_parser.add_argument('--symprec', default=1e-3, type=float,
923        help="""\
924symprec (float): Tolerance for symmetry finding. Defaults to 1e-3,
925which is fairly strict and works well for properly refined structures with atoms in the proper symmetry coordinates.
926For structures with slight deviations from their proper atomic positions (e.g., structures relaxed with electronic structure
927codes), a looser tolerance of 0.1 (the value used in Materials Project) is often needed.""")
928    spgopt_parser.add_argument('--angle-tolerance', default=5.0, type=float,
929        help="angle_tolerance (float): Angle tolerance for symmetry finding. Default: 5.0")
930    #spgopt_parser.add_argument("--no-time-reversal", default=False, action="store_true", help="Don't use time-reversal.")
931
932    # Parent parser for commands operating on pandas dataframes
933    pandas_parser = argparse.ArgumentParser(add_help=False)
934    pandas_parser.add_argument("-c", '--clipboard', default=False, action="store_true",
935            help="Copy dataframe to the system clipboard. This can be pasted into Excel, for example")
936
937    # Parent parser for commands supporting (ipython/jupyter)
938    ipy_parser = argparse.ArgumentParser(add_help=False)
939    ipy_parser.add_argument('-nb', '--notebook', default=False, action="store_true", help='Generate jupyter notebook.')
940    ipy_parser.add_argument('--classic-notebook', action='store_true', default=False,
941                            help="Use classic notebook instead of jupyterlab.")
942    ipy_parser.add_argument('--no-browser', action='store_true', default=False,
943                            help=("Start the jupyter server to serve the notebook "
944                                  "but don't open the notebook in the browser.\n"
945                                  "Use this option to connect remotely from localhost to the machine running the kernel"))
946    ipy_parser.add_argument('--foreground', action='store_true', default=False,
947        help="Run jupyter notebook in the foreground.")
948    ipy_parser.add_argument('-ipy', '--ipython', default=False, action="store_true", help='Invoke ipython terminal.')
949
950    # Parent parser for commands supporting (jupyter notebooks)
951    nb_parser = argparse.ArgumentParser(add_help=False)
952    nb_parser.add_argument('-nb', '--notebook', default=False, action="store_true", help='Generate jupyter notebook.')
953    nb_parser.add_argument('--foreground', action='store_true', default=False,
954        help="Run jupyter notebook in the foreground.")
955
956    # Parent parser for commands supporting expose
957    expose_parser = argparse.ArgumentParser(add_help=False)
958    expose_parser.add_argument("-e", '--expose', default=False, action="store_true",
959            help='Execute robot.expose to produce a pre-defined list of matplotlib figures.')
960    expose_parser.add_argument("-s", "--slide-mode", default=False, action="store_true",
961            help="Used if --expose to iterate over figures. Expose all figures at once if not given on the CLI.")
962    expose_parser.add_argument("-t", "--slide-timeout", type=int, default=None,
963            help="Close figure after slide-timeout seconds (only if slide-mode). Block if not specified.")
964
965    # Build the main parser.
966    parser = argparse.ArgumentParser(epilog=get_epilog() if with_epilog else "",
967                                     formatter_class=argparse.RawDescriptionHelpFormatter)
968    parser.add_argument('-V', '--version', action='version', version=abilab.__version__)
969
970    # Create the parsers for the sub-commands
971    subparsers = parser.add_subparsers(dest='command', help='sub-command help', description="Valid subcommands")
972
973    # Subparser for structure command.
974    p_struct = subparsers.add_parser('structure', parents=[copts_parser, ipy_parser, spgopt_parser, pandas_parser],
975            help=abicomp_structure.__doc__)
976    p_struct.add_argument("-g", "--group", default=False, action="store_true",
977        help="Compare a set of structures for similarity.")
978    p_struct.add_argument("-a", "--anonymous", default=False, action="store_true",
979        help="Whether to use anonymous mode in StructureMatcher. Default False")
980
981    # Subparser for spg command.
982    p_spg = subparsers.add_parser('spg', parents=[copts_parser, spgopt_parser, pandas_parser],
983            help=abicomp_spg.__doc__)
984    p_spg.add_argument("-t", "--tolsym", type=float, default=None, help="""\
985Gives the tolerance on the atomic positions (reduced coordinates), primitive vectors, or magnetization,
986to be considered equivalent, thanks to symmetry operations. This is used in the recognition of the set
987of symmetries of the system, or the application of the symmetry operations to generate from a reduced set of atoms,
988the full set of atoms. Note that a value larger than 0.01 is considered to be unacceptable.""")
989
990    # Subparser for mp_structure command.
991    p_mpstruct = subparsers.add_parser('mp_structure', parents=[copts_parser, nb_parser],
992        help=abicomp_mp_structure.__doc__)
993    p_mpstruct.add_argument("--same-spgnum", default=False, action="store_true",
994        help="Select only MP structures with same space group number as input structure.")
995
996    # Subparser for cod_structure command.
997    p_codstruct = subparsers.add_parser('cod_structure', parents=[copts_parser, nb_parser],
998        help=abicomp_cod_structure.__doc__)
999    #p_codstruct.add_argument("--same-spgnum", default=False, action="store_true",
1000    #    help="Select only COD structures with same space group number as input structure.")
1001
1002    # Subparser for xrd.
1003    p_xrd = subparsers.add_parser('xrd', parents=[copts_parser], help="Compare X-ray diffraction plots.")
1004    p_xrd.add_argument("-w", "--wavelength", default="CuKa", type=str, help=(
1005        "The wavelength can be specified as a string. It must be one of the "
1006        "supported definitions in the WAVELENGTHS dict declared in pymatgen/analysis/diffraction/xrd.py."
1007        "Defaults to 'CuKa', i.e, Cu K_alpha radiation."))
1008    p_xrd.add_argument("-s", "--symprec", default=0, type=float, help=(
1009        "Symmetry precision for structure refinement. "
1010        "If set to 0, no refinement is done. Otherwise, refinement is performed using spglib with provided precision."))
1011    p_xrd.add_argument("-t", "--two-theta-range", default=(0, 90), nargs=2, help=(
1012        "Tuple for range of two_thetas to calculate in degrees. Defaults to (0, 90)."))
1013    p_xrd.add_argument("-nap", "--no-annotate-peaks", default=False, action="store_true",
1014        help="Whether to annotate the peaks with plane information.")
1015
1016    # Subparser for data command.
1017    p_data = subparsers.add_parser('data', parents=[copts_parser, expose_parser], help=abicomp_data.__doc__)
1018    p_data.add_argument("-i", "--use-index", default=False, action="store_true",
1019        help="Use the row index as x-value in the plot. By default the plotter uses the first column as x-values")
1020
1021    # Subparser for ebands command.
1022    p_ebands = subparsers.add_parser('ebands', parents=[copts_parser, ipy_parser, pandas_parser],
1023            help=abicomp_ebands.__doc__)
1024    p_ebands.add_argument("-p", "--plot-mode", default="gridplot",
1025        choices=["gridplot", "combiplot", "boxplot", "combiboxplot", "plot_band_edges", "animate", "None"],
1026        help="Plot mode e.g. `-p combiplot` to plot bands on the same figure. Default is `gridplot`.")
1027    p_ebands.add_argument("-e0", default="fermie", choices=["fermie", "None"],
1028        help="Option used to define the zero of energy in the band structure plot. Default is `fermie`.")
1029
1030    # Subparser for edos command.
1031    p_edos = subparsers.add_parser('edos', parents=[copts_parser, ipy_parser, expose_parser],
1032        help=abicomp_edos.__doc__)
1033    p_edos.add_argument("-p", "--plot-mode", default="gridplot",
1034        choices=["gridplot", "combiplot", "None"],
1035        help="Plot mode e.g. `-p combiplot` to plot DOSes on the same figure. Default is `gridplot`.")
1036    p_edos.add_argument("-e0", default="fermie", choices=["fermie", "None"],
1037        help="Option used to define the zero of energy in the DOS plot. Default is `fermie`.")
1038
1039    # Subparser for phbands command.
1040    p_phbands = subparsers.add_parser('phbands', parents=[copts_parser, ipy_parser, expose_parser],
1041        help=abicomp_phbands.__doc__)
1042    p_phbands.add_argument("-p", "--plot-mode", default="gridplot",
1043        choices=["gridplot", "combiplot", "boxplot", "combiboxplot", "animate", "panel", "None"],
1044        help="Plot mode e.g. `-p combiplot` to plot bands on the same figure."
1045             "Use `panel` for GUI in web browser. Default is `gridplot`.")
1046
1047    # Subparser for phdos command.
1048    p_phdos = subparsers.add_parser('phdos', parents=[copts_parser, ipy_parser, expose_parser],
1049        help=abicomp_phdos.__doc__)
1050    p_phdos.add_argument("-p", "--plot-mode", default="gridplot",
1051        choices=["gridplot", "combiplot", "None"],
1052        help="Plot mode e.g. `-p combiplot` to plot DOSes on the same figure. Default is `gridplot`.")
1053
1054    # Subparser for getattr command.
1055    p_getattr = subparsers.add_parser('getattr', parents=[copts_parser], help=abicomp_getattr.__doc__)
1056    p_getattr.add_argument('--plot', default=False, action="store_true", help="Plot data with matplotlib (requires floats).")
1057    p_getattr.add_argument('--list', default=False, action="store_true", help="Print attributes available in file")
1058
1059    # Subparser for robot commands
1060    # Use own version of ipy_parser with different default values.
1061    robot_ipy_parser = argparse.ArgumentParser(add_help=False)
1062    robot_ipy_parser.add_argument('-nb', '--notebook', default=False, action="store_true", help='Generate jupyter notebook.')
1063    robot_ipy_parser.add_argument('--foreground', action='store_true', default=False,
1064        help="Run jupyter notebook in the foreground.")
1065    #robot_ipy_parser.add_argument('-ipy', '--ipython', default=True, action="store_true", help='Invoke ipython terminal.')
1066    robot_ipy_parser.add_argument('-p', '--print', default=False, action="store_true", help='Print robot and return.')
1067
1068    # Parent parser for *robot* commands
1069    robot_parser = argparse.ArgumentParser(add_help=False)
1070    robot_parser.add_argument('--no-walk', default=False, action="store_true", help="Don't enter subdirectories.")
1071    robot_parser.add_argument('--panel', default=False, action="store_true",
1072                              help="Open GUI in web browser, requires panel package. WARNING: Experimental")
1073
1074    robot_parents = [copts_parser, robot_ipy_parser, robot_parser, expose_parser, pandas_parser]
1075    p_gsr = subparsers.add_parser('gsr', parents=robot_parents, help=abicomp_gsr.__doc__)
1076    p_hist = subparsers.add_parser('hist', parents=robot_parents, help=abicomp_hist.__doc__)
1077    p_ddb = subparsers.add_parser('ddb', parents=robot_parents, help=abicomp_ddb.__doc__)
1078    p_anaddb = subparsers.add_parser('anaddb', parents=robot_parents, help=abicomp_anaddb.__doc__)
1079    p_phbst = subparsers.add_parser('phbst', parents=robot_parents, help=abicomp_phbst.__doc__)
1080    p_sigres = subparsers.add_parser('sigres', parents=robot_parents, help=abicomp_sigres.__doc__)
1081    p_mdf = subparsers.add_parser('mdf', parents=robot_parents, help=abicomp_mdf.__doc__)
1082    p_optic = subparsers.add_parser('optic', parents=robot_parents, help=abicomp_optic.__doc__)
1083    p_a2f = subparsers.add_parser('a2f', parents=robot_parents, help=abicomp_a2f.__doc__)
1084    p_sigeph = subparsers.add_parser('sigeph', parents=robot_parents, help=abicomp_sigeph.__doc__)
1085    p_rta = subparsers.add_parser('rta', parents=robot_parents, help=abicomp_rta.__doc__)
1086    p_gkq = subparsers.add_parser('gkq', parents=robot_parents, help=abicomp_gkq.__doc__)
1087    p_gkq.add_argument('-d', '--diff', default=False, action="store_true", help='Plot difference between eph matrix elements.')
1088    p_v1qavg = subparsers.add_parser('v1qavg', parents=robot_parents, help=abicomp_v1qavg.__doc__)
1089    #p_wrmax = subparsers.add_parser('wrmax', parents=robot_parents, help=abicomp_wrmax.__doc__)
1090    p_abiwan = subparsers.add_parser('abiwan', parents=robot_parents, help=abicomp_abiwan.__doc__)
1091
1092    # Subparser for pseudos command.
1093    p_pseudos = subparsers.add_parser('pseudos', parents=[copts_parser], help=abicomp_pseudos.__doc__)
1094
1095    # Subparser for time command.
1096    p_time = subparsers.add_parser('time', parents=[copts_parser, ipy_parser], help=abicomp_time.__doc__)
1097    p_time.add_argument("-e", "--ext", default=".abo", help=("File extension for Abinit output files. "
1098                        "Used when the first argument is a directory. Default is `.abo`"))
1099
1100    # Subparser for gs_scf command.
1101    p_gs_scf = subparsers.add_parser('gs_scf', parents=[copts_parser], help=abicomp_gs_scf.__doc__)
1102
1103    # Subparser for dfpt2_scf command.
1104    p_dftp2_scf = subparsers.add_parser('dfpt2_scf', parents=[copts_parser], help=abicomp_dfpt2_scf.__doc__)
1105
1106    # Subparser for text command.
1107    p_text = subparsers.add_parser('text', parents=[copts_parser], help=abicomp_text.__doc__)
1108    p_text.add_argument("-d", "--diffmode", default="difflib", help=("Select diff application. "
1109        "Possible values: difflib (default), pygmentize (requires package)."))
1110
1111    return parser
1112
1113
1114@prof_main
1115def main():
1116
1117    def show_examples_and_exit(err_msg=None, error_code=1):
1118        """Display the usage of the script."""
1119        sys.stderr.write(get_epilog())
1120        if err_msg:
1121            sys.stderr.write("Fatal Error\n" + err_msg + "\n")
1122        sys.exit(error_code)
1123
1124    parser = get_parser(with_epilog=True)
1125
1126    # Parse the command line.
1127    try:
1128        options = parser.parse_args()
1129    except Exception:
1130        show_examples_and_exit(error_code=1)
1131
1132    # loglevel is bound to the string value obtained from the command line argument.
1133    # Convert to upper case to allow the user to specify --loglevel=DEBUG or --loglevel=debug
1134    import logging
1135    numeric_level = getattr(logging, options.loglevel.upper(), None)
1136    if not isinstance(numeric_level, int):
1137        raise ValueError('Invalid log level: %s' % options.loglevel)
1138    logging.basicConfig(level=numeric_level)
1139
1140    if options.mpl_backend is not None:
1141        # Set matplotlib backend
1142        import matplotlib
1143        matplotlib.use(options.mpl_backend)
1144
1145    if options.seaborn:
1146        # Use seaborn settings.
1147        import seaborn as sns
1148        sns.set(context=options.seaborn, style='darkgrid', palette='deep',
1149                font='sans-serif', font_scale=1, color_codes=False, rc=None)
1150
1151    if options.verbose > 2:
1152        print(options)
1153
1154    # Dispatch
1155    return globals()["abicomp_" + options.command](options)
1156
1157
1158if __name__ == "__main__":
1159    sys.exit(main())
1160