"""Store the results of a vibration calculation.
This module defines the :class:`~autojob.tasks.vibration.Vibration`,
:class:`~autojob.tasks.vibration.VibrationInputs`, and
:class:`~autojob.tasks.vibration.VibrationOutputs` classes. Instances
of these classes represent the results of a vibration calculation, its inputs,
and its outputs, respectively.
For building the respective documents from a folder, the
:class:`~autojob.tasks.vibration.Vibration` and
:class:`~autojob.tasks.vibration.VibrationnOutputs` classes are
:class:`PathLoadable`.
In addition, this module defines the :func:`freeze_atoms` function-a utility
function for fixing the atoms of an :class:`.Atoms` object by tag.
Examples:
Create a vibration calculation with frozen atoms
.. code-block:: python
from ase.build import add_adsorbate, fcc100, molecule
from autojob.tasks.task import TaskInputs
from autojob.tasks.vibration import freeze_atoms, Vibration
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 vibration
# analysis
freeze_atoms(slab, [tag])
# Create the vibration calculation
tinputs = TaskInputs(atoms=slab)
vib = Vibration(task_inputs=tinputs)
Create a vibration calculation with frozen atoms
.. code-block:: python
from ase.build import add_adsorbate, fcc100, molecule
from autojob.tasks.task import TaskInputs
from autojob.tasks.vibration import Vibration, VibrationInputs
slab = fcc100("Cu")
co = molecule("CO")
add_adsorbate(slab, co, 1.8)
# Create the vibration calculation
tinputs = TaskInputs(atoms=slab)
indices = [a.index for a in slab if a.symbol != "Cu"]
vinputs = VibrationInputs(indices=indices)
vib = Vibration(task_inputs=tinputs, vibration_inputs=vinputs)
Load the results of vibration calculation
.. code-block:: python
from autojob.tasks.vibration import Vibration
src = "path/to/calculation/directory"
results = Vibration.from_directory(src)
"""
import logging
from pathlib import Path
from typing import Any
from typing import ClassVar
from typing import Literal
from typing import Self
from ase import Atoms
from ase.constraints import FixAtoms
import ase.io
from ase.vibrations.data import VibrationsData
from ase.vibrations.vibrations import Vibrations
from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import Field
from autojob import SETTINGS
from autojob.tasks.calculation import Calculation
from autojob.utils.schemas import PydanticAtoms
logger = logging.getLogger(__name__)
def _construct_vib_data(dir_name: Path, name: 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
vib = Vibrations(atoms, name=cache, nfree=nfree)
vib_json = Path(dir_name, f"{name}.json")
vib.get_vibrations().write(vib_json)
[docs]
class VibrationOutputs(BaseModel):
"""The outputs of a vibrational calculation.
This class is intended to wrap the
:class:`ase.vibrations.data.VibrationsData` class.
Example:
Observe vibrational modes
.. code:: python
import ase.io
from ase.vibrations.data import VibrationsData
from autojob.calculation.vibration import VibrationOutputs
atoms = ase.io.read("atoms.traj")
outputs = VibrationOutputs.from_directory(...)
vib_data = VibrationsData(
atoms=atoms,
hessian=outputs.hessian,
indices=outputs.indices,
)
for i, _ in enumerate(vib.get_energies()):
vib.write_mode(i)
Example:
Calculate thermodynamic quantities
.. code:: python
from ase.build import molecule
from ase.calculators.emt import EMT
from ase.vibrations.data import IdealGasThermo
from ase.vibrations.vibrations import Vibrations
from autojob.calculation.vibration import VibrationOutputs
# Run a Vibration Task
atoms = molecule("CO2")
atoms.calc = EMT()
vib = Vibrations(atoms=atoms, nfree=4, name=label, delta=0.015)
vib.run()
# Dump the results to a file
with Path('vib.json').open(mode="r", encoding="utf-8") as f:
vib.write(f)
# ... Then come back to the directory ...
# This provides validation and fall backs to reading other files
outputs = VibrationOutputs.from_directory(...)
vib_data = VibrationsData(
atoms=atoms,
hessian=outputs.hessian,
indices=outputs.indices,
)
# Get G, H, S, ZPE
ig_thermo = IdealGasThermo(
atoms=atoms,
vib_energies=vib_data.get_energies(),
geometry="linear",
potentialenergy=0.0,
symmetrynumber=1,
spin=0,
)
ig_thermo.get_gibbs_energy()
ig_thermo.get_entropy()
ig_thermo.get_enthalpy()
ig_thermo.get_ZPE_correction()
.. seealso:: :class:`ase.vibrations.data.VibrationsData`
"""
atoms: PydanticAtoms | None = Field(
default=None,
description="The :class:`ase.atoms.Atoms` object used for the Vibration task.",
)
hessian: list[list[list[list[float]]]] | None = Field(
default=None,
description="The calculated hessian for the Vibration task. An array with shape (indices, 3, indices, 3)",
)
indices: list[int] | None = Field(
default=None,
description="indices of atoms which are included "
"in Hessian. None includes all freely "
"moving atoms (i.e. not fixed ones).",
)
vib_energies: list[complex] | None = Field(
default=None, description="The energies of each vibrational mode"
)
zero_point_energy: float | None = Field(
default=None,
description="The zero-point energy",
alias="Zero-Point Energy",
)
model_config: ClassVar[ConfigDict] = ConfigDict(populate_by_name=True)
[docs]
@classmethod
def from_directory(
cls,
src: str | Path,
*,
name: str = "vib",
strict_mode: bool | None = None,
) -> "VibrationOutputs":
"""Extract thermo data from the input structure of the directory.
Args:
src: The directory of the completed 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`.
strict_mode: Whether or not to require all outputs. If True,
errors will be thrown on missing outputs.
Returns:
The thermodynamic data as `VibrationOutputs`. If no data is found,
and `strict_mode = False` every value will be None.
"""
if strict_mode is None:
strict_mode = SETTINGS.STRICT_MODE
logger.info("Loading vibration outputs from directory: %s", src)
logger.debug("Strict mode: %sabled", "en" if strict_mode else "dis")
try:
vib_json = Path(src, f"{name}.json")
if not vib_json.exists():
_construct_vib_data(Path(src), name)
with vib_json.open(mode="r", encoding="utf-8") as file:
vib_data: VibrationsData = VibrationsData.read(file)
vib_outputs = vib_data.todict()
vib_outputs["vib_energies"] = vib_data.get_energies()
vib_outputs["zero_point_energy"] = vib_data.get_zero_point_energy()
logger.info(
"Successfully vibration outputs from directory: %s", src
)
except (FileNotFoundError, KeyError, RuntimeError):
logger.exception(
"Unable to load vibration outputs from directory: %s", src
)
vib_outputs = {}
if strict_mode:
raise
return cls(**vib_outputs)
[docs]
class Vibration(Calculation):
"""A record representing a vibrational calculation."""
vibration_inputs: VibrationInputs = Field(
default_factory=VibrationInputs,
description="The inputs of the vibration calculation",
)
vibration_outputs: VibrationOutputs | None = Field(
default=None, description="The vibration calculation outputs"
)
[docs]
@classmethod
def from_directory(cls, src: str | Path, **kwargs) -> Self:
"""Generate a ``Vibration`` 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:`Vibration` or a subclass of a :class:`Vibration`.
.. seealso::
:meth:`.calculation.Calculation.from_directory`
"""
strict_mode = kwargs.get("strict_mode", SETTINGS.STRICT_MODE)
magic_mode = kwargs.get("magic_mode", False)
logger.debug("Loading vibration 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)
calculation = Calculation.from_directory(
src=src, strict_mode=strict_mode, magic_mode=False
)
data = calculation.calculation_inputs.model_extra.pop(
"vibration_inputs", {}
)
vib_inputs = VibrationInputs(**data)
vib_outputs = VibrationOutputs.from_directory(
src=src,
strict_mode=strict_mode,
name=vib_inputs.name,
)
logger.debug(
"Successfully loaded vibration calculation from directory: %s", src
)
return cls(
**calculation.model_dump(),
vibration_inputs=vib_inputs,
vibration_outputs=vib_outputs,
)
[docs]
def freeze_atoms(atoms: Atoms, tags_to_unfreeze: list[int]) -> None:
"""Freeze atoms of an input structure.
This function modifies ``atoms`` in place.
Args:
atoms: An :class:`.Atoms` object.
tags_to_unfreeze: A list of integers. Atoms with tags in this list
**will not** be frozen.
"""
indices = [a.index for a in atoms if a.tag not in tags_to_unfreeze]
c = FixAtoms(indices=indices)
atoms.set_constraint(c)