"""Decoupled ED-SOAP centre descriptor with absolute radii (experimental, opt-in).
Why
---
Production ED-SOAP (:func:`atom_hifi.soap_ML.soap.generate_ED_SOAP`, called by
``runner._build_descriptor_fn``) ties both channels to one ``lengthscale``: decaf forms
``r = lengthscale * cut`` and ``sigma = lengthscale * (cut / 8)`` for the long and the short
channel. A cutoff study needs the four values set independently.
What
----
:class:`EDSoapSpec` holds the ABSOLUTE long/short cutoffs and Gaussian widths (Angstrom), the
radial/angular maxima and the species, all explicit (nothing is defaulted, nothing is
system-specific). :func:`center_fn` returns a ``center_descriptor_fn(structures,
atom_ids_per_structure)`` that calls ``decaf.get_SOAP(structure, 1.0, SOAP_L_cut=r_long,
SOAP_L_sigma=sigma_long, SOAP_S_cut=r_short, SOAP_S_sigma=sigma_short, ...)``. With
``NNDistance = 1.0`` decaf's products ``1.0 * r`` and ``1.0 * sigma`` are exact, so the
values handed to DScribe are exactly the given ones. Output layout and normalisation are
decaf's: ``[long | short]`` per row, each channel L2-normalised per row.
Identity with production: :meth:`EDSoapSpec.from_production` forms the absolute values as
the production float products (``lengthscale * cut`` and ``lengthscale * (cut / 8)``), in
which case the rows are bitwise the production rows. Anything else is a different
descriptor, by construction.
The function composes with :func:`atom_hifi.experimental.periodic.corrected_center_fn`
(pass :func:`support_radii` as its ``radii``). Importing this module does no IO and does
not import decaf or DScribe; both are imported when a descriptor is evaluated.
"""
from __future__ import annotations
import contextlib
import io
import math
from dataclasses import dataclass
from typing import Callable, Sequence
import numpy as np
__all__ = ['EDSoapSpec', 'center_fn', 'ed_soap_rows', 'support_radii']
def _positive_float(name, value):
if isinstance(value, bool) or not isinstance(value, (int, float, np.integer, np.floating)):
raise TypeError(f'{name} must be a real number, got {value!r}')
value = float(value)
if not (math.isfinite(value) and value > 0.0):
raise ValueError(f'{name} must be positive and finite, got {value!r}')
return value
def _positive_int(name, value, minimum):
if isinstance(value, bool) or not isinstance(value, (int, np.integer)) or int(value) < minimum:
raise ValueError(f'{name} must be an int >= {minimum}, got {value!r}')
return int(value)
[docs]
@dataclass(frozen=True)
class EDSoapSpec:
"""Absolute ED-SOAP parameters (Angstrom). Every field is explicit.
``species`` are atomic numbers in the order handed to DScribe (as production passes them);
``periodic`` is the master switch, honoured per structure only where ``pbc`` is set (the
production rule).
"""
r_long: float
sigma_long: float
r_short: float
sigma_short: float
n_long: int
l_long: int
n_short: int
l_short: int
species: tuple
periodic: bool = True
def __post_init__(self):
for name in ('r_long', 'sigma_long', 'r_short', 'sigma_short'):
object.__setattr__(self, name, _positive_float(name, getattr(self, name)))
for name, minimum in (('n_long', 1), ('l_long', 0), ('n_short', 1), ('l_short', 0)):
object.__setattr__(self, name, _positive_int(name, getattr(self, name), minimum))
raw = tuple(self.species)
if any(isinstance(z, (bool, np.bool_)) or not isinstance(z, (int, np.integer)) for z in raw):
raise TypeError(f'species must be integer atomic numbers, got {self.species!r}')
species = tuple(int(z) for z in raw)
if not species or any(z < 1 for z in species) or len(set(species)) != len(species):
raise ValueError(f'species must be distinct positive atomic numbers, got {self.species!r}')
object.__setattr__(self, 'species', species)
if not isinstance(self.periodic, (bool, np.bool_)):
raise TypeError(f'periodic must be a bool, got {self.periodic!r}')
object.__setattr__(self, 'periodic', bool(self.periodic))
[docs]
@classmethod
def from_production(cls, *, lengthscale, l_cut, s_cut, l_nmax, l_lmax, s_nmax, s_lmax, species, periodic=True):
"""The production descriptor as absolute values: ``r = lengthscale * cut`` and
``sigma = lengthscale * (cut / 8)`` per channel (the exact float products decaf forms)."""
ls = float(lengthscale)
return cls(r_long=ls * l_cut, sigma_long=ls * (l_cut / 8), r_short=ls * s_cut, sigma_short=ls * (s_cut / 8),
n_long=l_nmax, l_long=l_lmax, n_short=s_nmax, l_short=s_lmax, species=tuple(species),
periodic=periodic)
[docs]
def identity(self):
"""Stable, exact identity record (floats as ``repr``, i.e. round-trip exact)."""
return {'kind': 'ed_soap_absolute', 'r_long': repr(self.r_long), 'sigma_long': repr(self.sigma_long),
'r_short': repr(self.r_short), 'sigma_short': repr(self.sigma_short), 'n_long': self.n_long,
'l_long': self.l_long, 'n_short': self.n_short, 'l_short': self.l_short,
'species': list(self.species), 'periodic': self.periodic, 'layout': 'long|short',
'normalisation': 'per-row L2 per channel (decaf)'}
[docs]
def ed_soap_rows(structures: Sequence, atom_ids_per_structure, spec: EDSoapSpec) -> np.ndarray:
"""ED-SOAP rows at the requested centres of every structure, stacked in structure order.
``atom_ids_per_structure`` is ``None`` (every atom of every structure) or one index array
per structure (an empty array means every atom, as in production).
"""
import decaf
if not isinstance(spec, EDSoapSpec):
raise TypeError('spec must be an EDSoapSpec')
structures = list(structures)
if atom_ids_per_structure is None:
atom_ids_per_structure = [np.array([], dtype=int) for _ in structures]
elif len(atom_ids_per_structure) != len(structures):
raise ValueError('atom_ids_per_structure must contain one index array per structure')
species = np.array(spec.species)
with contextlib.redirect_stdout(io.StringIO()):
rows = [decaf.get_SOAP(struct, 1.0, species=species,
SOAP_S_cut=spec.r_short, SOAP_S_sigma=spec.sigma_short,
SOAP_S_nmax=spec.n_short, SOAP_S_lmax=spec.l_short,
SOAP_L_cut=spec.r_long, SOAP_L_sigma=spec.sigma_long,
SOAP_L_nmax=spec.n_long, SOAP_L_lmax=spec.l_long,
periodic=bool(spec.periodic and struct.pbc.any()),
atomID=np.asarray(ids, dtype=int))
for struct, ids in zip(structures, atom_ids_per_structure)]
return np.concatenate(rows)
[docs]
def center_fn(spec: EDSoapSpec) -> Callable:
"""An Atom-HiFi ``center_descriptor_fn(structures, atom_ids_per_structure)`` for ``spec``."""
if not isinstance(spec, EDSoapSpec):
raise TypeError('spec must be an EDSoapSpec')
def fn(structures, atom_ids_per_structure):
return ed_soap_rows(structures, atom_ids_per_structure, spec)
fn.spec = spec
return fn
[docs]
def support_radii(spec: EDSoapSpec) -> tuple:
"""DScribe's own neighbour-search radius per channel, ``(long, short)``: the cutoff plus
DScribe's Gaussian padding, read from DScribe objects built with the spec's values (the
radii :func:`atom_hifi.experimental.periodic.corrected_center_fn` certifies)."""
from dscribe.descriptors import SOAP
out = []
for r, sigma, n, lmax in ((spec.r_long, spec.sigma_long, spec.n_long, spec.l_long),
(spec.r_short, spec.sigma_short, spec.n_short, spec.l_short)):
d = SOAP(species=list(spec.species), periodic=True, r_cut=r, sigma=sigma, n_max=n, l_max=lmax,
rbf='gto', average='off')
out.append(float(d._r_cut + d.get_cutoff_padding()))
return tuple(out)