"""Atomate2 is a library of computational materials science workflows."""

from atomate2._version import __version__
from atomate2.settings import Atomate2Settings

SETTINGS = Atomate2Settings()

src/atomate2_temp/__init__.py
tpurcell90/atomate2-fhi-aims
"""Functions dealing with FHI-aims files"""
from __future__ import annotations

import logging
from glob import glob
from pathlib import Path
from typing import Sequence

from atomate2.common.files import copy_files, get_zfile, gunzip_files
from atomate2.utils.file_client import FileClient, auto_fileclient
from atomate2.utils.path import strip_hostname

from atomate2_temp.aims.sets.base import AimsInputGenerator
from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

logger = logging.getLogger(__name__)

__all__ = ["copy_aims_outputs", "write_aims_input_set", "cleanup_aims_outputs"]


@auto_fileclient
def copy_aims_outputs(
    src_dir: Path | str,
    src_host: str | None = None,
    additional_aims_files: list[str] | None = None,
    restart_to_input: bool = True,
    file_client: FileClient | None = None,
):
    """
    Copy FHI-aims output files to the current directory (inspired by CP2K plugin).

    Parameters
    ----------
    src_dir : str or Path
        The source directory.
    src_host : str or None
        The source hostname used to specify a remote filesystem. Can be given as
        either "username@remote_host" or just "remote_host" in which case the username
        will be inferred from the current user. If ``None``, the local filesystem will
        be used as the source.
    additional_aims_files : list of str
        Additional files to copy
    restart_to_input : bool
        Move the aims restart files to by the aims input in the new directory
    file_client : .FileClient
        A file client to use for performing file operations.
    """
    src_dir = strip_hostname(src_dir)
    logger.info(f"Copying FHI-aims inputs from {src_dir}")
    directory_listing = file_client.listdir(src_dir, host=src_host)
    # additional files like bands, DOS, *.cube, whatever
    additional_files = additional_aims_files if additional_aims_files else []

    if restart_to_input:
        additional_files += ("hessian.aims", "geometry.in.next_step", "*.csc")

    # copy files
    files = ["aims.out", "*.json"]

    for pattern in set(additional_files):
        for f in glob((Path(src_dir) / pattern).as_posix()):
            files.append(Path(f).name)

    all_files = [get_zfile(directory_listing, r, allow_missing=True) for r in files]
    all_files = [f for f in all_files if f]

    copy_files(
        src_dir,
        src_host=src_host,
        include_files=all_files,
        file_client=file_client,
    )

    zipped_files = [f for f in all_files if f.name.endswith("gz")]

    gunzip_files(
        include_files=zipped_files,
        allow_missing=True,
        file_client=file_client,
    )

    logger.info("Finished copying inputs")


def write_aims_input_set(
    structure: MSONableAtoms,
    input_set_generator: AimsInputGenerator,
    directory: str | Path = ".",
    prev_dir: str | Path | None = None,
    **kwargs,
):
    """
    Write FHI-aims input set.

    Parameters
    ----------
    structure : .MSONableAtoms
        A structure.
    input_set_generator : .AimsInputGenerator
        An GHI-aims input set generator.
    directory : str or Path
        The directory to write the input files to.
    prev_dir : str or Path or None
        If the input set is to be initialized from a previous calculation,
        the previous calc directory
    **kwargs
        Keyword arguments to pass to :obj:`.AimsInputSet.write_input`.
    """
    properties = kwargs.get("properties", [])
    aims_is = input_set_generator.get_input_set(
        structure, prev_dir=prev_dir, properties=properties
    )

    logger.info("Writing FHI-aims input set.")
    aims_is.write_input(directory, **kwargs)


@auto_fileclient
def cleanup_aims_outputs(
    directory: Path | str,
    host: str | None = None,
    file_patterns: Sequence[str] = (),
    file_client: FileClient | None = None,
):
    """
    Remove unnecessary files.

    Parameters
    ----------
    directory:
        Directory containing files
    host:
        File client host
    file_patterns:
        Glob patterns to find files for deletion.
    file_client:
        A file client to use for performing file operations.
    """
    files_to_delete = []
    for pattern in file_patterns:
        files_to_delete.extend(file_client.glob(Path(directory) / pattern, host=host))

    for file in files_to_delete:
        file_client.remove(file)

src/atomate2_temp/aims/files.py
tpurcell90/atomate2-fhi-aims
""" (Work)flows for FHI-aims
"""
from copy import deepcopy
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict

from jobflow import Flow, Maker
from pymatgen.core import Molecule, Structure

from atomate2_temp.aims.jobs.base import BaseAimsMaker
from atomate2_temp.aims.jobs.core import RelaxMaker
from atomate2_temp.aims.sets.core import RelaxSetGenerator
from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

__all__ = [
    "DoubleRelaxMaker",
]


@dataclass
class DoubleRelaxMaker(Maker):
    """
    A maker to perform a double relaxation in FHI-aims (first with light,
    and then with tight species_defaults)

    Parameters
    ----------
    name : str
        A name for the flow
    relax_maker1: .RelaxMaker
        A maker that generates the first relaxation
    relax_maker2: .RelaxMaker
        A maker that generates the second relaxation
    """

    name: str = "Double relaxation"
    relax_maker1: BaseAimsMaker = field(default_factory=RelaxMaker)
    relax_maker2: BaseAimsMaker = field(default_factory=RelaxMaker)

    def make(
        self,
        structure: MSONableAtoms | Structure | Molecule,
        prev_dir: str | Path | None = None,
    ) -> Flow:
        """
        Create a flow with two chained relaxations.

        Parameters
        ----------
        structure : .MSONableAtoms, Structure, Molecule
            The structure to relax.
        prev_dir : str or Path or None
            A previous FHI-aims calculation directory to copy output files from.
        """
        relax1 = self.relax_maker1.make(structure, prev_dir=prev_dir)
        relax1.name += " 1"

        relax2 = self.relax_maker2.make(
            relax1.output.structure, prev_dir=relax1.output.dir_name
        )
        relax2.name += " 2"

        return Flow([relax1, relax2], relax2.output, name=self.name)

    @classmethod
    def from_parameters(
        cls,
        parameters: Dict[str, any],
        species_defaults: list | tuple = ("light", "tight"),
    ):
        """
        Creates a DoubleRelaxFlow for the same parameters with two different
        species defaults.

        Parameters
        ----------
        parameters : dict
            a dictionary with calculation parameters
        species_defaults: list | tuple
            paths for species defaults to use relative to the given `species_dir`
            in parameters
        """
        # various checks
        if len(species_defaults) != 2:
            raise ValueError(
                "Two species defaults directories should be provided for DoubleRelaxFlow"
            )
        if "species_dir" not in parameters:
            raise KeyError("Provided parameters do not include species_dir")
        species_dir = Path(parameters["species_dir"])
        for basis_set in species_defaults:
            if not (species_dir / basis_set).exists():
                raise OSError(
                    f"The species defaults directory {(species_dir / basis_set).as_posix()} does not exist"
                )

        # now the actual work begins
        makers = []
        for basis_set in species_defaults:
            parameters["species_dir"] = (species_dir / basis_set).as_posix()
            input_set = RelaxSetGenerator(user_parameters=deepcopy(parameters))
            makers.append(RelaxMaker(input_set_generator=input_set))
        return cls(relax_maker1=makers[0], relax_maker2=makers[1])

src/atomate2_temp/aims/flows/core.py
tpurcell90/atomate2-fhi-aims
"""
GW workflows for FHI-aims with automatic convergence
"""
from copy import deepcopy
from dataclasses import dataclass, field
from pathlib import Path

from jobflow import Flow
from pymatgen.core import Molecule, Structure

from atomate2_temp.aims.jobs.base import BaseAimsMaker, ConvergenceMaker
from atomate2_temp.aims.jobs.core import BandStructureMaker, GWMaker, StaticMaker
from atomate2_temp.aims.sets.bs import BandStructureSetGenerator, GWSetGenerator
from atomate2_temp.aims.sets.core import StaticSetGenerator
from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

__all__ = [
    "PeriodicGWConvergenceMaker",
]


@dataclass
class PeriodicGWConvergenceMaker(BaseAimsMaker):
    """
    A maker to perform a GW workflow with automatic convergence in FHI-aims.

    Parameters
    ----------
    name : str
        A name for the job
    criterion_name: str
        A name for the convergence criterion. Must be in the run results
    epsilon: float
        A difference in criterion value for subsequent runs
    convergence_field: str
        An input parameter that changes to achieve convergence
    convergence_steps: list | tuple
        An iterable of the possible values for the convergence field.
        If the iterable is depleted and the convergence is not reached,
        that the job is failed
    """

    name: str = "GW convergence"
    criterion_name: str = "bandgap"
    epsilon: float = 0.001
    convergence_field: str = field(default_factory=str)
    convergence_steps: list = field(default_factory=list)

    def make(
        self,
        structure: MSONableAtoms | Structure | Molecule,
        prev_dir: str | Path | None = None,
    ) -> Flow:
        """
        Create a flow from the DFT ground state and subsequent GW calculation.

        Parameters
        ----------
        structure : .MSONableAtoms, Structure, or Molecule
            The strcuture to calculate
        prev_dir : str or Path or None
            A previous FHI-aims calculation directory to copy output files from.
        """

        parameters = self.input_set_generator.user_parameters
        parameters["elsi_restart"] = ["read_and_write", 1]
        # the first calculation
        if all(structure.pbc):
            input_set = BandStructureSetGenerator(user_parameters=deepcopy(parameters))
            static_maker = BandStructureMaker(
                input_set_generator=input_set, run_aims_kwargs=self.run_aims_kwargs
            )
        else:
            input_set = StaticSetGenerator(user_parameters=deepcopy(parameters))
            static_maker = StaticMaker(
                input_set_generator=input_set, run_aims_kwargs=self.run_aims_kwargs
            )

        static = static_maker.make(structure, prev_dir=prev_dir)

        gw_input_set = GWSetGenerator(user_parameters=parameters)
        gw_maker = GWMaker(
            input_set_generator=gw_input_set, run_aims_kwargs=self.run_aims_kwargs
        )
        convergence = ConvergenceMaker(
            maker=gw_maker,
            epsilon=self.epsilon,
            criterion_name=self.criterion_name,
            convergence_field=self.convergence_field,
            convergence_steps=self.convergence_steps,
        )
        gw = convergence.make(static.output.structure, prev_dir=static.output.dir_name)
        return Flow([static, gw], gw.output, name=self.name)

src/atomate2_temp/aims/flows/gw.py
tpurcell90/atomate2-fhi-aims
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from atomate2_temp.aims.jobs.base import BaseAimsMaker
from atomate2_temp.aims.jobs.core import StaticMaker, RelaxMaker
from atomate2_temp.aims.jobs.phonons import (
    PhononDisplacementMaker,
    PhononDisplacementMakerSocket,
)
from atomate2_temp.common.flows.phonons import BasePhononMaker


@dataclass
class PhononMaker(BasePhononMaker):
    """
    Maker to calculate harmonic phonons with FHI-aims and Phonopy.

    Overwrites the default Makers for the common PhononMaker

    Parameters
    ----------
    bulk_relax_maker : .BaseAimsMaker or None
        A maker to perform a tight relaxation on the bulk.
        Set to ``None`` to skip the
        bulk relaxation
    static_energy_maker : .BaseAimsMaker or None
        A maker to perform the computation of the DFT energy on the bulk.
        Set to ``None`` to skip the
        static energy computation
    born_maker: .BaseAimsMaker or None
        Maker to compute the BORN charges.
    phonon_displacement_maker : .BaseAimsMaker or None
        Maker used to compute the forces for a supercell.
    """

    code: str = "aims"
    bulk_relax_maker: BaseAimsMaker | None = field(
        default_factory=lambda: RelaxMaker.full_relaxation()
    )
    static_energy_maker: BaseAimsMaker | None = field(default_factory=StaticMaker)
    born_maker: BaseAimsMaker | None = None
    phonon_displacement_maker: BaseAimsMaker | None = None

    def __post_init__(self):
        if self.phonon_displacement_maker is None:
            if self.socket:
                self.phonon_displacement_maker = PhononDisplacementMakerSocket()
            else:
                self.phonon_displacement_maker = PhononDisplacementMaker()

src/atomate2_temp/aims/flows/phonons.py
tpurcell90/atomate2-fhi-aims
"""A representation of FHI-aims output (based on ASE output parser)"""
from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, Sequence, List

from emmet.core.math import Matrix3D, Vector3D

from monty.json import MontyDecoder, MSONable

from atomate2_temp.aims.io.parsers import read_aims_header_info, read_aims_output
from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms


class AimsOutput(MSONable):
    """The main output file for FHI-aims"""

    def __init__(
        self,
        results: Sequence[MSONableAtoms],
        metadata: Dict[str, Any],
        atoms_summary: Dict[str, Any],
    ):
        """AimsOutput object constructor

        Parameters
        ----------
        results: Sequence[MSONableAtoms]
            A list of all images in an output file
        metadata: Dict[str, Any]
            The metadata of the executable used to preform the calculation
        atoms_summary: Dict[str, Any]
            The summary of the starting atomic structure
        """
        self._results = results
        self._metadata = metadata
        self._atoms_summary = atoms_summary

    def as_dict(self) -> Dict:
        """Create a dict representation of the outputs for MSONable"""
        d = {"@module": self.__class__.__module__, "@class": self.__class__.__name__}

        d["results"] = self._results
        d["metadata"] = self._metadata
        d["atoms_summary"] = self._atoms_summary
        return d

    @classmethod
    def from_outfile(cls, outfile: str | Path):
        """
        AimsOutput object constructor.

        Parameters
        ----------
        outfile: str or Path
            The aims.out file to parse
        """
        metadata, atoms_summary = read_aims_header_info(outfile)
        results = read_aims_output(outfile, index=slice(0, None))

        return cls(results, metadata, atoms_summary)

    @classmethod
    def from_dict(cls, d: Dict[str, Any]):
        """Constructor from a dictionary"""
        decoded = {
            k: MontyDecoder().process_decoded(v)
            for k, v in d.items()
            if not k.startswith("@")
        }
        return cls(decoded["results"], decoded["metadata"], decoded["atoms_summary"])

    def get_results_for_image(
        self, image_ind: int | slice
    ) -> MSONableAtoms | Sequence[MSONableAtoms]:
        """Get the results dictionary for a particular image or slice of images

        Parameters
        ----------
        image_ind: int or slice
            The index of the images to get the results for

        Returns
        -------
        The results for those images
        """
        return self._results[image_ind]

    @property
    def atoms_summary(self) -> Dict[str, Any]:
        """The summary of the material/molecule that the calculations represent"""
        return self._atoms_summary

    @property
    def metadata(self) -> Dict[str, Any]:
        """The system metadata"""
        return self._metadata

    @property
    def n_images(self) -> int:
        """The number of images in results"""
        return len(self._results)

    @property
    def initial_structure(self) -> MSONableAtoms:
        """The initial structure for the calculations"""
        return self._atoms_summary["initial_atoms"]

    @property
    def final_structure(self) -> MSONableAtoms:
        """The final structure for the calculation"""
        return self._results[-1]

    @property
    def structures(self) -> Sequence[MSONableAtoms]:
        """All images in the output file"""
        return self._results

    @property
    def fermi_energy(self) -> float:
        """The Fermi energy for the final structure in the calculation"""
        return self.get_results_for_image(-1).calc.results["fermi_energy"]

    @property
    def homo(self) -> float:
        """The HOMO level for the final structure in the calculation"""
        return self.get_results_for_image(-1).calc.results["homo"]

    @property
    def lumo(self) -> float:
        """The LUMO level for the final structure in the calculation"""
        return self.get_results_for_image(-1).calc.results["lumo"]

    @property
    def band_gap(self) -> float:
        """The band gap for the final structure in the calculation"""
        return self.get_results_for_image(-1).calc.results["gap"]

    @property
    def direct_band_gap(self) -> float:
        """The direct band gap for the final structure in the calculation"""
        return self.get_results_for_image(-1).calc.results["direct_gap"]

    @property
    def final_energy(self) -> float:
        """The total energy for the final structure in the calculation"""
        return self.get_results_for_image(-1).calc.results["energy"]

    @property
    def completed(self) -> bool:
        """Did the calculation complete"""
        return len(self._results) > 0

    @property
    def aims_version(self) -> str:
        """The version of FHI-aims used for the calculation"""
        return self._metadata["version_number"]

    @property
    def forces(self) -> List[Vector3D]:
        """The forces for the final image of the calculation"""
        return self.get_results_for_image(-1).calc.results.get("forces", None)

    @property
    def stress(self) -> Matrix3D:
        """The stress for the final image of the calculation"""
        return self.get_results_for_image(-1).calc.results.get("stress", None)

    @property
    def stresses(self) -> List[Matrix3D]:
        """The virial stresses for the final image of the calculation"""
        return self.get_results_for_image(-1).calc.results.get("stresses", None)

    @property
    def all_forces(self) -> List[List[Vector3D]]:
        """The forces for all images in the calculation"""
        return [res.calc.results.get("forces", None) for res in self._results]

src/atomate2_temp/aims/io/aims_output.py
tpurcell90/atomate2-fhi-aims
"""AIMS output parser, taken from ASE with modifications"""
from __future__ import annotations

import warnings
from pathlib import Path
from typing import Any, Dict, Iterable

import numpy as np
from ase import Atom
from ase.calculators.singlepoint import SinglePointDFTCalculator
from ase.constraints import FixAtoms, FixCartesian
from ase.io import ParseError
from ase.utils import lazymethod, lazyproperty, reader

from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

LINE_NOT_FOUND = object()


class AimsParseError(Exception):
    """Exception raised if an error occurs when parsing an Aims output file"""

    def __init__(self, message):
        self.message = message
        super().__init__(self.message)


# Read aims.out files
scalar_property_to_line_key = {
    "free_energy": ["| Electronic free energy"],
    "number_of_iterations": ["| Number of self-consistency cycles"],
    "magnetic_moment": ["N_up - N_down"],
    "n_atoms": ["| Number of atoms"],
    "n_bands": [
        "Number of Kohn-Sham states",
        "Reducing total number of  Kohn-Sham states",
        "Reducing total number of Kohn-Sham states",
    ],
    "n_electrons": ["The structure contains"],
    "n_kpts": ["| Number of k-points"],
    "n_spins": ["| Number of spin channels"],
    "electronic_temp": ["Occupation type:"],
    "fermi_energy": ["| Chemical potential (Fermi level)"],
}


class AimsOutChunk:
    """Base class for AimsOutChunks"""

    def __init__(self, lines: Iterable[str]):
        """Constructor

        Parameters
        ----------
        lines: list of str
            The set of lines from the output file the encompasses either
            a single structure within a trajectory or
            general information about the calculation (header)
        """
        self.lines = lines

    def reverse_search_for(self, keys: Iterable[str], line_start: int = 0):
        """Find the last time one of the keys appears in self.lines

        Parameters
        ----------
        keys: list of str
            The key strings to search for in self.lines
        line_start: int
            The lowest index to search for in self.lines

        Returns
        -------
        int
            The last time one of the keys appears in self.lines
        """
        for ll, line in enumerate(self.lines[line_start:][::-1]):
            if any([key in line for key in keys]):
                return len(self.lines) - ll - 1

        return LINE_NOT_FOUND

    def search_for_all(self, key: str, line_start: int = 0, line_end: int = -1):
        """Find the all times the key appears in self.lines

        Parameters
        ----------
        key: str
            The key string to search for in self.lines
        line_start: int
            The first line to start the search from
        line_end: int
            The last line to end the search at

        Returns
        -------
        list of ints
            All times the key appears in the lines
        """
        line_index = []
        for ll, line in enumerate(self.lines[line_start:line_end]):
            if key in line:
                line_index.append(ll + line_start)
        return line_index

    def parse_scalar(self, property: str):
        """Parse a scalar property from the chunk

        Parameters
        ----------
        property: str
            The property key to parse

        Returns
        -------
        float
            The scalar value of the property
        """
        line_start = self.reverse_search_for(scalar_property_to_line_key[property])

        if line_start == LINE_NOT_FOUND:
            return None

        line = self.lines[line_start]
        return float(line.split(":")[-1].strip().split()[0])


class AimsOutHeaderChunk(AimsOutChunk):
    """The header of the aims.out file containint general information"""

    def __init__(self, lines: Iterable[str]):
        """Constructor

        Parameters
        ----------
        lines: list of str
            The lines inside the aims.out header
        """
        super().__init__(lines)
        self._k_points = None
        self._k_point_weights = None

    @property
    def commit_hash(self):
        """Get the commit hash for the FHI-aims version"""
        line_start = self.reverse_search_for(["Commit number"])
        if line_start == LINE_NOT_FOUND:
            raise AimsParseError("This file does not appear to be an aims-output file")

        return self.lines[line_start].split(":")[1].strip()

    @property
    def aims_uuid(self):
        """Get the aims-uuid for the calculation"""
        line_start = self.reverse_search_for(["aims_uuid"])
        if line_start == LINE_NOT_FOUND:
            raise AimsParseError("This file does not appear to be an aims-output file")

        return self.lines[line_start].split(":")[1].strip()

    @property
    def version_number(self):
        """Get the commit hash for the FHI-aims version"""
        line_start = self.reverse_search_for(["FHI-aims version"])
        if line_start == LINE_NOT_FOUND:
            raise AimsParseError("This file does not appear to be an aims-output file")

        return self.lines[line_start].split(":")[1].strip()

    @property
    def fortran_compiler(self):
        """Get the fortran compiler used to make FHI-aims"""
        line_start = self.reverse_search_for(["Fortran compiler      :"])
        if line_start == LINE_NOT_FOUND:
            raise AimsParseError("This file does not appear to be an aims-output file")

        return self.lines[line_start].split(":")[1].split("/")[-1].strip()

    @property
    def c_compiler(self):
        """Get the C compiler used to make FHI-aims"""
        line_start = self.reverse_search_for(["C compiler            :"])
        if line_start == LINE_NOT_FOUND:
            return None

        return self.lines[line_start].split(":")[1].split("/")[-1].strip()

    @property
    def fortran_compiler_flags(self):
        """Get the fortran compiler flags used to make FHI-aims"""
        line_start = self.reverse_search_for(["Fortran compiler flags"])
        if line_start == LINE_NOT_FOUND:
            raise AimsParseError("This file does not appear to be an aims-output file")

        return self.lines[line_start].split(":")[1].strip()

    @property
    def c_compiler_flags(self):
        """Get the C compiler flags used to make FHI-aims"""
        line_start = self.reverse_search_for(["C compiler flags"])
        if line_start == LINE_NOT_FOUND:
            return None

        return self.lines[line_start].split(":")[1].strip()

    @property
    def build_type(self):
        """Get the optional build flags passed to cmake"""
        line_end = self.reverse_search_for(["Linking against:"])
        line_inds = self.search_for_all("Using", line_end=line_end)

        build_type_flags = [
            " ".join(self.lines[ind].split()[1:]).strip() for ind in line_inds
        ]
        return build_type_flags

    @property
    def linked_against(self):
        """Get all libraries used to link the FHI-aims executable"""
        line_start = self.reverse_search_for(["Linking against:"])
        if line_start == LINE_NOT_FOUND:
            return []

        linked_libs = [self.lines[line_start].split(":")[1].strip()]
        line_start += 1
        while "lib" in self.lines[line_start]:
            linked_libs.append(self.lines[line_start].strip())
            line_start += 1

        return linked_libs

    @property
    def constraints(self):
        """Parse the constraints from the aims.out file

        Constraints for the lattice vectors are not supported.
        """

        line_inds = self.search_for_all("Found relaxation constraint for atom")
        if len(line_inds) == 0:
            return []

        fix = []
        fix_cart = []
        for ll in line_inds:
            line = self.lines[ll]
            xyz = [0, 0, 0]
            ind = int(line.split()[5][:-1]) - 1
            if "All coordinates fixed" in line:
                if ind not in fix:
                    fix.append(ind)
            if "coordinate fixed" in line:
                coord = line.split()[6]
                if coord == "x":
                    xyz[0] = 1
                elif coord == "y":
                    xyz[1] = 1
                elif coord == "z":
                    xyz[2] = 1
                keep = True
                for n, c in enumerate(fix_cart):
                    if ind == c.index:
                        keep = False
                        break
                if keep:
                    fix_cart.append(FixCartesian(ind, xyz))
                else:
                    fix_cart[n].mask[xyz.index(1)] = 0
        if len(fix) > 0:
            fix_cart.append(FixAtoms(indices=fix))

        return fix_cart

    @property
    def initial_cell(self):
        """Parse the initial cell from the aims.out file"""
        line_start = self.reverse_search_for(["| Unit cell:"])
        if line_start == LINE_NOT_FOUND:
            return None

        return [
            [float(inp) for inp in line.split()[-3:]]
            for line in self.lines[line_start + 1 : line_start + 4]
        ]

    @property
    def initial_atoms(self):
        """Create an atoms object for the initial geometry.in structure
        from the aims.out file"""
        line_start = self.reverse_search_for(["Atomic structure:"])
        if line_start == LINE_NOT_FOUND:
            raise AimsParseError("No information about the structure in the chunk.")

        line_start += 2

        cell = self.initial_cell
        positions = np.zeros((self.n_atoms, 3))
        symbols = [""] * self.n_atoms
        for ll, line in enumerate(self.lines[line_start : line_start + self.n_atoms]):
            inp = line.split()
            positions[ll, :] = [float(pos) for pos in inp[4:7]]
            symbols[ll] = inp[3]

        atoms = MSONableAtoms(symbols=symbols, positions=positions)

        if cell:
            atoms.set_cell(cell)
            atoms.set_pbc([True, True, True])
        atoms.set_constraint(self.constraints)

        return atoms

    @property
    def is_md(self):
        """Determine if calculation is a molecular dynamics calculation"""
        return LINE_NOT_FOUND != self.reverse_search_for(
            ["Complete information for previous time-step:"]
        )

    @property
    def is_relaxation(self):
        """Determine if the calculation is a geometry optimization or not"""
        return LINE_NOT_FOUND != self.reverse_search_for(["Geometry relaxation:"])

    def _parse_k_points(self):
        """Get the list of k-points used in the calculation"""
        n_kpts = self.parse_scalar("n_kpts")
        if n_kpts is None:
            return {
                "k_points": None,
                "k_point_weights": None,
            }
        n_kpts = int(n_kpts)

        line_start = self.reverse_search_for(["| K-points in task"])
        line_end = self.reverse_search_for(["| k-point:"])
        if (
            (line_start == LINE_NOT_FOUND)
            or (line_end == LINE_NOT_FOUND)
            or (line_end - line_start != n_kpts)
        ):
            return {
                "k_points": None,
                "k_point_weights": None,
            }

        k_points = np.zeros((n_kpts, 3))
        k_point_weights = np.zeros((n_kpts))
        for kk, line in enumerate(self.lines[line_start + 1 : line_end + 1]):
            k_points[kk] = [float(inp) for inp in line.split()[4:7]]
            k_point_weights[kk] = float(line.split()[-1])

        return {
            "k_points": k_points,
            "k_point_weights": k_point_weights,
        }

    @property
    def n_atoms(self) -> int:
        """The number of atoms for the material"""
        n_atoms = self.parse_scalar("n_atoms")
        if n_atoms is None:
            raise AimsParseError(
                "No information about the number of atoms in the header."
            )
        return int(n_atoms)

    @property
    def n_bands(self):
        """The number of Kohn-Sham states for the chunk"""
        line_start = self.reverse_search_for(scalar_property_to_line_key["n_bands"])

        if line_start == LINE_NOT_FOUND:
            raise AimsParseError(
                "No information about the number of Kohn-Sham states " "in the header."
            )

        line = self.lines[line_start]
        if "| Number of Kohn-Sham states" in line:
            return int(line.split(":")[-1].strip().split()[0])

        return int(line.split()[-1].strip()[:-1])

    @property
    def n_electrons(self):
        """The number of electrons for the chunk"""
        line_start = self.reverse_search_for(scalar_property_to_line_key["n_electrons"])

        if line_start == LINE_NOT_FOUND:
            raise AimsParseError(
                "No information about the number of electrons in the header."
            )

        line = self.lines[line_start]
        return int(float(line.split()[-2]))

    @property
    def n_k_points(self):
        """The number of k_ppoints for the calculation"""
        n_kpts = self.parse_scalar("n_kpts")
        if n_kpts is None:
            return None

        return int(n_kpts)

    @property
    def n_spins(self):
        """The number of spin channels for the chunk"""
        n_spins = self.parse_scalar("n_spins")
        if n_spins is None:
            raise AimsParseError(
                "No information about the number of spin " "channels in the header."
            )
        return int(n_spins)

    @property
    def electronic_temperature(self):
        """The electronic temperature for the chunk"""
        line_start = self.reverse_search_for(
            scalar_property_to_line_key["electronic_temp"]
        )
        if line_start == LINE_NOT_FOUND:
            return 0.10

        line = self.lines[line_start]
        return float(line.split("=")[-1].strip().split()[0])

    @property
    def k_points(self):
        """All k-points listed in the calculation"""
        return self._parse_k_points()["k_points"]

    @property
    def k_point_weights(self):
        """The k-point weights for the calculation"""
        return self._parse_k_points()["k_point_weights"]

    @property
    def header_summary(self):
        """Dictionary summarizing the information inside the header"""
        return {
            "initial_atoms": self.initial_atoms,
            "initial_cell": self.initial_cell,
            "constraints": self.constraints,
            "is_relaxation": self.is_relaxation,
            "is_md": self.is_md,
            "n_atoms": self.n_atoms,
            "n_bands": self.n_bands,
            "n_electrons": self.n_electrons,
            "n_spins": self.n_spins,
            "electronic_temperature": self.electronic_temperature,
            "n_k_points": self.n_k_points,
            "k_points": self.k_points,
            "k_point_weights": self.k_point_weights,
        }

    @property
    def metadata_summary(self) -> Dict[str, str]:
        """Dictionary containing all metadata for FHI-aims build"""
        return {
            "commit_hash": self.commit_hash,
            "aims_uuid": self.aims_uuid,
            "version_number": self.version_number,
            "fortran_compiler": self.fortran_compiler,
            "c_compiler": self.c_compiler,
            "fortran_compiler_flags": self.fortran_compiler_flags,
            "c_compiler_flags": self.c_compiler_flags,
            "build_type": self.build_type,
            "linked_against": self.linked_against,
        }


class AimsOutCalcChunk(AimsOutChunk):
    """A part of the aims.out file correponding to a single structure"""

    def __init__(self, lines, header):
        """Constructor

        Parameters
        ----------
        lines: list of str
            The lines used for the structure
        header: AimsOutHeaderChunk
            A summary of the relevant information from the aims.out header
        """
        super().__init__(lines)
        self._header = header.header_summary

    def _parse_atoms(self):
        """Create an atoms object for the subsequent structures
        calculated in the aims.out file"""
        start_keys = [
            "Atomic structure (and velocities) as used in the preceding " "time step",
            "Updated atomic structure",
            "Atomic structure that was used in the preceding time step of "
            "the wrapper",
        ]
        line_start = self.reverse_search_for(start_keys)
        if line_start == LINE_NOT_FOUND:
            return self.initial_atoms

        line_start += 1

        line_end = self.reverse_search_for(
            ['Writing the current geometry to file "geometry.in.next_step"'], line_start
        )
        if line_end == LINE_NOT_FOUND:
            line_end = len(self.lines)

        cell = []
        velocities = []
        atoms = MSONableAtoms()
        for line in self.lines[line_start:line_end]:
            if "lattice_vector   " in line:
                cell.append([float(inp) for inp in line.split()[1:]])
            elif "atom   " in line:
                line_split = line.split()
                atoms.append(
                    Atom(line_split[4], tuple([float(inp) for inp in line_split[1:4]]))
                )
            elif "velocity   " in line:
                velocities.append([float(inp) for inp in line.split()[1:]])

        if len(atoms) != self.n_atoms:
            raise AimsParseError(
                "The number of atoms is inconsistent with the initial structure"
            )

        if (len(velocities) != self.n_atoms) and (len(velocities) != 0):
            raise AimsParseError(
                "The number of velocities is inconsistent with the number of atoms"
            )

        if len(cell) == 3:
            atoms.set_cell(np.array(cell))
            atoms.set_pbc([True, True, True])
        elif len(cell) != 0:
            raise AimsParseError(
                "Parsed geometry has incorrect number of lattice vectors."
            )

        if len(velocities) > 0:
            atoms.set_velocities(np.array(velocities))
        atoms.set_constraint(self.constraints)

        return atoms

    @property
    def forces(self):
        """Parse the forces from the aims.out file"""
        line_start = self.reverse_search_for(["Total atomic forces"])
        if line_start == LINE_NOT_FOUND:
            return

        line_start += 1

        return np.array(
            [
                [float(inp) for inp in line.split()[-3:]]
                for line in self.lines[line_start : line_start + self.n_atoms]
            ]
        )

    @property
    def stresses(self):
        """Parse the stresses from the aims.out file"""
        line_start = self.reverse_search_for(
            ["Per atom stress (eV) used for heat flux calculation"]
        )
        if line_start == LINE_NOT_FOUND:
            return None
        line_start += 3
        stresses = []
        for line in self.lines[line_start : line_start + self.n_atoms]:
            xx, yy, zz, xy, xz, yz = [float(d) for d in line.split()[2:8]]
            stresses.append([xx, yy, zz, yz, xz, xy])

        return np.array(stresses)

    @property
    def stress(self):
        """Parse the stress from the aims.out file"""
        from ase.stress import full_3x3_to_voigt_6_stress

        line_start = self.reverse_search_for(
            [
                "Analytical stress tensor - Symmetrized",
                "Numerical stress tensor",
            ]
        )  # Offest to relevant lines
        if line_start == LINE_NOT_FOUND:
            return

        stress = [
            [float(inp) for inp in line.split()[2:5]]
            for line in self.lines[line_start + 5 : line_start + 8]
        ]
        return full_3x3_to_voigt_6_stress(stress)

    @property
    def is_metallic(self):
        """Checks the outputfile to see if the chunk corresponds
        to a metallic system"""
        line_start = self.reverse_search_for(
            [
                "material is metallic within the approximate finite "
                "broadening function (occupation_type)"
            ]
        )
        return line_start != LINE_NOT_FOUND

    @property
    def energy(self):
        """Parse the energy from the aims.out file"""
        atoms = self._parse_atoms()

        if np.all(atoms.pbc) and self.is_metallic:
            line_ind = self.reverse_search_for(["Total energy corrected"])
        else:
            line_ind = self.reverse_search_for(["Total energy uncorrected"])
        if line_ind == LINE_NOT_FOUND:
            raise AimsParseError("No energy is associated with the structure.")

        return float(self.lines[line_ind].split()[5])

    @lazyproperty
    def dipole(self):
        """Parse the electric dipole moment from the aims.out file."""
        line_start = self.reverse_search_for(["Total dipole moment [eAng]"])
        if line_start == LINE_NOT_FOUND:
            return

        line = self.lines[line_start]
        return np.array([float(inp) for inp in line.split()[6:9]])

    @lazyproperty
    def dielectric_tensor(self):
        """Parse the dielectric tensor from the aims.out file"""
        line_start = self.reverse_search_for(["PARSE DFPT_dielectric_tensor"])
        if line_start == LINE_NOT_FOUND:
            return

        # we should find the tensor in the next three lines:
        lines = self.lines[line_start + 1 : line_start + 4]

        # make ndarray and return
        return np.array([np.fromstring(line, sep=" ") for line in lines])

    @lazyproperty
    def polarization(self):
        """Parse the polarization vector from the aims.out file"""
        line_start = self.reverse_search_for(["| Cartesian Polarization"])
        if line_start == LINE_NOT_FOUND:
            return
        line = self.lines[line_start]
        return np.array([float(s) for s in line.split()[-3:]])

    @lazymethod
    def _parse_homo_lumo(self):
        """Parse the HOMO/LUMO values and get band gap if periodic"""
        line_start = self.reverse_search_for(["Highest occupied state (VBM)"])
        homo = float(self.lines[line_start].split(" at ")[1].split("eV")[0].strip())

        line_start = self.reverse_search_for(["Lowest unoccupied state (CBM)"])
        lumo = float(self.lines[line_start].split(" at ")[1].split("eV")[0].strip())

        line_start = self.reverse_search_for(["verall HOMO-LUMO gap"])
        homo_lumo_gap = float(
            self.lines[line_start].split(":")[1].split("eV")[0].strip()
        )

        line_start = self.reverse_search_for(["Smallest direct gap"])
        if line_start == LINE_NOT_FOUND:
            return {
                "homo": homo,
                "lumo": lumo,
                "gap": homo_lumo_gap,
                "direct_gap": homo_lumo_gap,
            }

        direct_gap = float(self.lines[line_start].split(":")[1].split("eV")[0].strip())
        return {
            "homo": homo,
            "lumo": lumo,
            "gap": homo_lumo_gap,
            "direct_gap": direct_gap,
        }

    @lazymethod
    def _parse_hirshfeld(self):
        """Parse the Hirshfled charges volumes, and dipole moments from the
        ouput"""
        atoms = self._parse_atoms()

        line_start = self.reverse_search_for(
            ["Performing Hirshfeld analysis of fragment charges and moments."]
        )
        if line_start == LINE_NOT_FOUND:
            return {
                "charges": None,
                "volumes": None,
                "atomic_dipoles": None,
                "dipole": None,
            }

        line_inds = self.search_for_all("Hirshfeld charge", line_start, -1)
        hirshfeld_charges = np.array(
            [float(self.lines[ind].split(":")[1]) for ind in line_inds]
        )

        line_inds = self.search_for_all("Hirshfeld volume", line_start, -1)
        hirshfeld_volumes = np.array(
            [float(self.lines[ind].split(":")[1]) for ind in line_inds]
        )

        line_inds = self.search_for_all("Hirshfeld dipole vector", line_start, -1)
        hirshfeld_atomic_dipoles = np.array(
            [
                [float(inp) for inp in self.lines[ind].split(":")[1].split()]
                for ind in line_inds
            ]
        )

        if not np.any(atoms.pbc):
            hirshfeld_dipole = np.sum(
                hirshfeld_charges.reshape((-1, 1)) * atoms.get_positions(),
                axis=1,
            )
        else:
            hirshfeld_dipole = None
        return {
            "charges": hirshfeld_charges,
            "volumes": hirshfeld_volumes,
            "atomic_dipoles": hirshfeld_atomic_dipoles,
            "dipole": hirshfeld_dipole,
        }

    @lazymethod
    def _parse_eigenvalues(self):
        """Parse the eigenvalues and occupancies of the system. If eigenvalue
        for a particular k-point is not present in the output file
        then set it to np.nan
        """

        atoms = self._parse_atoms()

        line_start = self.reverse_search_for(["Writing Kohn-Sham eigenvalues."])
        if line_start == LINE_NOT_FOUND:
            return {"eigenvalues": None, "occupancies": None}

        line_end_1 = self.reverse_search_for(
            ["Self-consistency cycle converged."], line_start
        )
        line_end_2 = self.reverse_search_for(
            [
                "What follows are estimated values for band gap, " "HOMO, LUMO, etc.",
                "Current spin moment of the entire structure :",
                "Highest occupied state (VBM)",
            ],
            line_start,
        )
        if line_end_1 == LINE_NOT_FOUND:
            line_end = line_end_2
        elif line_end_2 == LINE_NOT_FOUND:
            line_end = line_end_1
        else:
            line_end = min(line_end_1, line_end_2)

        n_kpts = self.n_k_points if np.all(atoms.pbc) else 1
        if n_kpts is None:
            return {"eigenvalues": None, "occupancies": None}

        eigenvalues = np.full((n_kpts, self.n_bands, self.n_spins), np.nan)
        occupancies = np.full((n_kpts, self.n_bands, self.n_spins), np.nan)

        occupation_block_start = self.search_for_all(
            "State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]",
            line_start,
            line_end,
        )
        kpt_def = self.search_for_all("K-point: ", line_start, line_end)

        if len(kpt_def) > 0:
            kpt_inds = [int(self.lines[ll].split()[1]) - 1 for ll in kpt_def]
        elif (self.n_k_points is None) or (self.n_k_points == 1):
            kpt_inds = [0]
        else:
            raise ParseError("Cannot find k-point definitions")

        if len(kpt_inds) != len(occupation_block_start):
            raise AimsParseError(
                "Number of k-points is not the same of the number of occupation blocks that exist."
            )
        spins = [0] * len(occupation_block_start)

        if self.n_spins == 2:
            spin_def = self.search_for_all("Spin-", line_start, line_end)
            if len(spin_def) != len(occupation_block_start):
                raise AimsParseError(
                    "The number of spin definitions is not the same of the number of occupations"
                )

            spins = [int("Spin-down eigenvalues:" in self.lines[ll]) for ll in spin_def]

        for occ_start, kpt_ind, spin in zip(occupation_block_start, kpt_inds, spins):
            for ll, line in enumerate(
                self.lines[occ_start + 1 : occ_start + self.n_bands + 1]
            ):
                if "***" in line:
                    warn_msg = f"The {ll+1}th eigenvalue for the "
                    "{kpt_ind+1}th k-point and {spin}th channels could "
                    "not be read (likely too large to be printed "
                    "in the output file)"
                    warnings.warn(warn_msg)
                    continue
                split_line = line.split()
                eigenvalues[kpt_ind, ll, spin] = float(split_line[3])
                occupancies[kpt_ind, ll, spin] = float(split_line[1])
        return {"eigenvalues": eigenvalues, "occupancies": occupancies}

    @lazyproperty
    def atoms(self):
        """Convert AimsOutChunk to Atoms object and add all non-standard
        outputs to atoms.info"""
        atoms = self._parse_atoms()

        atoms.calc = SinglePointDFTCalculator(atoms)
        atoms.calc.results = self.results
        return atoms

    @property
    def results(self):
        """Convert an AimsOutChunk to a Results Dictionary"""
        results = {
            "energy": self.energy,
            "free_energy": self.free_energy,
            "forces": self.forces,
            "stress": self.stress,
            "stresses": self.stresses,
            "magmom": self.magmom,
            "dipole": self.dipole,
            "fermi_energy": self.E_f,
            "n_iter": self.n_iter,
            "hirshfeld_charges": self.hirshfeld_charges,
            "hirshfeld_dipole": self.hirshfeld_dipole,
            "hirshfeld_volumes": self.hirshfeld_volumes,
            "hirshfeld_atomic_dipoles": self.hirshfeld_atomic_dipoles,
            "dielectric_tensor": self.dielectric_tensor,
            "polarization": self.polarization,
            "homo": self.homo,
            "lumo": self.lumo,
            "gap": self.gap,
            "direct_gap": self.direct_gap,
        }

        return {key: value for key, value in results.items() if value is not None}

    # Properties from the aims.out header
    @lazyproperty
    def initial_atoms(self):
        """The initial structure defined in the geoemtry.in file"""
        return self._header["initial_atoms"]

    @lazyproperty
    def initial_cell(self):
        """The initial lattice vectors defined in the geoemtry.in file"""
        return self._header["initial_cell"]

    @lazyproperty
    def constraints(self):
        """The relaxation constraints for the calculation"""
        return self._header["constraints"]

    @lazyproperty
    def n_atoms(self):
        """The number of atoms for the material"""
        return self._header["n_atoms"]

    @lazyproperty
    def n_bands(self):
        """The number of Kohn-Sham states for the chunk"""
        return self._header["n_bands"]

    @lazyproperty
    def n_electrons(self):
        """The number of electrons for the chunk"""
        return self._header["n_electrons"]

    @lazyproperty
    def n_spins(self):
        """The number of spin channels for the chunk"""
        return self._header["n_spins"]

    @lazyproperty
    def electronic_temperature(self):
        """The electronic temperature for the chunk"""
        return self._header["electronic_temperature"]

    @lazyproperty
    def n_k_points(self):
        """The number of electrons for the chunk"""
        return self._header["n_k_points"]

    @lazyproperty
    def k_points(self):
        """The number of spin channels for the chunk"""
        return self._header["k_points"]

    @lazyproperty
    def k_point_weights(self):
        """k_point_weights electronic temperature for the chunk"""
        return self._header["k_point_weights"]

    @lazyproperty
    def free_energy(self):
        """The free energy for the chunk"""
        return self.parse_scalar("free_energy")

    @lazyproperty
    def n_iter(self):
        """The number of SCF iterations needed to converge the SCF cycle for
        the chunk"""
        return self.parse_scalar("number_of_iterations")

    @lazyproperty
    def magmom(self):
        """The magnetic moment for the chunk"""
        return self.parse_scalar("magnetic_moment")

    @lazyproperty
    def E_f(self):
        """The Fermi energy for the chunk"""
        return self.parse_scalar("fermi_energy")

    @lazyproperty
    def converged(self):
        """True if the chunk is a fully converged final structure"""
        return (len(self.lines) > 0) and ("Have a nice day." in self.lines[-5:])

    @lazyproperty
    def hirshfeld_charges(self):
        """The Hirshfeld charges for the chunk"""
        return self._parse_hirshfeld()["charges"]

    @lazyproperty
    def hirshfeld_atomic_dipoles(self):
        """The Hirshfeld atomic dipole moments for the chunk"""
        return self._parse_hirshfeld()["atomic_dipoles"]

    @lazyproperty
    def hirshfeld_volumes(self):
        """The Hirshfeld volume for the chunk"""
        return self._parse_hirshfeld()["volumes"]

    @lazyproperty
    def hirshfeld_dipole(self):
        """The Hirshfeld systematic dipole moment for the chunk"""
        atoms = self._parse_atoms()

        if not np.any(atoms.pbc):
            return self._parse_hirshfeld()["dipole"]

        return None

    @lazyproperty
    def eigenvalues(self):
        """All outputted eigenvalues for the system"""
        return self._parse_eigenvalues()["eigenvalues"]

    @lazyproperty
    def occupancies(self):
        """All outputted occupancies for the system"""
        return self._parse_eigenvalues()["occupancies"]

    @lazyproperty
    def homo(self):
        """The HOMO (CBM) of the calculation"""
        return self._parse_homo_lumo()["homo"]

    @lazyproperty
    def lumo(self):
        """The LUMO (VBM) of the calculation"""
        return self._parse_homo_lumo()["lumo"]

    @lazyproperty
    def gap(self):
        """The HOMO-LUMO gap (band gap) of the calculation"""
        return self._parse_homo_lumo()["gap"]

    @lazyproperty
    def direct_gap(self):
        """The direct band gap of the calculation"""
        return self._parse_homo_lumo()["direct_gap"]


def get_header_chunk(fd):
    """Returns the header information from the aims.out file"""
    header = []
    line = ""

    # Stop the header once the first SCF cycle begins
    while (
        "Convergence:    q app. |  density  | eigen (eV) | Etot (eV)" not in line
        and "Begin self-consistency iteration #" not in line
    ):
        try:
            line = next(fd).strip()  # Raises StopIteration on empty file
        except StopIteration:
            raise ParseError("No SCF steps present, calculation failed at setup.")

        header.append(line)
    return AimsOutHeaderChunk(header)


def get_aims_out_chunks(fd, header_chunk):
    """Yield unprocessed chunks (header, lines) for each AimsOutChunk image."""
    try:
        line = next(fd).strip()  # Raises StopIteration on empty file
    except StopIteration:
        return

    # If the calculation is relaxation the updated structural information
    # occurs before the re-initialization
    if header_chunk.is_relaxation:
        chunk_end_line = (
            "Geometry optimization: Attempting to predict improved coordinates."
        )
    else:
        chunk_end_line = "Begin self-consistency loop: Re-initialization"

    # If SCF is not converged then do not treat the next chunk_end_line as a
    # new chunk until after the SCF is re-initialized
    ignore_chunk_end_line = False
    while True:
        try:
            line = next(fd).strip()  # Raises StopIteration on empty file
        except StopIteration:
            break

        lines = []
        while chunk_end_line not in line or ignore_chunk_end_line:
            lines.append(line)
            # If SCF cycle not converged or numerical stresses are requested,
            # don't end chunk on next Re-initialization
            patterns = [
                (
                    "Self-consistency cycle not yet converged -"
                    " restarting mixer to attempt better convergence."
                ),
                (
                    "Components of the stress tensor (for mathematical "
                    "background see comments in numerical_stress.f90)."
                ),
                "Calculation of numerical stress completed",
            ]
            if any([pattern in line for pattern in patterns]):
                ignore_chunk_end_line = True
            elif "Begin self-consistency loop: Re-initialization" in line:
                ignore_chunk_end_line = False

            try:
                line = next(fd).strip()
            except StopIteration:
                break
        yield AimsOutCalcChunk(lines, header_chunk)


def check_convergence(chunks, non_convergence_ok=False):
    """Check if the aims output file is for a converged calculation

    Parameters
    ----------
    chunks: list of AimsOutChunks
        The list of chunks for the aims calculations
    non_convergence_ok: bool
        True if it is okay for the calculation to not be converged

    Returns
    -------
    bool
        True if the calculation is converged
    """
    if not non_convergence_ok and not chunks[-1].converged:
        raise ParseError("The calculation did not complete successfully")
    return True


@reader
def read_aims_header_info(
    fd: str | Path,
) -> tuple[Dict[str, str], Dict[str, Any]]:
    """Read the FHI-aims header information"""
    header_chunk = get_header_chunk(fd)

    system_summary = header_chunk.header_summary
    metadata = header_chunk.metadata_summary
    return metadata, system_summary


@reader
def read_aims_output(fd, index=-1, non_convergence_ok=False):
    """Import FHI-aims output files with all data available, i.e.
    relaxations, MD information, force information etc etc etc."""
    header_chunk = get_header_chunk(fd)
    chunks = list(get_aims_out_chunks(fd, header_chunk))
    check_convergence(chunks, non_convergence_ok)

    # Relaxations have an additional footer chunk due to how it is split
    if header_chunk.is_relaxation:
        images = [chunk.atoms for chunk in chunks[:-1]]
    else:
        images = [chunk.atoms for chunk in chunks]
    return images[index]

src/atomate2_temp/aims/io/parsers.py
tpurcell90/atomate2-fhi-aims
"""
A definition of base FHI-aims job Maker
(closely resembling that of VASP and CP2K in atomate2)
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from pathlib import Path

from jobflow import Flow, Maker, Response, job
from monty.serialization import dumpfn
from pymatgen.core import Molecule, Structure

from typing import Dict, Any

from atomate2_temp.aims.files import (
    cleanup_aims_outputs,
    copy_aims_outputs,
    write_aims_input_set,
)
from atomate2_temp.aims.run import run_aims, should_stop_children
from atomate2_temp.aims.schemas.task import AimsTaskDocument, ConvergenceSummary
from atomate2_temp.aims.sets.base import AimsInputGenerator
from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

logger = logging.getLogger(__name__)
CONVERGENCE_FILE_NAME = "convergence.json"  # make it a constant?


@dataclass
class BaseAimsMaker(Maker):
    """
    Base FHI-aims job maker.

    Parameters
    ----------
    name : str
        The job name.
    input_set_generator : AimsInputGenerator
        A generator used to make the input set.
    write_input_set_kwargs : dict
        Keyword arguments that will get passed to :obj:`.write_aims_input_set`.
    copy_aims_kwargs : dict
        Keyword arguments that will get passed to :obj:`.copy_aims_outputs`.
    run_aims_kwargs : dict
        Keyword arguments that will get passed to :obj:`.run_aims`.
    task_document_kwargs : dict
        Keyword arguments that will get passed to :obj:`.TaskDocument.from_directory`.
    stop_children_kwargs : dict
        Keyword arguments that will get passed to :obj:`.should_stop_children`.
    write_additional_data : dict
        Additional data to write to the current directory. Given as a dict of
        {filename: data}. Note that if using FireWorks, dictionary keys cannot contain
        the "." character which is typically used to denote file extensions. To avoid
        this, use the ":" character, which will automatically be converted to ".". E.g.
        ``{"my_file:txt": "contents of the file"}``.
    store_output_data: bool
        Whether the job output (TaskDocument) should be stored in the JobStore through
        the response.
    """

    name: str = "base"
    input_set_generator: AimsInputGenerator = field(default_factory=AimsInputGenerator)
    write_input_set_kwargs: Dict[str, Any] = field(default_factory=dict)
    copy_aims_kwargs: Dict[str, Any] = field(default_factory=dict)
    run_aims_kwargs: Dict[str, Any] = field(default_factory=dict)
    task_document_kwargs: Dict[str, Any] = field(default_factory=dict)
    stop_children_kwargs: Dict[str, Any] = field(default_factory=dict)
    write_additional_data: Dict[str, Any] = field(default_factory=dict)
    store_output_data: bool = True

    @job
    def make(
        self,
        structure: MSONableAtoms | Structure | Molecule,
        prev_dir: str | Path | None = None,
    ):
        """
        Run an FHI-aims calculation.

        Parameters
        ----------
        structure : MSONableAtoms or Structure or Molecule
            An ASE Atoms or pymatgen Structure object to create the calculation for.
        prev_dir : str or Path or None
            A previous FHI-aims calculation directory to copy output files from.
        """
        # the structure transformation part was deleted; can be reinserted when needed
        if isinstance(structure, Structure) or isinstance(structure, Molecule):
            atoms = MSONableAtoms.from_pymatgen(structure)
        else:
            atoms = structure.copy()

        # copy previous inputs
        if prev_dir is not None:
            copy_aims_outputs(prev_dir, **self.copy_aims_kwargs)

        # write aims input files
        self.write_input_set_kwargs["prev_dir"] = prev_dir
        write_aims_input_set(
            atoms, self.input_set_generator, **self.write_input_set_kwargs
        )

        # write any additional data
        for filename, data in self.write_additional_data.items():
            dumpfn(data, filename.replace(":", "."))

        # run FHI-aims
        run_aims(**self.run_aims_kwargs)

        # parse FHI-aims outputs
        task_doc = AimsTaskDocument.from_directory(
            Path.cwd(), **self.task_document_kwargs
        )
        task_doc.task_label = self.name

        # decide whether child jobs should proceed
        stop_children = should_stop_children(task_doc, **self.stop_children_kwargs)

        # cleanup files to save disk space
        cleanup_aims_outputs(directory=Path.cwd())

        return Response(
            stop_children=stop_children,
            output=task_doc if self.store_output_data else None,
        )


@dataclass
class ConvergenceMaker(Maker):
    """A job that performs convergence run for a given number of steps. Stops either
    when all steps are done, or when the convergence criterion is reached, that is when
    the absolute difference between the subsequent values of the convergence field is
    less than a given epsilon.

    Parameters
    ----------
    name : str
        A name for the job
    maker: .BaseAimsMaker
        A maker for the run
    criterion_name: str
        A name for the convergence criterion. Must be in the run results
    epsilon: float
        A difference in criterion value for subsequent runs
    convergence_field: str
        An input parameter that changes to achieve convergence
    convergence_steps: list | tuple
        An iterable of the possible values for the convergence field.
        If the iterable is depleted and the convergence is not reached,
        that the job is failed
    """

    name: str = "Convergence job"
    maker: BaseAimsMaker = field(default_factory=BaseAimsMaker)
    criterion_name: str = "energy_per_atom"
    epsilon: float = 0.001
    convergence_field: str = field(default_factory=str)
    convergence_steps: list = field(default_factory=list)

    def __post_init__(self):
        self.last_idx = len(self.convergence_steps)

    def make(self, atoms):
        """A top-level flow controlling convergence iteration

        Parameters
        ----------
            atoms : MSONableAtoms
                a structure to run a job
        """
        convergence_job = self.convergence_iteration(atoms)
        return Flow([convergence_job], output=convergence_job.output)

    @job
    def convergence_iteration(
        self,
        structure: MSONableAtoms | Structure | Molecule,
        prev_dir: str | Path = None,
    ) -> Response:
        """
        Runs several jobs with changing inputs consecutively to investigate
        convergence in the results

        Parameters
        ----------
        atoms : MSONableAtoms or Structure or Molecule
            The structure to run the job for
        prev_dir: str | None
            An FHI-aims calculation directory in which previous run contents are stored

        Returns
        -------
        The output response for the job
        """
        if isinstance(structure, Structure) or isinstance(structure, Molecule):
            atoms = MSONableAtoms.from_pymatgen(structure)
        else:
            atoms = structure.copy()

        # getting the calculation index
        idx = 0
        converged = False
        if prev_dir is not None:
            prev_dir = prev_dir.split(":")[-1]
            convergence_file = Path(prev_dir) / CONVERGENCE_FILE_NAME
            idx += 1
            if convergence_file.exists():
                with open(convergence_file) as f:
                    data = json.load(f)
                    idx = data["idx"] + 1
                    # check for convergence
                    converged = data["converged"]

        if idx < self.last_idx and not converged:
            # finding next jobs
            next_base_job = self.maker.make(atoms, prev_dir=prev_dir)
            next_base_job.update_maker_kwargs(
                {
                    "_set": {
                        f"input_set_generator->user_parameters->"
                        f"{self.convergence_field}": self.convergence_steps[idx]
                    }
                },
                dict_mod=True,
            )
            next_base_job.append_name(append_str=f" {idx}")

            update_file_job = self.update_convergence_file(
                prev_dir=prev_dir,
                job_dir=next_base_job.output.dir_name,
                output=next_base_job.output,
            )

            next_job = self.convergence_iteration(
                atoms, prev_dir=next_base_job.output.dir_name
            )

            replace_flow = Flow(
                [next_base_job, update_file_job, next_job], output=next_base_job.output
            )
            return Response(detour=replace_flow, output=replace_flow.output)
        else:
            task_doc = AimsTaskDocument.from_directory(prev_dir)
            summary = ConvergenceSummary.from_aims_calc_doc(task_doc)
            return summary

    @job(name="Writing a convergence file")
    def update_convergence_file(
        self, prev_dir: str | Path, job_dir: str | Path, output
    ):
        """Write a convergence file

        Parameters
        ----------
        TO DO: fill out
        """
        idx = 0
        if prev_dir is not None:
            prev_dir = prev_dir.split(":")[-1]
            convergence_file = Path(prev_dir) / CONVERGENCE_FILE_NAME
            if convergence_file.exists():
                with open(convergence_file) as f:
                    convergence_data = json.load(f)
                    idx = convergence_data["idx"] + 1
        else:
            idx = 0
            convergence_data = {
                "criterion_name": self.criterion_name,
                "criterion_values": [],
                "convergence_field_name": self.convergence_field,
                "convergence_field_values": [],
                "converged": False,
            }
        convergence_data["convergence_field_values"].append(self.convergence_steps[idx])
        convergence_data["criterion_values"].append(
            getattr(output.output, self.criterion_name)
        )
        convergence_data["idx"] = idx

        if len(convergence_data["criterion_values"]) > 1:
            # checking for convergence
            convergence_data["converged"] = (
                abs(
                    convergence_data["criterion_values"][-1]
                    - convergence_data["criterion_values"][-2]
                )
                < self.epsilon
            )

        job_dir = job_dir.split(":")[-1]
        convergence_file = Path(job_dir) / CONVERGENCE_FILE_NAME
        with open(convergence_file, "w") as f:
            json.dump(convergence_data, f)

    @job(name="Getting the results")
    def get_results(self, prev_dir: Path | str) -> Dict[str, Any]:
        """Get the results for a calculation from a given directory

        Parameters
        ----------
        prev_dir: Path or str
            The calculation directory to get the results for

        Results
        -------
        The results dictionary loaded from the JSON file
        """
        convergence_file = Path(prev_dir) / CONVERGENCE_FILE_NAME
        with open(convergence_file) as f:
            return json.load(f)

src/atomate2_temp/aims/jobs/base.py
tpurcell90/atomate2-fhi-aims
from __future__ import annotations

import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Sequence

from jobflow import Response, job
from monty.serialization import dumpfn
from pymatgen.core import Structure

from atomate2_temp.aims.files import (
    cleanup_aims_outputs,
    copy_aims_outputs,
    write_aims_input_set,
)
from atomate2_temp.aims.io.parsers import read_aims_output
from atomate2_temp.aims.jobs.base import BaseAimsMaker
from atomate2_temp.aims.run import run_aims_socket, should_stop_children
from atomate2_temp.aims.schemas.task import AimsTaskDocument
from atomate2_temp.aims.sets.base import AimsInputGenerator
from atomate2_temp.aims.sets.bs import BandStructureSetGenerator, GWSetGenerator
from atomate2_temp.aims.sets.core import (
    RelaxSetGenerator,
    SocketIOSetGenerator,
    StaticSetGenerator,
)
from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

logger = logging.getLogger(__name__)
"""Core job makers for FHI-aims workflows"""


@dataclass
class StaticMaker(BaseAimsMaker):
    """Maker to create FHI-aims SCF jobs

    Parameters
    ----------
    calc_type: str
        The type key for the calculation
    name: str
        The job name
    input_set_generator: AimsInputGenerator
        The InputGenerator for the calculation
    """

    calc_type: str = "scf"
    name: str = "SCF Calculation"
    input_set_generator: AimsInputGenerator = field(default_factory=StaticSetGenerator)


@dataclass
class RelaxMaker(BaseAimsMaker):
    """Maker to create relaxation calculations.

    Parameters
    ----------
    calc_type: str
        The type key for the calculation
    name: str
        The job name
    input_set_generator: AimsInputGenerator
        The InputGenerator for the calculation
    """

    calc_type: str = "relax"
    name: str = "Relaxation calculation"
    input_set_generator: AimsInputGenerator = field(default_factory=RelaxSetGenerator)

    @classmethod
    def fixed_cell_relaxation(cls, *args, **kwargs):
        """Create a fixed cell relaxation maker."""
        return cls(
            input_set_generator=RelaxSetGenerator(relax_cell=False, **kwargs),
            name=cls.name + " (fixed cell)",
        )

    @classmethod
    def full_relaxation(cls, *args, **kwargs):
        """Create a full relaxation maker."""
        return cls(input_set_generator=RelaxSetGenerator(relax_cell=True, **kwargs))


@dataclass
class SocketIOStaticMaker(BaseAimsMaker):
    """Maker for the SocketIO calculator in FHI-aims

    Parameters
    ----------
    calc_type: str
        The type key for the calculation
    name: str
        The job name
    host: str
        The name of the host to maitain the socket server on
    port: int
        The port number the socket server will listen on
    input_set_generator: SocketIOSetGenerator
        The InputGenerator for the calculation
    """

    calc_type: str = "multi_scf"
    name: str = "SCF Calculations Socket"
    host: str = "localhost"
    port: int = 12345
    input_set_generator: AimsInputGenerator = field(
        default_factory=SocketIOSetGenerator
    )

    @job
    def make(
        self,
        structure: Sequence[MSONableAtoms | Structure],
        prev_dir: str | Path | None = None,
    ) -> Response:
        """
        Run an FHI-aims calculation on multiple atoms object using the socket
        communicator.

        Parameters
        ----------
        structure : Sequence[MSONableAtoms | Structure]
            The list of atoms objects to run FHI-aims on
        prev_dir : str or Path or None
            A previous FHI-aims calculation directory to copy output files from.

        Returns
        -------
        The output response for the calculations
        """
        # copy previous inputs
        if not isinstance(structure, Sequence):
            structure = [MSONableAtoms(structure)]
        atoms = [
            st.copy()
            if isinstance(st, MSONableAtoms)
            else MSONableAtoms.from_pymatgen(st)
            for st in structure
        ]

        from_prev = prev_dir is not None
        if from_prev:
            copy_aims_outputs(prev_dir, **self.copy_aims_kwargs)

            dest_dir = self.copy_aims_kwargs.get("dest_dir", None)
            if dest_dir is None:
                dest_dir = Path.cwd()

            images = read_aims_output(f"{dest_dir}/aims.out")
            if not isinstance(images, Sequence):
                images = [images]

            for img in images:
                img.calc = None

            for ii in range(-1 * len(atoms), 0, -1):
                if atoms[ii] in images:
                    del atoms[ii]

        # write aims input files
        self.write_input_set_kwargs["prev_dir"] = prev_dir
        write_aims_input_set(
            atoms[0], self.input_set_generator, **self.write_input_set_kwargs
        )

        # write any additional data
        for filename, data in self.write_additional_data.items():
            dumpfn(data, filename.replace(":", "."))

        # run FHI-aims
        run_aims_socket(atoms, **self.run_aims_kwargs)

        # parse FHI-aims outputs
        task_doc = AimsTaskDocument.from_directory(
            Path.cwd(), **self.task_document_kwargs
        )
        task_doc.task_label = self.name

        # decide whether child jobs should proceed
        stop_children = should_stop_children(task_doc, **self.stop_children_kwargs)

        # cleanup files to save disk space
        cleanup_aims_outputs(directory=Path.cwd())

        return Response(
            stop_children=stop_children,
            output=task_doc if self.store_output_data else None,
        )


@dataclass
class BandStructureMaker(BaseAimsMaker):
    """A job Maker for a band structure calculation

    Parameters
    ----------
    calc_type: str
        The type key for the calculation
    name: str
        The job name
    input_set_generator: BandStructureSetGenerator
        The InputGenerator for the calculation
    """

    calc_type = "band_structure"
    name: str = "bands"
    input_set_generator: BandStructureSetGenerator = field(
        default_factory=BandStructureSetGenerator
    )


@dataclass
class GWMaker(BaseAimsMaker):
    """A job Maker for a GW band structure calculation

    Parameters
    ----------
    calc_type: str
        The type key for the calculation
    name: str
        The job name
    input_set_generator: GWSetGenerator
        The InputGenerator for the calculation
    """

    calc_type = "gw"
    name: str = "GW"
    input_set_generator: GWSetGenerator = field(default_factory=GWSetGenerator)

src/atomate2_temp/aims/jobs/core.py
tpurcell90/atomate2-fhi-aims
from dataclasses import dataclass, field

from atomate2_temp.aims.jobs.base import BaseAimsMaker
from atomate2_temp.aims.jobs.core import SocketIOStaticMaker
from atomate2_temp.aims.sets.base import AimsInputGenerator
from atomate2_temp.aims.sets.core import SocketIOSetGenerator, StaticSetGenerator


@dataclass
class PhononDisplacementMaker(BaseAimsMaker):
    """
    Maker to perform a static calculation as a part of the finite displacement method.

    The input set is for a static run with tighter convergence parameters.
    Both the k-point mesh density and convergence parameters
    are stricter than a normal relaxation.

    Parameters
    ----------
    name: str
        The job name.
    input_set_generator: AimsInputGenerator
        A generator used to make the input set.
    """

    name: str = "phonon static aims"

    input_set_generator: AimsInputGenerator = field(
        default_factory=lambda: StaticSetGenerator(
            user_parameters={"compute_forces": True},
            user_kpoints_settings={"density": 5.0, "even": True},
        )
    )


@dataclass
class PhononDisplacementMakerSocket(SocketIOStaticMaker):
    """
    Maker to perform a static calculation as a part of the finite displacement method.

    The input set is for a static run with tighter convergence parameters.
    Both the k-point mesh density and convergence parameters
    are stricter than a normal relaxation.

    Parameters
    ----------
    name: str
        The job name.
    input_set_generator: AimsInputGenerator
        A generator used to make the input set.
    """

    name: str = "phonon static aims socket"

    input_set_generator: AimsInputGenerator = field(
        default_factory=lambda: SocketIOSetGenerator(
            user_parameters={"compute_forces": True},
            user_kpoints_settings={"density": 5.0, "even": True},
        )
    )

src/atomate2_temp/aims/jobs/phonons.py
tpurcell90/atomate2-fhi-aims
"""An FHI-aims jobflow runner"""
from __future__ import annotations

import json
import logging
import os
import subprocess
from os.path import expandvars
from typing import Iterable

from ase.calculators.aims import Aims
from ase.calculators.socketio import SocketIOCalculator
from monty.json import MontyDecoder

from atomate2_temp.aims.schemas.task import AimsTaskDocument
from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

logger = logging.getLogger(__name__)


def run_aims(
    aims_cmd: str = None,
):
    """
    Run FHI-aims.

    Parameters
    ----------
    aims_cmd : str
        The command used to run FHI-aims (defaults to ASE_AIMS_COMMAND env variable).
    """
    if aims_cmd is None:
        aims_cmd = os.getenv("ASE_AIMS_COMMAND", "aims.x")

    aims_cmd = expandvars(aims_cmd)

    logger.info(f"Running command: {aims_cmd}")
    return_code = subprocess.call(["/bin/bash", "-c", aims_cmd], env=os.environ)
    logger.info(f"{aims_cmd} finished running with return code: {return_code}")


def should_stop_children(
    task_document: AimsTaskDocument,
    handle_unsuccessful: bool | str = True,
) -> bool:
    """
    Decide whether child jobs should continue.

    Parameters
    ----------
    task_document : .TaskDocument
        An FHI-aims task document.
    handle_unsuccessful : bool or str
        This is a three-way toggle on what to do if your job looks OK, but is actually
        not converged (either electronic or ionic):

        - `True`: Mark job as completed, but stop children.
        - `False`: Do nothing, continue with workflow as normal.
        - `"error"`: Throw an error.

    Returns
    -------
    bool
        Whether to stop child jobs.
    """
    if task_document.state == "successful":
        return False

    if isinstance(handle_unsuccessful, bool):
        return handle_unsuccessful

    if handle_unsuccessful == "error":
        raise RuntimeError("Job was not successful (not converged)!")

    raise RuntimeError(f"Unknown option for handle_unsuccessful: {handle_unsuccessful}")


def run_aims_socket(atoms_to_calculate: Iterable[MSONableAtoms], aims_cmd: str = None):
    """Uses the ASE interface to run FHI-aims from the socket

    Parameters
    ----------
    atoms_to_calculate
        The list of structures to run scf calculations on
    aims_cmd:
        The aims command to use
    """

    parameters = json.load(open("parameters.json", "rt"), cls=MontyDecoder)
    if aims_cmd:
        parameters["aims_command"] = aims_cmd
    elif "aims_command" not in parameters:
        aims_cmd = os.getenv("ASE_AIMS_COMMAND", "aims.x")

    calculator = Aims(**parameters)
    port = parameters["use_pimd_wrapper"][1]
    atoms = atoms_to_calculate[0].copy()

    with calculator.socketio(port=port) as calc:
        for atoms_calc in atoms_to_calculate:
            # Delete prior calculation results
            calc.results.clear()

            # Reset atoms information to the new cell
            atoms.info = atoms_calc.info
            atoms.cell = atoms_calc.cell
            atoms.positions = atoms_calc.positions

            calc.calculate(atoms, system_changes=["positions", "cell"])

        calc.close()

src/atomate2_temp/aims/run.py
tpurcell90/atomate2-fhi-aims
"""Schemas for FHI-aims calculation objects"""

import os
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
from ase.spectrum.band_structure import BandStructure
from ase.stress import voigt_6_to_full_3x3_stress
from emmet.core.math import Matrix3D, Vector3D
from jobflow.utils import ValueEnum
from pydantic import BaseModel, Field
from pymatgen.core.trajectory import Trajectory
from pymatgen.electronic_structure.dos import Dos
from pymatgen.io.common import VolumetricData

from atomate2_temp.aims.io.aims_output import AimsOutput
from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

STORE_VOLUMETRIC_DATA = ("total_density",)


__all__ = [
    "Status",
    "AimsObject",
    "Calculation",
]


class Status(ValueEnum):
    """FHI-aims calculation state."""

    SUCCESS = "successful"
    FAILED = "failed"


class AimsObject(ValueEnum):
    """Types of FHI-aims data objects."""

    DOS = "dos"
    BAND_STRUCTURE = "band_structure"
    ELECTRON_DENSITY = "electron_density"  # e_density
    WFN = "wfn"  # Wavefunction file
    TRAJECTORY = "trajectory"


class CalculationInput(BaseModel):
    """Summary of inputs for an FHI-aims calculation."""

    structure: MSONableAtoms = Field(None, description="The input Atoms object")

    species_info: Dict = Field(
        None, description="Description of parameters used for each atomic species"
    )

    aims_input: Dict = Field(
        None, description="The aims input parameters used for this task"
    )

    @classmethod
    def from_aims_output(cls, output: AimsOutput):
        """Initialize from AimsOutput object."""
        return cls(
            structure=output.initial_structure,
            species_info=output.data.get("species_info", None),
            aims_input=output.input.as_dict(),
        )


class CalculationOutput(BaseModel):
    """Document defining FHI-aims calculation outputs."""

    energy: float = Field(
        None, description="The final total DFT energy for the calculation"
    )
    energy_per_atom: float = Field(
        None, description="The final DFT energy per atom for the calculation"
    )
    structure: MSONableAtoms = Field(
        None, description="The final structure from the calculation"
    )
    efermi: float = Field(
        None, description="The Fermi level from the calculation in eV"
    )

    forces: List[Vector3D] = Field(None, description="Forces acting on each atom")
    all_forces: List[List[Vector3D]] = Field(
        None,
        description="Forces acting on each atom for each structure in the output file",
    )
    stress: Matrix3D = Field(None, description="The stress on the cell")
    stresses: List[Matrix3D] = Field(None, description="The stress on the cell")

    all_forces: List[List[Vector3D]] = Field(
        None, description="Forces acting on each atom for all images in the calculation"
    )

    is_metal: bool = Field(None, description="Whether the system is metallic")
    bandgap: float = Field(None, description="The band gap from the calculation in eV")
    cbm: float = Field(
        None,
        description="The conduction band minimum in eV (if system is not metallic)",
    )
    vbm: float = Field(
        None, description="The valence band maximum in eV (if system is not metallic)"
    )
    atomic_steps: List[MSONableAtoms] = Field(
        None, description="Structures for each ionic step"
    )

    @classmethod
    def from_aims_output(
        cls,
        output: AimsOutput,  # Must use auto_load kwarg when passed
        store_trajectory: bool = False,
    ) -> "CalculationOutput":
        """
        Create an FHI-aims output document from FHI-aims outputs.

        Parameters
        ----------
        output
            An AimsOutput object.
        store_trajectory
            A flag setting to store output trajectory

        Returns
        -------
            The FHI-aims calculation output document.
        """

        structure = output.final_structure

        electronic_output = {
            "efermi": output.fermi_energy,
            "homo": output.homo,
            "lumo": output.lumo,
            "bandgap": output.band_gap,
            "direct_bandgap": output.direct_band_gap,
        }

        forces = None
        if output.forces is not None:
            forces = output.forces.tolist()

        stress = None
        if output.stress is not None:
            stress = voigt_6_to_full_3x3_stress(output.stress).tolist()

        stresses = None
        if output.stresses is not None:
            stresses = [
                voigt_6_to_full_3x3_stress(st).tolist() for st in output.stresses
            ]

        all_forces = None
        if not any([ff is None for ff in output.all_forces]):
            all_forces = [
                f.tolist() if (f is not None) else None for f in output.all_forces
            ]

        return cls(
            structure=structure,
            energy=output.final_energy,
            energy_per_atom=output.final_energy / len(structure),
            **electronic_output,
            atomic_steps=output.structures,
            forces=forces,
            stress=stress,
            stresses=stresses,
            all_forces=all_forces,
        )


class Calculation(BaseModel):
    """Full FHI-aims calculation inputs and outputs."""

    dir_name: str = Field(
        None, description="The directory for this FHI-aims calculation"
    )
    aims_version: str = Field(
        None, description="FHI-aims version used to perform the calculation"
    )
    has_aims_completed: Status = Field(
        None, description="Whether FHI-aims completed the calculation successfully"
    )
    output: CalculationOutput = Field(
        None, description="The FHI-aims calculation output"
    )
    completed_at: str = Field(
        None, description="Timestamp for when the calculation was completed"
    )
    output_file_paths: Dict[str, str] = Field(
        None,
        description="Paths (relative to dir_name) of the FHI-aims output files "
        "associated with this calculation",
    )

    @classmethod
    def from_aims_files(
        cls,
        dir_name: Union[Path, str],
        task_name: str,
        aims_output_file: Union[Path, str] = "aims.out",
        volumetric_files: List[str] = None,
        parse_dos: Union[str, bool] = False,
        parse_bandstructure: Union[str, bool] = False,
        store_trajectory: bool = False,
        store_scf: bool = False,
        store_volumetric_data: Optional[Tuple[str]] = STORE_VOLUMETRIC_DATA,
    ) -> Tuple["Calculation", Dict[AimsObject, Dict]]:
        """
        Create an FHI-aims calculation document from a directory and file paths.

        Parameters
        ----------
        dir_name
            The directory containing the calculation outputs.
        task_name
            The task name.
        aims_output_file
            Path to the main output of aims job, relative to dir_name.
        volumetric_files
            Path to volumetric (Cube) files, relative to dir_name.
        parse_dos
            Whether to parse the DOS. Can be:

            - "auto": Only parse DOS if there are no ionic steps.
            - True: Always parse DOS.
            - False: Never parse DOS.

        parse_bandstructure
            How to parse the bandstructure. Can be:

            - "auto": Parse the bandstructure with projections for NSCF calculations
              and decide automatically if it's line or uniform mode.
            - "line": Parse the bandstructure as a line mode calculation with
              projections
            - True: Parse the bandstructure as a uniform calculation with
              projections .
            - False: Parse the band structure without projects and just store
              vbm, cbm, band_gap, is_metal and efermi rather than the full
              band structure object.

        store_trajectory:
            Whether to store the ionic steps as a pmg trajectory object, which can be
            pushed, to a bson data store, instead of as a list od dicts. Useful for
            large trajectories.
        store_scf:
            Whether to store the SCF convergence data.
        store_volumetric_data
            Which volumetric files to store.

        Returns
        -------
        Calculation
            An FHI-aims calculation document.
        """
        dir_name = Path(dir_name)
        aims_output_file = dir_name / aims_output_file

        volumetric_files = [] if volumetric_files is None else volumetric_files
        aims_output = AimsOutput.from_outfile(aims_output_file)

        completed_at = str(datetime.fromtimestamp(os.stat(aims_output_file).st_mtime))

        output_file_paths = _get_output_file_paths(volumetric_files)
        aims_objects: Dict[AimsObject, Any] = _get_volumetric_data(
            dir_name, output_file_paths, store_volumetric_data
        )

        dos = _parse_dos(parse_dos, aims_output)
        if dos is not None:
            aims_objects[AimsObject.DOS] = dos  # type: ignore

        bandstructure = _parse_bandstructure(parse_bandstructure, aims_output)
        if bandstructure is not None:
            aims_objects[AimsObject.BANDSTRUCTURE] = bandstructure  # type: ignore

        output_doc = CalculationOutput.from_aims_output(aims_output)

        has_aims_completed = Status.SUCCESS if aims_output.completed else Status.FAILED

        if store_trajectory:
            traj = _parse_trajectory(aims_output=aims_output)
            aims_objects[AimsObject.TRAJECTORY] = traj  # type: ignore

        return (
            cls(
                dir_name=str(dir_name),
                task_name=task_name,
                aims_version=aims_output.aims_version,
                has_aims_completed=has_aims_completed,
                completed_at=completed_at,
                output=output_doc,
                output_file_paths={
                    k.name.lower(): v for k, v in output_file_paths.items()
                },
            ),
            aims_objects,
        )


def _get_output_file_paths(volumetric_files: List[str]) -> Dict[AimsObject, str]:
    """
    Get the output file paths for FHI-aims output files from the list
    of volumetric files.

    Parameters
    ----------
    volumetric_files
        A list of volumetric files associated with the calculation.

    Returns
    -------
    Dict[AimsObject, str]
        A mapping between the Aims object type and the file path.
    """
    output_file_paths = {}
    for aims_object in AimsObject:  # type: ignore
        for volumetric_file in volumetric_files:
            if aims_object.name in str(volumetric_file):
                output_file_paths[aims_object] = str(volumetric_file)
    return output_file_paths


def _get_volumetric_data(
    dir_name: Path,
    output_file_paths: Dict[AimsObject, str],
    store_volumetric_data: Optional[Tuple[str]],
) -> Dict[AimsObject, VolumetricData]:
    """
    Load volumetric data files from a directory.

    Parameters
    ----------
    dir_name
        The directory containing the files.
    output_file_paths
        A dictionary mapping the data type to file path relative to dir_name.
    store_volumetric_data
        The volumetric data files to load.

    Returns
    -------
    Dict[AimsObject, VolumetricData]
        A dictionary mapping the FHI-aims object data type (`AimsObject.total_density`,
        `AimsObject.electron_density`, etc) to the volumetric data object.
    """
    if store_volumetric_data is None or len(store_volumetric_data) == 0:
        return {}

    volumetric_data = {}
    for file_type, file in output_file_paths.items():
        if file_type.name not in store_volumetric_data:
            continue
        try:
            volumetric_data[file_type] = VolumetricData.from_cube(
                (dir_name / file).as_posix()
            )
        except Exception as err:
            raise ValueError(f"Failed to parse {file_type} at {file}.") from err

    return volumetric_data


def _parse_dos(parse_dos: Union[str, bool], aims_output: AimsOutput) -> Optional[Dos]:
    """
    Parse DOS outputs from FHI-aims calculation.

    Parameters
    ----------
    parse_dos: Whether to parse the DOS. Can be:
        - "auto": Only parse DOS if there are no ionic steps.
        - True: Always parse DOS.
        - False: Never parse DOS.
    aims_output: AimsOutput object for the calculation being parsed.

    Returns
    -------
    A Dos object if parse_dos is set accordingly.
    """
    if parse_dos == "auto":
        if len(aims_output.ionic_steps) == 0:
            return aims_output.complete_dos
        return None
    if parse_dos:
        return aims_output.complete_dos
    return None


def _parse_bandstructure(
    parse_bandstructure: Union[str, bool], aims_output: AimsOutput
) -> Optional[BandStructure]:
    """
    Get the band structure.

    Parameters
    ----------
         parse_bandstructure (bool): Whether to parse. Does not support the
            auto/line distinction currently.
    """
    if parse_bandstructure:
        return aims_output.band_structure
    return None


def _parse_trajectory(aims_output: AimsOutput) -> Optional[Trajectory]:
    """
    Grab a Trajectory object given an FHI-aims output object.
    """

    return aims_output.structures

src/atomate2_temp/aims/schemas/calculation.py
tpurcell90/atomate2-fhi-aims
"""A definition of a MSON document representing an FHI-aims task"""

import json
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union

import numpy as np
from emmet.core.math import Matrix3D, Vector3D
from emmet.core.structure import MoleculeMetadata, StructureMetadata
from emmet.core.tasks import get_uri
from pydantic import BaseModel, Field
from pymatgen.entries.computed_entries import ComputedEntry

from atomate2_temp.aims.schemas.calculation import AimsObject, Calculation, Status
from atomate2_temp.aims.utils import datetime_str
from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

_T = TypeVar("_T", bound="AimsTaskDocument")
_VOLUMETRIC_FILES = ("total_density", "spin_density", "eigenstate_density")
logger = logging.getLogger(__name__)


class AnalysisSummary(BaseModel):
    """Calculation relaxation summary."""

    delta_volume: float = Field(None, description="Absolute change in volume")
    delta_volume_as_percent: float = Field(
        None, description="Percentage change in volume"
    )
    max_force: float = Field(None, description="Maximum force on the atoms")
    errors: List[str] = Field(None, description="Errors from the FHI-aims output")

    @classmethod
    def from_aims_calc_docs(cls, calc_docs: List[Calculation]) -> "AnalysisSummary":
        """
        Create analysis summary from FHI-aims calculation documents.

        Parameters
        ----------
        calc_docs
            FHI-aims calculation documents.

        Returns
        -------
        AnalysisSummary
            Summary object
        """

        errors = []

        delta_vol = None
        percent_delta_vol = None

        final_calc = calc_docs[-1]
        max_force = None
        if final_calc.has_aims_completed == Status.SUCCESS:
            max_force = _get_max_force(final_calc)

        return cls(
            delta_volume=delta_vol,
            delta_volume_as_percent=percent_delta_vol,
            max_force=max_force,
            errors=errors,
        )


class Species(BaseModel):
    """A representation of the most important information about each type of species."""

    element: str = Field(None, description="Element assigned to this atom kind")
    species_defaults: str = Field(None, description="Basis set for this atom kind")


class SpeciesSummary(BaseModel):
    """A summary of species defaults."""

    species_defaults: Dict[str, Species] = Field(
        None, description="Dictionary mapping atomic kind labels to their info"
    )

    @classmethod
    def from_species_info(cls, species_info: dict):
        """Initialize from the atomic_kind_info dictionary."""
        d: Dict[str, Dict[str, Any]] = {"species_defaults": {}}
        for kind, info in species_info.items():
            d["species_defaults"][kind] = {
                "element": info["element"],
                "species_defaults": info["species_defaults"],
            }
        return cls(**d)


class InputSummary(BaseModel):
    """Summary of inputs for an FHI-aims calculation."""

    structure: MSONableAtoms = Field(None, description="The input structure object")

    species_info: SpeciesSummary = Field(
        None, description="Summary of the species defaults used for each atom kind"
    )
    xc: str = Field(
        None, description="Exchange-correlation functional used if not the default"
    )

    @classmethod
    def from_aims_calc_doc(cls, calc_doc: Calculation) -> "InputSummary":
        """
        Create calculation input summary from a calculation document.

        Parameters
        ----------
        calc_doc
            An FHI-aims calculation document.

        Returns
        -------
        InputSummary
            A summary of the input structure and parameters.
        """
        summary = SpeciesSummary.from_species_info(calc_doc.input.species_info)

        return cls(
            structure=calc_doc.input.structure,
            atomic_kind_info=summary,
            xc=str(calc_doc.run_type),
        )


class OutputSummary(BaseModel):
    """Summary of the outputs for an FHI-aims calculation."""

    structure: MSONableAtoms = Field(None, description="The output structure object")
    trajectory: List[MSONableAtoms] = Field(
        None, description="The trajectory of output structures"
    )
    energy: float = Field(
        None, description="The final total DFT energy for the last calculation"
    )
    energy_per_atom: float = Field(
        None, description="The final DFT energy per atom for the last calculation"
    )
    bandgap: float = Field(None, description="The DFT bandgap for the last calculation")
    cbm: float = Field(None, description="CBM for this calculation")
    vbm: float = Field(None, description="VBM for this calculation")
    forces: List[Vector3D] = Field(
        None, description="Forces on atoms from the last calculation"
    )
    stress: Matrix3D = Field(
        None, description="Stress on the unit cell from the last calculation"
    )
    all_forces: List[List[Vector3D]] = Field(
        None, description="Forces on atoms from all calculations."
    )

    @classmethod
    def from_aims_calc_doc(cls, calc_doc: Calculation) -> "OutputSummary":
        """
        Create a summary of FHI-aims calculation outputs from an FHI-aims
        calculation document.

        Parameters
        ----------
        calc_doc
            An FHI-aims calculation document.

        Returns
        -------
        OutputSummary
            The calculation output summary.
        """

        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,
            cbm=calc_doc.output.cbm,
            vbm=calc_doc.output.vbm,
            forces=calc_doc.output.forces,
            stress=calc_doc.output.stress,
            all_forces=calc_doc.output.all_forces,
            trajectory=calc_doc.output.atomic_steps,
        )


class ConvergenceSummary(BaseModel):
    """Summary of the outputs for an FHI-aims convergence calculation."""

    structure: MSONableAtoms = Field(None, description="The output structure object")
    converged: bool = Field(None, description="Is convergence achieved?")

    convergence_criterion_name: str = Field(
        None, description="The output name of the convergence criterion"
    )
    convergence_field_name: str = Field(
        None, description="The name of the input setting to study convergence against"
    )
    convergence_criterion_value: float = Field(
        None, description="The output value of the convergence criterion"
    )
    convergence_field_value: Any = Field(
        None,
        description="The last value of the input setting to study convergence against",
    )
    asked_epsilon: float = Field(
        None,
        description="The difference in the values for the convergence criteria that was asked for",
    )
    actual_epsilon: float = Field(
        None, description="The actual difference in the convergence criteria values"
    )

    @classmethod
    def from_aims_calc_doc(cls, calc_doc: Calculation) -> "ConvergenceSummary":
        """
        Create a summary of FHI-aims calculation outputs from an FHI-aims calculation document.

        Parameters
        ----------
        calc_doc
            An FHI-aims calculation document.

        Returns
        -------
        :ConvergenceSummary
            The summary for convergence runs.
        """

        from fhi_aims_workflows.jobs.base import CONVERGENCE_FILE_NAME

        job_dir = calc_doc.dir_name.split(":")[-1]

        convergence_file = Path(job_dir) / CONVERGENCE_FILE_NAME
        if not convergence_file.exists():
            raise ValueError(
                f"Did not find the convergence json file {CONVERGENCE_FILE_NAME} in {calc_doc.dir_name}"
            )

        with open(convergence_file) as f:
            convergence_data = json.load(f)

        return cls(
            structure=calc_doc.output.structure,
            converged=convergence_data["converged"],
            convergence_criterion_name=convergence_data["criterion_name"],
            convergence_field_name=convergence_data["convergence_field_name"],
            convergence_criterion_value=convergence_data["criterion_values"][-1],
            convergence_field_value=convergence_data["convergence_field_values"][-1],
            asked_epsilon=None,
            actual_epsilon=abs(
                convergence_data["criterion_values"][-2]
                - convergence_data["criterion_values"][-1]
            ),
        )


class AimsTaskDocument(StructureMetadata, MoleculeMetadata):
    """Definition of FHI-aims task document."""

    dir_name: str = Field(None, description="The directory for this FHI-aims task")
    last_updated: str = Field(
        default_factory=datetime_str,
        description="Timestamp for this task document was last updated",
    )
    completed_at: str = Field(
        None, description="Timestamp for when this task was completed"
    )
    input: InputSummary = Field(None, description="The input to the first calculation")
    output: OutputSummary = Field(
        None, description="The output of the final calculation"
    )
    structure: MSONableAtoms = Field(
        None, description="Final output structure from the task"
    )
    state: Status = Field(None, description="State of this task")
    included_objects: List[AimsObject] = Field(
        None, description="List of FHI-aims objects included with this task document"
    )
    aims_objects: Dict[AimsObject, Any] = Field(
        None, description="FHI-aims objects associated with this task"
    )
    entry: ComputedEntry = Field(
        None, description="The ComputedEntry from the task doc"
    )
    analysis: AnalysisSummary = Field(
        None, description="Summary of structural relaxation and forces"
    )
    task_label: str = Field(None, description="A description of the task")
    tags: List[str] = Field(None, description="Metadata tags for this task document")
    author: str = Field(None, description="Author extracted from transformations")
    icsd_id: str = Field(
        None, description="International crystal structure database id of the structure"
    )
    calcs_reversed: List[Calculation] = Field(
        None, description="The inputs and outputs for all FHI-aims runs in this task."
    )
    transformations: Dict[str, Any] = Field(
        None,
        description="Information on the structural transformations, parsed from a "
        "transformations.json file",
    )
    custodian: Any = Field(
        None,
        description="Information on the custodian settings used to run this "
        "calculation, parsed from a custodian.json file",
    )
    additional_json: Dict[str, Any] = Field(
        None, description="Additional json loaded from the calculation directory"
    )

    @classmethod
    def from_directory(
        cls: Type[_T],
        dir_name: Union[Path, str],
        volumetric_files: Tuple[str, ...] = _VOLUMETRIC_FILES,
        additional_fields: Dict[str, Any] = None,
        **aims_calculation_kwargs,
    ) -> _T:
        """
        Create a task document from a directory containing FHi-aims files.

        Parameters
        ----------
        dir_name
            The path to the folder containing the calculation outputs.
        volumetric_files
            A volumetric files to search for.
        additional_fields
            Dictionary of additional fields to add to output document.
        **aims_calculation_kwargs
            Additional parsing options that will be passed to the
            :obj:`.Calculation.from_aims_files` function.

        Returns
        -------
        AimsTaskDoc
            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_aims_files(dir_name, volumetric_files=volumetric_files)

        if len(task_files) == 0:
            raise FileNotFoundError("No FHI-aims files found!")

        calcs_reversed = []
        all_aims_objects = []
        for task_name, files in task_files.items():
            calc_doc, aims_objects = Calculation.from_aims_files(
                dir_name, task_name, **files, **aims_calculation_kwargs
            )
            calcs_reversed.append(calc_doc)
            all_aims_objects.append(aims_objects)

        analysis = AnalysisSummary.from_aims_calc_docs(calcs_reversed)
        tags = additional_fields.get("tags")

        dir_name = get_uri(dir_name)  # convert to full uri path

        # only store objects from last calculation
        # TODO: make this an option
        aims_objects = all_aims_objects[-1]
        included_objects = None
        if aims_objects:
            included_objects = list(aims_objects.keys())

        # rewrite the original structure save!

        data = {
            "structure": calcs_reversed[-1].output.structure,
            "meta_structure": calcs_reversed[-1].output.structure,
            "dir_name": dir_name,
            "calcs_reversed": calcs_reversed,
            "analysis": analysis,
            "tags": tags,
            "completed_at": calcs_reversed[-1].completed_at,
            "output": OutputSummary.from_aims_calc_doc(calcs_reversed[-1]),
            "state": _get_state(calcs_reversed, analysis),
            "entry": cls.get_entry(calcs_reversed),
            "aims_objects": aims_objects,
            "included_objects": included_objects,
        }
        doc = cls(**data)
        return doc.copy(update=additional_fields)

    @staticmethod
    def get_entry(
        calc_docs: List[Calculation], job_id: Optional[str] = None
    ) -> ComputedEntry:
        """
        Get a computed entry from a list of FHI-aims calculation documents.

        Parameters
        ----------
        calc_docs
            A list of FHI-aims calculation documents.
        job_id
            The job identifier.

        Returns
        -------
        ComputedEntry
            A computed entry.
        """
        entry_dict = {
            "correction": 0.0,
            "entry_id": job_id,
            "composition": calc_docs[-1].output.structure.get_chemical_formula(),
            "energy": calc_docs[-1].output.energy,
            "parameters": {
                # Required to be compatible with MontyEncoder for the ComputedEntry
                # "run_type": str(calc_docs[-1].run_type),
                "run_type": "AIMS run"
            },
            "data": {
                "last_updated": datetime_str(),
            },
        }
        return ComputedEntry.from_dict(entry_dict)


def _find_aims_files(
    path: Union[str, Path],
    volumetric_files: Tuple[str, ...] = _VOLUMETRIC_FILES,
) -> Dict[str, Any]:
    """
    Find FHI-aims 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., aims.out) will be returned.

    CP2K 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 FHI-aims task,
        given as a ordered dictionary of::

            {
                task_name: {
                    "aims_output_file": aims_output_filename,
                    "volumetric_files": [v_hartree file, e_density file, etc],
    """
    task_names = ["precondition"] + [f"relax{i}" for i in range(9)]
    path = Path(path)
    task_files = {}

    def _get_task_files(files, suffix=""):
        aims_files = {}
        vol_files = []
        for file in files:
            # Here we make assumptions about the output file naming
            if file.match(f"*aims.out{suffix}*"):
                aims_files["aims_output_file"] = Path(file).name
        for vol in volumetric_files:
            _files = [f.name for f in files if f.match(f"*{vol}*cube{suffix}*")]
            if _files:
                vol_files.append(_files[0])

        if len(vol_files) > 0:
            # add volumetric files if some were found or other cp2k files were found
            aims_files["volumetric_files"] = vol_files

        return aims_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


def _get_max_force(calc_doc: Calculation) -> Optional[float]:
    """Get max force acting on atoms from a calculation document."""
    forces = calc_doc.output.forces
    if forces is not None:
        forces = np.array(forces)
        return max(np.linalg.norm(forces, axis=1))
    return None


def _get_state(calc_docs: List[Calculation], analysis: AnalysisSummary) -> Status:
    """Get state from calculation documents and relaxation analysis."""
    all_calcs_completed = all(c.has_aims_completed == Status.SUCCESS for c in calc_docs)
    if len(analysis.errors) == 0 and all_calcs_completed:
        return Status.SUCCESS  # type: ignore
    return Status.FAILED  # type: ignore

src/atomate2_temp/aims/schemas/task.py
tpurcell90/atomate2-fhi-aims
"""Module defining base FHI-aims input set and generator."""
from __future__ import annotations

import copy
import json
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Iterable, List, Sequence, Tuple
from warnings import warn

import numpy as np
from ase.cell import Cell
from ase.calculators.aims import AimsTemplate
from monty.json import MontyDecoder, MontyEncoder
from pymatgen.io.core import InputFile, InputGenerator, InputSet

from atomate2_temp.aims.io.parsers import AimsParseError, read_aims_output
from atomate2_temp.aims.utils.common import (
    CONTROL_FILE_NAME,
    GEOMETRY_FILE_NAME,
    PARAMS_JSON_FILE_NAME,
    TMPDIR_NAME,
    cwd,
)
from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

DEFAULT_AIMS_PROPERTIES = [
    "energy",
    "free_energy",
    "forces",
    "stress",
    "stresses",
    "dipole",
    "magmom",
]

__all__ = ["AimsInputSet", "AimsInputGenerator"]

logger = logging.getLogger(__name__)


@dataclass
class AimsInputFile(InputFile):
    """The input file for an FHI-aims calculation

    Parameters
    ----------
    _content_str: str
        The contents of the input file as a string
    """

    _content_str: str = ""

    def get_string(self) -> str:
        """The contents of the input file

        Returns
        -------
        The contents of the input file
        """
        return self._content_str

    def get_str(self) -> str:
        """The contents of the input file

        Returns
        -------
        The contents of the input file
        """
        return self._content_str

    @classmethod
    def from_string(cls, contents: str):
        """Create an input file from the contents string

        Parameters
        ----------
        contents: str
            The contents of the input file
        """
        return cls(contents)

    @classmethod
    def from_str(cls, contents: str):
        """Create an input file from the contents string

        Parameters
        ----------
        contents: str
            The contents of the input file
        """
        return cls(contents)


class AimsInputSet(InputSet):
    """
    A class to represent a set of Aims inputs.
    """

    def __init__(
        self,
        parameters: Dict[str, Any],
        atoms: MSONableAtoms,
        properties: Sequence[str] = ("energy", "free_energy"),
    ):
        """Constructor

        Parameters
        ----------
        parameters: Dict[str, Any]
            The ASE parameters object for the calculation
        atoms: MSONableAtoms
            The atoms objects to create the inputs for
        properties: Sequence[str]
            The properties to calculate for the calculation
        """
        self._parameters = parameters
        self._atoms = MSONableAtoms(atoms)
        self._properties = properties

        aims_control_in, aims_geometry_in = self.get_input_files()
        super().__init__(
            inputs={
                CONTROL_FILE_NAME: aims_control_in,
                GEOMETRY_FILE_NAME: aims_geometry_in,
                PARAMS_JSON_FILE_NAME: json.dumps(
                    self._parameters, indent=2, cls=MontyEncoder
                ),
            }
        )

    def get_input_files(self) -> Tuple(str, str):
        """Get the input file contents for the calculation

        Returns
        -------
        The contents of the control.in and geometry.in file
        """
        with cwd(TMPDIR_NAME, mkdir=True, rmdir=True):
            aims_template = AimsTemplate()
            aims_template.write_input(
                Path("./"), self._atoms, self._parameters, self._properties
            )
            aims_control_in = AimsInputFile.from_file("control.in")
            aims_geometry_in = AimsInputFile.from_file("geometry.in")
        return aims_control_in, aims_geometry_in

    @property
    def control_in(self) -> str:
        """Get the AimsInput object."""
        return self[CONTROL_FILE_NAME]

    @property
    def geometry_in(self) -> str:
        """Get the AimsInput object."""
        return self[GEOMETRY_FILE_NAME]

    @property
    def parameters_json(self) -> str:
        """Get the AimsInput object."""
        return self[PARAMS_JSON_FILE_NAME]

    def set_parameters(self, *args, **kwargs) -> Dict[str, Any]:
        """Set the parameters object for the AimsTemplate

        This sets the parameters object that is passed to an AimsTempalte and
        resets the control.in file

        One can pass a dictionary mapping the aims variables to their values or
        the aims variables as keyword arguments. A combination of the two
        options is also allowed.

        Returns
        -------
        dict
            dictionary with the variables that have been added.
        """
        self._parameters.clear()
        for arg in args:
            self._parameters.update(arg)

        self._parameters.update(kwargs)

        aims_control_in, _ = self.get_input_files()
        self.inputs[CONTROL_FILE_NAME] = aims_control_in
        self.inputs[PARAMS_JSON_FILE_NAME] = json.dumps(
            self._parameters, indent=2, cls=MontyEncoder
        )
        self.__dict__.update(self.inputs)

        return self._parameters

    def remove_parameters(
        self, keys: Iterable[str] | str, strict: bool = True
    ) -> Dict[str, Any]:
        """Remove the aims parameters listed in keys.

        This removes the aims variables from the parameters object.

        Parameters
        ----------
        keys
            string or list of strings with the names of the aims parameters
            to be removed.
        strict
            whether to raise a KeyError if one of the aims parameters to be
            removed is not present.

        Returns
        -------
        Dictionary with the variables that have been removed.
        """
        if isinstance(keys, str):
            keys = [keys]
        for key in keys:
            if strict and key not in self._parameters:
                raise ValueError(f"The key ({key}) is not in self._parameters")
            elif key not in self._parameters:
                continue
            del self._parameters[key]

        return self.set_parameters(**self._parameters)

    def set_atoms(self, atoms: MSONableAtoms):
        """Set the atoms object for this input set."""
        self._atoms = MSONableAtoms(atoms)

        _, aims_geometry_in = self.get_input_files()
        self.inputs[GEOMETRY_FILE_NAME] = aims_geometry_in
        self.__dict__.update(self.inputs)

    def deepcopy(self):
        """Deep copy of the input set."""
        return copy.deepcopy(self)


@dataclass
class AimsInputGenerator(InputGenerator):
    """
    A class to generate Aims input sets.

    Parameters
    ----------
    user_parameters: Dict[str, Any]
        Updates the default parameters for the FHI-aims calculator
    user_kpoints_settings: Dict[str, Any]
        The settings used to create the k-grid parameters for FHI-aims
    """

    user_parameters: Dict[str, Any] = field(default_factory=dict)
    user_kpoints_settings: Dict[str, Any] = field(default_factory=dict)

    def get_input_set(  # type: ignore
        self,
        atoms: MSONableAtoms = None,
        prev_dir: str | Path = None,
        properties: List[str] | Tuple[str] = None,
    ) -> AimsInputSet:
        """Generate an AimsInputSet object.

        Parameters
        ----------
        atoms : MSONableAtoms
            ASE Atoms object to generate the input set for.
        prev_dir: str or Path
            Path to the previous working directory
        properties: list or tuple of str
            System properties that are being calculated

        Returns
        -------
        The input set for the calculation of atoms
        """
        prev_atoms, prev_parameters, prev_results = self._read_previous(prev_dir)
        atoms = atoms if atoms is not None else prev_atoms

        parameters = self._get_input_parameters(atoms, prev_parameters)
        properties = self._get_properties(properties, parameters, prev_results)

        return AimsInputSet(parameters=parameters, atoms=atoms, properties=properties)

    def _read_previous(
        self, prev_dir: str | Path = None
    ) -> tuple[MSONableAtoms, Dict[str, Any], Dict[str, Iterable[float]]]:
        """Read in previous results

        Parameters
        ----------
        prev_dir: str or Path
            The previous directory for the calculation
        """
        prev_atoms = None
        prev_parameters = {}
        prev_results = {}

        if prev_dir:
            # strip hostname from the directory (not good, works only with run_locally.
            # Should be checked with Fireworks, will not for sure work with
            # jobflow_remote)
            prev_dir = prev_dir.split(":")[-1]
            prev_parameters = json.load(
                open(f"{prev_dir}/parameters.json", "rt"), cls=MontyDecoder
            )
            try:
                prev_atoms = read_aims_output(f"{prev_dir}/aims.out")
                prev_results = prev_atoms.calc.results
            except (IndexError, AimsParseError):
                pass

        return prev_atoms, prev_parameters, prev_results

    def _get_properties(
        self,
        properties: Iterable[str] = None,
        parameters: Dict[str, Any] = None,
        prev_results: Dict[str, Iterable[float]] = None,
    ) -> Iterable[str]:
        """Get the properties to calculate

        Parameters
        ----------
        properties
            The currently requested properties
        parameters
            The parameters for this calculation
        prev_results
            The previous calculation results

        Returns
        -------
        The list of properties to calculate
        """
        if properties is None:
            properties = ["energy", "free_energy"]

        for key in prev_results.keys():
            if key not in properties and key in DEFAULT_AIMS_PROPERTIES:
                properties.append(key)

        if "compute_forces" in parameters and "forces" not in properties:
            properties.append("forces")
        if "compute_heat_flux" in parameters and "stresses" not in properties:
            properties.append("stress")
            properties.append("stresses")
        if "stress" not in properties and (
            ("compute_analytical_stress" in parameters)
            or ("compute_numerical_stress" in parameters)
            or ("compute_heat_flux" in parameters)
        ):
            properties.append("stress")

        return properties

    def _get_input_parameters(
        self, atoms: MSONableAtoms, prev_parameters: Dict[str, Any] = None
    ) -> Dict[str, Any]:
        """Create the input parameters

        Parameters
        ----------
        atoms: MSONableAtoms
            The atoms object for the structures
        prev_parameters
            The previous calculation's calculation parameters

        Returns
        -------
        The input object
        """

        # Get the default configuration
        # FHI-aims recommends using their defaults so bare-bones default parameters
        parameters = {
            "xc": "pbe",
            "relativistic": "atomic_zora scalar",
        }

        # Override default parameters with previous parameters
        prev_parameters = (
            {} if prev_parameters is None else copy.deepcopy(prev_parameters)
        )
        prev_parameters.pop("relax_geometry", None)
        prev_parameters.pop("relax_unit_cell", None)

        kpt_settings = copy.deepcopy(self.user_kpoints_settings)
        if "k_grid" in prev_parameters:
            density = self.k2d(atoms, prev_parameters.pop("k_grid"))
            if "density" not in kpt_settings:
                kpt_settings["density"] = density

        parameter_updates = self.get_parameter_updates(atoms, prev_parameters)
        parameters = recursive_update(parameters, parameter_updates)

        # Override default parameters with user_parameters
        parameters = recursive_update(parameters, self.user_parameters)
        if ("k_grid" in parameters) and ("density" in kpt_settings):
            warn(
                "WARNING: the k_grid is set in user_parameters and in the kpt_settings, using the one passed in user_parameters."
            )
        elif np.any(atoms.pbc) and ("k_grid" not in parameters):
            density = kpt_settings.get("density", 5.0)
            even = kpt_settings.get("even", True)
            parameters["k_grid"] = self.d2k(atoms, density, even)
        elif not np.any(atoms.pbc) and "k_grid" in parameters:
            warn("WARNING: removing unnecessary k_grid information")
            del parameters["k_grid"]

        return parameters

    def get_parameter_updates(
        self, atoms: MSONableAtoms, prev_parameters: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Updates the parameters for a given calculation type

        Parameters
        ----------
        atoms : MSONableAtoms
            ASE Atoms object.
        prev_parameters: Dict[str, Any]
            Previous calculation parameters.

        Returns
        -------
        A dictionary of updates to apply.
        """
        return prev_parameters

    def d2k(
        self,
        atoms: MSONableAtoms,
        kptdensity: float | Iterable[float] = 5.0,
        even: bool = True,
    ) -> Iterable[float]:
        """Convert k-point density to Monkhorst-Pack grid size.

        inspired by [ase.calculators.calculator.kptdensity2monkhorstpack]

        Parameters
        ----------
        atoms: MSONableAtoms
            Contains unit cell and information about boundary conditions.
        kptdensity: float or list of floats
            Required k-point density.  Default value is 5.0 point per Ang^-1.
        even: bool
            Round up to even numbers.

        Returns
        -------
        Monkhorst-Pack grid size in all directions
        """
        recipcell = atoms.cell.reciprocal()
        return self.d2k_recipcell(recipcell, atoms.pbc, kptdensity, even)

    def k2d(self, atoms: MSONableAtoms, k_grid: Iterable[int]):
        """Generate the kpoint density in each direction from given k_grid.

        Parameters
        ----------
        atoms: MSONableAtoms
            Atoms object of interest.
        k_grid: Iterable[int]
            k_grid that was used.

        Returns
        -------
        Density of kpoints in each direction. result.mean() computes average density
        """
        recipcell = atoms.cell.reciprocal()
        densities = k_grid / (2 * np.pi * np.sqrt((recipcell**2).sum(axis=1)))
        return np.array(densities)

    @staticmethod
    def d2k_recipcell(
        recipcell: Cell,
        pbc: List[bool],
        kptdensity: float | Iterable[float] = 5.0,
        even: bool = True,
    ) -> Iterable[float]:
        """Convert k-point density to Monkhorst-Pack grid size.

        Parameters
        ----------
        recipcell: ASE Cell object
            The reciprocal cell
        pbc: list of Bools
            If element of pbc is True then system is periodic in that direction
        kptdensity: float or list of floats
            Required k-point density.  Default value is 3.5 point per Ang^-1.
        even: bool
            Round up to even numbers.

        Returns
        -------
        Monkhorst-Pack grid size in all directions
        """
        if not isinstance(kptdensity, Iterable):
            kptdensity = 3 * [float(kptdensity)]
        kpts = []
        for i in range(3):
            if pbc[i]:
                k = (
                    2
                    * np.pi
                    * np.sqrt((recipcell[i] ** 2).sum())
                    * float(kptdensity[i])
                )
                if even:
                    kpts.append(2 * int(np.ceil(k / 2)))
                else:
                    kpts.append(int(np.ceil(k)))
            else:
                kpts.append(1)
        return kpts


def recursive_update(d: Dict, u: Dict) -> Dict:
    """
    Update a dictionary recursively and return it.

    Parameters
    ----------
        d: Dict
            Input dictionary to modify
        u: Dict
            Dictionary of updates to apply

    Returns
    -------
    The updated dictionary.

    Example
    -------
        d = {'activate_hybrid': {"hybrid_functional": "HSE06"}}
        u = {'activate_hybrid': {"cutoff_radius": 8}}

        yields {'activate_hybrid': {"hybrid_functional": "HSE06", "cutoff_radius": 8}}}
    """
    for k, v in u.items():
        if isinstance(v, dict):
            d[k] = recursive_update(d.get(k, {}), v)
        elif k == "output" and isinstance(
            v, list
        ):  # for all other keys the list addition is not needed (I guess)
            old_v = d.get(k, [])
            d[k] = old_v + v
        else:
            d[k] = v
    return d

src/atomate2_temp/aims/sets/base.py
tpurcell90/atomate2-fhi-aims
"""Input sets for band structure calculations"""

from dataclasses import dataclass
from typing import Any, Dict, Sequence

from atomate2_temp.aims.sets.base import AimsInputGenerator
from atomate2_temp.aims.utils.bands import prepare_band_input
from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms


@dataclass
class BandStructureSetGenerator(AimsInputGenerator):
    """A generator for the band structure calculation input set

    Parameters
    ----------
    calc_type: str
        The type of calculations
    k_point_density: float
        The number of k_points per angstrom
    """

    calc_type: str = "bands"
    k_point_density: float = 20

    def get_parameter_updates(
        self, atoms: MSONableAtoms, prev_parameters: Dict[str, Any]
    ) -> Dict[str, Sequence[str]]:
        """Get the parameter updates for the calculation

        Parameters
        ----------
        atoms: MSONableAtoms
            The structure to calculate the bands for
        prev_parameters: Dict[str, Any]
            The previous parameters

        Returns
        -------
        The updated for the parameters for the output section of FHI-aims
        """
        updated_outputs = prev_parameters.get("output", list())
        updated_outputs += prepare_band_input(atoms.cell, self.k_point_density)
        return {"output": updated_outputs}


@dataclass
class GWSetGenerator(AimsInputGenerator):
    """
    A generator for the input set for calculations employing GW self-energy correction

    Parameters
    ----------
    calc_type: str
        The type of calculations
    k_point_density: float
        The number of k_points per angstrom
    """

    calc_type: str = "GW"
    k_point_density: float = 20

    def get_parameter_updates(
        self, atoms: MSONableAtoms, prev_parameters: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Get the parameter updates for the calculation

        Parameters
        ----------
        atoms: MSONableAtoms
            The structure to calculate the bands for
        prev_parameters: Dict[str, Any]
            The previous parameters

        Returns
        -------
        The updated for the parameters for the output section of FHI-aims
        """
        updates = {"anacon_type": "two-pole"}
        current_output = prev_parameters.get("output", list())
        if all(atoms.pbc):
            updates.update(
                {
                    "qpe_calc": "gw_expt",
                    "output": current_output
                    + prepare_band_input(atoms.cell, self.k_point_density),
                }
            )
        else:
            updates.update(
                {
                    "qpe_calc": "gw",
                }
            )
        return updates

src/atomate2_temp/aims/sets/bs.py
tpurcell90/atomate2-fhi-aims
"""Module defining core FHI-aims input set generators."""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict

from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

from atomate2_temp.aims.sets.base import AimsInputGenerator

__all__ = [
    "StaticSetGenerator",
    "RelaxSetGenerator",
]


@dataclass
class StaticSetGenerator(AimsInputGenerator):
    """Common class for ground-state generators.

    Parameters
    ----------
    calc_type: str
        The type of calculation
    """

    calc_type: str = "static"

    def get_parameter_updates(
        self, atoms: MSONableAtoms, prev_parameters: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Get the parameter updates for the calculation

        Parameters
        ----------
        atoms: MSONableAtoms
            The structure to calculate the bands for
        prev_parameters: Dict[str, Any]
            The previous parameters

        Returns
        -------
        The updated for the parameters for the output section of FHI-aims
        """
        return prev_parameters


@dataclass
class RelaxSetGenerator(AimsInputGenerator):
    """
    Class to generate FHI-aims relax sets.

    I.e., sets for optimization of internal and lattice coordinates.

    Parameters
    ----------
    calc_type: str
        The type of calculation
    relax_cell: bool
        If True then relax the unit cell from the structure
    max_force: float
        Maximum allowed force in the calculation
    method: str
        Method used for the geometry optimization
    """

    calc_type: str = "relaxation"
    relax_cell: bool = True
    max_force: float = 1e-3
    method: str = "trm"

    def get_parameter_updates(
        self, atoms: MSONableAtoms, prev_parameters: Dict[str, Any]
    ) -> dict:
        """Get the parameter updates for the calculation

        Parameters
        ----------
        atoms: MSONableAtoms
            The structure to calculate the bands for
        prev_parameters: Dict[str, Any]
            The previous parameters

        Returns
        -------
        The updated for the parameters for the output section of FHI-aims
        """
        updates = {"relax_geometry": f"{self.method} {self.max_force:e}"}
        if any(atoms.pbc) and self.relax_cell:
            updates["relax_unit_cell"] = "full"
        elif any(atoms.pbc):
            updates["relax_unit_cell"] = "none"

        return updates


@dataclass
class SocketIOSetGenerator(AimsInputGenerator):
    """Class to generate FHI-aims input sets for running with the socket

    Parameters
    ----------
    calc_type: str
        The type of calculation
    host: str
        The hostname for the server the socket is on
    port: int
        The port the socket server is listening on
    """

    calc_type: str = "multi_scf"
    host: str = "localhost"
    port: int = 12345

    def get_parameter_updates(
        self, atoms: MSONableAtoms, prev_parameters: Dict[str, Any]
    ) -> dict:
        """Get the parameter updates for the calculation

        Parameters
        ----------
        atoms: MSONableAtoms
            The structure to calculate the bands for
        prev_parameters: Dict[str, Any]
            The previous parameters

        Returns
        -------
        The updated for the parameters for the output section of FHI-aims
        """
        updates = {"use_pimd_wrapper": (self.host, self.port)}

        return updates

src/atomate2_temp/aims/sets/core.py
tpurcell90/atomate2-fhi-aims
"""A collection of helper utils found in atomate2 package"""


from datetime import datetime


def datetime_str() -> str:
    """
    Get a string representation of the current time.

    Returns
    -------
    str
        The current time.
    """
    return str(datetime.utcnow())

src/atomate2_temp/aims/utils/__init__.py
tpurcell90/atomate2-fhi-aims
"""Some utilities of dealing with bands. Copied from GIMS as of now; should be in its
own dedicated FHI-aims python package (with the parsers, plotters etc.)"""

import numpy as np
from ase.cell import Cell
from ase.dft.kpoints import kpoint_convert, resolve_kpt_path_string

# TODO add the same procedures but using pymatgen routines


def prepare_band_input(cell: Cell, density=20):
    """
    Prepares the band information needed for the FHI-aims control.in file.

    Parameters:

    cell: object
        ASE cell object
    density: float
        Number of kpoints per Angstrom. Default: 20
    """
    bp = cell.bandpath()
    r_kpts = resolve_kpt_path_string(bp.path, bp.special_points)

    lines_and_labels = []
    for labels, coords in zip(*r_kpts):
        dists = coords[1:] - coords[:-1]
        lengths = [np.linalg.norm(d) for d in kpoint_convert(cell, skpts_kc=dists)]
        points = np.int_(np.round(np.asarray(lengths) * density))
        # I store it here for now. Might be needed to get global info.
        lines_and_labels.append(
            [points, labels[:-1], labels[1:], coords[:-1], coords[1:]]
        )

    bands = []
    for segment in lines_and_labels:
        for points, lstart, lend, start, end in zip(*segment):
            bands.append(
                "band {:9.5f}{:9.5f}{:9.5f} {:9.5f}{:9.5f}{:9.5f} {:4} {:3}{:3}".format(
                    *start, *end, points, lstart, lend
                )
            )

    return bands

src/atomate2_temp/aims/utils/bands.py
tpurcell90/atomate2-fhi-aims
"""Module with common file names and classes used for FHI-aims flows."""

import contextlib
import os
import shutil

TMPDIR_NAME = "tmpdir"
OUTPUT_FILE_NAME: str = "aims.out"
CONTROL_FILE_NAME: str = "control.in"
PARAMS_JSON_FILE_NAME: str = "parameters.json"
GEOMETRY_FILE_NAME: str = "geometry.in"


@contextlib.contextmanager
def cwd(path: str, mkdir: bool = False, rmdir: bool = False):
    """Change cwd intermediately

    Example
    -------
    >>> with cwd(some_path):
    >>>     do so some stuff in some_path
    >>> do so some other stuff in old cwd

    Parameters
    ----------
    path: str or Path
        Path to change working directory to
    mkdir: bool
        If True make path if it does not exist
    rmdir: bool
        If True remove the working directory upon exiting
    """
    CWD = os.getcwd()

    if os.path.exists(path) is False and mkdir:
        os.makedirs(path)

    os.chdir(path)
    yield

    os.chdir(CWD)
    if rmdir:
        shutil.rmtree(path)

src/atomate2_temp/aims/utils/common.py
tpurcell90/atomate2-fhi-aims
import numpy as np
from ase.atoms import Atoms
from ase.calculators.singlepoint import SinglePointDFTCalculator
from monty.json import MontyDecoder, MSONable
from pymatgen.core import Molecule, Structure
from pymatgen.io.ase import AseAtomsAdaptor

from typing import Dict, Any

ASE_ADAPTOR = AseAtomsAdaptor()


class MSONableAtoms(Atoms, MSONable):
    def as_dict(self) -> Dict[str, Any]:
        """Represent an ASE Atoms object as a dict

        Returns
        -------
        The dictionary representation of the Atoms object
        """
        d = {"@module": self.__class__.__module__, "@class": self.__class__.__name__}

        for key, val in self.todict().items():
            d[key] = val

        if self.calc:
            d["calculated_results"] = self.calc.results

        return d

    @classmethod
    def from_dict(cls, d: Dict[str, Any]):
        """Create the Atoms object from the dictionary

        Parameters
        ----------
        d: Dict[str, Any]
            The dictionary representation of the Atoms object
        """
        decoded = {
            k: MontyDecoder().process_decoded(v)
            for k, v in d.items()
            if not k.startswith("@")
        }
        calculated_results = decoded.pop("calculated_results", None)

        atoms = Atoms.fromdict(decoded)

        calculator = SinglePointDFTCalculator(atoms)
        calculator.results = calculated_results

        return cls(atoms, calculator=calculator)

    @classmethod
    def from_pymatgen(cls, structure: Structure | Molecule):
        """Create an Atoms object from a pymatgen object"""
        return ASE_ADAPTOR.get_atoms(structure)

    @property
    def structure(self) -> Structure:
        """The pymatgen Structure of the Atoms object"""
        return ASE_ADAPTOR.get_structure(self)

    @property
    def pymatgen(self) -> Structure | Molecule:
        """The pymatgen Structure or Molecule of the Atoms object"""
        if np.any(self.pbc):
            return ASE_ADAPTOR.get_structure(self)

        return ASE_ADAPTOR.get_molecule(self)

src/atomate2_temp/aims/utils/msonable_atoms.py
tpurcell90/atomate2-fhi-aims
from ase import units as ase_units
from numpy import pi

PI = pi
EV = ase_units.CODATA["2002"]["_e"]  # = 1.602176634e-19  # [J]
THZ = 1e12
AMU = ase_units.CODATA["2002"]["_amu"]  # = 1.66053906660e-27  # [kg]
AA = 1e-10

omega_to_THz = (EV / AA**2 / AMU) ** 0.5 / THZ / 2 / PI
omega_to_THz

src/atomate2_temp/aims/utils/units.py
tpurcell90/atomate2-fhi-aims
"""Flows for calculating phonons."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from atomate2.common.jobs.utils import structure_to_conventional, structure_to_primitive
from jobflow import Flow, Maker

from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms
from atomate2_temp.common.jobs.phonons import (
    generate_frequencies_eigenvectors,
    generate_phonon_displacements,
    get_supercell_size,
    get_total_energy_per_cell,
    run_phonon_displacements,
)

if TYPE_CHECKING:
    from pathlib import Path

    from atomate2.vasp.jobs.base import BaseVaspMaker
    from emmet.core.math import Matrix3D
    from pymatgen.core.structure import Structure

    from atomate2_temp.aims.jobs.base import BaseAimsMaker

__all__ = ["BasePhononMaker"]

SUPPORTED_CODES = ["vasp", "aims"]


@dataclass
class BasePhononMaker(Maker):
    """
    Maker to calculate harmonic phonons with VASP or FHI-aims and Phonopy.

    Calculate the harmonic phonons of a material. Initially, a tight structural
    relaxation is performed to obtain a structure without forces on the atoms.
    Subsequently, supercells with one displaced atom are generated and accurate
    forces are computed for these structures. With the help of phonopy, these
    forces are then converted into a dynamical matrix. To correct for polarization
    effects, a correction of the dynamical matrix based on BORN charges can
    be performed. Finally, phonon densities of states, phonon band structures
    and thermodynamic properties are computed.

    .. Note::
        It is heavily recommended to symmetrize the structure before passing it to
        this flow. Otherwise, a different space group might be detected and too
        many displacement calculations will be generated.
        It is recommended to check the convergence parameters here and
        adjust them if necessary. The default might not be strict enough
        for your specific case.

    Parameters
    ----------
    name : str
        Name of the flows produced by this maker.
    sym_reduce : bool
        Whether to reduce the number of deformations using symmetry.
    symprec : float
        Symmetry precision to use in the
        reduction of symmetry to find the primitive/conventional cell
        (use_primitive_standard_structure, use_conventional_standard_structure)
        and to handle all symmetry-related tasks in phonopy
    displacement: float
        displacement distance for phonons
    min_length: float
        min length of the supercell that will be built
    prefer_90_degrees: bool
        if set to True, supercell algorithm will first try to find a supercell
        with 3 90 degree angles
    get_supercell_size_kwargs: dict
        kwargs that will be passed to get_supercell_size to determine supercell size
    use_symmetrized_structure: str
        allowed strings: "primitive", "conventional", None

        - "primitive" will enforce to start the phonon computation
          from the primitive standard structure
          according to Setyawan, W., & Curtarolo, S. (2010).
          High-throughput electronic band structure calculations:
          Challenges and tools. Computational Materials Science,
          49(2), 299-312. doi:10.1016/j.commatsci.2010.05.010.
          This makes it possible to use certain k-path definitions
          with this workflow. Otherwise, we must rely on seekpath
        - "conventional" will enforce to start the phonon computation
          from the conventional standard structure
          according to Setyawan, W., & Curtarolo, S. (2010).
          High-throughput electronic band structure calculations:
          Challenges and tools. Computational Materials Science,
          49(2), 299-312. doi:10.1016/j.commatsci.2010.05.010.
          We will however use seekpath and primitive structures
          as determined by from phonopy to compute the phonon band structure
    bulk_relax_maker : .BaseAimsMaker or None
        A maker to perform a tight relaxation on the bulk.
        Set to ``None`` to skip the
        bulk relaxation
    static_energy_maker : .BaseAimsMaker or None
        A maker to perform the computation of the DFT energy on the bulk.
        Set to ``None`` to skip the
        static energy computation
    born_maker: .BaseAimsMaker or None
        Maker to compute the BORN charges.
    phonon_displacement_maker : .BaseAimsMaker or None
        Maker used to compute the forces for a supercell.
    generate_frequencies_eigenvectors_kwargs : dict
        Keyword arguments passed to :obj:`generate_frequencies_eigenvectors`.
    create_thermal_displacements: bool
        Bool that determines if thermal_displacement_matrices are computed
    kpath_scheme: str
        scheme to generate kpoints. Please be aware that
        you can only use seekpath with any kind of cell
        Otherwise, please use the standard primitive structure
        Available schemes are:
        "seekpath", "hinuma", "setyawan_curtarolo", "latimer_munro".
        "seekpath" and "hinuma" are the same definition but
        seekpath can be used with any kind of unit cell as
        it relies on phonopy to handle the relationship
        to the primitive cell and not pymatgen
    code: str
        determines the dft code. currently only aims is implemented.
        This keyword might enable the implementation of other codes
        in the future
    store_force_constants: bool
        if True, force constants will be stored
    socket: bool
        If True, use the socket for the calculation
    """

    name: str = "phonon"
    sym_reduce: bool = True
    symprec: float = 1e-4
    displacement: float = 0.01
    min_length: float | None = 20.0
    prefer_90_degrees: bool = True
    get_supercell_size_kwargs: dict = field(default_factory=dict)
    use_symmetrized_structure: str | None = None
    bulk_relax_maker: BaseVaspMaker | BaseAimsMaker | None = None
    static_energy_maker: BaseVaspMaker | BaseAimsMaker | None = None
    born_maker: BaseVaspMaker | None = None
    phonon_displacement_maker: BaseVaspMaker | BaseAimsMaker = None
    create_thermal_displacements: bool = True
    generate_frequencies_eigenvectors_kwargs: dict = field(default_factory=dict)
    kpath_scheme: str = "seekpath"
    code: str = None
    store_force_constants: bool = True
    socket: bool = False

    def make(
        self,
        structure: Structure | MSONableAtoms,
        prev_dir: str | Path | None = None,
        born: list[Matrix3D] | None = None,
        epsilon_static: Matrix3D | None = None,
        total_dft_energy_per_formula_unit: float | None = None,
        supercell_matrix: Matrix3D | None = None,
    ):
        """
        Make flow to calculate the phonon properties.

        Parameters
        ----------
        structure : .Structure
            A pymatgen structure. Please start with a structure
            that is nearly fully optimized as the internal optimizers
            have very strict settings!
        prev_dir : str or Path or None
            A previous vasp calculation directory to use for copying outputs.
        born: Matrix3D
            Instead of recomputing born charges and epsilon,
            these values can also be provided manually.
            if born, epsilon_static are provided, the born
            run will be skipped
            it can be provided in the VASP convention with information for
            every atom in unit cell. Please be careful when converting
            structures within in this workflow as this could lead to errors
        epsilon_static: Matrix3D
            The high-frequency dielectric constant
            Instead of recomputing born charges and epsilon,
            these values can also be provided.
            if born, epsilon_static are provided, the born
            run will be skipped
        total_dft_energy_per_formula_unit: float
            It has to be given per formula unit (as a result in corresponding Doc)
            Instead of recomputing the energy of the bulk structure every time,
            this value can also be provided in eV. If it is provided,
            the static run will be skipped. This energy is the typical
            output dft energy of the dft workflow. No conversion needed.
        supercell_matrix: list
            instead of min_length, also a supercell_matrix can
            be given, e.g. [[1.0,0.0,0.0],[0.0,1.0,0.0],[0.0,0.0,1.0]
        """
        if self.code is None:
            raise ValueError(
                f"The code variable must be passed upon construction. Supported codes are: {SUPPORTED_CODES}"
            )

        if isinstance(structure, MSONableAtoms):
            structure = structure.pymatgen

        if self.use_symmetrized_structure not in [None, "primitive", "conventional"]:
            raise ValueError(
                "use_symmetrized_structure can only be primitive, conventional, None"
            )

        if (
            not self.use_symmetrized_structure == "primitive"
            and self.kpath_scheme != "seekpath"
        ):
            raise ValueError(
                "You can only use other kpath schemes with the primitive standard "
                "structure"
            )

        if self.kpath_scheme not in [
            "seekpath",
            "hinuma",
            "setyawan_curtarolo",
            "latimer_munro",
        ]:
            raise ValueError("kpath scheme is not implemented")

        jobs = []

        # TODO: should this be after or before structural
        # optimization as the optimization could change
        # the symmetry
        # we could add a tutorial and point out that the structure
        # should be nearly optimized before the phonon workflow
        if self.use_symmetrized_structure == "primitive":
            # These structures are compatible with many
            # of the kpath algorithms that are used for Materials Project
            prim_job = structure_to_primitive(structure, self.symprec)
            jobs.append(prim_job)
            structure = prim_job.output
        elif self.use_symmetrized_structure == "conventional":
            # it could be beneficial to use conventional
            # standard structures to arrive faster at supercells with right
            # angles
            conv_job = structure_to_conventional(structure, self.symprec)
            jobs.append(conv_job)
            structure = conv_job.output

        if self.bulk_relax_maker is not None:
            # optionally relax the structure
            bulk = self.bulk_relax_maker.make(structure, prev_dir=prev_dir)
            jobs.append(bulk)
            structure = bulk.output.structure.pymatgen
            optimization_run_job_dir = bulk.output.dir_name
            optimization_run_uuid = bulk.output.uuid
        else:
            optimization_run_job_dir = None
            optimization_run_uuid = None

        # if supercell_matrix is None, supercell size will be determined
        # after relax maker to ensure that cell lengths are really larger
        # than threshold
        if supercell_matrix is None:
            supercell_job = get_supercell_size(
                structure,
                self.min_length,
                self.prefer_90_degrees,
                **self.get_supercell_size_kwargs,
            )
            jobs.append(supercell_job)
            supercell_matrix = supercell_job.output

        # Computation of static energy
        total_dft_energy = None
        static_run_job_dir = None
        static_run_uuid = None
        if (self.static_energy_maker is not None) and (
            total_dft_energy_per_formula_unit is None
        ):
            static_job = self.static_energy_maker.make(
                structure=structure, prev_dir=prev_dir
            )
            jobs.append(static_job)
            total_dft_energy = static_job.output.output.energy
            static_run_job_dir = static_job.output.dir_name
            static_run_uuid = static_job.output.uuid
            prev_dir = static_job.output.dir_name
        elif total_dft_energy_per_formula_unit is not None:
            # to make sure that one can reuse results from Doc
            compute_total_energy_job = get_total_energy_per_cell(
                total_dft_energy_per_formula_unit, structure
            )
            jobs.append(compute_total_energy_job)
            total_dft_energy = compute_total_energy_job.output

        # get a phonon object from phonopy
        displacements = generate_phonon_displacements(
            structure=structure,
            supercell_matrix=supercell_matrix,
            displacement=self.displacement,
            sym_reduce=self.sym_reduce,
            symprec=self.symprec,
            use_symmetrized_structure=self.use_symmetrized_structure,
            kpath_scheme=self.kpath_scheme,
            code=self.code,
        )
        jobs.append(displacements)

        # perform the phonon displacement calculations
        displacement_calcs = run_phonon_displacements(
            displacements=displacements.output,
            structure=structure,
            supercell_matrix=supercell_matrix,
            phonon_maker=self.phonon_displacement_maker,
            socket=self.socket,
            prev_dir=prev_dir,
        )
        jobs.append(displacement_calcs)

        # Computation of BORN charges
        born_run_job_dir = None
        born_run_uuid = None
        if self.born_maker is not None and (born is None or epsilon_static is None):
            born_job = self.born_maker.make(structure, prev_dir=prev_dir)
            jobs.append(born_job)

            # I am not happy how we currently access "born" charges
            # This is very vasp specific code
            epsilon_static = born_job.output.calcs_reversed[0].output.epsilon_static
            born = born_job.output.calcs_reversed[0].output.outcar["born"]
            born_run_job_dir = born_job.output.dir_name
            born_run_uuid = born_job.output.uuid

        phonon_collect = generate_frequencies_eigenvectors(
            supercell_matrix=supercell_matrix,
            displacement=self.displacement,
            sym_reduce=self.sym_reduce,
            symprec=self.symprec,
            use_symmetrized_structure=self.use_symmetrized_structure,
            kpath_scheme=self.kpath_scheme,
            code=self.code,
            structure=structure,
            displacement_data=displacement_calcs.output,
            epsilon_static=epsilon_static,
            born=born,
            total_dft_energy=total_dft_energy,
            static_run_job_dir=static_run_job_dir,
            static_run_uuid=static_run_uuid,
            born_run_job_dir=born_run_job_dir,
            born_run_uuid=born_run_uuid,
            optimization_run_job_dir=optimization_run_job_dir,
            optimization_run_uuid=optimization_run_uuid,
            create_thermal_displacements=self.create_thermal_displacements,
            store_force_constants=self.store_force_constants,
            **self.generate_frequencies_eigenvectors_kwargs,
        )

        jobs.append(phonon_collect)

        # create a flow including all jobs for a phonon computation
        return Flow(jobs, phonon_collect.output)

src/atomate2_temp/common/flows/phonons.py
tpurcell90/atomate2-fhi-aims
"""Jobs for running phonon calculations."""

from __future__ import annotations

import contextlib
import logging
from typing import TYPE_CHECKING

import numpy as np
from emmet.core.math import Matrix3D
from jobflow import Flow, Response, job
from phonopy import Phonopy
from pymatgen.core import Structure
from pymatgen.io.phonopy import get_phonopy_structure, get_pmg_structure
from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine
from pymatgen.phonon.dos import PhononDos
from pymatgen.transformations.advanced_transformations import (
    CubicSupercellTransformation,
)

from atomate2_temp.common.schemas.phonons import ForceConstants, PhononBSDOSDoc
from atomate2_temp.common.utils.phonons import get_factor


if TYPE_CHECKING:
    from pathlib import Path

    from atomate2.forcefields.jobs import ForceFieldStaticMaker
    from atomate2.vasp.jobs.base import BaseVaspMaker

    from atomate2_temp.aims.jobs.base import BaseAimsMaker


logger = logging.getLogger(__name__)

__all__ = [
    "get_total_energy_per_cell",
    "get_supercell_size",
    "generate_phonon_displacements",
    "run_phonon_displacements",
    "generate_frequencies_eigenvectors",
]


@job
def get_total_energy_per_cell(
    total_dft_energy_per_formula_unit: float, structure: Structure
):
    """
    Job that computes total dft energy of the cell.

    Parameters
    ----------
    total_dft_energy_per_formula_unit: float
        Total DFT energy in eV per formula unit.
    structure: Structure object
        Corresponding structure object.
    """
    formula_units = (
        structure.composition.num_atoms
        / structure.composition.reduced_composition.num_atoms
    )

    return total_dft_energy_per_formula_unit * formula_units


@job
def get_supercell_size(
    structure: Structure, min_length: float, prefer_90_degrees: bool, **kwargs
):
    """
    Determine supercell size with given min_length.

    Parameters
    ----------
    structure: Structure Object
        Input structure that will be used to determine supercell
    min_length: float
        minimum length of cell in Angstrom
    prefer_90_degrees: bool
        if True, the algorithm will try to find a cell with 90 degree angles first
    **kwargs:
        Additional parameters that can be set.
    """
    kwargs.setdefault("min_atoms", None)
    kwargs.setdefault("force_diagonal", False)

    if not prefer_90_degrees:
        kwargs.setdefault("max_atoms", None)
        transformation = CubicSupercellTransformation(
            min_length=min_length,
            min_atoms=kwargs["min_atoms"],
            max_atoms=kwargs["max_atoms"],
            force_diagonal=kwargs["force_diagonal"],
            force_90_degrees=False,
        )
        transformation.apply_transformation(structure=structure)
    else:
        max_atoms = kwargs.get("max_atoms", 1000)
        kwargs.setdefault("angle_tolerance", 1e-2)
        try:
            transformation = CubicSupercellTransformation(
                min_length=min_length,
                min_atoms=kwargs["min_atoms"],
                max_atoms=max_atoms,
                force_diagonal=kwargs["force_diagonal"],
                force_90_degrees=True,
                angle_tolerance=kwargs["angle_tolerance"],
            )
            transformation.apply_transformation(structure=structure)

        except AttributeError:
            kwargs.setdefault("max_atoms", None)

            transformation = CubicSupercellTransformation(
                min_length=min_length,
                min_atoms=kwargs["min_atoms"],
                max_atoms=kwargs["max_atoms"],
                force_diagonal=kwargs["force_diagonal"],
                force_90_degrees=False,
            )
            transformation.apply_transformation(structure=structure)

    return transformation.transformation_matrix.tolist()


@job
def generate_phonon_displacements(
    structure: Structure,
    supercell_matrix: np.array,
    displacement: float,
    sym_reduce: bool,
    symprec: float,
    use_symmetrized_structure: str | None,
    kpath_scheme: str,
    code: str,
):
    """
    Generate displaced structures with phonopy.

    Parameters
    ----------
    structure: Structure object
        Fully optimized input structure for phonon run
    supercell_matrix: np.array
        array to describe supercell matrix
    displacement: float
        displacement in Angstrom
    sym_reduce: bool
        if True, symmetry will be used to generate displacements
    symprec: float
        precision to determine symmetry
    use_symmetrized_structure: str or None
        primitive, conventional or None
    kpath_scheme: str
        scheme to generate kpath
    code:
        code to perform the computations
    """
    cell = get_phonopy_structure(structure)
    factor = get_factor(code)

    # a bit of code repetition here as I currently
    # do not see how to pass the phonopy object?
    if use_symmetrized_structure == "primitive" and kpath_scheme != "seekpath":
        primitive_matrix: list[list[float]] | str = [
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [0.0, 0.0, 1.0],
        ]
    else:
        primitive_matrix = "auto"

    phonon = Phonopy(
        cell,
        supercell_matrix,
        primitive_matrix=primitive_matrix,
        factor=factor,
        symprec=symprec,
        is_symmetry=sym_reduce,
    )
    phonon.generate_displacements(distance=displacement)

    supercells = phonon.supercells_with_displacements

    return [get_pmg_structure(cell) for cell in supercells]


@job(
    output_schema=PhononBSDOSDoc,
    data=[PhononDos, PhononBandStructureSymmLine, ForceConstants],
)
def generate_frequencies_eigenvectors(
    structure: Structure,
    supercell_matrix: np.array,
    displacement: float,
    sym_reduce: bool,
    symprec: float,
    use_symmetrized_structure: str | None,
    kpath_scheme: str,
    code: str,
    displacement_data: dict[str, list],
    total_dft_energy: float,
    epsilon_static: Matrix3D = None,
    born: Matrix3D = None,
    **kwargs,
):
    """
    Analyze the phonon runs and summarize the results.

    Parameters
    ----------
    structure: Structure object
        Fully optimized structure used for phonon runs
    supercell_matrix: np.array
        array to describe supercell
    displacement: float
        displacement in Angstrom used for supercell computation
    sym_reduce: bool
        if True, symmetry will be used in phonopy
    symprec: float
        precision to determine symmetry
    use_symmetrized_structure: str
        primitive, conventional, None are allowed
    kpath_scheme: str
        kpath scheme for phonon band structure computation
    code: str
        code to run computations
    displacement_data: dict
        outputs from displacements
    total_dft_energy: float
        total dft energy in eV per cell
    epsilon_static: Matrix3D
        The high-frequency dielectric constant
    born: Matrix3D
        Born charges
    kwargs: dict
        Additional parameters that are passed to PhononBSDOSDoc.from_forces_born

    """
    phonon_doc = PhononBSDOSDoc.from_forces_born(
        structure=structure,
        supercell_matrix=supercell_matrix,
        displacement=displacement,
        sym_reduce=sym_reduce,
        symprec=symprec,
        use_symmetrized_structure=use_symmetrized_structure,
        kpath_scheme=kpath_scheme,
        code=code,
        displacement_data=displacement_data,
        total_dft_energy=total_dft_energy,
        epsilon_static=epsilon_static,
        born=born,
        **kwargs,
    )

    return phonon_doc


@job(data=["forces", "displaced_structures"])
def run_phonon_displacements(
    displacements,
    structure: Structure,
    supercell_matrix: Matrix3D,
    phonon_maker: BaseVaspMaker | ForceFieldStaticMaker | BaseAimsMaker,
    prev_dir: str | Path = None,
    socket: bool = False,
):
    """
    Run phonon displacements.

    Note, this job will replace itself with N displacement calculations,
    or a single socket calculation for all displacments.

    Parameters
    ----------
    displacements: Iterable
        All displacements to calculate
    structure: Structure object
        Fully optimized structure used for phonon computations
    supercell_matrix: Matrix3D
        supercell matrix for meta data
    phonon_maker : BaseVaspMaker or ForceFieldStaticMaker or BaseAimsMaker
        A maker to use to generate dispacement calculations
    prev_dir: str or Path
        The previous working directory
    socket: bool
        If True use the socket-io interface to increase performance
    """
    phonon_jobs = []
    outputs: dict[str, list] = {
        "displacement_number": [],
        "forces": [],
        "uuids": [],
        "dirs": [],
    }

    if socket:
        if prev_dir is not None:
            phonon_job = phonon_maker.make(displacements, prev_dir=prev_dir)
        else:
            phonon_job = phonon_maker.make(displacements)
        info = {
            "original_structure": structure,
            "supercell_matrix": supercell_matrix,
            "displaced_structures": displacements,
        }
        phonon_job.update_maker_kwargs(
            {"_set": {"write_additional_data->phonon_info:json": info}}, dict_mod=True
        )
        phonon_jobs.append(phonon_job)
        outputs["displacement_number"] = list(range(len(displacements)))
        outputs["uuids"] = [phonon_job.output.uuid] * len(displacements)
        outputs["dirs"] = [phonon_job.output.dir_name] * len(displacements)
        outputs["forces"] = phonon_job.output.output.all_forces
    else:
        for i, displacement in enumerate(displacements):
            if prev_dir is not None:
                phonon_job = phonon_maker.make(displacement, prev_dir=prev_dir)
            else:
                phonon_job = phonon_maker.make(displacement)
            phonon_job.append_name(f" {i + 1}/{len(displacements)}")

            # we will add some meta data
            info = {
                "displacement_number": i,
                "original_structure": structure,
                "supercell_matrix": supercell_matrix,
                "displaced_structure": displacement,
            }
            with contextlib.suppress(Exception):
                phonon_job.update_maker_kwargs(
                    {"_set": {"write_additional_data->phonon_info:json": info}},
                    dict_mod=True,
                )
            phonon_jobs.append(phonon_job)
            outputs["displacement_number"].append(i)
            outputs["uuids"].append(phonon_job.output.uuid)
            outputs["dirs"].append(phonon_job.output.dir_name)
            outputs["forces"].append(phonon_job.output.output.forces)

    displacement_flow = Flow(phonon_jobs, outputs)
    return Response(replace=displacement_flow)

src/atomate2_temp/common/jobs/phonons.py
tpurcell90/atomate2-fhi-aims
"""Schemas for phonon documents."""

import copy
import logging
from typing import Dict, List, Optional, Union

import numpy as np
from emmet.core.math import Matrix3D
from emmet.core.structure import StructureMetadata
from monty.json import MSONable
from phonopy import Phonopy
from phonopy.phonon.band_structure import get_band_qpoints_and_path_connections
from phonopy.structure.symmetry import symmetrize_borns_and_epsilon
from pydantic import BaseModel, Field
from pymatgen.core import Structure
from pymatgen.io.phonopy import (
    get_ph_bs_symm_line,
    get_ph_dos,
    get_phonopy_structure,
    get_pmg_structure,
)
from pymatgen.io.vasp import Kpoints
from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine
from pymatgen.phonon.dos import PhononDos
from pymatgen.phonon.plotter import PhononBSPlotter, PhononDosPlotter
from pymatgen.symmetry.bandstructure import HighSymmKpath
from pymatgen.symmetry.kpath import KPathSeek

from atomate2_temp.common.utils.phonons import get_factor

logger = logging.getLogger(__name__)

__all__ = [
    "PhononBSDOSDoc",
    "PhononComputationalSettings",
    "PhononUUIDs",
    "PhononJobDirs",
    "ThermalDisplacementData",
]


class PhononComputationalSettings(BaseModel):
    """Collection to store computational settings for the phonon computation."""

    # could be optional and implemented at a later stage?
    npoints_band: int = Field("number of points for band structure computation")
    kpath_scheme: str = Field("indicates the kpath scheme")
    kpoint_density_dos: int = Field(
        "number of points for computation of free energies and densities of states",
    )


class ThermalDisplacementData(BaseModel):
    """Collection to store information on the thermal displacement matrices."""

    freq_min_thermal_displacements: float = Field(
        "cutoff frequency in THz to avoid numerical issues in the "
        "computation of the thermal displacement parameters"
    )
    thermal_displacement_matrix_cif: List[List[Matrix3D]] = Field(
        None, description="field including thermal displacement matrices in CIF format"
    )
    thermal_displacement_matrix: List[List[Matrix3D]] = Field(
        None,
        description="field including thermal displacement matrices in Cartesian "
        "coordinate system",
    )
    temperatures_thermal_displacements: List[int] = Field(
        None,
        description="temperatures at which the thermal displacement matrices"
        "have been computed",
    )


class PhononUUIDs(BaseModel):
    """Collection to save all uuids connected to the phonon run."""

    optimization_run_uuid: str = Field(None, description="optimization run uuid")
    displacements_uuids: List[str] = Field(
        None, description="The uuids of the displacement jobs."
    )
    static_run_uuid: str = Field(None, description="static run uuid")
    born_run_uuid: str = Field(None, description="born run uuid")


class ForceConstants(MSONable):
    """A force constants class."""

    def __init__(self, force_constants: List[List[Matrix3D]]):
        self.force_constants = force_constants


class PhononJobDirs(BaseModel):
    """Collection to save all job directories relevant for the phonon run."""

    displacements_job_dirs: List[str] = Field(
        None, description="The directories where the displacement jobs were run."
    )
    static_run_job_dir: str = Field(
        None, description="Directory where static run was performed."
    )
    born_run_job_dir: str = Field(
        None, description="Directory where born run was performed."
    )
    optimization_run_job_dir: str = Field(
        None, description="Directory where optimization run was performed."
    )


class PhononBSDOSDoc(StructureMetadata):
    """Collection of all data produced by the phonon workflow."""

    structure: Structure = Field(
        None,
        description="Structure of Materials Project.",
    )

    phonon_bandstructure: PhononBandStructureSymmLine = Field(
        None,
        description="Phonon band structure object.",
    )

    phonon_dos: PhononDos = Field(
        None,
        description="Phonon density of states object.",
    )

    free_energies: List[float] = Field(
        None,
        description="vibrational part of the free energies in J/mol per "
        "formula unit for temperatures in temperature_list",
    )

    heat_capacities: List[float] = Field(
        None,
        description="heat capacities in J/K/mol per "
        "formula unit for temperatures in temperature_list",
    )

    internal_energies: List[float] = Field(
        None,
        description="internal energies in  J/mol per "
        "formula unit for temperatures in temperature_list",
    )
    entropies: List[float] = Field(
        None,
        description="entropies in J/(K*mol) per formula unit"
        "for temperatures in temperature_list ",
    )

    temperatures: List[int] = Field(
        None,
        description="temperatures at which the vibrational"
        " part of the free energies"
        " and other properties have been computed",
    )

    total_dft_energy: Optional[float] = Field("total DFT energy per formula unit in eV")

    has_imaginary_modes: bool = Field(
        None, description="if true, structure has imaginary modes"
    )

    # needed, e.g. to compute Grueneisen parameter etc
    force_constants: ForceConstants = Field(
        None, description="Force constants between every pair of atoms in the structure"
    )

    born: List[Matrix3D] = Field(
        None,
        description="born charges as computed from phonopy. Only for symmetrically "
        "different atoms",
    )

    epsilon_static: Matrix3D = Field(
        None, description="The high-frequency dielectric constant"
    )

    supercell_matrix: Matrix3D = Field("matrix describing the supercell")
    primitive_matrix: Matrix3D = Field(
        "matrix describing relationship to primitive cell"
    )

    code: str = Field("String describing the code for the computation")

    phonopy_settings: PhononComputationalSettings = Field(
        "Field including settings for Phonopy"
    )

    thermal_displacement_data: Optional[ThermalDisplacementData] = Field(
        "Includes all data of the computation of the thermal displacements"
    )

    jobdirs: PhononJobDirs = Field("Field including all relevant job directories")

    uuids: PhononUUIDs = Field("Field including all relevant uuids")

    @classmethod
    def from_forces_born(
        cls,
        structure: Structure,
        supercell_matrix: np.array,
        displacement: float,
        sym_reduce: bool,
        symprec: float,
        use_symmetrized_structure: Union[str, None],
        kpath_scheme: str,
        code: str,
        displacement_data: Dict[str, list],
        total_dft_energy: float,
        epsilon_static: Matrix3D = None,
        born: Matrix3D = None,
        **kwargs,
    ):
        """
        Generate collection of phonon data.

        Parameters
        ----------
        structure: Structure object
        supercell_matrix: numpy array describing the supercell
        displacement: float
            size of displacement in angstrom
        sym_reduce: bool
            if True, phonopy will use symmetry
        symprec: float
            precision to determine kpaths,
            primitive cells and symmetry in phonopy and pymatgen
        use_symmetrized_structure: str
            primitive, conventional or None
        kpath_scheme: str
            kpath scheme to generate phonon band structure
        code: str
            which code was used for computation
        displacement_data:
            output of the VASP displacement data
        total_dft_energy: float
            total energy in eV per cell
        epsilon_static: Matrix3D
            The high-frequency dielectric constant
        born: Matrix3D
            born charges
        **kwargs:
            additional arguments
        """
        factor = get_factor(code)
        # This opens the opportunity to add support for other codes
        # that are supported by phonopy

        cell = get_phonopy_structure(structure)

        if use_symmetrized_structure == "primitive" and kpath_scheme != "seekpath":
            primitive_matrix: Union[List[List[float]], str] = [
                [1.0, 0.0, 0.0],
                [0.0, 1.0, 0.0],
                [0.0, 0.0, 1.0],
            ]
        else:
            primitive_matrix = "auto"
        phonon = Phonopy(
            cell,
            supercell_matrix,
            primitive_matrix=primitive_matrix,
            factor=factor,
            symprec=symprec,
            is_symmetry=sym_reduce,
        )
        phonon.generate_displacements(distance=displacement)
        set_of_forces = [np.array(forces) for forces in displacement_data["forces"]]

        if born is not None and epsilon_static is not None:
            if len(structure) == len(born):
                borns, epsilon = symmetrize_borns_and_epsilon(
                    ucell=phonon.unitcell,
                    borns=np.array(born),
                    epsilon=np.array(epsilon_static),
                    symprec=symprec,
                    primitive_matrix=phonon.primitive_matrix,
                    supercell_matrix=phonon.supercell_matrix,
                    is_symmetry=kwargs.get("symmetrize_born", True),
                )
            else:
                raise ValueError(
                    "Number of born charges does not agree with number of atoms"
                )
            if code == "vasp" and not np.all(np.isclose(borns, 0.0)):
                phonon.nac_params = {
                    "born": borns,
                    "dielectric": epsilon,
                    "factor": 14.399652,
                }
            # Other codes could be added here
        else:
            borns = None
            epsilon = None

        # Produces all force constants
        phonon.produce_force_constants(forces=set_of_forces)

        # with phonon.load("phonopy.yaml") the phonopy API can be used
        phonon.save("phonopy.yaml")

        # get phonon band structure
        kpath_dict, kpath_concrete = cls.get_kpath(
            structure=get_pmg_structure(phonon.primitive),
            kpath_scheme=kpath_scheme,
            symprec=symprec,
        )

        npoints_band = kwargs.get("npoints_band", 101)
        qpoints, connections = get_band_qpoints_and_path_connections(
            kpath_concrete, npoints=kwargs.get("npoints_band", 101)
        )

        # phonon band structures will always be cmouted
        filename_band_yaml = "phonon_band_structure.yaml"

        # TODO: potentially add kwargs to avoid computation of eigenvectors
        phonon.run_band_structure(
            qpoints,
            path_connections=connections,
            with_eigenvectors=kwargs.get("band_structure_eigenvectors", False),
            is_band_connection=kwargs.get("band_structure_eigenvectors", False),
        )
        phonon.write_yaml_band_structure(filename=filename_band_yaml)
        bs_symm_line = get_ph_bs_symm_line(
            filename_band_yaml, labels_dict=kpath_dict, has_nac=born is not None
        )
        new_plotter = PhononBSPlotter(bs=bs_symm_line)

        new_plotter.save_plot(
            "phonon_band_structure.eps",
            img_format=kwargs.get("img_format", "eps"),
            units=kwargs.get("units", "THz"),
        )

        # will determine if imaginary modes are present in the structure
        imaginary_modes = bs_symm_line.has_imaginary_freq(
            tol=kwargs.get("tol_imaginary_modes", 1e-5)
        )

        # gets data for visualization on website - yaml is also enough
        if kwargs.get("band_structure_eigenvectors", False):
            bs_symm_line.write_phononwebsite("phonon_website.json")

        # get phonon density of states
        filename_dos_yaml = "phonon_dos.yaml"

        kpoint_density_dos = kwargs.get("kpoint_density_dos", 7000)
        kpoint = Kpoints.automatic_density(
            structure=get_pmg_structure(phonon.primitive),
            kppa=kpoint_density_dos,
            force_gamma=True,
        )
        phonon.run_mesh(kpoint.kpts[0])
        phonon.run_total_dos()
        phonon.write_total_dos(filename=filename_dos_yaml)
        dos = get_ph_dos(filename_dos_yaml)
        new_plotter_dos = PhononDosPlotter()
        new_plotter_dos.add_dos(label="total", dos=dos)
        new_plotter_dos.save_plot(
            filename="phonon_dos.eps",
            img_format=kwargs.get("img_format", "eps"),
            units=kwargs.get("units", "THz"),
        )

        # compute vibrational part of free energies per formula unit
        temperature_range = np.arange(
            kwargs.get("tmin", 0), kwargs.get("tmax", 500), kwargs.get("tstep", 10)
        )

        free_energies = [
            dos.helmholtz_free_energy(
                structure=get_pmg_structure(phonon.primitive), t=temperature
            )
            for temperature in temperature_range
        ]

        entropies = [
            dos.entropy(structure=get_pmg_structure(phonon.primitive), t=temperature)
            for temperature in temperature_range
        ]

        internal_energies = [
            dos.internal_energy(
                structure=get_pmg_structure(phonon.primitive), t=temperature
            )
            for temperature in temperature_range
        ]

        heat_capacities = [
            dos.cv(structure=get_pmg_structure(phonon.primitive), t=temperature)
            for temperature in temperature_range
        ]

        # will compute thermal displacement matrices
        # for the primitive cell (phonon.primitive!)
        # only this is available in phonopy
        if kwargs.get("create_thermal_displacements", False):
            phonon.run_mesh(
                kpoint.kpts[0], with_eigenvectors=True, is_mesh_symmetry=False
            )
            freq_min_thermal_displacements = kwargs.get(
                "freq_min_thermal_displacements", 0.0
            )
            phonon.run_thermal_displacement_matrices(
                t_min=kwargs.get("tmin_thermal_displacements", 0),
                t_max=kwargs.get("tmax_thermal_displacements", 500),
                t_step=kwargs.get("tstep_thermal_displacements", 100),
                freq_min=freq_min_thermal_displacements,
            )

            temperature_range_thermal_displacements = np.arange(
                kwargs.get("tmin_thermal_displacements", 0),
                kwargs.get("tmax_thermal_displacements", 500),
                kwargs.get("tstep_thermal_displacements", 100),
            )
            for i, temp in enumerate(temperature_range_thermal_displacements):
                phonon.thermal_displacement_matrices.write_cif(
                    phonon.primitive,
                    i,
                    filename="tdispmat_" + str(temp) + "K.cif",
                )
            tdisp_mat = (
                phonon._thermal_displacement_matrices.thermal_displacement_matrices.tolist()
            )

            tdisp_mat_cif = (
                phonon._thermal_displacement_matrices.thermal_displacement_matrices_cif.tolist()
            )

        else:
            tdisp_mat = None
            tdisp_mat_cif = None

        formula_units = (
            structure.composition.num_atoms
            / structure.composition.reduced_composition.num_atoms
        )

        total_dft_energy_per_formula_unit = (
            total_dft_energy / formula_units if total_dft_energy is not None else None
        )

        return cls.from_structure(
            structure=structure,
            meta_structure=structure,
            phonon_bandstructure=bs_symm_line,
            phonon_dos=dos,
            free_energies=free_energies,
            internal_energies=internal_energies,
            heat_capacities=heat_capacities,
            entropies=entropies,
            temperatures=temperature_range.tolist(),
            total_dft_energy=total_dft_energy_per_formula_unit,
            has_imaginary_modes=imaginary_modes,
            force_constants={"force_constants": phonon.force_constants.tolist()}
            if kwargs["store_force_constants"]
            else None,
            born=borns.tolist() if borns is not None else None,
            epsilon_static=epsilon.tolist() if epsilon is not None else None,
            supercell_matrix=phonon.supercell_matrix.tolist(),
            primitive_matrix=phonon.primitive_matrix.tolist(),
            code=code,
            thermal_displacement_data={
                "temperatures_thermal_displacements": temperature_range_thermal_displacements.tolist(),  # noqa: E501
                "thermal_displacement_matrix_cif": tdisp_mat_cif,
                "thermal_displacement_matrix": tdisp_mat,
                "freq_min_thermal_displacements": freq_min_thermal_displacements,
            }
            if kwargs.get("create_thermal_displacements", False)
            else None,
            jobdirs={
                "displacements_job_dirs": displacement_data["dirs"],
                "static_run_job_dir": kwargs["static_run_job_dir"],
                "born_run_job_dir": kwargs["born_run_job_dir"],
                "optimization_run_job_dir": kwargs["optimization_run_job_dir"],
            },
            uuids={
                "displacements_uuids": displacement_data["uuids"],
                "born_run_uuid": kwargs["born_run_uuid"],
                "optimization_run_uuid": kwargs["optimization_run_uuid"],
                "static_run_uuid": kwargs["static_run_uuid"],
            },
            phonopy_settings={
                "npoints_band": npoints_band,
                "kpath_scheme": kpath_scheme,
                "kpoint_density_dos": kpoint_density_dos,
            },
        )

    @staticmethod
    def get_kpath(
        structure: Structure, kpath_scheme: str, symprec: float, **kpath_kwargs
    ):
        """
        Get high-symmetry points in k-space in phonopy format.

        Parameters
        ----------
        structure: Structure Object
        kpath_scheme: str
            string describing kpath
        symprec: float
            precision for symmetry determination
        **kpath_kwargs:
            additional parameters that can be passed to this method as a dict
        """
        if kpath_scheme in [
            "setyawan_curtarolo",
            "latimer_munro",
            "hinuma",
        ]:
            highsymmkpath = HighSymmKpath(
                structure, path_type=kpath_scheme, symprec=symprec, **kpath_kwargs
            )
            kpath = highsymmkpath.kpath
        elif kpath_scheme == "seekpath":
            highsymmkpath = KPathSeek(structure, symprec=symprec, **kpath_kwargs)
            kpath = highsymmkpath._kpath

        path = copy.deepcopy(kpath["path"])

        for ilabelset, labelset in enumerate(kpath["path"]):
            for ilabel, label in enumerate(labelset):
                path[ilabelset][ilabel] = kpath["kpoints"][label]
        return kpath["kpoints"], path

src/atomate2_temp/common/schemas/phonons.py
tpurcell90/atomate2-fhi-aims
from phonopy.units import VaspToTHz

from atomate2_temp.aims.utils.units import omega_to_THz


def get_factor(code: str):
    """
    Get the frequency conversion factor to THz for each code

    Parameters
    ----------
    code: str
        The code to get the conversion factor for

    Returns
    -------
    float
        The correct conversion factor

    Raises
    ------
    ValueError
        If code is not defined
    """
    if code == "vasp":
        return VaspToTHz
    elif code == "aims":
        return omega_to_THz  # Based on CODATA 2002
    else:
        raise ValueError(f"Frequency conversion factor for code ({code}) not defined.")

src/atomate2_temp/common/utils/phonons.py
tpurcell90/atomate2-fhi-aims
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from atomate2.vasp.flows.core import DoubleRelaxMaker
from atomate2.vasp.jobs.core import DielectricMaker, StaticMaker, TightRelaxMaker
from atomate2.vasp.sets.core import StaticSetGenerator

from atomate2_temp.common.flows.phonons import BasePhononMaker
from atomate2_temp.vasp.jobs.phonons import PhononDisplacementMaker

if TYPE_CHECKING:
    from atomate2.vasp.jobs.base import BaseVaspMaker


@dataclass
class PhononMaker(BasePhononMaker):
    """
    Maker to calculate harmonic phonons with VASP and Phonopy.

    Overwrites the default Makers for the common PhononMaker

    Parameters
    ----------
    bulk_relax_maker : .BaseVaspMaker or None
        A maker to perform a tight relaxation on the bulk.
        Set to ``None`` to skip the
        bulk relaxation
    static_energy_maker : .BaseVaspMaker or None
        A maker to perform the computation of the DFT energy on the bulk.
        Set to ``None`` to skip the
        static energy computation
    born_maker: .BaseVaspMaker or None
        Maker to compute the BORN charges.
    phonon_displacement_maker : .BaseVaspMaker or None
        Maker used to compute the forces for a supercell.
    """

    code: str = "vasp"
    bulk_relax_maker: BaseVaspMaker | None = field(
        default_factory=lambda: DoubleRelaxMaker.from_relax_maker(TightRelaxMaker())
    )
    static_energy_maker: BaseVaspMaker | None = field(
        default_factory=lambda: StaticMaker(
            input_set_generator=StaticSetGenerator(auto_ispin=True)
        )
    )
    born_maker: BaseVaspMaker | None = field(default_factory=DielectricMaker)
    phonon_displacement_maker: BaseVaspMaker = field(
        default_factory=PhononDisplacementMaker
    )

src/atomate2_temp/vasp/flows/phonons.py
tpurcell90/atomate2-fhi-aims
from dataclasses import dataclass, field

from atomate2.vasp.jobs.base import BaseVaspMaker
from atomate2.vasp.sets.base import VaspInputGenerator
from atomate2.vasp.sets.core import StaticSetGenerator


@dataclass
class PhononDisplacementMaker(BaseVaspMaker):
    """
    Maker to perform a static calculation as a part of the finite displacement method.

    The input set is for a static run with tighter convergence parameters.
    Both the k-point mesh density and convergence parameters
    are stricter than a normal relaxation.

    Parameters
    ----------
    name : str
        The job name.
    input_set_generator : .VaspInputGenerator
        A generator used to make the input set.
    write_input_set_kwargs : dict
        Keyword arguments that will get passed to :obj:`.write_vasp_input_set`.
    copy_vasp_kwargs : dict
        Keyword arguments that will get passed to :obj:`.copy_vasp_outputs`.
    run_vasp_kwargs : dict
        Keyword arguments that will get passed to :obj:`.run_vasp`.
    task_document_kwargs : dict
        Keyword arguments that will get passed to :obj:`.TaskDoc.from_directory`.
    stop_children_kwargs : dict
        Keyword arguments that will get passed to :obj:`.should_stop_children`.
    write_additional_data : dict
        Additional data to write to the current directory. Given as a dict of
        {filename: data}. Note that if using FireWorks, dictionary keys cannot contain
        the "." character which is typically used to denote file extensions. To avoid
        this, use the ":" character, which will automatically be converted to ".". E.g.
        ``{"my_file:txt": "contents of the file"}``.
    """

    name: str = "phonon static"
    input_set_generator: VaspInputGenerator = field(
        default_factory=lambda: StaticSetGenerator(
            user_kpoints_settings={"reciprocal_density": 100},
            user_incar_settings={
                "IBRION": 2,
                "ISIF": 3,
                "ENCUT": 700,
                "EDIFF": 1e-7,
                "LAECHG": False,
                "LREAL": False,
                "ALGO": "Normal",
                "NSW": 0,
                "LCHARG": False,
            },
            auto_ispin=True,
        )
    )

src/atomate2_temp/vasp/jobs/phonons.py
tpurcell90/atomate2-fhi-aims
"""Utils or testing"""

from glob import glob
import json
from pathlib import Path


def compare_files(test_name, work_dir, ref_dir):
    for file in glob(f"{work_dir / test_name}/*in"):
        assert (
            open(file, "rt").readlines()[4:]
            == open(f"{ref_dir / test_name / Path(file).name}", "rt").readlines()[4:]
        )
    ref = json.load(open(f"{ref_dir / test_name}/parameters.json"))
    ref.pop("species_dir", None)
    check = json.load(open(f"{work_dir / test_name}/parameters.json"))
    check.pop("species_dir", None)

    assert ref == check

tests/__init__.py
tpurcell90/atomate2-fhi-aims
import logging
import os
from pathlib import Path
from typing import Union, Sequence, Literal

import pytest

from ase.build import bulk, molecule

import atomate2_temp.aims.jobs.base
import atomate2_temp.aims.run
from atomate2_temp.aims.sets.base import AimsInputGenerator

_REF_PATHS = {}
_FAKE_RUN_AIMS_KWARGS = {}
_VFILES = "control.in"


logger = logging.getLogger(__name__)


def pytest_addoption(parser):
    parser.addoption(
        "--generate-test-data",
        action="store_true",
        help="Runs FHI-aims to create test data; runs tests against the created outputs",
    )


@pytest.fixture
def Si():
    return bulk("Si")


@pytest.fixture
def O2():
    return molecule("O2")


@pytest.fixture(scope="session")
def species_dir():
    return Path(__file__).resolve().parent / "species_dir"


@pytest.fixture(scope="session")
def ref_path():
    from pathlib import Path

    module_dir = Path(__file__).resolve().parent
    test_dir = module_dir / "test_data"
    return test_dir.resolve()


@pytest.fixture
def should_mock_aims(request):
    return not request.config.getoption("--generate-test-data")


@pytest.fixture()
def mock_aims(monkeypatch, ref_path, should_mock_aims):
    """
    This fixture allows one to mock (fake) running FHI-aims.

    To use the fixture successfully, the following steps must be followed:
    1. "mock_aims" should be included as an argument to any test that would like to use
       its functionally.
    2. For each job in your workflow, you should prepare a reference directory
       containing two folders "inputs" (containing the reference input files expected
       to be produced by write_aims_input_set) and "outputs" (containing the expected
       output files to be produced by run_aims). These files should reside in a
       subdirectory of "tests/test_data/aims".
    3. Create a dictionary mapping each job name to its reference directory. Note that
       you should supply the reference directory relative to the "tests/test_data/aims"
       folder. For example, if your calculation has one job named "static" and the
       reference files are present in "tests/test_data/aims/Si_static", the dictionary
       would look like: ``{"static": "Si_static"}``.
    4. Optional (does not work yet): create a dictionary mapping each job name to
       custom keyword arguments that will be supplied to fake_run_aims.
       This way you can configure which control.in settings are expected for each job.
       For example, if your calculation has one job named "static" and you wish to validate
       that "xc" is set correctly in the control.in, your dictionary would look like
       ``{"static": {"input_settings": {"relativistic": "atomic_zora scalar"}}``.
    5. Inside the test function, call `mock_aims(ref_paths, fake_aims_kwargs)`, where
       ref_paths is the dictionary created in step 3 and fake_aims_kwargs is the
       dictionary created in step 4.
    6. Run your aims job after calling `mock_aims`.

    For examples, see the tests in tests/aims/jobs/core.py.
    """

    def mock_run_aims(*args, **kwargs):
        from jobflow import CURRENT_JOB

        name = CURRENT_JOB.job.name
        ref_dir = ref_path / _REF_PATHS[name]
        fake_run_aims(ref_dir, **_FAKE_RUN_AIMS_KWARGS.get(name, {}))

    get_input_set_orig = AimsInputGenerator.get_input_set

    def generate_test_data(*args, **kwargs):
        """A monkey patch for atomate2_temp.aims.run.run_aims that runs the actual executable and copies
        inputs and outputs to the test data directory"""
        import shutil
        from jobflow import CURRENT_JOB

        input_files = ["control.in", "geometry.in", "parameters.json"]
        name = CURRENT_JOB.job.name
        ref_dir = ref_path / _REF_PATHS[name]
        # running aims
        atomate2_temp.aims.run.run_aims()
        # copy output files
        output_files = [f for f in Path.cwd().glob("*") if f.name not in input_files]
        shutil.rmtree(ref_dir, ignore_errors=True)
        os.makedirs(ref_dir / "inputs")
        os.makedirs(ref_dir / "outputs")
        for f in input_files:
            shutil.copy(Path.cwd() / f, ref_dir / "inputs")
        for f in output_files:
            shutil.copy(f, ref_dir / "outputs")

    def mock_get_input_set(self, *args, **kwargs):
        return get_input_set_orig(self, *args, **kwargs)

    if should_mock_aims:
        monkeypatch.setattr(atomate2_temp.aims.run, "run_aims", mock_run_aims)
        monkeypatch.setattr(atomate2_temp.aims.jobs.base, "run_aims", mock_run_aims)
        monkeypatch.setattr(AimsInputGenerator, "get_input_set", mock_get_input_set)
    else:
        monkeypatch.setattr(
            atomate2_temp.aims.jobs.base, "run_aims", generate_test_data
        )

    def _run(ref_paths, fake_run_aims_kwargs=None):
        if fake_run_aims_kwargs is None:
            fake_run_aims_kwargs = {}

        _REF_PATHS.update(ref_paths)
        _FAKE_RUN_AIMS_KWARGS.update(fake_run_aims_kwargs)

    yield _run

    monkeypatch.undo()
    _REF_PATHS.clear()
    _FAKE_RUN_AIMS_KWARGS.clear()


def fake_run_aims(
    ref_path: Union[str, Path],
    input_settings: Sequence[str] = (),
    check_inputs: Sequence[Literal["control.in"]] = _VFILES,
    clear_inputs: bool = False,
):
    """
    Emulate running aims and validate aims input files.

    Parameters
    ----------
    ref_path
        Path to reference directory with aims input files in the folder named 'inputs'
        and output files in the folder named 'outputs'.
    input_settings
        A list of input settings to check.
    check_inputs
        A list of aims input files to check. Supported options are "aims.inp"
    clear_inputs
        Whether to clear input files before copying in the reference aims outputs.
    """
    logger.info("Running fake aims.")

    ref_path = Path(ref_path)

    logger.info("Verified inputs successfully")

    if clear_inputs:
        clear_aims_inputs()

    copy_aims_outputs(ref_path)

    # pretend to run aims by copying pre-generated outputs from reference dir
    logger.info("Generated fake aims outputs")


def clear_aims_inputs():
    for aims_file in ("control.in", "geometry.in", "parameters.json"):
        if Path(aims_file).exists():
            Path(aims_file).unlink()
    logger.info("Cleared aims inputs")


def copy_aims_outputs(ref_path: Union[str, Path]):
    import shutil

    output_path = ref_path / "outputs"
    for output_file in output_path.iterdir():
        if output_file.is_file():
            shutil.copy(output_file, ".")

tests/conftest.py
tpurcell90/atomate2-fhi-aims
"""Tests for file manipulation"""

import pytest
from pathlib import Path

TEST_DIR = Path(__file__).parent


@pytest.fixture()
def tmp_dir():
    """Same as clean_dir but is fresh for every test"""
    import os
    import shutil
    import tempfile

    old_cwd = os.getcwd()
    newpath = tempfile.mkdtemp()
    os.chdir(newpath)
    yield
    os.chdir(old_cwd)
    shutil.rmtree(newpath)


def test_copy_aims_outputs(tmp_dir):
    from atomate2_temp.aims.files import copy_aims_outputs

    files = ["aims.out"]
    restart_files = ["geometry.in.next_step", "D_spin_01_kpt_000001.csc"]

    path = TEST_DIR / "outputs"
    copy_aims_outputs(src_dir=path)

    for f in files + restart_files:
        assert Path(f).exists()

tests/test_files/test_files.py
tpurcell90/atomate2-fhi-aims
"""Test core FHI-aims workflows"""
import pytest
import os

from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

cwd = os.getcwd()


def test_double_relax(mock_aims, tmp_path, Si, species_dir):
    """A test for the double relaxation flow"""

    from jobflow import run_locally

    from atomate2_temp.aims.flows.core import DoubleRelaxMaker
    from atomate2_temp.aims.schemas.task import AimsTaskDocument

    # mapping from job name to directory containing test files
    ref_paths = {
        "Relaxation calculation 1": "double-relax-si/relax-1",
        "Relaxation calculation 2": "double-relax-si/relax-2",
    }

    parameters = {"k_grid": [2, 2, 2], "species_dir": species_dir.as_posix()}

    # settings passed to fake_run_aims
    fake_run_kwargs = {}

    # automatically use fake AIMS
    mock_aims(ref_paths, fake_run_kwargs)

    # generate flow
    flow = DoubleRelaxMaker.from_parameters(parameters).make(MSONableAtoms(Si))

    # Run the flow or job and ensure that it finished running successfully
    os.chdir(tmp_path)
    responses = run_locally(flow, create_folders=True, ensure_success=True)
    os.chdir(cwd)

    # validate output
    output1 = responses[flow.jobs[0].uuid][1].output
    output2 = responses[flow.jobs[1].uuid][1].output

    assert isinstance(output1, AimsTaskDocument)
    assert output1.output.energy == pytest.approx(-15800.22554)
    assert output2.output.energy == pytest.approx(-15800.25855)

tests/test_flows/test_core.py
tpurcell90/atomate2-fhi-aims
"""Test various makers"""
import os

from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

cwd = os.getcwd()


def test_phonon_flow(Si, tmp_path, mock_aims, species_dir):
    import numpy as np

    from jobflow import run_locally

    from atomate2_temp.aims.jobs.core import StaticMaker, RelaxMaker
    from atomate2_temp.aims.jobs.phonons import PhononDisplacementMaker

    from atomate2_temp.aims.sets.core import StaticSetGenerator

    from atomate2_temp.aims.flows.phonons import PhononMaker

    # mapping from job name to directory containing test files
    ref_paths = {
        "Relaxation calculation": "phonon-relax-si",
        "phonon static aims 1/1": "phonon-disp-si",
        "SCF Calculation": "phonon-energy-si",
    }

    # settings passed to fake_run_aims; adjust these to check for certain input settings
    fake_run_aims_kwargs = {}

    # automatically use fake FHI-aims
    mock_aims(ref_paths, fake_run_aims_kwargs)

    parameters = {"k_grid": [2, 2, 2], "species_dir": species_dir.as_posix()}
    # generate job

    parameters_phonon_disp = dict(compute_forces=True, **parameters)
    maker = PhononMaker(
        bulk_relax_maker=RelaxMaker.full_relaxation(user_parameters=parameters),
        static_energy_maker=StaticMaker(
            input_set_generator=StaticSetGenerator(user_parameters=parameters)
        ),
        use_symmetrized_structure="primitive",
        phonon_displacement_maker=PhononDisplacementMaker(
            input_set_generator=StaticSetGenerator(
                user_parameters=parameters_phonon_disp,
                user_kpoints_settings={"density": 5.0, "even": True},
            )
        ),
    )
    maker.name = "phonons"
    flow = maker.make(
        MSONableAtoms(Si),
        supercell_matrix=np.array([-1, 1, 1, 1, -1, 1, 1, 1, -1]).reshape((3, 3)),
    )

    # run the flow or job and ensure that it finished running successfully
    os.chdir(tmp_path)
    responses = run_locally(flow, create_folders=True, ensure_success=True)
    os.chdir(cwd)

    # validation the outputs of the job
    output = responses[flow.job_uuids[-1]][1].output
    assert output.code == "aims"
    assert output.born is None
    assert not output.has_imaginary_modes

    assert output.temperatures == list(range(0, 500, 10))
    assert output.heat_capacities[0] == 0.0
    assert np.round(output.heat_capacities[-1], 2) == 23.06
    assert (
        output.phonopy_settings.schema_json()
        == '{"title": "PhononComputationalSettings", "description": "Collection to store computational settings for the phonon computation.", "type": "object", "properties": {"npoints_band": {"title": "Npoints Band", "default": "number of points for band structure computation", "type": "integer"}, "kpath_scheme": {"title": "Kpath Scheme", "default": "indicates the kpath scheme", "type": "string"}, "kpoint_density_dos": {"title": "Kpoint Density Dos", "default": "number of points for computation of free energies and densities of states", "type": "integer"}}}'
    )
    assert np.round(output.phonon_bandstructure.bands[-1, 0], 2) == 14.41


def test_phonon_socket_flow(Si, tmp_path, mock_aims, species_dir):
    import numpy as np

    from jobflow import run_locally

    from atomate2_temp.aims.jobs.core import StaticMaker, RelaxMaker
    from atomate2_temp.aims.jobs.phonons import PhononDisplacementMakerSocket
    from atomate2_temp.aims.sets.core import StaticSetGenerator

    from atomate2_temp.aims.flows.phonons import PhononMaker

    # mapping from job name to directory containing test files
    ref_paths = {
        "Relaxation calculation": "phonon-relax-si",
        "phonon static aims 1/1": "phonon-disp-si",
        "SCF Calculation": "phonon-energy-si",
    }

    # settings passed to fake_run_aims; adjust these to check for certain input settings
    fake_run_aims_kwargs = {}

    # automatically use fake FHI-aims
    mock_aims(ref_paths, fake_run_aims_kwargs)

    parameters = {"k_grid": [2, 2, 2], "species_dir": species_dir.as_posix()}
    parameters_phonon_disp = dict(
        compute_forces=True, use_pimd_wrapper=("localhost", 12345), **parameters
    )

    # generate job

    maker = PhononMaker(
        bulk_relax_maker=RelaxMaker.full_relaxation(user_parameters=parameters),
        static_energy_maker=StaticMaker(
            input_set_generator=StaticSetGenerator(user_parameters=parameters)
        ),
        use_symmetrized_structure="primitive",
        socket=True,
        phonon_displacement_maker=PhononDisplacementMakerSocket(
            input_set_generator=StaticSetGenerator(
                user_parameters=parameters_phonon_disp,
                user_kpoints_settings={"density": 5.0, "even": True},
            )
        ),
    )
    maker.name = "phonons"
    flow = maker.make(
        MSONableAtoms(Si),
        supercell_matrix=np.array([-1, 1, 1, 1, -1, 1, 1, 1, -1]).reshape((3, 3)),
    )

    # run the flow or job and ensure that it finished running successfully
    os.chdir(tmp_path)
    responses = run_locally(flow, create_folders=True, ensure_success=True)
    os.chdir(cwd)

    # validation the outputs of the job
    output = responses[flow.job_uuids[-1]][1].output
    assert output.code == "aims"
    assert output.born is None
    assert not output.has_imaginary_modes

    assert output.temperatures == list(range(0, 500, 10))
    assert output.heat_capacities[0] == 0.0
    assert np.round(output.heat_capacities[-1], 2) == 23.06
    assert (
        output.phonopy_settings.schema_json()
        == '{"title": "PhononComputationalSettings", "description": "Collection to store computational settings for the phonon computation.", "type": "object", "properties": {"npoints_band": {"title": "Npoints Band", "default": "number of points for band structure computation", "type": "integer"}, "kpath_scheme": {"title": "Kpath Scheme", "default": "indicates the kpath scheme", "type": "string"}, "kpoint_density_dos": {"title": "Kpoint Density Dos", "default": "number of points for computation of free energies and densities of states", "type": "integer"}}}'
    )
    assert np.round(output.phonon_bandstructure.bands[-1, 0], 2) == 14.41


def test_phonon_default_flow(Si, tmp_path, mock_aims, species_dir):
    import numpy as np

    from jobflow import run_locally

    from atomate2_temp.aims.jobs.core import StaticMaker, RelaxMaker
    from atomate2_temp.aims.jobs.phonons import PhononDisplacementMaker

    from atomate2_temp.aims.sets.core import StaticSetGenerator

    from atomate2_temp.aims.flows.phonons import PhononMaker

    # mapping from job name to directory containing test files
    ref_paths = {
        "Relaxation calculation": "phonon-relax-default-si",
        "phonon static aims 1/1": "phonon-disp-default-si",
        "SCF Calculation": "phonon-energy-default-si",
    }

    # settings passed to fake_run_aims; adjust these to check for certain input settings
    fake_run_aims_kwargs = {}

    # automatically use fake FHI-aims
    mock_aims(ref_paths, fake_run_aims_kwargs)

    aims_sd = os.environ.get("AIMS_SPECIES_DIR", None)
    os.environ["AIMS_SPECIES_DIR"] = str(species_dir)

    maker = PhononMaker()
    maker.name = "phonons"
    flow = maker.make(
        MSONableAtoms(Si),
        supercell_matrix=np.array([-1, 1, 1, 1, -1, 1, 1, 1, -1]).reshape((3, 3)),
    )

    # run the flow or job and ensure that it finished running successfully
    os.chdir(tmp_path)
    responses = run_locally(flow, create_folders=True, ensure_success=True)
    os.chdir(cwd)

    # validation the outputs of the job
    output = responses[flow.job_uuids[-1]][1].output
    assert output.code == "aims"
    assert output.born is None
    assert not output.has_imaginary_modes

    assert output.temperatures == list(range(0, 500, 10))
    assert output.heat_capacities[0] == 0.0
    assert np.round(output.heat_capacities[-1], 2) == 22.85
    assert (
        output.phonopy_settings.schema_json()
        == '{"title": "PhononComputationalSettings", "description": "Collection to store computational settings for the phonon computation.", "type": "object", "properties": {"npoints_band": {"title": "Npoints Band", "default": "number of points for band structure computation", "type": "integer"}, "kpath_scheme": {"title": "Kpath Scheme", "default": "indicates the kpath scheme", "type": "string"}, "kpoint_density_dos": {"title": "Kpoint Density Dos", "default": "number of points for computation of free energies and densities of states", "type": "integer"}}}'
    )
    assert np.round(output.phonon_bandstructure.bands[-1, 0], 2) == 15.02

    if aims_sd is not None:
        os.environ["AIMS_SPECIES_DIR"] = aims_sd


def test_phonon_default_socket_flow(Si, tmp_path, mock_aims, species_dir):
    import numpy as np

    from jobflow import run_locally

    from atomate2_temp.aims.jobs.core import StaticMaker, RelaxMaker
    from atomate2_temp.aims.jobs.phonons import PhononDisplacementMakerSocket
    from atomate2_temp.aims.sets.core import StaticSetGenerator

    from atomate2_temp.aims.flows.phonons import PhononMaker

    aims_sd = os.environ.get("AIMS_SPECIES_DIR", None)
    os.environ["AIMS_SPECIES_DIR"] = str(species_dir)

    # mapping from job name to directory containing test files
    ref_paths = {
        "Relaxation calculation": "phonon-relax-default-si",
        "phonon static aims 1/1": "phonon-disp-default-si",
        "SCF Calculation": "phonon-energy-default-si",
    }

    # settings passed to fake_run_aims; adjust these to check for certain input settings
    fake_run_aims_kwargs = {}

    # automatically use fake FHI-aims
    mock_aims(ref_paths, fake_run_aims_kwargs)

    # generate job

    maker = PhononMaker(socket=True)
    maker.name = "phonons"
    flow = maker.make(
        MSONableAtoms(Si),
        supercell_matrix=np.array([-1, 1, 1, 1, -1, 1, 1, 1, -1]).reshape((3, 3)),
    )

    # run the flow or job and ensure that it finished running successfully
    os.chdir(tmp_path)
    responses = run_locally(flow, create_folders=True, ensure_success=True)
    os.chdir(cwd)

    # validation the outputs of the job
    output = responses[flow.job_uuids[-1]][1].output
    assert output.code == "aims"
    assert output.born is None
    assert not output.has_imaginary_modes

    assert output.temperatures == list(range(0, 500, 10))
    assert output.heat_capacities[0] == 0.0
    assert np.round(output.heat_capacities[-1], 2) == 22.85
    assert (
        output.phonopy_settings.schema_json()
        == '{"title": "PhononComputationalSettings", "description": "Collection to store computational settings for the phonon computation.", "type": "object", "properties": {"npoints_band": {"title": "Npoints Band", "default": "number of points for band structure computation", "type": "integer"}, "kpath_scheme": {"title": "Kpath Scheme", "default": "indicates the kpath scheme", "type": "string"}, "kpoint_density_dos": {"title": "Kpoint Density Dos", "default": "number of points for computation of free energies and densities of states", "type": "integer"}}}'
    )
    assert np.round(output.phonon_bandstructure.bands[-1, 0], 2) == 15.02

    if aims_sd is not None:
        os.environ["AIMS_SPECIES_DIR"] = aims_sd

tests/test_flows/test_phonon_workflow.py
tpurcell90/atomate2-fhi-aims
"""Tests the band structure input set generator"""

from atomate2_temp.aims.sets.bs import BandStructureSetGenerator
from tests import compare_files


def comp_system(atoms, user_params, test_name, work_path, ref_path):
    k_point_density = user_params.pop("k_point_density", 20)
    generator = BandStructureSetGenerator(
        user_parameters=user_params, k_point_density=k_point_density
    )
    input_set = generator.get_input_set(atoms)
    input_set.write_input(work_path / test_name)
    compare_files(test_name, work_path, ref_path)


def test_si_bs(Si, species_dir, tmp_path, ref_path):
    parameters = {"species_dir": str(species_dir), "k_grid": [8, 8, 8]}
    comp_system(Si, parameters, "static-si-bs", tmp_path, ref_path)


def test_si_bs_output(Si, species_dir, tmp_path, ref_path):
    parameters = {
        "species_dir": str(species_dir),
        "k_grid": [8, 8, 8],
        "output": [
            "json_log",
        ],
    }
    comp_system(Si, parameters, "static-si-bs-output", tmp_path, ref_path)


def test_si_bs_density(Si, species_dir, tmp_path, ref_path):
    parameters = {
        "species_dir": str(species_dir),
        "k_grid": [8, 8, 8],
        "k_point_density": 40,
    }
    comp_system(Si, parameters, "static-si-bs-density", tmp_path, ref_path)

tests/test_input_generators/test_bs_generator.py
tpurcell90/atomate2-fhi-aims
"""Tests the GW input set generator"""
from pathlib import Path

from atomate2_temp.aims.sets.bs import GWSetGenerator
from tests import compare_files


def comp_system(atoms, user_params, test_name, work_path, ref_path):
    k_point_density = user_params.pop("k_point_density", 20)
    generator = GWSetGenerator(
        user_parameters=user_params, k_point_density=k_point_density
    )
    input_set = generator.get_input_set(atoms)
    input_set.write_input(work_path / test_name)
    compare_files(test_name, work_path, ref_path)


def test_si_gw(Si, species_dir, tmp_path, ref_path):
    parameters = {
        "species_dir": str(species_dir),
        "k_grid": [2, 2, 2],
        "k_point_density": 10,
    }
    comp_system(Si, parameters, "static-si-gw", tmp_path, ref_path)

tests/test_input_generators/test_gw_generator.py
tpurcell90/atomate2-fhi-aims
from atomate2_temp.aims.sets.core import RelaxSetGenerator

import os
from tests import compare_files
from pathlib import Path


def comp_system(atoms, user_params, test_name, work_path, ref_path):
    generator = RelaxSetGenerator(user_parameters=user_params)
    input_set = generator.get_input_set(atoms)
    input_set.write_input(work_path / test_name)
    compare_files(test_name, work_path, ref_path)


def test_relax_si(Si, species_dir, tmp_path, ref_path):
    parameters = {"species_dir": str(species_dir), "k_grid": [2, 2, 2]}
    comp_system(Si, parameters, "relax-si/inputs", tmp_path, ref_path)


def test_relax_si_no_kgrid(Si, species_dir, tmp_path, ref_path):
    parameters = {"species_dir": str(species_dir)}
    comp_system(Si, parameters, "relax-no-kgrid-si", tmp_path, ref_path)


def test_relax_default_species_dir(Si, species_dir, tmp_path, ref_path):
    sd_def = os.getenv("AIMS_SPECIES_DIR", None)
    os.environ["AIMS_SPECIES_DIR"] = str(species_dir)
    parameters = {"k_grid": [2, 2, 2]}

    comp_system(Si, parameters, "relax-default-sd-si", tmp_path, ref_path)

    if sd_def:
        os.environ["AIMS_SPECIES_DIR"] = sd_def
    else:
        os.unsetenv("AIMS_SPECIES_DIR")


def test_relax_o2(O2, species_dir, tmp_path, ref_path):
    parameters = {"species_dir": str(species_dir)}
    comp_system(O2, parameters, "relax-o2", tmp_path, ref_path)


def test_relax_default_species_dir_o2(O2, species_dir, tmp_path, ref_path):
    sd_def = os.getenv("AIMS_SPECIES_DIR", None)
    os.environ["AIMS_SPECIES_DIR"] = str(species_dir)
    parameters = {"k_grid": [2, 2, 2]}

    comp_system(O2, parameters, "relax-default-sd-o2", tmp_path, ref_path)

    if sd_def:
        os.environ["AIMS_SPECIES_DIR"] = sd_def
    else:
        os.unsetenv("AIMS_SPECIES_DIR")

tests/test_input_generators/test_relax_generator.py
tpurcell90/atomate2-fhi-aims
from atomate2_temp.aims.sets.core import StaticSetGenerator
from tests import compare_files

import os


def comp_system(atoms, user_params, test_name, work_path, ref_path):
    generator = StaticSetGenerator(user_parameters=user_params)
    input_set = generator.get_input_set(atoms)
    input_set.write_input(work_path / test_name)
    compare_files(test_name, work_path, ref_path)


def test_static_si(Si, species_dir, tmp_path, ref_path):
    parameters = {"species_dir": str(species_dir), "k_grid": [2, 2, 2]}
    comp_system(Si, parameters, "static-si", tmp_path, ref_path)


def test_static_si_no_kgrid(Si, species_dir, tmp_path, ref_path):
    parameters = {"species_dir": str(species_dir)}
    comp_system(Si, parameters, "static-no-kgrid-si", tmp_path, ref_path)


def test_static_default_species_dir(Si, species_dir, tmp_path, ref_path):
    sd_def = os.getenv("AIMS_SPECIES_DIR", None)
    os.environ["AIMS_SPECIES_DIR"] = str(species_dir)
    parameters = {"k_grid": [2, 2, 2]}

    comp_system(Si, parameters, "static-default-sd-si", tmp_path, ref_path)

    if sd_def:
        os.environ["AIMS_SPECIES_DIR"] = sd_def
    else:
        os.unsetenv("AIMS_SPECIES_DIR")


def test_static_o2(O2, species_dir, tmp_path, ref_path):
    parameters = {"species_dir": str(species_dir)}
    comp_system(O2, parameters, "static-o2", tmp_path, ref_path)


def test_static_default_species_dir_o2(O2, species_dir, tmp_path, ref_path):
    sd_def = os.getenv("AIMS_SPECIES_DIR", None)
    os.environ["AIMS_SPECIES_DIR"] = str(species_dir)
    parameters = {"k_grid": [2, 2, 2]}

    comp_system(O2, parameters, "static-default-sd-o2", tmp_path, ref_path)

    if sd_def:
        os.environ["AIMS_SPECIES_DIR"] = sd_def
    else:
        os.unsetenv("AIMS_SPECIES_DIR")

tests/test_input_generators/test_static_generator.py
tpurcell90/atomate2-fhi-aims
"""The test of input sets generating from restart information"""

import json
import os
import shutil
from functools import wraps
from pathlib import Path

import pytest

from atomate2_temp.aims.sets.core import StaticSetGenerator
from tests import compare_files


def comp_system(atoms, prev_dir, test_name, work_path, ref_path, species_dir):
    generator = StaticSetGenerator(user_parameters={})
    # adjust species dir in the prev_dir
    params_file = Path(prev_dir) / "parameters.json"
    shutil.copy(params_file, Path(prev_dir) / "~parameters.json")
    with open(params_file) as f:
        params = json.load(f)
    params["species_dir"] = species_dir.as_posix()
    with open(params_file, "w") as f:
        json.dump(params, f)

    input_set = generator.get_input_set(
        atoms, prev_dir, properties=["energy", "forces", "stress"]
    )
    input_set.write_input(work_path / test_name)
    compare_files(test_name, work_path, ref_path)
    shutil.move(Path(prev_dir) / "~parameters.json", params_file)


def test_static_from_relax_si(Si, species_dir, tmp_path, ref_path):
    comp_system(
        Si,
        f"{ref_path}/relax-si/outputs",
        "static-from-prev-si",
        tmp_path,
        ref_path,
        species_dir,
    )


def test_static_from_relax_si_no_kgrid(Si, species_dir, tmp_path, ref_path):
    comp_system(
        Si,
        f"{ref_path}/relax-no-kgrid-si/",
        "static-from-prev-no-kgrid-si",
        tmp_path,
        ref_path,
        species_dir,
    )


def test_static_from_relax_default_species_dir(Si, species_dir, tmp_path, ref_path):
    sd_def = os.getenv("AIMS_SPECIES_DIR", None)
    os.environ["AIMS_SPECIES_DIR"] = str(species_dir)

    comp_system(
        Si,
        f"{ref_path}/relax-default-sd-si/",
        "static-from-prev-default-sd-si",
        tmp_path,
        ref_path,
        species_dir,
    )

    if sd_def:
        os.environ["AIMS_SPECIES_DIR"] = sd_def
    else:
        os.unsetenv("AIMS_SPECIES_DIR")


def test_static_from_relax_o2(O2, species_dir, tmp_path, ref_path):
    comp_system(
        O2,
        f"{ref_path}/relax-o2/",
        "static-from-prev-o2",
        tmp_path,
        ref_path,
        species_dir,
    )


def test_static_from_relax_default_species_dir_o2(O2, species_dir, tmp_path, ref_path):
    sd_def = os.getenv("AIMS_SPECIES_DIR", None)
    os.environ["AIMS_SPECIES_DIR"] = str(species_dir)

    comp_system(
        O2,
        f"{ref_path}/relax-default-sd-o2/",
        "static-from-prev-default-sd-o2",
        tmp_path,
        ref_path,
        species_dir,
    )

    if sd_def:
        os.environ["AIMS_SPECIES_DIR"] = sd_def
    else:
        os.unsetenv("AIMS_SPECIES_DIR")

tests/test_input_generators/test_static_restart_from_relax_generator.py
tpurcell90/atomate2-fhi-aims
from atomate2_temp.aims.sets.base import AimsInputFile

def test_input_file():
    input_file = AimsInputFile.from_str("testing")
    assert input_file.get_str() == "testing"

    input_file = AimsInputFile.from_string("testing")
    assert input_file.get_string() == "testing"

tests/test_input_set/test_input_file.py
tpurcell90/atomate2-fhi-aims
from atomate2_temp.aims.sets.base import AimsInputSet
import pytest
import numpy as np

control_in_str = '#===============================================================================\n# Created using the Atomic Simulation Environment (ASE)\n\n# Thu Oct  5 12:27:49 2023\n\n#===============================================================================\nxc                                 pbe\nk_grid                             2 2 2\ncompute_forces                     .true.\n#===============================================================================\n\n################################################################################\n#\n#  FHI-aims code project\n#  VB, Fritz-Haber Institut, 2009\n#\n#  Suggested "light" defaults for Si atom (to be pasted into control.in file)\n#  Be sure to double-check any results obtained with these settings for post-processing,\n#  e.g., with the "tight" defaults and larger basis sets.\n#\n#  2020/09/08 Added f function to "light" after reinspection of Delta test outcomes.\n#             This was done for all of Al-Cl and is a tricky decision since it makes\n#             "light" calculations measurably more expensive for these elements.\n#             Nevertheless, outcomes for P, S, Cl (and to some extent, Si) appear\n#             to justify this choice.\n#\n################################################################################\n  species        Si\n#     global species definitions\n    nucleus             14\n    mass                28.0855\n#\n    l_hartree           4\n#\n    cut_pot             3.5          1.5  1.0\n    basis_dep_cutoff    1e-4\n#\n    radial_base         42 5.0\n    radial_multiplier   1\n    angular_grids       specified\n      division   0.5866   50\n      division   0.9616  110\n      division   1.2249  194\n      division   1.3795  302\n#      division   1.4810  434\n#      division   1.5529  590\n#      division   1.6284  770\n#      division   1.7077  974\n#      division   2.4068 1202\n#      outer_grid   974\n      outer_grid 302\n################################################################################\n#\n#  Definition of "minimal" basis\n#\n################################################################################\n#     valence basis states\n    valence      3  s   2.\n    valence      3  p   2.\n#     ion occupancy\n    ion_occ      3  s   1.\n    ion_occ      3  p   1.\n################################################################################\n#\n#  Suggested additional basis functions. For production calculations, \n#  uncomment them one after another (the most important basis functions are\n#  listed first).\n#\n#  Constructed for dimers: 1.75 A, 2.0 A, 2.25 A, 2.75 A, 3.75 A\n#\n################################################################################\n#  "First tier" - improvements: -571.96 meV to -37.03 meV\n     hydro 3 d 4.2\n     hydro 2 p 1.4\n     hydro 4 f 6.2\n     ionic 3 s auto\n#  "Second tier" - improvements: -16.76 meV to -3.03 meV\n#     hydro 3 d 9\n#     hydro 5 g 9.4\n#     hydro 4 p 4\n#     hydro 1 s 0.65\n#  "Third tier" - improvements: -3.89 meV to -0.60 meV\n#     ionic 3 d auto\n#     hydro 3 s 2.6\n#     hydro 4 f 8.4\n#     hydro 3 d 3.4\n#     hydro 3 p 7.8\n#  "Fourth tier" - improvements: -0.33 meV to -0.11 meV\n#     hydro 2 p 1.6\n#     hydro 5 g 10.8\n#     hydro 5 f 11.2\n#     hydro 3 d 1\n#     hydro 4 s 4.5\n#  Further basis functions that fell out of the optimization - noise\n#  level... < -0.08 meV\n#     hydro 4 d 6.6\n#     hydro 5 g 16.4\n#     hydro 4 d 9\n'
control_in_str_rel = '#===============================================================================\n# Created using the Atomic Simulation Environment (ASE)\n\n# Thu Oct  5 12:33:50 2023\n\n#===============================================================================\nxc                                 pbe\nk_grid                             2 2 2\nrelax_geometry                     trm 1e-3\ncompute_forces                     .true.\n#===============================================================================\n\n################################################################################\n#\n#  FHI-aims code project\n#  VB, Fritz-Haber Institut, 2009\n#\n#  Suggested "light" defaults for Si atom (to be pasted into control.in file)\n#  Be sure to double-check any results obtained with these settings for post-processing,\n#  e.g., with the "tight" defaults and larger basis sets.\n#\n#  2020/09/08 Added f function to "light" after reinspection of Delta test outcomes.\n#             This was done for all of Al-Cl and is a tricky decision since it makes\n#             "light" calculations measurably more expensive for these elements.\n#             Nevertheless, outcomes for P, S, Cl (and to some extent, Si) appear\n#             to justify this choice.\n#\n################################################################################\n  species        Si\n#     global species definitions\n    nucleus             14\n    mass                28.0855\n#\n    l_hartree           4\n#\n    cut_pot             3.5          1.5  1.0\n    basis_dep_cutoff    1e-4\n#\n    radial_base         42 5.0\n    radial_multiplier   1\n    angular_grids       specified\n      division   0.5866   50\n      division   0.9616  110\n      division   1.2249  194\n      division   1.3795  302\n#      division   1.4810  434\n#      division   1.5529  590\n#      division   1.6284  770\n#      division   1.7077  974\n#      division   2.4068 1202\n#      outer_grid   974\n      outer_grid 302\n################################################################################\n#\n#  Definition of "minimal" basis\n#\n################################################################################\n#     valence basis states\n    valence      3  s   2.\n    valence      3  p   2.\n#     ion occupancy\n    ion_occ      3  s   1.\n    ion_occ      3  p   1.\n################################################################################\n#\n#  Suggested additional basis functions. For production calculations, \n#  uncomment them one after another (the most important basis functions are\n#  listed first).\n#\n#  Constructed for dimers: 1.75 A, 2.0 A, 2.25 A, 2.75 A, 3.75 A\n#\n################################################################################\n#  "First tier" - improvements: -571.96 meV to -37.03 meV\n     hydro 3 d 4.2\n     hydro 2 p 1.4\n     hydro 4 f 6.2\n     ionic 3 s auto\n#  "Second tier" - improvements: -16.76 meV to -3.03 meV\n#     hydro 3 d 9\n#     hydro 5 g 9.4\n#     hydro 4 p 4\n#     hydro 1 s 0.65\n#  "Third tier" - improvements: -3.89 meV to -0.60 meV\n#     ionic 3 d auto\n#     hydro 3 s 2.6\n#     hydro 4 f 8.4\n#     hydro 3 d 3.4\n#     hydro 3 p 7.8\n#  "Fourth tier" - improvements: -0.33 meV to -0.11 meV\n#     hydro 2 p 1.6\n#     hydro 5 g 10.8\n#     hydro 5 f 11.2\n#     hydro 3 d 1\n#     hydro 4 s 4.5\n#  Further basis functions that fell out of the optimization - noise\n#  level... < -0.08 meV\n#     hydro 4 d 6.6\n#     hydro 5 g 16.4\n#     hydro 4 d 9\n'
geometry_in_str = "#===============================================================================\n# Created using the Atomic Simulation Environment (ASE)\n\n# Thu Oct  5 12:27:49 2023\n\n#=======================================================\nlattice_vector 0.0000000000000000 2.7149999999999999 2.7149999999999999 \nlattice_vector 2.7149999999999999 0.0000000000000000 2.7149999999999999 \nlattice_vector 2.7149999999999999 2.7149999999999999 0.0000000000000000 \natom_frac 0.0000000000000000 0.0000000000000000 -0.0000000000000000 Si\natom_frac 0.2500000000000000 0.2500000000000000 0.2500000000000000 Si\n"
geometry_in_str_new = "#===============================================================================\n# Created using the Atomic Simulation Environment (ASE)\n\n# Thu Oct  5 12:27:49 2023\n\n#=======================================================\nlattice_vector 0.0000000000000000 2.7149999999999999 2.7149999999999999 \nlattice_vector 2.7149999999999999 0.0000000000000000 2.7149999999999999 \nlattice_vector 2.7149999999999999 2.7149999999999999 0.0000000000000000 \natom_frac -0.0100000000000000 0.0000000000000000 -0.0000000000000000 Si\natom_frac 0.2500000000000000 0.2500000000000000 0.2500000000000000 Si\n"


def test_input_set(Si, species_dir):
    parameters_json_str = (
        "{"
        + f'\n  "xc": "pbe",\n  "species_dir": "{species_dir}",\n  "k_grid": [\n    2,\n    2,\n    2\n  ]\n'
        + "}"
    )
    parameters_json_str_rel = (
        "{"
        + f'\n  "xc": "pbe",\n  "species_dir": "{species_dir}",\n  "k_grid": [\n    2,\n    2,\n    2\n  ],\n  "relax_geometry": "trm 1e-3"\n'
        + "}"
    )

    parameters = {"xc": "pbe", "species_dir": str(species_dir), "k_grid": [2, 2, 2]}
    properties = ("energy", "free_energy", "forces")

    in_set = AimsInputSet(parameters, Si, properties)
    assert geometry_in_str[175:] == in_set.geometry_in.get_str()[175:]
    assert control_in_str[175:] == in_set.control_in.get_str()[175:]
    print(parameters_json_str, "\n\n\n", in_set.parameters_json)
    assert parameters_json_str == in_set.parameters_json

    in_set_copy = in_set.deepcopy()
    assert geometry_in_str[175:] == in_set_copy.geometry_in.get_str()[175:]
    assert control_in_str[175:] == in_set_copy.control_in.get_str()[175:]
    assert parameters_json_str == in_set_copy.parameters_json

    in_set.set_parameters(**parameters, relax_geometry="trm 1e-3")
    assert control_in_str_rel[175:] == in_set.control_in.get_str()[175:]
    assert parameters_json_str_rel == in_set.parameters_json
    assert control_in_str[175:] == in_set_copy.control_in.get_str()[175:]
    assert parameters_json_str == in_set_copy.parameters_json

    in_set.remove_parameters(keys=["relax_geometry"])
    assert control_in_str[175:] == in_set.control_in.get_str()[175:]
    assert parameters_json_str == in_set.parameters_json

    in_set.remove_parameters(keys=["relax_geometry"], strict=False)
    assert control_in_str[175:] == in_set.control_in.get_str()[175:]
    assert parameters_json_str == in_set.parameters_json

    with pytest.raises(ValueError):
        in_set.remove_parameters(keys=["relax_geometry"], strict=True)

    new_atoms = Si.copy()
    new_atoms.set_scaled_positions(np.array([[-0.01, 0, 0], [0.25, 0.25, 0.25]]))
    in_set.set_atoms(new_atoms)
    assert geometry_in_str_new[175:] == in_set.geometry_in.get_str()[175:]
    assert geometry_in_str[175:] == in_set_copy.geometry_in.get_str()[175:]

tests/test_input_set/test_input_set.py
tpurcell90/atomate2-fhi-aims
from atomate2_temp.aims.io.aims_output import AimsOutput
from pathlib import Path

from monty.json import MontyDecoder
import json

outfile_dir = Path(__file__).parent / "aims_output_files"


def test_output_si():
    si = AimsOutput.from_outfile(f"{outfile_dir}/si.out")
    si_ref = json.load(open(f"{outfile_dir}/si_out.json", "rb"), cls=MontyDecoder)

    assert si_ref.metadata == si.metadata
    assert si_ref.atoms_summary == si.atoms_summary

    assert si_ref.n_images == si.n_images
    for ii in range(si.n_images):
        assert si_ref.get_results_for_image(ii) == si.get_results_for_image(ii)

    assert si.as_dict() == si_ref.as_dict()


def test_output_h2o():
    h2o = AimsOutput.from_outfile(f"{outfile_dir}/h2o.out")
    h2o_ref = json.load(open(f"{outfile_dir}/h2o_out.json", "rb"), cls=MontyDecoder)

    assert h2o_ref.metadata == h2o.metadata
    assert h2o_ref.atoms_summary == h2o.atoms_summary

    assert h2o_ref.n_images == h2o.n_images
    for ii in range(h2o.n_images):
        assert h2o_ref.get_results_for_image(ii) == h2o.get_results_for_image(ii)

    assert h2o.as_dict() == h2o_ref.as_dict()

tests/test_io/test_output.py
tpurcell90/atomate2-fhi-aims
""" A test for AIMS convergence maker (used for GW, for instance)
"""

import pytest
import os

from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

cwd = os.getcwd()


def test_convergence(mock_aims, tmp_path, Si, species_dir):
    """A test for the convergence job"""

    from jobflow import run_locally

    from atomate2_temp.aims.jobs.core import StaticMaker, StaticSetGenerator
    from atomate2_temp.aims.jobs.base import ConvergenceMaker
    from atomate2_temp.aims.schemas.task import ConvergenceSummary

    # mapping from job name to directory containing test files
    ref_paths = {
        "SCF Calculation 0": "k-grid-convergence-si/static-1",
        "SCF Calculation 1": "k-grid-convergence-si/static-2",
        "SCF Calculation 2": "k-grid-convergence-si/static-3",
    }

    input_set_parameters = {"species_dir": species_dir.as_posix()}

    parameters = {
        "maker": StaticMaker(
            input_set_generator=StaticSetGenerator(user_parameters=input_set_parameters)
        ),
        "criterion_name": "energy_per_atom",
        "epsilon": 0.2,
        "convergence_field": "k_grid",
        "convergence_steps": [[3, 3, 3], [4, 4, 4], [5, 5, 5], [6, 6, 6]],
    }

    # settings passed to fake_run_aims
    fake_run_kwargs = {}

    # automatically use fake AIMS
    mock_aims(ref_paths, fake_run_kwargs)

    # generate job
    flow = ConvergenceMaker(**parameters).make(MSONableAtoms(Si))

    # Run the job and ensure that it finished running successfully
    os.chdir(tmp_path)
    responses = run_locally(flow, create_folders=True, ensure_success=True)
    os.chdir(cwd)

    # a very nasty hack!
    # but otherwise I do not know how to get the uuid of the last job in a
    # dynamic workflow

    job_uuid = flow.all_uuids[0]
    while responses[job_uuid][1].detour:
        job_uuid = responses[job_uuid][1].detour.all_uuids[-1]

    output = responses[job_uuid][1].output

    # validate output
    assert isinstance(output, ConvergenceSummary)
    assert output.converged
    assert output.convergence_field_value == [5, 5, 5]
    assert output.actual_epsilon == pytest.approx(0.0614287)

tests/test_makers/test_convergence.py
tpurcell90/atomate2-fhi-aims
import os
import pytest

from jobflow import run_locally

from atomate2_temp.aims.jobs.core import RelaxMaker
from atomate2_temp.aims.schemas.task import AimsTaskDocument
from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms


cwd = os.getcwd()


def test_base_maker(tmp_path, species_dir, mock_aims, Si):
    # mapping from job name to directory containing test files
    ref_paths = {"relax_si": "relax-si"}

    # settings passed to fake_run_aims; adjust these to check for certain input settings
    fake_run_aims_kwargs = {}

    # automatically use fake FHI-aims
    mock_aims(ref_paths, fake_run_aims_kwargs)

    parameters = {"k_grid": [2, 2, 2], "species_dir": species_dir.as_posix()}
    # generate job
    maker = RelaxMaker.full_relaxation(user_parameters=parameters)
    maker.name = "relax_si"
    job = maker.make(MSONableAtoms(Si))

    # run the flow or job and ensure that it finished running successfully
    os.chdir(tmp_path)
    responses = run_locally(job, create_folders=True, ensure_success=True)
    os.chdir(cwd)

    # validation the outputs of the job
    output1 = responses[job.uuid][1].output
    assert isinstance(output1, AimsTaskDocument)
    assert output1.output.energy == pytest.approx(-15800.2255448846)


def test_relax_fixed_cell_maker(tmp_path, species_dir, mock_aims, Si):
    # mapping from job name to directory containing test files
    ref_paths = {"relax_fixed_cell_si": "relax-fixed-cell-si"}

    # settings passed to fake_run_aims; adjust these to check for certain input settings
    fake_run_aims_kwargs = {}

    # automatically use fake FHI-aims
    mock_aims(ref_paths, fake_run_aims_kwargs)

    parameters = {"k_grid": [2, 2, 2], "species_dir": species_dir.as_posix()}
    # generate job
    maker = RelaxMaker.fixed_cell_relaxation(user_parameters=parameters)
    maker.name = "relax_fixed_cell_si"
    atoms = Si.copy()
    atoms.positions[0, 0] += 0.25
    job = maker.make(MSONableAtoms(atoms))

    # run the flow or job and ensure that it finished running successfully
    os.chdir(tmp_path)
    responses = run_locally(job, create_folders=True, ensure_success=True)
    os.chdir(cwd)

    # validation the outputs of the job
    output1 = responses[job.uuid][1].output
    assert isinstance(output1, AimsTaskDocument)
    assert output1.output.energy == pytest.approx(-15800.099741042)

tests/test_makers/test_relax.py
tpurcell90/atomate2-fhi-aims
import pytest

from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms
import os


cwd = os.getcwd()


def test_static_socket_maker(Si, species_dir, mock_aims, tmp_path):
    from jobflow import run_locally

    from atomate2_temp.aims.jobs.core import SocketIOStaticMaker
    from atomate2_temp.aims.schemas.task import AimsTaskDocument
    from atomate2_temp.aims.sets.core import SocketIOSetGenerator

    atoms = MSONableAtoms(Si)
    atoms_list = [atoms, atoms.copy(), atoms.copy()]
    atoms_list[1].positions[0, 0] += 0.02
    atoms_list[2].cell[:, :] *= 1.02

    # mapping from job name to directory containing test files
    ref_paths = {"socket": "socket_tests"}

    # settings passed to fake_run_aims; adjust these to check for certain input settings
    fake_run_aims_kwargs = {}

    # automatically use fake FHI-aims
    mock_aims(ref_paths, fake_run_aims_kwargs)

    parameters = {"k_grid": [2, 2, 2], "species_dir": species_dir.as_posix()}
    # generate job
    maker = SocketIOStaticMaker(
        input_set_generator=SocketIOSetGenerator(user_parameters=parameters)
    )
    maker.name = "socket"
    job = maker.make(atoms_list)

    # run the flow or job and ensure that it finished running successfully
    os.chdir(tmp_path)
    responses = run_locally(job, create_folders=True, ensure_success=True)
    os.chdir(cwd)

    # validation the outputs of the job
    outputs = responses[job.uuid][1].output
    assert isinstance(outputs, AimsTaskDocument)
    assert len(outputs.output.trajectory) == 3
    assert outputs.output.trajectory[0].get_potential_energy() == pytest.approx(
        -15800.0997410132
    )
    assert outputs.output.trajectory[1].get_potential_energy() == pytest.approx(
        -15800.0962356206
    )
    assert outputs.output.trajectory[2].get_potential_energy() == pytest.approx(
        -15800.1847237514
    )
    # assert output1.output.energy == pytest.approx(-15800.099740991)

tests/test_makers/test_socket_calc.py
tpurcell90/atomate2-fhi-aims
"""Test various makers"""
import os

import pytest

from atomate2_temp.aims.utils.msonable_atoms import MSONableAtoms

cwd = os.getcwd()


def test_static_maker(Si, tmp_path, mock_aims, species_dir):
    from jobflow import run_locally

    from atomate2_temp.aims.jobs.core import StaticMaker
    from atomate2_temp.aims.schemas.task import AimsTaskDocument
    from atomate2_temp.aims.sets.core import StaticSetGenerator

    # mapping from job name to directory containing test files
    ref_paths = {"base": "static-si"}

    # settings passed to fake_run_aims; adjust these to check for certain input settings
    fake_run_aims_kwargs = {}

    # automatically use fake FHI-aims
    mock_aims(ref_paths, fake_run_aims_kwargs)

    parameters = {"k_grid": [2, 2, 2], "species_dir": species_dir.as_posix()}
    # generate job
    maker = StaticMaker(
        input_set_generator=StaticSetGenerator(user_parameters=parameters)
    )
    maker.name = "base"
    job = maker.make(MSONableAtoms(Si))

    # run the flow or job and ensure that it finished running successfully
    os.chdir(tmp_path)
    responses = run_locally(job, create_folders=True, ensure_success=True)
    os.chdir(cwd)

    # validation the outputs of the job
    output1 = responses[job.uuid][1].output
    assert isinstance(output1, AimsTaskDocument)
    assert output1.output.energy == pytest.approx(-15800.099740991)

tests/test_makers/test_static.py
tpurcell90/atomate2-fhi-aims
# flake8: noqa
import numpy as np
import pytest

from ase.io import read, ParseError
from ase.stress import full_3x3_to_voigt_6_stress
from numpy.linalg import norm
from pathlib import Path

parent = Path(__file__).parents[1]


def test_parse_socketio(tmpdir):
    traj = read(parent / "test_data/aims_out/socket.out", ":", format="aims-output")
    assert len(traj) == 6
    p0 = [[0.0, 0.0, 0.0], [0.9584, 0.0, 0.0], [-0.24, 0.9279, 0.0]]
    p1 = [
        [-0.00044436, 0.00021651, 0.00068957],
        [0.96112981, -0.00029923, 0.00096836],
        [-0.24091781, 0.93010946, 0.00061317],
    ]
    p_end = [
        [-0.00156048, -0.00072446, 0.00045281],
        [0.98615072, -0.00962614, -0.00053732],
        [-0.25646779, 0.95117586, 0.00820183],
    ]
    assert np.allclose(traj[0].get_positions(), p0)
    assert np.allclose(traj[1].get_positions(), p1)
    assert np.allclose(traj[-1].get_positions(), p_end)

    f0 = [
        [-0.481289284665163e00, -0.615051370384412e00, 0.811297123282653e-27],
        [0.762033585727896e00, -0.942008578636939e-01, -0.973556547939183e-27],
        [-0.280744301062733e00, 0.709252228248106e00, -0.649037698626122e-27],
    ]
    f1 = [
        [-0.346210275412861e00, -0.520615919604426e00, -0.966369462150621e-04],
        [0.587866333819113e00, -0.830442530429637e-01, 0.171037714240380e-03],
        [-0.241656058406252e00, 0.603660172647390e00, -0.744007680253175e-04],
    ]
    f_end = [
        [0.492882061544499e00, 0.499117230159087e00, 0.347959116743205e-02],
        [-0.724281788245024e00, 0.800633239635954e-01, 0.130633777464187e-02],
        [0.231399726700525e00, -0.579180554122683e00, -0.478592894207392e-02],
    ]
    assert np.allclose(traj[0].get_forces(), f0)
    assert np.allclose(traj[1].get_forces(), f1)
    assert np.allclose(traj[-1].get_forces(), f_end)


def test_parse_md(tmpdir):
    traj = read(parent / "test_data/aims_out/md.out", ":", format="aims-output")
    assert len(traj) == 5
    p0 = [[0.0, 0.0, 0.0], [0.9584, 0.0, 0.0], [-0.24, 0.9279, 0.0]]
    p1 = [
        [0.00247722, -0.00200215, 0.00000000],
        [0.93156204, -0.00330135, 0.00000000],
        [-0.25248383, 0.96298223, 0.00000000],
    ]
    p_end = [
        [-0.00044308, -0.00190646, 0.00000000],
        [0.98936333, -0.01341746, -0.00000000],
        [-0.26393022, 0.97157934, 0.00000000],
    ]
    assert np.allclose(traj[0].get_positions(), p0)
    assert np.allclose(traj[1].get_positions(), p1)
    assert np.allclose(traj[-1].get_positions(), p_end)

    f0 = [
        [-0.481289284665163e00, -0.615051370384412e00, 0.811297123282653e-27],
        [0.762033585727896e00, -0.942008578636939e-01, -0.973556547939183e-27],
        [-0.280744301062733e00, 0.709252228248106e00, -0.649037698626122e-27],
    ]
    f1 = [
        [-0.284519402890037e01, 0.121286349030924e01, 0.691733365155783e-17],
        [0.257911758656866e01, -0.471469245294899e-01, -0.730166238143266e-18],
        [0.266076442331705e00, -0.116571656577975e01, -0.618716741081841e-17],
    ]
    f_end = [
        [0.266848800470869e00, 0.137113486710510e01, 0.717235335588107e-17],
        [-0.812308728045051e00, 0.142880785873554e00, 0.874856850626564e-18],
        [0.545459927574182e00, -0.151401565297866e01, -0.804721020748119e-17],
    ]
    assert np.allclose(traj[0].get_forces(), f0)
    assert np.allclose(traj[1].get_forces(), f1)
    assert np.allclose(traj[-1].get_forces(), f_end)


def test_parse_relax(tmpdir):
    traj = read(parent / "test_data/aims_out/relax.out", ":", format="aims-output")
    assert len(traj) == 8
    p0 = [[0.0, 0.0, 0.0], [0.25, 0.25, 0.25]]
    assert all([np.allclose(at.get_scaled_positions(), p0) for at in traj])
    assert all([np.allclose(at.get_forces(), np.zeros((2, 3))) for at in traj])

    s0 = full_3x3_to_voigt_6_stress(
        [
            [0.07748922, 0.0, 0.0],
            [0.0, 0.07748923, 0.0],
            [0.0, 0.0, 0.07748923],
        ]
    )
    s_end = full_3x3_to_voigt_6_stress(
        [
            [-0.00007093, 0.00001488, 0.00001488],
            [0.00001488, -0.00007093, 0.00001488],
            [0.00001488, 0.00001488, -0.00007093],
        ]
    )
    assert np.allclose(traj[0].get_stress(), s0)
    assert np.allclose(traj[-1].get_stress(), s_end)

    cell_0 = [
        [0.000, 2.903, 2.903],
        [2.903, 0.000, 2.903],
        [2.903, 2.903, 0.000],
    ]
    cell_end = [
        [0.00006390, 2.70950938, 2.70950938],
        [2.70950941, 0.00006393, 2.70950939],
        [2.70950941, 2.70950939, 0.00006393],
    ]
    assert np.allclose(traj[0].get_cell(), cell_0)
    assert np.allclose(traj[-1].get_cell(), cell_end)


def test_parse_singlepoint(tmpdir):
    atoms = read(parent / "test_data/aims_out/singlepoint.out", format="aims-output")
    p0 = [[0.0, 0.0, 0.0], [0.9584, 0.0, 0.0], [-0.24, 0.9279, 0.0]]
    assert np.allclose(atoms.get_positions(), p0)

    f0 = [
        [-0.478299830661005e01, -0.616960341437854e01, 0.162259424656531e-27],
        [0.692109878170016e01, -0.419925410034428e00, 0.000000000000000e00],
        [-0.213810047509011e01, 0.658952882441297e01, 0.243389136984796e-27],
    ]
    assert np.allclose(atoms.get_forces(), f0)


def test_parse_dfpt_dielectric(tmpdir):
    outfile = parent / "test_data/aims_out/DFPT_dielectric.out"
    atoms = read(outfile, format="aims-output")

    diel = atoms.calc.results["dielectric_tensor"]

    diel_0 = [
        [7.18759265e00, -1.0000000e-15, 1.9000000e-14],
        [-1.000000e-15, 7.18759284e00, 2.59000000e-13],
        [2.0000000e-14, 2.58000000e-13, 7.1875928e00],
    ]

    assert np.allclose(diel, diel_0)

def test_parse_polarization(tmpdir):
    outfile = parent / "test_data/aims_out/polarization.out"
    atoms = read(outfile, format="aims-output")

    polar = atoms.calc.results["polarization"]

    polar_0 = [-51.045557E-03, -51.045557E-03, -51.458008E-03]

    assert np.allclose(polar, polar_0)

def test_preamble_failed(tmpdir):
    outfile = parent / "test_data/aims_out/preamble_fail.out"
    with pytest.raises(ParseError, match='No SCF steps'):
        read(outfile, format="aims-output")

def test_numerical_stress(tmpdir):
    outfile = parent / "test_data/aims_out/numerical_stress.out"

    atoms = read(outfile, format="aims-output")
    stress = atoms.get_stress()
    stress_actual = [
        0.00244726, 0.00267442, 0.00258710, 0.00000005, -0.00000026, -0.00000007
    ]

    assert np.allclose(stress, stress_actual)

tests/test_parsers/test_aims_out.py
tpurcell90/atomate2-fhi-aims
# flake8: noqa
import numpy as np
from ase.io import read
from atomate2_temp.aims.io.parsers import (
    AimsParseError,
    AimsOutChunk,
    AimsOutHeaderChunk,
    AimsOutCalcChunk,
    LINE_NOT_FOUND,
)
from ase.stress import full_3x3_to_voigt_6_stress

from numpy.linalg import norm

import pytest

eps_hp = 1e-15  # The espsilon value used to compare numbers that are high-precision
eps_lp = 1e-7  # The espsilon value used to compare numbers that are low-precision


@pytest.fixture
def default_chunk():
    lines = ["TEST", "A", "TEST", "| Number of atoms: 200 atoms"]
    return AimsOutChunk(lines)


def test_reverse_search_for(default_chunk):
    assert default_chunk.reverse_search_for(["TEST"]) == 2
    assert default_chunk.reverse_search_for(["TEST"], 1) == 2

    assert default_chunk.reverse_search_for(["TEST A"]) == LINE_NOT_FOUND

    assert default_chunk.reverse_search_for(["A"]) == 1
    assert default_chunk.reverse_search_for(["A"], 2) == LINE_NOT_FOUND


def test_search_for_all(default_chunk):
    assert default_chunk.search_for_all("TEST") == [0, 2]
    assert default_chunk.search_for_all("TEST", 0, 1) == [0]
    assert default_chunk.search_for_all("TEST", 1, -1) == [2]


def test_search_parse_scalar(default_chunk):
    assert default_chunk.parse_scalar("n_atoms") == 200
    assert default_chunk.parse_scalar("n_electrons") is None


@pytest.fixture
def empty_header_chunk():
    return AimsOutHeaderChunk([])


@pytest.mark.parametrize(
    "attrname",
    [
        "n_atoms",
        "n_bands",
        "n_electrons",
        "n_spins",
        "initial_atoms",
    ],
)
def test_missing_parameter(attrname, empty_header_chunk):
    with pytest.raises(AimsParseError, match="No information about"):
        getattr(empty_header_chunk, attrname)


def test_default_header_electronic_temperature(empty_header_chunk):
    assert empty_header_chunk.electronic_temperature == 0.1


def test_default_header_constraints(empty_header_chunk):
    assert empty_header_chunk.constraints == []


def test_default_header_initial_cell(empty_header_chunk):
    assert empty_header_chunk.initial_cell is None


def test_default_header_is_md(empty_header_chunk):
    assert not empty_header_chunk.is_md


def test_default_header_is_relaxation(empty_header_chunk):
    assert not empty_header_chunk.is_relaxation


def test_default_header_n_k_points(empty_header_chunk):
    assert empty_header_chunk.n_k_points is None


def test_default_header_k_points(empty_header_chunk):
    assert empty_header_chunk.k_points is None


def test_default_header_k_point_weights(empty_header_chunk):
    assert empty_header_chunk.k_point_weights is None


@pytest.fixture
def initial_cell():
    return np.array(
        [
            [1.00000000, 2.70300000, 3.70300000],
            [4.70300000, 2.00000000, 6.70300000],
            [8.70300000, 7.70300000, 3.00000000],
        ]
    )


@pytest.fixture
def initial_positions():
    return (np.array([[0.000, 1.000, 2.000], [2.703, 3.703, 4.703]]),)


@pytest.fixture
def header_chunk():
    lines = """
        | Number of atoms                   :        2
        | Number of Kohn-Sham states (occupied + empty):       3
        The structure contains        2 atoms,  and a total of         28.000 electrons.
        | Number of spin channels           :        2
        Occupation type: Gaussian broadening, width =   0.500000E-01 eV.
        Found relaxation constraint for atom      1: x coordinate fixed.
        Found relaxation constraint for atom      1: y coordinate fixed.
        Found relaxation constraint for atom      2: All coordinates fixed.
        Input geometry:
        | Unit cell:
        |        1.00000000        2.70300000        3.70300000
        |        4.70300000        2.00000000        6.70300000
        |        8.70300000        7.70300000        3.00000000
        | Atomic structure:
        |       Atom                x [A]            y [A]            z [A]
        |    1: Species Na            0.00000000        1.00000000        2.00000000
        |    2: Species Cl            2.70300000        3.70300000        4.70300000
        Initializing the k-points
        Using symmetry for reducing the k-points
        | k-points reduced from:        8 to        8
        | Number of k-points                             :         8
        The eigenvectors in the calculations are REAL.
        | K-points in task   0:         2
        | K-points in task   1:         2
        | K-points in task   2:         2
        | K-points in task   3:         2
        | k-point: 1 at     0.000000    0.000000    0.000000  , weight:   0.12500000
        | k-point: 2 at     0.000000    0.000000    0.500000  , weight:   0.12500000
        | k-point: 3 at     0.000000    0.500000    0.000000  , weight:   0.12500000
        | k-point: 4 at     0.000000    0.500000    0.500000  , weight:   0.12500000
        | k-point: 5 at     0.500000    0.000000    0.000000  , weight:   0.12500000
        | k-point: 6 at     0.500000    0.000000    0.500000  , weight:   0.12500000
        | k-point: 7 at     0.500000    0.500000    0.000000  , weight:   0.12500000
        | k-point: 8 at     0.500000    0.500000    0.500000  , weight:   0.12500000
        Geometry relaxation: A file "geometry.in.next_step" is written out by default after each step.
        Complete information for previous time-step:
    """
    lines = lines.splitlines()
    for ll, line in enumerate(lines):
        lines[ll] = line.strip()

    return AimsOutHeaderChunk(lines)


def test_header_n_atoms(header_chunk):
    assert header_chunk.n_atoms == 2


def test_header_n_bands(header_chunk):
    assert header_chunk.n_bands == 3


def test_header_n_electrons(header_chunk):
    assert header_chunk.n_electrons == 28


def test_header_n_spins(header_chunk):
    assert header_chunk.n_spins == 2


def test_header_constraints(header_chunk):
    assert len(header_chunk.constraints) == 2
    assert header_chunk.constraints[0].index == 0
    assert header_chunk.constraints[1].index == 1
    assert np.all(header_chunk.constraints[0].mask == [False, False, True])


def test_header_initial_atoms(header_chunk, initial_cell, initial_positions):
    assert len(header_chunk.initial_atoms) == 2
    assert np.allclose(
        header_chunk.initial_atoms.cell,
        initial_cell,
    )
    assert np.allclose(header_chunk.initial_atoms.positions, initial_positions)
    assert np.all(["Na", "Cl"] == header_chunk.initial_atoms.symbols)
    assert all(
        [
            str(const_1) == str(const_2)
            for const_1, const_2 in zip(
                header_chunk.constraints, header_chunk.initial_atoms.constraints
            )
        ]
    )


def test_header_initial_cell(header_chunk, initial_cell):
    assert np.allclose(header_chunk.initial_cell, initial_cell)


def test_header_electronic_temperature(header_chunk):
    assert header_chunk.electronic_temperature == 0.05


def test_header_is_md(header_chunk):
    assert header_chunk.is_md


def test_header_is_relaxation(header_chunk):
    assert header_chunk.is_relaxation


def test_header_n_k_points(header_chunk):
    assert header_chunk.n_k_points == 8


@pytest.fixture
def k_points():
    return np.array(
        [
            [0.000, 0.000, 0.000],
            [0.000, 0.000, 0.500],
            [0.000, 0.500, 0.000],
            [0.000, 0.500, 0.500],
            [0.500, 0.000, 0.000],
            [0.500, 0.000, 0.500],
            [0.500, 0.500, 0.000],
            [0.500, 0.500, 0.500],
        ]
    )


@pytest.fixture
def k_point_weights():
    return np.full((8), 0.125)


def test_header_k_point_weights(header_chunk, k_point_weights):
    assert np.allclose(header_chunk.k_point_weights, k_point_weights)


def test_header_k_points(header_chunk, k_points):
    assert np.allclose(header_chunk.k_points, k_points)


def test_header_header_summary(header_chunk, k_points, k_point_weights):
    header_summary = {
        "initial_atoms": header_chunk.initial_atoms,
        "initial_cell": header_chunk.initial_cell,
        "constraints": header_chunk.constraints,
        "is_relaxation": True,
        "is_md": True,
        "n_atoms": 2,
        "n_bands": 3,
        "n_electrons": 28,
        "n_spins": 2,
        "electronic_temperature": 0.05,
        "n_k_points": 8,
        "k_points": k_points,
        "k_point_weights": k_point_weights,
    }
    for key, val in header_chunk.header_summary.items():
        if key == "constraints":
            continue
        if isinstance(val, np.ndarray):
            assert np.allclose(val, header_summary[key])
        else:
            assert val == header_summary[key]


@pytest.fixture
def empty_calc_chunk(header_chunk):
    return AimsOutCalcChunk([], header_chunk)


def test_header_transfer_n_atoms(empty_calc_chunk):
    assert empty_calc_chunk.n_atoms == 2


def test_header_transfer_n_bands(empty_calc_chunk):
    assert empty_calc_chunk.n_bands == 3


def test_header_transfer_n_electrons(empty_calc_chunk):
    assert empty_calc_chunk.n_electrons == 28


def test_header_transfer_n_spins(empty_calc_chunk):
    assert empty_calc_chunk.n_spins == 2


def test_header_transfer_constraints(empty_calc_chunk):
    assert len(empty_calc_chunk.constraints) == 2
    assert empty_calc_chunk.constraints[0].index == 0
    assert empty_calc_chunk.constraints[1].index == 1
    assert np.all(empty_calc_chunk.constraints[0].mask == [False, False, True])


def test_header_transfer_initial_cell(empty_calc_chunk, initial_cell):
    assert np.allclose(empty_calc_chunk.initial_cell, initial_cell)


def test_header_transfer_initial_atoms(
    empty_calc_chunk, initial_cell, initial_positions
):
    assert np.allclose(
        empty_calc_chunk.initial_atoms.cell, empty_calc_chunk.initial_cell
    )
    assert len(empty_calc_chunk.initial_atoms) == 2
    assert np.allclose(empty_calc_chunk.initial_atoms.cell, initial_cell)
    assert np.allclose(
        empty_calc_chunk.initial_atoms.positions,
        initial_positions)
    assert np.all(["Na", "Cl"] == empty_calc_chunk.initial_atoms.symbols)
    assert all(
        [
            str(const_1) == str(const_2)
            for const_1, const_2 in zip(
                empty_calc_chunk.constraints, empty_calc_chunk.initial_atoms.constraints
            )
        ]
    )


def test_header_transfer_electronic_temperature(empty_calc_chunk):
    assert empty_calc_chunk.electronic_temperature == 0.05


def test_header_transfer_n_k_points(empty_calc_chunk):
    assert empty_calc_chunk.n_k_points == 8


def test_header_transfer_k_point_weights(empty_calc_chunk, k_point_weights):
    assert np.allclose(empty_calc_chunk.k_point_weights, k_point_weights)


def test_header_transfer_k_points(empty_calc_chunk, k_points):
    assert np.allclose(empty_calc_chunk.k_points, k_points)


def test_default_calc_energy_raises_error(empty_calc_chunk):
    with pytest.raises(
        AimsParseError, match="No energy is associated with the structure."
    ):
        getattr(empty_calc_chunk, "energy")


@pytest.mark.parametrize(
    "attrname",
    [
        "forces",
        "stresses",
        "stress",
        "free_energy",
        "n_iter",
        "magmom",
        "E_f",
        "dipole",
        "hirshfeld_charges",
        "hirshfeld_volumes",
        "hirshfeld_atomic_dipoles",
        "hirshfeld_dipole",
    ],
)
def test_chunk_defaults_none(attrname, empty_calc_chunk):
    assert getattr(empty_calc_chunk, attrname) is None


def test_default_calc_is_metallic(empty_calc_chunk):
    assert not empty_calc_chunk.is_metallic


def test_default_calc_converged(empty_calc_chunk):
    assert not empty_calc_chunk.converged


@pytest.fixture
def calc_chunk(header_chunk):
    lines = """
        | Number of self-consistency cycles          :           58
        | N = N_up - N_down (sum over all k points):         0.00000
        | Chemical potential (Fermi level):    -8.24271207 eV
        Total atomic forces (unitary forces were cleaned, then relaxation constraints were applied) [eV/Ang]:
        |    1          1.000000000000000E+00          2.000000000000000E+00          3.000000000000000E+00
        |    2          6.000000000000000E+00          5.000000000000000E+00          4.000000000000000E+00
        - Per atom stress (eV) used for heat flux calculation:
        Atom   | Stress components (1,1), (2,2), (3,3), (1,2), (1,3), (2,3)
        -------------------------------------------------------------------
        1 |    -1.0000000000E+01   -2.0000000000E+01   -3.0000000000E+01   -4.0000000000E+01   -5.0000000000E+01   -6.0000000000E+01
        2 |     1.0000000000E+01    2.0000000000E+01    3.0000000000E+01    4.0000000000E+01    5.0000000000E+01    6.0000000000E+01
        -------------------------------------------------------------------
        +-------------------------------------------------------------------+
        |              Analytical stress tensor - Symmetrized               |
        |                  Cartesian components [eV/A**3]                   |
        +-------------------------------------------------------------------+
        |                x                y                z                |
        |                                                                   |
        |  x         1.00000000       2.00000000       3.00000000           |
        |  y         2.00000000       5.00000000       6.00000000           |
        |  z         3.00000000       6.00000000       7.00000000           |
        |                                                                   |
        |  Pressure:       0.00383825   [eV/A**3]                           |
        |                                                                   |
        +-------------------------------------------------------------------+
        Energy and forces in a compact form:
        | Total energy uncorrected      :         -0.169503986610555E+05 eV
        | Total energy corrected        :         -2.169503986610555E+05 eV  <-- do not rely on this value for anything but (periodic) metals
        | Electronic free energy        :         -3.169503986610555E+05 eV
        Performing Hirshfeld analysis of fragment charges and moments.
        ----------------------------------------------------------------------
        | Atom     1: Na
        |   Hirshfeld charge        :      0.20898543
        |   Free atom volume        :     82.26734086
        |   Hirshfeld volume        :     73.39467444
        |   Hirshfeld spin moment   :      0.00000000
        |   Hirshfeld dipole vector :      0.00000000       0.00000000      -0.00000000
        |   Hirshfeld dipole moment :      0.00000000
        |   Hirshfeld second moments:      0.19955428       0.00000002       0.00000002
        |                                  0.00000002       0.19955473      -0.00000005
        |                                  0.00000002      -0.00000005       0.19955473
        ----------------------------------------------------------------------
        | Atom     2: Cl
        |   Hirshfeld charge        :     -0.20840994
        |   Free atom volume        :     65.62593663
        |   Hirshfeld volume        :     62.86011074
        |   Hirshfeld spin moment   :      0.00000000
        |   Hirshfeld dipole vector :      0.00000000       0.00000000       0.00000000
        |   Hirshfeld dipole moment :      0.00000000
        |   Hirshfeld second moments:      0.01482616      -0.00000001      -0.00000001
        |                                 -0.00000001       0.01482641       0.00000001
        |                                 -0.00000001       0.00000001       0.01482641
        ----------------------------------------------------------------------
        Writing Kohn-Sham eigenvalues.

        Spin-up eigenvalues:
        K-point:       1 at    0.000000    0.000000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             0.036749           1.000000
        2        1             0.183747           5.000000
        3        0             0.330744           9.000000

        Spin-down eigenvalues:
        K-point:       1 at    0.000000    0.000000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             0.110248           3.000000
        2        1             0.257245           7.000000
        3        0             0.404243           11.000000


        Spin-up eigenvalues:
        K-point:       2 at    0.000000    0.000000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             0.477741           13.000000
        2        1             0.624739           17.000000
        3        0             0.771736           21.000000

        Spin-down eigenvalues:
        K-point:       2 at    0.000000    0.000000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             0.551240           15.000000
        2        1             0.698237           19.000000
        3        0             0.845234           23.000000


        Spin-up eigenvalues:
        K-point:       3 at    0.000000    0.500000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             0.918733           25.000000
        2        1             1.065730           29.000000
        3        0             1.212728           33.000000

        Spin-down eigenvalues:
        K-point:       3 at    0.000000    0.500000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             0.992232           27.000000
        2        1             1.139229           31.000000
        3        0             1.286226           35.000000


        Spin-up eigenvalues:
        K-point:       4 at    0.000000    0.500000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             1.359725           37.000000
        2        1             1.506722           41.000000
        3        0             1.653720           45.000000

        Spin-down eigenvalues:
        K-point:       4 at    0.000000    0.500000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             1.433224           39.000000
        2        1             1.580221           43.000000
        3        0             1.727218           47.000000


        Spin-up eigenvalues:
        K-point:       5 at    0.500000    0.000000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             1.800717           49.000000
        2        1             1.947714           53.000000
        3        0             2.094711           57.000000

        Spin-down eigenvalues:
        K-point:       5 at    0.500000    0.000000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             1.874216           51.000000
        2        1             2.021213           55.000000
        3        0             2.168210           59.000000


        Spin-up eigenvalues:
        K-point:       6 at    0.500000    0.000000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             2.241709           61.000000
        2        1             2.388706           65.000000
        3        0             2.535703           69.000000

        Spin-down eigenvalues:
        K-point:       6 at    0.500000    0.000000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             2.315207           63.000000
        2        1             2.462205           67.000000
        3        0             2.609202           71.000000


        Spin-up eigenvalues:
        K-point:       7 at    0.500000    0.500000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             2.682701           73.000000
        2        1             2.829698           77.000000
        3        0             2.976695           81.000000

        Spin-down eigenvalues:
        K-point:       7 at    0.500000    0.500000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             2.756199           75.000000
        2        1             2.903197           79.000000
        3        0             3.050194           83.000000


        Spin-up eigenvalues:
        K-point:       8 at    0.500000    0.500000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             3.123693           85.000000
        2        1             3.270690           89.000000
        3        0             3.417687           93.000000

        Spin-down eigenvalues:
        K-point:       8 at    0.500000    0.500000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             3.197191           87.000000
        2        1             3.344189           91.000000
        3        0             3.491186           95.000000

        Current spin moment of the entire structure :
        | N = N_up - N_down (sum over all k points):         0.00000
        | S (sum over all k points)                :         0.00000

        What follows are estimated values for band gap, HOMO, LUMO, etc.
        | They are estimated on a discrete k-point grid and not necessarily exact.
        | For converged numbers, create a DOS and/or band structure plot on a denser k-grid.

        Highest occupied state (VBM) at     -8.19345940 eV (relative to internal zero)
        | Occupation number:      1.00000000
        | K-point:       1 at    0.000000    0.000000    0.000000 (in units of recip. lattice)
        | Spin channel:        1

        Lowest unoccupied state (CBM) at    -3.62542909 eV (relative to internal zero)
        | Occupation number:      0.00000000
        | K-point:       1 at    0.000000    0.000000    0.000000 (in units of recip. lattice)
        | Spin channel:        1

        ESTIMATED overall HOMO-LUMO gap:      4.56803031 eV between HOMO at k-point 1 and LUMO at k-point 1
        | This appears to be a direct band gap.
        The gap value is above 0.2 eV. Unless you are using a very sparse k-point grid
        this system is most likely an insulator or a semiconductor.
        | Chemical Potential                          :    -7.44914181 eV

        Self-consistency cycle converged.
        material is metallic within the approximate finite broadening function (occupation_type)
        Have a nice day.
        ------------------------------------------------------------
    """
    lines = lines.splitlines()
    for ll, line in enumerate(lines):
        lines[ll] = line.strip()
    return AimsOutCalcChunk(lines, header_chunk)

@pytest.fixture
def numerical_stress_chunk(header_chunk):
    lines = """
        | Number of self-consistency cycles          :           58
        | N = N_up - N_down (sum over all k points):         0.00000
        | Chemical potential (Fermi level):    -8.24271207 eV
        Total atomic forces (unitary forces were cleaned, then relaxation constraints were applied) [eV/Ang]:
        |    1          1.000000000000000E+00          2.000000000000000E+00          3.000000000000000E+00
        |    2          6.000000000000000E+00          5.000000000000000E+00          4.000000000000000E+00
        - Per atom stress (eV) used for heat flux calculation:
        Atom   | Stress components (1,1), (2,2), (3,3), (1,2), (1,3), (2,3)
        -------------------------------------------------------------------
        1 |    -1.0000000000E+01   -2.0000000000E+01   -3.0000000000E+01   -4.0000000000E+01   -5.0000000000E+01   -6.0000000000E+01
        2 |     1.0000000000E+01    2.0000000000E+01    3.0000000000E+01    4.0000000000E+01    5.0000000000E+01    6.0000000000E+01
        -------------------------------------------------------------------
        +-------------------------------------------------------------------+
        |                       Numerical stress tensor                     |
        |                    Cartesian components [eV/A**3]                 |
        +-------------------------------------------------------------------+
        |                x                y                z                |
        |                                                                   |
        |  x         1.00000000       2.00000000       3.00000000           |
        |  y         2.00000000       5.00000000       6.00000000           |
        |  z         3.00000000       6.00000000       7.00000000           |
        |                                                                   |
        |  Pressure:       0.00383825   [eV/A**3]                           |
        |                                                                   |
        +-------------------------------------------------------------------+
        Energy and forces in a compact form:
        | Total energy uncorrected      :         -0.169503986610555E+05 eV
        | Total energy corrected        :         -2.169503986610555E+05 eV  <-- do not rely on this value for anything but (periodic) metals
        | Electronic free energy        :         -3.169503986610555E+05 eV

        Writing Kohn-Sham eigenvalues.

        Spin-up eigenvalues:
        K-point:       1 at    0.000000    0.000000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             0.036749           1.000000
        2        1             0.183747           5.000000
        3        0             0.330744           9.000000

        Spin-down eigenvalues:
        K-point:       1 at    0.000000    0.000000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             0.110248           3.000000
        2        1             0.257245           7.000000
        3        0             0.404243           11.000000


        Spin-up eigenvalues:
        K-point:       2 at    0.000000    0.000000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             0.477741           13.000000
        2        1             0.624739           17.000000
        3        0             0.771736           21.000000

        Spin-down eigenvalues:
        K-point:       2 at    0.000000    0.000000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             0.551240           15.000000
        2        1             0.698237           19.000000
        3        0             0.845234           23.000000


        Spin-up eigenvalues:
        K-point:       3 at    0.000000    0.500000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             0.918733           25.000000
        2        1             1.065730           29.000000
        3        0             1.212728           33.000000

        Spin-down eigenvalues:
        K-point:       3 at    0.000000    0.500000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             0.992232           27.000000
        2        1             1.139229           31.000000
        3        0             1.286226           35.000000


        Spin-up eigenvalues:
        K-point:       4 at    0.000000    0.500000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             1.359725           37.000000
        2        1             1.506722           41.000000
        3        0             1.653720           45.000000

        Spin-down eigenvalues:
        K-point:       4 at    0.000000    0.500000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             1.433224           39.000000
        2        1             1.580221           43.000000
        3        0             1.727218           47.000000


        Spin-up eigenvalues:
        K-point:       5 at    0.500000    0.000000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             1.800717           49.000000
        2        1             1.947714           53.000000
        3        0             2.094711           57.000000

        Spin-down eigenvalues:
        K-point:       5 at    0.500000    0.000000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             1.874216           51.000000
        2        1             2.021213           55.000000
        3        0             2.168210           59.000000


        Spin-up eigenvalues:
        K-point:       6 at    0.500000    0.000000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             2.241709           61.000000
        2        1             2.388706           65.000000
        3        0             2.535703           69.000000

        Spin-down eigenvalues:
        K-point:       6 at    0.500000    0.000000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             2.315207           63.000000
        2        1             2.462205           67.000000
        3        0             2.609202           71.000000


        Spin-up eigenvalues:
        K-point:       7 at    0.500000    0.500000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             2.682701           73.000000
        2        1             2.829698           77.000000
        3        0             2.976695           81.000000

        Spin-down eigenvalues:
        K-point:       7 at    0.500000    0.500000    0.000000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             2.756199           75.000000
        2        1             2.903197           79.000000
        3        0             3.050194           83.000000


        Spin-up eigenvalues:
        K-point:       8 at    0.500000    0.500000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             3.123693           85.000000
        2        1             3.270690           89.000000
        3        0             3.417687           93.000000

        Spin-down eigenvalues:
        K-point:       8 at    0.500000    0.500000    0.500000 (in units of recip. lattice)

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1        1             3.197191           87.000000
        2        1             3.344189           91.000000
        3        0             3.491186           95.000000

        Current spin moment of the entire structure :
        | N = N_up - N_down (sum over all k points):         0.00000
        | S (sum over all k points)                :         0.00000

        What follows are estimated values for band gap, HOMO, LUMO, etc.
        | They are estimated on a discrete k-point grid and not necessarily exact.
        | For converged numbers, create a DOS and/or band structure plot on a denser k-grid.

        Highest occupied state (VBM) at     -8.19345940 eV (relative to internal zero)
        | Occupation number:      1.00000000
        | K-point:       1 at    0.000000    0.000000    0.000000 (in units of recip. lattice)
        | Spin channel:        1

        Lowest unoccupied state (CBM) at    -3.62542909 eV (relative to internal zero)
        | Occupation number:      0.00000000
        | K-point:       1 at    0.000000    0.000000    0.000000 (in units of recip. lattice)
        | Spin channel:        1

        ESTIMATED overall HOMO-LUMO gap:      4.56803031 eV between HOMO at k-point 1 and LUMO at k-point 1
        | This appears to be a direct band gap.
        The gap value is above 0.2 eV. Unless you are using a very sparse k-point grid
        this system is most likely an insulator or a semiconductor.
        | Chemical Potential                          :    -7.44914181 eV

        Self-consistency cycle converged.
        material is metallic within the approximate finite broadening function (occupation_type)
        Have a nice day.
        ------------------------------------------------------------
    """
    lines = lines.splitlines()
    for ll, line in enumerate(lines):
        lines[ll] = line.strip()
    return AimsOutCalcChunk(lines, header_chunk)


def test_calc_atoms(calc_chunk, initial_cell, initial_positions):
    assert len(calc_chunk.atoms) == 2
    assert np.allclose(calc_chunk.atoms.cell, initial_cell)
    assert np.allclose(calc_chunk.atoms.positions, initial_positions)
    assert np.all(["Na", "Cl"] == calc_chunk.atoms.symbols)
    assert all(
        [
            str(const_1) == str(const_2)
            for const_1, const_2 in zip(
                calc_chunk.constraints, calc_chunk.atoms.constraints
            )
        ]
    )


def test_calc_forces(calc_chunk):
    forces = np.array([[1.0, 2.0, 3.0], [6.0, 5.0, 4.0]])
    assert np.allclose(calc_chunk.forces, forces)

    # Different because of the constraints
    assert np.allclose(
        calc_chunk.atoms.get_forces(), np.array(
            [[0.0, 0.0, 3.0], [0.0, 0.0, 0.0]])
    )
    assert np.allclose(
        calc_chunk.atoms.get_forces(
            apply_constraint=False), forces)
    assert np.allclose(calc_chunk.results["forces"], forces)


def test_calc_stresses(calc_chunk):
    stresses = np.array(
        [
            [-10.0, -20.0, -30.0, -60.0, -50.0, -40.0],
            [10.0, 20.0, 30.0, 60.0, 50.0, 40.0],
        ]
    )
    assert np.allclose(calc_chunk.stresses, stresses)
    assert np.allclose(calc_chunk.atoms.get_stresses(), stresses)
    assert np.allclose(calc_chunk.results["stresses"], stresses)


def test_calc_stress(calc_chunk):
    stress = full_3x3_to_voigt_6_stress(
        np.array(
            [
                [1.00000000, 2.00000000, 3.00000000],
                [2.00000000, 5.00000000, 6.00000000],
                [3.00000000, 6.00000000, 7.00000000],
            ]
        )
    )
    assert np.allclose(calc_chunk.stress, stress)
    assert np.allclose(calc_chunk.atoms.get_stress(), stress)
    assert np.allclose(calc_chunk.results["stress"], stress)

def test_calc_num_stress(numerical_stress_chunk):
    stress = full_3x3_to_voigt_6_stress(
        np.array(
            [
                [1.00000000, 2.00000000, 3.00000000],
                [2.00000000, 5.00000000, 6.00000000],
                [3.00000000, 6.00000000, 7.00000000],
            ]
        )
    )
    assert np.allclose(numerical_stress_chunk.stress, stress)
    assert np.allclose(numerical_stress_chunk.atoms.get_stress(), stress)
    assert np.allclose(numerical_stress_chunk.results["stress"], stress)


def test_calc_free_energy(calc_chunk):
    free_energy = -3.169503986610555e05
    assert np.abs(calc_chunk.free_energy - free_energy) < eps_hp
    assert (
        np.abs(calc_chunk.atoms.calc.get_property(
            "free_energy") - free_energy) < eps_hp
    )
    assert np.abs(calc_chunk.results["free_energy"] - free_energy) < eps_hp


def test_calc_energy(calc_chunk):
    energy = -2.169503986610555e05
    assert np.abs(calc_chunk.energy - energy) < eps_hp
    assert np.abs(calc_chunk.atoms.get_potential_energy() - energy) < eps_hp
    assert np.abs(calc_chunk.results["energy"] - energy) < eps_hp


def test_calc_magnetic_moment(calc_chunk):
    magmom = 0
    assert calc_chunk.magmom == magmom
    assert calc_chunk.atoms.get_magnetic_moment() == magmom
    assert calc_chunk.results["magmom"] == magmom


def test_calc_n_iter(calc_chunk):
    n_iter = 58
    assert calc_chunk.n_iter == n_iter
    assert calc_chunk.results["n_iter"] == n_iter


def test_calc_fermi_energy(calc_chunk):
    Ef = -8.24271207
    assert np.abs(calc_chunk.E_f - Ef) < eps_lp
    assert np.abs(calc_chunk.results["fermi_energy"] - Ef) < eps_lp


def test_calc_dipole(calc_chunk):
    assert calc_chunk.dipole is None


def test_calc_is_metallic(calc_chunk):
    assert calc_chunk.is_metallic


def test_calc_converged(calc_chunk):
    assert calc_chunk.converged


def test_calc_hirshfeld_charges(calc_chunk):
    hirshfeld_charges = [0.20898543, -0.20840994]
    assert np.allclose(calc_chunk.hirshfeld_charges, hirshfeld_charges)
    assert np.allclose(
        calc_chunk.results["hirshfeld_charges"],
        hirshfeld_charges)


def test_calc_hirshfeld_volumes(calc_chunk):
    hirshfeld_volumes = [73.39467444, 62.86011074]
    assert np.allclose(calc_chunk.hirshfeld_volumes, hirshfeld_volumes)
    assert np.allclose(
        calc_chunk.results["hirshfeld_volumes"],
        hirshfeld_volumes)


def test_calc_hirshfeld_atomic_dipoles(calc_chunk):
    hirshfeld_atomic_dipoles = np.zeros((2, 3))
    assert np.allclose(
        calc_chunk.hirshfeld_atomic_dipoles,
        hirshfeld_atomic_dipoles)
    assert np.allclose(
        calc_chunk.results["hirshfeld_atomic_dipoles"], hirshfeld_atomic_dipoles
    )


def test_calc_hirshfeld_dipole(calc_chunk):
    assert calc_chunk.hirshfeld_dipole is None


@pytest.fixture
def molecular_header_chunk():
    lines = """
        | Number of atoms                   :        3
        | Number of spin channels           :        1
        The structure contains        3 atoms  and a total of         10.000 electrons.
        Input geometry:
        | Atomic structure:
        |       Atom                x [A]            y [A]            z [A]
        |    1: Species O             0.00000000        0.00000000        0.00000000
        |    2: Species H             0.95840000        0.00000000        0.00000000
        |    3: Species H            -0.24000000        0.92790000        0.00000000
        'Geometry relaxation: A file geometry.in.next_step is written out by default after each step.'
        | Maximum number of basis functions            :        7
        | Number of Kohn-Sham states (occupied + empty):       11
        Reducing total number of  Kohn-Sham states to        7.
    """

    lines = lines.splitlines()
    for ll, line in enumerate(lines):
        lines[ll] = line.strip()

    return AimsOutHeaderChunk(lines)


@pytest.mark.parametrize(
    "attrname",
    [
        "k_points",
        "k_point_weights",
        "initial_cell",
        "n_k_points",
    ],
)
def test_chunk_molecular_header_defaults_none(attrname, molecular_header_chunk):
    assert getattr(molecular_header_chunk, attrname) is None


def test_molecular_header_constraints(molecular_header_chunk):
    assert molecular_header_chunk.constraints == []


def test_molecular_header_n_bands(molecular_header_chunk):
    assert molecular_header_chunk.n_bands == 7


def test_molecular_header_initial_atoms(
        molecular_header_chunk, molecular_positions):
    assert len(molecular_header_chunk.initial_atoms) == 3
    assert np.all(["O", "H", "H"] ==
                  molecular_header_chunk.initial_atoms.symbols)
    assert np.allclose(
        molecular_header_chunk.initial_atoms.positions,
        np.array(
            [
                [0.00000000, 0.00000000, 0.00000000],
                [0.95840000, 0.00000000, 0.00000000],
                [-0.24000000, 0.92790000, 0.00000000],
            ]
        ),
    )


@pytest.fixture
def molecular_calc_chunk(molecular_header_chunk):
    lines = """
        | Number of self-consistency cycles          :           7
        | Chemical Potential                          :    -0.61315483 eV
        Updated atomic structure:
        x [A]             y [A]             z [A]
        atom        -0.00191785       -0.00243279        0.00000000  O
        atom         0.97071531       -0.00756333        0.00000000  H
        atom        -0.25039746        0.93789612       -0.00000000  H
        | Total dipole moment [eAng]          :          0.260286493869765E+00         0.336152447755231E+00         0.470003778119121E-15
        Energy and forces in a compact form:
        | Total energy uncorrected      :         -0.206778551123339E+04 eV
        | Total energy corrected        :         -5.206778551123339E+04 eV  <-- do not rely on this value for anything but (periodic) metals
        | Electronic free energy        :         -2.206778551123339E+04 eV
        Total atomic forces (unitary forces cleaned) [eV/Ang]:
        |    1          0.502371357164392E-03          0.518627676606471E-03          0.000000000000000E+00
        |    2         -0.108826758257187E-03         -0.408128912334209E-03         -0.649037698626122E-27
        |    3         -0.393544598907207E-03         -0.110498764272267E-03         -0.973556547939183E-27
        Performing Hirshfeld analysis of fragment charges and moments.
        ----------------------------------------------------------------------
        | Atom     1: O
        |   Hirshfeld charge        :     -0.32053200
        |   Free atom volume        :     23.59848617
        |   Hirshfeld volume        :     21.83060659
        |   Hirshfeld dipole vector :      0.04249319       0.05486053       0.00000000
        |   Hirshfeld dipole moment :      0.06939271
        |   Hirshfeld second moments:      0.04964380      -0.04453278      -0.00000000
        |                                 -0.04453278       0.02659295       0.00000000
        |                                 -0.00000000       0.00000000      -0.05608173
        ----------------------------------------------------------------------
        | Atom     2: H
        |   Hirshfeld charge        :      0.16022630
        |   Free atom volume        :     10.48483941
        |   Hirshfeld volume        :      6.07674041
        |   Hirshfeld dipole vector :      0.13710134      -0.00105126       0.00000000
        |   Hirshfeld dipole moment :      0.13710537
        |   Hirshfeld second moments:      0.12058896      -0.01198026      -0.00000000
        |                                 -0.01198026       0.14550360       0.00000000
        |                                 -0.00000000       0.00000000       0.10836357
        ----------------------------------------------------------------------
        | Atom     3: H
        |   Hirshfeld charge        :      0.16020375
        |   Free atom volume        :     10.48483941
        |   Hirshfeld volume        :      6.07684447
        |   Hirshfeld dipole vector :     -0.03534982       0.13248706       0.00000000
        |   Hirshfeld dipole moment :      0.13712195
        |   Hirshfeld second moments:      0.14974686      -0.00443579      -0.00000000
        |                                 -0.00443579       0.11633028      -0.00000000
        |                                 -0.00000000      -0.00000000       0.10836209
        ----------------
        Writing Kohn-Sham eigenvalues.

        State    Occupation    Eigenvalue [Ha]    Eigenvalue [eV]
        1       2.00000         -18.640915         -507.24511
        2       2.00000          -0.918449          -24.99226
        3       2.00000          -0.482216          -13.12175
        4       2.00000          -0.338691           -9.21626
        5       2.00000          -0.264427           -7.19543
        6       0.00000          -0.000414           -0.01127
        7       0.00000           0.095040            2.58616

        Highest occupied state (VBM) at     -7.19542820 eV
        | Occupation number:      2.00000000

        Lowest unoccupied state (CBM) at    -0.01126981 eV
        | Occupation number:      0.00000000

        Overall HOMO-LUMO gap:      7.18415839 eV.
        | Chemical Potential                          :    -0.61315483 eV

        Self-consistency cycle converged.
        Have a nice day.
        ------------------------------------------------------------

    """
    lines = lines.splitlines()
    for ll, line in enumerate(lines):
        lines[ll] = line.strip()
    return AimsOutCalcChunk(lines, molecular_header_chunk)


@pytest.fixture
def molecular_positions():
    return np.array(
        [
            [-0.00191785, -0.00243279, 0.00000000],
            [0.97071531, -0.00756333, 0.00000000],
            [-0.25039746, 0.93789612, 0.00000000],
        ]
    )


def test_molecular_calc_atoms(molecular_calc_chunk, molecular_positions):
    assert len(molecular_calc_chunk.atoms) == 3
    assert np.allclose(
        molecular_calc_chunk.atoms.positions,
        molecular_positions)
    assert np.all(["O", "H", "H"] == molecular_calc_chunk.atoms.symbols)


def test_molecular_calc_forces(molecular_calc_chunk):
    forces = np.array(
        [
            [0.502371357164392e-03, 0.518627676606471e-03, 0.000000000000000e00],
            [-0.108826758257187e-03, -0.408128912334209e-03, -0.649037698626122e-27],
            [-0.393544598907207e-03, -0.110498764272267e-03, -0.973556547939183e-27],
        ]
    )
    assert np.allclose(molecular_calc_chunk.forces, forces)
    assert np.allclose(molecular_calc_chunk.atoms.get_forces(), forces)
    assert np.allclose(molecular_calc_chunk.results["forces"], forces)


@pytest.mark.parametrize(
    "attrname",
    [
        "stresses",
        "stress",
        "magmom",
        "E_f",
    ],
)
def test_chunk_molecular_defaults_none(attrname, molecular_calc_chunk):
    assert getattr(molecular_calc_chunk, attrname) is None


def test_molecular_calc_free_energy(molecular_calc_chunk):
    free_energy = -2.206778551123339e04
    assert np.abs(molecular_calc_chunk.free_energy - free_energy) < eps_hp
    assert np.abs(
        molecular_calc_chunk.results["free_energy"] -
        free_energy) < eps_hp
    assert (
        np.abs(
            molecular_calc_chunk.atoms.calc.get_property(
                "free_energy") - free_energy
        )
        < eps_hp
    )


def test_molecular_calc_energy(molecular_calc_chunk):
    energy = -0.206778551123339e04
    assert np.abs(molecular_calc_chunk.energy - energy) < eps_hp
    assert np.abs(
        molecular_calc_chunk.atoms.get_potential_energy() -
        energy) < eps_hp
    assert np.abs(molecular_calc_chunk.results["energy"] - energy) < eps_hp


def test_molecular_calc_n_iter(molecular_calc_chunk):
    n_iter = 7
    assert molecular_calc_chunk.n_iter == n_iter
    assert molecular_calc_chunk.results["n_iter"] == n_iter


def test_molecular_calc_dipole(molecular_calc_chunk):
    dipole = [0.260286493869765, 0.336152447755231, 0.470003778119121e-15]
    assert np.allclose(molecular_calc_chunk.dipole, dipole)
    assert np.allclose(molecular_calc_chunk.atoms.get_dipole_moment(), dipole)
    assert np.allclose(molecular_calc_chunk.results["dipole"], dipole)


def test_molecular_calc_is_metallic(molecular_calc_chunk):
    assert not molecular_calc_chunk.is_metallic


def test_molecular_calc_converged(molecular_calc_chunk):
    assert molecular_calc_chunk.converged


@pytest.fixture
def molecular_hirshfeld_charges():
    return np.array([-0.32053200, 0.16022630, 0.16020375])


def test_molecular_calc_hirshfeld_charges(
    molecular_calc_chunk, molecular_hirshfeld_charges
):
    assert np.allclose(
        molecular_calc_chunk.hirshfeld_charges, molecular_hirshfeld_charges
    )
    assert np.allclose(
        molecular_calc_chunk.results["hirshfeld_charges"], molecular_hirshfeld_charges
    )


def test_molecular_calc_hirshfeld_volumes(molecular_calc_chunk):
    hirshfeld_volumes = np.array([21.83060659, 6.07674041, 6.07684447])
    assert np.allclose(
        molecular_calc_chunk.hirshfeld_volumes,
        hirshfeld_volumes)
    assert np.allclose(
        molecular_calc_chunk.results["hirshfeld_volumes"], hirshfeld_volumes
    )


def test_molecular_calc_hirshfeld_atomic_dipoles(molecular_calc_chunk):
    hirshfeld_atomic_dipoles = np.array(
        [
            [0.04249319, 0.05486053, 0.00000000],
            [0.13710134, -0.00105126, 0.00000000],
            [-0.03534982, 0.13248706, 0.00000000],
        ]
    )
    assert np.allclose(
        molecular_calc_chunk.hirshfeld_atomic_dipoles, hirshfeld_atomic_dipoles
    )
    assert np.allclose(
        molecular_calc_chunk.results["hirshfeld_atomic_dipoles"],
        hirshfeld_atomic_dipoles,
    )


def test_molecular_calc_hirshfeld_dipole(
    molecular_calc_chunk, molecular_hirshfeld_charges, molecular_positions
):
    hirshfeld_dipole = np.sum(
        molecular_hirshfeld_charges.reshape((-1, 1)) * molecular_positions, axis=1
    )
    assert np.allclose(molecular_calc_chunk.hirshfeld_dipole, hirshfeld_dipole)
    assert np.allclose(
        molecular_calc_chunk.results["hirshfeld_dipole"], hirshfeld_dipole
    )


tests/test_parsers/test_chunks.py
tpurcell90/atomate2-fhi-aims