"""Periodic PAO, pair-orbit, and PNO primitives for AICCM2026DEV-B.
The functions here are deliberately independent of the molecular DLPNO
driver. They define the finite Born-von Karman (BvK) objects consumed by
that driver and expose invariants that can be tested before a correlation
energy is trusted.
For an AO overlap ``S`` and an ``S``-orthonormal occupied coefficient matrix
``C_o``, the coefficient-space projector is
``Q = I - C_o C_o^H S``.
Columns of ``Q`` selected by a pair domain are redundant PAOs. Canonical
orthogonalization of their Gram matrix removes the null space without
changing their span. For a pair amplitude matrix ``T_ij`` the Hermitian
model pair density is
``D_ij = (T T^H + T^H T) / (1 + delta_ij)``.
Its eigenvectors are the PNOs and its eigenvalues are non-negative pair
occupations. Tightening the occupation threshold therefore gives nested
ranks and the zero-threshold limit retains the complete PAO pair space.
The MP2 energy itself is not variational, so strict monotonicity of the
energy error remains a numerical validation criterion rather than a theorem.
The pair-density convention matches the real closed-shell construction in
Pinski et al., J. Chem. Phys. 143, 034108 (2015). The finite-BvK pair-orbit
bookkeeping follows Nejad et al., J. Chem. Phys. 163, 214107 (2025),
doi:10.1063/5.0290816: simultaneous translations act on both occupied
indices, and optional space-group permutations enlarge those orbits.
"""
from __future__ import annotations
from dataclasses import dataclass
from itertools import product
from typing import Iterable, Sequence
import numpy as np
__all__ = [
"AICCM2026DevBPAOSpace",
"AICCM2026DevBPNOSpace",
"AICCM2026DevBPairOrbit",
"build_pair_natural_orbitals",
"enumerate_pair_orbits",
"full_space_noncanonical_mp2_energy",
"projected_atomic_orbitals",
"translation_permutations",
]
@dataclass(frozen=True)
class AICCM2026DevBPAOSpace:
"""An orthonormal PAO basis for one finite-torus domain."""
coefficients: np.ndarray
orbital_energies: np.ndarray
projector: np.ndarray
gram_eigenvalues: np.ndarray
source_ao_indices: np.ndarray
orthonormality_error: float
occupied_leakage_error: float
discarded_rank: int
@property
def rank(self) -> int:
return int(self.coefficients.shape[1])
@dataclass(frozen=True)
class AICCM2026DevBPNOSpace:
"""Pair natural orbitals obtained from one PAO pair space."""
coefficients: np.ndarray
occupations: np.ndarray
retained_occupations: np.ndarray
orbital_energies: np.ndarray | None
retained_rank: int
discarded_occupation: float
hermiticity_error: float
minimum_occupation: float
@dataclass(frozen=True)
class AICCM2026DevBPairOrbit:
"""One orbit of unordered occupied pairs under finite permutations."""
representative: tuple[int, int]
members: tuple[tuple[int, int], ...]
n_cells: int
@property
def multiplicity(self) -> int:
return len(self.members)
@property
def per_cell_weight(self) -> float:
"""Coefficient multiplying a representative pair energy per cell."""
return self.multiplicity / self.n_cells
def _validate_occupied_metric(
occupied: np.ndarray,
overlap: np.ndarray,
*,
tolerance: float,
) -> None:
metric = occupied.conj().T @ overlap @ occupied
error = float(np.max(np.abs(metric - np.eye(metric.shape[0]))))
if error > tolerance:
raise ValueError(
"AICCM2026DEV-B PAOs require S-orthonormal occupied orbitals; "
f"maximum error is {error:.3e}"
)
[docs]
def projected_atomic_orbitals(
occupied_coefficients: np.ndarray,
overlap: np.ndarray,
fock: np.ndarray,
*,
ao_indices: Sequence[int] | np.ndarray | None = None,
linear_dependence_threshold: float = 1e-8,
occupied_tolerance: float = 1e-8,
) -> AICCM2026DevBPAOSpace:
"""Project occupieds from selected AOs and canonically orthogonalize.
``ao_indices`` is the union of the atom-centred AOs in a pair domain.
Passing ``None`` selects the full AO space and provides the exact virtual
limit. Eigenvalues are retained relative to the largest PAO Gram
eigenvalue, which makes the threshold invariant to a uniform AO scaling.
"""
occupied = np.asarray(occupied_coefficients)
S = np.asarray(overlap)
F = np.asarray(fock)
if S.ndim != 2 or S.shape[0] != S.shape[1]:
raise ValueError("AICCM2026DEV-B PAO overlap must be square")
if occupied.ndim != 2 or occupied.shape[0] != S.shape[0]:
raise ValueError("AICCM2026DEV-B PAO occupied/overlap shape mismatch")
if F.shape != S.shape:
raise ValueError("AICCM2026DEV-B PAO Fock/overlap shape mismatch")
if linear_dependence_threshold <= 0.0:
raise ValueError("PAO linear-dependence threshold must be positive")
_validate_occupied_metric(
occupied,
S,
tolerance=occupied_tolerance,
)
n_ao = S.shape[0]
if ao_indices is None:
selected = np.arange(n_ao, dtype=int)
else:
selected = np.unique(np.asarray(ao_indices, dtype=int))
if selected.ndim != 1 or selected.size == 0:
raise ValueError("AICCM2026DEV-B PAO domain must select at least one AO")
if selected[0] < 0 or selected[-1] >= n_ao:
raise ValueError("AICCM2026DEV-B PAO domain AO index out of range")
projector = np.eye(n_ao, dtype=np.result_type(occupied, S)) - (
occupied @ (occupied.conj().T @ S)
)
raw = projector[:, selected]
gram = raw.conj().T @ S @ raw
gram = 0.5 * (gram + gram.conj().T)
eigenvalues, eigenvectors = np.linalg.eigh(gram)
order = np.argsort(eigenvalues)[::-1]
eigenvalues = np.asarray(eigenvalues[order], dtype=float)
eigenvectors = eigenvectors[:, order]
scale = max(float(eigenvalues[0]), np.finfo(float).eps)
keep = eigenvalues > linear_dependence_threshold * scale
if not np.any(keep):
raise RuntimeError("AICCM2026DEV-B PAO domain has zero virtual rank")
coefficients = raw @ (
eigenvectors[:, keep] / np.sqrt(eigenvalues[keep])[None, :]
)
pao_metric = coefficients.conj().T @ S @ coefficients
orthogonality_error = float(
np.max(np.abs(pao_metric - np.eye(pao_metric.shape[0])))
)
leakage = float(np.max(np.abs(occupied.conj().T @ S @ coefficients)))
if max(orthogonality_error, leakage) > 10.0 * occupied_tolerance:
raise RuntimeError(
"AICCM2026DEV-B PAO invariant failure: "
f"orthonormality={orthogonality_error:.3e}, "
f"occupied leakage={leakage:.3e}"
)
fock_pao = coefficients.conj().T @ F @ coefficients
fock_pao = 0.5 * (fock_pao + fock_pao.conj().T)
orbital_energies, rotation = np.linalg.eigh(fock_pao)
coefficients = coefficients @ rotation
return AICCM2026DevBPAOSpace(
coefficients=np.asarray(np.real_if_close(coefficients, tol=1000)),
orbital_energies=np.asarray(orbital_energies, dtype=float),
projector=projector,
gram_eigenvalues=eigenvalues,
source_ao_indices=selected,
orthonormality_error=orthogonality_error,
occupied_leakage_error=leakage,
discarded_rank=int(np.count_nonzero(~keep)),
)
[docs]
def build_pair_natural_orbitals(
pair_amplitudes: np.ndarray,
virtual_coefficients: np.ndarray,
*,
diagonal_pair: bool,
occupation_threshold: float,
fock: np.ndarray | None = None,
minimum_rank: int = 1,
negative_occupation_tolerance: float = 1e-10,
) -> AICCM2026DevBPNOSpace:
"""Diagonalize the Hermitian MP2 model pair density.
A zero threshold retains the complete supplied virtual space, including
exactly zero-occupation directions, and is therefore the algebraic
no-truncation limit. Positive thresholds retain occupations strictly
larger than the threshold, subject to ``minimum_rank``.
"""
amplitudes = np.asarray(pair_amplitudes)
virtuals = np.asarray(virtual_coefficients)
if amplitudes.ndim != 2 or amplitudes.shape[0] != amplitudes.shape[1]:
raise ValueError("AICCM2026DEV-B pair amplitudes must be square")
if virtuals.ndim != 2 or virtuals.shape[1] != amplitudes.shape[0]:
raise ValueError("AICCM2026DEV-B PNO virtual/amplitude shape mismatch")
if occupation_threshold < 0.0:
raise ValueError("PNO occupation threshold must be non-negative")
if minimum_rank < 0 or minimum_rank > amplitudes.shape[0]:
raise ValueError("PNO minimum rank is out of range")
adjoint = amplitudes.conj().T
divisor = 2.0 if diagonal_pair else 1.0
pair_density = (amplitudes @ adjoint + adjoint @ amplitudes) / divisor
hermiticity_error = float(
np.max(np.abs(pair_density - pair_density.conj().T))
)
pair_density = 0.5 * (pair_density + pair_density.conj().T)
occupations, eigenvectors = np.linalg.eigh(pair_density)
order = np.argsort(occupations)[::-1]
occupations = np.asarray(occupations[order], dtype=float)
eigenvectors = eigenvectors[:, order]
minimum_occupation = float(occupations[-1])
if minimum_occupation < -negative_occupation_tolerance:
raise RuntimeError(
"AICCM2026DEV-B pair density is not positive semidefinite; "
f"minimum occupation is {minimum_occupation:.3e}"
)
occupations = np.maximum(occupations, 0.0)
if occupation_threshold == 0.0:
keep = np.ones_like(occupations, dtype=bool)
else:
keep = occupations > occupation_threshold
if int(np.count_nonzero(keep)) < minimum_rank:
keep[:minimum_rank] = True
retained_vectors = eigenvectors[:, keep]
coefficients = virtuals @ retained_vectors
orbital_energies: np.ndarray | None = None
if fock is not None:
F = np.asarray(fock)
if F.shape != (virtuals.shape[0], virtuals.shape[0]):
raise ValueError("AICCM2026DEV-B PNO Fock matrix shape mismatch")
fock_pno = coefficients.conj().T @ F @ coefficients
fock_pno = 0.5 * (fock_pno + fock_pno.conj().T)
orbital_energies, rotation = np.linalg.eigh(fock_pno)
coefficients = coefficients @ rotation
return AICCM2026DevBPNOSpace(
coefficients=np.asarray(np.real_if_close(coefficients, tol=1000)),
occupations=occupations,
retained_occupations=occupations[keep],
orbital_energies=(
None
if orbital_energies is None
else np.asarray(orbital_energies, dtype=float)
),
retained_rank=int(np.count_nonzero(keep)),
discarded_occupation=float(np.sum(occupations[~keep])),
hermiticity_error=hermiticity_error,
minimum_occupation=minimum_occupation,
)
def full_space_noncanonical_mp2_energy(
occupied_coefficients: np.ndarray,
virtual_coefficients: np.ndarray,
fock: np.ndarray,
three_center_factors: np.ndarray,
*,
denominator_tolerance: float = 1e-12,
) -> float:
"""Return the closed-shell DF-MP2 energy in rotation-invariant form.
The occupied and virtual inputs may be any real orthonormal gauges. We
diagonalize the two Fock blocks, rotate the fitted three-center factors,
and evaluate the ordinary ordered-pair MP2 expression. This is the
complete-domain, zero-PNO-threshold limit used to audit the iterative
local-pair equations; it is not used for a truncated calculation.
"""
occupied = np.asarray(occupied_coefficients)
virtual = np.asarray(virtual_coefficients)
F = np.asarray(fock)
factors = np.asarray(three_center_factors)
if any(np.iscomplexobj(value) for value in (occupied, virtual, F, factors)):
maximum_imaginary = max(
float(np.max(np.abs(np.asarray(value).imag)))
for value in (occupied, virtual, F, factors)
)
if maximum_imaginary > 1e-10:
raise ValueError(
"AICCM2026DEV-B real-torus MP2 audit requires real tensors"
)
occupied = occupied.real
virtual = virtual.real
F = F.real
factors = factors.real
if occupied.ndim != 2 or virtual.ndim != 2:
raise ValueError("AICCM2026DEV-B MP2 audit coefficients must be matrices")
if occupied.shape[0] != virtual.shape[0] or F.shape != (
occupied.shape[0],
occupied.shape[0],
):
raise ValueError("AICCM2026DEV-B MP2 audit one-particle shape mismatch")
if factors.ndim != 3 or factors.shape[1:] != F.shape:
raise ValueError("AICCM2026DEV-B MP2 audit factor shape mismatch")
f_occ = occupied.T @ F @ occupied
f_vir = virtual.T @ F @ virtual
eps_occ, rotation_occ = np.linalg.eigh(0.5 * (f_occ + f_occ.T))
eps_vir, rotation_vir = np.linalg.eigh(0.5 * (f_vir + f_vir.T))
occupied_canonical = occupied @ rotation_occ
virtual_canonical = virtual @ rotation_vir
factors_ov = np.einsum(
"mi,Pmn,na->Pia",
occupied_canonical,
factors,
virtual_canonical,
optimize=True,
)
integrals = np.einsum(
"Pia,Pjb->ijab",
factors_ov,
factors_ov,
optimize=True,
)
denominator = (
eps_occ[:, None, None, None]
+ eps_occ[None, :, None, None]
- eps_vir[None, None, :, None]
- eps_vir[None, None, None, :]
)
if np.any(denominator >= -denominator_tolerance):
raise RuntimeError(
"AICCM2026DEV-B MP2 audit encountered a non-negative denominator"
)
amplitudes = integrals / denominator
return float(
np.einsum(
"ijab,ijab->",
amplitudes,
2.0 * integrals - integrals.swapaxes(2, 3),
optimize=True,
)
)
def translation_permutations(
n_bands: int,
mesh: Sequence[int],
) -> tuple[np.ndarray, ...]:
"""Return occupied-index permutations for every finite translation.
Occupied indices use ``index = cell_index * n_bands + band`` with cells
in lexicographic product order. Each permutation acts simultaneously on
both indices of an occupied pair.
"""
mesh_tuple = tuple(int(value) for value in mesh)
if len(mesh_tuple) != 3 or any(value < 1 for value in mesh_tuple):
raise ValueError("AICCM2026DEV-B pair mesh must contain three positive values")
if n_bands < 1:
raise ValueError("AICCM2026DEV-B pair orbit needs at least one band")
cells = tuple(product(*(range(value) for value in mesh_tuple)))
cell_index = {cell: index for index, cell in enumerate(cells)}
permutations: list[np.ndarray] = []
for shift in cells:
permutation = np.empty(len(cells) * n_bands, dtype=int)
for source_index, cell in enumerate(cells):
target = tuple(
(cell[axis] + shift[axis]) % mesh_tuple[axis]
for axis in range(3)
)
target_index = cell_index[target]
for band in range(n_bands):
permutation[source_index * n_bands + band] = (
target_index * n_bands + band
)
permutations.append(permutation)
return tuple(permutations)
def _canonical_pair(i: int, j: int) -> tuple[int, int]:
return (i, j) if i <= j else (j, i)
def _validate_permutations(
permutations: Iterable[Sequence[int] | np.ndarray],
) -> tuple[np.ndarray, ...]:
output = tuple(np.asarray(value, dtype=int) for value in permutations)
if not output:
raise ValueError("AICCM2026DEV-B pair orbits require at least one permutation")
n_orbitals = int(output[0].size)
expected = np.arange(n_orbitals)
for permutation in output:
if permutation.shape != (n_orbitals,):
raise ValueError("AICCM2026DEV-B pair permutation shape mismatch")
if not np.array_equal(np.sort(permutation), expected):
raise ValueError("AICCM2026DEV-B pair mapping is not a permutation")
return output
[docs]
def enumerate_pair_orbits(
permutations: Iterable[Sequence[int] | np.ndarray],
*,
n_cells: int,
) -> tuple[AICCM2026DevBPairOrbit, ...]:
"""Partition all unordered occupied pairs under supplied generators.
The breadth-first closure makes the result correct whether callers pass
the full finite symmetry group or only a generating set. The hard gate
verifies that every unordered pair occurs in exactly one orbit.
"""
group_generators = _validate_permutations(permutations)
if n_cells < 1:
raise ValueError("AICCM2026DEV-B pair orbit n_cells must be positive")
n_orbitals = int(group_generators[0].size)
all_pairs = {
(i, j)
for i in range(n_orbitals)
for j in range(i, n_orbitals)
}
unassigned = set(all_pairs)
orbits: list[AICCM2026DevBPairOrbit] = []
while unassigned:
seed = min(unassigned)
members = {seed}
frontier = [seed]
while frontier:
i, j = frontier.pop()
for permutation in group_generators:
mapped = _canonical_pair(
int(permutation[i]),
int(permutation[j]),
)
if mapped not in members:
members.add(mapped)
frontier.append(mapped)
if not members <= all_pairs:
raise RuntimeError("AICCM2026DEV-B pair orbit escaped pair space")
unassigned -= members
ordered = tuple(sorted(members))
orbits.append(
AICCM2026DevBPairOrbit(
representative=ordered[0],
members=ordered,
n_cells=n_cells,
)
)
flattened = [member for orbit in orbits for member in orbit.members]
if len(flattened) != len(all_pairs) or set(flattened) != all_pairs:
raise RuntimeError("AICCM2026DEV-B pair orbits do not partition pair space")
return tuple(sorted(orbits, key=lambda orbit: orbit.representative))