"""Phase 15b: multi-k periodic UHF SCF driver using the composed
EWALD_3D Coulomb dispatch.
Open-shell counterpart of :mod:`vibeqc.periodic_rhf_multi_k_ewald`.
Carries two density matrices ``D_α(g), D_β(g)`` per real-space cell
and constructs the per-spin 2e Fock at every k via the composed
Ewald split
F^{2e,σ}(g) = J(D_total, ω)(g) − K(D_σ)(g)
with ``D_total = D_α + D_β`` (one-particle convention, no factor 2)
and full-range exchange. The Ewald J is split into a short-range
real-space ``J_SR`` and a long-range FFT-Poisson ``J_LR`` exactly
as in the closed-shell multi-k Fock builder; per-spin K is extracted
from the closed-shell builder via
K(D_σ) = 2 · (J_full(D_σ) − F_full(D_σ))
where ``J_full(D)`` and ``F_full(D) = J_full(D) − ½ K(D)`` are two
separate ``build_fock_2e_real_space`` calls at ω = 0. Total cost per
iteration: 6 lattice-ERI builds (1 J_SR(D_total) + 1 J_LR(D_total)
+ 2 J_full(D_σ) + 2 F_full(D_σ)) + one FFT Poisson, vs 3 + 1 for
closed-shell multi-k RHF.
Bloch-sums each per-spin F^{2e,σ}(g) to F^{2e,σ}(k), adds Hcore(k),
diagonalises per spin per k. DIIS is per-spin (separate Pulay history
α vs β). Inverse-Bloch fold rebuilds D_α_real and D_β_real from the
per-k MOs each iteration.
Scope
-----
- Multi-k periodic UHF with EWALD_3D. The Γ-only path (single
[1,1,1] mesh) reproduces :func:`run_uhf_periodic_gamma_ewald3d`
by construction (the multi-k builder reduces to the molecular-
limit J convention at [1,1,1] mesh — same caveat as the RHF
case).
- Closed-shell case (multiplicity = 1) reproduces the multi-k
Ewald RHF energy to ~µHa.
- Standard HF (no DFT, no smearing). Smearing wired into multi-k
UHF lands with Phase 15c periodic UKS where it's the more
common use case.
"""
from __future__ import annotations
import math
from dataclasses import dataclass, field
from typing import List, Optional, Sequence, Tuple, Union
import numpy as np
from ._vibeqc_core import (
BasisSet,
BlochKMesh,
CoulombMethod,
LatticeMatrixSet,
LatticeSumOptions,
PeriodicRHFOptions,
PeriodicSystem,
SCFIteration,
bloch_sum,
build_fock_2e_real_space,
compute_kinetic_lattice,
compute_nuclear_lattice,
compute_overlap_lattice,
nuclear_repulsion_per_cell,
real_space_density_from_kpoints,
real_space_density_from_kpoints_fractional,
)
from .ewald_j import auto_grid
from .madelung import (
madelung_energy_correction_for_lat as _madelung_energy_correction_for_lat,
)
from .periodic_density import build_j_long_range_periodic
from .periodic_rhf_multi_k_ewald import (
_canonical_orthogonalizer_complex,
_damp_lattice_matrix,
_diag_in_orth_basis,
_g0_block,
)
from .periodic_scf_accelerators import (
DynamicDamping,
MultiKPeriodicUHFAccelerator,
)
from .periodic_uhf_ewald import _spin_squared
from .progress import ProgressLogger, resolve_progress
from .scf_divergence import check_scf_divergence
from .smearing._support import reject_unsupported_smearing_temperature
__all__ = [
"PeriodicUHFMultiKEwaldResult",
"run_uhf_periodic_multi_k_ewald3d",
]
[docs]
@dataclass
class PeriodicUHFMultiKEwaldResult:
"""Result of :func:`run_uhf_periodic_multi_k_ewald3d`.
Per-cell quantities (``energy``, ``e_electronic``, ``e_nuclear``)
plus per-k matrices for each spin (``mo_energies_alpha``,
``mo_coeffs_alpha``, ``fock_alpha``, ``mo_energies_beta``, ...)
and the converged real-space densities ``density_alpha``,
``density_beta`` (each a :class:`LatticeMatrixSet`). Reports
``s_squared`` evaluated at the home-cell (Γ-block) MOs.
"""
energy: float
e_electronic: float
e_nuclear: float
n_iter: int
converged: bool
s_squared: float
s_squared_ideal: float
# α spin
mo_energies_alpha: List[np.ndarray]
mo_coeffs_alpha: List[np.ndarray]
fock_alpha: List[np.ndarray]
density_alpha: LatticeMatrixSet
# β spin
mo_energies_beta: List[np.ndarray]
mo_coeffs_beta: List[np.ndarray]
fock_beta: List[np.ndarray]
density_beta: LatticeMatrixSet
# Per-k overlap + Hcore (shared between α/β).
overlap: List[np.ndarray]
hcore: List[np.ndarray]
scf_trace: List[SCFIteration] = field(default_factory=list)
omega: float = 0.0
grid_shape: Tuple[int, int, int] = (0, 0, 0)
def _build_uhf_fock_blocks_ewald3d(
basis: BasisSet,
system: PeriodicSystem,
D_alpha_real: LatticeMatrixSet,
D_beta_real: LatticeMatrixSet,
omega: float,
lat_opts: LatticeSumOptions,
grid_shape_t: Tuple[int, int, int],
origin: Optional[Sequence[float]],
spacing_bohr: float,
) -> Tuple[List[np.ndarray], List[np.ndarray]]:
"""Compute per-cell ``(F^{2e,α}(g), F^{2e,β}(g))`` blocks.
Six lattice-ERI calls per invocation:
- J_SR(D_total, ω): real-space short-range Hartree.
- J_LR(D_total, ω): FFT Poisson long-range Hartree.
- J_full(D_α): full-range Hartree of D_α (for K extraction).
- F_full(D_α): J_full(D_α) − ½ K_full(D_α).
- J_full(D_β): same for β.
- F_full(D_β): same for β.
Then ``K(D_σ) = 2 · (J_full(D_σ) − F_full(D_σ))`` and
F^{2e,σ}(g) = J_SR_total(g) + J_LR_total(g) − K(D_σ)(g).
"""
n_cells = len(D_alpha_real.cells)
# Total density D_total = D_α + D_β. Build via the overlap template
# and mutate the C++ storage in place via ``set_block``.
D_total_real = compute_overlap_lattice(basis, system, lat_opts)
for g in range(n_cells):
D_total_real.set_block(
g,
np.asarray(D_alpha_real.blocks[g], dtype=float)
+ np.asarray(D_beta_real.blocks[g], dtype=float),
)
# J(D_total, ω) — Ewald split.
J_SR_total_lms = build_fock_2e_real_space(
basis,
system,
lat_opts,
D_total_real,
0.0,
float(omega),
)
J_LR_total_blocks = build_j_long_range_periodic(
basis,
system,
D_total_real,
omega=float(omega),
grid_shape=grid_shape_t,
origin=origin,
spacing_bohr=spacing_bohr,
output_cells=list(range(n_cells)),
)
# Per-spin J_full(D_σ) and F_full(D_σ) at ω = 0.
J_full_alpha = build_fock_2e_real_space(
basis,
system,
lat_opts,
D_alpha_real,
0.0,
0.0,
)
F_full_alpha = build_fock_2e_real_space(
basis,
system,
lat_opts,
D_alpha_real,
1.0,
0.0,
)
J_full_beta = build_fock_2e_real_space(
basis,
system,
lat_opts,
D_beta_real,
0.0,
0.0,
)
F_full_beta = build_fock_2e_real_space(
basis,
system,
lat_opts,
D_beta_real,
1.0,
0.0,
)
F_alpha_blocks: List[np.ndarray] = []
F_beta_blocks: List[np.ndarray] = []
for g in range(n_cells):
J_SR_total = np.asarray(J_SR_total_lms.blocks[g], dtype=float)
J_LR_total = np.asarray(J_LR_total_blocks[g], dtype=float)
J_total = J_SR_total + J_LR_total
# K(D_σ) = 2 · (J_full(D_σ) − F_full(D_σ)) since
# F_full(D) = J_full(D) − ½ K(D) ⇒ K(D) = 2(J_full − F_full).
J_full_a = np.asarray(J_full_alpha.blocks[g], dtype=float)
F_full_a = np.asarray(F_full_alpha.blocks[g], dtype=float)
K_a = 2.0 * (J_full_a - F_full_a)
J_full_b = np.asarray(J_full_beta.blocks[g], dtype=float)
F_full_b = np.asarray(F_full_beta.blocks[g], dtype=float)
K_b = 2.0 * (J_full_b - F_full_b)
F_alpha_blocks.append(J_total - K_a)
F_beta_blocks.append(J_total - K_b)
return F_alpha_blocks, F_beta_blocks
def _bloch_sum_blocks(
blocks: Sequence[np.ndarray],
cells,
k_cart: np.ndarray,
) -> np.ndarray:
"""F(k) = Σ_g e^{i k·R_g} F(g)."""
k = np.asarray(k_cart, dtype=float).reshape(3)
F_k = np.zeros_like(blocks[0], dtype=complex)
for g_idx, block in enumerate(blocks):
R_g = np.asarray(cells[g_idx].r_cart, dtype=float)
phase = np.exp(1j * float(np.dot(k, R_g)))
F_k = F_k + phase * block
return F_k
[docs]
def run_uhf_periodic_multi_k_ewald3d(
system: PeriodicSystem,
basis: BasisSet,
kmesh: BlochKMesh,
options=None,
*,
omega: float = 0.0,
grid_shape: Optional[Union[Tuple[int, int, int], int]] = None,
origin: Optional[Sequence[float]] = None,
spacing_bohr: float = 0.3,
linear_dep_threshold: float = 1e-7,
canonical_orth_normalize_diag_first: bool = True,
auto_optimize_truncation: bool = True,
progress: Union[bool, ProgressLogger, None] = None,
verbose: Optional[int] = None,
) -> PeriodicUHFMultiKEwaldResult:
"""Multi-k open-shell UHF SCF with EWALD_3D Coulomb.
Mirror of :func:`run_rhf_periodic_multi_k_ewald3d` for open-shell
systems. Same DIIS / damping / orthogonaliser conventions, but
everything runs per-spin: two Pulay-DIIS instances, two
diagonalisations per k, two density matrices folded back via
inverse Bloch.
Occupations follow the molecule's multiplicity (assumed equal at
every k — closed-shell-insulator-flavoured setup, matching the
multi-k RHF driver).
"""
from ._vibeqc_core import PeriodicRHFOptions as _Opts
opts = options if options is not None else _Opts()
reject_unsupported_smearing_temperature(
opts,
"run_uhf_periodic_multi_k_ewald3d",
detail=(
"UHF needs spin-resolved fractional occupations and chemical "
"potentials; use closed-shell RHF/RKS smearing or leave "
"smearing_temperature at 0.0."
),
)
lat_opts: LatticeSumOptions = opts.lattice_opts
plog = resolve_progress(progress, verbose=verbose)
# ω must match the nuclear Ewald α (auto-selected from
# nuclear_cutoff_bohr in the C++ ewald engine) so that the
# jellium background terms cancel exactly — see the matching
# block in run_rhf_periodic_gamma_ewald3d. User can override via
# opts.ewald_omega (rarely needed). The driver kwarg ``omega`` is
# retained for signature parity but is overridden here.
_ewald_tol = getattr(opts, "ewald_tolerance", 1e-12)
_cutoff = getattr(opts, "ewald_cutoff_bohr", lat_opts.nuclear_cutoff_bohr)
if omega <= 0.0:
_user_omega = getattr(opts, "ewald_omega", None)
if _user_omega is not None and float(_user_omega) > 0.0:
omega = float(_user_omega)
else:
from .bipole_ext_el_pole import crystal_default_ewald_alpha
V_cell = float(abs(np.linalg.det(np.asarray(system.lattice, dtype=float))))
omega = crystal_default_ewald_alpha(V_cell)
lat = np.asarray(system.lattice, dtype=float)
if grid_shape is None:
grid_shape_t = auto_grid(lat, spacing_bohr)
elif isinstance(grid_shape, int):
grid_shape_t = (grid_shape, grid_shape, grid_shape)
else:
grid_shape_t = tuple(int(x) for x in grid_shape)
plog.info(
f"UHF multi-k EWALD_3D / omega = {float(omega):.3f}, "
f"FFT grid {grid_shape_t[0]}x{grid_shape_t[1]}x{grid_shape_t[2]}"
)
plog.info(f"basis: {basis.name} ({basis.nbasis} BFs / {basis.nshells} shells)")
from .options_dump import dump_active_settings
dump_active_settings(
plog,
[
("PeriodicRHFOptions", opts),
("LatticeSumOptions", lat_opts),
(
"Driver kwargs",
{
"omega": float(omega),
"grid_shape": grid_shape_t,
"origin": origin,
"spacing_bohr": float(spacing_bohr),
"linear_dep_threshold": float(linear_dep_threshold),
"canonical_orth_normalize_diag_first": canonical_orth_normalize_diag_first,
"auto_optimize_truncation": auto_optimize_truncation,
},
),
],
)
if plog.level >= 5:
from .scf_log import format_basis_summary
plog.write_raw(format_basis_summary(basis))
# Open-shell occupations.
n_elec = int(system.n_electrons())
mult = int(system.multiplicity)
if mult < 1:
raise ValueError(
f"run_uhf_periodic_multi_k_ewald3d: multiplicity must be ≥ 1, got {mult}"
)
if (n_elec + mult - 1) % 2 != 0 or (n_elec - mult + 1) % 2 != 0:
raise ValueError(
f"run_uhf_periodic_multi_k_ewald3d: (n_electrons={n_elec}, "
f"multiplicity={mult}) cannot be split into integer α/β."
)
n_alpha = (n_elec + mult - 1) // 2
n_beta = (n_elec - mult + 1) // 2
k_points = list(kmesh.kpoints)
weights = np.asarray(kmesh.weights, dtype=float)
n_k = len(k_points)
if n_k == 0:
raise ValueError("kmesh has no k-points")
if not np.isclose(weights.sum(), 1.0):
raise ValueError(f"kmesh.weights must sum to 1; got {weights.sum():.6f}")
plog.info(
f"k-mesh: {n_k} k-point{'s' if n_k != 1 else ''}, "
f"weights sum = {weights.sum():.4f}; "
f"n_alpha = {n_alpha}, n_beta = {n_beta}"
)
# ---- Auto-optimise lattice truncation (default ON) -------------------
if auto_optimize_truncation and lat_opts.coulomb_method == CoulombMethod.EWALD_3D:
from .eigs_preflight import (
format_truncation_optimization_report,
optimize_truncation,
)
k_arr = [np.asarray(k, dtype=float) for k in k_points]
opt_rep = optimize_truncation(
system,
basis,
lattice_opts=lat_opts,
k_points_cart=k_arr,
)
if (
opt_rep.n_evaluations > 1
or opt_rep.optimized_lattice_opts.cutoff_bohr != lat_opts.cutoff_bohr
):
plog.write_raw(format_truncation_optimization_report(opt_rep))
if not opt_rep.converged:
plog.warning("auto_optimize_truncation did not converge.")
lat_opts = opt_rep.optimized_lattice_opts
# ---- Real-space one-electron integrals -------------------------------
with plog.stage(
"integrals_lattice", detail=f"S/T/V at cutoff {lat_opts.cutoff_bohr:.2f} bohr"
):
S_lat = compute_overlap_lattice(basis, system, lat_opts)
T_lat = compute_kinetic_lattice(basis, system, lat_opts)
from .periodic_v_ne import compute_nuclear_lattice_dispatch
V_lat = compute_nuclear_lattice_dispatch(basis, system, lat_opts)
cells = list(S_lat.cells)
n_cells = len(cells)
# Per-k S(k), Hcore(k), orthogonaliser X(k).
S_k_list: List[np.ndarray] = []
Hcore_k_list: List[np.ndarray] = []
X_k_list: List[np.ndarray] = []
# Per-k linear-dependence preflight; see periodic_rhf_multi_k_ewald
# for the rationale (Searle et al., ARCHER eCSE04-16, 2017).
from .linear_dependence import scf_preflight_overlap_check
for k_idx, k in enumerate(k_points):
k_arr = np.asarray(k, dtype=float).reshape(3)
S_k = np.asarray(bloch_sum(S_lat, k_arr))
T_k = np.asarray(bloch_sum(T_lat, k_arr))
V_k = np.asarray(bloch_sum(V_lat, k_arr))
H_k = T_k + V_k
S_k = 0.5 * (S_k + S_k.conj().T)
H_k = 0.5 * (H_k + H_k.conj().T)
scf_preflight_overlap_check(
S_k,
plog=plog,
label=f"S(k={k_idx}, k_cart={k_arr.round(4).tolist()})",
basis=basis,
)
X_k, n_kept = _canonical_orthogonalizer_complex(
S_k,
linear_dep_threshold,
normalize_diag_first=canonical_orth_normalize_diag_first,
)
if max(n_alpha, n_beta) > n_kept:
raise RuntimeError(
f"run_uhf_periodic_multi_k_ewald3d: orth dropped too many "
f"directions (n_α={n_alpha}, n_β={n_beta}, "
f"n_kept={n_kept}) at k = {k_arr}"
)
S_k_list.append(S_k)
Hcore_k_list.append(H_k)
X_k_list.append(X_k)
e_nuc = float(nuclear_repulsion_per_cell(system, lat_opts))
# ---- Initial guess: diagonalize Hcore(k) for both spins.
C_alpha_per_k: List[np.ndarray] = []
eps_alpha_per_k: List[np.ndarray] = []
C_beta_per_k: List[np.ndarray] = []
eps_beta_per_k: List[np.ndarray] = []
for H_k, X_k in zip(Hcore_k_list, X_k_list):
C_a, eps_a = _diag_in_orth_basis(H_k, X_k)
C_b, eps_b = _diag_in_orth_basis(H_k, X_k)
C_alpha_per_k.append(C_a.astype(complex))
eps_alpha_per_k.append(eps_a)
C_beta_per_k.append(C_b.astype(complex))
eps_beta_per_k.append(eps_b)
# Build initial real-space densities. UHF densities are per-spin
# one-particle (no factor 2). Use the fractional-occupation C++
# density builder with occupations of 1.0 (rather than 2.0) to
# build the spin density directly — sidesteps the
# LatticeMatrixSet.blocks-is-a-Python-list-copy mutation footgun
# the closed-shell-then-halve approach hit on first try.
def _spin_density(C_per_k_local, n_occ_each):
"""Build a one-particle (no factor 2) real-space spin
density via the C++ fractional-occupation builder."""
nbf = C_per_k_local[0].shape[1]
occ_per_k = []
for _ in range(n_k):
occ = np.zeros(nbf, dtype=float)
occ[:n_occ_each] = 1.0
occ_per_k.append(occ)
return real_space_density_from_kpoints_fractional(
C_per_k_local,
occ_per_k,
kmesh,
cells,
)
D_alpha_real = _spin_density(C_alpha_per_k, n_alpha)
D_beta_real = _spin_density(C_beta_per_k, n_beta)
D_alpha_prev: Optional[LatticeMatrixSet] = None
D_beta_prev: Optional[LatticeMatrixSet] = None
damping = float(opts.damping)
if not (0.0 <= damping < 1.0):
raise ValueError(
f"run_uhf_periodic_multi_k_ewald3d: damping must be in "
f"[0, 1); got {damping}"
)
use_diis = bool(opts.use_diis)
diis_start_iter = int(opts.diis_start_iter)
accel: Optional[MultiKPeriodicUHFAccelerator] = (
MultiKPeriodicUHFAccelerator(opts) if use_diis else None
)
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)),
)
level_shift = float(getattr(opts, "level_shift", 0.0))
# Phase C1c — quadratic SCF fallback (per-spin per-k Newton step).
quadratic_fallback_iter = int(getattr(opts, "quadratic_fallback_iter", 0))
quadratic_fallback_shift = float(getattr(opts, "quadratic_fallback_shift", 0.1))
quadratic_fallback_max_step = float(
getattr(opts, "quadratic_fallback_max_step", 0.1)
)
# ---- SCF loop --------------------------------------------------------
plog.banner("SCF (UHF multi-k, EWALD_3D)")
plog.info(" iter energy (Ha) dE ||[F,DS]|| DIIS")
scf_trace: List[SCFIteration] = []
E_prev = 0.0
F_alpha_k_list: List[np.ndarray] = [np.zeros_like(H) for H in Hcore_k_list]
F_beta_k_list: List[np.ndarray] = [np.zeros_like(H) for H in Hcore_k_list]
converged = False
iter_idx = 0
for iter_idx in range(1, int(opts.max_iter) + 1):
if damper is not None:
damping = damper.alpha
diis_active = use_diis and iter_idx >= diis_start_iter
# Density damping (per-spin) when DIIS not yet active.
if iter_idx > 1 and damping > 0.0 and not diis_active:
D_alpha_used = _damp_lattice_matrix(
D_alpha_real,
D_alpha_prev,
damping,
)
D_beta_used = _damp_lattice_matrix(
D_beta_real,
D_beta_prev,
damping,
)
else:
D_alpha_used = D_alpha_real
D_beta_used = D_beta_real
# 2e Fock blocks for both spins.
F_alpha_blocks, F_beta_blocks = _build_uhf_fock_blocks_ewald3d(
basis,
system,
D_alpha_used,
D_beta_used,
omega,
lat_opts,
grid_shape_t,
origin,
spacing_bohr,
)
# Bloch-sum to per-k F^{2e,σ}(k) and add Hcore(k).
F_alpha_k_list = []
F_beta_k_list = []
for k_idx, k_cart in enumerate(k_points):
F_a = _bloch_sum_blocks(F_alpha_blocks, cells, np.asarray(k_cart))
F_b = _bloch_sum_blocks(F_beta_blocks, cells, np.asarray(k_cart))
F_a = F_a + Hcore_k_list[k_idx]
F_b = F_b + Hcore_k_list[k_idx]
F_alpha_k_list.append(F_a)
F_beta_k_list.append(F_b)
# Per-cell electronic energy + per-k commutator errors.
# E_elec = ½ tr[D_total · Hcore] + ½ tr[D_α · F_α] + ½ tr[D_β · F_β]
E_elec = 0.0
grad_norm_sum = 0.0
error_alpha_k_list: List[np.ndarray] = []
error_beta_k_list: List[np.ndarray] = []
for idx in range(n_k):
C_a = C_alpha_per_k[idx]
C_b = C_beta_per_k[idx]
C_a_occ = C_a[:, :n_alpha] if n_alpha > 0 else C_a[:, :0]
C_b_occ = C_b[:, :n_beta] if n_beta > 0 else C_b[:, :0]
D_a_k = C_a_occ @ C_a_occ.conj().T
D_b_k = C_b_occ @ C_b_occ.conj().T
H_k = Hcore_k_list[idx]
F_a_k = F_alpha_k_list[idx]
F_b_k = F_beta_k_list[idx]
w = float(weights[idx])
E_elec += w * (
0.5 * np.real(np.trace((D_a_k + D_b_k) @ H_k))
+ 0.5 * np.real(np.trace(D_a_k @ F_a_k))
+ 0.5 * np.real(np.trace(D_b_k @ F_b_k))
)
S_k = S_k_list[idx]
FDS_a = F_a_k @ D_a_k @ S_k
FDS_b = F_b_k @ D_b_k @ S_k
err_a = FDS_a - FDS_a.conj().T
err_b = FDS_b - FDS_b.conj().T
error_alpha_k_list.append(err_a)
error_beta_k_list.append(err_b)
grad_norm_sum += w * float(
np.sqrt(np.linalg.norm(err_a) ** 2 + np.linalg.norm(err_b) ** 2)
)
# Madelung-leak correction (v0.6.1).
_D_g0 = np.asarray(_g0_block(D_alpha_real)) + np.asarray(_g0_block(D_beta_real))
_S_g0 = np.asarray(_g0_block(S_lat))
E_madelung_fix = _madelung_energy_correction_for_lat(
_D_g0, _S_g0, system, lat_opts
)
E_total = float(E_elec) + e_nuc + E_madelung_fix
dE = E_total - E_prev
# Divergence detection (v0.6.2).
check_scf_divergence(
"run_uhf_periodic_multi_k_ewald3d",
iter_idx,
E_total,
grad_norm_sum,
dE,
)
diis_sub = accel.subspace_size if accel is not None else 0
scf_trace.append(
SCFIteration(
iter=iter_idx,
energy=float(E_total),
delta_e=float(dE if iter_idx > 1 else 0.0),
grad_norm=float(grad_norm_sum),
diis_subspace=diis_sub,
)
)
plog.iteration(
iter_idx,
energy=float(E_total),
dE=float(dE if iter_idx > 1 else 0.0),
grad=float(grad_norm_sum),
diis=diis_sub,
)
converged = (
iter_idx > 1
and abs(dE) < float(opts.conv_tol_energy)
and grad_norm_sum < float(opts.conv_tol_grad)
)
# Phase C1c gate. When the quadratic fallback is active, take
# per-spin per-k Newton steps in MO space — bypass DIIS and
# level shift since the Newton step is its own update.
in_quadratic_phase = (
quadratic_fallback_iter > 0 and iter_idx > quadratic_fallback_iter
)
new_C_alpha: List[np.ndarray] = []
new_eps_alpha: List[np.ndarray] = []
new_C_beta: List[np.ndarray] = []
new_eps_beta: List[np.ndarray] = []
if in_quadratic_phase:
from .quadratic_scf import quadratic_step
for idx in range(n_k):
C_a, eps_a = quadratic_step(
F_alpha_k_list[idx],
C_alpha_per_k[idx],
eps_alpha_per_k[idx],
n_alpha,
shift=quadratic_fallback_shift,
max_step=quadratic_fallback_max_step,
)
C_b, eps_b = quadratic_step(
F_beta_k_list[idx],
C_beta_per_k[idx],
eps_beta_per_k[idx],
n_beta,
shift=quadratic_fallback_shift,
max_step=quadratic_fallback_max_step,
)
new_C_alpha.append(C_a)
new_eps_alpha.append(eps_a)
new_C_beta.append(C_b)
new_eps_beta.append(eps_b)
else:
# SCF-accelerator extrapolation. DIIS runs per-spin
# per-k Pulay internally; KDIIS / EDIIS / ADIIS use a
# single spin-coupled history (one coefficient set
# applied to both spins). For the bridged paths
# (EDIIS / ADIIS / EDIIS_DIIS) we Bloch-sum the per-cell
# D_α / D_β densities to per-k for input.
if accel is not None:
density_alpha_k_list = [
_bloch_sum_blocks(
D_alpha_used.blocks,
D_alpha_used.cells,
np.asarray(k),
)
for k in k_points
]
density_beta_k_list = [
_bloch_sum_blocks(
D_beta_used.blocks,
D_beta_used.cells,
np.asarray(k),
)
for k in k_points
]
F_a_ex, F_b_ex = accel.extrapolate_uhf(
F_alpha_k_list,
F_beta_k_list,
error_alpha_k_list=error_alpha_k_list,
error_beta_k_list=error_beta_k_list,
density_alpha_k_list=density_alpha_k_list,
density_beta_k_list=density_beta_k_list,
energy=E_total,
mo_coeffs_alpha_k_list=C_alpha_per_k,
mo_coeffs_beta_k_list=C_beta_per_k,
n_alpha=n_alpha,
n_beta=n_beta,
weights=list(weights),
cells=cells,
kpoints=list(k_points),
)
if diis_active:
F_alpha_k_list = F_a_ex
F_beta_k_list = F_b_ex
# Optional Saunders-Hillier level shift per spin per k.
if level_shift != 0.0:
for idx in range(n_k):
S_k = S_k_list[idx]
C_a = C_alpha_per_k[idx]
C_b = C_beta_per_k[idx]
C_a_occ = C_a[:, :n_alpha] if n_alpha > 0 else C_a[:, :0]
C_b_occ = C_b[:, :n_beta] if n_beta > 0 else C_b[:, :0]
# UHF densities are 1·C_occ·C_occ† (no factor 2).
D_a_k = C_a_occ @ C_a_occ.conj().T
D_b_k = C_b_occ @ C_b_occ.conj().T
F_alpha_k_list[idx] = (
F_alpha_k_list[idx]
+ level_shift * S_k
- level_shift * (S_k @ D_a_k @ S_k)
)
F_beta_k_list[idx] = (
F_beta_k_list[idx]
+ level_shift * S_k
- level_shift * (S_k @ D_b_k @ S_k)
)
F_alpha_k_list[idx] = 0.5 * (
F_alpha_k_list[idx] + F_alpha_k_list[idx].conj().T
)
F_beta_k_list[idx] = 0.5 * (
F_beta_k_list[idx] + F_beta_k_list[idx].conj().T
)
for idx in range(n_k):
C_a, eps_a = _diag_in_orth_basis(
F_alpha_k_list[idx],
X_k_list[idx],
)
C_b, eps_b = _diag_in_orth_basis(
F_beta_k_list[idx],
X_k_list[idx],
)
new_C_alpha.append(C_a)
new_eps_alpha.append(eps_a)
new_C_beta.append(C_b)
new_eps_beta.append(eps_b)
C_alpha_per_k = new_C_alpha
eps_alpha_per_k = new_eps_alpha
C_beta_per_k = new_C_beta
eps_beta_per_k = new_eps_beta
# Rebuild D_α, D_β real-space densities (one-particle conv).
D_alpha_new = _spin_density(C_alpha_per_k, n_alpha)
D_beta_new = _spin_density(C_beta_per_k, n_beta)
D_alpha_prev = D_alpha_used
D_beta_prev = D_beta_used
D_alpha_real = D_alpha_new
D_beta_real = D_beta_new
if damper is not None:
damper.update(E_total)
E_prev = E_total
if converged:
break
# ---- Final pass on converged D's --------------------------------------
if converged:
F_alpha_blocks, F_beta_blocks = _build_uhf_fock_blocks_ewald3d(
basis,
system,
D_alpha_real,
D_beta_real,
omega,
lat_opts,
grid_shape_t,
origin,
spacing_bohr,
)
F_alpha_k_list = []
F_beta_k_list = []
E_elec = 0.0
for k_idx, k_cart in enumerate(k_points):
F_a = _bloch_sum_blocks(F_alpha_blocks, cells, np.asarray(k_cart))
F_b = _bloch_sum_blocks(F_beta_blocks, cells, np.asarray(k_cart))
F_a = F_a + Hcore_k_list[k_idx]
F_b = F_b + Hcore_k_list[k_idx]
F_alpha_k_list.append(F_a)
F_beta_k_list.append(F_b)
final_C_alpha: List[np.ndarray] = []
final_C_beta: List[np.ndarray] = []
final_eps_alpha: List[np.ndarray] = []
final_eps_beta: List[np.ndarray] = []
for idx in range(n_k):
C_a, eps_a = _diag_in_orth_basis(
F_alpha_k_list[idx],
X_k_list[idx],
)
C_b, eps_b = _diag_in_orth_basis(
F_beta_k_list[idx],
X_k_list[idx],
)
final_C_alpha.append(C_a)
final_C_beta.append(C_b)
final_eps_alpha.append(eps_a)
final_eps_beta.append(eps_b)
C_a_occ = C_a[:, :n_alpha] if n_alpha > 0 else C_a[:, :0]
C_b_occ = C_b[:, :n_beta] if n_beta > 0 else C_b[:, :0]
D_a_k = C_a_occ @ C_a_occ.conj().T
D_b_k = C_b_occ @ C_b_occ.conj().T
w = float(weights[idx])
E_elec += w * (
0.5 * np.real(np.trace((D_a_k + D_b_k) @ Hcore_k_list[idx]))
+ 0.5 * np.real(np.trace(D_a_k @ F_alpha_k_list[idx]))
+ 0.5 * np.real(np.trace(D_b_k @ F_beta_k_list[idx]))
)
C_alpha_per_k = final_C_alpha
C_beta_per_k = final_C_beta
eps_alpha_per_k = final_eps_alpha
eps_beta_per_k = final_eps_beta
# Madelung-leak correction (v0.6.1).
_D_g0_f = np.asarray(_g0_block(D_alpha_real)) + np.asarray(
_g0_block(D_beta_real)
)
_S_g0_f = np.asarray(_g0_block(S_lat))
E_madelung_fix_f = _madelung_energy_correction_for_lat(
_D_g0_f, _S_g0_f, system, lat_opts
)
E_total = float(E_elec) + e_nuc + E_madelung_fix_f
# ⟨S²⟩ from the home-cell (Γ-block) MOs as a representative
# diagnostic. Multi-k ⟨S²⟩ is more involved; this is the standard
# quick check used by PySCF.pbc / CRYSTAL.
if n_alpha == 0 or n_beta == 0:
s2 = 0.25 * (n_alpha - n_beta) * (n_alpha - n_beta + 2) + n_beta
else:
# Use k=0 if present, else first k.
k0_idx = 0
for i, k in enumerate(k_points):
if np.allclose(np.asarray(k), 0.0):
k0_idx = i
break
S_real = np.real(S_k_list[k0_idx])
s2 = _spin_squared(
n_alpha,
n_beta,
np.real(C_alpha_per_k[k0_idx]),
np.real(C_beta_per_k[k0_idx]),
S_real,
)
plog.converged(n_iter=iter_idx, energy=E_total, converged=converged)
return PeriodicUHFMultiKEwaldResult(
energy=E_total,
e_electronic=float(E_elec),
e_nuclear=e_nuc,
n_iter=iter_idx,
converged=converged,
s_squared=float(s2),
s_squared_ideal=0.25 * (mult - 1) * (mult + 1),
mo_energies_alpha=eps_alpha_per_k,
mo_coeffs_alpha=C_alpha_per_k,
fock_alpha=F_alpha_k_list,
density_alpha=D_alpha_real,
mo_energies_beta=eps_beta_per_k,
mo_coeffs_beta=C_beta_per_k,
fock_beta=F_beta_k_list,
density_beta=D_beta_real,
overlap=S_k_list,
hcore=Hcore_k_list,
scf_trace=scf_trace,
omega=float(omega),
grid_shape=grid_shape_t,
)