# mypy: ignore-errors import logging import re from collections import OrderedDict from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Mapping, Optional, Tuple, Type, TypeVar, Union import numpy as np from emmet.core.common import convert_datetime from emmet.core.mpid import MPID from emmet.core.structure import StructureMetadata from emmet.core.utils import utcnow from emmet.core.vasp.calc_types import ( CalcType, calc_type, TaskType, run_type, RunType, task_type, ) from emmet.core.vasp.calculation import ( CalculationInput, Calculation, PotcarSpec, RunStatistics, VaspObject, ) from emmet.core.vasp.task_valid import TaskState from monty.json import MontyDecoder from monty.serialization import loadfn from pydantic import ( BaseModel, ConfigDict, Field, field_validator, model_validator, ) from pymatgen.analysis.structure_analyzer import oxide_type from pymatgen.core.structure import Structure from pymatgen.core.trajectory import Trajectory from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry from pymatgen.io.vasp import Incar, Kpoints, Poscar from pymatgen.io.vasp import Potcar as VaspPotcar monty_decoder = MontyDecoder() logger = logging.getLogger(__name__) _T = TypeVar("_T", bound="TaskDoc") _VOLUMETRIC_FILES = ("CHGCAR", "LOCPOT", "AECCAR0", "AECCAR1", "AECCAR2") class Potcar(BaseModel): pot_type: Optional[str] = Field(None, description="Pseudo-potential type, e.g. PAW") functional: Optional[str] = Field( None, description="Functional type use in the calculation." ) symbols: Optional[List[str]] = Field( None, description="List of VASP potcar symbols used in the calculation." ) class OrigInputs(CalculationInput): poscar: Optional[Poscar] = Field( None, description="Pymatgen object representing the POSCAR file.", ) potcar: Optional[Union[Potcar, VaspPotcar, List[Any]]] = Field( None, description="Pymatgen object representing the POTCAR file.", ) @field_validator("potcar", mode="before") @classmethod def potcar_ok(cls, v): """Check that the POTCAR meets type requirements.""" if isinstance(v, list): return list(v) return v @field_validator("potcar", mode="after") @classmethod def parse_potcar(cls, v): """Check that potcar attribute is not a pymatgen POTCAR.""" if isinstance(v, VaspPotcar): # The user should not mix potential types, but account for that here # Using multiple potential types will be caught in validation pot_typ = "_".join(set(p.potential_type for p in v)) return Potcar(pot_type=pot_typ, functional=v.functional, symbols=v.symbols) return v model_config = ConfigDict(arbitrary_types_allowed=True) class OutputDoc(BaseModel): structure: Optional[Structure] = Field( None, title="Output Structure", description="Output Structure from the VASP calculation.", ) density: Optional[float] = Field(None, description="Density of in units of g/cc.") energy: Optional[float] = Field(None, description="Total Energy in units of eV.") forces: Optional[List[List[float]]] = Field( None, description="The force on each atom in units of eV/A." ) stress: Optional[List[List[float]]] = Field( None, description="The stress on the cell in units of kB." ) energy_per_atom: Optional[float] = Field( None, description="The final DFT energy per atom for the last calculation" ) bandgap: Optional[float] = Field( None, description="The DFT bandgap for the last calculation" ) @model_validator(mode="before") def set_density_from_structure(cls, values): # Validator to automatically set density from structure if not already # specified. This might happen when importing an older atomate2-format # TaskDocument. if not values.get("density", None): if isinstance(values["structure"], dict): values["density"] = values["structure"].get("density", None) else: values["density"] = values["structure"].density return values @classmethod def from_vasp_calc_doc( cls, calc_doc: Calculation, trajectory: Optional[Trajectory] = None ) -> "OutputDoc": """ Create a summary of VASP calculation outputs from a VASP calculation document. This will first look for ionic steps in the calculation document. If found, will use it and ignore the trajectory. I not, will get ionic steps from the trajectory. Parameters ---------- calc_doc A VASP calculation document. trajectory A pymatgen Trajectory. Returns ------- OutputDoc The calculation output summary. """ if calc_doc.output.ionic_steps: forces = calc_doc.output.ionic_steps[-1].forces stress = calc_doc.output.ionic_steps[-1].stress elif trajectory: ionic_steps = trajectory.frame_properties forces = ionic_steps[-1]["forces"] stress = ionic_steps[-1]["stress"] else: raise RuntimeError("Unable to find ionic steps.") return cls( structure=calc_doc.output.structure, energy=calc_doc.output.energy, energy_per_atom=calc_doc.output.energy_per_atom, bandgap=calc_doc.output.bandgap, forces=forces, stress=stress, ) class InputDoc(CalculationInput): """Light wrapper around `CalculationInput` with a few extra fields. pseudo_potentials (Potcar) : summary of the POTCARs used in the calculation xc_override (str) : the exchange-correlation functional used if not the one specified by POTCAR is_lasph (bool) : how the calculation set LASPH (aspherical corrections) magnetic_moments (list of floats) : on-site magnetic moments """ pseudo_potentials: Optional[Potcar] = Field( None, description="Summary of the pseudo-potentials used in this calculation" ) xc_override: Optional[str] = Field( None, description="Exchange-correlation functional used if not the default" ) is_lasph: Optional[bool] = Field( None, description="Whether the calculation was run with aspherical corrections" ) magnetic_moments: Optional[List[float]] = Field( None, description="Magnetic moments for each atom" ) @field_validator("parameters", mode="after") @classmethod def parameter_keys_should_not_contain_spaces(cls, parameters: Optional[Dict]): # A change in VASP introduced whitespace into some parameters, # for example `PE` was observed in # VASP 6.4.3. This will lead to an incorrect return value from RunType. # This validator will ensure that any already-parsed documents are fixed. if parameters: return {k.strip(): v for k, v in parameters.items()} @classmethod def from_vasp_calc_doc(cls, calc_doc: Calculation) -> "InputDoc": """ Create calculation input summary from a calculation document. Parameters ---------- calc_doc A VASP calculation document. Returns ------- InputDoc A summary of the input structure and parameters. """ xc = calc_doc.input.incar.get("GGA") or calc_doc.input.incar.get("METAGGA") if xc: xc = xc.upper() pot_type, func = calc_doc.input.potcar_type[0].split("_") func = "lda" if len(pot_type) == 1 else "_".join(func) pps = Potcar(pot_type=pot_type, functional=func, symbols=calc_doc.input.potcar) return cls( **calc_doc.input.model_dump(), pseudo_potentials=pps, xc_override=xc, is_lasph=calc_doc.input.parameters.get("LASPH", False), magnetic_moments=calc_doc.input.parameters.get("MAGMOM"), ) class CustodianDoc(BaseModel): corrections: Optional[List[Any]] = Field( None, title="Custodian Corrections", description="List of custodian correction data for calculation.", ) job: Optional[Any] = Field( None, title="Custodian Job Data", description="Job data logged by custodian.", ) class AnalysisDoc(BaseModel): delta_volume: Optional[float] = Field( None, title="Volume Change", description="Volume change for the calculation.", ) delta_volume_percent: Optional[float] = Field( None, title="Volume Change Percent", description="Percent volume change for the calculation.", ) max_force: Optional[float] = Field( None, title="Max Force", description="Maximum force on any atom at the end of the calculation.", ) warnings: Optional[List[str]] = Field( None, title="Calculation Warnings", description="Warnings issued after analysis.", ) errors: Optional[List[str]] = Field( None, title="Calculation Errors", description="Errors issued after analysis.", ) @classmethod def from_vasp_calc_docs( cls, calcs_reversed: List[Calculation], volume_change_warning_tol: float = 0.2, ) -> "AnalysisDoc": """ Create analysis summary from VASP calculation documents. Parameters ---------- calcs_reversed A list of VASP calculation documents in reverse order . volume_change_warning_tol Maximum volume change allowed in VASP relaxations before the calculation is tagged with a warning. Returns ------- AnalysisDoc The relaxation analysis. """ initial_vol = calcs_reversed[-1].input.structure.lattice.volume final_vol = calcs_reversed[0].output.structure.lattice.volume delta_vol = final_vol - initial_vol percent_delta_vol = 100 * delta_vol / initial_vol warnings = [] errors = [] if abs(percent_delta_vol) > volume_change_warning_tol * 100: warnings.append(f"Volume change > {volume_change_warning_tol * 100}%") final_calc = calcs_reversed[0] max_force = None if final_calc.has_vasp_completed == TaskState.SUCCESS: # max force and valid structure checks structure = final_calc.output.structure # do not check max force for MD run if calcs_reversed[0].input.parameters.get("IBRION", -1) != 0: max_force = _get_max_force(final_calc) warnings.extend(_get_drift_warnings(final_calc)) if not structure.is_valid(): errors.append("Bad structure (atoms are too close!)") return cls( delta_volume=delta_vol, delta_volume_percent=percent_delta_vol, max_force=max_force, warnings=warnings, errors=errors, ) class TaskDoc(StructureMetadata, extra="allow"): """Calculation-level details about VASP calculations that power Materials Project.""" tags: Union[List[str], None] = Field( [], title="tag", description="Metadata tagged to a given task." ) dir_name: Optional[str] = Field( None, description="The directory for this VASP task" ) state: Optional[TaskState] = Field(None, description="State of this calculation") calcs_reversed: Optional[List[Calculation]] = Field( None, title="Calcs reversed data", description="Detailed data for each VASP calculation contributing to the task document.", ) structure: Optional[Structure] = Field( None, description="Final output structure from the task" ) task_type: Optional[Union[TaskType, CalcType]] = Field( None, description="The type of calculation." ) run_type: Optional[RunType] = Field( None, description="The functional used in the calculation." ) calc_type: Optional[CalcType] = Field( None, description="The functional and task type used in the calculation." ) task_id: Optional[Union[MPID, str]] = Field( None, description="The (task) ID of this calculation, used as a universal reference across property documents." "This comes in the form: mp-******.", ) orig_inputs: Optional[OrigInputs] = Field( None, description="The exact set of input parameters used to generate the current task document.", ) input: Optional[InputDoc] = Field( None, description="The input structure used to generate the current task document.", ) output: Optional[OutputDoc] = Field( None, description="The exact set of output parameters used to generate the current task document.", ) included_objects: Optional[List[VaspObject]] = Field( None, description="List of VASP objects included with this task document" ) vasp_objects: Optional[Dict[VaspObject, Any]] = Field( None, description="Vasp objects associated with this task" ) entry: Optional[ComputedEntry] = Field( None, description="The ComputedEntry from the task doc" ) task_label: Optional[str] = Field(None, description="A description of the task") author: Optional[str] = Field( None, description="Author extracted from transformations" ) icsd_id: Optional[Union[str, int]] = Field( None, description="Inorganic Crystal Structure Database id of the structure" ) transformations: Optional[Any] = Field( None, description="Information on the structural transformations, parsed from a " "transformations.json file", ) additional_json: Optional[Dict[str, Any]] = Field( None, description="Additional json loaded from the calculation directory" ) custodian: Optional[List[CustodianDoc]] = Field( None, title="Calcs reversed data", description="Detailed custodian data for each VASP calculation contributing to the task document.", ) analysis: Optional[AnalysisDoc] = Field( None, title="Calculation Analysis", description="Some analysis of calculation data after collection.", ) last_updated: Optional[datetime] = Field( utcnow(), description="Timestamp for the most recent calculation for this task document", ) batch_id: Optional[str] = Field( None, description="Identifier for this calculation; should provide rough information about the calculation origin and purpose.", ) # Note that private fields are needed because TaskDoc permits extra info # added to the model, unlike TaskDocument. Because of this, when pydantic looks up # attrs on the model, it searches for them in the model extra dict first, and if it # can't find them, throws an AttributeError. It does this before looking to see if the # class has that attr defined on it. # _structure_entry: Optional[ComputedStructureEntry] = PrivateAttr(None) def model_post_init(self, __context: Any) -> None: # Always refresh task_type, calc_type, run_type # See, e.g. https://github.com/materialsproject/emmet/issues/960 # where run_type's were set incorrectly in older versions of TaskDoc # only run if attributes containing input sets are available attrs = ["calcs_reversed", "input", "orig_inputs"] if not any(hasattr(self, attr) and getattr(self, attr) for attr in attrs): return # To determine task and run type, we search for input sets in this order # of precedence: calcs_reversed, inputs, orig_inputs inp_set = None inp_sets_to_check = [self.input, self.orig_inputs] if (calcs_reversed := getattr(self, "calcs_reversed", None)) is not None: inp_sets_to_check = [calcs_reversed[0].input] + inp_sets_to_check for inp_set in inp_sets_to_check: if inp_set is not None: self.task_type = task_type(inp_set) break # calcs_reversed needed below if calcs_reversed is not None: self.run_type = self._get_run_type(calcs_reversed) if inp_set is not None: self.calc_type = self._get_calc_type(calcs_reversed, inp_set) # TODO: remove after imposing TaskDoc schema on older tasks in collection if self.structure is None: self.structure = calcs_reversed[0].output.structure # Make sure that the datetime field is properly formatted # (Unclear when this is not the case, please leave comment if observed) @field_validator("last_updated", mode="before") @classmethod def last_updated_dict_ok(cls, v) -> datetime: return convert_datetime(cls, v) @field_validator("batch_id", mode="before") @classmethod def _validate_batch_id(cls, v) -> str: if v is not None: invalid_chars = set( char for char in v if (not char.isalnum()) and (char not in {"-", "_"}) ) if len(invalid_chars) > 0: raise ValueError( f"Invalid characters in batch_id: {' '.join(invalid_chars)}" ) return v @model_validator(mode="after") def set_entry(self) -> datetime: if ( not self.entry and self.calcs_reversed and getattr(self.calcs_reversed[0].output, "structure", None) ): self.entry = self.get_entry(self.calcs_reversed, self.task_id) return self @classmethod def from_directory( cls: Type[_T], dir_name: Union[Path, str], volumetric_files: Tuple[str, ...] = _VOLUMETRIC_FILES, store_additional_json: bool = True, additional_fields: Optional[Dict[str, Any]] = None, volume_change_warning_tol: float = 0.2, task_names: Optional[list[str]] = None, **vasp_calculation_kwargs, ) -> _T: """ Create a task document from a directory containing VASP files. Parameters ---------- dir_name The path to the folder containing the calculation outputs. store_additional_json Whether to store additional json files found in the calculation directory. volumetric_files Volumetric files to search for. additional_fields Dictionary of additional fields to add to output document. volume_change_warning_tol Maximum volume change allowed in VASP relaxations before the calculation is tagged with a warning. task_names Naming scheme for multiple calculations in on folder e.g. ["relax1","relax2"]. Can be subfolder or extension. **vasp_calculation_kwargs Additional parsing options that will be passed to the :obj:`.Calculation.from_vasp_files` function. Returns ------- TaskDoc A task document for the calculation. """ logger.info(f"Getting task doc in: {dir_name}") additional_fields = {} if additional_fields is None else additional_fields dir_name = Path(dir_name) task_files = _find_vasp_files( dir_name, volumetric_files=volumetric_files, task_names=task_names ) if len(task_files) == 0: raise FileNotFoundError("No VASP files found!") calcs_reversed = [] all_vasp_objects = [] for task_name, files in task_files.items(): calc_doc, vasp_objects = Calculation.from_vasp_files( dir_name, task_name, **files, **vasp_calculation_kwargs ) calcs_reversed.append(calc_doc) all_vasp_objects.append(vasp_objects) # Reverse the list of calculations in the order: newest calc is the first # To match with calcs_reversed, all_vasp_objects is also reversed. calcs_reversed.reverse() all_vasp_objects.reverse() analysis = AnalysisDoc.from_vasp_calc_docs( calcs_reversed, volume_change_warning_tol=volume_change_warning_tol ) transformations, icsd_id, tags, author = _parse_transformations(dir_name) custodian = _parse_custodian(dir_name) orig_inputs = _parse_orig_inputs(dir_name) additional_json = None if store_additional_json: additional_json = _parse_additional_json(dir_name) dir_name = get_uri(dir_name) # convert to full uri path # only store objects from last calculation # TODO: make this an option vasp_objects = all_vasp_objects[0] included_objects = None if vasp_objects: included_objects = list(vasp_objects.keys()) doc = cls.from_structure( structure=calcs_reversed[0].output.structure, meta_structure=calcs_reversed[0].output.structure, include_structure=True, dir_name=dir_name, calcs_reversed=calcs_reversed, analysis=analysis, transformations=transformations, custodian=custodian, orig_inputs=orig_inputs, additional_json=additional_json, icsd_id=icsd_id, tags=tags, author=author, completed_at=calcs_reversed[0].completed_at, input=InputDoc.from_vasp_calc_doc(calcs_reversed[-1]), output=OutputDoc.from_vasp_calc_doc( calcs_reversed[0], vasp_objects.get(VaspObject.TRAJECTORY), # type: ignore ), state=_get_state(calcs_reversed, analysis), run_stats=_get_run_stats(calcs_reversed), vasp_objects=vasp_objects, included_objects=included_objects, task_type=calcs_reversed[0].task_type, ) return doc.model_copy(update=additional_fields) @classmethod def from_vasprun( cls: Type[_T], path: Union[str, Path], additional_fields: Optional[Dict[str, Any]] = None, volume_change_warning_tol: float = 0.2, **vasp_calculation_kwargs, ) -> _T: """ Create a task document from a vasprun.xml file. This is not recommended and will raise warnings, since some necessary information is absent from the vasprun.xml file, such as per-atom magnetic moments. However, the majority of the TaskDoc will be complete. Parameters ---------- path The path to the vasprun.xml. additional_fields: Dict[str, Any] = None, volume_change_warning_tol Maximum volume change allowed in VASP relaxations before the calculation is tagged with a warning. **vasp_calculation_kwargs Additional parsing options that will be passed to the :obj:`.Calculation.from_vasp_files` function. Returns ------- TaskDoc A task document for the calculation. """ logger.info(f"Getting vasprun.xml at: {path}") path = Path(path) dir_name = path.resolve().parent calc = Calculation.from_vasprun(path, **vasp_calculation_kwargs) calcs_reversed = [calc] analysis = AnalysisDoc.from_vasp_calc_docs( calcs_reversed, volume_change_warning_tol=volume_change_warning_tol ) # assume orig_inputs are those stated in vasprun.xml orig_inputs = OrigInputs( incar=calc.input.incar, poscar=Poscar(calc.input.structure), kpoints=calc.input.kpoints, potcar=calc.input.potcar, ) doc = cls.from_structure( structure=calcs_reversed[0].output.structure, meta_structure=calcs_reversed[0].output.structure, include_structure=True, dir_name=get_uri(dir_name), calcs_reversed=calcs_reversed, analysis=analysis, orig_inputs=orig_inputs, completed_at=calcs_reversed[0].completed_at, input=InputDoc.from_vasp_calc_doc(calcs_reversed[-1]), output=OutputDoc.from_vasp_calc_doc(calcs_reversed[0]), state=_get_state(calcs_reversed, analysis), run_stats=None, vasp_objects={}, included_objects=[], task_type=calcs_reversed[0].task_type, ) if additional_fields: doc = doc.model_copy(update=additional_fields) return doc @staticmethod def get_entry( calcs_reversed: List[Calculation], task_id: Optional[Union[MPID, str]] = None ) -> ComputedEntry: """ Get a computed entry from a list of VASP calculation documents. Parameters ---------- calcs_reversed A list of VASP calculation documents in a reverse order. task_id The job identifier. Returns ------- ComputedEntry A computed entry. """ entry_dict = { "correction": 0.0, "entry_id": task_id, "composition": calcs_reversed[0].output.structure.composition, "energy": calcs_reversed[0].output.energy, "parameters": { # Cannot be PotcarSpec document, pymatgen expects a dict # Note that `potcar_spec` is optional "potcar_spec": ( [dict(d) for d in calcs_reversed[0].input.potcar_spec] if calcs_reversed[0].input.potcar_spec else [] ), # Required to be compatible with MontyEncoder for the ComputedEntry "run_type": str(calcs_reversed[0].run_type), "is_hubbard": calcs_reversed[0].input.is_hubbard, "hubbards": calcs_reversed[0].input.hubbards, }, "data": { "oxide_type": oxide_type(calcs_reversed[0].output.structure), "aspherical": calcs_reversed[0].input.parameters.get("LASPH", False), "last_updated": str(utcnow()), }, } return ComputedEntry.from_dict(entry_dict) @staticmethod def _get_calc_type( calcs_reversed: list[Calculation], orig_inputs: OrigInputs ) -> CalcType: """Get the calc type from calcs_reversed. Returns -------- CalcType The type of calculation. """ inputs = ( calcs_reversed[0].input.model_dump() if len(calcs_reversed) > 0 else orig_inputs ) params = calcs_reversed[0].input.parameters incar = calcs_reversed[0].input.incar return calc_type(inputs, {**params, **incar}) @staticmethod def _get_run_type(calcs_reversed: list[Calculation]) -> RunType: """Get the run type from calcs_reversed. Returns -------- RunType The type of calculation. """ params = calcs_reversed[0].input.parameters incar = calcs_reversed[0].input.incar return run_type({**params, **incar}) @property def structure_entry(self) -> ComputedStructureEntry: """ Retrieve a ComputedStructureEntry for this TaskDoc. Returns ------- ComputedStructureEntry The TaskDoc.entry with corresponding TaskDoc.structure added. """ return ComputedStructureEntry( structure=self.structure, energy=self.entry.energy, correction=self.entry.correction, composition=self.entry.composition, energy_adjustments=self.entry.energy_adjustments, parameters=self.entry.parameters, data=self.entry.data, entry_id=self.entry.entry_id, ) class TrajectoryDoc(BaseModel): """Model for task trajectory data.""" task_id: Optional[str] = Field( None, description="The (task) ID of this calculation, used as a universal reference across property documents." "This comes in the form: mp-******.", ) trajectories: Optional[List[Trajectory]] = Field( None, description="Trajectory data for calculations associated with a task doc.", ) class EntryDoc(BaseModel): """Model for task entry data.""" task_id: Optional[str] = Field( None, description="The (task) ID of this calculation, used as a universal reference across property documents." "This comes in the form: mp-******.", ) entry: Optional[ComputedStructureEntry] = Field( None, description="Computed structure entry for the calculation associated with the task doc.", ) class DeprecationDoc(BaseModel): """Model for task deprecation data.""" task_id: Optional[str] = Field( None, description="The (task) ID of this calculation, used as a universal reference across property documents." "This comes in the form: mp-******.", ) deprecated: Optional[bool] = Field( None, description="Whether the ID corresponds to a deprecated calculation.", ) deprecation_reason: Optional[str] = Field( None, description="Reason for deprecation.", ) def get_uri(dir_name: Union[str, Path]) -> str: """ Return the URI path for a directory. This allows files hosted on different file servers to have distinct locations. Parameters ---------- dir_name : str or Path A directory name. Returns ------- str Full URI path, e.g., "fileserver.host.com:/full/path/of/dir_name". """ import socket fullpath = Path(dir_name).absolute() hostname = socket.gethostname() try: hostname = socket.gethostbyaddr(hostname)[0] except (socket.gaierror, socket.herror): pass return f"{hostname}:{fullpath}" def _parse_transformations( dir_name: Path, ) -> Tuple[Dict, Optional[int], Optional[List[str]], Optional[str]]: """Parse transformations.json file.""" transformations = {} filenames = tuple(dir_name.glob("transformations.json*")) icsd_id = None if len(filenames) >= 1: transformations = loadfn(filenames[0], cls=None) try: match = re.match(r"(\d+)-ICSD", transformations["history"][0]["source"]) if match: icsd_id = int(match.group(1)) except (KeyError, IndexError): pass # We don't want to leave tags or authors in the # transformations file because they'd be copied into # every structure generated after this one. other_parameters = transformations.get("other_parameters", {}) new_tags = other_parameters.pop("tags", None) new_author = other_parameters.pop("author", None) if "other_parameters" in transformations and not other_parameters: # if dict is now empty remove it transformations.pop("other_parameters") return transformations, icsd_id, new_tags, new_author def _parse_custodian(dir_name: Path) -> Optional[Dict]: """ Parse custodian.json file. Calculations done using custodian have a custodian.json file which tracks the makers performed and any errors detected and fixed. Parameters ---------- dir_name Path to calculation directory. Returns ------- Optional[dict] The information parsed from custodian.json file. """ filenames = tuple(dir_name.glob("custodian.json*")) if len(filenames) >= 1: return loadfn(filenames[0], cls=None) return None def _parse_orig_inputs( dir_name: Path, ) -> Dict[str, Union[Kpoints, Poscar, PotcarSpec, Incar]]: """ Parse original input files. Calculations using custodian generate a *.orig file for the inputs. This is useful to know how the calculation originally started. Parameters ---------- dir_name Path to calculation directory. Returns ------- Dict[str, Union[Kpints, Poscar, PotcarSpec, Incar]] The original POSCAR, KPOINTS, POTCAR, and INCAR data. """ orig_inputs = {} input_mapping = { "INCAR": Incar, "KPOINTS": Kpoints, "POTCAR": VaspPotcar, "POSCAR": Poscar, } for filename in dir_name.glob("*.orig*"): for name, vasp_input in input_mapping.items(): if f"{name}.orig" in str(filename): if name == "POTCAR": # can't serialize POTCAR orig_inputs[name.lower()] = PotcarSpec.from_potcar( vasp_input.from_file(filename) ) else: orig_inputs[name.lower()] = vasp_input.from_file(filename) return orig_inputs def _parse_additional_json(dir_name: Path) -> Dict[str, Any]: """Parse additional json files in the directory.""" additional_json = {} for filename in dir_name.glob("*.json*"): key = filename.name.split(".")[0] # ignore FW.json(.gz) so jobflow doesn't try to parse prev_vasp_dir OutputReferences # was causing atomate2 MP workflows to fail with ValueError: Could not resolve reference # 7f5a7f14-464c-4a5b-85f9-8d11b595be3b not in store or cache # contact @janosh in case of questions if key not in ("custodian", "transformations", "FW"): additional_json[key] = loadfn(filename, cls=None) return additional_json def _get_max_force(calc_doc: Calculation) -> Optional[float]: """Get max force acting on atoms from a calculation document.""" if calc_doc.output.ionic_steps: forces: Optional[Union[np.ndarray, List]] = None if calc_doc.output.ionic_steps: forces = calc_doc.output.ionic_steps[-1].forces structure = calc_doc.output.structure if forces: forces = np.array(forces) sdyn = structure.site_properties.get("selective_dynamics") if sdyn: forces[np.logical_not(sdyn)] = 0 return max(np.linalg.norm(forces, axis=1)) return None def _get_drift_warnings(calc_doc: Calculation) -> List[str]: """Get warnings of whether the drift on atoms is too large.""" warnings = [] if calc_doc.input.parameters.get("NSW", 0) > 0: drift = calc_doc.output.outcar.get("drift", [[0, 0, 0]]) max_drift = max(np.linalg.norm(d) for d in drift) ediffg = calc_doc.input.parameters.get("EDIFFG", None) max_force = -float(ediffg) if ediffg and float(ediffg) < 0 else np.inf if max_drift > max_force: warnings.append( f"Drift ({drift}) > desired force convergence ({max_force}), structure " "likely not converged to desired accuracy." ) return warnings def _get_state(calcs_reversed: List[Calculation], analysis: AnalysisDoc) -> TaskState: """Get state from calculation documents and relaxation analysis.""" all_calcs_completed = all( c.has_vasp_completed == TaskState.SUCCESS for c in calcs_reversed ) if len(analysis.errors) == 0 and all_calcs_completed: return TaskState.SUCCESS # type: ignore return TaskState.FAILED # type: ignore def _get_run_stats(calcs_reversed: List[Calculation]) -> Dict[str, RunStatistics]: """Get summary of runtime statistics for each calculation in this task.""" run_stats = {} total = dict( average_memory=0.0, max_memory=0.0, elapsed_time=0.0, system_time=0.0, user_time=0.0, total_time=0.0, cores=0, ) for calc_doc in calcs_reversed: stats = calc_doc.output.run_stats run_stats[calc_doc.task_name] = stats total["average_memory"] = max(total["average_memory"], stats.average_memory) total["max_memory"] = max(total["max_memory"], stats.max_memory) total["cores"] = max(total["cores"], stats.cores) total["elapsed_time"] += stats.elapsed_time total["system_time"] += stats.system_time total["user_time"] += stats.user_time total["total_time"] += stats.total_time run_stats["overall"] = RunStatistics(**total) return run_stats def _find_vasp_files( path: Union[str, Path], volumetric_files: Tuple[str, ...] = _VOLUMETRIC_FILES, task_names: Optional[list[str]] = None, ) -> Dict[str, Any]: """ Find VASP files in a directory. Only files in folders with names matching a task name (or alternatively files with the task name as an extension, e.g., vasprun.relax1.xml) will be returned. VASP files in the current directory will be given the task name "standard". Parameters ---------- path Path to a directory to search. volumetric_files Volumetric files to search for. Returns ------- dict[str, Any] The filenames of the calculation outputs for each VASP task, given as a ordered dictionary of:: { task_name: { "vasprun_file": vasprun_filename, "outcar_file": outcar_filename, "contcar_file": contcar_filename, "volumetric_files": [CHGCAR, LOCPOT, etc] "elph_poscars": [POSCAR.T=300, POSCAR.T=400, etc] }, ... } """ task_names = ["precondition"] + [f"relax{i}" for i in range(9)] path = Path(path) task_files = OrderedDict() def _get_task_files(files, suffix=""): vasp_files = {} vol_files = [] elph_poscars = [] for file in files: file_no_path = file.relative_to(path) if file.match(f"*vasprun.xml{suffix}*"): vasp_files["vasprun_file"] = file_no_path elif file.match(f"*OUTCAR{suffix}*"): vasp_files["outcar_file"] = file_no_path elif file.match(f"*CONTCAR{suffix}*"): vasp_files["contcar_file"] = file_no_path elif any(file.match(f"*{f}{suffix}*") for f in volumetric_files): vol_files.append(file_no_path) elif file.match(f"*POSCAR.T=*{suffix}*"): elph_poscars.append(file_no_path) elif file.match(f"*OSZICAR{suffix}*"): vasp_files["oszicar_file"] = file_no_path if len(vol_files) > 0: # add volumetric files if some were found or other vasp files were found vasp_files["volumetric_files"] = vol_files if len(elph_poscars) > 0: # add elph displaced poscars if they were found or other vasp files found vasp_files["elph_poscars"] = elph_poscars return vasp_files for task_name in task_names: subfolder_match = list(path.glob(f"{task_name}/*")) suffix_match = list(path.glob(f"*.{task_name}*")) if len(subfolder_match) > 0: # subfolder match task_files[task_name] = _get_task_files(subfolder_match) elif len(suffix_match) > 0: # try extension schema task_files[task_name] = _get_task_files( suffix_match, suffix=f".{task_name}" ) if len(task_files) == 0: # get any matching file from the root folder standard_files = _get_task_files(list(path.glob("*"))) if len(standard_files) > 0: task_files["standard"] = standard_files return task_files