"""Store the results of an Infrared calculation.
This module defines the :class:`~autojob.tasks.infrared.Infrared`,
:class:`~autojob.tasks.infrared.InfraredInputs`, and
:class:`~autojob.tasks.infrared.InfraredOutputs` classes. Instances
of these classes represent the results of an infrared calculation, its inputs,
and its outputs, respectively.
For building the respective documents from a folder, the
:class:`~autojob.tasks.infrared.InfraredInputs` and
:class:`~autojob.tasks.infrared.InfraredOutputs` classes are
:class:`PathLoadable`.
Examples:
Create an infrared calculation
.. code-block:: python
from ase.build import add_adsorbate, fcc100, molecule
from autojob.tasks.task import TaskInputs
from autojob.tasks.infrared import Infrared
slab = fcc100("Cu")
co = molecule("CO")
# All atoms with this tag will NOT be frozen
tag = -99
co.set_tags([tag] * len(co))
add_adsorbate(slab, co, 1.8)
# Freeze the atoms that should not be displaced during the IR
# analysis
freeze_atoms(slab, [tag])
# Create the infrared calculation
tinputs = TaskInputs(atoms=slab)
ir = Infrared(task_inputs=tinputs)
Load the results of infrared calculation
.. code-block:: python
from autojob.tasks.infrared import Infrared
src = "path/to/calculation/directory"
results = Infrared.from_directory(src)
"""
import json
import logging
from pathlib import Path
import sys
from typing import Any
from typing import ClassVar
from typing import Literal
from typing import Self
if sys.version_info < (3, 12):
from typing_extensions import TypedDict
else:
from typing import TypedDict
import ase.io
from ase.vibrations.infrared import Infrared as Infrared_ASE
from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import Field
from autojob import SETTINGS
from autojob.tasks.vibration import Vibration
logger = logging.getLogger(__name__)
[docs]
class SpectraParams(TypedDict):
"""Keyword arguments for :meth:`.Infrared.write_spectra`.
Keys:
out: The name of the file in which to save the spectra.
start: The wavenumber at which to start plotting the spectra.
end: The wavenumber at which to end plotting the spectra.
npts: The number of points to use to plot the spectra.
width: The widths used to represent peaks.
type: The type of folding method. One of 'Gaussian', 'Lorentzian'
method: The finite-difference method. One of 'standard', 'frederiksen'.
direction: The finite-difference scheme in which to calculate the
Hessian. One of 'central', 'forward', 'backward'.
intensity_unit: The unit for the intensities.
normalize: Whether or not to normalize the spectrum.
"""
out: str
start: float
end: float
npts: int | None
width: float
type: str | Literal["Gaussian"]
method: str | Literal["standard", "frederiksen"]
direction: str | Literal["central"]
intensity_unit: str | Literal["(D/A)2/amu"]
normalize: bool
# TODO: Implement loading dipole moments from caches
def _construct_ir_data(
dir_name: Path, name: str, intensities_file: str
) -> None:
atoms = ase.io.read(Path(dir_name, SETTINGS.OUTPUT_ATOMS_FILE))
atoms = atoms[-1] if isinstance(atoms, list) else atoms
cache = Path(dir_name, name)
# ASE names forces caches cache-.iXN.json where i is an atom index,
# X is direction (x/y/z) and N is number of steps
nfree = 4 if list(cache.glob("cache.*--.json")) else 2
run_ir = Infrared_ASE(atoms, name=cache, nfree=nfree)
ir_frequencies = run_ir.get_frequencies().real
ir_intensities = run_ir.intensities
ir_intensities = ir_intensities[ir_frequencies.argsort()]
ir_frequencies.sort()
ir_json = Path(dir_name, intensities_file)
data = {
"ir_frequencies": ir_frequencies.tolist(),
"ir_intensities": ir_intensities.tolist(),
}
with ir_json.open(mode="w", encoding="utf-8") as file:
_ = json.dump(data, file, indent=4)
[docs]
class InfraredOutputs(BaseModel):
"""The outputs of an infrared calculation."""
ir_frequencies: list[float] | None = Field(
default=None, description="The infrared spectrum frequencies in cm^-1."
)
ir_intensities: list[float] | None = Field(
default=None,
description="The absolute intensities of the infrared spectrum.",
)
model_config: ClassVar[ConfigDict] = ConfigDict(populate_by_name=True)
[docs]
@classmethod
def from_directory(
cls,
src: str | Path,
*,
name: str = "ir",
intensities_file: str = "intensities.json",
strict_mode: bool | None = None,
) -> "InfraredOutputs":
"""Extract infrared outputs from a task directory.
Args:
src: The directory of the infrared calculation.
name: The name used to write vibration calculation results. This
should either be the name of the directory containing the cache
files or the filename stem for a JSON containing the output of
a vibration calculation such as that produced by
:meth:`~VibrationData.write`.
intensities_file: The name of the JSON file used to write infrared
intensities.
strict_mode: Whether or not to require all outputs. If True,
errors will be thrown on missing outputs.
Returns:
An instance of :class:`.InfraredOutputs`.
"""
if strict_mode is None:
strict_mode = SETTINGS.STRICT_MODE
logger.info("Loading infrared outputs from directory: %s", src)
logger.debug("Strict mode: %sabled", "en" if strict_mode else "dis")
ir_outputs = {}
try:
logger.info(
"Successfully infrared outputs from directory: %s", src
)
ir_file = Path(src, f"{intensities_file}")
if not ir_file.exists():
_construct_ir_data(src, name, intensities_file)
with ir_file.open(mode="r", encoding="utf-8") as file:
ir_outputs = json.load(file)
except (FileNotFoundError, KeyError, RuntimeError, TypeError):
logger.exception(
"Unable to load infrared outputs from directory: %s", src
)
if strict_mode:
raise
return cls(**ir_outputs)
[docs]
class Infrared(Vibration):
"""A record representing an infrared calculation."""
infrared_inputs: InfraredInputs = Field(
default_factory=InfraredInputs,
description="The inputs of the infrared calculation",
)
infrared_outputs: InfraredOutputs | None = Field(
default=None, description="The infrared calculation outputs"
)
[docs]
@classmethod
def from_directory(cls, src: str | Path, **kwargs) -> Self:
"""Generate an ``Infrared`` document from a calculation directory.
Args:
src: The directory of a calculation.
kwargs: Additional keyword arguments:
- strict_mode : Whether or not to fail on any error. Defaults to
`SETTINGS.STRICT_MODE`.
- magic_mode : Whether or not to instantiate subclasses. If
True, the task returned must be an instance determined by
metadata in the directory. Defaults to False.
Returns:
A :class:`Infrared` or a subclass of a :class:`Infrared`.
.. seealso::
:meth:`.vibration.Vibration.from_directory`
"""
strict_mode = kwargs.get("strict_mode", SETTINGS.STRICT_MODE)
magic_mode = kwargs.get("magic_mode", False)
logger.debug("Loading infrared calculation from directory: %s", src)
logger.debug("Magic mode: %sabled", "en" if magic_mode else "dis")
logger.debug("Strict mode: %sabled", "en" if strict_mode else "dis")
if magic_mode:
return cls.load_magic(src, strict_mode=strict_mode)
vibration = Vibration.from_directory(
src=src, strict_mode=strict_mode, magic_mode=False
)
data = vibration.task_inputs.model_extra.pop("infrared_inputs", {})
ir_inputs = InfraredInputs(**data)
ir_outputs = InfraredOutputs.from_directory(
src=src,
name=vibration.vibration_inputs.name,
intensities_file=ir_inputs.intensities_file,
strict_mode=strict_mode,
)
logger.debug(
"Successfully loaded infrared calculation from directory: %s", src
)
return cls(
**vibration.model_dump(),
infrared_inputs=ir_inputs,
infrared_outputs=ir_outputs,
)