Source code for autojob.calculation.charge

"""Harvest charge analysis results from completed task directory."""

import logging
from pathlib import Path
from typing import Annotated
from typing import TypeVar

import ase
from ase.geometry import geometry
from pydantic import PlainSerializer
from pymatgen.command_line.chargemol_caller import ChargemolAnalysis
from pymatgen.core.periodic_table import Element
from pymatgen.io.vasp import Chgcar
from typing_extensions import TypedDict

from autojob.utils.files import extract_structure_name

logger = logging.getLogger(__name__)


DDEC_NET_ATOMIC_CHARGE_FILE = "DDEC6_even_tempered_net_atomic_charges.xyz"


logger = logging.getLogger(__name__)


_T = TypeVar("_T")

StringSerialized = Annotated[
    _T, PlainSerializer(lambda x: str(x), return_type=str)
]


class _BondedToDict(TypedDict):
    index: int
    element: StringSerialized[Element]
    bond_order: float
    direction: tuple[float, float, float]
    spin_polarization: float


class _BondOrders(TypedDict):
    element: StringSerialized[Element]
    bonded_to: list[_BondedToDict]
    bond_order_sum: float


BondOrderDict = dict[int, _BondOrders]




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

    atomic_vol: float
    min_dist: float
    charge: float
    x: float
    y: float
    z: float





class BaderDensity(TypedDict):
    """Bader density data."""

    data: list[list[float]]
    shift: list[float]
    dim: list[float]





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

    data: BaderData
    vacuum_volume: float
    vacuum_charge: float
    nelectrons: int
    chgcar: Chgcar
    atomic_densities: list[BaderDensity]





class DDEC6Analysis(TypedDict):
    """DDEC6 analysis data."""

    partial_charges: list[float]
    spin_moments: list[float]
    dipoles: list[float]
    rsquared_moments: list[float]
    rcubed_moments: list[float]
    rfourth_moments: list[float]
    bond_order_dict: BondOrderDict





def get_ddec6_index_map(dir_name: str | Path) -> list[int]:
    """Return a list of integers mapping DDEC6 indices to ASE indices.

    The ASE index of the atom at index i in the DDEC6 structure can be found
    as follows:

        index_map = get_ddec6_index_map(dir_name)
        index = index_map[i]

    Args:
        dir_name: The directory containing a calculation.
    """
    with (
        Path(dir_name)
        .joinpath("run.py")
        .open(mode="r", encoding="utf-8") as python_script
    ):
        input_traj = extract_structure_name(python_script).removeprefix("./")

    ase_atoms = ase.io.read(Path(dir_name).joinpath(input_traj))
    ase_atoms.center()

    xyz_file = Path(dir_name).joinpath(DDEC_NET_ATOMIC_CHARGE_FILE)
    ddec_atoms = ase.io.read(xyz_file)
    ddec_atoms.cell = ase_atoms.cell
    ddec_atoms.center()

    ddec6_index_map = []

    for atom in ddec_atoms:
        closest_ase_atom = ase_atoms[0]
        _, min_distance = geometry.get_distances(
            [atom.position],
            [ase_atoms[0].position],
            cell=ddec_atoms.cell,
            pbc=True,
        )

        for ase_atom in ase_atoms:
            _, distance = geometry.get_distances(
                [atom.position],
                [ase_atom.position],
                cell=ddec_atoms.cell,
                pbc=True,
            )
            if distance <= min_distance:
                closest_ase_atom = ase_atom
                min_distance = distance

        ddec6_index_map.append(closest_ase_atom.index)

    if len(ddec6_index_map) != len(set(ddec6_index_map)):
        msg = (
            f"DDEC6 map is incomplete for calculation in directory: {dir_name}"
        )
        raise RuntimeError(msg)

    return ddec6_index_map





def load_ddec6_data(dir_name: str | Path) -> DDEC6Analysis:
    """Extract the DDEC6 data from the job directory.

    Args:
        dir_name: The directory of the completed calculation.

    Returns:
        The DDEC6 data.  If no data is found, every value will be None.
    """
    logger.info(f"Loading DDEC6 data for {dir_name!s}")
    rsquared_moments = rcubed_moments = rfourth_moments = None
    dipoles = charges = spin_densities = bond_orders = None
    try:
        analysis = ChargemolAnalysis(path=dir_name, run_chargemol=False)
        ddec6_index_map = get_ddec6_index_map(dir_name)
        charges = [0.0] * len(ddec6_index_map)
        spin_densities = [0.0] * len(ddec6_index_map)
        bond_orders: BondOrderDict = {}
        dipoles = [0.0] * len(ddec6_index_map)

        for i_ddec, i_ase in enumerate(ddec6_index_map):
            charges[i_ase] = analysis.ddec_charges[i_ddec]
            spin_densities[i_ase] = analysis.ddec_spin_moments[i_ddec]
            bond_orders[i_ase] = analysis.bond_order_dict[i_ddec]
            dipoles[i_ase] = analysis.dipoles[i_ddec]
            rsquared_moments[i_ase] = analysis.ddec_rsquared_moments[i_ddec]
            rcubed_moments[i_ase] = analysis.ddec_rcubed_moments[i_ddec]
            rfourth_moments[i_ase] = analysis.ddec_rfourth_moments[i_ddec]

        for bond_order in bond_orders.values():
            for bonded_to in bond_order["bonded_to"]:
                bonded_to["index"] = ddec6_index_map[bonded_to["index"]]

        logger.info(f"Successfully loaded DDEC6 data for {dir_name!s}")

    except FileNotFoundError:
        logger.warning(f"Unable to load DDEC6 data for {dir_name!s}")

    return DDEC6Analysis(
        partial_charges=charges,
        spin_moments=spin_densities,
        dipoles=dipoles,
        rsquared_moments=rsquared_moments,
        rcubed_moments=rcubed_moments,
        rfourth_moments=rfourth_moments,
        bond_order_dict=bond_orders,
    )