"""Module for sampling snapshots from ASE MD simulations."""
from __future__ import annotations
import logging
import traceback
from collections.abc import Callable
from numbers import Number
from typing import TYPE_CHECKING
import ase
import numpy as np
from ase.md import MDLogger
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
from directupsampling.parallel import DummyMPI
from directupsampling.snapshots import SnapshotContainer
from directupsampling.utils import get_calcs_names, make_single_point_calculator
from .settings import MDLogSetting, SamplingSetting, ThermostatSetting
if TYPE_CHECKING:
from collections.abc import Iterator
from ase.atoms import Atoms
from ase.calculators.calculator import Calculator
logger = logging.getLogger(__name__)
[docs]
class MDHandler:
"""Handler class for ASE Thermostat and velocity initialization."""
def __init__(
self,
atoms: Atoms,
temperature: float,
setting: SamplingSetting,
) -> None:
"""Handle ASE Thermostat and velocity initialization.
Parameters
----------
atoms : ase.Atoms
Atoms object to run MD on. Must have a `calc` attribute, with an
instance of an ASE `Calculator` derived class.
temperature : float
Temperature in K.
setting : SamplingSetting
Settings for the sampling.
Attributes
----------
atoms : ase.Atoms
The atoms object to run MD on.
rng : np.random.Generator
Random number generator, used for initialization and dynamics.
thermostat : ASE Thermostat
The thermostat instance to run MD with. Must be set with
`set_thermostat` before running.
Properties
----------
nsteps : int
Number of MD steps taken so far.
"""
self.atoms = atoms
self.rng = np.random.default_rng(setting.seed)
self.initialize(temperature)
self.set_thermostat(temperature, setting.thermostat_setting)
self.attach(setting.attach)
log_md = setting.log_md
if log_md:
if isinstance(log_md, Number) and not isinstance(log_md, bool):
setting.mdlogger_setting.interval = log_md
self.set_logger(setting.mdlogger_setting)
[docs]
def initialize(self, temperature: float) -> None:
"""Initialize atoms with the Maxwell-Boltzmann distribution of ASE."""
MaxwellBoltzmannDistribution(
self.atoms,
temperature_K=temperature,
rng=self.rng,
communicator="serial", # NOTE: syntax differs from Langevin
)
[docs]
def set_thermostat(
self,
temperature: float,
setting: ThermostatSetting | None = None,
) -> None:
"""Set the thermostat to be used for MD."""
setting = setting or ThermostatSetting()
thermostat_class = setting.class_type
# Because we always run in serial for each process, we need to give
# DummyMPI() to the thermostat. Otherwise the random numbers will be
# broadcasted across the different independent runs, as well as
# interfere with the gathering of snapshots for writing.
if ase.__version__ >= "3.25.0":
kwargs = {"comm": DummyMPI()}
else:
kwargs = {"communicator": None}
kwargs.update(setting.get_kwargs())
self.thermostat = thermostat_class(
atoms=self.atoms,
temperature_K=temperature,
**kwargs,
)
self.thermostat.rng = self.rng
[docs]
def attach(self, functions: list[Callable | tuple]) -> None:
"""Attach functions to the thermostat."""
for func in functions if isinstance(functions, list) else [functions]:
if isinstance(func, Callable):
self.thermostat.attach(func)
else:
func, interval, args, kwargs = func
self.thermostat.attach(func, interval, *args, **kwargs)
[docs]
def set_logger(self, setting: MDLogSetting | None = None) -> None:
"""Set the MDLogger to log the MD simulation."""
setting = setting or MDLogSetting()
kwargs = setting.todict()
interval = kwargs.pop("interval")
logger = MDLogger(dyn=self.thermostat, atoms=self.atoms, **kwargs)
self.thermostat.attach(logger, interval=interval)
@property
def nsteps(self) -> int:
"""Number of MD steps taken so far."""
return self.thermostat.nsteps
[docs]
def irun(self, steps=50) -> Iterator[bool]:
"""Run molecular dynamics algorithm as a generator.
Parameters
----------
steps : int
Number of molecular dynamics steps to be run.
Yields
------
converged : bool
True if the maximum number of steps are reached.
"""
return self.thermostat.irun(steps=steps)
[docs]
def run(self, steps: int = 50):
"""Run molecular dynamics algorithm.
Parameters
----------
steps : int
Number of molecular dynamics steps to be run.
Returns
-------
converged : bool
True if the maximum number of steps are reached.
"""
return self.thermostat.run(steps=steps)
[docs]
class MDSampler:
def __init__(
self,
atoms: Atoms,
temperature: float,
setting: SamplingSetting,
snapshot_container: SnapshotContainer | None = None,
comm=DummyMPI(),
) -> None:
"""Sampler class to run MD and take samples at intervals.
Parameters
----------
atoms: Atoms
Initial structure for MD. Must have a `calc` attribute, with an
instance of an ASE Calculator derived class.
temperature: float
In K.
setting: SamplingSetting
Settings for the sampling.
snapshot_container: SnapshotContainer, optional
Container to store the sampled snapshots. If None, a new container
is created with index names ("seed", "step").
"""
self.atoms = atoms
self.setting = setting
self.comm = comm
self.buffer_size = setting.buffer_for_writing
self._write_buffer = []
index_names = ["seed", "step"]
if snapshot_container is not None:
self.snapshot_container = snapshot_container
else:
self.snapshot_container = SnapshotContainer(index_names)
mdh = MDHandler(self.atoms, temperature, setting)
self.md_handler = mdh
[docs]
def is_sampling_time(self) -> bool:
step = self.md_handler.nsteps
offset = self.setting.sampling_offset
interval = self.setting.sampling_interval
return step >= offset and (((step - offset) % interval) == 0)
[docs]
def sample(self) -> None:
"""Take a sample of the MD simulation and save it to the container."""
atoms = self.atoms
data = {}
# Sampling of energy
ke = atoms.get_kinetic_energy()
pe = atoms.get_potential_energy()
data["kinetic_energy"] = ke
data["potential_energy"] = pe
data[f"potential_energy_{atoms.calc.name}"] = pe
# Special treatment of mixing calculators
calcs: list[Calculator] = []
if hasattr(atoms.calc, "calcs"): # For compatibility with ASE <3.23.0
calcs = atoms.calc.calcs
elif hasattr(atoms.calc, "mixer"):
calcs = atoms.calc.mixer.calcs
if len(calcs) > 0:
names = get_calcs_names(*calcs)
for i, (c, n) in enumerate(zip(calcs, names)):
pe = c.get_potential_energy(atoms)
data[f"potential_energy_{i}"] = pe
data[f"potential_energy_{n}"] = pe
# For easier plotting with `ase db`
data["potential_energy_difference"] = (
data[f"potential_energy_{i}"] - data["potential_energy_0"]
)
if self.setting.sample_atoms:
atoms.get_forces() # Needed to save forces for some thermostats
atoms_copy = atoms.copy()
atoms_copy.calc = make_single_point_calculator(atoms)
data["atoms"] = atoms_copy
# Save prepared data
current_index = {"seed": self.setting.seed, "step": self.md_handler.nsteps}
sc = self.snapshot_container
full_index = sc.insert(current_index, **data)
if self.setting.file:
self._write_buffer.append(full_index)
if len(self._write_buffer) >= self.buffer_size:
self.write_buffer_to_database()
[docs]
def write_buffer_to_database(self) -> None:
"""Write the buffered snapshots to the database and clear the buffer.
Notes
-----
This method also discard the written snapshots from the container, which
keeps memory usage low.
"""
# Make temporary container to gather only what should be written
tmp_sc = self.snapshot_container[self._write_buffer]
tmp_sc.gather_container(self.comm)
tmp_sc.write(self.setting.file, append=True)
self.snapshot_container.discard_indices(self._write_buffer)
self._write_buffer.clear()
self._nwrites -= 1
def _handle_value_error(self) -> None:
msg = (
"MD aborted due to ValueError at step "
f"{self.md_handler.nsteps} (rank {self.comm.rank}):\n"
f"{traceback.format_exc()}\n"
" continuing..."
)
logger.error(msg)
if self.setting.file:
# Dummy write buffer to avoid parallel deadlock
while self._nwrites > 0:
self.write_buffer_to_database()
[docs]
def run(self, ntimesteps: int) -> SnapshotContainer:
"""Run dynamics with ASE and returns snapshots.
Returns
-------
snapshots
list of ase.Atoms
"""
logger.debug(f"Running ASE MD with {self.atoms.calc} on rank {self.comm.rank}")
# Count down number of writes, which contains a collective operation,
# for later dummy write in case of crash
offset = self.setting.sampling_offset
interval = self.setting.sampling_interval
nsamples = (ntimesteps - offset) / interval + 1
self._nwrites = int(nsamples / self.buffer_size)
# Try and except to not crash (and lock) because of ValueError
try:
for _ in self.md_handler.irun(ntimesteps):
if self.is_sampling_time():
self.sample()
except ValueError:
self._handle_value_error()
# Finish with writing buffer
if self.setting.file:
self.write_buffer_to_database()
logger.debug(f"Completed ASE MD on rank {self.comm.rank}")
return self.snapshot_container
[docs]
class DummyMDSampler:
"""Dummy sampler class for parallel runs, to avoid deadlock."""
def __init__(
self,
setting: SamplingSetting,
comm: DummyMPI,
) -> None:
"""Initialize the dummy sampler."""
self.setting = setting
self.comm = comm
self._write_buffer = []
index_names = ["seed", "step"]
self.snapshot_container = SnapshotContainer(index_names)
[docs]
def write_buffer_to_database(self) -> None:
"""Pretentd to write the buffered snapshots and clear the buffer."""
# Make temporary container to gather only what should be written
tmp_sc = self.snapshot_container[self._write_buffer]
tmp_sc.gather_container(self.comm)
tmp_sc.write(self.setting.file, append=True)
self.snapshot_container.discard_indices(self._write_buffer)
self._write_buffer.clear()
self._nwrites -= 1
[docs]
def run(self, ntimesteps: int) -> None:
"""Run the dummy sampler, which only performs dummy writes to avoid deadlock."""
# Count down number of writes, which contains the collective operation
sampling_offset = self.setting.sampling_offset
sampling_interval = self.setting.sampling_interval
nsamples = (ntimesteps - sampling_offset) / sampling_interval + 1
self._nwrites = int(nsamples / self.setting.buffer_for_writing)
if self.setting.file:
# Here is the dummy write
while self._nwrites > 0:
self.write_buffer_to_database()
self.write_buffer_to_database()