Source code for autojob.harvest.harvesters.gpaw

"""GPAW calculation harvesting utilities.

This module defines the :func:`.harvest_gpaw_results` function which
can be used to harvest GPAW results from a directory.

To ensure that outputs can be harvested from a directory, the GPAW log file
should be named with :data:`.GPAW_LOG` or :data:`.GPAW_LOG` should be set to
the value of GPAW output file.

.. code-block:: python

    >>> from autojob.harvest.harvesters.gpaw import GPAW_LOG
    >>> from autojob.tasks.calculation import CalculationInputs

    >>> calc_inputs = CalculationInputs(
    ...     calculator="gpaw",
    ...     calc_params={
    ...         "mode": "pw",
    ...         "txt": GPAW_LOG,
    ...     },
    ... )

Example:
    from pathlib import Path
    from autojob.harvest.harvesters.gpaw import harvest_gpaw_results

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

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

from ase.io.gpaw_out import read_gpaw_out

if TYPE_CHECKING:
    from ase import Atoms

logger = logging.getLogger(__name__)

# ! This is not a standard fileaname
# TODO: Replace with default GPAW log filename
GPAW_LOG = "gpaw-out.log"


# TODO: Set converged correctly


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

    Args:
        src: The directory containing the GPAW output file.

    Returns:
        A dictionary containing GPAW 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 GPAW calculation in directory: %s",
        src,
    )
    log_file = Path(src, GPAW_LOG)
    results: dict[str, Any] = {"forces": None, "energy": None}

    if log_file.exists():
        with log_file.open(mode="r", encoding="utf-8") as file:
            atoms: Atoms = read_gpaw_out(file, -1)  # type: ignore[no-untyped-call]

        data: dict[str, Any] = atoms.calc.results
        results["energy"] = data.pop("energy", None)
        results["forces"] = data.pop("forces", None)
        results["converged"] = False
        results["calculator_results"] = {**data, "atoms": atoms}

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

    return results