"""Explicit corrected periodic-image adapter for descriptor evaluation (experimental, opt-in).
Why
---
DScribe's periodic SOAP (used by ED-SOAP through decaf) builds periodic copies of the *raw*
positions ``m`` lattice vectors away, with ``m_a`` in ``[-n_a, n_a]`` and
``n_a = ceil(R / h_a)`` (``R`` = cutoff + Gaussian padding, ``h_a`` = perpendicular cell
height). It never wraps positions. When coordinates are unwrapped (e.g. an MD trajectory
that lets atoms drift several cells), a neighbour whose required lattice image lies outside
that copy range is silently dropped and contributes exactly zero.
What
----
:func:`wrap_atoms` returns a *copy* whose atoms (and optional coordinate centres) are
translated by integer lattice vectors into the unit cell, ``r' = r + n @ C`` (row vectors).
Atoms already inside the cell are left bitwise unchanged. With every atom and centre inside
one cell, the required image of any neighbour satisfies ``|k_a| <= ceil(R/h_a)`` whenever
``R/h_a`` is not an integer, so DScribe's own copy range is complete -- including multi-image
cells with ``h_a < R`` (proof sketch: for an image vector ``v = (s_j - s_i + k) @ C`` the
projection onto the unit normal of the face spanned by the other two cell vectors is
``(s_j - s_i + k)_a h_a``, so ``|v| <= R`` implies ``|s_j - s_i + k_a| <= R/h_a``; with
``s`` in ``[0, 1]`` this bounds ``|k_a|`` by ``floor(R/h_a) + 1 = ceil(R/h_a)``).
Image convention (used by every integer comparison): row vectors, cell ``C`` with lattice
vectors as rows, ``displacement_ij(k) = r_j + k @ C - r_i``. Under translations
``r'_j = r_j + n_j @ C`` and ``r'_i = r_i + n_i @ C`` the same physical neighbour has
``k' = k + n_i - n_j``; :func:`canonical_image` maps observed offsets back.
Scope: only fully periodic cells (``pbc == (True, True, True)``) are supported in this
stage; anything else, non-finite positions, a degenerate cell or a radius whose ``R/h`` is
(numerically) an integer raises. Nothing here changes any Atom-HiFi default.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Callable, Optional, Sequence
import numpy as np
__all__ = ['WrapResult', 'cell_heights', 'legacy_copy_counts', 'check_supported', 'wrap_atoms',
'canonical_image', 'enumerate_images', 'corrected_center_fn']
[docs]
@dataclass(frozen=True)
class WrapResult:
"""A wrapped copy and the integer translations that produced it (``r' = r + shifts @ cell``)."""
atoms: object # ase.Atoms copy
shifts: np.ndarray # (n_atoms, 3) int64
centers: Optional[np.ndarray] # (n_centers, 3) float64 translated coordinate centres, or None
center_shifts: Optional[np.ndarray]
[docs]
def cell_heights(cell) -> np.ndarray:
"""Perpendicular heights ``h_a = V / |a_b x a_c|`` of a row-vector cell."""
c = np.asarray(cell, dtype=np.float64)
vol = abs(float(np.linalg.det(c)))
return np.array([vol / np.linalg.norm(np.cross(c[1], c[2])),
vol / np.linalg.norm(np.cross(c[2], c[0])),
vol / np.linalg.norm(np.cross(c[0], c[1]))])
[docs]
def legacy_copy_counts(cell, radius: float) -> np.ndarray:
"""Copies per side DScribe generates on each periodic axis: ``ceil(radius / h_a)``."""
return np.ceil(float(radius) / cell_heights(cell)).astype(np.int64)
[docs]
def check_supported(atoms, radii: Sequence[float] = (), tol: float = 1e-6) -> None:
"""Raise on configurations this adapter does not certify."""
pbc = np.asarray(atoms.get_pbc(), dtype=bool)
if not pbc.all():
raise ValueError(f'only fully periodic cells are supported, got pbc={pbc.tolist()}')
pos = np.asarray(atoms.get_positions(), dtype=np.float64)
if not np.isfinite(pos).all():
raise ValueError('non-finite atomic positions')
cell = np.asarray(atoms.get_cell(), dtype=np.float64)
if not np.isfinite(cell).all() or abs(float(np.linalg.det(cell))) <= 0.0:
raise ValueError('degenerate or non-finite cell')
h = cell_heights(cell)
for r in radii:
q = float(r) / h
if np.any(np.abs(q - np.rint(q)) < tol):
raise ValueError(f'radius {r!r} is within {tol} of an integer multiple of a cell height {h.tolist()}; '
'the copy-range proof needs a non-integer R/h')
def _integer_wrap(points: np.ndarray, cell: np.ndarray, inv: np.ndarray):
frac = points @ inv
n = -np.floor(frac).astype(np.int64)
out = points.copy()
moved = np.any(n != 0, axis=1)
out[moved] = points[moved] + n[moved] @ cell
return out, n
[docs]
def wrap_atoms(atoms, radii: Sequence[float] = (), centers=None, tol: float = 1e-6, frac_guard: float = 1e-9) -> WrapResult:
"""Return a copy translated into the unit cell by integer lattice vectors.
``centers`` may be an ``(m, 3)`` array of coordinate centres; they are translated by the
same rule (their own integer shifts) and returned. Atom-index centres need no handling:
an index into the wrapped copy *is* the wrapped centre. ASE constraints (e.g. ``FixAtoms``)
are kept on the copy but do not block the lattice translation: a periodic image is the same
physical configuration, not a move.
"""
check_supported(atoms, radii, tol)
cell = np.asarray(atoms.get_cell(), dtype=np.float64)
inv = np.linalg.inv(cell)
pos = np.asarray(atoms.get_positions(), dtype=np.float64)
new, n = _integer_wrap(pos, cell, inv)
out = atoms.copy()
out.set_positions(new, apply_constraint=False)
got = np.asarray(out.get_positions(), dtype=np.float64)
if got.tobytes() != new.tobytes():
raise RuntimeError('the wrapped copy does not hold the translated positions')
s = got @ inv
if np.any(s < -frac_guard) or np.any(s > 1.0 + frac_guard):
raise RuntimeError('wrapped fractional coordinates left [0, 1] beyond the guard')
c_new = c_n = None
if centers is not None:
cpos = np.atleast_2d(np.asarray(centers, dtype=np.float64))
if cpos.shape[1] != 3 or not np.isfinite(cpos).all():
raise ValueError('coordinate centres must be a finite (m, 3) array')
c_new, c_n = _integer_wrap(cpos, cell, inv)
return WrapResult(atoms=out, shifts=n, centers=c_new, center_shifts=c_n)
[docs]
def canonical_image(k_observed, n_center, n_neighbor) -> np.ndarray:
"""Map an image offset observed on translated coordinates back to the source convention.
``k = k' - n_i + n_j`` for centre translation ``n_i`` and neighbour translation ``n_j``.
"""
return (np.asarray(k_observed, dtype=np.int64) - np.asarray(n_center, dtype=np.int64)
+ np.asarray(n_neighbor, dtype=np.int64))
[docs]
def enumerate_images(positions, cell, center, radius: float, guard: float = 1e-9):
"""Independent complete enumeration of all (atom, image) pairs within ``radius`` of ``center``.
Works on arbitrary (unwrapped) coordinates: for each atom ``j`` the fractional relative
displacement ``delta = (r_j - c) @ inv(C)`` bounds the admissible images per axis by
``|delta_a + k_a| <= R / h_a``; candidates are widened by one on each side and filtered by
the Cartesian distance. Pairs with ``|d - R| < guard`` are flagged rather than decided.
Returns ``(j, k, d, in_guard_band)`` with ``k`` in the convention of the given positions.
"""
pos = np.asarray(positions, dtype=np.float64)
c = np.asarray(cell, dtype=np.float64)
inv = np.linalg.inv(c)
h = cell_heights(c)
q = float(radius) / h
delta = (pos - np.asarray(center, dtype=np.float64)) @ inv
lo = np.ceil(-delta - q).astype(np.int64) - 1
hi = np.floor(-delta + q).astype(np.int64) + 1
js, ks, ds, gs = [], [], [], []
span = (hi - lo + 1)
for a in range(int(span[:, 0].max())):
for b in range(int(span[:, 1].max())):
for g in range(int(span[:, 2].max())):
k = lo + np.array([a, b, g])
ok = np.all(k <= hi, axis=1)
if not ok.any():
continue
jj = np.flatnonzero(ok)
v = pos[jj] + k[jj] @ c - np.asarray(center, dtype=np.float64)
d = np.sqrt(np.einsum('ij,ij->i', v, v))
keep = d <= float(radius) + guard
js.append(jj[keep]); ks.append(k[jj][keep]); ds.append(d[keep])
gs.append(np.abs(d[keep] - float(radius)) < guard)
if not js:
return (np.zeros(0, np.int64), np.zeros((0, 3), np.int64), np.zeros(0), np.zeros(0, bool))
j = np.concatenate(js); k = np.concatenate(ks); d = np.concatenate(ds); g = np.concatenate(gs)
order = np.lexsort((k[:, 2], k[:, 1], k[:, 0], j))
return j[order], k[order], d[order], g[order]
[docs]
def corrected_center_fn(center_fn: Callable, radii: Sequence[float]) -> Callable:
"""Wrap an Atom-HiFi ``center_descriptor_fn(structures, atom_ids_per_structure)``.
Each structure is replaced by its integer-wrapped copy before the unchanged descriptor
function runs; atom-index centres keep their identity and row order. ``radii`` (the
descriptor's cutoff-plus-padding radii) are required: the copy-range check needs them.
"""
radii = tuple(float(r) for r in radii)
if not radii or not all(np.isfinite(radii)) or min(radii) <= 0:
raise ValueError(f'corrected_center_fn needs the positive finite descriptor radii (cutoff plus Gaussian '
f'padding) to check the copy range; got {radii!r}')
def fn(structures, atom_ids_per_structure):
wrapped = [wrap_atoms(s, radii).atoms for s in structures]
return center_fn(wrapped, atom_ids_per_structure)
fn.__wrapped__ = center_fn
return fn