"""Phase 15c-3a: Γ-point periodic UKS SCF driver using the composed
EWALD_3D Coulomb dispatch.
Open-shell DFT counterpart of :func:`run_rks_periodic_gamma_ewald3d`
(15c-1) and :func:`run_uhf_periodic_gamma_ewald3d` (15a). Per SCF
iteration:
F_α = H_core + J_ewald(ω, D_α + D_β) − α·K(D_α) + V_xc^α
F_β = H_core + J_ewald(ω, D_α + D_β) − α·K(D_β) + V_xc^β
where ``α = func.hf_exchange_fraction``. Compared to UHF Ewald (15a)
which uses ``-K_σ`` (full HF exchange per spin), here we scale by α
for hybrids and skip K entirely for pure DFT.
V_xc comes from the new :func:`build_xc_periodic_uks` (Phase 15c-3
C++ addition), which returns separate V_α / V_β / E_xc from
spin-resolved densities. The functional is constructed with
``spin=2`` to take the libxc polarized path.
Density flow. Γ-only molecular-limit cells have all electron density
in the home cell, so per-spin densities are wrapped in degenerate
LatticeMatrixSets (block 0 = D_σ, others zero) before being handed
to ``build_xc_periodic_uks``. Same trick as the Γ-only RKS Ewald
driver — exact in the molecular limit, the regime ``build_j_ewald_3d``
itself targets at Γ.
Energy formula. From the molecular UKS pattern (uks.cpp):
E_elec = E_xc + tr((D_α + D_β)·H_core)
+ ½ tr(D_α·F_HF_α) + ½ tr(D_β·F_HF_β)
where ``F_HF_σ = J(D_total) − α·K(D_σ)`` is the HF-exchange piece
of F_σ (V_xc reported through E_xc rather than a trace).
Scope.
* Γ-only single k-point.
* Open-shell, ``multiplicity ≥ 1``.
* Pure DFT, hybrid, and HF (``α = 1``, equivalent to UHF Ewald).
* DIIS on the block-diagonal [α, β] error vector with one rolling
history per spin (matches UHF Ewald and the molecular UKS C++
driver's convention).
* Saunders-Hillier level shift via ``options.level_shift``.
* Periodic Becke partition selectable via
``options.use_periodic_becke``.
The multi-k UKS Ewald driver lands in
:mod:`vibeqc.periodic_uks_multi_k_ewald` (15c-3b); it shares the
helpers with :mod:`vibeqc.periodic_uhf_multi_k_ewald` and adds the
periodic UKS XC.
"""
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,
CoulombMethod,
Functional,
InitialGuess,
LatticeMatrixSet,
LatticeSumOptions,
PeriodicKSOptions,
PeriodicSystem,
SCFIteration,
bloch_sum,
build_grid,
build_jk_gamma_molecular_limit,
build_xc_periodic_uks,
compute_kinetic_lattice,
compute_nuclear_lattice,
compute_overlap_lattice,
nuclear_repulsion_per_cell,
)
from .ewald_composed import build_j_ewald_3d
from .ewald_j import auto_grid
from .guess import initial_densities_open_shell
from .madelung import (
madelung_energy_correction_for_lat as _madelung_energy_correction_for_lat,
)
from .periodic_grid import build_periodic_becke_grid
from .periodic_rhf_ewald import _canonical_orthogonalizer
from .periodic_scf_accelerators import DynamicDamping, PeriodicSCFAccelerator
from .periodic_uhf_ewald import _spin_squared
from .progress import ProgressLogger, resolve_progress
from .scf_divergence import check_scf_divergence
__all__ = [
"PeriodicUKSEwaldResult",
"run_uks_periodic_gamma_ewald3d",
]
[docs]
@dataclass
class PeriodicUKSEwaldResult:
"""Result of :func:`run_uks_periodic_gamma_ewald3d`.
Spin-resolved counterpart of :class:`PeriodicRKSEwaldResult` with
UHF-style ``s_squared`` diagnostic. All matrices are at Γ.
"""
energy: float
e_electronic: float
e_nuclear: float
e_xc: float
e_coulomb: float
e_hf_exchange: float
n_iter: int
converged: bool
s_squared: float
s_squared_ideal: float
# α spin
mo_energies_alpha: np.ndarray
mo_coeffs_alpha: np.ndarray
density_alpha: np.ndarray
fock_alpha: np.ndarray
# β spin
mo_energies_beta: np.ndarray
mo_coeffs_beta: np.ndarray
density_beta: np.ndarray
fock_beta: np.ndarray
overlap: np.ndarray
functional: str = ""
scf_trace: List[SCFIteration] = field(default_factory=list)
omega: float = 0.0
grid_shape: Tuple[int, int, int] = (0, 0, 0)
def _density_set_gamma(
template: LatticeMatrixSet,
D: np.ndarray,
) -> LatticeMatrixSet:
"""Wrap the Γ-folded density ``D`` in a degenerate
:class:`LatticeMatrixSet` matching the cell layout of ``template``:
block 0 = D, every other block = zeros. Mutates ``template`` in
place (LatticeMatrixSet's ``.blocks`` list is a transient copy;
real mutation goes through ``set_block``)."""
n_bf = D.shape[0]
if n_bf != template.nbf:
raise ValueError(
f"_density_set_gamma: D has nbf={n_bf} but template has nbf={template.nbf}"
)
zero = np.zeros_like(D)
for i in range(len(template)):
template.set_block(i, D if i == 0 else zero)
return template
[docs]
def run_uks_periodic_gamma_ewald3d(
system: PeriodicSystem,
basis: BasisSet,
options: Optional[PeriodicKSOptions] = 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,
) -> PeriodicUKSEwaldResult:
"""Γ-point open-shell periodic Kohn-Sham SCF with Ewald-3D Coulomb.
Parameters mirror :func:`run_rks_periodic_gamma_ewald3d`.
Multiplicity is read from ``system.multiplicity``; the driver
handles closed-shell (mult=1, special case where α=β occupation)
and open-shell systems uniformly.
Returns
-------
:class:`PeriodicUKSEwaldResult`.
"""
opts = options if options is not None else PeriodicKSOptions()
lat_opts = 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 the jellium
# background terms cancel exactly. Mirrors the override block in
# run_rks_periodic_gamma_ewald3d (commit 49f8ae91). The driver
# kwarg ``omega`` is retained for signature parity but overridden
# here; users override via ``opts.ewald_omega`` when needed.
_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"UKS Gamma EWALD_3D / functional={opts.functional!r}, "
f"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,
[
("PeriodicKSOptions", 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))
n_elec = system.n_electrons()
mult = int(system.multiplicity)
if mult < 1:
raise ValueError(
f"run_uks_periodic_gamma_ewald3d: multiplicity must be ≥ 1; got {mult}"
)
n_alpha = (n_elec + mult - 1) // 2
n_beta = (n_elec - mult + 1) // 2
if n_alpha + n_beta != n_elec:
raise ValueError(
"run_uks_periodic_gamma_ewald3d: n_electrons and multiplicity "
f"are inconsistent (n_e={n_elec}, mult={mult})"
)
if n_beta < 0:
raise ValueError(
f"run_uks_periodic_gamma_ewald3d: n_beta < 0; n_e={n_elec}, mult={mult}"
)
# ---- Functional + DFT grid ------------------------------------------
func = Functional(opts.functional, 2) # spin-polarized
alpha = float(func.hf_exchange_fraction)
if opts.use_periodic_becke:
grid = build_periodic_becke_grid(
system,
grid_options=opts.grid,
image_radius_bohr=float(opts.becke_image_radius_bohr),
)
else:
grid = build_grid(system.unit_cell_molecule(), opts.grid)
plog.info(
f"UKS Gamma occupations: n_alpha = {n_alpha}, "
f"n_beta = {n_beta} (multiplicity = {mult})"
)
# ---- 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,
)
opt_rep = optimize_truncation(system, basis, lattice_opts=lat_opts)
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; SCF will run "
"with the last evaluated settings."
)
lat_opts = opt_rep.optimized_lattice_opts
# ---- One-electron integrals at Γ ------------------------------------
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)
k_gamma = np.zeros(3)
S = np.real(bloch_sum(S_lat, k_gamma))
T = np.real(bloch_sum(T_lat, k_gamma))
V = np.real(bloch_sum(V_lat, k_gamma))
Hcore = T + V
S = 0.5 * (S + S.T)
Hcore = 0.5 * (Hcore + Hcore.T)
# Linear-dependence preflight on S(Γ); see periodic_rhf_ewald.py.
from .linear_dependence import scf_preflight_overlap_check
scf_preflight_overlap_check(
S,
plog=plog,
label="S(Γ)",
basis=basis,
)
X, n_kept = _canonical_orthogonalizer(
S,
linear_dep_threshold,
normalize_diag_first=canonical_orth_normalize_diag_first,
)
if max(n_alpha, n_beta) > n_kept:
raise RuntimeError(
"run_uks_periodic_gamma_ewald3d: canonical orthogonalisation "
f"kept {n_kept} directions; need ≥ {max(n_alpha, n_beta)} "
f"(n_α={n_alpha}, n_β={n_beta})."
)
e_nuc = float(nuclear_repulsion_per_cell(system, lat_opts))
def diagonalise(F: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
Fp = X.T @ F @ X
Fp = 0.5 * (Fp + Fp.T)
eps, Cp = np.linalg.eigh(Fp)
return X @ Cp, eps
# ---- Initial guess via the unified engine.
# Bug fix (v0.9.x): this Γ-only UKS path previously ignored
# ``opts.initial_guess`` and always used HCore. The engine call
# respects the request; HCore reproduces the prior behaviour.
guess = getattr(opts, "initial_guess", InitialGuess.HCORE)
# Seed the per-spin MO frame from Hcore for the quadratic-step
# tracking (used even when SAD provides the start density).
C_alpha, eps_alpha = diagonalise(Hcore)
C_beta, eps_beta = diagonalise(Hcore)
split = initial_densities_open_shell(
system.unit_cell_molecule(),
basis,
n_alpha,
n_beta,
guess,
is_periodic=True,
)
if split is not None:
D_alpha, D_beta = split
else:
D_alpha = (
C_alpha[:, :n_alpha] @ C_alpha[:, :n_alpha].T
if n_alpha > 0
else np.zeros_like(Hcore)
)
D_beta = (
C_beta[:, :n_beta] @ C_beta[:, :n_beta].T
if n_beta > 0
else np.zeros_like(Hcore)
)
D_alpha = 0.5 * (D_alpha + D_alpha.T)
D_beta = 0.5 * (D_beta + D_beta.T)
D_alpha_prev = D_alpha.copy()
D_beta_prev = D_beta.copy()
damping = float(opts.damping)
if not (0.0 <= damping < 1.0):
raise ValueError(
f"run_uks_periodic_gamma_ewald3d: 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[PeriodicSCFAccelerator] = (
PeriodicSCFAccelerator(opts) if use_diis else None
)
level_shift = float(getattr(opts, "level_shift", 0.0))
# Phase C1c — quadratic SCF fallback (per-spin 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)
)
# Track per-spin MO basis between iterations for the C1c step.
C_alpha_prev_mo = C_alpha
eps_alpha_prev_mo = eps_alpha
C_beta_prev_mo = C_beta
eps_beta_prev_mo = eps_beta
# Pre-allocate degenerate LMS templates for D_α and D_β (V_xc input).
D_alpha_set = compute_overlap_lattice(basis, system, lat_opts)
D_beta_set = compute_overlap_lattice(basis, system, lat_opts)
# ---- SCF loop -------------------------------------------------------
scf_trace: List[SCFIteration] = []
s_squared_ideal = 0.25 * (mult - 1) * (mult + 1)
result = PeriodicUKSEwaldResult(
energy=0.0,
e_electronic=0.0,
e_nuclear=e_nuc,
e_xc=0.0,
e_coulomb=0.0,
e_hf_exchange=0.0,
n_iter=0,
converged=False,
s_squared=0.0,
s_squared_ideal=s_squared_ideal,
mo_energies_alpha=np.empty(0),
mo_coeffs_alpha=np.empty((0, 0)),
density_alpha=D_alpha.copy(),
fock_alpha=np.empty((0, 0)),
mo_energies_beta=np.empty(0),
mo_coeffs_beta=np.empty((0, 0)),
density_beta=D_beta.copy(),
fock_beta=np.empty((0, 0)),
overlap=S,
functional=str(opts.functional),
scf_trace=scf_trace,
omega=float(omega),
grid_shape=grid_shape_t,
)
E_prev = 0.0
plog.banner(f"SCF (UKS Gamma {opts.functional!r}, EWALD_3D)")
plog.info(" iter energy (Ha) dE ||[F,DS]|| DIIS")
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
if iter_idx == 1 or damping == 0.0 or diis_active:
D_alpha_used = D_alpha
D_beta_used = D_beta
else:
D_alpha_used = damping * D_alpha_prev + (1.0 - damping) * D_alpha
D_beta_used = damping * D_beta_prev + (1.0 - damping) * D_beta
D_total = D_alpha_used + D_beta_used
# Hartree J via Ewald — same convention as UHF Ewald.
J = build_j_ewald_3d(
basis,
system,
D_total,
omega=float(omega),
lattice_opts=lat_opts,
grid_shape=grid_shape_t,
origin=origin,
spacing_bohr=spacing_bohr,
)
# Per-spin K (only when α > 0 — pure DFT skips this).
if alpha != 0.0:
jk_alpha = build_jk_gamma_molecular_limit(
basis,
system,
lat_opts,
2.0 * D_alpha_used,
0.0,
)
jk_beta = build_jk_gamma_molecular_limit(
basis,
system,
lat_opts,
2.0 * D_beta_used,
0.0,
)
# K(D_σ) = K(2·D_σ) / 2 (RHF convention inside the builder).
K_alpha = 0.5 * np.asarray(jk_alpha.K)
K_beta = 0.5 * np.asarray(jk_beta.K)
else:
K_alpha = None
K_beta = None
# V_xc via the periodic UKS XC kernel.
_density_set_gamma(D_alpha_set, D_alpha_used)
_density_set_gamma(D_beta_set, D_beta_used)
xc = build_xc_periodic_uks(
basis,
system,
grid,
func,
D_alpha_set,
D_beta_set,
lat_opts,
)
V_xc_alpha = np.real(bloch_sum(xc.V_alpha, k_gamma))
V_xc_beta = np.real(bloch_sum(xc.V_beta, k_gamma))
V_xc_alpha = 0.5 * (V_xc_alpha + V_xc_alpha.T)
V_xc_beta = 0.5 * (V_xc_beta + V_xc_beta.T)
E_xc = float(xc.e_xc)
# F_σ = Hcore + J − α·K_σ + V_xc^σ
if K_alpha is not None:
F_HF_alpha = J - alpha * K_alpha
F_HF_beta = J - alpha * K_beta
else:
F_HF_alpha = J
F_HF_beta = J
F_alpha = Hcore + F_HF_alpha + V_xc_alpha
F_beta = Hcore + F_HF_beta + V_xc_beta
F_alpha = 0.5 * (F_alpha + F_alpha.T)
F_beta = 0.5 * (F_beta + F_beta.T)
# Energy decomposition.
E_core_trace = float(np.einsum("ij,ij->", D_total, Hcore))
E_HF_alpha_trace = 0.5 * float(np.einsum("ij,ij->", D_alpha_used, F_HF_alpha))
E_HF_beta_trace = 0.5 * float(np.einsum("ij,ij->", D_beta_used, F_HF_beta))
E_elec = E_xc + E_core_trace + E_HF_alpha_trace + E_HF_beta_trace
# Madelung-leak correction (v0.6.1).
E_madelung_fix = _madelung_energy_correction_for_lat(
D_alpha_used + D_beta_used, S, system, lat_opts
)
E_total = E_elec + e_nuc + E_madelung_fix
# Reporting decomposition: J vs HF-K.
E_coulomb = 0.5 * float(np.einsum("ij,ij->", D_total, J))
if K_alpha is not None:
E_hf_K = -0.5 * alpha * float(
np.einsum("ij,ij->", D_alpha_used, K_alpha)
) - 0.5 * alpha * float(np.einsum("ij,ij->", D_beta_used, K_beta))
else:
E_hf_K = 0.0
# Per-spin orbital-gradient norms.
FDS_alpha = F_alpha @ D_alpha_used @ S
FDS_beta = F_beta @ D_beta_used @ S
grad_alpha = FDS_alpha - FDS_alpha.T
grad_beta = FDS_beta - FDS_beta.T
grad_norm = float(
np.sqrt(np.linalg.norm(grad_alpha) ** 2 + np.linalg.norm(grad_beta) ** 2)
)
dE = E_total - E_prev
# Divergence detection (v0.6.2).
check_scf_divergence(
"run_uks_periodic_gamma_ewald3d",
iter_idx,
E_total,
grad_norm,
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),
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),
diis=diis_sub,
)
# Per-iter energy decomposition for cross-code parity comparison.
E_kin_iter = float(np.einsum("ij,ij->", D_total, T))
E_ne_iter = float(np.einsum("ij,ij->", D_total, V))
plog.energy_decomposition(
iter_idx,
E_kin=E_kin_iter,
E_ne=E_ne_iter,
E_J=E_coulomb,
E_xc=float(E_xc),
E_K=float(E_hf_K),
E_nuc=float(e_nuc),
E_madelung=float(E_madelung_fix),
)
converged = (
iter_idx > 1
and abs(dE) < float(opts.conv_tol_energy)
and grad_norm < float(opts.conv_tol_grad)
)
# Phase C1c gate.
in_quadratic_phase = (
quadratic_fallback_iter > 0 and iter_idx > quadratic_fallback_iter
)
if in_quadratic_phase:
from .quadratic_scf import quadratic_step
C_alpha_new, eps_alpha_new = quadratic_step(
F_alpha,
C_alpha_prev_mo,
eps_alpha_prev_mo,
n_alpha,
shift=quadratic_fallback_shift,
max_step=quadratic_fallback_max_step,
)
C_beta_new, eps_beta_new = quadratic_step(
F_beta,
C_beta_prev_mo,
eps_beta_prev_mo,
n_beta,
shift=quadratic_fallback_shift,
max_step=quadratic_fallback_max_step,
)
else:
# SCF-accelerator extrapolation (DIIS per-spin Pulay,
# KDIIS / EDIIS / ADIIS spin-coupled — matches the C++
# UHF dispatch).
if accel is not None:
F_alpha_ex, F_beta_ex = accel.extrapolate_uhf(
F_alpha,
F_beta,
error_alpha=grad_alpha,
error_beta=grad_beta,
density_alpha=D_alpha_used,
density_beta=D_beta_used,
energy=E_total,
mo_coeffs_alpha=C_alpha_prev_mo,
mo_coeffs_beta=C_beta_prev_mo,
mo_energies_alpha=eps_alpha_prev_mo,
mo_energies_beta=eps_beta_prev_mo,
n_alpha=n_alpha,
n_beta=n_beta,
)
if diis_active:
F_alpha = F_alpha_ex
F_beta = F_beta_ex
# Saunders-Hillier level shift per spin.
if level_shift != 0.0:
F_alpha_diag = (
F_alpha
+ level_shift * S
- (level_shift / 2.0) * (S @ D_alpha_used @ S)
)
F_beta_diag = (
F_beta
+ level_shift * S
- (level_shift / 2.0) * (S @ D_beta_used @ S)
)
F_alpha_diag = 0.5 * (F_alpha_diag + F_alpha_diag.T)
F_beta_diag = 0.5 * (F_beta_diag + F_beta_diag.T)
else:
F_alpha_diag = F_alpha
F_beta_diag = F_beta
C_alpha_new, eps_alpha_new = diagonalise(F_alpha_diag)
C_beta_new, eps_beta_new = diagonalise(F_beta_diag)
C_alpha_prev_mo = C_alpha_new
eps_alpha_prev_mo = eps_alpha_new
C_beta_prev_mo = C_beta_new
eps_beta_prev_mo = eps_beta_new
D_alpha_prev = D_alpha_used
D_beta_prev = D_beta_used
D_alpha = (
C_alpha_new[:, :n_alpha] @ C_alpha_new[:, :n_alpha].T
if n_alpha > 0
else np.zeros_like(Hcore)
)
D_beta = (
C_beta_new[:, :n_beta] @ C_beta_new[:, :n_beta].T
if n_beta > 0
else np.zeros_like(Hcore)
)
D_alpha = 0.5 * (D_alpha + D_alpha.T)
D_beta = 0.5 * (D_beta + D_beta.T)
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 = iter_idx
result.mo_energies_alpha = eps_alpha_new
result.mo_coeffs_alpha = C_alpha_new
result.density_alpha = D_alpha_used
result.fock_alpha = F_alpha
result.mo_energies_beta = eps_beta_new
result.mo_coeffs_beta = C_beta_new
result.density_beta = D_beta_used
result.fock_beta = F_beta
if damper is not None:
damper.update(E_total)
E_prev = E_total
if converged:
# Final consistency pass on the fresh densities.
D_total_f = D_alpha + D_beta
J_f = build_j_ewald_3d(
basis,
system,
D_total_f,
omega=float(omega),
lattice_opts=lat_opts,
grid_shape=grid_shape_t,
origin=origin,
spacing_bohr=spacing_bohr,
)
if alpha != 0.0:
jk_a_f = build_jk_gamma_molecular_limit(
basis,
system,
lat_opts,
2.0 * D_alpha,
0.0,
)
jk_b_f = build_jk_gamma_molecular_limit(
basis,
system,
lat_opts,
2.0 * D_beta,
0.0,
)
K_alpha_f = 0.5 * np.asarray(jk_a_f.K)
K_beta_f = 0.5 * np.asarray(jk_b_f.K)
else:
K_alpha_f = None
K_beta_f = None
_density_set_gamma(D_alpha_set, D_alpha)
_density_set_gamma(D_beta_set, D_beta)
xc_f = build_xc_periodic_uks(
basis,
system,
grid,
func,
D_alpha_set,
D_beta_set,
lat_opts,
)
V_xc_a_f = np.real(bloch_sum(xc_f.V_alpha, k_gamma))
V_xc_b_f = np.real(bloch_sum(xc_f.V_beta, k_gamma))
V_xc_a_f = 0.5 * (V_xc_a_f + V_xc_a_f.T)
V_xc_b_f = 0.5 * (V_xc_b_f + V_xc_b_f.T)
E_xc_f = float(xc_f.e_xc)
if K_alpha_f is not None:
F_HF_a_f = J_f - alpha * K_alpha_f
F_HF_b_f = J_f - alpha * K_beta_f
else:
F_HF_a_f = J_f
F_HF_b_f = J_f
F_alpha_f = Hcore + F_HF_a_f + V_xc_a_f
F_beta_f = Hcore + F_HF_b_f + V_xc_b_f
F_alpha_f = 0.5 * (F_alpha_f + F_alpha_f.T)
F_beta_f = 0.5 * (F_beta_f + F_beta_f.T)
C_a_f, eps_a_f = diagonalise(F_alpha_f)
C_b_f, eps_b_f = diagonalise(F_beta_f)
E_core_f = float(np.einsum("ij,ij->", D_alpha + D_beta, Hcore))
E_HF_a_f = 0.5 * float(np.einsum("ij,ij->", D_alpha, F_HF_a_f))
E_HF_b_f = 0.5 * float(np.einsum("ij,ij->", D_beta, F_HF_b_f))
E_elec_f = E_xc_f + E_core_f + E_HF_a_f + E_HF_b_f
E_coulomb_f = 0.5 * float(np.einsum("ij,ij->", D_alpha + D_beta, J_f))
if K_alpha_f is not None:
E_hf_K_f = -0.5 * alpha * float(
np.einsum("ij,ij->", D_alpha, K_alpha_f)
) - 0.5 * alpha * float(np.einsum("ij,ij->", D_beta, K_beta_f))
else:
E_hf_K_f = 0.0
E_madelung_fix_f = _madelung_energy_correction_for_lat(
D_alpha + D_beta, S, system, lat_opts
)
result.energy = E_elec_f + e_nuc + E_madelung_fix_f
result.e_electronic = E_elec_f
result.e_xc = E_xc_f
result.e_coulomb = E_coulomb_f
result.e_hf_exchange = E_hf_K_f
result.mo_energies_alpha = eps_a_f
result.mo_coeffs_alpha = C_a_f
result.density_alpha = D_alpha
result.fock_alpha = F_alpha_f
result.mo_energies_beta = eps_b_f
result.mo_coeffs_beta = C_b_f
result.density_beta = D_beta
result.fock_beta = F_beta_f
result.s_squared = _spin_squared(
n_alpha,
n_beta,
C_a_f,
C_b_f,
S,
)
result.converged = True
plog.converged(
n_iter=result.n_iter,
energy=result.energy,
converged=True,
)
return result
# Did not converge — populate ⟨S²⟩ on the final iterate.
result.s_squared = _spin_squared(
n_alpha,
n_beta,
result.mo_coeffs_alpha,
result.mo_coeffs_beta,
S,
)
plog.converged(
n_iter=result.n_iter,
energy=result.energy,
converged=False,
)
return result