Source code for atomate2.vasp.drones

"""Drones for parsing VASP calculations and related outputs."""

from __future__ import annotations

import logging
import os
from pathlib import Path

from emmet.core.tasks import TaskDoc
from pymatgen.apps.borg.hive import AbstractDrone

logger = logging.getLogger(__name__)


[docs] class VaspDrone(AbstractDrone): """ A VASP drone to parse VASP outputs. Parameters ---------- **task_document_kwargs Additional keyword args passed to :obj:`.TaskDoc.from_directory`. """ def __init__(self, **task_document_kwargs) -> None: self.task_document_kwargs = task_document_kwargs
[docs] def assimilate(self, path: str | Path | None = None) -> TaskDoc: """ Parse VASP output files and return the output document. Parameters ---------- path : str or Path or None Path to the directory containing vasprun.xml and OUTCAR files. Returns ------- TaskDoc A VASP task document. """ if path is None: path = Path.cwd() try: doc = TaskDoc.from_directory(path, **self.task_document_kwargs) except Exception: import traceback logger.exception( f"Error in {Path(path).absolute()}\n{traceback.format_exc()}" ) raise return doc
[docs] def get_valid_paths(self, path: tuple[str, list[str], list[str]]) -> list[str]: """ Get valid paths to assimilate. There are some restrictions on the valid directory structures: 1. There can be only one vasprun in each directory. Nested directories are fine. 2. Directories designated "relax1"..."relax9" are considered to be parts of a multiple-optimization run. 3. Directories containing VASP output with ".relax1"...".relax9" are also considered as parts of a multiple-optimization run. Parameters ---------- path : tuple of (str, list of str, list of str) Input path as a tuple generated from ``os.walk``, i.e., (parent, subdirs, files). Returns ------- list of str A list of paths to assimilate. """ parent, subdirs, _ = path task_names = ["precondition"] + [f"relax{i}" for i in range(9)] if set(task_names).intersection(subdirs): return [parent] if ( not any(parent.endswith(os.sep + r) for r in task_names) and len(list(Path(parent).glob("vasprun.xml*"))) > 0 ): return [parent] return []