"""Voxel library substitution for super-large structures ("piecing").
Why this exists
---------------
Atom-HiFi selects a minimal representative subset of *structures* from a library
by DECAF-clustering atomic environments and running a greedy set-cover, scoring
each candidate subset with the fidelity metric ``F = L / R``. This works when
the library holds many small-to-medium structures.
It **breaks** when the library is a handful of very large structures (e.g. ~11
cells of ~1000 atoms). Each giant structure already contains nearly *every*
microstate, so a single structure covers the whole motif vocabulary: the
set-cover returns ``|S| = 1``, the bandwidth factor ``c*`` inflates, and ``F``
saturates. This is an artefact of **unit granularity**, not chemistry — the
selection unit (a whole 1000-atom cell) is far coarser than the thing being
covered (a local atomic environment).
The fix: cut each large structure into many small, spatially local **pieces**
and run the selection with pieces as the library units. A bulk piece carries
bulk motifs, a defect piece carries defect motifs — the units are now
heterogeneous, set-cover regains a meaningful gradient, and ``F`` behaves.
Pieces are a *statistics-layer cut only*
----------------------------------------
**Descriptors are never recomputed.** They are computed once on the intact
periodic cell (the descriptor's own ``r_cut`` already encodes the local
environment), and a piece simply *selects rows* of that pre-computed matrix.
Physically carving atoms out would change the chemical environment of boundary
atoms (severed bonds, fake surfaces) and corrupt their descriptors — so we do
not do it. A piece is pure index bookkeeping plus a row permutation; the
numbers fed to DECAF are bit-identical to the intact-cell numbers.
Because the cut touches only which *unit* owns each atom (never the descriptor
values), the whole downstream stack — ``Judge``, ``HiFiWorkflow``, the 1-D /
grid / NSGA-II scans, the set-cover — runs unmodified and unaware that pieces
exist. This module is the container boundary: the runner calls
:func:`make_pieces` to wrap, runs the normal workflow, then calls
:func:`assign_back` to unwrap the per-atom group labels back onto the original
structures.
How a structure is cut: the voxel grid
---------------------------------------
Pieces are disjoint sub-cells of a **fractional-coordinate grid**:
1. Estimate the structure's atom density ``ρ = N / V`` (``V`` = cell volume).
2. Target voxel edge ``L = (target_atoms / ρ) ** (1/3)`` — the cube edge that
would, at this density, hold ``target_atoms`` atoms.
3. Divide each cell vector *i* into ``n_i = max(1, round(|cell_i| / L))`` parts.
The grid lives in fractional coordinates, so it is automatically PBC-correct
and triclinic-safe (sub-cells are parallelepipeds congruent with the cell,
not axis-aligned cubes).
4. Each atom's wrapped scaled position lands it in exactly one voxel. Every
non-empty voxel becomes one piece; empty voxels are dropped; sparse voxels
(a few atoms near a surface) are kept as small pieces — harmless, they just
cover little.
``target_atoms`` is therefore a *granularity dial*, not a hard size: actual
piece populations vary with local density. Exposing atom-count rather than a
voxel edge in Å keeps the knob system-size- and density-independent.
Spatial locality is mandatory. The partition **must** be spatial — chopping
the atom list into arbitrary index chunks would put a uniform random sample of
microstates in every chunk, re-creating the very artefact this module fixes.
Pass-through
------------
Structures with ``len(atoms) <= 4 * target_atoms`` are already fine-grained and
pass through uncut (one piece == the whole structure). Mixed libraries (a few
giants among many small structures) work transparently.
Reporting note
--------------
When piecing is active, ``|S|`` and fidelity reports count **pieces, not
structures**. Representative output contains unique selected original parents,
with every focus atom annotated. :func:`assign_back` returns annotated copies
of every original structure (``motif_id``, ``piece_id``, ``in_selection``); when
a fitted Judge is available the runner writes these copies to ``diagnosed.xyz``
and their three per-atom arrays to ``diagnosed.npz``.
"""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
from ase import Atoms
from ase.calculators.singlepoint import SinglePointCalculator
__all__ = [
'PieceMap', 'make_pieces', 'original_representatives', 'assign_back',
]
[docs]
@dataclass
class PieceMap:
"""Bookkeeping to invert a :func:`make_pieces` cut.
Attributes
----------
parent_idx : list[int]
``parent_idx[p]`` = index of the original structure piece *p* came from.
atom_idx : list[np.ndarray]
``atom_idx[p]`` = atom indices (into the parent structure's full-atom
array, ascending) that belong to piece *p*. Piece-local atom *j*
corresponds to parent atom ``atom_idx[p][j]``.
n_orig : int
Number of original structures.
target_atoms : int
The granularity dial used for the cut (recorded for provenance).
"""
parent_idx: list
atom_idx: list
n_orig: int
target_atoms: int
@property
def n_pieces(self) -> int:
return len(self.parent_idx)
def _voxel_labels(atoms: Atoms, target_atoms: int) -> np.ndarray:
"""Return an integer voxel label per atom (spatial fractional-coord grid)."""
n = len(atoms)
vol = float(abs(atoms.get_volume())) if atoms.cell.rank == 3 else 0.0
if vol > 0.0:
rho = n / vol
edge = (target_atoms / rho) ** (1.0 / 3.0)
lengths = atoms.cell.lengths()
n_div = np.maximum(1, np.round(lengths / edge).astype(int))
scaled = atoms.get_scaled_positions(wrap=True) # [0, 1)
else:
# No 3-D cell: fall back to a Cartesian bounding-box grid.
pos = atoms.get_positions()
span = pos.max(axis=0) - pos.min(axis=0)
span = np.where(span > 0, span, 1.0)
bbox_vol = float(np.prod(span))
rho = n / bbox_vol
edge = (target_atoms / rho) ** (1.0 / 3.0)
n_div = np.maximum(1, np.round(span / edge).astype(int))
scaled = (pos - pos.min(axis=0)) / span # [0, 1]
ijk = np.floor(scaled * n_div).astype(int)
ijk = np.clip(ijk, 0, n_div - 1)
# Flatten the 3-D voxel index to a single label.
return (ijk[:, 0] * n_div[1] + ijk[:, 1]) * n_div[2] + ijk[:, 2]
[docs]
def make_pieces(lib, desc_list, target_atoms):
"""Cut large structures into disjoint voxel pieces (bookkeeping only).
Parameters
----------
lib : list[ase.Atoms]
Original structures (full atoms, intact cells).
desc_list : list[np.ndarray]
Per-structure descriptor arrays, ``desc_list[i]`` of shape
``(len(lib[i]), D)`` — one row per atom, in structure order. Computed
on the intact cell; never recomputed here.
target_atoms : int
Granularity dial (approx. atoms per voxel).
Returns
-------
pieces_lib : list[ase.Atoms]
One ``Atoms`` per non-empty piece (subset of a parent; parent cell/pbc
retained). Drop-in replacement for ``lib``.
pieces_desc_list : list[np.ndarray]
Matching per-piece descriptor rows (selected from ``desc_list``, never
recomputed). Drop-in replacement for the per-structure descriptor list.
piece_map : PieceMap
Inverse-mapping bookkeeping for :func:`assign_back`.
"""
if len(lib) != len(desc_list):
raise ValueError(
f'lib ({len(lib)}) and desc_list ({len(desc_list)}) length mismatch')
pieces_lib = []
pieces_desc_list = []
parent_idx = []
atom_idx = []
for si, (atoms, desc) in enumerate(zip(lib, desc_list)):
n = len(atoms)
if desc.shape[0] != n:
raise ValueError(
f'desc_list[{si}] has {desc.shape[0]} rows but structure has '
f'{n} atoms')
if n <= 4 * target_atoms:
# Pass-through: the whole structure is one piece.
idx_groups = [np.arange(n)]
else:
labels = _voxel_labels(atoms, target_atoms)
idx_groups = [np.where(labels == lab)[0]
for lab in np.unique(labels)]
for idx in idx_groups:
if len(idx) == 0:
continue
pieces_lib.append(atoms[idx])
pieces_desc_list.append(desc[idx])
parent_idx.append(si)
atom_idx.append(idx)
piece_map = PieceMap(parent_idx=parent_idx, atom_idx=atom_idx,
n_orig=len(lib), target_atoms=target_atoms)
return pieces_lib, pieces_desc_list, piece_map
def _piece_focus_positions(parent, atom_idx, element):
"""Return piece-local positions clustered for one focus element."""
symbols = np.asarray(parent.get_chemical_symbols())[atom_idx]
mask = symbols == element
if 'is_focus' in parent.arrays:
mask &= np.asarray(parent.arrays['is_focus'], dtype=bool)[atom_idx]
return np.flatnonzero(mask)
def _copy_with_calculator_results(source):
"""Copy an Atoms object without sharing its calculator."""
copied = source.copy()
if source.calc is not None and source.calc.results:
results = {
key: np.array(value, copy=True)
if isinstance(value, np.ndarray) else value
for key, value in source.calc.results.items()
}
copied.calc = SinglePointCalculator(copied, **results)
return copied
[docs]
def original_representatives(
piece_map, lib, selected_piece_idx, per_el_group_libs, species,
meta_by_group=None, noise_by_group=None):
"""Return selected full parents with every focus atom annotated.
Selection remains piece-wise. Parent structures are deduplicated in first
selected-piece order, while group ids are gathered from *all* sibling
pieces belonging to each selected parent.
"""
if len(lib) != piece_map.n_orig:
raise ValueError(
f'lib ({len(lib)}) and piece_map.n_orig ({piece_map.n_orig}) mismatch')
for element in species:
groups = per_el_group_libs.get(element)
if groups is None or len(groups) != piece_map.n_pieces:
actual = 'missing' if groups is None else len(groups)
raise ValueError(
f'per_el_group_libs[{element!r}] has {actual} piece entries; '
f'expected {piece_map.n_pieces}')
selected_parents = []
seen = set()
for piece_idx in (int(i) for i in selected_piece_idx):
if not 0 <= piece_idx < piece_map.n_pieces:
raise IndexError(
f'selected piece index {piece_idx} outside '
f'[0, {piece_map.n_pieces})')
parent_idx = int(piece_map.parent_idx[piece_idx])
if parent_idx not in seen:
seen.add(parent_idx)
selected_parents.append(parent_idx)
representatives = {
parent_idx: _copy_with_calculator_results(lib[parent_idx])
for parent_idx in selected_parents
}
group_ids = {
parent_idx: np.full(len(lib[parent_idx]), -1, dtype=int)
for parent_idx in selected_parents
}
for piece_idx in range(piece_map.n_pieces):
parent_idx = int(piece_map.parent_idx[piece_idx])
if parent_idx not in representatives:
continue
atom_idx = np.asarray(piece_map.atom_idx[piece_idx], dtype=int)
for element in species:
values = np.asarray(
per_el_group_libs[element][piece_idx], dtype=int)
focus_local = _piece_focus_positions(
lib[parent_idx], atom_idx, element)
if len(values) != len(focus_local):
raise ValueError(
f'piece {piece_idx} element {element!r}: '
f'{len(focus_local)} focus atoms but {len(values)} group ids')
group_ids[parent_idx][atom_idx[focus_local]] = values
result = []
for parent_idx in selected_parents:
full = representatives[parent_idx]
gids = group_ids[parent_idx]
full.set_array('group_id', gids)
if noise_by_group is not None:
# issue #63 provenance: -1 = group promoted from the noise slot.
# Non-focus atoms keep gid -1 and are reported 0 (not rare, just
# not clustered), matching _build_reps_with_group_id.
full.set_array(
'noise_derived',
np.asarray(
[-1 if noise_by_group.get(int(gid), False) else 0
for gid in gids],
dtype=int,
),
)
if meta_by_group is not None:
full.set_array(
'meta_group_id',
np.asarray(
[meta_by_group.get(int(gid), -1) for gid in gids],
dtype=int,
),
)
result.append(full)
return result
def _element_offsets(per_el_group_libs, species):
"""Joint integer encoding offsets so motif ids are unique across elements.
``offset[el]`` is the start of element *el*'s group-id block; element block
width = ``1 + max group id`` observed for that element.
"""
offsets = {}
cur = 0
for el in species:
offsets[el] = cur
gid_lists = per_el_group_libs.get(el, [])
max_gid = -1
for gids in gid_lists:
if len(gids) > 0:
max_gid = max(max_gid, int(np.max(gids)))
cur += max_gid + 1
return offsets
[docs]
def assign_back(piece_map, lib, per_el_group_libs, species, selected_piece_idx,
ids_are_global=False):
"""Map per-piece group labels back onto the original structures.
Pure inverse bookkeeping — no inference, no OOD. Each piece's per-element
DECAF group ids (in piece-local ascending atom order, exactly as the focus
filter produced them) are scattered onto the parent structure's atoms.
Parameters
----------
piece_map : PieceMap
From :func:`make_pieces`.
lib : list[ase.Atoms]
The **original** structures (not pieces).
per_el_group_libs : dict[str, list]
``per_el_group_libs[el][p]`` = list of group ids for element *el*'s
atoms of piece *p* (e.g. ``judge.per_el_group_libs``).
species : list[str]
Focus elements, in the order used to build the joint group space.
selected_piece_idx : iterable[int]
Piece indices chosen by set-cover (the representative subset).
ids_are_global : bool
True when ``per_el_group_libs`` already carries refined/sorted global
group ids (``Judge.apply_refined_group_ids``, issue #40) — element
offsets are skipped so ids are not offset twice.
Returns
-------
list[ase.Atoms]
Copies of the original structures, each carrying integer ``arrays``:
``motif_id`` (joint-encoded group id; -1 for non-focus atoms),
``piece_id`` (global piece index owning the atom), and
``in_selection`` (1 if the owning piece was selected, else 0).
"""
selected = set(int(i) for i in selected_piece_idx)
offsets = ({el: 0 for el in species} if ids_are_global
else _element_offsets(per_el_group_libs, species))
diagnosed = [a.copy() for a in lib]
motif = [np.full(len(a), -1, dtype=int) for a in diagnosed]
piece = [np.full(len(a), -1, dtype=int) for a in diagnosed]
insel = [np.zeros(len(a), dtype=int) for a in diagnosed]
for p in range(piece_map.n_pieces):
parent = piece_map.parent_idx[p]
idx = piece_map.atom_idx[p] # parent atom indices, piece-local order
piece[parent][idx] = p
if p in selected:
insel[parent][idx] = 1
for el in species:
gids = per_el_group_libs.get(el, [])
if p >= len(gids) or len(gids[p]) == 0:
continue
el_local = _piece_focus_positions(
diagnosed[parent], idx, el)
if len(el_local) != len(gids[p]):
raise ValueError(
f'piece {p} el {el!r}: {len(el_local)} atoms but '
f'{len(gids[p])} group ids')
motif[parent][idx[el_local]] = np.asarray(gids[p], int) + offsets[el]
for parent, a in enumerate(diagnosed):
a.set_array('motif_id', motif[parent])
a.set_array('piece_id', piece[parent])
a.set_array('in_selection', insel[parent])
return diagnosed