Source code for pymatgen.cli.pmg_plot

#!/usr/bin/env python
# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

"""
Implementation for `pmg plot` CLI.
"""

from collections import OrderedDict
from pymatgen import Structure
from pymatgen.electronic_structure.plotter import DosPlotter
from pymatgen.io.vasp import Vasprun, Chgcar
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.analysis.diffraction.xrd import XRDCalculator


def get_dos_plot(args):
    """
    Plot DOS.

    Args:
        args (dict): Args from argparse.
    """
    v = Vasprun(args.dos_file)
    dos = v.complete_dos

    all_dos = OrderedDict()
    all_dos["Total"] = dos

    structure = v.final_structure

    if args.site:
        for i in range(len(structure)):
            site = structure[i]
            all_dos["Site " + str(i) + " " + site.specie.symbol] = \
                dos.get_site_dos(site)

    if args.element:
        syms = [tok.strip() for tok in args.element[0].split(",")]
        all_dos = {}
        for el, dos in dos.get_element_dos().items():
            if el.symbol in syms:
                all_dos[el] = dos
    if args.orbital:
        all_dos = dos.get_spd_dos()

    plotter = DosPlotter()
    plotter.add_dos_dict(all_dos)
    return plotter.get_plot()


def get_chgint_plot(args):
    """
    Plot integrated charge.

    Args:
        args (dict): args from argparse.
    """
    chgcar = Chgcar.from_file(args.chgcar_file)
    s = chgcar.structure

    if args.inds:
        atom_ind = [int(i) for i in args.inds[0].split(",")]
    else:
        finder = SpacegroupAnalyzer(s, symprec=0.1)
        sites = [sites[0] for sites in
                 finder.get_symmetrized_structure().equivalent_sites]
        atom_ind = [s.sites.index(site) for site in sites]

    from pymatgen.util.plotting import pretty_plot
    plt = pretty_plot(12, 8)
    for i in atom_ind:
        d = chgcar.get_integrated_diff(i, args.radius, 30)
        plt.plot(d[:, 0], d[:, 1],
                 label="Atom {} - {}".format(i, s[i].species_string))
    plt.legend(loc="upper left")
    plt.xlabel("Radius (A)")
    plt.ylabel("Integrated charge (e)")
    plt.tight_layout()
    return plt


def get_xrd_plot(args):
    """
    Plot XRD

    Args:
        args (dict): Args from argparse
    """
    s = Structure.from_file(args.xrd_structure_file)
    c = XRDCalculator()
    return c.get_plot(s)


def plot(args):
    """
    Master control method calling other plot methods based on args.

    Args:
        args (dict): Args from argparse.
    """
    plt = None
    if args.chgcar_file:
        plt = get_chgint_plot(args)
    elif args.xrd_structure_file:
        plt = get_xrd_plot(args)
    elif args.dos_file:
        plt = get_dos_plot(args)

    if plt:
        if args.out_file:
            plt.savefig(args.out_file)
        else:
            plt.show()