"""
This module contains the DistributedStructs class, which is used to handle distributed structure dictionaries.
This source code is licensed under the BSD-3-Clause license found in the
LICENSE file in the root directory of this source tree.
"""
from __future__ import annotations
__author__ = ["Yi Yang", "Rithwik Tom"]
__email__ = "yiy5@andrew.cmu.edu"
__group__ = "https://www.noamarom.com/"
import logging
import os
from itertools import chain
from pathlib import Path
import numpy as np
from ase.atoms import Atoms
import gnrs.output as gout
import gnrs.parallel as gp
from gnrs.parallel.io import decode_struct, encode_struct
logger = logging.getLogger("DistributedStructs")
[docs]
class DistributedStructs:
"""
Contains functions for handling distributed structure dictionaries.
"""
[docs]
def __init__(self, structs: dict):
"""
Initialize with a dictionary of structures.
Args:
structs: Dictionary mapping structure names to ASE Atoms objects
Raises:
ValueError: If structs is not a dictionary
"""
self.structs = structs
self.logger = logger
# Names completed by the current task, see ``done``
self._done: set[str] = set()
[docs]
def get_num_structs(self) -> int:
"""
Get the total number of structures in a distributed
structures dictionary.
Returns:
Total number of structures across all ranks
"""
# Handle structs = None case
if self.structs is not None:
num_each_rank = len(self.structs)
else:
num_each_rank = 0
total_num_list = gp.comm.gather(num_each_rank, root=0)
if gp.is_master:
total_num = sum(total_num_list)
logger.debug(f"xtal distribution across cores: {total_num_list}")
else:
total_num = None
total_num = gp.comm.bcast(total_num, root=0)
return total_num
[docs]
def find_matches(self, target: Atoms, settings: dict | None = None) -> list:
"""
Runs pymatgen duplicate checks on a distributed struct dictionary.
Args:
target: Target structure to be matched
settings: Settings for pymatgen StructureMatcher
Returns:
List of matching structure IDs
"""
from pymatgen.analysis.structure_matcher import StructureMatcher
from pymatgen.io.ase import AseAtomsAdaptor
pmg_target = AseAtomsAdaptor.get_structure(target)
if settings is None:
settings = {"stol": 0.5, "ltol": 0.5, "angle_tol": 10}
matcher = StructureMatcher(**settings)
match_list = []
for name, xtal in self.structs.items():
pmg_xtal = AseAtomsAdaptor.get_structure(xtal)
if matcher.fit(pmg_target, pmg_xtal):
match_list.append(name)
# Combine match list and flatten
match_list = gp.comm.allgather(match_list)
match_list = [item for sublist in match_list for item in sublist]
logger.info(f"Matched with {len(match_list)} structures")
logger.debug(f"Matched structures = {match_list}")
return match_list
[docs]
def collect_property(self, prpty: str, ptype: str = "info") -> list:
"""
Collects the property of all the structures into a list.
Args:
prpty: Property name to collect
ptype: Property type, either 'info' or 'method'
Returns:
List of property values on master rank, None on other ranks
"""
# Construct property list for each rank
prpty_list = []
for xtal in self.structs.values():
if ptype == "method":
prop = getattr(xtal, prpty)()
elif ptype == "info":
prop = xtal.info.get(prpty)
prpty_list.append(prop)
# Combine
prpty_list = gp.comm.gather(prpty_list)
if gp.is_master:
prpty_list = [item for sublist in prpty_list for item in sublist]
return prpty_list
[docs]
def get_statistics(self, prpty: str, ptype: str = "info") -> dict:
"""
Gets the statistics on a property of interest.
Args:
prpty: Property name to analyze
ptype: Get property from either 'info' or 'method'
Returns:
Dictionary with statistics on master rank, None on other ranks
"""
prpty_list = self.collect_property(prpty, ptype)
if not gp.is_master:
return None
prpty_array = np.array(prpty_list)
stats = {
"Minimum": np.min(prpty_array),
"Maximum": np.max(prpty_array),
"Average": np.average(prpty_array),
"Std": np.std(prpty_array),
}
return stats
[docs]
def find_spg(self, tol: float = 0.001) -> None:
"""
Finds the space group of all structures.
Space group number is stored in info["spg"]
Args:
tol: Tolerance for symmetry finding
"""
from ase.spacegroup.spacegroup import get_spacegroup
for struct in self.structs.values():
struct.info["spg"] = get_spacegroup(struct, symprec=tol).no
@property
def done(self) -> set[str]:
"""
Names of the structures already completed by the current task, i.e.
restored from a checkpoint or checkpointed on this rank.
"""
return self._done
[docs]
def checkpoint_save(self, path: str, name: str, xtal: Atoms) -> None:
"""
Appends one completed structure to this rank's checkpoint log.
Does not communicate, so ranks can call it independently.
The structure is written as one line in the same form used inside
structures.json (see ``gnrs.parallel.io.encode_struct``). A job
killed mid-write can therefore damage only the last line.
Args:
path: Directory for this rank's checkpoint log
name: Structure ID
xtal: The completed structure. Need not belong to this rank's
pool, e.g. in serial DFT mode rank 0 computes copies of every
rank's structures.
"""
if name in self._done:
return
with open(os.path.join(path, f"{gp.rank}.ckpt"), "ab") as chk:
# Leading newline terminates a line cut short by a job kill, so
# only the damaged entry is lost; the reader skips blank lines
chk.write(b"\n" + encode_struct(name, xtal).encode())
self._done.add(name)
[docs]
@staticmethod
def checkpoint_clear(path: str) -> int:
"""
Removes checkpoint files left under a task's calc directory.
Call on one rank only; does not communicate.
Args:
path: Directory containing rank_* checkpoint folders
Returns:
Number of files removed.
"""
n_removed = 0
for save_file in Path(path).glob("rank_*/*.ckpt"):
if save_file.is_file():
save_file.unlink(missing_ok=True)
n_removed += 1
return n_removed
@staticmethod
def _read_checkpoint(checkpoint: Path) -> tuple[dict, int]:
"""
Read the completed structures of one checkpoint log.
Args:
checkpoint: ``.ckpt`` log, one completed structure per line.
Returns:
Completed structures in the file and the number of damaged lines
skipped.
"""
structs, n_bad = {}, 0
with open(checkpoint) as chk:
for line in chk:
if not line.strip():
continue
try:
name, xtal = decode_struct(line)
except (ValueError, KeyError, TypeError):
n_bad += 1
continue
structs[name] = xtal
return structs, n_bad
[docs]
def checkpoint_load(self, path: str) -> int:
"""
Merges checkpoints of an interrupted run into the current structures
and spreads the remaining work evenly across ranks. Unlike save, load
is a collective and blocking operation.
A checkpointed copy replaces the current copy of the same structure
and is marked as done (see ``done``); structures that were never
checkpointed (e.g. held by a rank that was killed before it finished
its first calculation) are kept as they are. Checkpointed structures
that are not in the current pool are ignored. Damaged lines and
unreadable files are skipped with a warning; the affected structures
are simply recomputed.
Args:
path: Directory containing rank_* checkpoint folders
Returns:
Number of distinct structures restored from checkpoints.
"""
checkpoints = None
if gp.is_master:
checkpoints = sorted(Path(path).glob("rank_*/*.ckpt"))
checkpoints = gp.comm.bcast(checkpoints, root=0)
self._done = set()
if not checkpoints:
# Nothing to merge: keep every rank's structures where they are
return 0
# Each rank reads its share of the files
restored, problems = [], []
for idx in range(gp.rank, len(checkpoints), gp.size):
checkpoint = checkpoints[idx]
try:
structs, n_bad = self._read_checkpoint(checkpoint)
except (OSError, ValueError, KeyError, TypeError) as exc:
problems.append(f"{checkpoint} could not be read ({exc})")
continue
if n_bad:
problems.append(f"{checkpoint}: {n_bad} damaged line(s) skipped")
restored.extend(structs.items())
restored = gp.comm.gather(restored, root=0)
problems = gp.comm.gather(problems, root=0)
current = gp.comm.gather(self.structs or {}, root=0)
chunks, done = None, None
if gp.is_master:
for problem in chain.from_iterable(problems):
self.logger.error(f"Checkpoint {problem}")
gout.emit(
f"WARNING: Checkpoint {problem}. The affected structures "
"will be recomputed."
)
combined = {n: x for structs in current for n, x in structs.items()}
done = set()
for name, xtal in chain.from_iterable(restored):
if name in combined:
combined[name] = xtal
done.add(name)
self.logger.debug(f"Read {len(done)} structures from checkpoints")
chunks = self._balanced_chunks(combined, done, gp.size)
self.structs = gp.comm.scatter(chunks, root=0)
# Restored structures are done and already on disk; only new results
# get logged
done = gp.comm.bcast(done, root=0)
self._done = done & set(self.structs)
return len(done)
@staticmethod
def _balanced_chunks(structs: dict, done: set[str], size: int) -> list[dict]:
"""
Split structures into one chunk per rank so that every rank gets an
equal share of the remaining work, not just of the structures.
Args:
structs: All structures.
done: Names of the structures that are already completed.
size: Number of ranks.
Returns:
One structure dictionary per rank.
"""
remaining = [(n, x) for n, x in structs.items() if n not in done]
finished = [(n, x) for n, x in structs.items() if n in done]
return [dict(remaining[r::size] + finished[r::size]) for r in range(size)]
[docs]
def redistribute(self) -> None:
"""
Redistribute structures such that all ranks have almost
equal number of structures. Helpful for balancing load
across cores
"""
allstructs = gp.comm.gather(self.structs, root=0)
chunks = None
if gp.is_master:
combined = {n: x for structs in allstructs for n, x in structs.items()}
chunks = self._balanced_chunks(combined, set(), gp.size)
self.structs = gp.comm.scatter(chunks, root=0)