Added draft of vasp_util.py for VASP trajectories and OTF

flare/dft_interface/vasp_util.py

@@ -0,0 +1,128 @@
+from pymatgen.io.vasp.inputs import Poscar
+from pymatgen.io.vasp.outputs import Vasprun
+from subprocess import call
+import numpy as np
+from flare.struc import Structure
+from typing import List, Union
+from json import dump, load
+from flare.util import NumpyEncoder
+import os, shutil
+
+def check_vasprun(vasprun: Union[str, Vasprun]) -> Vasprun:
+    """
+    Helper utility to take a vasprun file name or Vasprun object
+    and return a vasprun object.
+    :param vasprun: vasprun filename or object
+    """
+    if type(vasprun) == str:
+        return Vasprun(vasprun)
+    elif type(vasprun) == Vasprun:
+        return vasprun
+    else:
+        raise ValueError('Vasprun argument is not a string or Vasprun instance!')
+
+def run_dft(calc_dir: str, dft_loc: str,
+    structure: Structure=None, en: bool=False, vasp_cmd="{}"):
+    """
+    Run a VASP calculation.
+    :param calc_dir: Name of directory to perform calculation in
+    :param dft_loc: Name of VASP command (i.e. executable location)
+    :param structure: flare structure object
+    :param en: whether to return the final energy along with the forces
+    :param vasp_cmd: Command to run VASP (leave a formatter "{}" to insert dft_loc);
+        this can be used to specify mpirun, srun, etc.
+
+    Returns:
+        forces on each atom (and energy if en=True)
+    """
+
+    currdir = os.getcwd()
+    if structure:
+        edit_dft_input_positions("POSCAR", structure)
+    try:
+        if currdir != calc_dir:
+            os.chdir(calc_dir)
+        call(vasp_cmd.format(dft_loc).split())
+
+        if en:
+            parse_func = parse_dft_forces_and_energy
+        else:
+            parse_func = parse_dft_forces
+
+        try:
+            forces = parse_func("vasprun.xml")
+        except FileNotFoundError:
+            raise FileNotFoundError("""Could not load vasprun.xml. The calculation may not have finished.
+                                        Current directory is %s""" % os.getcwd())
+
+        os.chdir(currdir)
+        return forces
+
+    except Exception as e:
+        os.chdir(currdir)
+        raise e
+
+
+def dft_input_to_structure(poscar: str):
+    """
+    Parse the DFT input in a directory.
+    :param vasp_input: directory of vasp input
+    """
+    return Structure.from_pmg_structure(Poscar.from_file(poscar).structure)
+
+
+def edit_dft_input_positions(poscar: str, structure: Structure):
+    """
+    Writes a VASP POSCAR file from structure with the name poscar .
+    WARNING: Destructively replaces the file with the name specified by poscar
+    :param poscar: Name of file
+    :param structure: structure to write to file
+    """
+    if os.path.isfile(poscar):
+        shutil.copyfile(poscar, poscar+'.bak')
+    f = open(poscar, 'w')
+    f.write(Poscar(structure.to_pmg_structure()).get_string(significant_figures=15))
+    f.close()
+    return poscar
+
+
+def parse_dft_forces(vasprun: Union[str, Vasprun]):
+    """
+    Parses the DFT forces from a VASP vasprun.xml file
+    :param vasprun: pymatgen Vasprun object or vasprun filename
+    """
+    vasprun = check_vasprun(vasprun)
+    istep = vasprun.ionic_steps[-1]
+    return np.array(istep['forces'])
+
+
+def parse_dft_forces_and_energy(vasprun: Union[str, Vasprun]):
+    """
+    Parses the DFT forces and energy from a VASP vasprun.xml file
+    :param vasprun: pymatgen Vasprun object or vasprun filename
+    """
+    vasprun = check_vasprun(vasprun)
+    istep = vasprun.ionic_steps[-1]
+    return np.array(istep['forces']), istep["electronic_steps"][-1]["e_0_energy"]
+
+
+def md_trajectory_from_vasprun(vasprun: Union[str, Vasprun], ionic_step_skips=1):
+    """
+    Returns a list of flare Structure objects decorated with forces, stress,
+    and total energy from a MD trajectory performed in VASP.
+    :param vasprun: pymatgen Vasprun object or vasprun filename
+    :param ionic_step_skips: if True, only samples the configuration every
+        ionic_skip_steps steps.
+    """
+    vasprun = check_vasprun(vasprun)
+
+    struc_lst = []
+    for step in vasprun.ionic_steps[::ionic_step_skips]:
+        structure = Structure.from_pmg_structure(step["structure"])
+        structure.energy = step["electronic_steps"][-1]["e_0_energy"]
+        #TODO should choose e_wo_entrp or e_fr_energy?
+        structure.forces = np.array(step["forces"])
+        structure.stress = np.array(step["stress"])
+        struc_lst.append(structure)
+
+    return struc_lst

flare/flare_io.py

@@ -0,0 +1,23 @@
+from flare.struc import Structure
+from typing import List
+from json import dump, load
+from flare.util import NumpyEncoder
+
+def md_trajectory_to_file(filename: str, structures: List[Structure]):
+    """
+    Take a list of structures and write them to a json file.
+    :param filename:
+    :param structures:
+    """
+    with open(filename, 'w') as f:
+        dump([s.as_dict() for s in structures], f, cls=NumpyEncoder)
+
+def md_trajectory_from_file(filename: str):
+    """
+    Read a list of structures from a json file, formatted as in md_trajectory_to_file.
+    :param filename:
+    """
+    with open(filename, 'r') as f:
+        structure_list = load(f)
+        structures = [Structure.from_dict(dictionary) for dictionary in structure_list]
+    return structures

flare/struc.py

@@ -171,6 +171,7 @@ def from_dict(dictionary):
                           positions=np.array(dictionary['positions']),
                           species=dictionary['coded_species'])
 
+        struc.energy = dictionary['energy']
         struc.forces = np.array(dictionary['forces'])
         struc.stress = dictionary['stress']
         struc.stds = np.array(dictionary['stds'])
@@ -221,9 +222,9 @@ def from_pmg_structure(structure):
         :return: FLARE Structure
         """
 
-        cell = structure.lattice.matrix
+        cell = structure.lattice.matrix.copy()
         species = [str(spec) for spec in structure.species]
-        positions = structure.cart_coords
+        positions = structure.cart_coords.copy()
 
         new_struc = Structure(cell=cell,species=species,
                               positions=positions)

tests/test_files/dummy_vasp.py

@@ -0,0 +1,12 @@
+import os, shutil, sys
+
+if "test_fail" in sys.argv:
+	exit(0)
+
+curr_path = os.path.dirname(os.path.abspath(__file__))
+vasprun_name = os.path.join(curr_path, 'test_vasprun.xml')
+poscar_name = os.path.join(curr_path, 'test_POSCAR')
+calc_dir = os.path.join(curr_path, '..')
+
+shutil.copyfile(poscar_name, os.path.join(calc_dir, "POSCAR"))
+shutil.copyfile(vasprun_name, os.path.join(calc_dir, "vasprun.xml"))

tests/test_files/test_POSCAR

@@ -0,0 +1,14 @@
+H4 C1 O1
+1.0
+10.638000 0.000000 0.000000
+0.000000 10.029603 0.000000
+0.000000 0.000000 20.000000
+C O H
+1 1 4
+direct
+0.066840 0.466900 0.684808 C
+0.997467 0.388247 0.636881 O
+0.058718 0.354817 0.604874 H
+0.059432 0.573624 0.674176 H
+0.168523 0.443819 0.681503 H
+0.024984 0.447208 0.735591 H