def test_from_directory(self):
     vi = VaspInput.from_directory(test_dir, optional_files={"CONTCAR.Li2O": Poscar})
     self.assertEqual(vi["INCAR"]["ALGO"], "Damped")
     self.assertIn("CONTCAR.Li2O", vi)
     d = vi.to_dict
     vinput = VaspInput.from_dict(d)
     self.assertIn("CONTCAR.Li2O", vinput)

    def setup(self):
        files = os.listdir(".")
        num_structures = 0
        if not set(files).issuperset(VASP_INPUT_FILES):
            for f in files:
                try:
                    struct = read_structure(f)
                    num_structures += 1
                except:
                    pass
            if num_structures != 1:
                raise RuntimeError("{} structures found. Unable to continue."
                                   .format(num_structures))
            else:
                self.default_vis.write_input(struct, ".")

        if self.backup:
            for f in VASP_INPUT_FILES:
                shutil.copy(f, "{}.orig".format(f))

        if self.auto_npar:
            try:
                vi = VaspInput.from_directory(".")
                incar = vi["INCAR"]
                #Only optimized NPAR for non-HF and non-RPA calculations.
                if (not incar.get("LHFCALC")) and (not incar.get("LRPA")):
                    import multiprocessing
                    ncores = multiprocessing.cpu_count()
                    for npar in range(int(round(math.sqrt(ncores))), ncores):
                        if ncores % npar == 0:
                            incar["NPAR"] = npar
                            break
                    incar.write_file("INCAR")
            except:
                pass

        if self.settings_override is not None:
            vi = VaspInput.from_directory(".")
            m = Modder([FileActions, DictActions])
            modified = []
            for a in self.settings_override:
                if "dict" in a:
                    modified.append(a["dict"])
                    vi[a["dict"]] = m.modify_object(a["action"],
                                                    vi[a["dict"]])
                elif "filename" in a:
                    m.modify(a["action"], a["filename"])
            for f in modified:
                vi[f].write_file(f)

class VaspInputTest(unittest.TestCase):

    def setUp(self):
        filepath = os.path.join(test_dir, 'INCAR')
        incar = Incar.from_file(filepath)
        filepath = os.path.join(test_dir, 'POSCAR')
        poscar = Poscar.from_file(filepath)
        if "VASP_PSP_DIR" not in os.environ:
            test_potcar_dir = os.path.abspath(
                os.path.join(os.path.dirname(__file__), "..", "..", "..", "..",
                             "test_files"))
            os.environ["VASP_PSP_DIR"] = test_potcar_dir
        filepath = os.path.join(test_dir, 'POTCAR')
        potcar = Potcar.from_file(filepath)
        filepath = os.path.join(test_dir, 'KPOINTS.auto')
        kpoints = Kpoints.from_file(filepath)
        self.vinput = VaspInput(incar, kpoints, poscar, potcar)

    def test_to_from_dict(self):
        d = self.vinput.as_dict()
        vinput = VaspInput.from_dict(d)
        comp = vinput["POSCAR"].structure.composition
        self.assertEqual(comp, Composition("Fe4P4O16"))

    def test_from_directory(self):
        vi = VaspInput.from_directory(test_dir,
                                      optional_files={"CONTCAR.Li2O": Poscar})
        self.assertEqual(vi["INCAR"]["ALGO"], "Damped")
        self.assertIn("CONTCAR.Li2O", vi)
        d = vi.as_dict()
        vinput = VaspInput.from_dict(d)
        self.assertIn("CONTCAR.Li2O", vinput)

 def check(self):
     vi = VaspInput.from_directory(".")
     nelm = vi["INCAR"].get("NELM", 60)
     try:
         oszicar = Oszicar(self.output_filename)
         esteps = oszicar.ionic_steps
         if len(esteps) > 10:
             return all([len(e) == nelm for e in esteps[-11:-1]])
     except:
         pass
     return False

    def correct(self):
        backup(["INCAR", "KPOINTS", "POSCAR", "OUTCAR",
                self.output_filename, "vasp.out"])
        vi = VaspInput.from_directory(".")
        ediff = float(vi["INCAR"].get("EDIFF", 1e-4))
        actions = [{"file": "CONTCAR",
                    "action": {"_file_copy": {"dest": "POSCAR"}}},
                   {"dict": "INCAR",
                    "action": {"_set": {"EDIFF": ediff*0.75}}}]
        VaspModder(vi=vi).apply_actions(actions)

        return {"errors": ["MaxForce"], "actions": actions}

    def run(self):
        cmd = list(self.vasp_cmd)
        if self.auto_gamma:
            vi = VaspInput.from_directory(".")
            kpts = vi["KPOINTS"]
            if kpts.style == "Gamma" and tuple(kpts.kpts[0]) == (1, 1, 1):
                if self.gamma_vasp_cmd is not None and os.path.exists(
                        self.gamma_vasp_cmd[-1]):
                    cmd = self.gamma_vasp_cmd
                elif os.path.exists(cmd[-1] + ".gamma"):
                    cmd[-1] += ".gamma"

        with open(self.output_file, 'w') as f:
            return subprocess.Popen(cmd, stdout=f)

 def correct(self):
     backup()
     vi = VaspInput.from_directory(".")
     potim = float(vi["INCAR"].get("POTIM", 0.5)) * 0.5
     actions = [{"dict": "INCAR",
                 "action": {"_set": {"POTIM": potim}}}]
     m = Modder()
     modified = []
     for a in actions:
         modified.append(a["dict"])
         vi[a["dict"]] = m.modify_object(a["action"], vi[a["dict"]])
     for f in modified:
         vi[f].write_file(f)
     return {"errors": ["POTIM"], "actions": actions}

 def setUp(self):
     filepath = os.path.join(test_dir, 'INCAR')
     incar = Incar.from_file(filepath)
     filepath = os.path.join(test_dir, 'POSCAR')
     poscar = Poscar.from_file(filepath)
     if "VASP_PSP_DIR" not in os.environ:
         test_potcar_dir = os.path.abspath(
             os.path.join(os.path.dirname(__file__), "..", "..", "..", "..",
                          "test_files"))
         os.environ["VASP_PSP_DIR"] = test_potcar_dir
     filepath = os.path.join(test_dir, 'POTCAR')
     potcar = Potcar.from_file(filepath)
     filepath = os.path.join(test_dir, 'KPOINTS.auto')
     kpoints = Kpoints.from_file(filepath)
     self.vinput = VaspInput(incar, kpoints, poscar, potcar)

    def run_task(self, fw_spec):
        chgcar_start = False

        vi = VaspInput.from_directory(".")  # read the VaspInput from the previous run

        # figure out what GGA+U values to use and override them
        mpvis = MPVaspInputSet()
        incar = mpvis.get_incar(vi['POSCAR'].structure).to_dict
        incar_updates = {k: incar[k] for k in incar.keys() if 'LDAU' in k}  # LDAU values to use
        vi['INCAR'].update(incar_updates)  # override the +U keys

        # start from the CHGCAR of previous run
        if os.path.exists('CHGCAR'):
            vi['INCAR']['ICHARG'] = 1
            chgcar_start = True

        vi['INCAR'].write_file('INCAR')  # write back the new INCAR to the current directory

        return FWAction(stored_data={'chgcar_start': chgcar_start})

 def check(self):
     msg = "Reciprocal lattice and k-lattice belong to different class of" \
           " lattices."
     try:
         v = Vasprun(self.output_vasprun)
         vi = VaspInput.from_directory('.')
         # According to VASP admins, you can disregard this error
         # if symmetry is off
         #Also disregard if automatic KPOINT generation is used
         if v.converged or (not v.incar.get('ISYM', True)) or\
                         vi["KPOINTS"].style == Kpoints.supported_modes.Automatic:
             return False
     except:
         pass
     with open(self.output_filename, "r") as f:
         for line in f:
             l = line.strip()
             if l.find(msg) != -1:
                 return True
     return False

    def run(self):
        """
        Perform the actual VASP run.

        Returns:
            (subprocess.Popen) Used for monitoring.
        """
        cmd = list(self.vasp_cmd)
        if self.auto_gamma:
            vi = VaspInput.from_directory(".")
            kpts = vi["KPOINTS"]
            if kpts.style == "Gamma" and tuple(kpts.kpts[0]) == (1, 1, 1):
                if self.gamma_vasp_cmd is not None and which(
                        self.gamma_vasp_cmd[-1]):
                    cmd = self.gamma_vasp_cmd
                elif which(cmd[-1] + ".gamma"):
                    cmd[-1] += ".gamma"
        logging.info("Running {}".format(" ".join(cmd)))
        with open(self.output_file, 'w') as f:
            p = subprocess.Popen(cmd, stdout=f)
        return p

    def correct(self): #pylint: disable=R0915,R0912,R0914
        """
        "brmix"にてINCARにADDGRIDを追加する
        """
        backup([self.output_filename, "INCAR", "KPOINTS", "POSCAR", "OUTCAR",
                "vasprun.xml"])
        actions = []
        vi = VaspInput.from_directory(".")

        if "tet" in self.errors or "dentet" in self.errors:
            actions.append({"dict": "INCAR",
                            "action": {"_set": {"ISMEAR": 0}}})

        if "inv_rot_mat" in self.errors:
            actions.append({"dict": "INCAR",
                            "action": {"_set": {"SYMPREC": 1e-8}}})

        if "brmix" in self.errors and "NELECT" in vi["INCAR"]:
            #brmix error always shows up after DAV steps if NELECT is specified
            self.errors.remove("brmix")
        if "brmix" in self.errors or "zpotrf" in self.errors:
            actions.append({"dict": "INCAR",
                            "action": {"_set": {"ISYM": 0}}})
            #140814 add
            actions.append({"dict": "INCAR",
                            "action": {"_set": {"ADDGRID": True}}})
            # Based on VASP forum's recommendation, you should delete the
            # CHGCAR and WAVECAR when dealing with these errors.
            actions.append({"file": "CHGCAR",
                            "action": {"_file_delete": {'mode': "actual"}}})
            actions.append({"file": "WAVECAR",
                            "action": {"_file_delete": {'mode': "actual"}}})

        if "subspacematrix" in self.errors or "rspher" in self.errors or \
                        "real_optlay" in self.errors:
            actions.append({"dict": "INCAR",
                            "action": {"_set": {"LREAL": False}}})

        if "tetirr" in self.errors or "incorrect_shift" in self.errors or \
                    "rot_matrix" in self.errors:
            actions.append({"dict": "KPOINTS",
                            "action": {"_set": {"generation_style": "Gamma"}}})

        if "amin" in self.errors:
            actions.append({"dict": "INCAR",
                            "action": {"_set": {"AMIN": "0.01"}}})

        if "triple_product" in self.errors:
            s = vi["POSCAR"].structure
            trans = SupercellTransformation(((1, 0, 0), (0, 0, 1), (0, 1, 0)))
            new_s = trans.apply_transformation(s)
            actions.append({"dict": "POSCAR",
                            "action": {"_set": {"structure": new_s.to_dict}},
                            "transformation": trans.to_dict})

        if "pricel" in self.errors:
            actions.append({"dict": "INCAR",
                            "action": {"_set": {"SYMPREC": 1e-8, "ISYM": 0}}})

        if "brions" in self.errors:
            potim = float(vi["INCAR"].get("POTIM", 0.5)) + 0.1
            actions.append({"dict": "INCAR",
                            "action": {"_set": {"POTIM": potim}}})

        if "zbrent" in self.errors:
            actions.append({"dict": "INCAR",
                            "action": {"_set": {"IBRION": 1}}})

        if "too_few_bands" in self.errors:
            if "NBANDS" in vi["INCAR"]:
                nbands = int(vi["INCAR"]["NBANDS"])
            else:
                with open("OUTCAR") as f:
                    for line in f:
                        if "NBANDS" in line:
                            try:
                                d = line.split("=")
                                nbands = int(d[-1].strip())
                                break
                            except (IndexError, ValueError):
                                pass
            actions.append({"dict": "INCAR",
                            "action": {"_set": {"NBANDS": int(1.1 * nbands)}}})

        if "aliasing" in self.errors:
            with open("OUTCAR") as f:
                grid_adjusted = False
                changes_dict = {}
                for line in f:
                    if "aliasing errors" in line:
                        try:
                            grid_vector = line.split(" NG", 1)[1]
                            value = [int(s) for s in grid_vector.split(" ")
                                     if s.isdigit()][0]

                            changes_dict["NG" + grid_vector[0]] = value
                            grid_adjusted = True
                        except (IndexError, ValueError):
                            pass
                    #Ensure that all NGX, NGY, NGZ have been checked
#.........这里部分代码省略.........

 def test_to_from_dict(self):
     d = self.vinput.to_dict
     vinput = VaspInput.from_dict(d)
     comp = vinput["POSCAR"].structure.composition
     self.assertEqual(comp, Composition.from_formula("Fe4P4O16"))

__author__ = 'Qimin'

from pymatgen.io.vaspio.vasp_output import Vasprun
from pymatgen.io.vaspio.vasp_input import VaspInput
from pymatgen.io.vaspio_set import MPNonSCFVaspInputSet
import os

vasp_dir = os.path.dirname(os.path.abspath(__file__))
vasp_run = Vasprun(os.path.join(vasp_dir, "vasprun.xml")).to_dict
nband = int(vasp_run['input']['parameters']['NBANDS'])
prec = str(vasp_run['input']['incar']['PREC'])
encut = int(vasp_run['input']['incar']['ENCUT'])

user_incar_settings={"PREC":prec,"ENCUT":encut,"EDIFF":1E-4,"NBANDS":nband,"NSW":0}
#user_incar_settings={"EDIFF":1E-4,"NBANDS":nband,"NSW":0}
mpvis = MPNonSCFVaspInputSet(user_incar_settings=user_incar_settings)
vi = VaspInput.from_directory(".")  # read the VaspInput from the previous run
mpvis.get_kpoints(vi['POSCAR'].structure).write_file('KPOINTS')
mpvis.get_incar(vi['POSCAR'].structure).write_file('INCAR')