"""BIPOLE-style periodic UKS driver in CRYSTAL's electrostatic gauge.
The open-shell DFT counterpart of :mod:`vibeqc.pbc_bipole_rks`.
Uses the same CRYSTAL-gauge Ewald J-split F^2e build for the Hartree
term and adds the spin-polarised libxc XC potential on the periodic
Becke grid. For hybrid functionals a fraction of HF exchange is
retained per spin channel via the native ``build_jk_2e_real_space``.
This driver keeps the same multi-k SCF scaffold as the UHF BIPOLE
driver: shared Ewald a across V_ne/E_nn/J^LR, analytic V_ne via
AO-pair Fourier transforms, real-space energy accounting, and
the full suite of convergence accelerators (DIIS, ODA, MOM,
level-shift schedule, damping).
"""
from __future__ import annotations
import warnings
from dataclasses import dataclass, field
from typing import List, Optional, Sequence, Tuple, Union
import numpy as np
from ._vibeqc_core import (
BasisSet,
BlochKMesh,
EwaldOptions,
Functional,
GridOptions,
InitialGuess,
LatticeMatrixSet,
LatticeSumOptions,
PeriodicKSOptions,
PeriodicSystem,
SCFIteration,
bloch_sum,
build_grid,
build_xc_periodic_uks,
compute_kinetic_lattice,
compute_overlap_lattice,
direct_lattice_cells,
ewald_nuclear_repulsion,
make_lattice_matrix_set,
nuclear_repulsion_per_cell,
real_space_density_from_kpoints_fractional,
)
from .bipole_ext_el_pole import compute_ext_el_spheropole
from .guess import initial_densities_open_shell, initial_density_closed_shell
from .level_shift_schedule import LevelShiftSchedule
from .mom import select_occupied_by_max_overlap as _mom_select
from .oda import compute_oda_lambda as _compute_oda_lambda
from .oda import oda_mix_densities as _oda_mix
from .pbc_bipole_common import (
PBCBipoleEnergyComponents,
_bloch_sum_blocks,
_combine_density_sets,
_compute_nuclear_lattice_ewald_reciprocal_ft,
_copy_lattice_with_blocks, # noqa: F401 (compat re-export, see tests/test_pbc_bipole_common.py)
_crystal_ewald_options,
_default_bipole_v_ne_grid_options,
_density_set_gamma_or_lattice,
_expand_ibz_kmesh_for_ewald_j,
_lattice_contract,
_spin_occupations,
_zero_cross_cell_density,
bvk_torus_density_matrices,
prepare_bipole_lattice_options,
resolve_bipole_fock_symmetry,
warn_bipole_charged_cell,
warn_bipole_legacy_multik_gauge,
home_cell_block,
smearing_basin_warning,
)
from .pbc_bipole_fock import (
BipoleFockContext,
build_bipole_unrestricted_fock,
)
from .periodic_grid import build_periodic_becke_grid
from .periodic_rhf_multi_k_ewald import (
_canonical_orthogonalizer_complex,
_damp_lattice_matrix,
_diag_in_orth_basis,
)
from .periodic_scf_accelerators import (
DynamicDamping,
MultiKPeriodicUHFAccelerator,
)
from .periodic_uhf_ewald import _spin_squared
from .periodic_v_ne import compute_nuclear_lattice_dispatch
from .periodic_screened_exchange import resolve_periodic_exchange
from .progress import ProgressLogger, resolve_progress
from .scf_divergence import check_scf_divergence
from .smearing import (
SmearingOptions,
)
from .smearing import (
closed_shell_periodic_occupations as _closed_shell_periodic_occupations,
)
from .symmetry_integrals_reduced import (
compute_kinetic_lattice_reduced,
compute_overlap_lattice_reduced,
)
__all__ = [
"PBCBipoleUKSResult",
"run_pbc_bipole_uks",
]
[docs]
@dataclass
class PBCBipoleUKSResult:
"""Result of :func:`run_pbc_bipole_uks`."""
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
mo_energies_alpha: List[np.ndarray]
mo_coeffs_alpha: List[np.ndarray]
fock_alpha: List[np.ndarray]
density_alpha: LatticeMatrixSet
mo_energies_beta: List[np.ndarray]
mo_coeffs_beta: List[np.ndarray]
fock_beta: List[np.ndarray]
density_beta: LatticeMatrixSet
overlap: List[np.ndarray]
hcore: List[np.ndarray]
scf_trace: List[SCFIteration] = field(default_factory=list)
e_ext_el_spheropole: Optional[float] = None
ewald_alpha_bohr_inv: Optional[float] = None
# Dudarev DFT+U contribution per unit cell (Hartree). 0 unless the
# caller passed ``dft_plus_u=[HubbardSite(...)]``. Sum of per-spin
# contributions for the open-shell UKS path.
e_dft_plus_u: float = 0.0
energy_components: List[PBCBipoleEnergyComponents] = field(
default_factory=list,
)
# Gauge provenance (option (b)): see run_pbc_bipole_uhf.
exchange_ewald_split: bool = False
exchange_exxdiv: Optional[str] = None
# Non-None when finite-T smearing straddled the HOMO-LUMO gap and may
# have selected a near-metallic basin (ionic-Γ basin trap).
basin_warning: Optional[str] = None
functional: str = ""
# Smearing diagnostics (zero when smearing_temperature == 0).
smearing_temperature: float = 0.0
fermi_level: float = 0.0
entropy: float = 0.0
free_energy: float = 0.0
occupations_alpha: List[np.ndarray] = field(default_factory=list)
occupations_beta: List[np.ndarray] = field(default_factory=list)
# Cartesian k-points (bohr^-1) and weights this result spans, in the
# same order as the per-k, per-spin ``mo_coeffs_*`` / ``mo_energies_*``
# lists. Lets optional single-k output writers locate Gamma instead of
# assuming the first k-point is Gamma. See PBCBipoleRHFResult.
kpoints_cart: Optional[np.ndarray] = None
kpoint_weights: Optional[np.ndarray] = None
@dataclass
class _PBCBipoleUKSFockBuild:
"""Internal Fock bundle for one UKS density pair."""
f2e_alpha_real: LatticeMatrixSet
f2e_beta_real: LatticeMatrixSet
f_alpha_k_list: List[np.ndarray]
f_beta_k_list: List[np.ndarray]
e_j_short_range: Optional[float] = None
e_j_long_range: Optional[float] = None
e_exchange: Optional[float] = None
e_j_multipole: Optional[float] = None
# k-space per-spin exchange correction (Ewald exchange split,
# a_HF-scaled); lives in F_s(k) -- caller adds to E_2e.
e_2e_k_correction: float = 0.0
# XC energy of the SAME grid quadrature whose per-spin V_xc entered
# f_alpha/beta_k_list (None for the HF-like PATOM seed build).
# Consuming this instead of re-running build_xc_periodic_uks keeps
# the iteration's V_xc and e_xc from one build.
e_xc: Optional[float] = None
[docs]
def run_pbc_bipole_uks(
system: PeriodicSystem,
basis: BasisSet,
kmesh: BlochKMesh,
options=None,
*,
functional: Optional[str] = None,
linear_dep_threshold: float = 1e-7,
canonical_orth_normalize_diag_first: bool = True,
level_shift_schedule: Optional[LevelShiftSchedule] = None,
use_mom: bool = False,
use_oda: bool = False,
oda_trust_lambda_max: float = 1.0,
use_incremental_fock: bool = True,
use_ewald_j_split: Optional[bool] = None,
ewald_omega: Optional[float] = None,
ewald_precision: float = 1e-8,
v_ne_grid_options: Optional[GridOptions] = None,
use_multipole_far_field: bool = False,
use_exchange_ewald_split: Optional[bool] = None,
exchange_exxdiv: str = "ewald",
multipole_l_max: int = 2,
use_fock_symmetry: Optional[bool] = None,
use_fock_symmetry_reduce: bool = False,
sr_image_extent_bohr: Optional[float] = None,
progress: Union[bool, ProgressLogger, None] = None,
verbose: Optional[int] = None,
init_alpha: Optional[Sequence[np.ndarray]] = None,
init_beta: Optional[Sequence[np.ndarray]] = None,
dft_plus_u: Optional[List["HubbardSite"]] = None,
) -> PBCBipoleUKSResult:
"""Multi-k open-shell UKS via the CRYSTAL-gauge BIPOLE scaffold.
``dft_plus_u``: optional list of :class:`HubbardSite`. Same
per-spin per-k convention as :func:`run_pbc_bipole_uhf` --
``F_s(k) += S(k) V_AO_s S(k)`` with
``V_AO_s = U_eff (1/2 - n_s)``; ``E_U_total`` is summed over
spins and surfaced via ``result.e_dft_plus_u``.
"""
from ._vibeqc_core import PeriodicKSOptions as _PKSOpts
opts = options if options is not None else _PKSOpts()
# SPINLOCK. SPIN_SCHEDULE (two-phase) is delegated before SCF setup;
# PATTERN_HOLD reuses the per-k MOM machinery below (gated to the
# spinlock window).
from ._vibeqc_core import SpinlockMode
from .spinlock_periodic import check_spinlock_support, run_spin_schedule
if (
getattr(opts, "spinlock_mode", SpinlockMode.OFF) == SpinlockMode.SPIN_SCHEDULE
and int(getattr(opts, "spinlock_iterations", 0)) > 0
):
return run_spin_schedule(
lambda sysx, o: run_pbc_bipole_uks(
sysx,
basis,
kmesh,
o,
functional=functional,
linear_dep_threshold=linear_dep_threshold,
canonical_orth_normalize_diag_first=canonical_orth_normalize_diag_first,
level_shift_schedule=level_shift_schedule,
use_mom=use_mom,
use_oda=use_oda,
oda_trust_lambda_max=oda_trust_lambda_max,
use_incremental_fock=use_incremental_fock,
use_ewald_j_split=use_ewald_j_split,
ewald_omega=ewald_omega,
ewald_precision=ewald_precision,
v_ne_grid_options=v_ne_grid_options,
use_multipole_far_field=use_multipole_far_field,
use_exchange_ewald_split=use_exchange_ewald_split,
exchange_exxdiv=exchange_exxdiv,
multipole_l_max=multipole_l_max,
use_fock_symmetry=use_fock_symmetry,
use_fock_symmetry_reduce=use_fock_symmetry_reduce,
progress=progress,
verbose=verbose,
dft_plus_u=dft_plus_u,
),
system,
opts,
)
check_spinlock_support(
opts,
{SpinlockMode.PATTERN_HOLD, SpinlockMode.SPIN_SCHEDULE},
"the BIPOLE UKS driver",
)
if functional is not None:
opts.functional = str(functional)
if not getattr(opts, "functional", None):
opts.functional = "pbe"
smearing_T = float(getattr(opts, "smearing_temperature", 0.0))
if smearing_T < 0.0:
raise ValueError("run_pbc_bipole_uks: smearing_temperature must be >= 0")
func = Functional(opts.functional, 2) # spin-polarised
# CAM exchange assembly (see run_pbc_bipole_rks): alpha_hf is the
# FULL-RANGE arm only; screened hybrids (hse06, c_full = 0) get
# their erfc short-range K inside build_bipole_unrestricted_fock
# (no seam, no far-field); LR-heavy RSH fails closed here.
exx = resolve_periodic_exchange(func, where="run_pbc_bipole_uks")
alpha_hf = float(exx.c_full)
fock_mixing_value = float(getattr(opts, "fock_mixing", 0.0))
if fock_mixing_value == 0.0 and alpha_hf < 1.0:
# CRYSTAL-style default FMIXING for pure DFT: 30% previous
# Fock matrix weight to damp SCF oscillations on tight
# ionic cells. Zero for HF (alpha_hf == 1.0).
fock_mixing_value = 0.30
if not (0.0 <= fock_mixing_value < 1.0):
raise ValueError(
f"run_pbc_bipole_uks: fock_mixing must be in [0, 1); "
f"got {fock_mixing_value}"
)
lat_opts: LatticeSumOptions = opts.lattice_opts
plog = resolve_progress(progress, verbose=verbose)
(
use_ewald_j_split,
use_ewald_j_split_auto,
lat_opts_2e,
lat_opts_1e,
) = prepare_bipole_lattice_options(system, lat_opts, use_ewald_j_split, plog)
n_elec = int(system.n_electrons())
n_alpha, n_beta = _spin_occupations(system)
mult = int(system.multiplicity)
plog.info(
f"PBC BIPOLE UKS (CRYSTAL-gauge) / cutoff {lat_opts.cutoff_bohr:.2f} bohr"
)
if exx.is_screened:
plog.info(
f" functional = {opts.functional}, screened exchange "
f"K = {exx.c_sr:g}*K_erfc(omega={exx.omega_screen:g} bohr^-1)"
)
else:
plog.info(
f" functional = {opts.functional}, hf_exchange_fraction = {alpha_hf}"
)
plog.info(f" n_alpha = {n_alpha}, n_beta = {n_beta}, multiplicity = {mult}")
_kmesh_ibz = kmesh
_ir_mapping = np.asarray(getattr(kmesh, "ir_mapping", []), dtype=int).reshape(-1)
k_points = list(_kmesh_ibz.kpoints)
weights = np.asarray(_kmesh_ibz.weights, dtype=float)
if use_ewald_j_split and _ir_mapping.size > 0:
# IBZ inputs run on the EXPANDED full mesh (correctness: the
# IBZ-native replication shortcut lacked the star AO rotations
# D(R.k) = P(R).D(k).P(R)ᵀ and broke non-trivial crystals -- MgO
# probe 2026-06-10, 8.25 Ha. See vibeqc.periodic_k_symmetry for
# the transport groundwork and pbc_bipole.py for the full note).
kmesh_full = _expand_ibz_kmesh_for_ewald_j(system, kmesh, plog)
if len(list(kmesh_full.kpoints)) > len(k_points):
kmesh = kmesh_full
k_points = list(kmesh.kpoints)
weights = np.asarray(kmesh.weights, dtype=float)
_ir_mapping = np.asarray([], dtype=int)
k_points_full = k_points
weights_full = weights
else:
k_points_full = k_points
weights_full = weights
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}")
if use_ewald_j_split and n_k > 1 and _ir_mapping.size == 0:
uniform_w = 1.0 / float(n_k)
if not np.allclose(weights, uniform_w, atol=1e-9):
raise ValueError(
"use_ewald_j_split at multi-k requires uniform weights or IBZ mesh"
)
plog.info(f"k-mesh: {n_k} k-points, weights sum = {weights.sum():.4f}")
# ---- Exchange/gauge resolution (option (b) Phase 4b, 2026-06-11) --
# Mirrors run_pbc_bipole_uhf: per-spin exchange split at 3D Γ,
# alpha_HF-scaled for hybrids; pure functionals get the gauge change
# only (no exchange machinery).
if exchange_exxdiv not in ("ewald", "none"):
raise ValueError(
f"run_pbc_bipole_uks: exchange_exxdiv must be 'ewald' or "
f"'none'; got {exchange_exxdiv!r}"
)
_x_split_auto = use_exchange_ewald_split is None
exchange_split_active = (
bool(use_ewald_j_split) if _x_split_auto else bool(use_exchange_ewald_split)
)
if exchange_split_active and not use_ewald_j_split:
raise ValueError(
"run_pbc_bipole_uks: use_exchange_ewald_split=True requires "
"the Ewald J split (use_ewald_j_split=True)."
)
# Multi-k split (option (b) Phase 3): needs the true Monkhorst-
# Pack dimensions (q-channel tables, BvK-torus fold, supercell
# Madelung); ad-hoc k-lists are rejected.
_bvk_mesh: Optional[Tuple[int, int, int]] = None
if exchange_split_active and n_k > 1:
_mesh_attr = tuple(int(x) for x in getattr(kmesh, "mesh", (1, 1, 1)))
if int(np.prod(_mesh_attr)) != n_k:
if _x_split_auto:
plog.info(
" multi-k corrected exchange gauge needs a "
"Monkhorst-Pack mesh (BvK-torus fold + supercell ξ_M); "
"this ad-hoc k-list has no mesh metadata -> falling back "
"to the legacy gauge. Pass a monkhorst_pack(...) mesh "
"for the corrected gauge."
)
exchange_split_active = False
else:
raise ValueError(
"run_pbc_bipole_uks: the Ewald exchange split at multi-k "
"requires a Monkhorst-Pack BlochKMesh carrying its mesh "
f"dimensions (got mesh={_mesh_attr} for {n_k} k-points). "
"Build the mesh via monkhorst_pack(...); ad-hoc k-point "
"lists are not supported on the corrected gauge."
)
else:
_bvk_mesh = _mesh_attr
warn_bipole_legacy_multik_gauge(system, exchange_split_active, n_k, plog)
warn_bipole_charged_cell(system, plog)
_xi_madelung = 0.0
if exchange_split_active and exchange_exxdiv == "ewald" and alpha_hf > 0.0:
if n_k > 1:
from .bipole_fock_ewald import probe_charge_madelung_supercell
assert _bvk_mesh is not None
_xi_madelung = probe_charge_madelung_supercell(system, _bvk_mesh)
else:
from .bipole_fock_ewald import probe_charge_madelung
_xi_madelung = probe_charge_madelung(system)
if exchange_split_active:
plog.info(
" Gauge: corrected (full Bloch density, no spheropole"
+ (
f"; per-spin exchange split exxdiv={exchange_exxdiv}, "
f"alpha_HF={alpha_hf:g}"
if alpha_hf > 0.0
else "; pure functional"
)
+ ")"
)
# Shared Ewald state
ewald_options_1e = None
omega_used = None
ewald_cell_volume = None
ewald_k_max: Optional[float] = None
if system.dim == 3:
from .bipole_ext_el_pole import (
crystal_default_ewald_alpha,
crystal_ewald_reciprocal_cutoff,
)
V_cell = float(abs(np.linalg.det(np.asarray(system.lattice, dtype=float))))
ewald_cell_volume = V_cell
omega_used = (
float(ewald_omega)
if ewald_omega is not None
else crystal_default_ewald_alpha(V_cell)
)
ewald_k_max = crystal_ewald_reciprocal_cutoff(V_cell)
ewald_options_1e = _crystal_ewald_options(
lat_opts_1e,
alpha_bohr_inv=omega_used,
tolerance=float(ewald_precision),
recip_cutoff_bohr_inv=ewald_k_max,
)
# DFT grid
if getattr(opts, "use_periodic_becke", False):
grid = build_periodic_becke_grid(
system,
grid_options=opts.grid,
image_radius_bohr=float(getattr(opts, "becke_image_radius_bohr", 10.0)),
)
else:
grid = build_grid(system.unit_cell_molecule(), opts.grid)
# One-electron integrals
with plog.stage(
"integrals_lattice", detail=f"S/T/V at cutoff {lat_opts.cutoff_bohr:.2f} bohr"
):
_use_sym = bool(getattr(getattr(system, "symmetry", None), "operations", None))
if _use_sym:
ops = system.symmetry.operations
cells = direct_lattice_cells(system, lat_opts_2e.cutoff_bohr)
plog.info(
f"S/T integrals: symmetry-reduced path "
f"(SG {system.symmetry.international_symbol}, "
f"{system.symmetry.order} ops, "
f"{len(cells)} lattice cells)"
)
_, S_blocks = compute_overlap_lattice_reduced(
basis,
system,
lat_opts_2e,
ops,
)
S_lat = make_lattice_matrix_set(
basis.nbasis, cells, [np.asarray(b, dtype=float) for b in S_blocks]
)
_, T_blocks = compute_kinetic_lattice_reduced(
basis,
system,
lat_opts_2e,
ops,
)
T_lat = make_lattice_matrix_set(
basis.nbasis, cells, [np.asarray(b, dtype=float) for b in T_blocks]
)
else:
S_lat = compute_overlap_lattice(basis, system, lat_opts_2e)
T_lat = compute_kinetic_lattice(basis, system, lat_opts_2e)
v_ne_lr_cache = None
if (
system.dim == 3
and ewald_options_1e is not None
and v_ne_grid_options is None
):
V_lat, v_ne_lr_cache = _compute_nuclear_lattice_ewald_reciprocal_ft(
basis,
system,
lat_opts_1e,
ewald_options_1e,
S_lat,
precision=ewald_precision,
K_max=ewald_k_max,
)
else:
v_ne_grid = (
v_ne_grid_options
if v_ne_grid_options is not None
else (_default_bipole_v_ne_grid_options() if system.dim == 3 else None)
)
V_lat = compute_nuclear_lattice_dispatch(
basis,
system,
lat_opts_1e,
grid_options=v_ne_grid,
ewald_options=ewald_options_1e,
)
cells = list(S_lat.cells)
# Lattice-fold guard + wide density list under the corrected gauge
# (rationale: run_pbc_bipole_rhf).
if exchange_split_active:
from .pbc_bipole_common import s_fold_truncation_drift
_s_drift = s_fold_truncation_drift(
basis,
system,
lat_opts_2e,
k_points=(k_points if n_k > 1 else None),
)
_s_label = (
"S(k) fold truncation (max over mesh)"
if n_k > 1
else "S(Γ) fold truncation"
)
if _s_drift > 1e-2:
plog.info(
f" WARNING: {_s_label} {_s_drift:.1e} -- lattice "
f"sums badly under-converged; absolute energies UNRELIABLE. "
f"Increase cutoff_bohr."
)
elif _s_drift > 1e-4:
plog.info(f" note: {_s_label} {_s_drift:.1e}")
cells_density = list(
direct_lattice_cells(system, 2.0 * float(lat_opts_2e.cutoff_bohr))
)
plog.info(f" density cell list: {len(cells_density)} cells (2x cutoff)")
else:
cells_density = cells
# Per-k matrices
from .linear_dependence import scf_preflight_overlap_check
S_k_list = []
Hcore_k_list = []
X_k_list = []
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})", 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_pbc_bipole_uks: canonical orth at k={k_idx} dropped too many directions"
)
S_k_list.append(S_k)
Hcore_k_list.append(H_k)
X_k_list.append(X_k)
# Nuclear repulsion
e_nuc = (
float(ewald_nuclear_repulsion(system, ewald_options_1e))
if ewald_options_1e is not None
else float(nuclear_repulsion_per_cell(system, lat_opts_1e))
)
# Initial guess
C_alpha_per_k = []
eps_alpha_per_k = []
C_beta_per_k = []
eps_beta_per_k = []
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)
def _spin_density(C_per_k_local, n_occ_each):
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)
result = real_space_density_from_kpoints_fractional(
C_per_k_local, occ_per_k, kmesh, cells_density
)
if not exchange_split_active:
_zero_cross_cell_density(result, nbf, n_k)
return result
D_alpha_real = _spin_density(C_alpha_per_k, n_alpha)
D_beta_real = _spin_density(C_beta_per_k, n_beta)
# --- Smearing occupation helpers (open-shell: per-spin mu) ---
from .smearing.fermi_dirac import fermi_dirac_occupations_per_k as _fd_per_k
def _occupations_per_spin(
eps_spin_per_k: Sequence[np.ndarray],
n_spin: int,
):
"""Compute per-spin fractional occupations via Fermi-Dirac.
Returns (occ_per_k, mu, entropy_spin) where occupations are in
[0, 1] per band (single-spin), mu is the spin chemical
potential in Hartree, and entropy_spin is S/k_B for that spin.
"""
if smearing_T <= 0.0 or n_spin == 0:
nbf = eps_spin_per_k[0].shape[0] if eps_spin_per_k else 0
occ = []
for eps in eps_spin_per_k:
o = np.zeros(eps.shape[0], dtype=float)
o[:n_spin] = 1.0
occ.append(o)
return occ, 0.0, 0.0
# The closed-shell FD returns occupations in [0, 2]; divide by 2
# for per-spin occupations in [0, 1].
occ_double, mu, entropy_double = _fd_per_k(
eps_spin_per_k,
weights,
float(2 * n_spin), # target: 2 electrons per occupied orbital
smearing_T,
)
occ = [np.asarray(o, dtype=float) * 0.5 for o in occ_double]
return occ, mu, entropy_double * 0.5
occ_alpha_per_k, mu_alpha, entropy_alpha = _occupations_per_spin(
eps_alpha_per_k, n_alpha
)
occ_beta_per_k, mu_beta, entropy_beta = _occupations_per_spin(
eps_beta_per_k, n_beta
)
entropy = entropy_alpha + entropy_beta
if smearing_T > 0.0:
D_alpha_real = real_space_density_from_kpoints_fractional(
C_alpha_per_k, occ_alpha_per_k, kmesh, cells_density
)
D_beta_real = real_space_density_from_kpoints_fractional(
C_beta_per_k, occ_beta_per_k, kmesh, cells_density
)
if not exchange_split_active:
_zero_cross_cell_density(D_alpha_real, D_alpha_real.blocks[0].shape[0], n_k)
_zero_cross_cell_density(D_beta_real, D_beta_real.blocks[0].shape[0], n_k)
# Caller-supplied warm-start spin densities take precedence over
# the SAD/Hcore guess engine -- see run_pbc_bipole_uhf for the
# same contract. Used by the NEB driver for within-image density
# warm-start for periodic NEB.
if (init_alpha is not None) != (init_beta is not None):
raise ValueError(
"run_pbc_bipole_uks: init_alpha and init_beta must be "
"provided together (both None or both populated)"
)
if init_alpha is not None and init_beta is not None:
blocks_a = list(init_alpha)
blocks_b = list(init_beta)
if len(blocks_a) != len(D_alpha_real.cells):
raise ValueError(
f"run_pbc_bipole_uks: init_alpha has {len(blocks_a)} "
f"blocks; expected {len(D_alpha_real.cells)}"
)
if len(blocks_b) != len(D_beta_real.cells):
raise ValueError(
f"run_pbc_bipole_uks: init_beta has {len(blocks_b)} "
f"blocks; expected {len(D_beta_real.cells)}"
)
for g_idx, (ba, bb) in enumerate(zip(blocks_a, blocks_b)):
D_alpha_real.set_block(g_idx, np.asarray(ba, dtype=float))
D_beta_real.set_block(g_idx, np.asarray(bb, dtype=float))
initial_density_is_local = True
density_from_c_per_k = False
else:
guess = getattr(opts, "initial_guess", InitialGuess.HCORE)
# ATOMSPIN: see pbc_bipole_uhf -- when a per-atom spin pattern is
# requested it must reach the open-shell engine, not the even-electron
# closed-shell shortcut (which would wash it out to a 50/50 split).
_atomic_spins = getattr(opts, "atomic_spins", None) or None
# READ carries the prior's per-spin g=0 densities -- skip the
# even-electron 50/50 shortcut for READ too (see pbc_bipole_uhf).
_is_read = guess == InitialGuess.READ
_needs_open_context = guess in {
InitialGuess.SAP,
InitialGuess.PATOM,
InitialGuess.HUECKEL,
InitialGuess.MINAO,
}
_patom_seed_pending = guess == InitialGuess.PATOM
D_guess = None
if (
n_elec % 2 == 0
and _atomic_spins is None
and not _is_read
and not _needs_open_context
):
D_total_guess = initial_density_closed_shell(
system.unit_cell_molecule(), basis, n_elec // 2, guess, is_periodic=True
)
if D_total_guess is not None:
D_guess = (
0.5 * np.asarray(D_total_guess, dtype=float),
0.5 * np.asarray(D_total_guess, dtype=float),
)
if D_guess is None:
_seed_guess = InitialGuess.SAD if _patom_seed_pending else guess
D_guess = initial_densities_open_shell(
system.unit_cell_molecule(),
basis,
n_alpha,
n_beta,
_seed_guess,
is_periodic=True,
periodic_system=system,
lattice_opts=lat_opts,
atomic_spins=_atomic_spins,
read_density_alpha=getattr(opts, "read_density_alpha", None),
read_density_beta=getattr(opts, "read_density_beta", None),
read_path=getattr(opts, "read_path", ""),
)
initial_density_is_local = D_guess is not None
if D_guess is not None:
D_a0, D_b0 = D_guess
# Multi-k cross-cell density fix (v0.14.0) -- same rationale as
# pbc_bipole_rks.py. For n_k > 1, keep the per-k fold's P(g!=0)
# blocks; they restore image-cell AO overlap on the DFT grid.
if n_k > 1:
plog.info(
f"initial guess: {guess.name} (g=0 from GuessEngine, "
f"g!=0 from per-k fold for cross-cell DFT grid density)"
)
for g_idx in range(len(D_alpha_real.cells)):
is_g0 = (
np.asarray(D_alpha_real.cells[g_idx].index, dtype=int)
== np.array([0, 0, 0])
).all()
if is_g0:
D_alpha_real.set_block(g_idx, D_a0)
D_beta_real.set_block(g_idx, D_b0)
else:
zero_a = np.zeros_like(D_a0, dtype=float)
zero_b = np.zeros_like(D_b0, dtype=float)
for g_idx in range(len(D_alpha_real.cells)):
is_g0 = (
np.asarray(D_alpha_real.cells[g_idx].index, dtype=int)
== np.array([0, 0, 0])
).all()
D_alpha_real.set_block(g_idx, D_a0 if is_g0 else zero_a)
D_beta_real.set_block(g_idx, D_b0 if is_g0 else zero_b)
density_from_c_per_k = not initial_density_is_local
if init_alpha is not None and init_beta is not None:
_patom_seed_pending = False
D_alpha_prev = None
D_beta_prev = None
# SCF aids
damping = float(opts.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
)
level_shift_static = float(getattr(opts, "level_shift", 0.0))
if level_shift_schedule is not None and not isinstance(
level_shift_schedule, LevelShiftSchedule
):
raise TypeError(f"level_shift_schedule must be LevelShiftSchedule")
# CRYSTAL-style FMIXING (per-spin, per-k).
# Value resolved above in the options setup; log it here near the
# SCF-loop preamble.
if fock_mixing_value != 0.0:
plog.info(
f"fock mixing: CRYSTAL FMIXING "
f"{100.0 * fock_mixing_value:.1f}% "
"(previous Fock matrix weight)"
)
F_alpha_prev_mixed: Optional[List[np.ndarray]] = None
F_beta_prev_mixed: Optional[List[np.ndarray]] = None
# SPINLOCK PATTERN_HOLD reuses the per-k MOM machinery (gated to cycles
# 2..spinlock_iterations) to hold the seeded broken-symmetry occupied set
# then release -- protecting an ATOMSPIN seed.
_spinlock_pattern_hold = (
getattr(opts, "spinlock_mode", SpinlockMode.OFF) == SpinlockMode.PATTERN_HOLD
and int(getattr(opts, "spinlock_iterations", 0)) > 0
)
_spinlock_iters = int(getattr(opts, "spinlock_iterations", 0))
if use_mom:
plog.info("MOM: ON")
elif _spinlock_pattern_hold:
plog.info(f"SPINLOCK PATTERN_HOLD: MOM-hold for {_spinlock_iters} cycles")
C_prev_occ_alpha_per_k = None
C_prev_occ_beta_per_k = None
if use_oda and use_diis:
raise ValueError("use_oda and use_diis are mutually exclusive")
if use_oda:
plog.info(f"ODA: ON (trust lambda_max = {oda_trust_lambda_max})")
# Ewald J-split cache
j_lr_cache = v_ne_lr_cache
if use_ewald_j_split:
if system.dim != 3:
raise ValueError(f"use_ewald_j_split requires dim=3")
from .bipole_fock_ewald import _build_j_long_range_cache
assert omega_used is not None
cells_r_cart_arr = np.array(
[np.asarray(c.r_cart, dtype=float) for c in cells], dtype=float
)
if j_lr_cache is None:
j_lr_cache = _build_j_long_range_cache(
basis,
system,
cells_r_cart_arr,
omega_used,
ewald_precision,
K_max=ewald_k_max,
)
# Multi-k split: per-(k,k′) q-channel tables for the per-spin LR
# exchange (hybrids only). Shares the J^LR w/K_max envelope.
x_lr_cache = None
if exchange_split_active and n_k > 1 and (alpha_hf > 0.0 or _patom_seed_pending):
from .bipole_fock_ewald import build_k_exchange_long_range_cache
assert j_lr_cache is not None and ewald_k_max is not None
x_lr_cache = build_k_exchange_long_range_cache(
basis,
system,
j_lr_cache,
K_max=ewald_k_max,
)
plog.info(
f" K^LR q-channels: {n_k} distinct q = k-k′ shifts on the "
f"shared K_max = {ewald_k_max:.2f} bohr⁻¹ envelope "
f"(per spin, a_HF = {alpha_hf:g})"
)
def _split_k_density_list(density) -> List[np.ndarray]:
"""Per-k (single-spin) density matrices for the split paths.
Home block at Γ; exact BvK-torus fold at multi-k (see
``bvk_torus_density_matrices``).
"""
if n_k == 1:
return [home_cell_block(density).astype(complex)]
assert _bvk_mesh is not None
return bvk_torus_density_matrices(density, k_points, _bvk_mesh)
# Incremental/differential J_SR(+K_SR) accumulators (opt-in; see
# run_pbc_bipole_rhf). Hybrids build per-spin J_SR+K_SR (two
# accumulators on D_a/D_b); pure functionals build J_SR on the total
# density (one accumulator). Corrected gauge + DIIS only.
incremental_jk_alpha = None
incremental_jk_beta = None
incremental_jk_total = None
if use_incremental_fock:
if exchange_split_active and not use_oda:
from .bipole_fock_ewald import IncrementalJK
if alpha_hf > 0.0:
incremental_jk_alpha = IncrementalJK()
incremental_jk_beta = IncrementalJK()
else:
incremental_jk_total = IncrementalJK()
plog.info(
" incremental Fock (differential J_SR/K_SR via ΔD "
"density-envelope screening): ON"
)
else:
plog.info(
" incremental Fock requested but inactive "
"(needs the corrected gauge + DIIS, not ODA)"
)
# (M2d) The UKS two-electron Fock assembly is routed through
# pbc_bipole_fock.build_bipole_unrestricted_fock; the thin wrapper
# below (defined after multipole + SYM3b resolution) adds the
# per-spin V_xc and assembles F(k) in the driver.
# ---- Multipole far-field config (resolve once before SCF loop) -----
from .bipole_fock_multipole import ( # noqa: E402
BipoleMultipoleConfig,
resolve_multipole_config,
)
if exchange_split_active and use_multipole_far_field:
raise NotImplementedError(
"run_pbc_bipole_uks: the multipole far-field J replacement has "
"not been re-validated under the corrected gauge. Pass "
"use_exchange_ewald_split=False or omit use_multipole_far_field."
)
_mp_config = resolve_multipole_config(
system,
basis,
lat_opts_2e,
user_enable=(False if exchange_split_active else use_multipole_far_field),
multipole_l_max=multipole_l_max,
)
if _mp_config.enabled:
plog.info(
f" BIPOLE multipole far-field: ENABLED "
f"(L_max={_mp_config.L_max}, R_bipole={_mp_config.R_bipole:.1f} bohr, "
f"n_cells={len(_mp_config.cache.cells) if _mp_config.cache else 0})"
)
else:
plog.info(
f" BIPOLE multipole far-field: off "
f"(R_bipole={_mp_config.R_bipole:.1f} bohr, "
f"cutoff={lat_opts_2e.cutoff_bohr:.1f} bohr)"
)
# SYM3b Fock symmetry enforcement is OPT-IN ONLY -- rationale
# (boundary truncation asymmetry) lives on the shared resolver.
_fock_sym_map, _rep_cell_indices = resolve_bipole_fock_symmetry(
system,
basis,
lat_opts_2e,
use_fock_symmetry,
use_fock_symmetry_reduce,
plog,
per_spin=True,
)
# ---- Shared unrestricted Fock-build context (M2d unification) -------
# Bundle the per-run invariants the inline UKS Fock build captured.
# Placed after multipole and symmetry resolution, before the PATOM
# in-field step and SCF loop first call the wrapper.
_fock_ctx = BipoleFockContext(
basis=basis,
system=system,
lat_opts_2e=lat_opts_2e,
use_ewald_j_split=use_ewald_j_split,
exchange_split_active=exchange_split_active,
n_k=n_k,
omega_used=omega_used,
ewald_precision=ewald_precision,
ewald_cell_volume=ewald_cell_volume,
n_elec=n_elec,
xi_madelung=_xi_madelung,
j_lr_cache=j_lr_cache,
x_lr_cache=x_lr_cache,
incremental_jk=None,
rep_cell_indices=_rep_cell_indices,
fock_sym_map=_fock_sym_map,
mp_config=_mp_config,
s_lat=S_lat,
s_k_list=S_k_list,
k_points=k_points,
weights=weights,
k_points_full=k_points_full,
weights_full=weights_full,
ir_mapping=_ir_mapping,
bvk_mesh=_bvk_mesh,
n_occ=n_alpha,
plog=plog,
sr_image_extent=sr_image_extent_bohr,
)
def _build_fock_for_density(
D_alpha,
D_beta,
*,
coeffs_alpha_for_rho,
coeffs_beta_for_rho,
use_incremental=True,
patom_hf_like=False,
):
fb = build_bipole_unrestricted_fock(
_fock_ctx,
D_alpha,
D_beta,
n_alpha=n_alpha,
n_beta=n_beta,
coeffs_alpha_for_rho=coeffs_alpha_for_rho,
coeffs_beta_for_rho=coeffs_beta_for_rho,
incremental_jk_alpha=incremental_jk_alpha,
incremental_jk_beta=incremental_jk_beta,
incremental_jk_total=incremental_jk_total,
use_incremental=use_incremental,
alpha_hf=alpha_hf,
exchange_assembly=exx,
patom_hf_like=patom_hf_like,
legacy_exchange_halved=True,
use_reduced_j_only=True,
)
# XC potential -- open-shell UKS API: per-spin V_xc(g) + total
# E_xc. PATOM is an HF-like in-field seed, independent of the
# target functional, so skip XC for that one pre-SCF Fock build.
xc_result = None
if not patom_hf_like:
D_xc_alpha = _density_set_gamma_or_lattice(S_lat, D_alpha)
D_xc_beta = _density_set_gamma_or_lattice(S_lat, D_beta)
xc_result = build_xc_periodic_uks(
basis, system, grid, func, D_xc_alpha, D_xc_beta, lat_opts_2e
)
# Per-k Fock: F_s(k) = F^2e_s(k) + V_xc,s(k) + Hcore(k) −
# K_corr,s(k), hermitised after V_xc so a non-Hermitian Bloch
# sum is symmetrised exactly as the pre-M2d inline build did.
f_alpha_k_list = []
f_beta_k_list = []
for k_idx, k in enumerate(k_points):
k_arr = np.asarray(k, dtype=float)
F_a_2e = _bloch_sum_blocks(
fb.f2e_alpha_real.blocks, fb.f2e_alpha_real.cells, k_arr
)
F_b_2e = _bloch_sum_blocks(
fb.f2e_beta_real.blocks, fb.f2e_beta_real.cells, k_arr
)
if xc_result is None:
Vxc_a_k = 0.0
Vxc_b_k = 0.0
else:
Vxc_a_k = np.asarray(bloch_sum(xc_result.V_alpha, k_arr))
Vxc_b_k = np.asarray(bloch_sum(xc_result.V_beta, k_arr))
F_a = F_a_2e + Vxc_a_k + np.asarray(Hcore_k_list[k_idx], dtype=complex)
F_b = F_b_2e + Vxc_b_k + np.asarray(Hcore_k_list[k_idx], dtype=complex)
if fb.k_corr_alpha_per_k is not None:
# Per-spin Ewald exchange split: one matrix per k (a
# single Γ entry at n_k = 1).
F_a = F_a - fb.k_corr_alpha_per_k[k_idx]
F_b = F_b - fb.k_corr_beta_per_k[k_idx]
f_alpha_k_list.append(0.5 * (F_a + F_a.conj().T))
f_beta_k_list.append(0.5 * (F_b + F_b.conj().T))
return _PBCBipoleUKSFockBuild(
f2e_alpha_real=fb.f2e_alpha_real,
f2e_beta_real=fb.f2e_beta_real,
f_alpha_k_list=f_alpha_k_list,
f_beta_k_list=f_beta_k_list,
e_j_short_range=fb.e_j_short_range,
e_j_long_range=fb.e_j_long_range,
e_exchange=fb.e_exchange,
e_j_multipole=fb.e_j_multipole,
e_2e_k_correction=fb.e_2e_k_correction,
e_xc=(None if xc_result is None else float(xc_result.e_xc)),
)
if _patom_seed_pending:
plog.info("initial guess: PATOM (SAD + one BIPOLE in-field step)")
patom_fock = _build_fock_for_density(
D_alpha_real,
D_beta_real,
coeffs_alpha_for_rho=None,
coeffs_beta_for_rho=None,
use_incremental=False,
patom_hf_like=True,
)
C_alpha_per_k = []
eps_alpha_per_k = []
C_beta_per_k = []
eps_beta_per_k = []
for idx in range(n_k):
C_a, eps_a = _diag_in_orth_basis(
patom_fock.f_alpha_k_list[idx],
X_k_list[idx],
)
C_b, eps_b = _diag_in_orth_basis(
patom_fock.f_beta_k_list[idx],
X_k_list[idx],
)
C_alpha_per_k.append(C_a)
eps_alpha_per_k.append(eps_a)
C_beta_per_k.append(C_b)
eps_beta_per_k.append(eps_b)
occ_alpha_per_k, mu_alpha, entropy_alpha = _occupations_per_spin(
eps_alpha_per_k, n_alpha
)
occ_beta_per_k, mu_beta, entropy_beta = _occupations_per_spin(
eps_beta_per_k, n_beta
)
entropy = entropy_alpha + entropy_beta
if smearing_T > 0.0:
D_alpha_real = real_space_density_from_kpoints_fractional(
C_alpha_per_k, occ_alpha_per_k, kmesh, cells_density
)
D_beta_real = real_space_density_from_kpoints_fractional(
C_beta_per_k, occ_beta_per_k, kmesh, cells_density
)
if not exchange_split_active:
_zero_cross_cell_density(
D_alpha_real, D_alpha_real.blocks[0].shape[0], n_k
)
_zero_cross_cell_density(
D_beta_real, D_beta_real.blocks[0].shape[0], n_k
)
else:
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 = None
D_beta_prev = None
initial_density_is_local = False
density_from_c_per_k = True
plog.banner("SCF (PBC BIPOLE UKS, direct-space)")
plog.info(" iter energy (Ha) dE ||[F,DS]|| DIIS")
# ---- DFT+U setup (Increment 4d-bipole, UKS) -------------------------
dft_plus_u_sites_cxx: List = []
dft_plus_u_ao_groups: List[List[int]] = []
if dft_plus_u:
from ._vibeqc_core import _HubbardSiteCxx
from .dft_plus_u import ao_group_indices
ao_groups_map = ao_group_indices(basis)
for site in dft_plus_u:
key = (site.atom_index, site.l)
if key not in ao_groups_map:
raise ValueError(
f"run_pbc_bipole_uks: HubbardSite{key} has no AOs "
f"in the basis. Available channels: "
f"{sorted(ao_groups_map.keys())}"
)
dft_plus_u_sites_cxx.append(
_HubbardSiteCxx(site.atom_index, site.l, site.U_eff_hartree)
)
dft_plus_u_ao_groups.append(ao_groups_map[key])
scf_trace = []
energy_components = []
E_prev = 0.0
E_elec = 0.0
e_xc = 0.0
e_dft_plus_u = 0.0
F_alpha_k_list = [np.zeros_like(H) for H in Hcore_k_list]
F_beta_k_list = [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
# SPINLOCK PATTERN_HOLD: the accelerator (DIIS / EDIIS / ADIIS /
# KDIIS -- whatever MultiKPeriodicUHFAccelerator resolved) is
# suspended (no history recorded, no extrapolation, damping stays
# live) while the hold is active. Fock extrapolation across
# held-window iterates steers the SCF toward the symmetric attractor
# by continuous orbital rotation -- a collapse the
# occupation-selecting MOM hold cannot see -- and poisons the
# post-release history with out-of-basin iterates. The history
# starts fresh at release.
hold_active = _spinlock_pattern_hold and iter_idx <= _spinlock_iters
diis_active = (
use_diis and iter_idx >= diis_start_iter and not hold_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
d_used_is_damped = iter_idx > 1 and damping > 0.0 and not diis_active
d_used_from_coeffs = (
density_from_c_per_k
and not (initial_density_is_local and iter_idx == 1)
and not d_used_is_damped
)
fock_build = _build_fock_for_density(
D_alpha_used,
D_beta_used,
coeffs_alpha_for_rho=(C_alpha_per_k if d_used_from_coeffs else None),
coeffs_beta_for_rho=(C_beta_per_k if d_used_from_coeffs else None),
)
F_alpha_k_list = fock_build.f_alpha_k_list
F_beta_k_list = fock_build.f_beta_k_list
# ---- DFT+U: per-spin per-k Fock contribution -----------------
# Same pattern as run_pbc_bipole_uhf (Increment 4d-bipole UHF).
# Bloch-sum the lattice per-spin density to P_s(k); kernel
# returns V_AO_s (real, k-independent) + per-spin E_s; per k
# add S(k) V_AO_s S(k) to F_s(k).
e_dft_plus_u = 0.0
if dft_plus_u_sites_cxx:
from ._vibeqc_core import (
_compute_dft_plus_u_multi_k_per_spin_cxx,
)
P_split_a_for_u: Optional[List[np.ndarray]] = None
P_split_b_for_u: Optional[List[np.ndarray]] = None
if exchange_split_active:
# Unprojected Bloch fold: BvK representatives (home
# block at Γ; exact torus fold at multi-k -- S_g would
# overcount).
P_split_a_for_u = _split_k_density_list(D_alpha_used)
P_split_b_for_u = _split_k_density_list(D_beta_used)
P_alpha_k_for_U: List[np.ndarray] = []
P_beta_k_for_U: List[np.ndarray] = []
for k_idx in range(n_k):
if P_split_a_for_u is not None:
Pa_k = P_split_a_for_u[k_idx]
Pb_k = P_split_b_for_u[k_idx]
else:
k_arr = np.asarray(k_points[k_idx], dtype=float)
Pa_k = _bloch_sum_blocks(
D_alpha_used.blocks,
D_alpha_used.cells,
k_arr,
)
Pb_k = _bloch_sum_blocks(
D_beta_used.blocks,
D_beta_used.cells,
k_arr,
)
Pa_k = 0.5 * (Pa_k + Pa_k.conj().T)
Pb_k = 0.5 * (Pb_k + Pb_k.conj().T)
P_alpha_k_for_U.append(Pa_k)
P_beta_k_for_U.append(Pb_k)
E_a, V_AO_a = _compute_dft_plus_u_multi_k_per_spin_cxx(
dft_plus_u_sites_cxx,
dft_plus_u_ao_groups,
S_k_list,
P_alpha_k_for_U,
list(weights),
)
E_b, V_AO_b = _compute_dft_plus_u_multi_k_per_spin_cxx(
dft_plus_u_sites_cxx,
dft_plus_u_ao_groups,
S_k_list,
P_beta_k_for_U,
list(weights),
)
e_dft_plus_u = float(E_a) + float(E_b)
V_AO_a_cmplx = np.asarray(V_AO_a, dtype=complex)
V_AO_b_cmplx = np.asarray(V_AO_b, dtype=complex)
for k_idx in range(n_k):
S_k = S_k_list[k_idx]
F_alpha_k_list[k_idx] = F_alpha_k_list[k_idx] + (
S_k @ V_AO_a_cmplx @ S_k
)
F_beta_k_list[k_idx] = F_beta_k_list[k_idx] + (S_k @ V_AO_b_cmplx @ S_k)
F_alpha_k_list[k_idx] = 0.5 * (
F_alpha_k_list[k_idx] + F_alpha_k_list[k_idx].conj().T
)
F_beta_k_list[k_idx] = 0.5 * (
F_beta_k_list[k_idx] + F_beta_k_list[k_idx].conj().T
)
# XC energy: from the SAME quadrature that produced this
# iteration's per-spin V_xc (inside _build_fock_for_density on
# D_alpha_used/D_beta_used) -- the loop used to run a second,
# redundant build_xc_periodic_uks here every iteration just to
# extract e_xc.
assert fock_build.e_xc is not None
e_xc = float(fock_build.e_xc)
# Energy
D_total_used = _combine_density_sets(
basis, system, lat_opts_2e, D_alpha_used, D_beta_used
)
E_kin = _lattice_contract(D_total_used, T_lat, operator_name="T")
E_ne = _lattice_contract(D_total_used, V_lat, operator_name="V_ne")
E_2e = fock_build.e_2e_k_correction + 0.5 * (
_lattice_contract(
D_alpha_used, fock_build.f2e_alpha_real, operator_name="F2e_alpha"
)
+ _lattice_contract(
D_beta_used, fock_build.f2e_beta_real, operator_name="F2e_beta"
)
)
E_elec = E_kin + E_ne + E_2e + e_xc
# Gradient
grad_norm_sum = 0.0
error_alpha_k_list = []
error_beta_k_list = []
D_alpha_k_list = []
D_beta_k_list = []
D_k_split_guess_a: Optional[List[np.ndarray]] = None
D_k_split_guess_b: Optional[List[np.ndarray]] = None
if exchange_split_active and initial_density_is_local and iter_idx == 1:
# BvK representatives (home block at Γ; exact torus fold at
# multi-k) for local AND wide warm-start storage.
D_k_split_guess_a = _split_k_density_list(D_alpha_used)
D_k_split_guess_b = _split_k_density_list(D_beta_used)
for idx in range(n_k):
if initial_density_is_local and iter_idx == 1:
if D_k_split_guess_a is not None and D_k_split_guess_b is not None:
D_a_k = D_k_split_guess_a[idx]
D_b_k = D_k_split_guess_b[idx]
else:
k_arr = np.asarray(k_points[idx], dtype=float)
D_a_k = _bloch_sum_blocks(
D_alpha_used.blocks, D_alpha_used.cells, k_arr
)
D_b_k = _bloch_sum_blocks(
D_beta_used.blocks, D_beta_used.cells, k_arr
)
D_a_k = 0.5 * (D_a_k + D_a_k.conj().T)
D_b_k = 0.5 * (D_b_k + D_b_k.conj().T)
else:
C_a = C_alpha_per_k[idx]
C_b = C_beta_per_k[idx]
if smearing_T <= 0.0:
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
else:
# Fractional-occupation per-spin density
occ_a = np.asarray(occ_alpha_per_k[idx], dtype=float)
occ_b = np.asarray(occ_beta_per_k[idx], dtype=float)
C_a_full = np.asarray(C_a, dtype=np.complex128)
C_b_full = np.asarray(C_b, dtype=np.complex128)
D_a_k = (C_a_full * occ_a[None, :]) @ C_a_full.conj().T
D_b_k = (C_b_full * occ_b[None, :]) @ C_b_full.conj().T
D_alpha_k_list.append(D_a_k)
D_beta_k_list.append(D_b_k)
S_k = S_k_list[idx]
F_a_k = F_alpha_k_list[idx]
F_b_k = F_beta_k_list[idx]
err_a = F_a_k @ D_a_k @ S_k - (F_a_k @ D_a_k @ S_k).conj().T
err_b = F_b_k @ D_b_k @ S_k - (F_b_k @ D_b_k @ S_k).conj().T
error_alpha_k_list.append(err_a)
error_beta_k_list.append(err_b)
grad_norm_sum += float(weights[idx]) * float(
np.sqrt(np.linalg.norm(err_a) ** 2 + np.linalg.norm(err_b) ** 2)
)
E_total = float(E_elec) + e_nuc + e_dft_plus_u
# EXT EL-SPHEROPOLE -- uses total (alpha+beta) density.
D_total_used = _combine_density_sets(
basis, system, lat_opts_2e, D_alpha_used, D_beta_used
)
# EXT EL-SPHEROPOLE -- a 3D-Ewald-gauge correction, identically
# zero in the direct (non-Ewald) gauge used for dim<3.
if system.dim == 3 and not exchange_split_active:
E_sphero = compute_ext_el_spheropole(D_total_used, basis, system, lat_opts)
E_total += E_sphero
else:
E_sphero = None
dE = E_total - E_prev if iter_idx > 1 else 0.0
free_energy = E_total - smearing_T * entropy
check_scf_divergence(
"run_pbc_bipole_uks", iter_idx, free_energy, grad_norm_sum, dE
)
diis_sub = max(
accel.subspace_size if accel is not None else 0,
accel.subspace_size if accel is not None else 0,
)
scf_trace.append(
SCFIteration(
iter=iter_idx,
energy=float(free_energy),
delta_e=float(dE),
grad_norm=float(grad_norm_sum),
diis_subspace=diis_sub,
)
)
plog.iteration(
iter_idx,
energy=float(free_energy),
dE=float(dE),
grad=float(grad_norm_sum),
diis=diis_sub,
)
energy_components.append(
PBCBipoleEnergyComponents(
iter=int(iter_idx),
e_total=float(E_total),
e_electronic=float(E_elec),
e_kinetic=float(E_kin),
e_nuclear_attraction=float(E_ne),
e_two_electron=float(E_2e),
e_nuclear_repulsion=float(e_nuc),
e_bielet_zone_ee=(None if use_ewald_j_split else float(E_2e)),
e_j_short_range=fock_build.e_j_short_range,
e_j_long_range=fock_build.e_j_long_range,
e_exchange=fock_build.e_exchange,
e_ext_el_spheropole=E_sphero,
)
)
plog.energy_decomposition(
iter_idx,
E_kin=float(E_kin),
E_ne=float(E_ne),
E_2e=float(E_2e),
E_elec=float(E_elec),
E_nuc=float(e_nuc),
)
converged = (
iter_idx > 1
and abs(dE) < float(opts.conv_tol_energy)
and grad_norm_sum < float(opts.conv_tol_grad)
)
# SCF-accelerator extrapolation. Skipped entirely (not even
# recorded) while the PATTERN_HOLD window is active; see the
# hold_active note at the loop head.
if accel is not None and not hold_active:
if exchange_split_active:
# Unprojected Bloch fold: BvK representatives per k
# (home block at Γ; exact torus fold at multi-k).
# Mirrors the RHF accelerator fold; only the bridged
# EDIIS/ADIIS modes consume these densities.
density_alpha_k_list = _split_k_density_list(D_alpha_used)
density_beta_k_list = _split_k_density_list(D_beta_used)
else:
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=free_energy,
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),
)
# On the converged iteration, diagonalise the *physical* Fock
# F(D) -- not the extrapolated one -- so the final orbitals/
# density stay on the converged fixed point. A near-machine-
# zero DIIS error history (SCF nailing the solution in one
# step) makes the Pulay B-matrix singular and its extrapolated
# Fock garbage. See the RHF twin in pbc_bipole.py and
# tests/test_pbc_bipole_diis_converged_basin.py.
if diis_active and not converged:
F_alpha_k_list = F_a_ex
F_beta_k_list = F_b_ex
# --- FMIXING (per-spin, per-k, after DIIS, before level-shift)
# Skipped on the converged iteration (see DIIS note above).
if fock_mixing_value != 0.0 and not converged:
if F_alpha_prev_mixed is not None:
for spin_list, prev_list in [
(F_alpha_k_list, F_alpha_prev_mixed),
(F_beta_k_list, F_beta_prev_mixed),
]:
for idx in range(n_k):
F_mixed = (1.0 - fock_mixing_value) * spin_list[
idx
] + fock_mixing_value * prev_list[idx]
spin_list[idx] = 0.5 * (F_mixed + F_mixed.conj().T)
F_alpha_prev_mixed = [
np.asarray(F, dtype=complex).copy() for F in F_alpha_k_list
]
F_beta_prev_mixed = [
np.asarray(F, dtype=complex).copy() for F in F_beta_k_list
]
# Level shift
if level_shift_schedule is not None:
level_shift_b = level_shift_schedule.at(iter_idx)
else:
level_shift_b = level_shift_static
if level_shift_b != 0.0 and not converged:
F_alpha_for_diag = []
F_beta_for_diag = []
for idx in range(n_k):
S_k = S_k_list[idx]
D_a_k = D_alpha_k_list[idx]
D_b_k = D_beta_k_list[idx]
F_a_shift = (
F_alpha_k_list[idx]
+ level_shift_b * S_k
- (level_shift_b / 2.0) * (S_k @ D_a_k @ S_k)
)
F_b_shift = (
F_beta_k_list[idx]
+ level_shift_b * S_k
- (level_shift_b / 2.0) * (S_k @ D_b_k @ S_k)
)
F_alpha_for_diag.append(0.5 * (F_a_shift + F_a_shift.conj().T))
F_beta_for_diag.append(0.5 * (F_b_shift + F_b_shift.conj().T))
else:
F_alpha_for_diag = F_alpha_k_list
F_beta_for_diag = F_beta_k_list
# Diagonalise
new_C_alpha = []
new_eps_alpha = []
new_C_beta = []
new_eps_beta = []
for idx in range(n_k):
C_a, eps_a = _diag_in_orth_basis(F_alpha_for_diag[idx], X_k_list[idx])
C_b, eps_b = _diag_in_orth_basis(F_beta_for_diag[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)
# MOM (use_mom: every cycle; SPINLOCK PATTERN_HOLD: cycles
# 2..spinlock_iterations, then release).
_mom_this_iter = use_mom or (
_spinlock_pattern_hold and 1 < iter_idx <= _spinlock_iters
)
if _mom_this_iter and C_prev_occ_alpha_per_k is not None:
for idx in range(n_k):
for spin, (C_k, eps_k, n_occ_spin, C_prev_occ_k) in enumerate(
[
(
new_C_alpha[idx],
new_eps_alpha[idx],
n_alpha,
C_prev_occ_alpha_per_k[idx],
),
(
new_C_beta[idx],
new_eps_beta[idx],
n_beta,
C_prev_occ_beta_per_k[idx],
),
]
):
if n_occ_spin == 0:
continue
S_k = S_k_list[idx]
sel = _mom_select(C_k, S_k, C_prev_occ_k, n_occ_spin, eps_new=eps_k)
n_kept_idx = C_k.shape[1]
virt_mask = np.ones(n_kept_idx, dtype=bool)
virt_mask[sel] = False
virt_sel = np.where(virt_mask)[0]
virt_sel = virt_sel[np.argsort(np.real(eps_k[virt_sel]))]
order = np.concatenate([sel, virt_sel])
if spin == 0:
new_C_alpha[idx] = C_k[:, order]
new_eps_alpha[idx] = eps_k[order]
else:
new_C_beta[idx] = C_k[:, order]
new_eps_beta[idx] = eps_k[order]
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
occ_alpha_per_k, mu_alpha, entropy_alpha = _occupations_per_spin(
eps_alpha_per_k, n_alpha
)
occ_beta_per_k, mu_beta, entropy_beta = _occupations_per_spin(
eps_beta_per_k, n_beta
)
entropy = entropy_alpha + entropy_beta
if smearing_T > 0.0:
D_alpha_new = real_space_density_from_kpoints_fractional(
C_alpha_per_k, occ_alpha_per_k, kmesh, cells_density
)
D_beta_new = real_space_density_from_kpoints_fractional(
C_beta_per_k, occ_beta_per_k, kmesh, cells_density
)
if not exchange_split_active:
_zero_cross_cell_density(
D_alpha_new, D_alpha_new.blocks[0].shape[0], n_k
)
_zero_cross_cell_density(D_beta_new, D_beta_new.blocks[0].shape[0], n_k)
else:
D_alpha_new = _spin_density(C_alpha_per_k, n_alpha)
D_beta_new = _spin_density(C_beta_per_k, n_beta)
# ODA
if use_oda:
fock_naive = _build_fock_for_density(
D_alpha_new,
D_beta_new,
coeffs_alpha_for_rho=C_alpha_per_k,
coeffs_beta_for_rho=C_beta_per_k,
use_incremental=False, # off the per-iter ΔD chain
)
oda_step = _compute_oda_lambda(
D_alpha_used,
D_alpha_new,
F_alpha_k_list,
fock_naive.f_alpha_k_list,
[np.asarray(k) for k in k_points],
weights,
trust_lambda_max=oda_trust_lambda_max,
)
_oda_mix(D_alpha_used, D_alpha_new, oda_step.lam)
_oda_mix(D_beta_used, D_beta_new, oda_step.lam)
D_alpha_prev = D_alpha_real
D_beta_prev = D_beta_real
D_alpha_real = D_alpha_used
D_beta_real = D_beta_used
density_from_c_per_k = oda_step.lam == 1.0
plog.info(
f" ODA: lambda = {oda_step.lam:.4f} (g0 = {oda_step.g0:+.3e}, g1 = {oda_step.g1:+.3e})"
)
else:
D_alpha_prev = D_alpha_used
D_beta_prev = D_beta_used
D_alpha_real = D_alpha_new
D_beta_real = D_beta_new
density_from_c_per_k = True
# Snapshot for next iter MOM (use_mom: every cycle; PATTERN_HOLD: while
# inside the hold window).
if use_mom or (_spinlock_pattern_hold and iter_idx <= _spinlock_iters):
C_prev_occ_alpha_per_k = [
np.asarray(C_alpha_per_k[idx][:, :n_alpha]).copy()
if n_alpha > 0
else np.zeros((C_alpha_per_k[idx].shape[0], 0), dtype=complex)
for idx in range(n_k)
]
C_prev_occ_beta_per_k = [
np.asarray(C_beta_per_k[idx][:, :n_beta]).copy()
if n_beta > 0
else np.zeros((C_beta_per_k[idx].shape[0], 0), dtype=complex)
for idx in range(n_k)
]
if damper is not None:
damper.update(free_energy)
E_prev = free_energy
if converged:
break
# S^2
if n_alpha == 0 or n_beta == 0:
s2 = 0.25 * (n_alpha - n_beta) * (n_alpha - n_beta + 2) + n_beta
else:
k0_idx = next(
idx
for idx, k in enumerate(k_points)
if np.allclose(np.asarray(k, dtype=float), 0.0)
)
s2 = _spin_squared(
n_alpha,
n_beta,
np.real(C_alpha_per_k[k0_idx]),
np.real(C_beta_per_k[k0_idx]),
np.real(S_k_list[k0_idx]),
)
plog.converged(n_iter=iter_idx, energy=E_total, converged=converged)
# ---- Ionic-Γ basin-health diagnostic (smearing straddle, per spin) -
# See run_pbc_bipole_rks / smearing_basin_warning. Both spin channels
# are checked; the minimum per-k gap across spins governs the trigger.
basin_warning = smearing_basin_warning(
smearing_T,
[
(eps_alpha_per_k, occ_alpha_per_k, n_alpha, 1.0),
(eps_beta_per_k, occ_beta_per_k, n_beta, 1.0),
],
entropy,
"run_pbc_bipole_uks",
)
if basin_warning is not None:
plog.info(" WARNING: " + basin_warning)
warnings.warn(basin_warning, UserWarning, stacklevel=2)
free_energy_final = E_total - smearing_T * entropy
# ---- Post-loop: recompute energy on final density for consistency
if converged:
_fb = _build_fock_for_density(
D_alpha_real,
D_beta_real,
coeffs_alpha_for_rho=C_alpha_per_k,
coeffs_beta_for_rho=C_beta_per_k,
use_incremental=False,
)
D_xc_a = _density_set_gamma_or_lattice(S_lat, D_alpha_real)
D_xc_b = _density_set_gamma_or_lattice(S_lat, D_beta_real)
_xc = build_xc_periodic_uks(
basis, system, grid, func, D_xc_a, D_xc_b, lat_opts_2e
)
D_tot = _combine_density_sets(
basis, system, lat_opts_2e, D_alpha_real, D_beta_real
)
E_kin_f = _lattice_contract(D_tot, T_lat, operator_name="T")
E_ne_f = _lattice_contract(D_tot, V_lat, operator_name="V_ne")
E_2e_f = _fb.e_2e_k_correction + 0.5 * (
_lattice_contract(D_alpha_real, _fb.f2e_alpha_real, operator_name="F2e")
+ _lattice_contract(D_beta_real, _fb.f2e_beta_real, operator_name="F2e")
)
E_elec = E_kin_f + E_ne_f + E_2e_f + float(_xc.e_xc)
E_total = float(E_elec) + e_nuc + e_dft_plus_u
# Fresh E_total doesn't include spheropole -- add it (3D only; the
# term is zero in the direct gauge used for dim<3).
if system.dim == 3 and not exchange_split_active:
E_sphero_final = compute_ext_el_spheropole(D_tot, basis, system, lat_opts)
E_total += E_sphero_final
else:
E_sphero_final = None
else:
# Non-converged: surface the last iterate's spheropole from
# the per-iteration breakdown (mirrors run_pbc_bipole_rks /
# run_pbc_bipole_uhf, whose non-converged branch does the same).
E_sphero_final = energy_components[-1].e_ext_el_spheropole
return PBCBipoleUKSResult(
energy=float(E_total),
e_electronic=float(E_elec),
e_nuclear=e_nuc,
e_ext_el_spheropole=E_sphero_final,
e_xc=float(e_xc),
e_coulomb=float(E_2e),
e_hf_exchange=float(fock_build.e_exchange or 0.0),
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,
ewald_alpha_bohr_inv=omega_used,
e_dft_plus_u=float(e_dft_plus_u),
energy_components=energy_components,
exchange_ewald_split=bool(exchange_split_active),
exchange_exxdiv=(exchange_exxdiv if exchange_split_active else None),
basin_warning=basin_warning,
functional=str(opts.functional),
smearing_temperature=smearing_T,
fermi_level=float(mu_alpha) if n_alpha > 0 else float(mu_beta),
entropy=float(entropy),
free_energy=float(free_energy_final),
occupations_alpha=[np.asarray(o, dtype=float) for o in occ_alpha_per_k],
occupations_beta=[np.asarray(o, dtype=float) for o in occ_beta_per_k],
kpoints_cart=np.asarray(k_points, dtype=float).reshape(-1, 3),
kpoint_weights=np.asarray(weights, dtype=float).reshape(-1),
)