read vasp calc

read_vasp.py

import io
from pathlib import Path

import ase.io
import subprocess as sp
import numpy as np
from ase.calculators.singlepoint import SinglePointCalculator

import re

def find_energies(input_text):
    pattern = re.compile(r"E0= ([+-]?(?=\.\d|\d)(?:\d+)?(?:\.?\d*)(?:[Ee][+-]?\d+))?", re.MULTILINE)
    return [float(x) for x in pattern.findall(input_text)]

def _maybe_list(x):
    if not isinstance(x, list):
        return [x]
    return x


def read_vasp(calcdir):
    calcdir = Path(calcdir)


    images = _maybe_list(ase.io.read(calcdir / "XDATCAR", ":"))
    natoms = len(images[0])

    # Get energies
    cmd = ["grep", "E0", str(calcdir / "OSZICAR")]
    out = sp.check_output(cmd).decode("utf-8")
    energies = find_energies(out)

    # Get forces
    cmd = ["grep", "--group-separator", "@", "TOTAL-FORCE", str(calcdir / "OUTCAR"), f"-A{natoms + 1}"]
    out = sp.check_output(cmd).decode("utf-8")
    force_blocks = out.split("@")
    forces_list = []
    for block in force_blocks:
        sio = io.StringIO(block.strip())
        sio.seek(0)
        arr = np.loadtxt(sio, skiprows=2)
        forces = arr[:, 3:]
        forces_list.append(forces)


    assert len(energies) == len(images)
    assert len(forces_list) == len(images)
    assert all(f.shape[0] == natoms for f in forces_list)

    for i in range(len(images)):
        calc = SinglePointCalculator(images[i], energy=energies[i], forces=forces[i])
        images[i].calc = calc
    if len(images) == 1:
        images = images[0]
    return images