from __future__ import annotations
import copy
import logging
import numbers
from typing import TYPE_CHECKING, Any
import numpy as np
from ase.calculators.singlepoint import SinglePointCalculator
from directupsampling.parallel import DummyMPI, world
from directupsampling.snapshots import io as _io
from directupsampling.snapshots.index import IndexPoint, IndexView
from directupsampling.utils import (
gather_chunks,
mlip2ase_stress,
)
if TYPE_CHECKING:
import pathlib
from collections.abc import Iterable
from ase.atoms import Atoms
try:
import pandas as pd
except ImportError:
pd = None # type: ignore[assignment]
logger = logging.getLogger(__name__)
[docs]
class SnapshotContainer:
"""Container for snapshots of atoms and their properties.
It is indexed by a flexible multi-field index.
"""
def __init__(
self,
index_names: Iterable[str],
index: Iterable[tuple] = [],
data: list[dict] | None = None,
*,
replace: bool = False,
) -> None:
"""Initialize the SnapshotContainer.
Parameters
----------
index_names: Iterable of str
Sequence of index names.
index: Iterable of tuple, optional
Indices corresponding to *data*.
data: list of dict, optional
Data dicts, one per entry in *index*.
replace: bool, optional
If True, inserting an existing index replaces its entry. (Default: False)
"""
self.replace: bool = replace
self._snapshots: dict[IndexPoint, Any] = {}
self.index: IndexView = IndexView(
self._snapshots,
tuple(np.atleast_1d(index_names)),
)
self._index_defaults: dict[str, Any] = {}
if data is None:
data = [{} for _ in index]
for ind, dct in zip(index, data, strict=True):
self.insert(ind, **dct)
[docs]
def set_index_defaults(self, **defaults: Any) -> None:
"""Set default values for index fields.
Parameters
----------
\*\*defaults
Keyword arguments with default values for the index fields.
The keys must be in `index_names`.
"""
self._index_defaults = {
name: defaults[name] for name in self.index.names if name in defaults
}
@property
def data(self) -> list[dict]:
"""List of the data dicts, one per snapshot."""
return list(self._snapshots.values())
[docs]
def get(self, key: str) -> np.ndarray | list:
"""Return all values stored under *key*, as an array if numeric.
Parameters
----------
key : str
Data field name, e.g. ``'electronic_free_energy'``.
Returns
-------
numpy.ndarray
If every stored value is numeric.
list
Otherwise.
"""
vals = [dct.get(key) for dct in self._snapshots.values()]
if all(isinstance(v, (int, float, complex, np.number)) for v in vals):
return np.array(vals)
return vals
def __getitem__(self, key: Any) -> dict | SnapshotContainer:
"""Get entry dict or a filtered container.
Parameters
----------
key : int or IndexPoint or list
- ``sc[i]`` — entry dict at positional index *i*
- ``sc[index_point]`` — entry dict for that IndexPoint
- ``sc[[True, False, ...]]`` — new container by boolean mask
- ``sc[[2, 4, 5]]`` — new container by positional indices
- ``sc[[idx1, idx2]]`` — new container by IndexPoints
Returns
-------
dict or SnapshotContainer
"""
if isinstance(key, numbers.Integral):
return self._snapshots[self.index[key]]
if isinstance(key, IndexPoint):
return self._snapshots[key]
indices = self.index[key]
return SnapshotContainer(
index_names=self.index.names,
index=indices,
data=[self._snapshots[i] for i in indices],
)
def __len__(self):
return len(self._snapshots)
def __repr__(self) -> str:
return f"SnapshotContainer(n={len(self)}, index_names={self.index.names})"
def __copy__(self) -> SnapshotContainer:
"""Return a shallow copy (shared data dicts, independent index).
Returns
-------
SnapshotContainer
"""
return self[list(self.index)]
def __deepcopy__(self, memo: dict) -> SnapshotContainer:
"""Return a deep copy.
Returns
-------
SnapshotContainer
"""
dicts = [copy.deepcopy(dct, memo) for dct in self._snapshots.values()]
return type(self)(list(self.index.names), list(self.index), dicts)
[docs]
def to_pandas(self) -> Any:
"""Convert the SnapshotContainer to a pandas DataFrame with a MultiIndex.
Returns
-------
pandas.DataFrame
DataFrame with MultiIndex.
Raises
------
ModuleNotFoundError
If pandas is not installed.
Notes
-----
All indices are casted to float.
"""
if pd is None:
msg = "pandas is required for `to_pandas()`"
raise ModuleNotFoundError(msg)
names = self.index.names
tuples = [tuple(getattr(idx, n) for n in names) for idx in self._snapshots]
ind = pd.MultiIndex.from_tuples(tuples, names=names)
ind = ind.set_levels([lvl.astype(float) for lvl in ind.levels])
return pd.DataFrame(self._snapshots.values(), index=ind)
[docs]
def filter(self, expr: str) -> SnapshotContainer:
"""Return a new SnapshotContainer with entries matching *expr*.
Parameters
----------
expr : str
Expression string, e.g. ``'volume > 15'`` or
``'volume == 12.0 & temperature == 800'``.
Returns
-------
SnapshotContainer
"""
new = type(self)(list(self.index.names))
for idx in self.index.filter(expr):
new.insert(idx, **self._snapshots[idx])
return new
[docs]
def gather_container(self, comm: DummyMPI = world) -> SnapshotContainer:
"""Gather snapshots from all MPI ranks into this instance on all ranks.
Parameters
----------
comm : DummyMPI, optional
MPI communicator. Defaults to the global ``world`` communicator.
Returns
-------
SnapshotContainer
Self, now containing all snapshots from all ranks.
"""
id_list = gather_chunks(comm.gather(list(self.index)))
da_list = gather_chunks(comm.gather(self.data))
tmp_c = SnapshotContainer(self.index.names, id_list, da_list)
self._snapshots = tmp_c._snapshots
self.index = tmp_c.index
return self
[docs]
def insert(
self,
index: dict[str, float] | Iterable[float] | IndexPoint,
**data_entries: Any,
) -> IndexPoint:
"""Insert or update a single row.
Parameters
----------
index : dict, iterable, or IndexPoint
Index identifying the row. Accepted forms:
- mapping of field name → value
- iterable of values in ``index_names`` order
- an existing :class:`IndexPoint`
\*\*data_entries
Data fields to store. Keys must not overlap with index field names.
Returns
-------
IndexPoint
The coerced index that was inserted.
Raises
------
RuntimeError
If the index already exists and ``replace=False``.
ValueError
If a data key conflicts with an index field name.
"""
index = self.index.coerce(index, self._index_defaults)
if not self.replace and index in self.index:
msg = f"Index {index} is already in container."
raise RuntimeError(msg)
dct = {}
for data_name, entry in data_entries.items():
if data_name in self.index.names:
msg = f"Entry {data_name} forbidden, exists in index."
raise ValueError(msg)
dct[data_name] = entry
self._snapshots.setdefault(index, {}).update(dct)
return index
[docs]
def discard_indices(self, indices: Iterable[IndexPoint]) -> None:
"""Remove entries by index, silently ignoring missing ones.
Parameters
----------
indices : iterable of IndexPoint
Indices to remove.
"""
for index in indices:
self._snapshots.pop(index, None)
[docs]
def extend(
self,
other: SnapshotContainer,
fill_index: list | None = None,
) -> None:
"""Extend this instance from another SnapshotContainer."""
fill_index = fill_index or []
pairs = zip(other.index, other.data, strict=True)
for ind, data in pairs:
fill = list(fill_index).copy()
ind_dict = ind.todict()
new_ind = []
for name in self.index.names:
if name in ind_dict:
new_ind.append(ind_dict[name])
else:
new_ind.append(fill.pop(0))
self.insert(new_ind, **data)
[docs]
def extend_from_db(self, file: str | pathlib.Path) -> None:
"""Extend this instance from an ASE .db file.
Parameters
----------
file: str or Path-like
The database to read from.
Notes
-----
Correspondingly to `write()`, the database volume_per_atom key is
set to index `volume` in the resulting SnapshotContainer.
"""
_io.extend_from_db(self, file)
[docs]
def write(
self,
file: str | pathlib.Path,
*,
append: bool = True,
overwrite: bool = False,
) -> None:
"""Write the SnapshotContainer to an ASE database file.
The index is written to key_value_pairs and works with query and
plotting. "volume" is written to "volume_per_atom" to avoid conflict
with the ASE special key. For example, to plot the energy vs step for
different lambdas at a volume and temperature, use `ase db ... -p` like
>>> ase db database.db 'volume_per_atom=12.0,temperature=800' -s step -p lambda_:step,potential_energy
where `-s step` is added to plot it sorted on steps.
Parameters
----------
file: str or Path-like
Name of the file to write to
append: bool, default True
If True, appends to the existing database.
overwrite: bool, default False
If True, and `append=False`, overwrites an existing database.
"""
_io.write(self, file, append=append, overwrite=overwrite)
[docs]
@classmethod
def read(cls, file: str | pathlib.Path) -> SnapshotContainer:
"""Read in a SnapshotContainer from an ASE database.
The database must contain a __snapshot_container_index__ data entry.
Parameters
----------
file: str or Path-like
The database to read from.
Returns
-------
SnapshotContainer
"""
return _io.read(file, cls)
[docs]
@classmethod
def from_list(cls, list_of_atoms: list[Atoms]) -> SnapshotContainer:
"""Create a SnapshotContainer from a list of atoms.
The atoms must either have a .calc attribute, or .energy, .forces, .stress.
This is mainly a legacy method for training with MLPTrainer and mlippy.
Parameters
----------
list_of_atoms: list of Atoms
List of Atoms objects to create the SnapshotContainer from.
Returns
-------
SnapshotContainer
"""
index_list = []
data = []
for i, atoms in enumerate(list_of_atoms):
volume = atoms.get_volume() / len(atoms)
index_list.append((volume, 0, i))
if atoms.calc is None:
results = {}
if hasattr(atoms, "energy"):
results["energy"] = atoms.energy
if hasattr(atoms, "forces"):
results["forces"] = atoms.forces
if hasattr(atoms, "stress"):
ase_stress = mlip2ase_stress(atoms.stress, atoms.get_volume())
results["stress"] = ase_stress
atoms.calc = SinglePointCalculator(atoms, **results)
data.append({"atoms": atoms})
return cls(["volume", "seed", "step"], index_list, data)