additional initial defect requirements for parsing in defect wf

Author: dbroberg

atomate/vasp/drones.py

@@ -20,6 +20,7 @@
 from monty.io import zopen
 from monty.json import jsanitize
 from monty.os.path import which
+from monty.shutil import decompress_file
 
 import numpy as np
 
@@ -28,7 +29,8 @@
 from pymatgen.core.operations import SymmOp
 from pymatgen.electronic_structure.bandstructure import BandStructureSymmLine
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
-from pymatgen.io.vasp import BSVasprun, Vasprun, Outcar, Locpot, Chgcar
+from pymatgen.io.vasp import BSVasprun, Vasprun, Outcar, Locpot, Chgcar, \
+    parse_defect_states, Procar, Wavecar
 from pymatgen.io.vasp.inputs import Poscar, Potcar, Incar, Kpoints
 from pymatgen.apps.borg.hive import AbstractDrone
 from pymatgen.command_line.bader_caller import bader_analysis_from_path
@@ -83,7 +85,8 @@ class VaspDrone(AbstractDrone):
 
     def __init__(self, runs=None, parse_dos="auto", bandstructure_mode="auto",
                  parse_locpot=True, additional_fields=None, use_full_uri=True,
-                 parse_bader=bader_exe_exists, parse_chgcar=False, parse_aeccar=False):
+                 parse_bader=bader_exe_exists, parse_chgcar=False, parse_aeccar=False,
+                 defect_wf_parsing=None):
         """
         Initialize a Vasp drone to parse vasp outputs
         Args:
@@ -106,6 +109,10 @@ def __init__(self, runs=None, parse_dos="auto", bandstructure_mode="auto",
             parse_bader (bool): Run and parse Bader charge data. Defaults to True if Bader is present
             parse_chgcar (bool): Run and parse CHGCAR file
             parse_aeccar (bool): Run and parse AECCAR0 and AECCAR2 files
+            defect_wf_parsing (Site): If Site is provided, drone considers Procar and
+                Wavecar parsing relative to the position of Site.
+                 Useful for consideration of defect localization
+                Defaults to None (no extra procar or wavecar parsing occurs)
         """
         self.parse_dos = parse_dos
         self.additional_fields = additional_fields or {}
@@ -117,6 +124,7 @@ def __init__(self, runs=None, parse_dos="auto", bandstructure_mode="auto",
         self.parse_bader = parse_bader
         self.parse_chgcar = parse_chgcar
         self.parse_aeccar = parse_aeccar
+        self.defect_wf_parsing = defect_wf_parsing
 
     def assimilate(self, path):
         """
@@ -210,6 +218,35 @@ def generate_doc(self, dir_name, vasprun_files, outcar_files):
                 logger.error("Bad run stats for {}.".format(fullpath))
             d["run_stats"] = run_stats
 
+            # store defect localization/band filling information
+            if self.defect_wf_parsing:
+                # need to make sure all procars and wavecars are zipped
+                for propat in self.filter_files(fullpath, file_pattern="PROCAR").values():
+                    decompress_file(os.path.join(fullpath, propat))
+                for wavepat in self.filter_files(fullpath, file_pattern="WAVECAR").values():
+                    decompress_file(os.path.join(fullpath, wavepat))
+                procar_paths = [os.path.join(fullpath, ppath) for ppath in
+                                self.filter_files(fullpath, file_pattern="PROCAR").values()]
+                wavecar_paths = [os.path.join(fullpath, wpath) for wpath in
+                                 self.filter_files(fullpath, file_pattern="WAVECAR").values()]
+
+                for i, d_calc in enumerate(d["calcs_reversed"]):
+                    if d_calc.get("output"):
+                        if len(procar_paths) and len(wavecar_paths):
+                            procar = Procar(procar_paths[i])
+                            wavecar = Wavecar(wavecar_paths[i])
+                            structure = Structure.from_dict(d_calc["output"]["structure"])
+                            defect_data = parse_defect_states(structure, self.defect_wf_parsing, wavecar, procar)
+                            d_calc["output"].update({"defect": defect_data})
+
+                        filename = list(vasprun_files.values())[i]
+                        vasprun_file = os.path.join(dir_name, filename)
+                        vrun = Vasprun(vasprun_file)
+                        eigenvalues = vrun.eigenvalues.copy()
+                        kpoint_weights = vrun.actual_kpoints_weights
+                        vr_eigenvalue_dict = {'eigenvalues': eigenvalues, 'kpoint_weights': kpoint_weights}
+                        d_calc["output"].update({"vr_eigenvalue_dict": vr_eigenvalue_dict})
+
             # reverse the calculations data order so newest calc is first
             d["calcs_reversed"].reverse()
 

atomate/vasp/firetasks/parse_outputs.py

@@ -91,7 +91,8 @@ def run_task(self, fw_spec):
                           parse_dos=self.get("parse_dos", False),
                           bandstructure_mode=self.get("bandstructure_mode", False),
                           parse_chgcar=self.get("parse_chgcar", False),
-                          parse_aeccar=self.get("parse_aeccar", False))
+                          parse_aeccar=self.get("parse_aeccar", False),
+                          defect_wf_parsing=self.get("defect_wf_parsing", None))
 
         # assimilate (i.e., parse)
         task_doc = drone.assimilate(calc_dir)

atomate/vasp/fireworks/core.py

@@ -538,7 +538,9 @@ def __init__(self, structure, transformations, transformation_params=None,
                  vasp_input_set=None, prev_calc_dir=None,
                  name="structure transmuter", vasp_cmd="vasp",
                  copy_vasp_outputs=True, db_file=None,
-                 parents=None, override_default_vasp_params=None, **kwargs):
+                 parents=None, bandstructure_mode=None,
+                 override_default_vasp_params=None,
+                 defect_wf_parsing=None, **kwargs):
         """
         Apply the transformations to the input structure, write the input set corresponding
         to the transformed structure, and run vasp on them.  Note that if a transformation yields 
@@ -559,7 +561,14 @@ def __init__(self, structure, transformations, transformation_params=None,
             prev_calc_dir (str): Path to a previous calculation to copy from
             db_file (string): Path to file specifying db credentials.
             parents (Firework): Parents of this particular Firework. FW or list of FWS.
+            bandstructure_mode (str): Set to "uniform" for uniform band structure.
+                Set to "line" for line mode. If not set, band structure will not
+                be parsed.
             override_default_vasp_params (dict): additional user input settings for vasp_input_set.
+            defect_wf_parsing (Site): If Site is provided, drone considers Procar and
+                Wavecar localization parsing relative to the position of Site.
+                 Useful for consideration of defect localization
+                Defaults to None (no extra procar or wavecar parsing occurs for VaspToDb)
             \*\*kwargs: Other kwargs that are passed to Firework.__init__.
         """
         fw_name = "{}-{}".format(structure.composition.reduced_formula, name)
@@ -603,7 +612,9 @@ def __init__(self, structure, transformations, transformation_params=None,
                               "task_label": name,
                               "transmuter": {"transformations": transformations,
                                              "transformation_params": transformation_params}
-                          }))
+                          },
+                          bandstructure_mode=bandstructure_mode,
+                          defect_wf_parsing=defect_wf_parsing))
 
         super(TransmuterFW, self).__init__(t, parents=parents,
                                            name=fw_name, **kwargs)