from __future__ import annotations
from typing import Iterator, List
import ase
from clease.datastructures import MCStep
from clease.calculator import CleaseCacheCalculator
from .mc_observer import MCObserver
__all__ = ("MoveObserver",)
[docs]class MoveObserver(MCObserver):
"""Store each step from an MC run to reconstruct the individual atoms objects later.
The interval must be set to 1 when attaching this observer,
as otherwise steps may be lost and the reconstruction may end up being incorrect.
Args:
base_atoms (ase.Atoms): The base atoms object which is run in the MC.
only_accept (bool, optional): Selects whether the only accepted moves or
all the attempted moves are saved. If False, every move will be saved.
Defaults to False.
"""
def __init__(self, base_atoms: ase.Atoms, only_accept: bool = False):
super().__init__()
self.base_atoms = base_atoms.copy()
self.only_accept = only_accept
self.steps = []
[docs] def observe_step(self, mc_step: MCStep) -> None:
"""Observe a single step"""
if self.only_accept is True and mc_step.move_accepted is False:
# Step was rejected, and we only want accepted steps
return
self.steps.append(mc_step)
[docs] def reset(self) -> None:
self.steps = []
[docs] def reconstruct(self) -> List[ase.Atoms]:
"""Rebuild the atoms objects as defined by the observed changes."""
return list(self.reconstruct_iter())
[docs] def reconstruct_iter(self) -> Iterator[ase.Atoms]:
"""Iterator which builds the atoms objects 1-by-1."""
atoms = self.base_atoms.copy()
for step in self.steps:
# Update the current atoms
atoms = _apply_step(atoms, step)
yield atoms
[docs] def interval_ok(self, interval: int) -> bool:
"""Missing steps will result in incorrect reconstructions"""
return interval == 1
def _apply_step(atoms: ase.Atoms, step: MCStep) -> ase.Atoms:
"""Helper function to apply a step to an atoms object.
Always a returns a copy of the atoms object."""
atoms = atoms.copy()
# Apply changes to a copy of the atoms
if step.move_accepted:
for change in step.last_move:
change.apply_change(atoms)
# Attach a calculator with the energies
calc = CleaseCacheCalculator(energy=step.energy)
atoms.calc = calc
return atoms