"""Native multi-k periodic RHF + RKS via Gaussian density fitting.
This module implements closed-shell multi-k SCF on arbitrary Bravais
lattices and Monkhorst-Pack k-meshes, pairing the cell-resolved C++
DF kernels (``compute_{2c,3c}_eri_lattice_blocks``) with Bloch
phase assembly (``build_lpq_bloch_native(q_cart)``) and a per-k
SCF loop reusing the multi-k DIIS machinery from the legacy Ewald
driver.
Coulomb is k-diagonal (built once from ``Lpq(q=0)`` against the
k-weighted density). Exchange is per-pair: ``K(k_i)`` accumulates
contributions from every ``k_j`` weighted by the mesh, contracted
through ``Lpq(q = k_j - k_i)`` and its Hermitian conjugate. The
``Lpq(q)`` tensors are pre-computed once per unique momentum
transfer at SCF setup; storage is currently dense in memory, with
streaming queued for paper-grade systems where the
``nkpts^2 x naux x nao^2`` tensor exceeds host RAM.
Convention notes
----------------
* HF / hybrid exchange carries the ``exxdiv='ewald'`` Madelung
correction: the multi-k branch applies the finite-k-mesh G=0
divergence shift via the BvK-supercell constant
(:func:`_madelung_for_kmesh`), and the Γ HF fast-path delegates to
:func:`vibeqc.run_pbc_gdf_rhf` (primitive-cell shift). PySCF parity
therefore uses ``exxdiv='ewald'`` (PySCF's default). The legacy
molecular-limit driver :func:`vibeqc.run_rhf_periodic_gamma_gdf`
(exxdiv=None semantics) remains the Γ fallback for RKS / dim<3 /
charged / smeared / use_compcell runs.
* Generic Bravais: the driver uses arbitrary lattice vectors via
the existing ``bloch_sum`` and Bloch-block helpers, with no
cubic-only assumption. 1D / 2D / 3D periodicity all supported.
* Closed-shell RHF and RKS only in this module. UHF / UKS multi-k
is a separate module (pending).
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Sequence, Tuple, Union
import numpy as np
from ._vibeqc_core import (
BasisSet,
BlochKMesh,
Functional,
GridOptions,
InitialGuess,
LatticeSumOptions,
PeriodicKSOptions,
PeriodicRHFOptions,
PeriodicSystem,
SCFIteration,
bloch_sum,
build_grid,
build_xc_periodic,
compute_kinetic_lattice,
compute_overlap_lattice,
direct_lattice_cells,
level_shift_at_iter,
nuclear_repulsion_per_cell,
)
from ._vibeqc_core import (
direct_lattice_cells as _direct_cells,
)
from ._vibeqc_core import (
monkhorst_pack as _mp_native,
)
from .aux_basis import (
build_lpq_bloch_compcell,
build_lpq_bloch_mdf,
build_lpq_bloch_native_fft,
default_aux_for,
make_aux_basis_set,
make_modrho_aux_basis,
)
from .kpoints import KPoints
from .lattice_screening import (
RcutStrategy,
make_lattice_opts,
)
from .linear_dependence import (
LinearDependenceError,
raise_if_severe,
scf_preflight_overlap_check,
)
from .madelung import apply_exxdiv_ewald_to_K
from .occupations import (
hartree_to_kelvin_temperature as _hartree_to_kelvin_temperature,
)
from .periodic_k_density import (
density_matrices_per_k as _density_from_orbitals,
)
from .periodic_k_density import (
real_space_density_from_per_k_density as _real_space_density_from_per_k_density,
)
from .smearing import (
SmearingOptions as _SmearingOptions,
apply_smearing_open_shell as _apply_smearing_open_shell,
closed_shell_periodic_occupations as _closed_shell_periodic_occupations,
)
from .options_dump import dump_active_settings
from .periodic_fock_multi_k import (
build_periodic_fock_ewald3d_k,
make_ewald_3d_lattice_j_cache,
)
from .periodic_grid import build_periodic_becke_grid
from .periodic_rhf_gdf import (
PeriodicRHFGDFResult,
_resolve_fock_mixing,
_resolve_level_shift_warmup_cycles,
run_rhf_periodic_gamma_gdf,
)
from .periodic_rhf_multi_k_ewald import (
_canonical_orthogonalizer_complex,
_diag_in_orth_basis,
)
from .periodic_scf_accelerators import (
DynamicDamping,
MultiKPeriodicSCFAccelerator,
)
from .periodic_screened_exchange import reject_unscreened_range_separated
from .progress import ProgressLogger, resolve_progress
__all__ = [
"PeriodicKRHFGDFResult",
"PeriodicKRKSGDFResult",
"run_krhf_periodic_gdf",
"run_krks_periodic_gdf",
]
# =====================================================================
# Basis-aware one-electron lattice-sum cutoff
# =====================================================================
def _is_pyscf_auto(strategy: object) -> bool:
"""True when ``strategy`` selects PySCF-style per-shell rcut tuning.
Accepts the :class:`RcutStrategy` enum member, the lowercase string
``"pyscf_auto"`` (the driver default), or its ``repr`` form.
"""
if strategy is None:
return False
if isinstance(strategy, RcutStrategy):
return strategy is RcutStrategy.PYSCF_AUTO
return str(strategy).strip().lower() in (
"pyscf_auto",
"rcutstrategy.pyscf_auto",
)
def _oneel_lattice_opts(
system: PeriodicSystem,
basis: BasisSet,
base_opts: LatticeSumOptions,
*,
rcut_strategy: object,
k_points_cart: "np.ndarray",
plog=None,
) -> LatticeSumOptions:
"""Lattice-sum options for the one-electron overlap/kinetic sums,
with ``cutoff_bohr`` grown until the Bloch overlap S(k) is no longer
in the non-PSD ("critical") tier at every requested k-point.
The flat default (``LatticeSumOptions.cutoff_bohr`` = 15 bohr) under-
converges the Bloch overlap sum S(k) = Σ_g e^{ik·g} S(g) for diffuse
molecular bases (small primitive exponents): the truncated sum is not
a Gram matrix and acquires spurious negative eigenvalues, tripping the
``scf_preflight_overlap_check`` non-PSD ("critical") abort even though
the image summation itself is correct. The GDF Lpq path already sizes
its cutoff per basis; this brings the *one-electron* S/T sum in line.
Rather than the a-priori PySCF rcut estimate (which targets the
AFT-corrected Lpq sum and systematically under-sizes a *bare* overlap
— e.g. it returns 11.6 bohr for a def2-SVP/MgO basis whose overlap
only turns PSD at ~17 bohr), this *measures*: it grows the cutoff via
:func:`optimize_truncation` until the worst k-point reaches "error"
severity or better. That is the important boundary: "critical" means
the lattice-truncated overlap has the wrong metric signature, while
"error" means the overlap is PSD within the negative-eigenvalue tolerance
but has near-null directions. The GDF driver already uses canonical
orthogonalisation per k-point, so those near-null directions are projected
explicitly instead of aborting before SCF. Overlap builds are milliseconds,
so even the 8-evaluation worst case is negligible against SCF cost.
Only ever *raises* the cutoff: tight/solid bases (sto-3g,
pob-tzvp-rev2) are already PSD at the caller's cutoff, so
optimisation returns it unchanged and every previously-validated job
is untouched. ``rcut_strategy="flat"`` opts out entirely.
A basis that is genuinely near-linearly-dependent on a dense lattice
(full def2-SVP on rocksalt MgO: converged S(k) λ_min ~1e-9) is still
allowed once the lattice-sum truncation artefact is removed; canonical
orthogonalisation drops the redundant near-null directions. If cutoff
growth cannot eliminate critical negative eigenvalues, this returns the
caller's opts unchanged and the per-k preflight aborts with the
actionable hint.
"""
if not _is_pyscf_auto(rcut_strategy):
return base_opts
# Lazy import: eigs_preflight does `import vibeqc` at module load,
# which re-enters the partially-initialised package while this module
# is itself being imported by vibeqc/__init__. Deferring to call time
# sidesteps that ordering hazard.
from .eigs_preflight import optimize_truncation
rep = optimize_truncation(
system,
basis,
lattice_opts=base_opts,
k_points_cart=k_points_cart,
target_severity="error",
cutoff_growth_factor=1.5,
joint_growth=False,
)
if not rep.converged:
# Still has critical negative directions: no cutoff below the cap fixed
# the corrupted metric. Leave the caller's opts so the per-k preflight
# aborts with the hint.
return base_opts
grown = float(rep.optimized_lattice_opts.cutoff_bohr)
if grown <= float(base_opts.cutoff_bohr) + 1e-9:
return base_opts # already PSD at caller's cutoff (tight basis)
widened = make_lattice_opts(
basis,
strategy=RcutStrategy.FLAT,
base_opts=base_opts,
cutoff_bohr=grown,
)
if plog is not None:
plog.info(
"S/T lattice cutoff grown "
f"{float(base_opts.cutoff_bohr):.2f} -> {grown:.2f} bohr "
f"(overlap non-critical at all k, severity={rep.final_severity}) "
f"for diffuse basis '{basis.name}'"
)
return widened
# Appended to the LinearDependenceError message when a periodic S(k)
# preflight aborts, replacing the generic "upstream bug" framing with
# actionable remedies. Passed to scf_preflight_overlap_check by the
# multi-k GDF drivers below.
_PERIODIC_OVERLAP_HINT = (
"Periodic S(k): the one-electron lattice cutoff is auto-widened for "
"diffuse bases and near-null directions are handled by canonical "
"orthogonalisation. A residual critical S(k) therefore means the metric "
"still has negative directions after cutoff growth. Filter diffuse "
"primitives with vq.make_basis(..., exp_to_discard=0.1), raise the "
"one-electron cutoff, or use a solid-state basis (pob-tzvp-rev2). "
"vibeqc.eigs_preflight."
"disambiguate_critical_overlap() reports which applies."
)
def _gdf_overlap_preflight(
S: np.ndarray,
*,
plog,
label: str,
basis: BasisSet,
):
"""GDF-specific overlap policy after cutoff auto-growth.
``critical`` still aborts: the Bloch overlap has negative eigenvalues and
the metric signature is wrong. ``error`` is allowed here because it is a
PSD near-linear-dependence case, and the multi-k GDF driver immediately
applies canonical orthogonalisation with ``linear_dep_threshold``.
"""
rep = scf_preflight_overlap_check(
S,
plog=plog,
label=label,
basis=basis,
remediation_hint=_PERIODIC_OVERLAP_HINT,
raise_on_severe=False,
)
if rep.severity == "critical":
try:
raise_if_severe(rep, allow_warn=True, allow_critical=False)
except LinearDependenceError as exc:
raise LinearDependenceError(
f"{exc.args[0]} {_PERIODIC_OVERLAP_HINT}",
rep,
) from None
if rep.severity == "error" and plog is not None:
plog.info(
f"[WARN] {label}: near-linear-dependent periodic overlap "
"will be handled by canonical orthogonalisation"
)
return rep
# =====================================================================
# Result types
# =====================================================================
[docs]
@dataclass
class PeriodicKRHFGDFResult:
"""Result of :func:`run_krhf_periodic_gdf`.
All per-k arrays are length-``nkpts`` Python lists holding
complex Hermitian (or real, for Γ-only) matrices in AO basis.
"""
energy: float
e_electronic: float
e_nuclear: float
n_iter: int
converged: bool
mo_energies: List[np.ndarray]
mo_coeffs: List[np.ndarray]
fock: List[np.ndarray]
overlap: List[np.ndarray]
hcore: List[np.ndarray]
density: List[np.ndarray]
kpoints_cart: np.ndarray
kpoint_weights: np.ndarray
scf_trace: List[SCFIteration] = field(default_factory=list)
functional: Optional[str] = None
e_xc: float = 0.0
e_coulomb: float = 0.0
e_hf_exchange: float = 0.0
e_dft_plus_u: float = 0.0
fock_mixing: float = 0.0
level_shift: float = 0.0
level_shift_warmup_cycles: int = 0
smearing_temperature: float = 0.0
fermi_level: float = 0.0
entropy: float = 0.0
free_energy: float = 0.0
occupations: List[np.ndarray] = field(default_factory=list)
aux_basis_name: str = ""
n_aux: int = 0
backend: str = "native-multi-k-gdf"
@property
def energy_per_cell_ha(self) -> float:
return float(self.energy)
[docs]
@dataclass
class PeriodicKRKSGDFResult(PeriodicKRHFGDFResult):
"""Result of :func:`run_krks_periodic_gdf`."""
# =====================================================================
# Setup helpers
# =====================================================================
def _options_or_default(options, *, is_ks: bool):
if options is not None:
return options
return PeriodicKSOptions() if is_ks else PeriodicRHFOptions()
@dataclass(frozen=True)
class _GammaKMeshInfo:
"""Single-Γ k-mesh metadata for the native Γ fast path."""
kpoints_cart: np.ndarray
weights: np.ndarray
input_n_kpoints: int
def _gamma_kmesh_info(
system: PeriodicSystem,
kmesh: Union[Sequence[int], KPoints, BlochKMesh],
) -> Optional[_GammaKMeshInfo]:
"""Return metadata when ``kmesh`` is exactly a single Γ point.
The KRHF/KRKS public entry points use multi-k-shaped APIs, but for
a single Γ point the answer is delegated to the (fully native and
well-tested) Γ-GDF driver. Anything else falls through into the
real multi-k SCF.
"""
if isinstance(kmesh, KPoints):
kpoints = np.asarray(kmesh.kpoints_cart, dtype=np.float64).reshape(-1, 3)
weights = np.asarray(kmesh.weights, dtype=np.float64).reshape(-1)
elif isinstance(kmesh, BlochKMesh):
kpoints = np.asarray(kmesh.kpoints, dtype=np.float64).reshape(-1, 3)
weights = np.asarray(kmesh.weights, dtype=np.float64).reshape(-1)
else:
mesh = _mesh_tuple_for_system(system, kmesh)
if mesh != (1, 1, 1):
return None
bm = _mp_native(system, [1, 1, 1], [0, 0, 0], False)
kpoints = np.asarray(bm.kpoints, dtype=np.float64).reshape(-1, 3)
weights = np.asarray(bm.weights, dtype=np.float64).reshape(-1)
if kpoints.shape != (1, 3) or weights.shape != (1,):
return None
if not np.allclose(kpoints[0], 0.0, atol=1e-12, rtol=0.0):
return None
if not np.isclose(float(weights[0]), 1.0, atol=1e-12, rtol=0.0):
return None
return _GammaKMeshInfo(
kpoints_cart=kpoints.copy(),
weights=weights.copy(),
input_n_kpoints=1,
)
def _wrap_gamma_gdf_result(
gamma: PeriodicRHFGDFResult,
info: _GammaKMeshInfo,
*,
functional: Optional[str],
result_cls,
):
"""Adapt the native Γ-GDF result to the KRHF/KRKS result shape."""
return result_cls(
energy=float(gamma.energy),
e_electronic=float(gamma.e_electronic),
e_nuclear=float(gamma.e_nuclear),
n_iter=int(gamma.n_iter),
converged=bool(gamma.converged),
mo_energies=[np.asarray(gamma.mo_energies)],
mo_coeffs=[np.asarray(gamma.mo_coeffs)],
fock=[np.asarray(gamma.fock)],
overlap=[np.asarray(gamma.overlap)],
hcore=[np.asarray(gamma.hcore)],
density=[np.asarray(gamma.density)],
kpoints_cart=np.asarray(info.kpoints_cart, dtype=np.float64),
kpoint_weights=np.asarray(info.weights, dtype=np.float64),
scf_trace=list(gamma.scf_trace),
functional=functional or str(getattr(gamma, "functional", "") or "") or None,
e_xc=float(getattr(gamma, "e_xc", 0.0)),
e_coulomb=float(getattr(gamma, "e_coulomb", 0.0)),
e_hf_exchange=float(getattr(gamma, "e_hf_exchange", 0.0)),
e_dft_plus_u=float(getattr(gamma, "e_dft_plus_u", 0.0)),
fock_mixing=float(getattr(gamma, "fock_mixing", 0.0)),
level_shift=float(getattr(gamma, "level_shift", 0.0)),
level_shift_warmup_cycles=int(getattr(gamma, "level_shift_warmup_cycles", 0)),
smearing_temperature=float(getattr(gamma, "smearing_temperature", 0.0)),
fermi_level=float(getattr(gamma, "fermi_level", 0.0)),
entropy=float(getattr(gamma, "entropy", 0.0)),
free_energy=float(getattr(gamma, "free_energy", gamma.energy)),
occupations=[np.asarray(getattr(gamma, "occupations", np.empty(0)))],
aux_basis_name=str(getattr(gamma, "aux_basis_name", "") or ""),
n_aux=int(getattr(gamma, "n_aux", 0)),
backend="native-gamma-gdf-via-k-gdf",
)
def _mesh_tuple_for_system(
system: PeriodicSystem,
mesh: Union[Sequence[int], KPoints, BlochKMesh],
) -> Tuple[int, int, int]:
dim = int(system.dim)
if dim not in (1, 2, 3):
raise ValueError(f"PeriodicSystem.dim must be 1, 2, or 3; got {dim}")
if isinstance(mesh, (KPoints, BlochKMesh)):
mesh_metadata = getattr(mesh, "mesh", None)
if mesh_metadata is None:
raise ValueError(
"periodic GDF: this operation requires structured k-mesh "
"metadata; explicit unstructured k-points are insufficient"
)
arr = list(mesh_metadata)
else:
arr = list(mesh)
if len(arr) == dim:
arr = arr + [1] * (3 - dim)
elif len(arr) != 3:
raise ValueError(
f"periodic GDF: kmesh tuple must have length {dim} for "
f"dim={dim} systems or length 3; got {arr!r}"
)
out = tuple(int(x) for x in arr)
if any(x < 1 for x in out):
raise ValueError(f"periodic GDF: kmesh entries must be >= 1; got {arr!r}")
return tuple(out[i] if i < dim else 1 for i in range(3))
def _kmesh_to_kpoints_weights(
system: PeriodicSystem,
kmesh: Union[Sequence[int], KPoints, BlochKMesh],
) -> Tuple[np.ndarray, np.ndarray]:
"""Normalise ``kmesh`` to ``(kpoints_cart, weights)`` arrays.
``kpoints_cart`` is shape ``(n_k, 3)`` in bohr⁻¹.
``weights`` is shape ``(n_k,)`` summing to 1.
"""
if isinstance(kmesh, KPoints):
kpts = np.asarray(kmesh.kpoints_cart, dtype=np.float64).reshape(-1, 3)
w = np.asarray(kmesh.weights, dtype=np.float64).reshape(-1)
elif isinstance(kmesh, BlochKMesh):
kpts = np.asarray(kmesh.kpoints, dtype=np.float64).reshape(-1, 3)
w = np.asarray(kmesh.weights, dtype=np.float64).reshape(-1)
else:
mesh = _mesh_tuple_for_system(system, kmesh)
bm = _mp_native(system, list(mesh), [0, 0, 0], False)
kpts = np.asarray(bm.kpoints, dtype=np.float64).reshape(-1, 3)
w = np.asarray(bm.weights, dtype=np.float64).reshape(-1)
if kpts.shape[0] == 0:
raise ValueError("periodic k-GDF: kmesh has zero k-points")
if not np.isclose(float(w.sum()), 1.0, atol=1e-9):
raise ValueError(
f"periodic k-GDF: kpoint weights must sum to 1; got {float(w.sum()):.6f}"
)
return kpts, w
# =====================================================================
def _occupations_per_k(
eps_per_k: Sequence[np.ndarray],
weights: np.ndarray,
n_elec_per_cell: int,
smearing_T: float,
n_occ_each: int,
*,
bz_integration: Optional[str] = None,
system: Optional[PeriodicSystem] = None,
kmesh: Optional[BlochKMesh] = None,
) -> Tuple[List[np.ndarray], float, float]:
"""Fermi-Dirac or hard-Aufbau occupations across the k-mesh.
Returns ``(occ_per_k, fermi_level, entropy_per_cell)``.
"""
if bz_integration is not None:
bz_integration = str(bz_integration).strip().lower()
if bz_integration not in ("smearing", "gilat"):
raise ValueError(
"run_krhf_periodic_gdf: bz_integration must be None, "
f"'smearing', or 'gilat'; got {bz_integration!r}"
)
if bz_integration == "gilat":
if smearing_T > 0.0:
raise NotImplementedError(
"run_krhf_periodic_gdf: bz_integration='gilat' is a "
"sharp-Fermi-surface occupation backend and cannot be "
"combined with finite-temperature smearing."
)
if system is None or kmesh is None:
raise ValueError(
"run_krhf_periodic_gdf: Gilat occupations require the "
"PeriodicSystem and BlochKMesh metadata."
)
from .bz_integration import gilat_occupations_for_kmesh
occ_gr, ef_gr = gilat_occupations_for_kmesh(
system,
kmesh,
eps_per_k,
float(n_elec_per_cell),
spin_degeneracy=2.0,
)
return occ_gr, float(ef_gr), 0.0
# Global BZ filling (one Fermi level across all k; PySCF KSCF
# get_occ convention) -- per-k-independent Aufbau silently
# mis-occupies band-overlap systems (the b4a6faba-regressed
# rebuild filled exactly n_occ at every k; on LiH FCC (2,2,2)
# that lands at -2.96 Ha instead of the documented -7.92 Ha).
return _closed_shell_periodic_occupations(
eps_per_k,
weights,
float(n_elec_per_cell),
int(n_occ_each),
float(smearing_T),
)
def _normalise_initial_density_k(
initial_density_k: Sequence[np.ndarray],
*,
n_k: int,
n_basis: int,
label: str,
) -> List[np.ndarray]:
"""Validate caller-supplied per-k density blocks."""
blocks = list(initial_density_k)
if len(blocks) != int(n_k):
raise ValueError(
f"{label}: initial_density_k has {len(blocks)} blocks; "
f"expected {int(n_k)} for the target k-mesh."
)
out: List[np.ndarray] = []
for ik, block in enumerate(blocks):
D = np.asarray(block, dtype=complex)
if D.shape != (int(n_basis), int(n_basis)):
raise ValueError(
f"{label}: initial_density_k[{ik}] has shape {D.shape}; "
f"expected {(int(n_basis), int(n_basis))}."
)
out.append(0.5 * (D + D.conj().T))
return out
def _madelung_for_kmesh(system: PeriodicSystem, mesh: Sequence[int]) -> float:
"""k-mesh-aware Ewald-Madelung constant ``ξ`` for the exxdiv shift.
The finite-k-mesh HF-exchange divergence correction (``exxdiv='ewald'``)
uses the Madelung constant of the **Born-von-Kármán supercell** implied
by the k-mesh -- the cell whose lattice vectors are the primitive ones
scaled by the per-direction mesh count -- NOT the primitive cell. For an
``(n1, n2, n3)`` Monkhorst-Pack mesh the supercell lattice is
``A . diag(n1, n2, n3)``.
Matches PySCF ``pyscf.pbc.tools.pbc.madelung(cell, kpts)`` (verified
out-of-process: LiH primitive FCC at (2,2,2) -> ξ = 0.297038, exactly
half the primitive-cell ξ = 0.594076). Using the primitive-cell value
over-counts the exxdiv K-shift by ``Nk^(1/3)`` and over-binds the
multi-k total energy -- LiH (2,2,2): -592 mHa (the bug behind the prior
-2495 Ha; the cderi-gauge fix exposes it).
"""
from .madelung import madelung_constant_for_cell
A = np.asarray(system.lattice, dtype=float)
A_super = A @ np.diag([float(n) for n in mesh])
# ξ depends only on the lattice geometry + volume; reuse one atom as a
# placeholder (madelung_constant_for_cell ignores atom Z / positions).
super_sys = PeriodicSystem(3, A_super, [system.unit_cell[0]])
return float(madelung_constant_for_cell(super_sys))
# Positive-energy slack for the multi-k energy-sanity guard (Ha). A bound
# neutral closed-shell cell has E_total < 0, but cramped/artificial Bravais
# smoke-test cells can converge just above zero (the 8-bohr hexagonal H₂
# coverage cell lands at +0.067 Ha); a positive energy beyond this slack is
# unphysical (the broken multi-k compcell+AFT LiH lands at +8.485 Ha).
POSITIVE_E_SLACK_HA = 1.0
def _check_energy_sanity(
result: PeriodicKRHFGDFResult,
system: PeriodicSystem,
plog: ProgressLogger,
) -> None:
"""Post-condition: the multi-k SCF total energy is physically sane.
A bound, neutral, closed-shell unit cell has a **negative** total
energy whose magnitude is of order ``S_atoms Z^2/2`` (the loose
hydrogenic bound on absolute binding). Two failure signatures are
rejected, both observed on the multi-k GDF path for tight ionic
crystals (LiH primitive FCC, kmesh=(2,2,2), def2-svp-jk aux):
* **Runaway** -- ``|E_total|`` orders of magnitude beyond
``max(10.SZ^2, 100)`` Ha. ``use_compcell=True``/``exxdiv='ewald'``
lands at ``E ≈ -2495 Ha`` vs PySCF ``-7.92 Ha``: the per-q
compcell ``Lpq`` fit is internally inconsistent and the SCF
"converges" to a numerical fixed point of a broken Fock.
* **Unbound** -- ``E_total`` positive beyond ``POSITIVE_E_SLACK_HA``.
A bound neutral cell has ``E_total < 0``; the
``apply_aft_correction=True`` variant lands at ``+8.485 Ha``, which
the runaway bound alone would miss (its magnitude is comparable to
the true ``-7.92``). A *small* positive energy is tolerated:
cramped/artificial Bravais smoke-test cells can converge just above
zero without being the catastrophic-garbage pattern (e.g. the
8-bohr hexagonal H₂ "coverage" cell at ``+0.067 Ha``). A fixed
slack (rather than a ``SZ^2``-scaled one) keeps the check strict for
heavy cells, where any sizeable positive energy is unphysical.
When a non-physical energy is reported as **converged**, this is the
silent-corruption pattern CLAUDE.md Sec.7 warns about -- a user gets
``converged=True`` with a meaningless number and may use it
downstream. We RAISE rather than return it (do NOT paper over with
damping/thresholds -- Sec.7). For a non-converged run (``converged=False``
already signals failure) we warn + tag the backend so partial state
stays inspectable.
Multi-k GDF parity landed 2026-06-02 (the per-(k_i,k_j)-resolved
Lpq cache + the BvK-supercell exxdiv Madelung; see
``handovers/HANDOVER_GDF_OUTSTANDING.md`` Sec. 1), so this guard firing means a
regression in that machinery -- a gauge mismatch, a wrong Madelung
convention, or an inconsistent cderi cache. H₂-style vacuum-box
cells are fine at multi-k -- the guard only trips on the genuinely
broken numbers, so it does not fire on a correct multi-k energy
(e.g. LiH ``-7.92 Ha``).
"""
E = float(result.energy)
z_sum_sq = sum(atom.Z**2 for atom in system.unit_cell)
sane_bound = max(10.0 * z_sum_sq, 100.0) # loose hydrogenic + floor
runaway = abs(E) > sane_bound
unbound = E > POSITIVE_E_SLACK_HA
if not (runaway or unbound):
return
reason = "runaway divergence" if runaway else "positive (unbound) total energy"
msg = (
f"run_krhf_periodic_gdf: SCF total energy {E:.6e} Ha is non-physical "
f"({reason}). A bound, neutral, closed-shell cell has E_total < 0 "
f"(a positive energy beyond {POSITIVE_E_SLACK_HA:g} Ha slack is "
f"unbound) and |E_total| < {sane_bound:.2e} Ha (loose hydrogenic "
"bound on S_atoms Z^2/2). The SCF has converged to a numerical fixed "
"point of a broken Fock -- a gauge mismatch, a wrong exxdiv Madelung "
"convention, or an internally-inconsistent Lpq cache (CLAUDE.md Sec.7 -- "
"not a convergence-aid problem). Multi-k GDF parity landed 2026-06-02 "
"(handovers/HANDOVER_GDF_OUTSTANDING.md Sec. 1), so this firing indicates a "
"regression in that machinery. Pass check_energy_sanity=False to "
"bypass this guard (diagnostics only)."
)
try:
result.backend = result.backend + "+SANITY_FAILED"
except Exception:
pass
if result.converged:
raise RuntimeError(msg)
plog.info(" WARNING: " + msg)
def _build_xc_k_from_density(
*,
basis: BasisSet,
system: PeriodicSystem,
grid: object,
func: Functional,
density_k: Sequence[np.ndarray],
kmesh_bloch: BlochKMesh,
cells: Sequence[object],
kpoints_cart: Sequence[np.ndarray],
lat_opts: LatticeSumOptions,
) -> tuple[float, List[np.ndarray]]:
"""Build periodic XC from the full finite-torus density.
Local and semilocal XC are primitive-cell functionals of the density, but
the AO matrix returned by :func:`build_xc_periodic` is a lattice matrix.
A multi-k finite-torus RKS Fock must therefore reconstruct the real-space
density blocks from all k-point density matrices and Bloch-fold the
resulting ``V_xc(R)`` to each k. Feeding only the home-cell density and
adding one Γ matrix to every k is exact only in the vacuum/molecular
limit; on tight crystals it can make the KS map non-variational.
"""
density_real = _real_space_density_from_per_k_density(
density_k,
kmesh_bloch,
cells,
)
xc_contrib = build_xc_periodic(
basis,
system,
grid,
func,
density_real,
lat_opts,
)
vxc_k: List[np.ndarray] = []
for k_cart in kpoints_cart:
vk = np.asarray(
bloch_sum(xc_contrib.V_xc, np.asarray(k_cart, dtype=float).reshape(3)),
dtype=complex,
)
vxc_k.append(0.5 * (vk + vk.conj().T))
return float(xc_contrib.e_xc), vxc_k
def _build_xc_k_from_density_uks(
*,
basis: BasisSet,
system: PeriodicSystem,
grid: object,
func: Functional,
density_alpha_k: Sequence[np.ndarray],
density_beta_k: Sequence[np.ndarray],
kmesh_bloch: BlochKMesh,
cells: Sequence[object],
kpoints_cart: Sequence[np.ndarray],
lat_opts: LatticeSumOptions,
) -> tuple[float, List[np.ndarray], List[np.ndarray]]:
"""Open-shell sibling of :func:`_build_xc_k_from_density`.
Folds each spin's per-k density to the full real-space finite-torus
density set (the inverse Bloch sum, ``P_s(g) = S_k w_k e^{-ik.g}
P_s(k)``), evaluates the spin-polarised periodic XC on the pair, and
Bloch-folds ``V_a(g)`` / ``V_b(g)`` back to every k. This is the
same convention the closed-shell multi-k branch uses -- the
BZ-averaged home-cell shortcut it replaces was exact only in the
vacuum/molecular limit (the 2026-07-09 KRKS finding's defect class,
commit b3f74aa9). With ``P_a = P_b = P/2`` the folded spin densities
reproduce the closed-shell density pointwise, so KUKS(mult=1) ==
KRKS holds by construction on any grid.
"""
from ._vibeqc_core import build_xc_periodic_uks
density_alpha_real = _real_space_density_from_per_k_density(
density_alpha_k,
kmesh_bloch,
cells,
)
density_beta_real = _real_space_density_from_per_k_density(
density_beta_k,
kmesh_bloch,
cells,
)
xc_contrib = build_xc_periodic_uks(
basis,
system,
grid,
func,
density_alpha_real,
density_beta_real,
lat_opts,
)
va_k: List[np.ndarray] = []
vb_k: List[np.ndarray] = []
for k_cart in kpoints_cart:
k_arr = np.asarray(k_cart, dtype=float).reshape(3)
va = np.asarray(bloch_sum(xc_contrib.V_alpha, k_arr), dtype=complex)
vb = np.asarray(bloch_sum(xc_contrib.V_beta, k_arr), dtype=complex)
va_k.append(0.5 * (va + va.conj().T))
vb_k.append(0.5 * (vb + vb.conj().T))
return float(xc_contrib.e_xc), va_k, vb_k
def _preflight_gdf_lpq_memory(
plog,
*,
n_basis: int,
n_aux: int,
n_kpoints: int,
need_k_pairs: bool,
open_shell: bool,
route_label: str,
options=None,
) -> None:
"""Estimate + gate the dense multi-k GDF ``Lpq`` cache before building it.
The per-pair ``Lpq`` cache is the multi-k GDF memory bottleneck and is
held dense in RAM (streaming is future work, see the module docstring), so
a paper-grade cell can be OOM-killed (exit 137) before SCF iter 1 with no
``.out``/``.system``/``.qvf`` artifacts -- the prompt-75 NiO/def2-SVP KUKS
``(4,4,4)`` case. This logs the peak estimate (so a `progress`-on rerun
localises the cost) and raises :class:`~vibeqc.memory.InsufficientMemoryError`
early with route-specific remedies when it cannot fit. Override with
``VIBEQC_GDF_MEMORY_OVERRIDE=1`` (or a truthy ``options.memory_override``).
"""
import os as _os
from .memory import check_periodic_gdf_memory, estimate_periodic_multik_gdf
est = estimate_periodic_multik_gdf(
n_basis=int(n_basis),
n_aux=int(n_aux),
n_kpoints=int(n_kpoints),
need_k_pairs=bool(need_k_pairs),
open_shell=bool(open_shell),
)
pair_kind = "k^2 exchange pairs" if need_k_pairs else "diagonal pairs"
plog.info(
f"GDF Lpq cache peak estimate ~{est.total_gb:.1f} GB "
f"({int(n_kpoints)} k-points, {pair_kind}, naux={int(n_aux)}, "
f"nao={int(n_basis)})"
)
allow = bool(getattr(options, "memory_override", False)) or bool(
_os.environ.get("VIBEQC_GDF_MEMORY_OVERRIDE")
)
check_periodic_gdf_memory(
est,
n_kpoints=int(n_kpoints),
route_label=route_label,
allow_exceed=allow,
)
def _reject_slab_dim(system: PeriodicSystem, entry: str) -> None:
"""Fail closed when a bulk GDF driver is handed a ``dim=2`` slab.
GDF's Bloch / AFT machinery assumes 3-D periodicity. Handed a ``dim=2``
slab it silently treats the layer as a 3-D crystal of sheets stacked
``a3`` apart (the normal-axis k-mesh pinned to a single point), so the
total energy depends on the bookkeeping ``|a3|`` *and* on the k-mesh --
the CLAUDE.md Sec. 7 "impossible / mesh-dependent energy" symptom (a
purely geometric Madelung sum cannot depend on the k-mesh). Slabs must use
the vacuum-free 2D gauge instead. The Gamma-only ``run_pbc_gdf_*`` drivers
already guard this; the multi-k entries did not, which is how a slab
reached the 3-D crystal path.
``dim == 1`` is deliberately NOT rejected. There is no rigorous 1-D Coulomb
gauge yet (``CoulombMethod.NEUTRALIZED_1D`` still raises), so the
cached-Lpq GDF Hartree is the *supported* polymer/wire route -- see
``5fe4d021`` ("Lpq Hartree for multi-k dim<3 KS"), which fixed the dim=1
H2-chain from -3.0897 to the variational -3.89 Ha/cell. Guarding ``dim != 3``
here would silently revert that fix. Only the slab has a rigorous vacuum-free
alternative to redirect users to, so only the slab fails closed.
"""
dim = int(system.dim)
if dim == 2:
raise NotImplementedError(
f"{entry}: GDF is a bulk (dim=3) Coulomb builder; got dim={dim} "
"(a slab). Running a slab through it would treat it as a 3-D "
"crystal of sheets a3 apart, giving a3- and k-mesh-dependent "
"energies. Use jk_method='auto' (or 'slab_ewald_2d') -- the "
"rigorous vacuum-free 2D Coulomb (SLAB_EWALD_2D). For a "
"3-D-with-vacuum reference cell instead, build it with "
"vibeqc.build.slab(..., periodic_z=True) so dim=3."
)
[docs]
def run_krhf_periodic_gdf(
system: PeriodicSystem,
basis: BasisSet,
kmesh: Union[Sequence[int], KPoints, BlochKMesh] = (1, 1, 1),
options: Optional[Union[PeriodicRHFOptions, PeriodicKSOptions]] = None,
*,
functional: Optional[str] = None,
aux_basis: Optional[str] = None,
aux_drop_eta: float = 0.0,
linear_dep_threshold: float = 1e-7,
gdf_linear_dep_threshold: float = 1e-9,
apply_modrho: bool = True,
fock_mixing: Optional[float] = None,
level_shift_warmup_cycles: Optional[int] = None,
use_compcell: bool = False,
compcell_eta: float = 1.0,
apply_aft_correction: bool = True,
aft_ft_convention: str = "libint",
aft_precision: float = 1e-10,
rcut_strategy: Optional[object] = "pyscf_auto",
rcut_precision: float = 1e-8,
k_exchange: str = "gdf",
gdf_method: str = "rsgdf",
rsgdf_ke_cutoff: float = 200.0,
rsgdf_tail_ke_cutoff: Optional[float] = None,
fit_screen_threshold: float = 0.0,
mdf_ke_cutoff: float = 40.0,
bz_integration: Optional[str] = None,
density_mixer: Optional[str] = None,
density_mixer_depth: int = 8,
density_mixer_beta: float = 0.5,
density_mixer_kerker: bool = False,
kerker_k0: float = 1.5,
kerker_strength: float = 1.0,
kerker_cutoff_ha: float = 120.0,
dft_plus_u_sites: Optional[Sequence[object]] = None,
initial_density_k: Optional[Sequence[np.ndarray]] = None,
check_energy_sanity: bool = True,
progress: Union[bool, ProgressLogger, None] = None,
verbose: Optional[int] = None,
) -> PeriodicKRHFGDFResult:
"""Run closed-shell periodic HF / KS multi-k SCF via native GDF.
For a single Γ point this delegates to
:func:`vibeqc.run_rhf_periodic_gamma_gdf` (kept in lock-step with
the multi-k path); for any other ``kmesh`` it runs the full
multi-k loop here.
Parameters
----------
system, basis
Periodic system and AO basis.
kmesh
``(n1, n2, n3)`` Monkhorst-Pack mesh, a :class:`KPoints`
instance, or a :class:`BlochKMesh`. Defaults to Γ-only.
options
:class:`PeriodicRHFOptions` (HF) or :class:`PeriodicKSOptions`
(KS).
functional
libxc functional name when running KS; ``None`` means HF.
aux_basis
Auxiliary basis name. Defaults to ``default_aux_for(basis.name)``.
aux_drop_eta
Auxiliary primitive cull threshold passed to
:func:`make_aux_basis_set`.
linear_dep_threshold
Per-k overlap eigenvalue floor for canonical orthogonalisation.
gdf_linear_dep_threshold
Auxiliary metric eigenvalue floor for ``Lpq`` Cholesky-style
fitting (forwarded to :func:`build_lpq_bloch_native`).
apply_modrho
Whether the auxiliary basis is renormalised via
:func:`aux_basis.modrho_renormalise` before fitting (default
on; matches the Γ-only driver).
fock_mixing
Override the resolver-resolved CRYSTAL FMIXING fraction.
level_shift_warmup_cycles
Override the resolver-resolved level-shift warm-up length.
k_exchange
Exchange backend on the ``use_compcell=True`` path:
``'gdf'`` (default) contracts the cached per-(k_i, k_j) Lpq
tensors -- O(N_k^2) pair contractions per iteration; ``'cosx'``
(**EXPERIMENTAL**) builds K via the real-space multi-k COSX
engine (:class:`vibeqc.periodic_cosx_k.KPointCosxK`, M3b-3):
one K(g) block build per iteration (mesh-size independent) +
Bloch folds, and the off-diagonal Lpq tensors are skipped at
setup (only the diagonal J pairs are built). The Coulomb J
stays on GDF either way; the exxdiv='ewald' correction
applies identically to both backends. Requires
``use_compcell=True`` and a multi-k mesh (the Γ-only fast
path ignores it, like ``use_compcell``).
Since M3b-4b the COSX K is the composed SR+LR exchange --
matrix-level validated at ~1e-4 against the independent RSGDF
route. Pair it with ``gdf_method='rsgdf'`` for a single-gauge
Fock (see below); a runtime warning documents the remaining
experimental status (dense-mesh SCF convergence -- M3b-4c).
gdf_method
Lpq builder for the ``use_compcell=True`` cache: ``'rsgdf'``
(default -- the (k_i,k_j)-ket-resolved all-FT Bloch-pair route,
:func:`vibeqc.aux_basis.build_lpq_bloch_native_fft`; validated
at µHa parity vs PySCF on LiH FCC (2,2,2),
handovers/HANDOVER_GDF_OUTSTANDING.md Sec. 1) or ``'compcell'`` (Sun-2017
compensated charges + AFT correction -- q-only cderi, exact
only in the vacuum-box limit; catastrophically wrong on tight
ionic cells, the M3b-5 finding #1 / the -2495 Ha class -- keep
it off the default until the compcell builder is
pair-resolved). The RSGDF route is also the consistent partner
for ``k_exchange='cosx'`` (single-gauge pairing, M3b-4c).
rsgdf_ke_cutoff
Dense-FFT-mesh kinetic-energy cutoff (Ha) for
``gdf_method='rsgdf'`` (default 200).
fit_screen_threshold
Cauchy-Schwarz pre-screen of the three-centre GDF fit on the
``gdf_method='rsgdf'`` path (default ``0.0`` = off, exact). AO
pairs whose Schwarz bound ``max_P sqrt((P|P)).sqrt((muν|muν))``
stays below the threshold are dropped from every per-(k_i,k_j)
fit tensor; kept/dropped pair counts are logged (no silent
truncation). ``1e-10``-class values reproduce the unscreened
energy to well below SCF accuracy. Forwarded to
:func:`vibeqc.aux_basis.build_lpq_bloch_native_fft` -- every
consumer of this driver (CCM ``run_ccm_rhf_gdf`` /
``run_ccm_rks_gdf``, RIJCOSX diagonal-J, KS) inherits it.
bz_integration
``None`` / ``"smearing"`` use the existing global Aufbau or
Fermi-Dirac occupation path. ``"gilat"`` selects the
parameter-free Gilat-Raubenheimer net at T = 0; it cannot be
combined with finite-temperature smearing.
dft_plus_u_sites
Optional Dudarev +U sites. For true multi-k meshes, the driver
builds the k-averaged per-spin occupation matrix and adds the
resulting ``S(k) V_U S(k)`` shift to each closed-shell Fock block.
check_energy_sanity
When ``True`` (default) a post-SCF guard rejects a non-physical
total energy: a converged run that lands at a runaway
(``|E| > max(10.SZ^2, 100)`` Ha) or positive (unbound) energy
RAISES ``RuntimeError`` instead of returning a converged garbage
number (CLAUDE.md Sec.7 silent-corruption). Set ``False`` to bypass
the guard for parity/debug scripts that want the raw value back;
see :func:`_check_energy_sanity`.
progress, verbose
Live progress logging passthrough.
"""
_reject_slab_dim(system, "run_krhf_periodic_gdf")
opts = _options_or_default(options, is_ks=functional is not None)
plog = resolve_progress(progress, verbose=verbose)
# ---------------- Γ fast path ----------------------------------
# The Γ-fast-path delegates to run_rhf_periodic_gamma_gdf which
# short-circuits J/K to Ewald-3D + molecular-limit-K on dim=3
# (it doesn't have a compcell option). For Γ-only compcell SCF,
# users should call ``vibeqc.run_pbc_gdf_rhf`` directly -- the
# Γ-only driver that uses Lpq for both J and K. We attempted to
# skip this fast path when use_compcell=True and fall through to
# the multi-k branch with n_k=1, but the multi-k SCF loop's
# per-k weighting doesn't degenerate cleanly to Γ-only
# (gives ~5x incorrect energy on H2). Keeping the fast path means
# ``use_compcell=True`` at ``kmesh=(1,1,1)`` is silently ignored
# -- surfaced via a warning so users know to switch drivers.
if k_exchange not in ("gdf", "cosx"):
raise ValueError(
f"run_krhf_periodic_gdf: k_exchange must be 'gdf' or "
f"'cosx'; got {k_exchange!r}"
)
if k_exchange == "cosx" and not use_compcell:
raise ValueError(
"run_krhf_periodic_gdf: k_exchange='cosx' requires "
"use_compcell=True (the COSX K rides the cached-Lpq GDF-J "
"path; the legacy Ewald-3D path has its own exchange)"
)
if gdf_method not in ("compcell", "rsgdf", "mdf"):
raise ValueError(
f"run_krhf_periodic_gdf: gdf_method must be 'compcell', "
f"'rsgdf', or 'mdf'; got {gdf_method!r}"
)
if float(fit_screen_threshold) < 0.0:
raise ValueError(
"run_krhf_periodic_gdf: fit_screen_threshold must be >= 0; "
f"got {fit_screen_threshold}"
)
if float(fit_screen_threshold) > 0.0 and gdf_method != "rsgdf":
# Loud, not silent: the Schwarz fit screen lives in the rsgdf
# builder (build_lpq_bloch_native_fft); a threshold on the
# compcell/mdf routes would be silently ignored otherwise.
raise NotImplementedError(
"run_krhf_periodic_gdf: fit_screen_threshold is implemented "
f"for gdf_method='rsgdf' only (got {gdf_method!r})."
)
gamma_info = _gamma_kmesh_info(system, kmesh)
if gamma_info is not None and dft_plus_u_sites:
raise NotImplementedError(
"run_krhf_periodic_gdf: dft_plus_u_sites is wired on the true "
"multi-k GDF loop only. Use a non-Gamma k-mesh such as "
"(1,1,2), or omit dft_plus_u_sites for the Gamma GDF fast path."
)
# Closed-shell Γ HF *and* KS in run_pbc_gdf_rhf's supported domain
# delegate to that PySCF-µHa-validated driver (exxdiv='ewald'), so
# kmesh=(1,1,1) is the Nk=1 limit of the multi-k exxdiv='ewald' path --
# no convention discontinuity vs (2,1,1)+, and consistent with
# run_periodic_job's default-Γ routing. KS delegation landed 2026-07-09
# (the Finding-§4 residual): the legacy molecular-limit gamma driver's
# exchange channel carries NO exxdiv convention at all -- its full-range
# real-space K left Γ-path HYBRID KS +8.32e-2 Ha off the exxdiv-matched
# real-Γ direct route / external PySCF KRKS on rocksalt LiH/STO-3G PBE0
# (not the strict-zero-mode offset a_x·ξ·N_e/2 = 0.297 Ha, a distinct
# truncated-lattice-sum gauge; see tests/test_ccm_rks_direct.py).
# The legacy gamma driver (below) stays the fallback for charged /
# open-shell / dim<3 cells, finite-T smearing, fock-mixing / level-shift
# convergence aids, and explicit use_compcell (kept on the historical
# warning path). The gate mirrors run_pbc_gdf_rhf's preconditions + the
# knobs it honours (DIIS + damping, not fock_mixing / level_shift /
# smearing).
if (
gamma_info is not None
and not use_compcell
and density_mixer is None
):
# Resolve KS-via-options too, so the result class + wrapper agree
# with what run_pbc_gdf_rhf will actually run.
_gamma_func = (
functional or str(getattr(opts, "functional", "") or "") or None
)
_gamma_q_nuc = float(sum(atom.Z for atom in system.unit_cell))
_gamma_n_elec = int(system.n_electrons())
_pure_gdf_gamma_ok = (
int(system.dim) == 3
and _gamma_n_elec % 2 == 0
and int(system.multiplicity) == 1
and abs(_gamma_q_nuc - _gamma_n_elec) <= 0.5
and float(getattr(opts, "smearing_temperature", 0.0) or 0.0) <= 0.0
and float(getattr(opts, "fock_mixing", 0.0) or 0.0) == 0.0
and float(getattr(opts, "level_shift", 0.0) or 0.0) == 0.0
and not fock_mixing
)
if _pure_gdf_gamma_ok:
from .pbc_gdf import run_pbc_gdf_rhf
gamma = run_pbc_gdf_rhf(
system,
basis,
opts,
functional=_gamma_func,
aux_basis=aux_basis,
aux_drop_eta=aux_drop_eta,
exxdiv="ewald",
gdf_method=gdf_method,
mdf_ke_cutoff=mdf_ke_cutoff,
rsgdf_ke_cutoff=rsgdf_ke_cutoff,
rsgdf_tail_ke_cutoff=rsgdf_tail_ke_cutoff,
linear_dep_threshold=linear_dep_threshold,
gdf_linear_dep_threshold=gdf_linear_dep_threshold,
fit_screen_threshold=fit_screen_threshold,
progress=plog,
verbose=verbose,
)
return _wrap_gamma_gdf_result(
gamma,
gamma_info,
functional=_gamma_func,
result_cls=(
PeriodicKRKSGDFResult
if _gamma_func is not None
else PeriodicKRHFGDFResult
),
)
if gamma_info is not None and rsgdf_tail_ke_cutoff is not None:
raise NotImplementedError(
"run_krhf_periodic_gdf: rsgdf_tail_ke_cutoff requires the pure "
"PBC-GDF Gamma fast path or a true multi-k mesh. This Gamma "
"k-mesh would fall back to the legacy molecular-limit GDF driver, "
"which has no high-|G| tail correction."
)
if gamma_info is not None and use_compcell and density_mixer is None:
plog.info(
" WARNING: use_compcell=True at kmesh=(1,1,1) is currently "
"ignored (the Γ-fastpath delegates to run_rhf_periodic_gamma_gdf "
"which doesn't support compcell). For Γ-only compcell SCF, "
"call vibeqc.run_pbc_gdf_rhf(...) directly. Continuing with "
"the legacy Ewald-3D + molecular-limit-K path. "
"(k_exchange is ignored on this path too.)"
)
if gamma_info is not None and float(fit_screen_threshold) > 0.0:
# Reaching here means the pure PBC-GDF Γ fast path (which
# supports the screen via run_pbc_gdf_rhf) was NOT taken --
# this Γ k-mesh falls back to the legacy molecular-limit GDF
# driver, which has no screened fit. Loud, not silent.
raise NotImplementedError(
"run_krhf_periodic_gdf: fit_screen_threshold at a Γ k-mesh "
"requires the pure PBC-GDF Γ fast path (rsgdf, closed-shell, "
"no density_mixer / fock_mixing / level_shift / smearing / "
"compcell); this run would fall back to the legacy "
"molecular-limit GDF driver, which has no screened fit. "
"Call vibeqc.run_pbc_gdf_rhf(..., fit_screen_threshold=...) "
"directly, adjust the conflicting options, or drop the "
"threshold."
)
if gamma_info is not None and density_mixer is None:
gamma = run_rhf_periodic_gamma_gdf(
system,
basis,
opts,
functional=functional,
aux_basis=aux_basis,
aux_drop_eta=aux_drop_eta,
linear_dep_threshold=linear_dep_threshold,
gdf_linear_dep_threshold=gdf_linear_dep_threshold,
apply_modrho=apply_modrho,
fock_mixing=fock_mixing,
level_shift_warmup_cycles=level_shift_warmup_cycles,
progress=plog,
verbose=verbose,
)
result_cls = (
PeriodicKRKSGDFResult if functional is not None else PeriodicKRHFGDFResult
)
return _wrap_gamma_gdf_result(
gamma,
gamma_info,
functional=functional,
result_cls=result_cls,
)
# ---------------- Multi-k branch ------------------------------
func_name = functional or str(getattr(opts, "functional", "") or "")
is_ks = bool(func_name)
func = Functional(func_name, 1) if is_ks else None
# The Lpq-contracted (and COSX) K is full-range only; screened
# hybrids must not silently run as their full-range twins.
reject_unscreened_range_separated(func, where="run_krhf_periodic_gdf")
alpha = float(func.hf_exchange_fraction) if func is not None else 1.0
fock_mixing_value = _resolve_fock_mixing(opts, fock_mixing)
level_shift = float(getattr(opts, "level_shift", 0.0))
max_iter = int(opts.max_iter)
warmup_cycles = _resolve_level_shift_warmup_cycles(
opts,
level_shift=level_shift,
max_iter=max_iter,
override=level_shift_warmup_cycles,
)
# Explicit per-iteration schedule (unified with the molecular
# drivers). Empty ⇒ the warm-up step function; non-empty ⇒ resolved
# per iteration by the shared C++ helper.
_ls_schedule = list(getattr(opts, "level_shift_schedule", None) or [])
smearing_T = float(getattr(opts, "smearing_temperature", 0.0))
if smearing_T < 0.0:
raise ValueError("run_krhf_periodic_gdf: smearing_temperature must be >= 0")
if bz_integration is not None:
bz_integration = str(bz_integration).strip().lower()
if bz_integration not in ("smearing", "gilat"):
raise ValueError(
"run_krhf_periodic_gdf: bz_integration must be None, "
f"'smearing', or 'gilat'; got {bz_integration!r}"
)
use_gilat = bz_integration == "gilat"
if use_gilat and smearing_T > 0.0:
raise NotImplementedError(
"run_krhf_periodic_gdf: bz_integration='gilat' is a "
"sharp-Fermi-surface occupation backend and cannot be combined "
"with finite-temperature smearing."
)
lat_opts: LatticeSumOptions = opts.lattice_opts
label = f"KRKS {func_name}" if is_ks else "KRHF"
n_elec = system.n_electrons()
if n_elec % 2 != 0:
raise ValueError(
"run_krhf_periodic_gdf: closed-shell RHF/RKS requires "
f"even electron count; got {n_elec}"
)
if system.multiplicity != 1:
raise ValueError(
"run_krhf_periodic_gdf: closed-shell RHF/RKS requires "
f"multiplicity=1; got {system.multiplicity}"
)
n_occ = n_elec // 2
kpoints_cart, weights = _kmesh_to_kpoints_weights(system, kmesh)
n_k = kpoints_cart.shape[0]
# Resolve kmesh to BlochKMesh + lattice cells for Ewald Fock builder.
if isinstance(kmesh, BlochKMesh):
kmesh_bloch = kmesh
elif isinstance(kmesh, KPoints):
kmesh_bloch = kmesh.to_bloch_kmesh()
else:
mesh = _mesh_tuple_for_system(system, kmesh)
kmesh_bloch = _mp_native(system, list(mesh), [0, 0, 0], False)
cells = _direct_cells(system, lat_opts.cutoff_bohr)
aux_name = aux_basis or default_aux_for(basis.name)
plog.banner(f"run_krhf_periodic_gdf {label} kmesh={n_k} k-points")
plog.info(
f"{label} multi-k native GDF / aux={aux_name}, "
f"cutoff={lat_opts.cutoff_bohr:.2f} bohr"
)
if use_gilat:
plog.info("bz integration: Gilat-Raubenheimer sharp-Fermi net")
plog.info(f"basis: {basis.name} ({basis.nbasis} BFs / {basis.nshells} shells)")
dim = int(system.dim)
active_lengths = [
float(np.linalg.norm(np.asarray(system.lattice, dtype=float)[:, i]))
for i in range(dim)
]
plog.info(
f"periodicity: dim={dim}D, active lengths="
+ ", ".join(f"{x:.3f}" for x in active_lengths)
+ " bohr"
)
n_int_cells = len(direct_lattice_cells(system, lat_opts.cutoff_bohr))
n_nuc_cells = len(direct_lattice_cells(system, lat_opts.nuclear_cutoff_bohr))
plog.info(
"lattice cells: "
f"one-electron/GDF cutoff -> {n_int_cells}, "
f"nuclear cutoff -> {n_nuc_cells}"
)
dump_active_settings(
plog,
[
("PeriodicKSOptions" if is_ks else "PeriodicRHFOptions", opts),
("LatticeSumOptions", lat_opts),
(
"k-GDF kwargs",
{
"functional": func_name or None,
"hf_exchange_fraction": alpha,
"fock_mixing": fock_mixing_value,
"fmixing_percent": 100.0 * fock_mixing_value,
"level_shift": level_shift,
"level_shift_warmup_cycles": warmup_cycles,
"smearing_temperature": smearing_T,
"aux_basis": aux_name,
"aux_drop_eta": float(aux_drop_eta),
"linear_dep_threshold": float(linear_dep_threshold),
"gdf_linear_dep_threshold": float(gdf_linear_dep_threshold),
"apply_modrho": bool(apply_modrho),
"n_kpoints": n_k,
},
),
],
)
# ---- Functional + grid ----------------------------------------
grid = None
if is_ks:
grid_options = getattr(opts, "grid", None)
if grid_options is None:
grid_options = GridOptions()
if bool(getattr(opts, "use_periodic_becke", False)):
grid = build_periodic_becke_grid(
system,
grid_options=grid_options,
image_radius_bohr=float(getattr(opts, "becke_image_radius_bohr", 0.0)),
)
else:
grid = build_grid(system.unit_cell_molecule(), grid_options)
# ---- Real-space one-electron integrals ------------------------
# Widen the S/T overlap-decay cutoff for diffuse bases so the Bloch
# overlap sum is basis-converged (see _oneel_lattice_opts); tight
# bases are left at the caller's cutoff. V_ne keeps its own (larger,
# 1/r-decay) nuclear cutoff via gauge_lat_opts below.
oneel_lat_opts = _oneel_lattice_opts(
system, basis, lat_opts,
rcut_strategy=rcut_strategy, k_points_cart=kpoints_cart, plog=plog,
)
with plog.stage(
"integrals_lattice",
detail=f"S/T at cutoff {oneel_lat_opts.cutoff_bohr:.2f} bohr, "
f"V at cutoff {lat_opts.cutoff_bohr:.2f} bohr",
):
S_lat = compute_overlap_lattice(basis, system, oneel_lat_opts)
T_lat = compute_kinetic_lattice(basis, system, oneel_lat_opts)
from .periodic_rhf_gdf import _gauge_lat_opts_for_v_ne_and_e_nuc
from .periodic_v_ne import compute_nuclear_lattice_dispatch
# Always use Ewald-3D gauge for V_ne and e_nuc -- J/K are
# built via build_periodic_fock_ewald3d_k (Ewald gauge).
gauge_lat_opts = _gauge_lat_opts_for_v_ne_and_e_nuc(lat_opts, system)
V_lat = compute_nuclear_lattice_dispatch(basis, system, gauge_lat_opts)
# ---- Per-k S(k), Hcore(k), canonical orthog X(k) -------------
S_k: List[np.ndarray] = []
Hcore_k: List[np.ndarray] = []
X_k: List[np.ndarray] = []
n_kept_k: List[int] = []
for k_idx in range(n_k):
k_arr = kpoints_cart[k_idx]
Sk = np.asarray(bloch_sum(S_lat, k_arr))
Tk = np.asarray(bloch_sum(T_lat, k_arr))
Vk = np.asarray(bloch_sum(V_lat, k_arr))
Sk = 0.5 * (Sk + Sk.conj().T)
Hk = 0.5 * ((Tk + Vk) + (Tk + Vk).conj().T)
_gdf_overlap_preflight(
Sk,
plog=plog,
label=f"S(k={k_idx}, k_cart={k_arr.round(4).tolist()})",
basis=basis,
)
Xk, n_kept = _canonical_orthogonalizer_complex(
Sk,
linear_dep_threshold,
normalize_diag_first=True,
)
if n_occ > n_kept:
raise RuntimeError(
"run_krhf_periodic_gdf: canonical orthogonalisation at "
f"k = {k_arr} dropped too many directions "
f"(n_occ={n_occ}, n_kept={n_kept}); loosen "
"linear_dep_threshold or pick a less redundant basis."
)
S_k.append(Sk)
Hcore_k.append(Hk)
X_k.append(Xk)
n_kept_k.append(n_kept)
dftu_sites = list(dft_plus_u_sites or ())
dftu_sites_cxx = []
dftu_ao_groups: List[List[int]] = []
if dftu_sites:
from ._vibeqc_core import _HubbardSiteCxx
from .dft_plus_u import ao_group_indices
ao_groups_map = ao_group_indices(basis)
for site in dftu_sites:
key = (int(site.atom_index), int(site.l))
if key not in ao_groups_map:
raise ValueError(
f"HubbardSite(atom_index={site.atom_index}, l={site.l}) "
"has no AOs in the basis."
)
dftu_sites_cxx.append(
_HubbardSiteCxx(
int(site.atom_index),
int(site.l),
float(site.U_eff_hartree),
)
)
dftu_ao_groups.append(list(ao_groups_map[key]))
plog.info(
"DFT+U: closed-shell multi-k GDF projector active "
f"({len(dftu_sites_cxx)} site(s))"
)
e_nuc = float(nuclear_repulsion_per_cell(system, gauge_lat_opts))
# Route HF/hybrid multi-k to the exxdiv-corrected compcell GDF path.
#
# The exxdiv Madelung exchange-divergence K-shift lives ENTIRELY inside
# the use_compcell=True branch below; the use_compcell=False Ewald-3D-K
# Fock builder does not apply it. For HF/hybrid (alpha > 0) that leaves
# the finite-k-mesh exchange divergence uncorrected, so the energy is
# wrong by the Madelung shift (catastrophically -- +17 Ha -- on sparse
# meshes). The compcell + exxdiv='ewald' path is PySCF-µHa-correct, so we
# route HF/hybrid there automatically rather than return wrong energies.
# Pure DFT (alpha == 0, no HF exchange) keeps the cheaper, correct
# Ewald-3D path. [Maintainer decision 2026-06-04: route to correct path.]
if not use_compcell and alpha > 0.0:
use_compcell = True
plog.info(
"multi-k HF/hybrid exchange: routing to the compcell GDF path "
"with exxdiv='ewald' (the use_compcell=False Ewald-3D-K path "
"omits the exxdiv Madelung exchange-divergence correction)."
)
# Route pure-DFT multi-k on dim<3 (vacuum-padded wire/slab) to the
# cached-Lpq GDF Hartree as well. The EWALD_3D J is 3D-only: its
# analytic-FT kernel raises on dim<3, so ewald_3d_j_blocks silently
# degrades to the diagnostic FFT-Poisson grid backend there -- wrong
# on vacuum-padded low-D cells (the 2026-07-09 Finding-§4 residual:
# dim=1 H2-chain (3,1,1) PBE reported -3.0897 Ha/cell where its own
# stored density evaluates to -3.7814 under the same functional and
# the variational minimum sits at -3.8947; the Lpq J lands at the
# real-Γ direct route's minimum to ~1e-10). HF/hybrids already ride
# the cached-Lpq path via the exxdiv routing above; this closes the
# same gap for the alpha == 0 KS branch. Pure DFT on dim=3 keeps the
# cheaper EWALD_3D path (validated: 29 uHa vs external PySCF at a
# non-TRIM mesh, tests/test_krks_gdf_xc_density.py).
if not use_compcell and is_ks and int(system.dim) != 3:
use_compcell = True
plog.info(
"multi-k pure-KS on dim<3: routing to the cached-Lpq GDF "
"Hartree (the EWALD_3D J falls back to the diagnostic "
"FFT-Poisson grid backend on dim<3, which mis-sums the "
"vacuum-padded Coulomb)."
)
# ---- Aux basis + Lpq(q) cache (GDF density-fitting path) -------
# When use_compcell=True, build per-(k_i,k_j) Lpq once and contract
# J/K from the cached cderi each iteration (like PySCF's GDF).
# When use_compcell=False, the aux is built but not used -- the
# per-iteration Fock goes through EWALD_3D real-space J/K instead.
mol = system.unit_cell_molecule()
with plog.stage("aux_basis", detail=aux_name):
aux = make_aux_basis_set(
mol,
aux_name=aux_name,
drop_eta=float(aux_drop_eta),
)
plog.info(f"aux basis: {aux_name} ({aux.nbasis} BFs / {aux.nshells} shells)")
# Per-pair Lpq cache: lpq_cache[(ki, kj)] = Lpq(k_i,k_j) tensor.
lpq_cache: Dict[Tuple[int, int], np.ndarray] = {}
if use_compcell:
# The off-diagonal (k_i != k_j) Lpq tensors exist solely for
# the GDF exchange contraction; the COSX exchange backend
# works in real space and never touches them -- only the
# diagonal J pairs are built then (the structural setup +
# memory saving of the COSX route).
need_k_pairs = alpha != 0.0 and k_exchange == "gdf"
# Per-pair Lpq builder, dispatched on gdf_method. 'rsgdf'
# is the production default: the (k_i,k_j)-ket-resolved all-FT
# Bloch-pair route (build_lpq_bloch_native_fft), validated at
# µHa parity vs PySCF (handovers/HANDOVER_GDF_OUTSTANDING.md Sec. 1), and
# the consistent partner for k_exchange='cosx' (RSGDF-J +
# COSX-K is single-gauge -- the LR complement shares the
# pair-FT conventions). 'compcell' (Sun-2017 compensated
# charges + AFT) remains selectable for its own development:
# its cderi is q-only -- exact only in the vacuum-box limit,
# catastrophically wrong on tight ionic cells (M3b-5 finding
# #1 / handovers/HANDOVER_RIJCOSX_M3A.md Sec. M3b-4b zone-edge deviation).
if gdf_method == "rsgdf":
aux_modrho = make_modrho_aux_basis(aux, mol)
def _build_pair_lpq(ki: np.ndarray, kj: np.ndarray):
return build_lpq_bloch_native_fft(
system,
basis,
aux_modrho,
ki,
kj,
ke_cutoff=float(rsgdf_ke_cutoff),
tail_ke_cutoff=(
float(rsgdf_tail_ke_cutoff)
if rsgdf_tail_ke_cutoff is not None
else None
),
lat_opts=lat_opts,
linear_dep_thr=float(gdf_linear_dep_threshold),
fit_screen_threshold=float(fit_screen_threshold),
progress=plog,
)
elif gdf_method == "mdf":
def _build_pair_lpq(ki: np.ndarray, kj: np.ndarray):
# Mixed Density Fitting: ket-resolved combined cderi
# [L_gauss; cderi_pw]; the multi-k J/K take the conjugate.
return build_lpq_bloch_mdf(
system,
basis,
aux,
ki,
kj,
molecule=mol,
lat_opts=lat_opts,
linear_dep_thr=float(gdf_linear_dep_threshold),
compcell_eta=float(compcell_eta),
mdf_ke_cutoff=float(mdf_ke_cutoff),
rcut_strategy=rcut_strategy,
rcut_precision=float(rcut_precision),
)
else:
def _build_pair_lpq(ki: np.ndarray, kj: np.ndarray):
return build_lpq_bloch_compcell(
system,
basis,
aux,
kj - ki, # momentum transfer q = k_j - k_i
molecule=mol,
lat_opts=lat_opts,
linear_dep_thr=float(gdf_linear_dep_threshold),
compcell_eta=float(compcell_eta),
apply_aft_correction=bool(apply_aft_correction),
aft_ft_convention=str(aft_ft_convention),
aft_precision=float(aft_precision),
rcut_strategy=rcut_strategy,
rcut_precision=float(rcut_precision),
)
# Fail early rather than OOM-killing a doomed run: gate the dense
# per-pair Lpq cache peak against available RAM (prompt-75 pattern).
_preflight_gdf_lpq_memory(
plog,
n_basis=basis.nbasis,
n_aux=aux.nbasis,
n_kpoints=n_k,
need_k_pairs=need_k_pairs,
open_shell=False,
route_label=label,
options=opts,
)
with plog.stage(
"gdf_cderi",
detail=f"per-pair Lpq for {n_k} k-points ({gdf_method})",
):
n_pairs = n_k * n_k if need_k_pairs else n_k
plog.info(
f"Building per-pair Lpq cache ({n_pairs} pairs, "
f"{gdf_method})..."
)
for i in range(n_k):
ki = kpoints_cart[i]
if need_k_pairs:
# Hybrid: build all (k_i, k_j) pairs for exchange.
for j in range(n_k):
lpq_cache[(i, j)] = _build_pair_lpq(
ki, kpoints_cart[j]
)
else:
# Pure DFT / J-only, or COSX exchange backend:
# only diagonal pairs needed.
lpq_cache[(i, i)] = _build_pair_lpq(ki, ki)
n_fit = lpq_cache[(0, 0)].shape[0] if lpq_cache else 0
plog.info(
f"Lpq cache built: {len(lpq_cache)} pairs, "
f"{n_fit} fit vectors, "
f"shape=({n_fit}, {basis.nbasis}, {basis.nbasis})"
)
# ---- Multi-k COSX exchange bridge (k_exchange='cosx') ----------
# SCF-invariant setup: truncated cell list, per-relative-shift
# analytic-integral caches, periodic Becke grid, Q-junction. Per
# iteration the bridge folds D(k) -> D(g), runs one real-space
# K(g) build (mesh-size independent), and Bloch-folds to K(k).
cosx_bridge = None
if k_exchange == "cosx" and use_compcell and alpha != 0.0:
from .periodic_cosx_k import KPointCosxK
# Range-separation parameter for the SR/LR exchange split
# (M3b-4): the real-space erfc-SR part must be dead at both
# the cell-list cutoff and the BvK half-super-period (alias
# boundary); the smooth LR-erf complement is built in
# reciprocal space and restores the full kernel exactly, so
# w only tunes the split -- erfc(5) ≈ 1.5e-12 sets the reach.
lat_np = np.asarray(system.lattice, dtype=float)
mesh_dims = _mesh_tuple_for_system(system, kmesh)
half_supers = [
0.5 * mesh_dims[i] * float(np.linalg.norm(lat_np[:, i]))
for i in range(int(system.dim))
]
sr_reach = min(min(half_supers), float(lat_opts.cutoff_bohr))
cosx_omega = 5.0 / sr_reach
with plog.stage("cosx_caches"):
cosx_bridge = KPointCosxK(
basis, system, lat_opts=lat_opts, omega=cosx_omega
)
plog.info(
"K backend: multi-k COSX, range-separated "
f"(SR: real-space erfc, w = {cosx_omega:.3f} bohr⁻¹; "
"LR: reciprocal-space erf complement; "
f"cells={len(cosx_bridge.cells)}, "
f"deltas={cosx_bridge.caches.n_deltas})"
)
if gdf_method == "rsgdf":
plog.info(
" K/J pairing: single-gauge (RSGDF J + COSX K share "
"the Bloch pair-FT conventions). Validated at "
"sub-mHa backend parity (0.024-0.027 mHa vs "
"k_exchange='gdf' on the chain anchor across "
"(1,1,2)-(1,1,6)) with clean SCF convergence through "
"(1,1,8) -- M3b-4c."
)
plog.info(
" Scope criterion: the SR exchange range (≈5/w) "
"PLUS the basis pair extent must fit inside the BvK "
"half-super-period. Diffuse-basis tight cells at "
"small meshes violate it (measured: LiH/sto-3g "
"(2,2,2): w-invariance broken at 2e-2, SCF stall, "
"534 mHa off -- Li 2sp extent ~15 bohr vs 5.46-bohr "
"half-super-period); the BvK-consistent SR cell "
"summation is M3b-6 (handovers/HANDOVER_RIJCOSX_M3A.md "
"Sec. M3b-5)."
)
else:
plog.info(
" WARNING: k_exchange='cosx' with "
"gdf_method='compcell' is a MIXED-GAUGE Fock: the "
"COSX K is matrix-level exact (~1e-4 vs the "
"independent RSGDF exchange), but on vacuum-padded "
"systems the compcell-J tensors carry a flagged "
"zone-edge deviation (handovers/HANDOVER_RIJCOSX_M3A.md "
"Sec. M3b-4b, escalated to the GDF route), and the "
"inconsistency degrades dense-k-mesh SCF convergence "
"(100-iter stalls at (1,1,6)+ on the chain anchor). "
"Pair with gdf_method='rsgdf' for the validated "
"single-gauge combination."
)
# ---- Initial guess: Hcore diagonalisation per k --------------
C_k: List[np.ndarray] = []
eps_k: List[np.ndarray] = []
for i in range(n_k):
Ci, ei = _diag_in_orth_basis(Hcore_k[i], X_k[i])
C_k.append(Ci.astype(complex))
eps_k.append(ei)
occ_k, fermi_level, entropy = _occupations_per_k(
eps_k,
weights,
n_elec,
smearing_T,
n_occ,
bz_integration=bz_integration,
system=system,
kmesh=kmesh_bloch,
)
D_k = _density_from_orbitals(C_k, occ_k)
if initial_density_k is not None:
D_k = _normalise_initial_density_k(
initial_density_k,
n_k=n_k,
n_basis=basis.nbasis,
label="run_krhf_periodic_gdf",
)
plog.info("initial guess: READ (caller-supplied per-k density)")
else:
plog.info("initial guess: HCORE (per-k Hcore diagonalisation)")
if smearing_T > 0.0:
plog.info(
"smearing: Fermi-Dirac kBT = "
f"{smearing_T:.6g} Ha "
f"({_hartree_to_kelvin_temperature(smearing_T):.1f} K)"
)
# ---- SCF setup -----------------------------------------------
damping = float(opts.damping)
if not (0.0 <= damping < 1.0):
raise ValueError(
f"run_krhf_periodic_gdf: damping must be in [0, 1); got {damping}"
)
if fock_mixing_value != 0.0:
plog.info(
"fock mixing: CRYSTAL FMIXING "
f"{100.0 * fock_mixing_value:.1f}% "
"(previous Fock/KS matrix weight, applied per k)"
)
damper: Optional[DynamicDamping] = None
if bool(getattr(opts, "dynamic_damping", False)):
damper = DynamicDamping(
initial_alpha=damping,
alpha_min=float(getattr(opts, "dynamic_damping_min", 0.0)),
alpha_max=float(getattr(opts, "dynamic_damping_max", 0.95)),
)
use_diis = bool(opts.use_diis)
diis_start_iter = int(opts.diis_start_iter)
accel: Optional[MultiKPeriodicSCFAccelerator] = (
MultiKPeriodicSCFAccelerator(opts) if use_diis else None
)
# ---- Density-space mixer (Anderson / Broyden [+ Kerker]) ----------
# Ported from the multi-k EWALD_3D RKS driver (same machinery in
# periodic_density_mixing): mix the per-k density MATRICES, with the
# Kerker filter optionally preconditioning the residual on a
# plane-wave grid. When selected it REPLACES Fock-DIIS, linear
# damping, dynamic damping, and fock mixing (one accelerator owns
# the update; see the Ewald driver + CLAUDE.md Sec.7).
from .periodic_density_mixing import (
AndersonMixer as _AndersonMixer,
BroydenMixer as _BroydenMixer,
KerkerPreconditioner as _KerkerPreconditioner,
per_k_density_to_vector as _per_k_density_to_vector,
vector_to_per_k_density as _vector_to_per_k_density,
)
_mixer_key = (
None if density_mixer is None else str(density_mixer).strip().lower()
)
if _mixer_key in (None, "", "none", "diis"):
density_space_mixer = None
elif _mixer_key == "anderson":
density_space_mixer = _AndersonMixer(
depth=int(density_mixer_depth), beta=float(density_mixer_beta)
)
elif _mixer_key == "broyden":
density_space_mixer = _BroydenMixer(
depth=int(density_mixer_depth), beta=float(density_mixer_beta)
)
else:
raise ValueError(
f"run_krhf_periodic_gdf: density_mixer={density_mixer!r} is not "
f"recognised; expected one of None, 'diis', 'anderson', 'broyden'."
)
kerker_precond: Optional[_KerkerPreconditioner] = None
if density_space_mixer is not None:
if density_mixer_kerker:
from .periodic_rhf_multi_k_ewald import _g0_block
kerker_precond = _KerkerPreconditioner(
basis,
system,
_g0_block(S_lat),
k0=float(kerker_k0),
strength=float(kerker_strength),
cutoff_ha=float(kerker_cutoff_ha),
)
use_diis = False
accel = None
damper = None
damping = 0.0
fock_mixing_value = 0.0
plog.info(
f"density mixer: {density_mixer!r} "
f"(depth={int(density_mixer_depth)}, "
f"beta={float(density_mixer_beta)}"
+ (
f"; Kerker k0={float(kerker_k0)}, "
f"strength={float(kerker_strength)}, "
f"cutoff={float(kerker_cutoff_ha)} Ha"
if kerker_precond is not None
else ""
)
+ ") -- Fock-DIIS, damping and fock mixing disabled"
)
elif density_mixer_kerker:
raise ValueError(
"run_krhf_periodic_gdf: density_mixer_kerker=True requires "
"density_mixer='anderson' or 'broyden' -- Kerker preconditions "
"the density mixer's residual, it is not a standalone "
"accelerator."
)
if level_shift != 0.0:
if warmup_cycles > 0:
cycle_word = "cycle" if warmup_cycles == 1 else "cycles"
plog.info(
f"level-shift warm-up: {warmup_cycles} {cycle_word} at "
f"{level_shift:.3f} Ha (per k); restart unshifted afterwards"
)
else:
plog.info(
f"level shift: {level_shift:.3f} Ha (per k) "
"applied at each diagonalization"
)
plog.banner(f"SCF ({label} multi-k, native GDF)")
plog.info(" iter energy (Ha) dE ||[F,DS]|| DIIS")
scf_trace: List[SCFIteration] = []
result = PeriodicKRHFGDFResult(
energy=0.0,
e_electronic=0.0,
e_nuclear=float(e_nuc),
n_iter=0,
converged=False,
mo_energies=[e.copy() for e in eps_k],
mo_coeffs=[C.copy() for C in C_k],
fock=[np.empty((0, 0), dtype=complex) for _ in range(n_k)],
overlap=[S.copy() for S in S_k],
hcore=[H.copy() for H in Hcore_k],
density=[D.copy() for D in D_k],
kpoints_cart=kpoints_cart.copy(),
kpoint_weights=weights.copy(),
scf_trace=scf_trace,
functional=func_name or None,
fock_mixing=fock_mixing_value,
level_shift=level_shift,
level_shift_warmup_cycles=warmup_cycles,
smearing_temperature=smearing_T,
fermi_level=float(fermi_level),
entropy=float(entropy),
occupations=[np.asarray(o, dtype=float) for o in occ_k],
aux_basis_name=aux_name,
n_aux=int(aux.nbasis),
backend=(
f"native-multi-k-gdf-{k_exchange}-"
f"{'rks' if is_ks else 'rhf'}"
),
)
F_prev_k: Optional[List[np.ndarray]] = None
D_prev_k: List[np.ndarray] = [D.copy() for D in D_k]
E_prev = 0.0
# ---- GDF J/K builders (cached Lpq contraction) -----------------
def _build_j_from_lpq(D_k_in: List[np.ndarray]) -> List[np.ndarray]:
"""Closed-shell Coulomb ``J(k_i)`` from the diagonal cderi blocks.
``J(k_i) = S_P L(k_i,k_i)_{P,muν} . r_P`` with the fitted
**total-density** coefficients
``r_P = S_j w_j S_{ls} L(k_j,k_j)*_{P,ls} D(k_j)_{sl}``.
The Hartree potential sees the k-summed (BZ-averaged) density --
building J(k_i) from D(k_i) alone is exact only when D(k) is
k-independent (vacuum-box limit; why H₂ passed and LiH landed
at -3.10 instead of -7.92 Ha before this was restored). The
diagonal (``q = 0``) cderi ``L(k,k)`` still depends on ``k`` for
tight cells, so the Hartree matrix is built per-k from the
shared ``r_P``.
"""
naux = lpq_cache[(0, 0)].shape[0]
rho = np.zeros(naux, dtype=complex)
for j in range(n_k):
L_jj = lpq_cache[(j, j)]
# BZ-summed fitted density r_P = S_k w_k tr(L(k,k).D(k)). The
# diagonal cderi L(k,k) is NOT conjugated here -- matching PySCF
# ``get_j_kpts`` (df_jk.py: ``rho_L = S_k L(k)_{L,pq} D(k)_{qp}``,
# no conjugate). L(k,k) is genuinely complex for k != 0 on tight
# cells (inter-cell R!=0 overlap; see build_lpq_bloch_native_fft),
# so a spurious ``L_jj.conj()`` is a no-op only for real
# (vacuum-box / cubic) cderi and otherwise mis-contracts the
# Coulomb -- over-binding E_J by Madelung-scale (uniform H-chain
# metal, (1,1,8): 0.54 Ha; cubic-box insulators were unharmed,
# which masked the bug behind the H₂ (2,1,1) gate).
rho = rho + float(weights[j]) * np.einsum(
"Pls,sl->P", L_jj, np.asarray(D_k_in[j]), optimize=True
)
J_k = []
for i in range(n_k):
Ji = np.einsum("P,Pmn->mn", rho, lpq_cache[(i, i)], optimize=True)
Ji = 0.5 * (Ji + Ji.conj().T)
J_k.append(Ji)
return J_k
def _build_k_from_lpq(D_k_in: List[np.ndarray]) -> List[np.ndarray]:
"""Build K matrices from cached Lpq. All (k_i, k_j) pairs."""
K_k = [
np.zeros((basis.nbasis, basis.nbasis), dtype=complex) for _ in range(n_k)
]
for i in range(n_k):
for j in range(n_k):
Lpq = lpq_cache[(i, j)]
Dj = np.asarray(D_k_in[j])
# K(k_i) += (1/N_k) * S_L L_L(k_i,k_j) @ D(k_j) @ L_L(k_i,k_j).H
# tmp[L,p,s] = S_r L[L,p,r] D[r,s]
tmp = np.einsum("Lpr,rs->Lps", Lpq, Dj)
K_k[i] += (1.0 / n_k) * np.einsum("Lps,Lqs->pq", tmp, Lpq.conj())
for i in range(n_k):
K_k[i] = 0.5 * (K_k[i] + K_k[i].conj().T)
return K_k
# Iteration-invariant EWALD_3D Hartree-J cache: the per-cell AO-pair
# FT dominates every build_periodic_fock_ewald3d_k call (~93 % of a
# multi-k GDF SCF iteration profiled on LiH (1,1,2): 557 s of 601 s
# total were pair-FT recomputation), and it depends only on
# (basis, cells, mesh) -- never on the density. Same mechanism the
# multi-k EWALD_3D RHF/RKS drivers already hoist
# (make_ewald_3d_lattice_j_cache; bit-identical contraction).
#
# Built ONLY when the EWALD_3D J/K branch below will actually run:
# the cached-Lpq GDF branch (use_compcell=True, every HF/hybrid run
# after the auto-routing above) never touches it, and the cache is
# the dominant memory of a small-cell multi-k GDF run -- the per-cell
# AO-pair FT is (n_cells, nbf, nbf, n_G) on its own VIBEQC_J_EWALD3D_KE
# mesh (c-diamond primitive sto-3g, 177 cells x 10^2 AO-pairs x 10417
# G-points = 2.95 GB, ~5.9 GB transient with the scale copy; measured
# 6.13 GB peak RSS -> 2026-07-09, HANDOVER_GDF_FIT_SCREENING.md). It
# is also ke-independent of rsgdf_ke_cutoff, which is why it was
# misattributed to the GDF fit tensors in the original finding.
_ewald_j_cache = (
make_ewald_3d_lattice_j_cache(
basis, system, cells, lattice_opts=lat_opts
)
if not (use_compcell and lpq_cache)
else None
)
for it in range(1, max_iter + 1):
if damper is not None:
damping = damper.alpha
if warmup_cycles > 0 and it == warmup_cycles + 1:
if accel is not None:
accel = MultiKPeriodicSCFAccelerator(opts)
F_prev_k = None
plog.info("restart: unshifted Fock with fresh DIIS history (per k)")
if _ls_schedule:
active_level_shift = level_shift_at_iter(
level_shift, warmup_cycles, _ls_schedule, max_iter, it
)
else:
active_level_shift = (
level_shift
if (level_shift != 0.0 and (warmup_cycles == 0 or it <= warmup_cycles))
else 0.0
)
diis_active = use_diis and it >= diis_start_iter
# Density damping (per k, in AO basis).
if it == 1 or damping == 0.0 or diis_active:
D_used = [D.copy() for D in D_k]
else:
D_used = [
damping * Dp + (1.0 - damping) * Dn for Dp, Dn in zip(D_prev_k, D_k)
]
# ---- J + K build -------------------------------------------
if use_compcell and lpq_cache:
# True GDF: contract cached Lpq per iteration (J always;
# K per the selected exchange backend).
J_k = _build_j_from_lpq(D_used)
if alpha != 0.0:
if cosx_bridge is not None:
K_k = cosx_bridge.k_matrices(
D_used, list(kpoints_cart),
lr_complement=True,
)
else:
K_k = _build_k_from_lpq(D_used)
else:
K_k = [np.zeros_like(J_k[0]) for _ in range(n_k)]
# Apply exxdiv='ewald' Madelung correction to K.
# k-mesh-aware (Born-von-Kármán supercell) Madelung -- NOT the
# primitive-cell ξ, which over-counts the exxdiv K-shift by
# Nk^(1/3) and over-binds the multi-k energy (LiH (2,2,2):
# -592 mHa). See _madelung_for_kmesh.
#
# The shift's energy contribution enters ONCE, through the
# shifted K inside ``E_elec = Tr[D.Hcore] + 1/2.Tr[D.F_2e]``
# below -- do NOT also add ``exxdiv_ewald_energy_shift`` to
# E_total (that double-counts the Madelung correction; on
# H₂/12-bohr (2,1,1) the double-count over-binds by
# ~150 mHa vs the published -1.12013988 Ha).
if alpha != 0.0:
madelung = _madelung_for_kmesh(system, kmesh_bloch.mesh)
K_k = list(apply_exxdiv_ewald_to_K(K_k, S_k, D_used, madelung))
F_k = [Jk - 0.5 * alpha * Kk for Jk, Kk in zip(J_k, K_k)]
F_2e_k = [np.asarray(f).copy() for f in F_k]
# XC on the full real-space finite-torus density, Bloch-folded
# to each k. A single Γ AO matrix is only a vacuum-limit shortcut.
V_xc_k = None
E_xc = 0.0
if is_ks:
E_xc, V_xc_k = _build_xc_k_from_density(
basis=basis,
system=system,
grid=grid,
func=func,
density_k=D_used,
kmesh_bloch=kmesh_bloch,
cells=cells,
kpoints_cart=kpoints_cart,
lat_opts=lat_opts,
)
E_coulomb = 0.5 * sum(
float(weights[i]) * float(np.real(np.trace(D_used[i] @ J_k[i])))
for i in range(n_k)
)
E_hf_K = (
-0.25
* alpha
* sum(
float(weights[i]) * float(np.real(np.trace(D_used[i] @ K_k[i])))
for i in range(n_k)
)
if alpha != 0.0
else 0.0
)
plog.info(
f"J backend: native multi-k GDF (cached Lpq, "
f"{len(lpq_cache)} pairs); K backend: "
+ ("multi-k COSX (real-space K(g) + Bloch fold)"
if cosx_bridge is not None
else "multi-k GDF (k-pair Lpq contraction)")
)
# Fold Hcore (+ V_xc for RKS) into the Fock that gets
# extrapolated + diagonalised. The M3b-4b/4c refactor moved
# XC into the branches and left this fold in the EWALD_3D
# branch only -- every compcell SCF then diagonalised the
# bare 2e Fock (J - 1/2aK), walking to a spurious fixed point
# (LiH (2,2,2): -2.96 Ha instead of -7.92, converged=True).
for i in range(n_k):
Fi = F_k[i] + Hcore_k[i]
if V_xc_k is not None:
Fi = Fi + V_xc_k[i]
Fi = 0.5 * (Fi + Fi.conj().T)
F_k[i] = Fi
else:
# ---- J + K via EWALD_3D gauge (legacy path) --------------
# Fold the exact density used for this iteration. This is
# essential for smearing and damping: rebuilding from ``C_k``
# and hard ``n_occ`` would silently feed an integer-Aufbau
# density to the Fock builder while reporting fractional
# occupations.
D_real = _real_space_density_from_per_k_density(
D_used,
kmesh_bloch,
cells,
)
F_k = build_periodic_fock_ewald3d_k(
basis,
system,
D_real,
omega=0.5,
k_points_cart=[np.asarray(k) for k in kpoints_cart],
Hcore_k=None,
lattice_opts=lat_opts,
exchange_scale=alpha,
j_cache=_ewald_j_cache,
)
# XC on the full real-space finite-torus density, Bloch-folded
# to each k.
V_xc_k = None
E_xc = 0.0
if is_ks:
E_xc, V_xc_k = _build_xc_k_from_density(
basis=basis,
system=system,
grid=grid,
func=func,
density_k=D_used,
kmesh_bloch=kmesh_bloch,
cells=cells,
kpoints_cart=kpoints_cart,
lat_opts=lat_opts,
)
# Save F_2e before adding V_xc (for energy decomposition).
F_2e_k = [np.asarray(f).copy() for f in F_k]
for i in range(n_k):
Fi = F_k[i] + Hcore_k[i]
if V_xc_k is not None:
Fi = Fi + V_xc_k[i]
Fi = 0.5 * (Fi + Fi.conj().T)
F_k[i] = Fi
# ---- Energy decomposition (J-only Fock for E_J, E_K). ------
E_coulomb = 0.0
E_hf_K = 0.0
if alpha != 0.0:
F_J_k = build_periodic_fock_ewald3d_k(
basis,
system,
D_real,
omega=0.5,
k_points_cart=[np.asarray(k) for k in kpoints_cart],
Hcore_k=None,
lattice_opts=lat_opts,
exchange_scale=0.0,
j_cache=_ewald_j_cache,
)
E_J_val = 0.0
E_K_val = 0.0
for i in range(n_k):
w = float(weights[i])
J_k_i = np.asarray(F_J_k[i])
K_k_i = 2.0 * (J_k_i - F_2e_k[i]) / alpha
K_k_i = 0.5 * (K_k_i + K_k_i.conj().T)
E_J_val += w * float(np.real(np.trace(D_used[i] @ J_k_i)))
E_K_val += w * float(np.real(np.trace(D_used[i] @ K_k_i)))
E_coulomb = 0.5 * E_J_val
E_hf_K = -0.25 * alpha * E_K_val
# ---- DFT+U closed-shell Fock contribution ------------------
# D_used is the total closed-shell density that built this
# iteration's Coulomb/XC Fock. The Dudarev kernel is per spin, so
# pass P_sigma(k) = 1/2 P_total(k), then add the same k-independent
# V_AO_s to alpha and beta through a single closed-shell Fock shift.
e_dft_plus_u = 0.0
if dftu_sites_cxx:
from ._vibeqc_core import (
_compute_dft_plus_u_multi_k_per_spin_cxx,
)
P_sigma_k = [
0.5 * np.asarray(D, dtype=np.complex128) for D in D_used
]
E_sigma, V_AO = _compute_dft_plus_u_multi_k_per_spin_cxx(
dftu_sites_cxx,
dftu_ao_groups,
[np.asarray(S, dtype=np.complex128) for S in S_k],
P_sigma_k,
list(weights),
)
e_dft_plus_u = 2.0 * float(E_sigma)
V_AO_c = np.asarray(V_AO, dtype=np.complex128)
for i in range(n_k):
F_u = S_k[i] @ V_AO_c @ S_k[i]
F_k[i] = F_k[i] + F_u
F_k[i] = 0.5 * (F_k[i] + F_k[i].conj().T)
# ---- Energy (E_elec = Tr[D.Hcore] + 0.5 Tr[D.F_2e]). -------
E_elec = 0.0
for i in range(n_k):
w = float(weights[i])
Di = D_used[i]
Hi = Hcore_k[i]
Fi = F_2e_k[i]
E_elec += w * float(
np.real(np.trace(Di @ Hi)) + 0.5 * np.real(np.trace(Di @ Fi))
)
E_total = E_elec + E_xc + float(e_nuc) + e_dft_plus_u
free_energy = E_total - smearing_T * entropy
# ---- Convergence -------------------------------------------
grad_k: List[np.ndarray] = []
grad_norm_sq = 0.0
for i in range(n_k):
FDS = F_k[i] @ D_used[i] @ S_k[i]
err = FDS - FDS.conj().T
grad_k.append(err)
grad_norm_sq += float(weights[i]) * float(np.linalg.norm(err) ** 2)
grad_norm = float(np.sqrt(grad_norm_sq))
dE = free_energy - E_prev
scf_trace.append(
SCFIteration(
iter=it,
energy=float(free_energy),
delta_e=float(dE if it > 1 else 0.0),
grad_norm=float(grad_norm),
diis_subspace=(accel.subspace_size if accel is not None else 0),
)
)
plog.iteration(
it,
energy=float(free_energy),
dE=float(dE if it > 1 else 0.0),
grad=float(grad_norm),
diis=(accel.subspace_size if accel is not None else 0),
)
converged = (
it > 1
and (warmup_cycles == 0 or it > warmup_cycles)
and abs(dE) < float(opts.conv_tol_energy)
and grad_norm < float(opts.conv_tol_grad)
)
# ---- SCF-accelerator extrapolation, FMIXING, level shift ----
# The full {DIIS, KDIIS, EDIIS, EDIIS_DIIS, ADIIS} family +
# dynamic_damping is wired here (M4); GDF keeps density per-k
# natively, so ``density_k_list`` is ``D_used`` with no Bloch
# sum. See ``MultiKPeriodicSCFAccelerator`` in
# ``periodic_scf_accelerators.py`` for the per-mode dispatch.
if accel is not None:
F_ex_list = accel.extrapolate_rhf(
F_k,
error_k_list=grad_k,
density_k_list=D_used,
energy=E_total,
mo_coeffs_k_list=C_k,
n_occ=n_occ,
weights=list(weights),
cells=cells,
kpoints=list(kpoints_cart),
)
if diis_active:
F_k = F_ex_list
if fock_mixing_value != 0.0:
if F_prev_k is not None:
F_mixed: List[np.ndarray] = []
for i in range(n_k):
Fmix = (1.0 - fock_mixing_value) * F_k[
i
] + fock_mixing_value * F_prev_k[i]
F_mixed.append(0.5 * (Fmix + Fmix.conj().T))
F_k = F_mixed
F_prev_k = [F.copy() for F in F_k]
# Per-k Saunders-Hillier level shift (only at diagonalization).
F_diag = []
for i in range(n_k):
if active_level_shift != 0.0:
Fi_shifted = (
F_k[i]
+ active_level_shift * S_k[i]
- (active_level_shift / 2.0) * (S_k[i] @ D_used[i] @ S_k[i])
)
Fi_shifted = 0.5 * (Fi_shifted + Fi_shifted.conj().T)
F_diag.append(Fi_shifted)
else:
F_diag.append(F_k[i])
# ---- Diagonalise per k + occupations + density -------------
C_new: List[np.ndarray] = []
eps_new: List[np.ndarray] = []
for i in range(n_k):
Ci, ei = _diag_in_orth_basis(F_diag[i], X_k[i])
C_new.append(Ci.astype(complex))
eps_new.append(ei)
occ_k, fermi_level, entropy = _occupations_per_k(
eps_new,
weights,
n_elec,
smearing_T,
n_occ,
bz_integration=bz_integration,
system=system,
kmesh=kmesh_bloch,
)
D_new = _density_from_orbitals(C_new, occ_k)
# ---- Density-space mixing (Anderson / Broyden [+ Kerker]) ----
# Same construction as the multi-k EWALD_3D driver: mix the per-k
# density matrices; the Kerker filter preconditions the residual
# (D_out' = D_in + K(D_out - D_in)), leaving the fixed point
# unchanged.
if density_space_mixer is not None:
if kerker_precond is not None:
_resid = [
np.asarray(Dn) - np.asarray(Du)
for Dn, Du in zip(D_new, D_used)
]
_resid = kerker_precond.precondition(_resid, weights)
_D_out_eff = [
np.asarray(Du) + r for Du, r in zip(D_used, _resid)
]
else:
_D_out_eff = D_new
_x_in = _per_k_density_to_vector(D_used)
_x_out = _per_k_density_to_vector(_D_out_eff)
_x_next = density_space_mixer.update(_x_in, _x_out)
D_new = _vector_to_per_k_density(_x_next, D_new)
D_prev_k = D_used
D_k = D_new
C_k = C_new
eps_k = eps_new
if damper is not None:
damper.update(free_energy)
E_prev = free_energy
# Update the result placeholder so partial-run callers see
# the last iter's state. We do this every iter (rather than
# only on converge) so a max_iter abort still yields useful
# numbers.
result.energy = E_total
result.e_electronic = E_elec
result.e_xc = E_xc
result.e_coulomb = E_coulomb
result.e_hf_exchange = E_hf_K
result.e_dft_plus_u = e_dft_plus_u
result.n_iter = it
result.mo_energies = [e.copy() for e in eps_new]
result.mo_coeffs = [C.copy() for C in C_new]
result.fock = [F.copy() for F in F_k]
result.density = [D.copy() for D in D_k]
result.fermi_level = float(fermi_level)
result.entropy = float(entropy)
result.free_energy = float(free_energy)
result.occupations = [np.asarray(o, dtype=float) for o in occ_k]
if converged:
result.converged = True
plog.converged(
n_iter=result.n_iter,
energy=result.energy,
converged=True,
)
if check_energy_sanity:
_check_energy_sanity(result, system, plog)
return result
result.converged = False
plog.converged(
n_iter=result.n_iter,
energy=result.energy,
converged=False,
)
if check_energy_sanity:
_check_energy_sanity(result, system, plog)
return result
@dataclass
class PeriodicKUHFGDFResult:
"""Result of :func:`run_kuhf_periodic_gdf` / :func:`run_kuks_periodic_gdf`.
Open-shell multi-k GDF: per-spin a/b per-k lists (length ``nkpts``,
complex Hermitian in AO basis) + the ``<S^2>`` spin-contamination
diagnostic. Mirrors :class:`PeriodicKRHFGDFResult` split by spin.
"""
energy: float
e_electronic: float
e_nuclear: float
n_iter: int
converged: bool
s_squared: float
s_squared_ideal: float
mo_energies_alpha: List[np.ndarray]
mo_coeffs_alpha: List[np.ndarray]
density_alpha: List[np.ndarray]
fock_alpha: List[np.ndarray]
mo_energies_beta: List[np.ndarray]
mo_coeffs_beta: List[np.ndarray]
density_beta: List[np.ndarray]
fock_beta: List[np.ndarray]
overlap: List[np.ndarray]
hcore: List[np.ndarray]
kpoints_cart: np.ndarray
kpoint_weights: np.ndarray
scf_trace: List[SCFIteration] = field(default_factory=list)
functional: Optional[str] = None
e_xc: float = 0.0
e_coulomb: float = 0.0
e_hf_exchange: float = 0.0
aux_basis_name: str = ""
n_aux: int = 0
backend: str = "native-multi-k-gdf-uhf"
# Smearing (open-shell: independent per-spin chemical potentials).
# All default to the no-smearing values so T = 0 runs are unchanged.
smearing_temperature: float = 0.0
fermi_level_alpha: float = 0.0
fermi_level_beta: float = 0.0
entropy: float = 0.0 # S/k_B per cell, summed over both spin channels
free_energy: float = 0.0 # Mermin A = E - T.(S_a + S_b)
occupations_alpha: List[np.ndarray] = field(default_factory=list)
occupations_beta: List[np.ndarray] = field(default_factory=list)
@property
def energy_per_cell_ha(self) -> float:
return float(self.energy)
def _multi_k_s_squared(
n_alpha: int,
n_beta: int,
C_alpha_k: List[np.ndarray],
C_beta_k: List[np.ndarray],
S_k: List[np.ndarray],
weights: Sequence[float],
*,
occ_alpha_k: Optional[Sequence[np.ndarray]] = None,
occ_beta_k: Optional[Sequence[np.ndarray]] = None,
) -> float:
"""<S^2> for a multi-k UHF/UKS determinant: the spin value
``Sz(Sz+1) + n_b`` minus the BZ-weighted a/b overlap
``S_k w_k S_ij n^a_i(k) n^b_j(k) |<a_i(k)|S(k)|b_j(k)>|^2``. At M=1 (a==b)
this is 0; for an integer-filled doublet it is 0.75.
``occ_alpha_k`` / ``occ_beta_k`` are the per-k fractional occupations (in
``[0, 1]``) from the smearing path. When supplied and non-empty (T > 0)
the a/b overlap is weighted by ``n^a_i n^b_j`` over **all** orbitals -- the
fractional-occupation (ensemble-UHF) generalisation of Szabo & Ostlund
Eq. 2.271. When absent or empty (T = 0 / integer filling) it falls back to
the first-``n_s`` hard cutoff with unit weights, which is the exact
integer-occupation value and stays **bit-identical** to the pre-smearing
path. ``n_alpha`` / ``n_beta`` are the conserved per-spin totals
(``S_k w_k S_i n^s_i = n_s`` exactly under the fixed-multiplicity
constraint), so the ``Sz(Sz+1) + n_b`` part is unchanged by fractional
occupation."""
diff = n_alpha - n_beta
s2 = 0.25 * diff * (diff + 2) + n_beta
smeared = (
occ_alpha_k is not None
and occ_beta_k is not None
and len(occ_alpha_k) > 0
)
if smeared:
for i in range(len(S_k)):
oa = np.asarray(occ_alpha_k[i], dtype=float)
ob = np.asarray(occ_beta_k[i], dtype=float)
# All-orbital a/b overlap weighted by the fractional per-spin
# occupations n^a_i n^b_j (reduces to the first-n_s unit-weight
# cutoff below when occ in {0, 1}).
M = C_alpha_k[i].conj().T @ S_k[i] @ C_beta_k[i]
s2 -= float(weights[i]) * float(
np.sum(np.outer(oa, ob) * np.abs(M) ** 2)
)
elif n_alpha > 0 and n_beta > 0:
for i in range(len(S_k)):
Ca = C_alpha_k[i][:, :n_alpha]
Cb = C_beta_k[i][:, :n_beta]
M = Ca.conj().T @ S_k[i] @ Cb
s2 -= float(weights[i]) * float(np.sum(np.abs(M) ** 2))
return float(s2)
def run_kuhf_periodic_gdf(
system: PeriodicSystem,
basis: BasisSet,
kmesh: Union[Sequence[int], KPoints, BlochKMesh] = (1, 1, 1),
options: Optional[Union[PeriodicRHFOptions, PeriodicKSOptions]] = None,
*,
functional: Optional[str] = None,
aux_basis: Optional[str] = None,
aux_drop_eta: float = 0.0,
linear_dep_threshold: float = 1e-7,
gdf_linear_dep_threshold: float = 1e-9,
compcell_eta: float = 1.0,
apply_aft_correction: bool = True,
aft_ft_convention: str = "libint",
aft_precision: float = 1e-10,
rcut_strategy: Optional[object] = "pyscf_auto",
rcut_precision: float = 1e-8,
gdf_method: str = "rsgdf",
rsgdf_ke_cutoff: float = 200.0,
fit_screen_threshold: float = 0.0,
mdf_ke_cutoff: float = 40.0,
k_exchange: str = "gdf",
check_energy_sanity: bool = True,
progress: Union[bool, ProgressLogger, None] = None,
verbose: Optional[int] = None,
) -> PeriodicKUHFGDFResult:
"""Open-shell (UHF / UKS) periodic multi-k SCF via native GDF.
The open-shell sibling of :func:`run_krhf_periodic_gdf`: spin-independent
per-``(kᵢ,kⱼ)`` ``Lpq`` cderi cache, Hartree ``J`` from the BZ-summed
**total** density ``S_j w_j (Da+Db)(k_j)``, per-spin exchange ``Ks``
from the same cache, and the ``exxdiv='ewald'`` supercell-Madelung
K-shift applied per spin. a/b occupations follow ``multiplicity`` via
per-k Aufbau (``na=(n_e+mult-1)//2``); at M=1 on a gapped cell this
reproduces :func:`run_krhf_periodic_gdf`. ``functional`` selects UKS
(native spin-polarised XC on the per-spin real-space finite-torus
density fold, Bloch-folded to every k -- the closed-shell multi-k
convention, :func:`_build_xc_k_from_density_uks` -- plus the
functional's HF-exchange fraction); ``None`` runs UHF. KS runs
default ``options`` to :class:`PeriodicKSOptions` (periodic-Becke
grid); UHF keeps :class:`PeriodicRHFOptions`.
Returns a :class:`PeriodicKUHFGDFResult`. ``smearing_temperature > 0``
enables per-spin Fermi-Dirac smearing with independent global chemical
potentials mu_a, mu_b across the BZ (:func:`apply_smearing_open_shell`),
Mermin free energy ``A = E - T(S_a + S_b)``, and fractional per-spin
occupations (``occupations_alpha/beta``); ``T = 0`` keeps the exact per-k
Aufbau, bit-identical to the pre-smearing driver.
"""
_reject_slab_dim(system, "run_kuhf_periodic_gdf")
from .periodic_scf_accelerators import MultiKPeriodicUHFAccelerator
from .periodic_uhf_ewald import _spin_squared # noqa: F401 (parity ref)
from .periodic_rhf_gdf import _gauge_lat_opts_for_v_ne_and_e_nuc
from .periodic_v_ne import compute_nuclear_lattice_dispatch
plog = resolve_progress(progress, verbose=verbose)
# KS runs must default to PeriodicKSOptions: PeriodicRHFOptions carries no
# ``use_periodic_becke``, so defaulting KUKS to it silently selected the
# molecular Becke grid while the closed-shell KRKS wrapper defaulted to the
# periodic-Becke grid -- the grid/density-convention split behind the
# -2.4 Ha KUKS(mult=1)-vs-KRKS failure of the first fold fix (2026-07-09).
opts = _options_or_default(options, is_ks=functional is not None)
lat_opts: LatticeSumOptions = opts.lattice_opts
func_name = functional or str(getattr(opts, "functional", "") or "")
is_ks = bool(func_name)
func = Functional(func_name, 2) if is_ks else None # spin-polarized
# Full-range-only K (see run_krhf_periodic_gdf).
reject_unscreened_range_separated(func, where="run_kuhf_periodic_gdf")
alpha = float(func.hf_exchange_fraction) if func is not None else 1.0
label = f"KUKS {func_name}" if is_ks else "KUHF"
if gdf_method not in ("compcell", "rsgdf", "mdf"):
raise ValueError(
f"run_kuhf_periodic_gdf: gdf_method must be 'compcell', "
f"'rsgdf', or 'mdf'; got {gdf_method!r}"
)
if float(fit_screen_threshold) < 0.0:
raise ValueError(
"run_kuhf_periodic_gdf: fit_screen_threshold must be >= 0; "
f"got {fit_screen_threshold}"
)
if float(fit_screen_threshold) > 0.0 and gdf_method != "rsgdf":
raise NotImplementedError(
"run_kuhf_periodic_gdf: fit_screen_threshold is implemented "
f"for gdf_method='rsgdf' only (got {gdf_method!r})."
)
if k_exchange not in ("gdf", "cosx"):
raise ValueError(
"run_kuhf_periodic_gdf: k_exchange must be 'gdf' or "
f"'cosx'; got {k_exchange!r}"
)
smearing_T = float(getattr(opts, "smearing_temperature", 0.0) or 0.0)
if smearing_T < 0.0:
raise ValueError(
"run_kuhf_periodic_gdf: smearing_temperature must be >= 0"
)
# Open-shell Fermi-Dirac smearing (M3): independent per-spin chemical
# potentials mu_a, mu_b via apply_smearing_open_shell. T = 0 keeps the
# exact per-k Aufbau path (bit-identical to the pre-smearing driver).
smear_opts = _SmearingOptions.from_legacy_kwarg(smearing_T)
n_elec = system.n_electrons()
mult = int(system.multiplicity)
if mult < 1:
raise ValueError(f"run_kuhf_periodic_gdf: multiplicity must be >= 1, got {mult}")
if (n_elec + mult - 1) % 2 != 0:
raise ValueError(
f"run_kuhf_periodic_gdf: n_electrons={n_elec} and multiplicity="
f"{mult} cannot be split into integer a/b occupations."
)
n_alpha = (n_elec + mult - 1) // 2
n_beta = (n_elec - mult + 1) // 2
kpoints_cart, weights = _kmesh_to_kpoints_weights(system, kmesh)
n_k = kpoints_cart.shape[0]
if isinstance(kmesh, BlochKMesh):
kmesh_bloch = kmesh
elif isinstance(kmesh, KPoints):
kmesh_bloch = kmesh.to_bloch_kmesh()
else:
mesh = _mesh_tuple_for_system(system, kmesh)
kmesh_bloch = _mp_native(system, list(mesh), [0, 0, 0], False)
cells = _direct_cells(system, lat_opts.cutoff_bohr)
aux_name = aux_basis or default_aux_for(basis.name)
plog.banner(f"run_kuhf_periodic_gdf {label} kmesh={n_k} k-points")
plog.info(
f"{label} multi-k GDF / aux={aux_name}, n_alpha={n_alpha}, "
f"n_beta={n_beta} (mult={mult}), alpha={alpha:g}"
)
# ---- Functional grid (UKS) ----------------------------------------
grid = None
if is_ks:
grid_options = getattr(opts, "grid", None) or GridOptions()
if bool(getattr(opts, "use_periodic_becke", False)):
grid = build_periodic_becke_grid(
system, grid_options=grid_options,
image_radius_bohr=float(getattr(opts, "becke_image_radius_bohr", 0.0)),
)
else:
grid = build_grid(system.unit_cell_molecule(), grid_options)
# ---- One-electron integrals (Ewald-3D gauge for V_ne/e_nuc) -------
# Basis-aware S/T overlap cutoff for diffuse bases (see
# _oneel_lattice_opts); tight bases unchanged. V_ne keeps the larger
# nuclear cutoff via gauge_lat_opts.
oneel_lat_opts = _oneel_lattice_opts(
system, basis, lat_opts,
rcut_strategy=rcut_strategy, k_points_cart=kpoints_cart, plog=plog,
)
with plog.stage(
"integrals_lattice",
detail=f"S/T cutoff {oneel_lat_opts.cutoff_bohr:.2f}, "
f"V cutoff {lat_opts.cutoff_bohr:.2f}",
):
S_lat = compute_overlap_lattice(basis, system, oneel_lat_opts)
T_lat = compute_kinetic_lattice(basis, system, oneel_lat_opts)
gauge_lat_opts = _gauge_lat_opts_for_v_ne_and_e_nuc(lat_opts, system)
V_lat = compute_nuclear_lattice_dispatch(basis, system, gauge_lat_opts)
S_k: List[np.ndarray] = []
Hcore_k: List[np.ndarray] = []
X_k: List[np.ndarray] = []
for k_idx in range(n_k):
k_arr = kpoints_cart[k_idx]
Sk = np.asarray(bloch_sum(S_lat, k_arr))
Tk = np.asarray(bloch_sum(T_lat, k_arr))
Vk = np.asarray(bloch_sum(V_lat, k_arr))
Sk = 0.5 * (Sk + Sk.conj().T)
Hk = 0.5 * ((Tk + Vk) + (Tk + Vk).conj().T)
_gdf_overlap_preflight(
Sk,
plog=plog,
label=f"S(k={k_idx})",
basis=basis,
)
Xk, n_kept = _canonical_orthogonalizer_complex(
Sk, linear_dep_threshold, normalize_diag_first=True
)
if max(n_alpha, n_beta) > n_kept:
raise RuntimeError(
f"run_kuhf_periodic_gdf: orthogonalisation at k={k_idx} kept "
f"{n_kept} directions; need >= {max(n_alpha, n_beta)}."
)
S_k.append(Sk)
Hcore_k.append(Hk)
X_k.append(Xk)
e_nuc = float(nuclear_repulsion_per_cell(system, gauge_lat_opts))
# ---- Per-(kᵢ,kⱼ) Lpq cderi cache (spin-independent) ---------------
mol = system.unit_cell_molecule()
aux = make_aux_basis_set(mol, aux_name=aux_name, drop_eta=float(aux_drop_eta))
if gdf_method == "rsgdf":
aux_modrho = make_modrho_aux_basis(aux, mol)
def _build_pair_lpq(ki: np.ndarray, kj: np.ndarray) -> np.ndarray:
return build_lpq_bloch_native_fft(
system, basis, aux_modrho, ki, kj,
ke_cutoff=float(rsgdf_ke_cutoff), lat_opts=lat_opts,
linear_dep_thr=float(gdf_linear_dep_threshold),
fit_screen_threshold=float(fit_screen_threshold),
progress=plog,
)
elif gdf_method == "mdf":
def _build_pair_lpq(ki: np.ndarray, kj: np.ndarray) -> np.ndarray:
# MDF: ket-resolved combined cderi [L_gauss; cderi_pw].
return build_lpq_bloch_mdf(
system, basis, aux, ki, kj, molecule=mol, lat_opts=lat_opts,
linear_dep_thr=float(gdf_linear_dep_threshold),
compcell_eta=float(compcell_eta),
mdf_ke_cutoff=float(mdf_ke_cutoff),
rcut_strategy=rcut_strategy, rcut_precision=float(rcut_precision),
)
else:
def _build_pair_lpq(ki: np.ndarray, kj: np.ndarray) -> np.ndarray:
return build_lpq_bloch_compcell(
system, basis, aux, kj - ki, molecule=mol, lat_opts=lat_opts,
linear_dep_thr=float(gdf_linear_dep_threshold),
compcell_eta=float(compcell_eta),
apply_aft_correction=bool(apply_aft_correction),
aft_ft_convention=str(aft_ft_convention),
aft_precision=float(aft_precision),
rcut_strategy=rcut_strategy, rcut_precision=float(rcut_precision),
)
lpq_cache: Dict[Tuple[int, int], np.ndarray] = {}
need_k_pairs = alpha != 0.0 and k_exchange == "gdf"
# Fail early rather than OOM-killing a doomed run: the dense per-pair Lpq
# cache built below is the multi-k open-shell GDF memory bottleneck. Abort
# with a route-specific diagnostic if it cannot fit (prompt-75 NiO KUKS).
_preflight_gdf_lpq_memory(
plog,
n_basis=basis.nbasis,
n_aux=aux.nbasis,
n_kpoints=n_k,
need_k_pairs=need_k_pairs,
open_shell=True,
route_label=label,
options=opts,
)
with plog.stage("gdf_cderi", detail=f"per-pair Lpq, {n_k} k ({gdf_method})"):
for i in range(n_k):
ki = kpoints_cart[i]
if need_k_pairs:
for j in range(n_k):
lpq_cache[(i, j)] = _build_pair_lpq(ki, kpoints_cart[j])
else:
lpq_cache[(i, i)] = _build_pair_lpq(ki, ki)
n_fit = lpq_cache[(0, 0)].shape[0]
plog.info(f"Lpq cache: {len(lpq_cache)} pairs, {n_fit} fit vectors")
# COSX is a spin-independent exchange operator builder: the same
# geometry/grid cache acts separately on D_alpha and D_beta. The
# returned matrices use the same G=0-dropped gauge as the GDF factors;
# the existing per-spin Ewald Madelung correction below is therefore
# applied exactly once for either exchange backend.
cosx_bridge = None
if k_exchange == "cosx" and alpha != 0.0:
from .periodic_cosx_k import KPointCosxK
lat_np = np.asarray(system.lattice, dtype=float)
mesh_dims = _mesh_tuple_for_system(system, kmesh)
half_supers = [
0.5 * mesh_dims[axis] * float(np.linalg.norm(lat_np[axis]))
for axis in range(int(system.dim))
]
sr_reach = min(min(half_supers), float(lat_opts.cutoff_bohr))
if sr_reach <= 0.0:
raise ValueError("run_kuhf_periodic_gdf: invalid COSX super-period")
cosx_omega = 5.0 / sr_reach
with plog.stage("cosx_caches"):
cosx_bridge = KPointCosxK(
basis, system, lat_opts=lat_opts, omega=cosx_omega
)
plog.info(
"K backend: open-shell multi-k COSX, range-separated "
f"omega={cosx_omega:.3f} bohr^-1"
)
# ---- GDF J/K builders (copied from run_krhf_periodic_gdf) ---------
def _build_j_from_lpq(D_k_in: List[np.ndarray]) -> List[np.ndarray]:
naux = lpq_cache[(0, 0)].shape[0]
rho = np.zeros(naux, dtype=complex)
for j in range(n_k):
L_jj = lpq_cache[(j, j)]
# No conjugate on L(k,k) -- see the matching note in
# run_krhf_periodic_gdf._build_j_from_lpq. r_P = S_k w_k
# tr(L(k,k).D(k)) (PySCF get_j_kpts convention); ``L_jj.conj()``
# over-binds E_J on tight cells with complex diagonal cderi.
rho = rho + float(weights[j]) * np.einsum(
"Pls,sl->P", L_jj, np.asarray(D_k_in[j]), optimize=True
)
J_k = []
for i in range(n_k):
Ji = np.einsum("P,Pmn->mn", rho, lpq_cache[(i, i)], optimize=True)
J_k.append(0.5 * (Ji + Ji.conj().T))
return J_k
def _build_k_from_lpq(D_k_in: List[np.ndarray]) -> List[np.ndarray]:
K_k = [
np.zeros((basis.nbasis, basis.nbasis), dtype=complex) for _ in range(n_k)
]
for i in range(n_k):
for j in range(n_k):
Lpq = lpq_cache[(i, j)]
Dj = np.asarray(D_k_in[j])
tmp = np.einsum("Lpr,rs->Lps", Lpq, Dj)
K_k[i] += (1.0 / n_k) * np.einsum("Lps,Lqs->pq", tmp, Lpq.conj())
for i in range(n_k):
K_k[i] = 0.5 * (K_k[i] + K_k[i].conj().T)
return K_k
nbf = basis.nbasis
def _spin_density_k(C_k_local: List[np.ndarray], n_occ_each: int) -> List[np.ndarray]:
out = []
for i in range(n_k):
if n_occ_each > 0:
Cocc = C_k_local[i][:, :n_occ_each]
out.append(Cocc @ Cocc.conj().T)
else:
out.append(np.zeros((nbf, nbf), dtype=complex))
return out
def _occupy_and_density(
C_alpha: List[np.ndarray], eps_alpha: List[np.ndarray],
C_beta: List[np.ndarray], eps_beta: List[np.ndarray],
):
"""Per-spin occupations + densities + entropy.
``T = 0`` -> exact per-k hard Aufbau (the pre-smearing path,
bit-identical). ``T > 0`` -> independent per-spin Fermi-Dirac with
a global mu_s across the BZ (:func:`apply_smearing_open_shell`).
Returns ``(D_a, D_b, occ_a, occ_b, mu_a, mu_b, S_total)`` where
``S_total = S_a + S_b`` is the dimensionless entropy per cell.
"""
if not smear_opts.enabled:
return (
_spin_density_k(C_alpha, n_alpha),
_spin_density_k(C_beta, n_beta),
[], [], 0.0, 0.0, 0.0,
)
a_res, b_res = _apply_smearing_open_shell(
eps_alpha, eps_beta, weights=list(weights),
n_alpha=n_alpha, n_beta=n_beta, smearing=smear_opts,
)
return (
_density_from_orbitals(C_alpha, a_res.occupations_per_k),
_density_from_orbitals(C_beta, b_res.occupations_per_k),
a_res.occupations_per_k, b_res.occupations_per_k,
float(a_res.mu), float(b_res.mu),
float(a_res.entropy + b_res.entropy),
)
# ---- Initial guess: per-k Hcore diag, per-spin densities ----------
C_alpha_k: List[np.ndarray] = []
eps_alpha_k: List[np.ndarray] = []
for i in range(n_k):
Ci, ei = _diag_in_orth_basis(Hcore_k[i], X_k[i])
C_alpha_k.append(Ci.astype(complex))
eps_alpha_k.append(ei)
C_beta_k = [C.copy() for C in C_alpha_k]
eps_beta_k = [e.copy() for e in eps_alpha_k]
(D_alpha_k, D_beta_k, occ_alpha_k, occ_beta_k,
fermi_alpha, fermi_beta, entropy_total) = _occupy_and_density(
C_alpha_k, eps_alpha_k, C_beta_k, eps_beta_k
)
if smearing_T > 0.0:
plog.info(
"smearing: per-spin Fermi-Dirac kBT = "
f"{smearing_T:.6g} Ha "
f"({_hartree_to_kelvin_temperature(smearing_T):.1f} K)"
)
D_alpha_prev = [D.copy() for D in D_alpha_k]
D_beta_prev = [D.copy() for D in D_beta_k]
# ---- SCF setup ----------------------------------------------------
damping = float(opts.damping)
if not (0.0 <= damping < 1.0):
raise ValueError(f"run_kuhf_periodic_gdf: damping must be in [0, 1); got {damping}")
damper: Optional[DynamicDamping] = None
if bool(getattr(opts, "dynamic_damping", False)):
damper = DynamicDamping(
initial_alpha=damping,
alpha_min=float(getattr(opts, "dynamic_damping_min", 0.0)),
alpha_max=float(getattr(opts, "dynamic_damping_max", 0.95)),
)
use_diis = bool(opts.use_diis)
diis_start_iter = int(opts.diis_start_iter)
accel: Optional[MultiKPeriodicUHFAccelerator] = (
MultiKPeriodicUHFAccelerator(opts) if use_diis else None
)
max_iter = int(opts.max_iter)
scf_trace: List[SCFIteration] = []
result = PeriodicKUHFGDFResult(
energy=0.0, e_electronic=0.0, e_nuclear=float(e_nuc), n_iter=0,
converged=False, s_squared=0.0,
s_squared_ideal=0.25 * (mult - 1) * (mult + 1),
mo_energies_alpha=[e.copy() for e in eps_alpha_k],
mo_coeffs_alpha=[C.copy() for C in C_alpha_k],
density_alpha=[D.copy() for D in D_alpha_k],
fock_alpha=[np.empty((0, 0), dtype=complex) for _ in range(n_k)],
mo_energies_beta=[e.copy() for e in eps_beta_k],
mo_coeffs_beta=[C.copy() for C in C_beta_k],
density_beta=[D.copy() for D in D_beta_k],
fock_beta=[np.empty((0, 0), dtype=complex) for _ in range(n_k)],
overlap=[S.copy() for S in S_k], hcore=[H.copy() for H in Hcore_k],
kpoints_cart=kpoints_cart.copy(), kpoint_weights=weights.copy(),
scf_trace=scf_trace, functional=func_name or None,
aux_basis_name=aux_name, n_aux=int(aux.nbasis),
backend=(
f"native-multi-k-gdf-{k_exchange}-"
f"{'uks' if is_ks else 'uhf'}"
),
smearing_temperature=smearing_T,
fermi_level_alpha=float(fermi_alpha),
fermi_level_beta=float(fermi_beta),
entropy=float(entropy_total),
occupations_alpha=[np.asarray(o, dtype=float) for o in occ_alpha_k],
occupations_beta=[np.asarray(o, dtype=float) for o in occ_beta_k],
)
plog.banner(f"SCF ({label} multi-k, native GDF)")
E_prev = 0.0
for it in range(1, max_iter + 1):
if damper is not None:
damping = damper.alpha
diis_active = use_diis and it >= diis_start_iter
if it == 1 or damping == 0.0 or diis_active:
Da_used = [D.copy() for D in D_alpha_k]
Db_used = [D.copy() for D in D_beta_k]
else:
Da_used = [damping * Dp + (1.0 - damping) * Dn
for Dp, Dn in zip(D_alpha_prev, D_alpha_k)]
Db_used = [damping * Dp + (1.0 - damping) * Dn
for Dp, Dn in zip(D_beta_prev, D_beta_k)]
D_total = [Da_used[i] + Db_used[i] for i in range(n_k)]
J_k = _build_j_from_lpq(D_total)
if alpha != 0.0:
if cosx_bridge is None:
Ka_k = _build_k_from_lpq(Da_used)
Kb_k = _build_k_from_lpq(Db_used)
else:
Ka_k = cosx_bridge.k_matrices(
Da_used, list(kpoints_cart), lr_complement=True
)
Kb_k = cosx_bridge.k_matrices(
Db_used, list(kpoints_cart), lr_complement=True
)
madelung = _madelung_for_kmesh(system, kmesh_bloch.mesh)
Ka_k = list(apply_exxdiv_ewald_to_K(Ka_k, S_k, Da_used, madelung))
Kb_k = list(apply_exxdiv_ewald_to_K(Kb_k, S_k, Db_used, madelung))
else:
Ka_k = [np.zeros_like(J_k[0]) for _ in range(n_k)]
Kb_k = [np.zeros_like(J_k[0]) for _ in range(n_k)]
Fa_2e = [J_k[i] - alpha * Ka_k[i] for i in range(n_k)]
Fb_2e = [J_k[i] - alpha * Kb_k[i] for i in range(n_k)]
E_xc = 0.0
Va_xc_k = Vb_xc_k = None
if is_ks:
# Spin-polarised XC on the full real-space finite-torus density
# (per-spin inverse Bloch fold), Bloch-folded back to every k --
# the same convention as the closed-shell multi-k branch
# (_build_xc_k_from_density). The historical BZ-averaged
# home-cell shortcut (one k-independent V_xc(Γ) added to every
# k) was exact only in the vacuum/molecular limit -- the
# 2026-07-09 KRKS finding's defect class (commit b3f74aa9).
E_xc, Va_xc_k, Vb_xc_k = _build_xc_k_from_density_uks(
basis=basis,
system=system,
grid=grid,
func=func,
density_alpha_k=Da_used,
density_beta_k=Db_used,
kmesh_bloch=kmesh_bloch,
cells=cells,
kpoints_cart=kpoints_cart,
lat_opts=lat_opts,
)
Fa_k: List[np.ndarray] = []
Fb_k: List[np.ndarray] = []
for i in range(n_k):
Fa = Fa_2e[i] + Hcore_k[i]
Fb = Fb_2e[i] + Hcore_k[i]
if Va_xc_k is not None:
Fa = Fa + Va_xc_k[i]
Fb = Fb + Vb_xc_k[i]
Fa_k.append(0.5 * (Fa + Fa.conj().T))
Fb_k.append(0.5 * (Fb + Fb.conj().T))
E_coulomb = 0.5 * sum(
float(weights[i]) * float(np.real(np.trace(D_total[i] @ J_k[i])))
for i in range(n_k)
)
E_hf_K = (
-0.5 * alpha * sum(
float(weights[i]) * (
float(np.real(np.trace(Da_used[i] @ Ka_k[i])))
+ float(np.real(np.trace(Db_used[i] @ Kb_k[i])))
)
for i in range(n_k)
)
if alpha != 0.0 else 0.0
)
E_elec = 0.0
for i in range(n_k):
w = float(weights[i])
E_elec += w * (
float(np.real(np.trace(Da_used[i] @ Hcore_k[i])))
+ float(np.real(np.trace(Db_used[i] @ Hcore_k[i])))
+ 0.5 * float(np.real(np.trace(Da_used[i] @ Fa_2e[i])))
+ 0.5 * float(np.real(np.trace(Db_used[i] @ Fb_2e[i])))
)
E_total = E_elec + E_xc + float(e_nuc)
# Mermin free energy A = E - T.(S_a + S_b); SCF converges on A and
# the trace reports it. At T = 0 entropy_total = 0 => A = E_total,
# so the no-smearing run is bit-identical to the pre-smearing driver.
free_energy = E_total - smearing_T * entropy_total
grad_a: List[np.ndarray] = []
grad_b: List[np.ndarray] = []
gnorm2 = 0.0
for i in range(n_k):
FDSa = Fa_k[i] @ Da_used[i] @ S_k[i]
FDSb = Fb_k[i] @ Db_used[i] @ S_k[i]
ea = FDSa - FDSa.conj().T
eb = FDSb - FDSb.conj().T
grad_a.append(ea)
grad_b.append(eb)
gnorm2 += float(weights[i]) * (
float(np.linalg.norm(ea) ** 2) + float(np.linalg.norm(eb) ** 2)
)
grad_norm = float(np.sqrt(gnorm2))
dE = free_energy - E_prev
converged = (
it > 1
and abs(dE) < float(opts.conv_tol_energy)
and grad_norm < float(opts.conv_tol_grad)
)
scf_trace.append(SCFIteration(
iter=it, energy=float(free_energy),
delta_e=float(dE if it > 1 else 0.0), grad_norm=float(grad_norm),
diis_subspace=(accel.subspace_size if accel is not None else 0),
))
plog.iteration(
it, energy=float(free_energy), dE=float(dE if it > 1 else 0.0),
grad=float(grad_norm),
diis=(accel.subspace_size if accel is not None else 0),
)
result.energy = E_total
result.e_electronic = E_elec
result.e_xc = E_xc
result.e_coulomb = E_coulomb
result.e_hf_exchange = E_hf_K
result.n_iter = it
result.mo_energies_alpha = [e.copy() for e in eps_alpha_k]
result.mo_coeffs_alpha = [C.copy() for C in C_alpha_k]
result.density_alpha = [D.copy() for D in Da_used]
result.fock_alpha = [F.copy() for F in Fa_k]
result.mo_energies_beta = [e.copy() for e in eps_beta_k]
result.mo_coeffs_beta = [C.copy() for C in C_beta_k]
result.density_beta = [D.copy() for D in Db_used]
result.fock_beta = [F.copy() for F in Fb_k]
result.free_energy = float(free_energy)
result.entropy = float(entropy_total)
result.fermi_level_alpha = float(fermi_alpha)
result.fermi_level_beta = float(fermi_beta)
result.occupations_alpha = [np.asarray(o, dtype=float) for o in occ_alpha_k]
result.occupations_beta = [np.asarray(o, dtype=float) for o in occ_beta_k]
if converged:
result.converged = True
result.s_squared = _multi_k_s_squared(
n_alpha, n_beta, C_alpha_k, C_beta_k, S_k, weights,
occ_alpha_k=occ_alpha_k, occ_beta_k=occ_beta_k,
)
plog.converged(n_iter=it, energy=E_total, converged=True)
return result
if accel is not None:
Fa_ex, Fb_ex = accel.extrapolate_uhf(
Fa_k, Fb_k,
error_alpha_k_list=grad_a, error_beta_k_list=grad_b,
density_alpha_k_list=Da_used, density_beta_k_list=Db_used,
energy=E_total,
mo_coeffs_alpha_k_list=C_alpha_k, mo_coeffs_beta_k_list=C_beta_k,
n_alpha=n_alpha, n_beta=n_beta, weights=list(weights),
cells=cells, kpoints=list(kpoints_cart),
)
if diis_active:
Fa_k, Fb_k = Fa_ex, Fb_ex
Ca_new: List[np.ndarray] = []
ea_new: List[np.ndarray] = []
Cb_new: List[np.ndarray] = []
eb_new: List[np.ndarray] = []
for i in range(n_k):
Ca, eai = _diag_in_orth_basis(Fa_k[i], X_k[i])
Cb, ebi = _diag_in_orth_basis(Fb_k[i], X_k[i])
Ca_new.append(Ca.astype(complex))
ea_new.append(eai)
Cb_new.append(Cb.astype(complex))
eb_new.append(ebi)
D_alpha_prev = Da_used
D_beta_prev = Db_used
C_alpha_k, eps_alpha_k = Ca_new, ea_new
C_beta_k, eps_beta_k = Cb_new, eb_new
(D_alpha_k, D_beta_k, occ_alpha_k, occ_beta_k,
fermi_alpha, fermi_beta, entropy_total) = _occupy_and_density(
C_alpha_k, eps_alpha_k, C_beta_k, eps_beta_k
)
if damper is not None:
damper.update(free_energy)
E_prev = free_energy
result.converged = False
result.s_squared = _multi_k_s_squared(
n_alpha, n_beta, C_alpha_k, C_beta_k, S_k, weights,
occ_alpha_k=occ_alpha_k, occ_beta_k=occ_beta_k,
)
plog.converged(n_iter=result.n_iter, energy=result.energy, converged=False)
return result
def run_kuks_periodic_gdf(
system: PeriodicSystem,
basis: BasisSet,
kmesh: Union[Sequence[int], KPoints, BlochKMesh] = (1, 1, 1),
options: Optional[PeriodicKSOptions] = None,
*,
functional: str,
**kwargs,
) -> PeriodicKUHFGDFResult:
"""Open-shell periodic UKS multi-k SCF via native GDF -- thin wrapper
over :func:`run_kuhf_periodic_gdf` with ``functional`` required."""
if not functional:
raise ValueError("run_kuks_periodic_gdf: a functional name is required.")
return run_kuhf_periodic_gdf(
system, basis, kmesh, options, functional=functional, **kwargs
)
[docs]
def run_krks_periodic_gdf(
system: PeriodicSystem,
basis: BasisSet,
kmesh: Union[Sequence[int], KPoints, BlochKMesh] = (1, 1, 1),
options: Optional[PeriodicKSOptions] = None,
*,
functional: Optional[str] = None,
aux_basis: Optional[str] = None,
aux_drop_eta: float = 0.0,
gdf_linear_dep_threshold: float = 1e-9,
apply_modrho: bool = True,
fock_mixing: Optional[float] = None,
level_shift_warmup_cycles: Optional[int] = None,
linear_dep_threshold: float = 1e-7,
use_compcell: bool = False,
compcell_eta: float = 1.0,
apply_aft_correction: bool = True,
aft_ft_convention: str = "libint",
aft_precision: float = 1e-10,
rcut_strategy: Optional[object] = "pyscf_auto",
rcut_precision: float = 1e-8,
k_exchange: str = "gdf",
gdf_method: str = "rsgdf",
rsgdf_ke_cutoff: float = 200.0,
rsgdf_tail_ke_cutoff: Optional[float] = None,
fit_screen_threshold: float = 0.0,
mdf_ke_cutoff: float = 40.0,
bz_integration: Optional[str] = None,
density_mixer: Optional[str] = None,
density_mixer_depth: int = 8,
density_mixer_beta: float = 0.5,
density_mixer_kerker: bool = False,
kerker_k0: float = 1.5,
kerker_strength: float = 1.0,
kerker_cutoff_ha: float = 120.0,
dft_plus_u_sites: Optional[Sequence[object]] = None,
initial_density_k: Optional[Sequence[np.ndarray]] = None,
check_energy_sanity: bool = True,
progress: Union[bool, ProgressLogger, None] = None,
verbose: Optional[int] = None,
) -> PeriodicKRKSGDFResult:
"""Run closed-shell periodic KS-DFT multi-k SCF via native GDF.
Thin wrapper around :func:`run_krhf_periodic_gdf` that asserts a
functional has been provided. Functional dispatch and exact-
exchange mixing happen inside the shared SCF loop.
"""
opts = _options_or_default(options, is_ks=True)
func = functional or getattr(opts, "functional", None)
if not func:
raise ValueError("run_krks_periodic_gdf requires functional=...")
return run_krhf_periodic_gdf(
system,
basis,
kmesh,
opts,
functional=str(func),
aux_basis=aux_basis,
aux_drop_eta=aux_drop_eta,
gdf_linear_dep_threshold=gdf_linear_dep_threshold,
apply_modrho=apply_modrho,
fock_mixing=fock_mixing,
level_shift_warmup_cycles=level_shift_warmup_cycles,
linear_dep_threshold=linear_dep_threshold,
use_compcell=use_compcell,
compcell_eta=compcell_eta,
apply_aft_correction=apply_aft_correction,
aft_ft_convention=aft_ft_convention,
aft_precision=aft_precision,
rcut_strategy=rcut_strategy,
rcut_precision=rcut_precision,
k_exchange=k_exchange,
gdf_method=gdf_method,
rsgdf_ke_cutoff=rsgdf_ke_cutoff,
rsgdf_tail_ke_cutoff=rsgdf_tail_ke_cutoff,
fit_screen_threshold=fit_screen_threshold,
mdf_ke_cutoff=mdf_ke_cutoff,
bz_integration=bz_integration,
density_mixer=density_mixer,
density_mixer_depth=density_mixer_depth,
density_mixer_beta=density_mixer_beta,
density_mixer_kerker=density_mixer_kerker,
kerker_k0=kerker_k0,
kerker_strength=kerker_strength,
kerker_cutoff_ha=kerker_cutoff_ha,
dft_plus_u_sites=dft_plus_u_sites,
initial_density_k=initial_density_k,
check_energy_sanity=check_energy_sanity,
progress=progress,
verbose=verbose,
)