"""
Utility functions to improve the efficiency of common
functions/workflows/calculations in ``doped``.
"""
import contextlib
import copy
import itertools
from collections import defaultdict
from collections.abc import Generator, Sequence
from fractions import Fraction
from functools import cached_property, lru_cache
from string import digits
from typing import TYPE_CHECKING, Any
import numpy as np
from numpy.typing import NDArray
from pymatgen.analysis.defects.generators import VacancyGenerator
from pymatgen.analysis.defects.utils import VoronoiPolyhedron, remove_collisions
from pymatgen.core.composition import Composition, DummySpecies
from pymatgen.core.lattice import Lattice
from pymatgen.core.periodic_table import Element, Species
from pymatgen.core.sites import PeriodicSite, Site
from pymatgen.core.structure import IStructure, Molecule, Structure
from pymatgen.core.structure_matcher import (
AbstractComparator,
ElementComparator,
FrameworkComparator,
StructureMatcher,
)
from pymatgen.io.vasp.sets import get_valid_magmom_struct
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer, SymmOp
from pymatgen.util.misc import is_np_dict_equal
from scipy.spatial import Voronoi
if TYPE_CHECKING:
from numpy.typing import ArrayLike
from doped.core import Vacancy
# Note: overrides of ``__eq__`` should also override ``__hash__`` (and vice versa), preserving the
# invariant that ``a == b`` implies ``hash(a) == hash(b)`` -- hashes may be coarser than equality
# (collisions are resolved by ``__eq__`` in ``dict``/``set``/``lru_cache`` lookups) but not finer.
# A perfectly eq-consistent hash is impossible for tolerance-based equality (not transitive), so the
# ``PeriodicSite``/``Structure`` hashes below are deliberately finer. Safe for caching (worst case a
# spurious miss & recompute) and for sets of bitwise-identical copies, but ``set``/``dict`` dedup of
# near-equal (eq-equal, not bitwise-identical) objects can silently fail -- dedup those with explicit
# ``==`` loops instead. These patches also make some mutable objects hashable: don't mutate hashed fields
# while keyed!
_TRANSLATE_REMOVE_CHARGE = str.maketrans("", "", digits + "+-.")
def _freeze(obj):
"""
Recursively convert ``obj`` (e.g. nested ``properties`` values) into a
hashable canonical form, whose equality mirrors ``__eq__`` (e.g. order-
insensitive dict equality) -- required for the eq/hash invariant.
Container images are type-tagged (tuples excepted; already-hashable
composites) so eq-unequal values of different types (``[1, 2]`` vs
``(1,2)``) don't collide into equal hashes (keeping hash equality a sound
proxy for equality in the ``_Structure__eq__`` hash-equal fast path).
"""
if isinstance(obj, dict): # frozenset of items; insertion-order independent (like dict eq)
return ("dict", frozenset((k, _freeze(v)) for k, v in obj.items()))
if isinstance(obj, list):
return ("list", tuple(_freeze(v) for v in obj))
if isinstance(obj, tuple):
return tuple(_freeze(v) for v in obj)
if isinstance(obj, set):
return ("set", frozenset(_freeze(v) for v in obj))
if isinstance(obj, np.ndarray):
return ("ndarray", obj.shape, obj.dtype.str, obj.tobytes())
# else assume it's already hashable (int, str, custom ...)
return obj
def _frozen_properties(properties: dict):
"""
Hashable image of a ``properties`` dict, for use in ``__hash__`` methods.
"""
try: # fast path: property values already hashable (floats/strings/tuples etc.)
return frozenset(properties.items())
except TypeError: # unhashable values (arrays/lists/dicts etc.); canonicalise with ``_freeze``
pass
try:
return _freeze(properties)
except Exception: # exotic unhashable values; repr fallback (worst case: benign cache misses)
return repr(sorted(properties.items(), key=repr))
def _properties_equal(props_1: dict, props_2: dict) -> bool:
"""
``properties`` dict equality, tolerating ``np.ndarray`` values (plain dict
``==`` raises ``ValueError`` on these, e.g. ``selective_dynamics``).
"""
try:
return props_1 == props_2
except ValueError: # array-aware (slower) comparison, only used when needed:
return is_np_dict_equal(props_1, props_2)
# Composition overrides:
# First, define operand slots for the cached equality functions below. Those ``lru_cache``\s are keyed on
# the operands' exact int fingerprints -- they can't take the objects themselves, as ``lru_cache`` keys
# args via their ``hash``/``==``, recursing into the very ``__eq__`` being implemented -- but hash
# fingerprints aren't invertible, so on a cache miss the function body gets the actual operands from these
# slots instead (set immediately before each cached call, read at function entry; cache hits never enter
# the function body). Avoids unbounded hash->instance registry (memory leak) and associated per-call
# insertion overhead.
_composition_eq_pair: list = [None, None]
[docs]
@lru_cache(maxsize=int(1e5)) # maxsize on the order of 30 Mb
def doped_Composition_eq_func(self_hash, other_hash):
r"""
Updated equality function for |Composition| instances, which breaks early
for mismatches and also uses caching, making it orders of magnitude faster
than ``pymatgen``\s equality function.
"""
self_comp, other_comp = _composition_eq_pair # read at entry; only runs on cache misses
return fast_Composition_eq(self_comp, other_comp)
[docs]
def fast_Composition_eq(self, other):
"""
Fast equality function for |Composition| instances, breaking early for
mismatches.
"""
if len(self) != len(other): # skip type check here, already checked upstream in ``_Composition__eq__``
return False
return all(abs(amt - other[el]) <= type(self).amount_tolerance for el, amt in self.items())
def _composition_fingerprint(composition):
"""
Exact-stoichiometry fingerprint for |Composition| instances, used as a fast
equality key below (exact amounts distinguish e.g. Fe2O3 vs Fe3O2).
"""
return hash(frozenset(composition._data.items()))
def _Composition__eq__(self, other):
"""
Custom ``__eq__`` method for |Composition| instances, using a cached
equality function to speed up comparisons.
"""
if not isinstance(other, type(self) | dict):
return NotImplemented
if not isinstance(other, type(self)): # plain dicts have no ``_data``/hash-cache support;
return fast_Composition_eq(self, other) # compare directly
_composition_eq_pair[:] = (self, other) # slice-assign to mutate the module-level slots
return doped_Composition_eq_func(_composition_fingerprint(self), _composition_fingerprint(other))
[docs]
class Hashabledict(dict):
"""
Hashable ``dict`` subclass, allowing dictionaries to be used as cache keys.
"""
def __hash__(self):
"""
Make the dictionary hashable by recursively "freezing" into only
hashable built-ins (with ``_freeze``), then hash that.
Handles nested dicts, lists, sets, tuples, arrays etc.
"""
return hash(_freeze(self))
def _get_hashable_dict(d: dict) -> Hashabledict:
if isinstance(d, Hashabledict):
return d
if isinstance(d, dict):
return Hashabledict(d) # convert to hashable dict for caching purposes
return d
def _fast_dict_deepcopy_max_two_levels(d: dict) -> dict:
"""
Fast deepcopy of a dict with at most two levels of nested dicts (i.e. d ->
dict -> dict -> values).
Implemented to allow fast deep-copying of nested chemical potential dicts,
avoiding the overhead of ``deepcopy`` when looping over many chemical
potential dicts.
"""
return {
k: (
{
k2: (v2.copy() if isinstance(v2, dict) else v2) # final level, shallow copy sufficient
for k2, v2 in v1.items()
}
if isinstance(v1, dict)
else v1
)
for k, v1 in d.items()
}
@lru_cache(maxsize=int(1e5))
def _cached_Composition_init(comp_input):
return Composition(comp_input)
def _cache_ready_Composition_init(comp_input):
# copy on every call (incl. cache hits) so caller mutation can't corrupt the cached object:
return _cached_Composition_init(_get_hashable_dict(comp_input)).copy()
def _fast_get_composition_from_sites(sites: Sequence[Site], assume_full_occupancy: bool = False):
"""
Helper function to quickly get the composition of a collection of sites,
faster than initializing a |Structure| object.
Used in initial drafts of defect stenciling code, but replaced by faster
methods.
"""
elem_map: dict[Species, float] = defaultdict(float)
for site in sites:
if assume_full_occupancy:
elem_map[next(iter(site._species))] += 1
else:
for species, occu in site._species.items():
elem_map[species] += occu
return Composition(elem_map)
Composition.__eq__ = _Composition__eq__
def _parse_site_species_str(site: Site, wout_charge: bool = False) -> str:
"""
Get the species string of a :class:`~pymatgen.core.sites.Site`, with or
without charge information.
Much faster than direct ``str(Site)``. Note that this is faster without
caching.
Args:
site (Site):
:class:`~pymatgen.core.sites.Site` to get the species string of.
wout_charge (bool):
Whether to remove charge information from the species string.
Returns:
str:
Species string of the :class:`~pymatgen.core.sites.Site`, with or
without charge information.
"""
return _parse_species_str(site._species, wout_charge=wout_charge)
def _parse_species_str(sp_el: Species | Element, wout_charge: bool = False) -> str:
"""
Get the string representation of a
:class:`~pymatgen.core.periodic_table.Species` or
:class:`~pymatgen.core.periodic_table.Element`, with or without charge
information.
Much faster than direct ``str(sp_el)``.
Args:
sp_el (Species | Element):
:class:`~pymatgen.core.periodic_table.Species` or
:class:`~pymatgen.core.periodic_table.Element` for which to get the
string representation.
wout_charge (bool):
Whether to remove charge information from the species string.
Default is ``False``.
Returns:
str:
String representation of ``sp_el``, with or without charge
information.
"""
if isinstance(sp_el, str):
species_string = sp_el
elif isinstance(sp_el, Species) and wout_charge:
return sp_el.element.symbol # no charge info, return element sybmol
elif isinstance(sp_el, Element):
return sp_el.symbol
elif isinstance(sp_el, Composition | Species):
species_string = str(sp_el)
else:
species_string = str(Composition(sp_el))
# remove all digits, +, - or . from species string, if `wout_charge` is True
return species_string.translate(_TRANSLATE_REMOVE_CHARGE) if wout_charge else species_string
def _patch_method(cls, name, patch):
"""
Set ``cls.name`` to ``patch``, returning the method it replaces (which the
patches here delegate to).
The replaced method is also recorded on the patch, and taken from there if
present, so that re-importing this module (e.g. ``importlib.reload``,
``%autoreload``) returns ``pymatgen``'s original rather than the previous
patch -- which would then delegate to itself, and recurse infinitely.
"""
method = getattr(cls, name)
patch._doped_replaced = original = getattr(method, "_doped_replaced", method)
setattr(cls, name, patch)
return original
# Species overrides:
def _species__str__(self):
"""
Memoized ``Species.__str__`` (immutable objects); avoids heavy string
formatting in the many millions of ``Species.__hash__`` (= ``hash(str)``)
calls from |Composition| ``dict`` operations.
We memoize the string, not
the hash: changing ``__hash__`` values breaks pre-built ``Species``-keyed
``dict``s in ``pymatgen`` (e.g. ``bond_valence``), which store key hashes
at insertion.
"""
try:
return self.__dict__["_doped_str"]
except KeyError:
self.__dict__["_doped_str"] = string = _orig_species__str__(self)
return string
_orig_species__str__ = _patch_method(Species, "__str__", _species__str__)
def _noise_rounded_bytes(arr) -> bytes:
"""
Byte-image of a float array with noise below ``1e-10`` collapsed (rounded
to 10 dp; ``+ 0.0`` normalises ``-0.0``), so coordinates differing only by
float noise (e.g. from symmetry operations / Cartesian <-> fractional
round-trips) share a hash.
``1e-10`` sits safely below the ``__eq__`` tolerances (``atol=1e-8``) for
|Structure|/|PeriodicSite|, so hash-merged values are always still
eq-equal -- required for the ``_Structure__eq__`` hash-equality fast path
to stay sound.
"""
return (arr.round(10) + 0.0).tobytes()
def _species_info(species: dict) -> tuple:
# avoid ``str(el)`` (``Species.__str__``/format machinery); equal species give equal
# ``(symbol, oxi, spin, amount)`` tuples, incl. amounts to distinguish partial occupancies. ``spin`` is
# included because ``Species.__eq__`` compares it, so omitting it made structures differing only in
# spin (e.g. ferro- vs antiferro-magnetic orderings of one lattice) compare equal and collide in...
return tuple( # ...structure-keyed caches
(el.symbol, getattr(el, "_oxi_state", None), getattr(el, "_spin", None), amt)
for el, amt in species.items()
)
# PeriodicSite overrides:
def _periodic_site__hash__(self):
"""
Custom ``__hash__`` method for |PeriodicSite| instances; deliberately finer
than the tolerance-based ``__eq__`` (see module comment above), though
noise-rounded (``_noise_rounded_bytes``) so float-noise twins share a hash.
All fields ``__eq__`` compares exactly (incl. amounts & properties) are
hashed, which the ``doped`` |Structure| ``__eq__`` hash-equality fast path
relies on.
"""
base_hash_tuple = (
_species_info(self._species),
_noise_rounded_bytes(self.lattice.matrix),
self.lattice.pbc,
_noise_rounded_bytes(self.frac_coords),
)
if not self.properties:
return hash(base_hash_tuple)
return hash((*base_hash_tuple, _frozen_properties(self.properties)))
@lru_cache(maxsize=int(1e3))
def _cached_lattice_eq(matrix_bytes_1: bytes, matrix_bytes_2: bytes, pbc_1: tuple, pbc_2: tuple) -> bool:
"""
Cached lattice equality from raw matrix bytes and pbc; ``Lattice.__eq__``
(``np.allclose``) semantics but ~10x faster in hot containment loops, where
the same lattice pair repeats.
"""
if pbc_1 != pbc_2:
return False
if matrix_bytes_1 == matrix_bytes_2:
return True
return bool(
np.allclose(
np.frombuffer(matrix_bytes_1).reshape(3, 3), np.frombuffer(matrix_bytes_2).reshape(3, 3)
)
)
[docs]
def cache_ready_Site__eq__(self, other):
"""
Custom ``__eq__`` method for ``Site`` and |PeriodicSite| instances, using
a cached equality function to speed up comparisons.
"""
if self is other:
return True
needed_attrs = ("_species", "coords", "properties")
if not all(hasattr(other, attr) for attr in needed_attrs):
return NotImplemented
if not (
self._species == other._species # should always work fine (and is faster) if ``Site`` initialised
and ( # without ``skip_checks`` (default)
self.coords is other.coords # if coords are the same object
or cached_allclose(tuple(self.coords), tuple(other.coords), atol=type(self).position_atol)
)
):
return False
# lattice checked only for otherwise-matching sites (cheap checks above dominate hot containment loops)
other_lattice = getattr(other, "lattice", None) # plain ``Site``s have no lattice -> skip:
if (
other_lattice is not None
and self.lattice is not other_lattice
and not _cached_lattice_eq(
self.lattice.matrix.tobytes(),
other_lattice.matrix.tobytes(),
self.lattice.pbc,
other_lattice.pbc,
)
):
return False
return _properties_equal(self.properties, other.properties)
[docs]
@lru_cache(maxsize=int(2e3)) # maxsize on the order of 1 Mb
def cached_allclose(a: tuple, b: tuple, rtol: float = 1e-05, atol: float = 1e-08):
"""
Cached version of ``np.allclose``, taking tuples as inputs (so that they
are hashable and thus cacheable).
Implemented directly in Python (same semantics as ``np.allclose`` for
finite inputs), as the ``numpy`` machinery has ~10x higher overhead for the
small (len-3 coordinate) tuples this is mostly used for (and full
vectorised allclose should be used otherwise, anyway).
"""
return all(abs(x - y) <= atol + rtol * abs(y) for x, y in zip(a, b, strict=True))
def _to_unit_cell(lattice: Lattice, frac_coords: np.ndarray) -> np.ndarray:
"""
Fold fractional coordinates into the unit cell, along the periodic
directions of ``lattice`` (as in ``PeriodicSite.to_unit_cell``, but faster
with ``numpy`` vectorisation).
"""
return np.where(lattice.pbc, np.mod(frac_coords, 1), frac_coords)
def _structure_from_sites_and_coords(
lattice: Lattice,
sites: Sequence[PeriodicSite],
frac_coords: np.ndarray,
keep_labels: bool = False,
charge: float | None = None,
properties: dict | None = None,
) -> Structure:
"""
Build a |Structure| with the given ``lattice`` and ``frac_coords``, taking
species and site properties from ``sites``.
One construction, rather than the site-by-site rebuild of
``Structure.from_sites`` (faster).
"""
property_keys = dict.fromkeys(key for site in sites for key in site.properties)
return Structure(
lattice,
[site.species for site in sites],
frac_coords,
charge=charge,
site_properties={key: [site.properties.get(key) for site in sites] for key in property_keys}
or None,
labels=[site.label for site in sites] if keep_labels else None,
properties=properties,
)
_direct_site_attrs = frozenset({"_lattice", "_frac_coords", "_species", "_coords", "properties", "_label"})
def _fast_site__setattr__(self, attr: str, value):
"""
Fast path for the attributes set during ``(Periodic)Site`` initialisation,
falling back to the original ``Site.__setattr__`` for all others.
``Site.__setattr__`` is an override whose purpose is to deprecate setting
arbitrary attributes on a site (redirecting them into ``site.properties``),
so it tests every assignment against several cases. The six attributes in
``_direct_site_attrs``, set by ``(Periodic)Site.__init__``, all match its
first case, which simply delegates to ``super().__setattr__`` (which
resolves to a plain ``object.__setattr__``). Calling it directly is thus
exactly equivalent, skipping the string tests and the ``super()`` lookup;
worth doing as it runs six times per site creation, which dominates cost
for most structure manipulations. Any other attribute falls through to the
original, so the deprecation is unaffected.
"""
if attr in _direct_site_attrs:
object.__setattr__(self, attr, value)
else:
_orig_site__setattr__(self, attr, value)
PeriodicSite.__eq__ = cache_ready_Site__eq__
Site.__eq__ = cache_ready_Site__eq__
PeriodicSite.__hash__ = _periodic_site__hash__
_orig_site__setattr__ = _patch_method(Site, "__setattr__", _fast_site__setattr__)
# Lattice overrides:
# (note: memoizing ``Lattice.__hash__`` was tested and found to give negligible speedup (~0.1% of
# ``DefectsGenerator`` runtime), as ``pymatgen`` already caches the lengths/angles used in its hash)
@lru_cache(maxsize=int(1e4)) # maxsize on the order of 20 Mb for typical use cases
def _cached_get_all_distances(self: Lattice, frac_coords1: tuple, frac_coords2: tuple):
return _orig_lattice_get_all_distances(self, np.array(frac_coords1), np.array(frac_coords2))
[docs]
def array_to_tuple(array: "ArrayLike | tuple") -> tuple:
"""
Convert an array-like input to tuple.
"""
array = np.array(array)
if array.ndim == 1:
return tuple(array)
return tuple(map(tuple, array))
[docs]
def get_all_distances(
self,
frac_coords1: "ArrayLike",
frac_coords2: "ArrayLike",
) -> NDArray[np.floating[Any]]:
"""
Get the distances between two lists of coordinates taking into account
periodic boundary conditions and the lattice.
See :meth:`~pymatgen.core.lattice.Lattice.get_all_distances`.
"""
return _cached_get_all_distances(
self, array_to_tuple(frac_coords1), array_to_tuple(frac_coords2)
).copy()
_orig_lattice_get_all_distances = _patch_method(Lattice, "get_all_distances", get_all_distances)
# Structure overrides:
def _structure__hash__(self):
"""
Custom ``__hash__`` method for |Structure| instances; deliberately finer
than the tolerance-based, site-order-independent ``__eq__`` (see module
comment above), though noise-rounded (``_noise_rounded_bytes``) so float-
noise twins share a hash.
Includes the full lattice matrix (the coarser ``Lattice.__hash__`` can't
distinguish rotated/reflected settings) and properties, so every field
``__eq__`` compares exactly contributes to the hash -- which the ``__eq__``
hash-equality fast path relies on.
"""
lattice_info = (_noise_rounded_bytes(self.lattice.matrix), self.lattice.pbc)
sites_info = tuple(
(
_species_info(site._species),
_frozen_properties(site.properties) if site.properties else None,
)
for site in self._sites
)
coords_info = _noise_rounded_bytes(self.frac_coords) # vectorised coord rounding, rather than per-site
if not self.properties:
return hash((lattice_info, coords_info, sites_info))
return hash((lattice_info, coords_info, sites_info, _frozen_properties(self.properties)))
[docs]
@contextlib.contextmanager
def cache_species(structure_cls):
"""
Context manager that makes ``Structure.species`` a cached property, which
significantly speeds up ``pydefect`` eigenvalue parsing in large structures
(due to repeated use of ``Structure.indices_from_symbol``.
"""
original_species = structure_cls.species
try:
cached = cached_property(original_species.fget)
cached.__set_name__(structure_cls, "species") # Explicit initialization
structure_cls.species = cached
yield
finally:
structure_cls.species = original_species
[docs]
def doped_Structure__eq__(self, other: IStructure) -> bool:
"""
Copied from ``pymatgen``, but updated to break early once a mis-matching
site is found, to speed up structure matching by ~2x.
"""
# skip matching object type check here, as already checked upstream in ``_Structure__eq__``
if other is self:
return True
if len(self) != len(other):
return False
if self.lattice != other.lattice:
return False
if not _properties_equal(self.properties, other.properties):
return False
for site in self: # noqa: SIM110
if site not in other:
return False # break early!
return True
# operand slots for the cached structure equality function; see ``_composition_eq_pair`` comment above:
_structure_eq_pair: list = [None, None]
[docs]
@lru_cache(maxsize=int(1e4))
def cached_Structure_eq_func(self_hash, other_hash):
"""
Cached equality function for |Structure| instances.
"""
self_struct, other_struct = _structure_eq_pair # read at entry; only runs on cache misses
return doped_Structure__eq__(self_struct, other_struct)
def _Structure__eq__(self, other):
"""
Custom ``__eq__`` method for |Structure|/``IStructure`` instances, using
both caching and an updated, faster equality function to speed up
comparisons.
"""
if self is other:
return True # identity fast-path, avoids hashing both structures below
needed_attrs = ("lattice", "sites", "properties")
if not all(hasattr(other, attr) for attr in needed_attrs):
return NotImplemented
self_hash = _structure__hash__(self)
other_hash = _structure__hash__(other)
if self_hash == other_hash:
return True
_structure_eq_pair[:] = (self, other) # slice-assign to mutate the module-level slots
return cached_Structure_eq_func(self_hash, other_hash)
def _structure__deepcopy__(self, memo):
"""
Fast ``__deepcopy__`` for |Structure|: shallow ``.copy()``, then deep-copy
only the mutable ``properties`` dicts (structure- and site-level) so the
copy shares no state with the original.
"""
new_structure = self.copy()
new_structure.properties = copy.deepcopy(self.properties, memo)
for new_site, site in zip(new_structure, self, strict=True):
new_site.properties = copy.deepcopy(site.properties, memo)
return new_structure
IStructure.__eq__ = _Structure__eq__
IStructure.__hash__ = _structure__hash__
Structure.__eq__ = _Structure__eq__
Structure.__hash__ = _structure__hash__
Structure.__deepcopy__ = _structure__deepcopy__
# Molecule overrides:
def _DopedMolecule__hash__(self):
"""
``__hash__`` for (mutable, unhashable-by-default) ``Molecule``, using only
fields compared _exactly_ by ``Molecule.__eq__`` (composition, charge,
spin) -- not tolerance-compared coordinates -- so eq-equal molecules always
share a hash (sets/dicts/caches dedup correctly) and collisions
(e.g. conformers) are resolved by ``__eq__``.
"""
return hash((self.composition, self.charge, self.spin_multiplicity))
Molecule.__hash__ = _DopedMolecule__hash__
# SpacegroupAnalyzer overrides:
def _sga__hash__(self):
"""
Custom ``__hash__`` for ``SpacegroupAnalyzer`` (e.g. for cache keys);
coarser than the (default, identity) ``__eq__``, so invariant-safe.
"""
return hash((self._cell, self._symprec, self._angle_tol))
def _get_symmetry(self) -> tuple[NDArray, NDArray]:
"""
Get the symmetry operations associated with the structure, memoised per-
instance and ``get_sga`` already caches SGA construction by structure.
For non-magnetic cells, the rotations/translations are extracted from the
symmetry dataset already computed at ``SpacegroupAnalyzer`` init, rather
than re-calling the (expensive) ``spglib.get_symmetry`` function.
The cached arrays are frozen so caller mutation raises loudly rather than
silently corrupting the shared values.
"""
try:
return self._doped_symmetry
except AttributeError:
dataset = getattr(self, "_space_group_data", None)
if dataset is not None and len(self._cell) == 3 and hasattr(dataset, "rotations"):
# non-magnetic cell (no magmoms in ``self._cell``): reuse the init dataset, replicating
# ``SpacegroupAnalyzer._get_symmetry``'s cleanup of small/unity translation values;
rotations = dataset.rotations.copy() # copy, to not freeze the shared dataset arrays below
translations = np.array(
[[float(Fraction(c).limit_denominator(1000)) for c in row] for row in dataset.translations]
)
translations[np.abs(translations) == 1] = 0 # fractional translations of 1 -> 0
else: # magnetic cell, or no/unrecognised dataset; use original (``spglib``-calling) method
rotations, translations = _original_get_symmetry(self)
rotations.flags.writeable = False # freeze mutatable arrays
translations.flags.writeable = False
self._doped_symmetry = (rotations, translations)
return self._doped_symmetry
def _get_symmetry_operations(self, cartesian: bool = False) -> list[SymmOp]:
"""
Get the symmetry operations associated with the structure, memoised per-
instance (as with ``_get_symmetry``).
A fresh ``list`` is returned on every call (incl. hits) so callers can
freely modify it; the shared ``SymmOp`` objects themselves should not be
mutated in-place.
"""
cache = self.__dict__.setdefault("_doped_symmetry_operations", {})
if cartesian not in cache:
cache[cartesian] = _original_get_symmetry_operations(self, cartesian=cartesian)
return list(cache[cartesian])
SpacegroupAnalyzer.__hash__ = _sga__hash__
_original_get_symmetry = _patch_method(SpacegroupAnalyzer, "_get_symmetry", _get_symmetry)
_original_get_symmetry_operations = _patch_method(
SpacegroupAnalyzer, "get_symmetry_operations", _get_symmetry_operations
)
def _get_symbol(element: Element | Species, comparator: AbstractComparator | None = None) -> str:
"""
Convenience function to get the symbol of an ``Element`` or ``Species`` as
a string, with charge information included or excluded depending on the
choice of ``comparator``.
By default, the returned symbol does not include any charge / oxidation
state information. If ``comparator`` is provided and is not
``ElementComparator`` / ``FrameworkComparator``, then the ``str(element)``
representation is returned (which will include charge information if
``element`` is a ``Species``).
Args:
element (Element | Species):
``Element`` or ``Species`` to get the symbol of.
comparator (AbstractComparator | None):
Comparator to check if we should return the ``str(element)``
representation (which includes charge information if ``element`` is
a ``Species``), or just the element symbol (i.e.
``element.symbol``, or ``element.element.symbol`` if ``element`` is
a ``Species`` object) -- which is the case when ``comparator`` is
``None`` (default) or ``ElementComparator`` /
``FrameworkComparator``.
Returns:
str: Symbol of the element as a string.
"""
if (
comparator is not None
and not isinstance(comparator, ElementComparator | FrameworkComparator)
and isinstance(element, Species)
):
return str(element)
return element.symbol if isinstance(element, Element | DummySpecies) else element.element.symbol
[docs]
def get_element_indices(
structure: Structure,
elements: list[Element | Species | str] | None = None,
comparator: AbstractComparator | None = None,
) -> dict[str, list[int]]:
"""
Convenience function to generate a dictionary of ``{element: [indices]}``
for a given |Structure|, where ``indices`` are the indices of the sites in
the structure corresponding to the given ``elements`` (default is all
elements in the structure).
Args:
structure (|Structure|):
|Structure| to get the indices from.
elements (list[Element | Species | str] | None):
List of elements to get the indices of. If ``None`` (default), all
elements in the structure are used.
comparator (AbstractComparator | None):
Comparator to check if we should return the ``str(element)``
representation (which includes charge information if ``element`` is
a ``Species``), or just the element symbol (i.e.
``element.element.symbol``) -- which is the case when
``comparator`` is ``None`` (default) or ``ElementComparator`` /
``FrameworkComparator``.
Returns:
dict[str, list[int]]:
Dictionary of ``{element: [indices]}`` for the given ``elements``
in the structure.
"""
if elements is None:
elements = _fast_get_composition_from_sites(structure).elements
if not all(isinstance(element, str) for element in elements):
elements = [_get_symbol(element, comparator) for element in elements]
species = np.array([_get_symbol(site.specie, comparator) for site in structure])
return {element: np.where(species == element)[0].tolist() for element in elements}
[docs]
def get_element_min_max_bond_length_dict(structure: Structure, **sm_kwargs) -> dict:
r"""
Get a dictionary of ``{element: (min_bond_length, max_bond_length)}`` for a
given |Structure|, where ``min_bond_length`` and ``max_bond_length`` are
the minimum and maximum `smallest` interatomic bond lengths for each
element in the structure.
Args:
structure (|Structure|):
|Structure| to calculate bond lengths for.
**sm_kwargs:
Additional keyword arguments to pass to ``StructureMatcher()``.
Just used to check if ``comparator`` has been set here (if
``ElementComparator``/``FrameworkComparator`` used, then we use
``Element``\s rather than ``Species`` as the keys), or if
``ignored_species`` is set (in which case these species are
ignored when calculating bond lengths).
Returns:
dict: Dictionary of ``{element: (min_bond_length, max_bond_length)}``.
"""
comparator = sm_kwargs.get("comparator")
if len(structure) == 1:
structure *= 2 # need at least two sites to calculate bond lengths
elif len(structure) == 0: # edge case of a 'defect structure' in a primitive cell with 1 atom
return {}
# get the distance matrix broken down by species:
element_idx_dict = get_element_indices(structure, comparator=comparator)
ignored_indices = [
idx for elt in sm_kwargs.get("ignored_species", []) for idx in element_idx_dict.get(elt, [])
]
# Note: The repeated distance_matrix call here, particularly when this function is called twice in
# ``apply_s2_to_s1_transformation()``, can become expensive for large structures (due to the NxN
# distance calculation). In ``apply_s2_to_s1_transformation()``, where this function is called
# twice, we sub-select sites in the structure to run the min/max bond lengths test, for this reason
distance_matrix = structure.distance_matrix
np.fill_diagonal(distance_matrix, np.inf) # set diagonal to np.inf to ignore self-distances of 0
distance_matrix[:, ignored_indices] = np.inf # set ignored indices to np.inf to ignore these distances
distance_matrix[ignored_indices, :] = np.inf # set ignored indices to np.inf to ignore these distances
element_min_max_bond_length_dict = {
elt: np.array([np.min(structure.lattice.abc), np.max(structure.lattice.abc)])
for elt in element_idx_dict
} # default to min/max lattice vectors (for cases where there are no other matching non-ignored atoms)
for elt, site_indices in element_idx_dict.items():
element_dist_matrix = distance_matrix[:, site_indices] # (N_of_that_element, N_sites) matrix
if element_dist_matrix.size != 0:
min_interatomic_distances_per_atom = np.min(element_dist_matrix, axis=0) # min along columns
if np.min(min_interatomic_distances_per_atom) != np.inf: # other non-ignored matching atoms
element_min_max_bond_length_dict[elt] = np.array(
[
np.min(min_interatomic_distances_per_atom),
np.max(min_interatomic_distances_per_atom),
]
)
return element_min_max_bond_length_dict
[docs]
def get_dist_equiv_stol(dist: float, structure: Structure) -> float:
"""
Get the equivalent ``stol`` value for a given Cartesian distance (``dist``)
in a given |Structure|.
``stol`` is a site tolerance parameter used in ``pymatgen``
|StructureMatcher| functions, defined as the fraction of the average free
length per atom := ( V / Nsites ) ** (1/3).
Args:
dist (float): Cartesian distance in Å.
structure (|Structure|): |Structure| to calculate ``stol`` for.
Returns:
float: Equivalent ``stol`` value for the given distance.
"""
return dist / (structure.volume / max(len(structure), 1)) ** (1 / 3) # max to ensure no divide-by-zero
[docs]
def get_min_stol_for_s1_s2(struct1: Structure, struct2: Structure, **sm_kwargs) -> float:
"""
Get the minimum possible ``stol`` value which will give a match between
``struct1`` and ``struct2`` using |StructureMatcher|, based on the ranges
of per-element minimum interatomic distances in the two structures.
Args:
struct1 (|Structure|): Initial structure.
struct2 (|Structure|): Final structure.
**sm_kwargs:
Additional keyword arguments to pass to ``StructureMatcher()``.
Just used to check if ``ignored_species`` or ``comparator`` has
been set here.
Returns:
float:
Minimum ``stol`` value for a match between ``struct1`` and
``struct2``. If a direct match is detected (corresponding to min
``stol`` = 0, then ``1e-4`` is returned).
"""
s1_min_max_bond_length_dict = get_element_min_max_bond_length_dict(struct1, **sm_kwargs)
s2_min_max_bond_length_dict = get_element_min_max_bond_length_dict(struct2, **sm_kwargs)
common_elts = set(s1_min_max_bond_length_dict.keys()) & set(s2_min_max_bond_length_dict.keys())
if not common_elts: # try without oxidation states
struct1_wout_oxi = struct1.copy()
struct2_wout_oxi = struct2.copy()
struct1_wout_oxi.remove_oxidation_states()
struct2_wout_oxi.remove_oxidation_states()
s1_min_max_bond_length_dict = get_element_min_max_bond_length_dict(struct1_wout_oxi, **sm_kwargs)
s2_min_max_bond_length_dict = get_element_min_max_bond_length_dict(struct2_wout_oxi, **sm_kwargs)
common_elts = set(s1_min_max_bond_length_dict.keys()) & set(s2_min_max_bond_length_dict.keys())
min_min_dist_change = 1e-4
with contextlib.suppress(Exception):
min_min_dist_change = (
max(
{
elt: max(np.abs(s1_min_max_bond_length_dict[elt] - s2_min_max_bond_length_dict[elt]))
for elt in common_elts
if elt not in sm_kwargs.get("ignored_species", [])
}.values()
)
/ 2
) # divide by two as sites may have displaced toward each other (so Δbond-length = 2*Δsite)
return max(get_dist_equiv_stol(min_min_dist_change, struct1), 1e-4)
def _sm_get_atomic_disps(sm: StructureMatcher, struct1: Structure, struct2: Structure):
"""
Get the root-mean-square displacement `and atomic displacements` between
two structures, normalized by the mean free length per atom:
``(Vol/Nsites)^(1/3)``.
These values are not directly returned by |StructureMatcher| methods.
This function replicates ``StructureMatcher.get_rms_dist()``, but changes
the return value from ``match[0], max(match[1])`` to ``match[0], match[1]``
to allow further analysis of displacements. Mainly intended for use by
|ShakeNBreak|.
Args:
sm (|StructureMatcher|): ``pymatgen`` |StructureMatcher| object.
struct1 (|Structure|): Initial structure.
struct2 (|Structure|): Final structure.
Returns:
tuple:
- float: Normalised RMS displacement between the two structures.
- np.ndarray: Normalised displacements between the two structures.
or ``None`` if no match is found.
"""
struct1, struct2 = sm._process_species([struct1, struct2])
struct1, struct2, fu, s1_supercell = sm._preprocess(struct1, struct2)
match = sm._match(struct1, struct2, fu, s1_supercell, use_rms=True, break_on_match=False)
return None if match is None else (match[0], match[1])
[docs]
def StructureMatcher_scan_stol(
struct1: Structure,
struct2: Structure,
func_name: str = "get_s2_like_s1",
min_stol: float | None = None,
max_stol: float = 0.3,
stol_factor: float = 0.5,
**sm_kwargs,
):
r"""
Utility function to scan through a range of ``stol`` values for
|StructureMatcher| until a match is found between ``struct1`` and
``struct2`` (i.e. ``StructureMatcher.{func_name}`` returns a result).
The ``StructureMatcher.match()`` function (used in most
|StructureMatcher| methods) speed is heavily dependent on ``stol``, with
smaller values being faster, so we can speed up evaluation by starting with
small values and increasing until a match is found (especially with the
``doped`` efficiency tools which implement caching (and other improvements)
to ensure no redundant work here).
Note that ``ElementComparator()`` is used by default here! (So sites with
different species but the same element (e.g. "S2-" & "S0+") will be
considered match-able). This can be controlled with
``sm_kwargs['comparator']``.
Note: If you know reduction to primitive cells is not possible/needed, then
setting ``primitive_cell=False`` in ``sm_kwargs`` can significantly speed
up matching here (by avoiding expensive reduction to primitive cells for
large structures).
Args:
struct1 (|Structure|): ``struct1`` for ``StructureMatcher.match()``.
struct2 (|Structure|): ``struct2`` for ``StructureMatcher.match()``.
func_name (str):
The name of the |StructureMatcher| method to return the result
of ``StructureMatcher.{func_name}(struct1, struct2)`` for, such
as:
- "get_s2_like_s1" (default)
- "get_rms_dist"
- "fit"
- "fit_anonymous"
- "get_rms_anonymous"
min_stol (float):
Minimum ``stol`` value to try. Default is to use ``doped``\s
``get_min_stol_for_s1_s2()`` function to estimate the minimum
``stol`` necessary, and start with 2x this value to achieve fast
structure-matching in most cases.
max_stol (float):
Maximum ``stol`` value to try. Default: 0.3 (matching
|StructureMatcher| default).
stol_factor (float):
Fractional increment to increase ``stol`` by each time (when a
match is not found). Default value of 0.5 increases ``stol`` by 50%
each time.
**sm_kwargs:
Additional keyword arguments to pass to ``StructureMatcher()``.
Returns:
Result of ``StructureMatcher.{func_name}(struct1, struct2)`` or
``None`` if no match is found.
"""
# use doped efficiency tools to make structure-matching as fast as possible:
if func_name == "_get_atomic_disps": # only used by SnB; add get-atomic-disps method:
StructureMatcher._get_atomic_disps = _sm_get_atomic_disps
if "comparator" not in sm_kwargs:
sm_kwargs["comparator"] = ElementComparator()
if min_stol is None:
min_stol = get_min_stol_for_s1_s2(struct1, struct2, **sm_kwargs) * 2
# here we cycle through a range of stols, because we just need to find the closest match so we could
# use a high ``stol`` from the start and it would give correct result, but higher ``stol``\s take
# much longer to run as it cycles through multiple possible matches. So we start with a low ``stol``
# and break once a match is found:
stol = min_stol
while stol <= max_stol:
if user_stol := sm_kwargs.pop("stol", False): # first run, try using user-provided stol first:
sm_full_user_custom = StructureMatcher(stol=user_stol, **sm_kwargs)
result = getattr(sm_full_user_custom, func_name)(struct1, struct2)
if result is not None:
return result
sm = StructureMatcher(stol=stol, **sm_kwargs)
result = getattr(sm, func_name)(struct1, struct2)
if (
result is not None
and result is not False
and not (isinstance(result, tuple) and result[0] is None) # for ``get_rms_anonymous()``
):
return result
if stol == max_stol: # failed with max_stol; break
break
stol = min(stol * (1 + stol_factor), max_stol)
# Note: this function could possibly be sped up if ``StructureMatcher._match()`` was updated to
# return the guessed ``best_match`` value (even if larger than ``stol``), which will always be
# >= the best possible match it seems, and then using this to determine the next ``stol`` value
# to trial. Seems like it could give a ~50% speedup in some cases? Not clear though,
# as once you're getting a reasonable guessed value out, the trial ``stol`` should be pretty
# close to the necessary value anyway.
return None
[docs]
class DopedTopographyAnalyzer:
"""
This is a modified version of the ``pymatgen-analysis-defects``
``TopographyAnalyzer`` class, to lean down the input options and make
initialisation far more efficient (~2 orders of magnitude faster).
The original code was written by Danny Broberg and colleagues
(10.1016/j.cpc.2018.01.004), which was then added to ``pymatgen`` before
being cut.
"""
def __init__(
self,
structure: Structure,
image_tol: float = 0.0001,
max_cell_range: int = 1,
constrained_c_frac: float = 0.5,
thickness: float = 0.5,
) -> None:
"""
Args:
structure (|Structure|):
|Structure| to analyse.
image_tol (float):
A tolerance distance for the analysis, used to determine if
sites are periodic images of each other. Default (of 1e-4) is
usually fine.
max_cell_range (int):
This is the range of periodic images to construct the Voronoi
tessellation. A value of 1 means that we include all points
from ``(x +- 1, y +- 1, z+- 1)`` in the Voronoi construction.
This is because the Voronoi polyhedra extend beyond the
standard unit cell because of PBC. Typically, the default value
of 1 works fine for most structures and is fast. But for very
small unit cells with high symmetry, this may need to be
increased to 2 or higher. If there are < 5 atoms in the input
structure and ``max_cell_range`` is 1, this will automatically
be increased to 2.
constrained_c_frac (float):
Constrain the region where topology analysis is performed.
Only sites with ``z`` fractional coordinates between
``constrained_c_frac +/- thickness`` are considered. Default of
0.5 (with ``thickness`` of 0.5) includes all sites in the unit
cell.
thickness (float):
Constrain the region where topology analysis is performed.
Only sites with ``z`` fractional coordinates between
``constrained_c_frac +/- thickness`` are considered. Default of
0.5 (with ``thickness`` of 0.5) includes all sites in the unit
cell.
"""
# if input cell is very small (< 5 atoms) and max cell range is 1 (default), bump to 2 for
# accurate Voronoi tessellation:
if len(structure) < 5 and max_cell_range == 1:
max_cell_range = 2
self.structure = structure.copy()
self.structure.remove_oxidation_states()
constrained_sites = []
for _i, site in enumerate(self.structure):
if (
site.frac_coords[2] >= constrained_c_frac - thickness
and site.frac_coords[2] <= constrained_c_frac + thickness
):
constrained_sites.append(site)
constrained_struct = Structure.from_sites(sites=constrained_sites)
lattice = constrained_struct.lattice
coords = []
cell_range = list(range(-max_cell_range, max_cell_range + 1))
for shift in itertools.product(cell_range, cell_range, cell_range):
for site in constrained_struct.sites:
shifted = site.frac_coords + shift
coords.append(lattice.get_cartesian_coords(shifted))
# Perform the voronoi tessellation.
voro = Voronoi(coords)
node_points_map: defaultdict[int, set] = defaultdict(set)
for pts, vs in voro.ridge_dict.items():
for v in vs:
node_points_map[v].update(pts)
vnodes: list[VoronoiPolyhedron] = []
def get_mapping(vnodes, poly: VoronoiPolyhedron):
"""
Check if a Voronoi Polyhedron is a periodic image of one of the
existing polyhedra.
Modified to avoid expensive ``np.allclose()`` calls.
"""
if not vnodes:
return None
distance_matrix = lattice.get_all_distances([v.frac_coords for v in vnodes], poly.frac_coords)
if np.any(distance_matrix < image_tol):
for v in vnodes:
if v.is_image(poly, image_tol):
return v
return None
# Filter all the voronoi polyhedra so that we only consider those
# which are within the unit cell:
for i, vertex in enumerate(voro.vertices):
if i == 0:
continue
fcoord = lattice.get_fractional_coords(vertex)
if np.all([-image_tol <= c < 1 + image_tol for c in fcoord]):
poly = VoronoiPolyhedron(lattice, fcoord, node_points_map[i], coords, i)
if get_mapping(vnodes, poly) is None:
vnodes.append(poly)
self.coords = coords
self.vnodes = vnodes
[docs]
def get_voronoi_nodes(structure: Structure) -> list[PeriodicSite]:
"""
Get the Voronoi nodes of a ``pymatgen`` |Structure|.
Maximises efficiency by mapping down to the primitive cell, doing Voronoi
analysis (with the efficient ``DopedTopographyAnalyzer`` class), and then
mapping back to the original structure (typically a supercell).
Args:
structure (|Structure|):
``pymatgen`` |Structure| object.
Returns:
list[PeriodicSite]:
List of |PeriodicSite| objects representing the Voronoi nodes.
"""
# fresh ``list`` on every call (incl. cache hits), so caller mutation cannot corrupt the cache:
return list(_hashable_get_voronoi_nodes(structure))
@lru_cache(maxsize=int(1e2))
def _hashable_get_voronoi_nodes(structure: Structure) -> list[PeriodicSite]:
from doped.utils.symmetry import _get_orientation_preserving_primitive, doped_cluster_frac_coords
# map all sites to the unit cell; 0 ≤ xyz < 1.
structure = Structure.from_sites(structure, to_unit_cell=True)
# get Voronoi nodes in the primitive structure and then map back to the supercell; using the
# orientation-preserving primitive, so that mapping back is just the integer supercell matrix expansion
# (no structure-matching or origin-shift handling needed):
try:
prim_and_matrix = _get_orientation_preserving_primitive(structure)
except ValueError: # non-integer supercell matrix (no other primitive); analyse structure directly
prim_and_matrix = None
prim_structure, supercell_matrix = prim_and_matrix or (structure, np.eye(3, dtype=int))
top_analyzer = DopedTopographyAnalyzer(prim_structure)
voronoi_coords = [v.frac_coords for v in top_analyzer.vnodes]
# remove nodes less than 0.5 Å from sites in the structure
voronoi_coords = remove_collisions(voronoi_coords, structure=prim_structure, min_dist=0.5)
# cluster nodes within 0.2 Å of each other:
prim_vnodes = doped_cluster_frac_coords(voronoi_coords, prim_structure, tol=0.2)
voronoi_struct = Structure.from_sites(
[PeriodicSite("X", fpos, prim_structure.lattice, skip_checks=True) for fpos in prim_vnodes]
) # Structure with Voronoi nodes as sites
voronoi_struct.make_supercell(supercell_matrix) # Map back to the supercell
return voronoi_struct.sites.copy() # copy() to help avoid mutability issues with cached outputs
[docs]
class DopedVacancyGenerator(VacancyGenerator):
"""
Vacancy defects generator, subclassed from ``pymatgen-analysis-defects`` to
improve efficiency (particularly when handling defect complexes).
"""
[docs]
def generate(
self,
structure: Structure,
rm_species: set[str | Species] | list[str | Species] | None = None,
**kwargs,
) -> Generator["Vacancy", None, None]:
"""
Generate vacancy defects.
Args:
structure (|Structure|):
The structure to generate vacancy defects in.
rm_species (set[str | Species] | list[str | Species] | None):
List/set of species to be removed (i.e. to consider for vacancy
generation). If ``None``, considers all species.
**kwargs:
Additional keyword arguments for the ``Vacancy`` constructor.
Returns:
Generator[Vacancy, None, None]:
Generator that yields a list of ``Vacancy`` objects.
"""
from doped.core import Vacancy
from doped.utils.symmetry import get_sga
# core difference is the removal of unnecessary `remove_oxidation_states` calls
structure = get_valid_magmom_struct(structure)
all_species = {elt.symbol for elt in structure.composition.elements}
rm_species = all_species if rm_species is None else {*map(str, rm_species)}
if not set(rm_species).issubset(all_species):
raise ValueError(
f"rm_species ({rm_species}) must be a subset of the structure's species ({all_species})."
)
sga = get_sga(structure)
sym_struct = sga.get_symmetrized_structure()
for site_group in sym_struct.equivalent_sites:
site = site_group[0]
if site.specie.symbol in rm_species:
yield Vacancy(
structure=structure, # note that we no longer remove oxi states here! or in get_sga
site=site,
equivalent_sites=site_group,
**kwargs,
)