Source code for autojob.harvest.harvesters.bader

"""Bader charge analysis harvesting utilities.

This module defines the :func:`.harvest_bader_results` and
:func:`.parse_acf` functions and the :class:`.BaderAnalysis` class.
:func:`.harvest_bader_results` and :func:`.parse_acf` can be used to parse Bader
analysis data. The :class:`.BaderAnalysis` class is the data model for Bader
analysis data.

Example:
    from pathlib import Path
    from autojob.harvest.harvesters.bader import harvest_bader_results

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

import logging
from pathlib import Path
from typing import TypedDict

from pymatgen.io.vasp.outputs import Vasprun

logger = logging.getLogger(__name__)




class BaderAnalysis(TypedDict):
    """Bader change analysis data.

    Keys:
        positions: The positions of each Bader volume.
        min_dist: A list of floats specifying the minimum distances between
            each Bader volume and its nearest neighbor.
        charge: A list of floats specifying the charge containing in each Bader
            volume.
        atomic_volume: A list of floats specifying the size (in cubic
            Angstroms) of each Bader volume.
        charge_transfer: A list of floats specifying the net number of
            electrons in each Bader volume. Positive values indicate net
            **negative** charge for the atom centered in the Bader volume.
        vacuum_charge: The system's vacuum charge.
        vacuum_volume: The system's vacuum volume.
        nelectrons: The total number of electrons in the system.
    """

    positions: list[list[float]]
    min_dist: list[float]
    charge: list[float]
    atomic_volume: list[float]
    charge_transfer: list[float] | None
    vacuum_charge: float
    vacuum_volume: float
    nelectrons: float





def parse_acf(src: str | Path) -> BaderAnalysis:
    """Parse the Bader output file (``ACF.dat``).

    Args:
        src: The path to a directory containing Bader calculation files.

    Returns:
        A dictionary containing the location, charge, volume, and minimum
        distance (to other Bader volumes) for each Bader volume found by the
        Bader program. The dictionary also contains the vacuum charge, vacuum
        volume, and total number of electrons in the charge density.
    """
    # us-ascii encoding is used for compatibility with Bader program
    with Path(src, "ACF.dat").open(mode="r", encoding="us-ascii") as file:
        lines = file.readlines()

    data = BaderAnalysis(
        positions=[],
        charge=[],
        min_dist=[],
        atomic_volume=[],
        vacuum_charge=0.0,
        vacuum_volume=0.0,
        nelectrons=0.0,
        charge_transfer=None,
    )
    # Skip header lines
    lines.pop(0)
    lines.pop(0)

    should_continue = True

    while should_continue:
        line = lines.pop(0).strip()
        if line.startswith("-"):
            should_continue = False
        else:
            vals = list(map(float, line.split()[1:]))

            data["positions"].append(vals[:3])
            data["charge"].append(vals[3])
            data["min_dist"].append(vals[4])
            data["atomic_volume"].append(vals[5])

    for line in lines:
        tokens = line.strip().split(":")
        if tokens[0] == "VACUUM CHARGE":
            data["vacuum_charge"] = float(tokens[1])
        elif tokens[0] == "VACUUM VOLUME":
            data["vacuum_volume"] = float(tokens[1])
        elif tokens[0] == "NUMBER OF ELECTRONS":
            data["nelectrons"] = float(tokens[1])

    return data



def _calculate_charge_transfer(
    src: str | Path, charge_data: list[float]
) -> list[float]:
    """Calculate the net number of electrons for each Bader volume.

    Args:
        src: The path to a directory containing Bader calculation files. The
            directory must contain the ``ACF.dat`` and vasprun.xml files.
        charge_data: A list of floats representing the number of electrons in
            each Bader volume.

    Returns:
        A list of floats representing the net number of electrons for each
        Bader volume. Positive values indicate net negative charge. Negative
        values indicate net positive charge.
    """
    vasprun = Vasprun(Path(src, "vasprun.xml"))
    symbol_to_elect: dict[str, float] = {}
    potcars = vasprun.get_potcars(True)

    if potcars:
        for potcar in potcars:
            symbol_to_elect[potcar.symbol] = potcar.nelectrons

        nelects = [symbol_to_elect[s] for s in vasprun.atomic_symbols]
        charge_transfer = [c - nelects[i] for i, c in enumerate(charge_data)]

        logger.info(f"Successfully loaded Bader data for {src!s}")

    return charge_transfer




def harvest_bader_results(src: str | Path) -> BaderAnalysis:
    """Harvest VASP-derived Bader charge analysis results from a directory.

    Args:
        src: The path to a directory containing Bader calculation files. The
            directory must contain the ``ACF.dat`` file and, if charge transfer
            is desired, the vasprun.xml file as well.

    Returns:
        A :class:`.BaderAnalysis` object.

    Warning:
        When calculating charge transfer, this method assumes that the order
        of charge volumes identified by the Bader program corresponds to that
        of the atomic positions found in the vasprun.xml file.
    """
    logger.info(f"Loading Bader data for {src!s}")
    data = parse_acf(src=src)
    try:
        data["charge_transfer"] = _calculate_charge_transfer(
            src,
            data["charge"],
        )
    except FileNotFoundError:
        logger.warning(f"Unable to load Bader data for {src!s}")
        raise

    return BaderAnalysis(**data)