"""Module for sampling snapshots from ASE MD simulations."""
from __future__ import annotations
import hashlib
import logging
from typing import TYPE_CHECKING
from directupsampling.folders.read import (
read_folders_to_sampler,
read_grid_from_folders,
)
from directupsampling.folders.write import write_folders_from_sampler, write_joblist
from directupsampling.grid import Grid, GridPoint
from directupsampling.parallel import world
from directupsampling.snapshots import SnapshotContainer
from directupsampling.utils import chunks, make_single_point_calculator, set_volume
from .md import DummyMDSampler, MDSampler
from .settings import SamplingSetting
if TYPE_CHECKING:
from collections.abc import Callable
from ase.atoms import Atoms
from ase.calculators.calculator import Calculator
from directupsampling.parallel import DummyMPI
logger = logging.getLogger(__name__)
def _parse_grid_and_container(
grid: Grid | list[tuple[float]] | None,
container: SnapshotContainer | None,
) -> tuple[Grid, SnapshotContainer]:
if grid is None and container is None:
msg = "At least one of 'grid' and 'snapshot_container' must be provided."
raise TypeError(msg)
if grid is None:
grid = container.index.get_grid()
grid = Grid.from_any(grid)
if container is None:
index_names = tuple(dict.fromkeys((*tuple(grid.axes), "seed", "step")))
container = SnapshotContainer(index_names)
return grid, container
[docs]
class Seeder:
"""Class for making consistent seeds."""
def __init__(
self,
base_seed: int,
min_seed: int = 10_000,
max_seed: int = 10_000_000,
) -> None:
"""Initialize seeder with base seed and seed range.
Parameters
----------
base_seed : int
The base seed to use for generating seeds.
min_seed : int, optional
The minimum seed value (inclusive). Default is 10_000.
max_seed : int, optional
The maximum seed value (exclusive). Default is 10_000_000.
"""
self.base_seed = base_seed
self.min_seed = min_seed
self.max_seed = max_seed
self.key_counts = {}
[docs]
def make_seed(self, key: str) -> int:
"""Make a new seed based on key.
Parameters
----------
key : str
A string key for which to make a seed. The seed will be based on
this key and the number of times this key has been used before.
Returns
-------
seed : int
A seed integer in the range [self.min_seed, self.max_seed) based on
the key and the number of times this key has been used before.
"""
n = self.key_counts.get(key, 0) + 1
self.key_counts[key] = n
numbered_key = key + f"_{n}"
h = hashlib.sha256(numbered_key.encode()).hexdigest()
seed_int = int(h[:8], 16)
min_seed = self.min_seed
max_seed = self.max_seed
return min_seed + seed_int % (max_seed - min_seed)
[docs]
class GridSeeder(Seeder):
"""Class for making consistent seeds based on grid."""
[docs]
def make_seed(self, grid_point: GridPoint) -> int:
"""Make a new seed based on grid point.
Parameters
----------
grid_point : tuple of float
A grid point for which to make a seed.
Returns
-------
seed : int
A seed integer in the range [self.min_seed, self.max_seed) based on
the grid point and the number of times this grid point has been used before.
"""
values = grid_point.todict().values()
key = f"{self.base_seed}" + "_".join(f"{_:.12f}" for _ in values)
return Seeder.make_seed(self, key)
[docs]
def make_grid_seeds(self, grid: Grid) -> list[int]:
"""Make seeds for all grid points.
Parameters
----------
grid : Grid
The grid for which to make seeds.
Returns
-------
seeds : list of int
A list of seed integers for each grid point in the grid.
"""
return [self.make_seed(gp) for gp in grid]
[docs]
def scatter_grid_seeds(self, grid: Grid, comm: DummyMPI) -> list[int]:
"""Scatter seeds for all grid points over MPI ranks.
Parameters
----------
grid : Grid
The grid for which to make seeds.
comm : DummyMPI
The MPI communicator to use for scattering the seeds.
Returns
-------
seeds : list of int
A list of seed integers scattered consistently with Grid.iscatter()
over the MPI ranks.
"""
all_seeds = []
if comm.rank == 0:
all_seeds = self.make_grid_seeds(grid)
seed_chunks = chunks(all_seeds, comm.size, fill_none=True)
return comm.scatter(seed_chunks, root=0)
[docs]
def allgather(self, comm: DummyMPI) -> None:
"""Gather key_counts from all ranks."""
all_key_counts = comm.gather(self.key_counts, root=0)
if comm.rank == 0:
for kc in all_key_counts:
for key, count in kc.items():
n = self.key_counts.get(key, 0) + count
self.key_counts[key] = n
comm.bcast(self.key_counts, root=0)
[docs]
class GridSampler:
"""Class for sampling snapshots on a grid."""
def __init__(
self,
atoms: Atoms = None,
grid: Grid | list[tuple[float, ...]] | None = None,
setting: SamplingSetting | dict | None = None,
snapshot_container: SnapshotContainer | None = None,
comm: DummyMPI = world,
) -> None:
"""Class for sampling snaphshots on a grid.
Parameters
----------
atoms : Atoms
Initial atoms from which the MD will start (with rescaled
volume). Must have an attribute `.calc` with an ASE `Calculator`.
grid : Grid or list of tuple, optional
Instance of Grid class or list of coordinates. If None, grid is
taken from `snapshot_container`.
setting : SamplingSetting, optional
Settings for the sampling. If None, default settings are used.
snapshot_container : SnapshotContainer, optional
Object to contain the sampled snapshots. The container index names
should at least contain the grid names plus ("seed", "step").
comm : DummyMPI, optional
MPI communicator to use for parallel sampling. Default is `world`.
For the arguments `sampling_offset`, `sampling_interval`,
`thermostat_kwargs` and `thermostat_class`, see the `Dynamics` class.
Notes
-----
Elements of `grid` can be repeated for parallel/multiple sampling with
different random seeds.
`setting.seed` will be regenerated for each grid point automatically,
based on the `self.setting.rng` random number generator, by default
initialized with `self.setting.seed`.
"""
self.atoms = atoms
self.setting = SamplingSetting.from_any(setting)
self.comm = comm
self._observers: list[Callable] = []
grid, snapshot_container = _parse_grid_and_container(grid, snapshot_container)
self.grid = grid
self.snapshot_container = snapshot_container
self.seed_maker = GridSeeder(self.setting.seed)
def __repr__(self) -> str:
return f"SnapshotSampler({self.__dict__!r})"
[docs]
def attach(self, observer: Callable) -> None:
"""Attach on observer to be evaluated each grid list step."""
self._observers.append(observer)
[docs]
def sample_snapshots(
self,
nsnapshots: int = 150,
**kwargs: dict,
) -> SnapshotContainer:
"""Sample `nsnapshots` on `self.grid` with `self.atoms`.
Parameters
----------
nsnapshots: int
The number of snapshots to sample.
\*\*kwargs
Additional keyword arguments to override those in `self.setting`.
Providing `seed` here will override (and set the same) seed for each
grid point.
Returns
-------
snapshot_container : SnapshotContainer
The sampled snapshots.
"""
sampling_offset = self.setting.sampling_offset
sampling_interval = self.setting.sampling_interval
ntimesteps = sampling_interval * (nsnapshots - 1) + sampling_offset
if self.comm.rank == 0:
logger.info("Starting MD sampling with %d steps", ntimesteps)
self.comm.barrier() # For nicer printing
atoms = self.atoms
initial_positions = atoms.get_scaled_positions()
tmp_snapshots = SnapshotContainer(self.snapshot_container.index.names)
seeds = self.seed_maker.scatter_grid_seeds(self.grid, self.comm)
for gp, seed in zip(self.grid.iscatter(), seeds, strict=True):
if gp is None:
self._run_dummy(ntimesteps, **kwargs)
continue
set_volume(atoms, gp.volume * len(atoms))
setting = self.setting.replace(**{"seed": seed} | kwargs)
snapshots = self._run_sampling(atoms, ntimesteps, setting)
tmp_snapshots.extend(snapshots)
self.snapshot_container.extend(snapshots)
# Reset positions to allow future sampling from initial positions
atoms.set_scaled_positions(initial_positions)
tmp_snapshots.gather_container(self.comm)
self.snapshot_container.gather_container(self.comm)
return tmp_snapshots
def _run_dummy(self, ntimesteps: int, **kwargs: dict) -> None:
"""Run dummy MD to avoid parallel deadlock."""
dyn = DummyMDSampler(self.setting.replace(**kwargs), self.comm)
dyn.run(ntimesteps)
def _run_sampling(
self,
atoms: Atoms,
ntimesteps: int,
setting: SamplingSetting,
) -> SnapshotContainer:
gp = self.grid.current
logger.info(gp)
self.call_observers()
snapshots = SnapshotContainer(self.snapshot_container.index.names)
snapshots.set_index_defaults(**gp.todict())
self.md = MDSampler(
atoms,
gp.temperature,
setting=setting,
snapshot_container=snapshots,
comm=self.comm,
)
self.md.run(ntimesteps)
return snapshots
[docs]
def call_observers(self) -> None:
"""Call observers."""
for observer in self._observers:
observer()
[docs]
def calculate_with(
self,
calc: Calculator,
properties: list[str] = ["energy", "forces"],
inplace: bool = False,
) -> SnapshotContainer:
"""Perform calculations for the snapshots in *self.snapshot_container*.
Parameters
----------
calc : ASE Calculator or list of ASE Calculator
Calculators to calculate with. Labeled with calc.name. If several
names are the same, a number is appended.
properties: list of str (optional)
List of which properties to calculate. Defaults to energy and forces.
Note: For some Calculators calculating for example energy will
automatically also calculate other properties like forces.
inplace: bool (optional)
If True, changes the Calculators of self.snapshot_container.atoms.
Defaul is False.
Returns
-------
snapshot_container : SnapshotContainer
A SnapshotContainer calculated with the provided calculator.
If inplace is True, this is self.snapshot_container.
"""
if not inplace:
new_sc = SnapshotContainer(index_names=self.snapshot_container.index.names)
else:
new_sc = None
for gp in self.grid:
self._icalc(calc, properties, new_sc)
if not inplace:
new_sc.gather_container(self.comm)
return new_sc
else:
self.snapshot_container.gather_container(self.comm)
return self.snapshot_container
def _icalc(
self,
calc: Calculator,
properties: list[str],
new_sc: SnapshotContainer | None,
) -> None:
gp = self.grid.current
self.call_observers()
logger.info(gp)
snapshots = self.snapshot_container.filter(
f"volume == {gp.volume} & temperature == {gp.temperature}"
)
for index, atoms in zip(snapshots.index, snapshots.get("atoms"), strict=True):
if new_sc is not None:
new_atoms = atoms.copy()
new_sc.insert(index, atoms=new_atoms)
else: # inplace
new_atoms = atoms
new_atoms.calc = calc
for prop in properties:
calc.get_property(prop, new_atoms, allow_calculation=True)
new_atoms.calc = make_single_point_calculator(new_atoms)
[docs]
def write_folders(
self,
joblist_file: str = "jobList",
*,
short: bool = False,
**kwargs: dict,
) -> list:
"""Write folders from sampler on the grid provided.
Parameters
----------
joblist_file : str (optional)
Name of the file to which to write the paths of each job. Defaults
to 'jobList'
short : bool (optional)
If True, writes only the path to each (V,T) grid point, else writes
path to every snapshot. Default False
root_path : str | Path (optional)
Where to write the folder structure. The default is the current
working directory.
kwargs
Arguments sent to write_folders_from_sampler.
Returns
-------
joblist : list
Notes
-----
If a Calculator is given with the parameter ismear = -1, the parameter
sigma will be set to correspond to the temperature.
"""
joblist, joblist_angk = write_folders_from_sampler(self, **kwargs)
if short:
joblist = joblist_angk
if joblist_file:
write_joblist(joblist, filename=joblist_file)
return joblist
def _read_folders(self, *args: tuple, **kwargs: dict) -> None:
read_folders_to_sampler(self, *args, **kwargs)
[docs]
@classmethod
def read_folders(cls, path: str, **kwargs: dict) -> GridSampler:
"""Make GridSampler from a given folder.
Returns
-------
GridSampler
"""
logger.info("Reading folders from %s ...", path)
grid = read_grid_from_folders(path)
logger.debug(" Grid: %s", grid)
new = cls(grid=grid)
logger.debug(" GridSampler created with %d grid points", len(new.grid))
new._read_folders(path, **kwargs)
logger.info(" %d snapshots read", len(new.snapshot_container))
return new