Source code for autojob.harvest.harvesters.espresso

"""Espresso calculation harvesting utilities.

This module defines the :func:`.harvest_espresso_results` function which
can be used to harvest Quantum Espresso results from a directory.

Example:
    from pathlib import Path
    from autojob.harvest.harvesters.espresso import harvest_espresso_results

    outputs = harvest_espresso_results(Path.cwd())
"""

import logging
from pathlib import Path
from typing import TYPE_CHECKING
from typing import Any

from ase.io.espresso import read_espresso_out

if TYPE_CHECKING:
    from ase import Atoms

logger = logging.getLogger(__name__)

FILES_TO_CARRYOVER = ["pwscf.save"]
PWSCF_LOG = "pwscf.save"


# TODO: Set converged correctly


def harvest_espresso_results(src: str | Path) -> dict[str, Any]:
    """Harvest Quantum Espresso calculation results from a directory.

    Args:
        src: The directory containing the Espresso (PWscf) output files.

    Returns:
        A dictionary containing Espresso calculation outputs.

    Warning:
        That a calculation has converged must be confirmed manually as the
        value of the ``"converged"`` key is always set to False.
    """
    logger.debug(
        "Loading calculation outputs for Espresso calculation in directory: %s",
        src,
    )
    log_file = Path(src, PWSCF_LOG)
    results: dict[str, Any] = {"forces": None, "energy": None}

    if log_file.exists():
        atoms: Atoms = read_espresso_out(str(log_file))
        data: dict[str, Any] = atoms.calc.results
        data["efermi"] = atoms.calc.efermi
        results["energy"] = data.pop("energy")
        results["forces"] = data.pop("forces")
        results["converged"] = False
        results["calculator_results"] = {**data, "atoms": atoms}

        logger.debug(
            "Successfully loaded calculation outputs for Espresso calculation "
            "in directory: %s",
            src,
        )
    else:
        logger.warning(
            "Espresso (PWscf) output file %s does not exist in directory: %s",
            PWSCF_LOG,
            src,
        )

    return results