Source code for vibeqc.symmetry_integrals

"""Phase SYM3: symmetry-reduced lattice-integral helpers.

Builds on the SYM2c atom-pair / cell-index orbit machinery
(:mod:`vibeqc.symmetry_lattice_c`) to provide compressed
representations of one-electron periodic integrals: overlap S(g),
kinetic T(g), nuclear V(g). For each operator-closed cell list the
orbit partition lets us store one ``(n_AO_a, n_AO_b)`` representative
sub-block per orbit instead of the full ``(n_bf × n_bf)`` block at
every cell — a ``|G|``-fold memory saving on high-symmetry crystals
(e.g. simple cubic He at Pm-3m gives orbit compression ~7×; NaCl in
Fm-3m gives the same up to atom-pair structure).

This module ships **storage** reduction. The integrals are still
computed in full by the existing C++ kernels and then compressed —
the |G|-fold *compute* reduction (skip non-representative shell-pair
× cell triples in the kernel itself) is SYM3b, deferred to a
follow-up commit. The compress / reconstruct round-trip provided
here is the validation substrate that SYM3b's compute-reduced kernel
must agree with to machine precision.

The other deliverable here is :func:`verify_lattice_matrix_set_symmetry`
— a debug witness that any user-supplied :class:`LatticeMatrixSet`
satisfies the SYM2c symmetry relations to within a tolerance. Useful
as a contract check on densities / Fock matrices reconstructed from
multi-k SCF, where numerical drift could otherwise accumulate
silently across iterations.

Convention note: SYM2c (and therefore everything here) works on
**symmorphic** symmetry operations only — ``translation = 0``. Many
common space groups (Im-3m for the BCC cubic Wigner-Seitz lattice,
Fd-3m for diamond, etc.) carry non-symmorphic glide-translations
and require additional bookkeeping deferred to a later phase.
:func:`symmorphic_operations` filters a full operator list down to
the symmorphic subset and is exposed as a helper for callers that
want to use whatever symmorphic operators happen to be present.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Sequence, Tuple

import numpy as np

from ._vibeqc_core import (
    BasisSet,
    LatticeMatrixSet,
    LatticeSumOptions,
    PeriodicSystem,
    SymmetryOp,
    compute_kinetic_lattice,
    compute_nuclear_lattice,
    compute_overlap_lattice,
)
from .symmetry_lattice_c import (
    AtomPairOrbits,
    compress_lattice_matrix_set_c,
    identify_atom_pair_orbits,
    reconstruct_lattice_matrix_set_c,
)


__all__ = [
    "OrbitReducedLatticeMatrix",
    "SymmetrisedOrbitIntegralResult",
    "compute_overlap_lattice_with_orbits",
    "compute_kinetic_lattice_with_orbits",
    "compute_nuclear_lattice_with_orbits",
    "compute_overlap_lattice_symmetrised_with_orbits",
    "compute_kinetic_lattice_symmetrised_with_orbits",
    "compute_nuclear_lattice_symmetrised_with_orbits",
    "symmorphic_operations",
    "verify_lattice_matrix_set_symmetry",
]


@dataclass
class OrbitReducedLatticeMatrix:
    """Symmetry-compressed view of a periodic lattice matrix.

    Storage shape:
      - ``orbits`` — :class:`AtomPairOrbits` partition of (atom_a,
        atom_b, cell_index) triples under the symmorphic operator set.
      - ``representatives`` — list of ``(n_AO_a, n_AO_b)`` sub-blocks,
        one per orbit, in orbit order.

    Reconstruction back to a full ``LatticeMatrixSet`` lives in
    :func:`vibeqc.reconstruct_lattice_matrix_set_c`. The
    ``compression_ratio`` property is the average orbit size — the
    factor by which storage shrinks vs the dense LatticeMatrixSet.
    """

    orbits: AtomPairOrbits
    representatives: List[np.ndarray]

    @property
    def compression_ratio(self) -> float:
        return self.orbits.compression_ratio

    @property
    def n_orbits(self) -> int:
        return self.orbits.n_orbits


[docs] @dataclass class SymmetrisedOrbitIntegralResult: """Result bundle for GS3 standardise-then-compress helpers. ``system`` is the standardised/idealised periodic system used for the integral build, ``basis`` is the rebuilt basis on that system, ``report`` is the :func:`vibeqc.symmetrise` metadata, ``operations`` are the symmorphic operations consumed by the orbit partition, and ``reduced`` / ``full`` are the compressed and original lattice matrix representations. """ system: PeriodicSystem basis: BasisSet report: "SymmetriseReport" operations: List[SymmetryOp] reduced: OrbitReducedLatticeMatrix full: LatticeMatrixSet def __iter__(self): """Allow ``reduced, full = compute_*_symmetrised_with_orbits(...)``.""" yield self.reduced yield self.full
def symmorphic_operations( operations: Sequence[SymmetryOp], *, tol: float = 1e-8, ) -> List[SymmetryOp]: """Return the subset of ``operations`` whose fractional translation is zero (symmorphic). Useful preprocessing step for SYM2c / SYM3 — non-symmorphic operations carry glide / screw translations that the current orbit machinery doesn't yet handle. Identity-with-zero-translation is always included (it's symmorphic). """ out: List[SymmetryOp] = [] for op in operations: t = np.asarray(op.translation, dtype=float) if np.all(np.abs(t) < tol): out.append(op) return out def _compress_with_orbits( full_lms: LatticeMatrixSet, basis: BasisSet, system: PeriodicSystem, operations: Sequence[SymmetryOp], *, require_closed: bool = True, ) -> Tuple[AtomPairOrbits, List[np.ndarray]]: """Internal helper: build orbits, compress an LMS, return both. Filters ``operations`` to the symmorphic subset before calling :func:`identify_atom_pair_orbits` so callers don't need to know about symmorphic-vs-non-symmorphic plumbing. ``require_closed=False`` is useful when the cell list comes from a spherical real-space cutoff that may clip some orbit partners near the boundary. The compress / reconstruct round-trip remains exact for the triples that *are* in the cell list; non-closed triples are simply not stored. """ sym_ops = symmorphic_operations(operations) orbits = identify_atom_pair_orbits( full_lms.cells, sym_ops, system, require_closed=require_closed, ) reps = compress_lattice_matrix_set_c(full_lms, orbits, basis) return orbits, reps def compute_overlap_lattice_with_orbits( basis: BasisSet, system: PeriodicSystem, options: LatticeSumOptions, operations: Sequence[SymmetryOp], *, require_closed: bool = True, ) -> Tuple[OrbitReducedLatticeMatrix, LatticeMatrixSet]: """Periodic overlap S(g) plus its SYM2c orbit-compressed view. Returns ``(reduced, full_lms)`` where ``reduced`` carries the orbit partition + per-orbit representative sub-blocks, and ``full_lms`` is the original :class:`LatticeMatrixSet` from the C++ ``compute_overlap_lattice``. The two are equivalent representations to machine precision. For a high-symmetry simple-cubic cell with the full Pm-3m point group (|G| = 48) and a single-atom basis at the origin, the compression ratio approaches |G| as the cell list grows. """ full = compute_overlap_lattice(basis, system, options) orbits, reps = _compress_with_orbits(full, basis, system, operations, require_closed=require_closed) return OrbitReducedLatticeMatrix(orbits=orbits, representatives=reps), full def compute_kinetic_lattice_with_orbits( basis: BasisSet, system: PeriodicSystem, options: LatticeSumOptions, operations: Sequence[SymmetryOp], *, require_closed: bool = True, ) -> Tuple[OrbitReducedLatticeMatrix, LatticeMatrixSet]: """Periodic kinetic T(g) plus its SYM2c orbit-compressed view.""" full = compute_kinetic_lattice(basis, system, options) orbits, reps = _compress_with_orbits(full, basis, system, operations, require_closed=require_closed) return OrbitReducedLatticeMatrix(orbits=orbits, representatives=reps), full def compute_nuclear_lattice_with_orbits( basis: BasisSet, system: PeriodicSystem, options: LatticeSumOptions, operations: Sequence[SymmetryOp], *, require_closed: bool = True, ) -> Tuple[OrbitReducedLatticeMatrix, LatticeMatrixSet]: """Periodic nuclear V(g) plus its SYM2c orbit-compressed view. Note: ``compute_nuclear_lattice`` dispatches on ``options.coulomb_method``. The orbit partition itself is operator-only and method-independent, so this works identically for ``DIRECT_TRUNCATED`` and ``EWALD_3D`` nuclear builds. """ full = compute_nuclear_lattice(basis, system, options) orbits, reps = _compress_with_orbits(full, basis, system, operations, require_closed=require_closed) return OrbitReducedLatticeMatrix(orbits=orbits, representatives=reps), full
[docs] def compute_overlap_lattice_symmetrised_with_orbits( system: PeriodicSystem, basis_name: str, options: LatticeSumOptions, *, symprec: float = 1.0e-4, to_primitive: bool = False, idealize: bool = True, require_closed: bool = True, ) -> SymmetrisedOrbitIntegralResult: """Standardise ``system`` and build orbit-compressed overlap S(g). This is the GS3 convenience path: imported/near-symmetric structures are first cleaned with :func:`vibeqc.symmetrise`, then the basis is rebuilt on the standardised cell, and the SYM3 storage-compressed lattice matrix is built from the attached symmorphic operations. """ clean, basis, report, ops = _standardised_orbit_inputs( system, basis_name, symprec=symprec, to_primitive=to_primitive, idealize=idealize, ) reduced, full = compute_overlap_lattice_with_orbits( basis, clean, options, ops, require_closed=require_closed, ) return SymmetrisedOrbitIntegralResult( system=clean, basis=basis, report=report, operations=ops, reduced=reduced, full=full, )
[docs] def compute_kinetic_lattice_symmetrised_with_orbits( system: PeriodicSystem, basis_name: str, options: LatticeSumOptions, *, symprec: float = 1.0e-4, to_primitive: bool = False, idealize: bool = True, require_closed: bool = True, ) -> SymmetrisedOrbitIntegralResult: """Standardise ``system`` and build orbit-compressed kinetic T(g).""" clean, basis, report, ops = _standardised_orbit_inputs( system, basis_name, symprec=symprec, to_primitive=to_primitive, idealize=idealize, ) reduced, full = compute_kinetic_lattice_with_orbits( basis, clean, options, ops, require_closed=require_closed, ) return SymmetrisedOrbitIntegralResult( system=clean, basis=basis, report=report, operations=ops, reduced=reduced, full=full, )
[docs] def compute_nuclear_lattice_symmetrised_with_orbits( system: PeriodicSystem, basis_name: str, options: LatticeSumOptions, *, symprec: float = 1.0e-4, to_primitive: bool = False, idealize: bool = True, require_closed: bool = True, ) -> SymmetrisedOrbitIntegralResult: """Standardise ``system`` and build orbit-compressed nuclear V(g).""" clean, basis, report, ops = _standardised_orbit_inputs( system, basis_name, symprec=symprec, to_primitive=to_primitive, idealize=idealize, ) reduced, full = compute_nuclear_lattice_with_orbits( basis, clean, options, ops, require_closed=require_closed, ) return SymmetrisedOrbitIntegralResult( system=clean, basis=basis, report=report, operations=ops, reduced=reduced, full=full, )
def _standardised_orbit_inputs( system: PeriodicSystem, basis_name: str, *, symprec: float, to_primitive: bool, idealize: bool, ) -> tuple[PeriodicSystem, BasisSet, "SymmetriseReport", List[SymmetryOp]]: from .periodic_symmetrize import SymmetriseReport, symmetrise result = symmetrise( system, symprec=symprec, to_primitive=to_primitive, idealize=idealize, ) clean = result.system basis = BasisSet(clean.unit_cell_molecule(), str(basis_name)) if clean.symmetry is None: raise RuntimeError( "compute_*_lattice_symmetrised_with_orbits: symmetrise did " "not attach symmetry to the standardised system" ) ops = symmorphic_operations(clean.symmetry.operations) if not ops: raise ValueError( "compute_*_lattice_symmetrised_with_orbits: the standardised " "system has no symmorphic symmetry operations available for " "orbit compression" ) report: SymmetriseReport = result.report return clean, basis, report, ops def verify_lattice_matrix_set_symmetry( lms: LatticeMatrixSet, basis: BasisSet, system: PeriodicSystem, operations: Sequence[SymmetryOp], *, atol: float = 1e-10, require_closed: bool = True, ) -> dict: """Verify ``lms`` respects the SYM2c symmetry relations. Compresses to per-orbit representatives and reconstructs; the reconstructed full LMS must equal the input to within ``atol``. Returns a diagnostics dict:: { "n_orbits": int, "compression_ratio": float, "max_residual": float, "passes": bool, } A failing ``passes`` flag indicates the LMS does **not** transform correctly under the (symmorphic) symmetry operations — useful for catching numerical drift in densities reconstructed from multi-k SCF, or for catching bugs in routines that should produce symmetry-respecting outputs. Raises if the operator list contains no symmorphic operators (no symmetry to check against). """ sym_ops = symmorphic_operations(operations) if not sym_ops: raise ValueError( "verify_lattice_matrix_set_symmetry: operator list contains " "no symmorphic operations; nothing to check." ) orbits = identify_atom_pair_orbits( lms.cells, sym_ops, system, require_closed=require_closed, ) reps = compress_lattice_matrix_set_c(lms, orbits, basis) recon = reconstruct_lattice_matrix_set_c( reps, orbits, basis, system, sym_ops, cells=list(lms.cells), ) max_residual = 0.0 for slot, recon_blk in enumerate(recon): full_blk = np.asarray(lms.blocks[slot], dtype=float) diff = float(np.linalg.norm(full_blk - recon_blk)) if diff > max_residual: max_residual = diff return { "n_orbits": orbits.n_orbits, "compression_ratio": orbits.compression_ratio, "max_residual": max_residual, "passes": max_residual < atol, }