"""BIPOLE periodic atomic gradient — G3.
Computes the per-cell nuclear gradient for BIPOLE SCF results.
Uses existing C++ functions for 1-electron and nuclear contributions.
The 2-electron contribution (Ewald J-split derivatives) is handled via
the direct-space ERI gradient for J^SR + K, with a finite-difference
fallback for the full gradient.
Status (2026-05-20): scaffold with 1-e + nuclear parts exact, 2-e
part via existing ``eri_lattice_gradient_contribution`` (approximate
for the Ewald J-split gauge — the J^LR reciprocal derivative is not
yet included). For production use, the finite-difference fallback
``compute_bipole_gradient_fd`` is exact but 6N× more expensive.
The full analytic BIPOLE gradient needs:
- Derivatives of J^SR (erfc-screened ERI gradients) — partially
captured by ``eri_lattice_gradient_contribution``
- Derivatives of J^LR (reciprocal-space AO-pair FT derivatives)
- Derivatives of the background potential V_bg·S (overlap derivatives)
"""
from __future__ import annotations
from typing import List, Optional
import numpy as np
from ._vibeqc_core import (
BasisSet,
LatticeMatrixSet,
LatticeSumOptions,
PeriodicSystem,
compute_overlap_lattice,
eri_lattice_gradient_contribution,
kinetic_lattice_gradient_contribution,
nuclear_lattice_gradient_contribution,
nuclear_repulsion_gradient_per_cell,
overlap_lattice_gradient_contribution,
)
from .pbc_bipole import PBCBipoleRHFResult
from .pbc_bipole_rks import PBCBipoleRKSResult
from .pbc_bipole_uhf import PBCBipoleUHFResult
from .pbc_bipole_uks import PBCBipoleUKSResult
__all__ = [
"compute_bipole_gradient_rhf",
"compute_bipole_gradient_uhf",
"compute_bipole_gradient_rks",
"compute_bipole_gradient_uks",
"compute_bipole_gradient_fd",
"compute_stress_tensor",
]
def _build_energy_weighted_density_closed(
mo_coeffs: List[np.ndarray],
mo_energies: List[np.ndarray],
n_occ: int,
) -> List[np.ndarray]:
"""Build per-k W(k) = 2·Σ_i ε_i·C_i·C_i†(k) for closed-shell."""
W_k_list: List[np.ndarray] = []
for C_k, eps_k in zip(mo_coeffs, mo_energies):
C = np.asarray(C_k)
eps = np.asarray(np.real(eps_k))
if n_occ > C.shape[1]:
raise ValueError(f"n_occ={n_occ} exceeds n_mo={C.shape[1]}")
C_occ = C[:, :n_occ]
W_k = 2.0 * (C_occ * eps[:n_occ][None, :]) @ C_occ.conj().T
W_k_list.append(np.real(W_k))
return W_k_list
def _build_energy_weighted_density_open(
mo_coeffs_alpha: List[np.ndarray],
mo_energies_alpha: List[np.ndarray],
mo_coeffs_beta: List[np.ndarray],
mo_energies_beta: List[np.ndarray],
n_alpha: int,
n_beta: int,
) -> List[np.ndarray]:
"""Build per-k total W(k) = W_α(k) + W_β(k) for open-shell."""
W_k_list: List[np.ndarray] = []
for C_a, eps_a, C_b, eps_b in zip(
mo_coeffs_alpha,
mo_energies_alpha,
mo_coeffs_beta,
mo_energies_beta,
):
C_a_arr = np.asarray(C_a)
C_b_arr = np.asarray(C_b)
eps_a_arr = np.asarray(np.real(eps_a))
eps_b_arr = np.asarray(np.real(eps_b))
W_k = np.zeros_like(C_a_arr, dtype=float)
if n_alpha > 0:
Ca_occ = C_a_arr[:, :n_alpha]
W_k += (Ca_occ * eps_a_arr[:n_alpha][None, :]) @ Ca_occ.conj().T
if n_beta > 0:
Cb_occ = C_b_arr[:, :n_beta]
W_k += (Cb_occ * eps_b_arr[:n_beta][None, :]) @ Cb_occ.conj().T
W_k_list.append(np.real(W_k))
return W_k_list
def _gamma_lattice_set(template: LatticeMatrixSet, M: np.ndarray) -> LatticeMatrixSet:
"""Fill a LatticeMatrixSet with M in every cell block."""
M_arr = np.asarray(M, dtype=np.float64)
for c in range(len(template.cells)):
template.set_block(c, M_arr)
return template
def _compute_bipole_gradient(
system: PeriodicSystem,
basis: BasisSet,
D_real: LatticeMatrixSet,
W_k_list: List[np.ndarray],
n_elec: int,
*,
lattice_opts: LatticeSumOptions,
alpha_hf: float,
ewald_alpha: Optional[float],
) -> np.ndarray:
"""Shared gradient computation for all BIPOLE methods."""
n_atoms = len(system.unit_cell)
W_gamma = W_k_list[0]
W_set = compute_overlap_lattice(basis, system, lattice_opts)
_gamma_lattice_set(W_set, W_gamma)
grad = np.zeros((n_atoms, 3), dtype=np.float64)
# 1-electron + nuclear (exact)
grad += np.asarray(nuclear_repulsion_gradient_per_cell(system, lattice_opts))
grad += np.asarray(
overlap_lattice_gradient_contribution(basis, system, W_set, lattice_opts)
)
grad += np.asarray(
kinetic_lattice_gradient_contribution(basis, system, D_real, lattice_opts)
)
grad += np.asarray(
nuclear_lattice_gradient_contribution(basis, system, D_real, lattice_opts)
)
# Background potential V_bg·S
if ewald_alpha is not None and ewald_alpha > 0 and system.dim == 3:
V_cell = float(abs(np.linalg.det(np.asarray(system.lattice, dtype=float))))
v_bg = -np.pi * float(n_elec) / (float(ewald_alpha) ** 2 * V_cell)
D_bg_set = compute_overlap_lattice(basis, system, lattice_opts)
for c in range(len(D_bg_set.cells)):
D_bg_set.set_block(c, v_bg * np.asarray(D_real.blocks[c], dtype=float))
grad += np.asarray(
overlap_lattice_gradient_contribution(basis, system, D_bg_set, lattice_opts)
)
# 2-electron
grad += np.asarray(
eri_lattice_gradient_contribution(
basis, system, D_real, lattice_opts, float(alpha_hf)
)
)
return grad
[docs]
def compute_stress_tensor(
system: PeriodicSystem,
gradient: np.ndarray,
) -> np.ndarray:
"""Force-based stress tensor from atomic gradients.
Computes the 3×3 stress tensor in Ha/bohr³:
σ_{ij} = -(1/V) · Σ_A R_{A,i} · F_{A,j}
where R_A are atomic positions and F_A = -dE/dR_A.
The energy derivative w.r.t. strain is: dE/dε_{ij} = V · σ_{ij}.
Parameters
----------
system : PeriodicSystem
The periodic system (provides lattice + atomic positions).
gradient : (n_atoms, 3) ndarray
Atomic gradient in Ha/bohr (negative of forces).
Returns
-------
ndarray shape (3, 3)
Stress tensor in Ha/bohr³.
"""
lattice = np.asarray(system.lattice, dtype=float)
V = float(abs(np.linalg.det(lattice)))
if V < 1e-14:
raise ValueError(f"Degenerate lattice (V={V})")
grad = np.asarray(gradient, dtype=float)
n_atoms = len(system.unit_cell)
if grad.shape != (n_atoms, 3):
raise ValueError(f"Gradient shape {grad.shape} != ({n_atoms}, 3)")
stress = np.zeros((3, 3), dtype=float)
for a in range(n_atoms):
R = np.asarray(system.unit_cell[a].xyz, dtype=float)
for i in range(3):
for j in range(3):
stress[i, j] -= R[i] * grad[a, j]
stress /= V
return stress
def compute_bipole_gradient_rhf(
system: PeriodicSystem,
basis: BasisSet,
result: PBCBipoleRHFResult,
*,
lattice_opts: Optional[LatticeSumOptions] = None,
alpha_hf: float = 1.0,
) -> np.ndarray:
"""BIPOLE RHF atomic gradient."""
if not result.converged:
import warnings
warnings.warn(
f"compute_bipole_gradient_rhf: result not converged "
f"(n_iter={result.n_iter}). Gradient may be inaccurate."
)
if lattice_opts is None:
lattice_opts = LatticeSumOptions()
n_elec = system.n_electrons()
n_occ = n_elec // 2
W_k_list = _build_energy_weighted_density_closed(
result.mo_coeffs, result.mo_energies, n_occ
)
return _compute_bipole_gradient(
system,
basis,
result.density,
W_k_list,
n_elec,
lattice_opts=lattice_opts,
alpha_hf=alpha_hf,
ewald_alpha=getattr(result, "ewald_alpha_bohr_inv", None),
)
def compute_bipole_gradient_uhf(
system: PeriodicSystem,
basis: BasisSet,
result: PBCBipoleUHFResult,
*,
lattice_opts: Optional[LatticeSumOptions] = None,
alpha_hf: float = 1.0,
) -> np.ndarray:
"""BIPOLE UHF atomic gradient."""
if not result.converged:
import warnings
warnings.warn("compute_bipole_gradient_uhf: not converged")
if lattice_opts is None:
lattice_opts = LatticeSumOptions()
n_elec = system.n_electrons()
n_alpha = (n_elec + system.multiplicity - 1) // 2
n_beta = (n_elec - system.multiplicity + 1) // 2
from .pbc_bipole_uhf import _combine_density_sets
D_total = _combine_density_sets(
basis, system, lattice_opts, result.density_alpha, result.density_beta
)
W_k_list = _build_energy_weighted_density_open(
result.mo_coeffs_alpha,
result.mo_energies_alpha,
result.mo_coeffs_beta,
result.mo_energies_beta,
n_alpha,
n_beta,
)
return _compute_bipole_gradient(
system,
basis,
D_total,
W_k_list,
n_elec,
lattice_opts=lattice_opts,
alpha_hf=alpha_hf,
ewald_alpha=getattr(result, "ewald_alpha_bohr_inv", None),
)
def compute_bipole_gradient_rks(
system: PeriodicSystem,
basis: BasisSet,
result: PBCBipoleRKSResult,
*,
lattice_opts: Optional[LatticeSumOptions] = None,
) -> np.ndarray:
"""BIPOLE RKS (DFT) atomic gradient. Uses alpha_hf from functional."""
if not result.converged:
import warnings
warnings.warn("compute_bipole_gradient_rks: not converged")
if lattice_opts is None:
lattice_opts = LatticeSumOptions()
n_elec = system.n_electrons()
n_occ = n_elec // 2
# Extract HF exchange fraction from functional name stored in result
alpha_hf = 0.0
func_name = getattr(result, "functional", "")
if func_name:
from ._vibeqc_core import Functional
try:
alpha_hf = float(Functional(func_name, 1).hf_exchange_fraction)
except Exception:
alpha_hf = 0.0
W_k_list = _build_energy_weighted_density_closed(
result.mo_coeffs, result.mo_energies, n_occ
)
return _compute_bipole_gradient(
system,
basis,
result.density,
W_k_list,
n_elec,
lattice_opts=lattice_opts,
alpha_hf=alpha_hf,
ewald_alpha=getattr(result, "ewald_alpha_bohr_inv", None),
)
def compute_bipole_gradient_uks(
system: PeriodicSystem,
basis: BasisSet,
result: PBCBipoleUKSResult,
*,
lattice_opts: Optional[LatticeSumOptions] = None,
) -> np.ndarray:
"""BIPOLE UKS (spin-DFT) atomic gradient."""
if not result.converged:
import warnings
warnings.warn("compute_bipole_gradient_uks: not converged")
if lattice_opts is None:
lattice_opts = LatticeSumOptions()
n_elec = system.n_electrons()
n_alpha = (n_elec + system.multiplicity - 1) // 2
n_beta = (n_elec - system.multiplicity + 1) // 2
from .pbc_bipole_uhf import _combine_density_sets
D_total = _combine_density_sets(
basis,
system,
lattice_opts,
result.density_alpha,
result.density_beta,
)
W_k_list = _build_energy_weighted_density_open(
result.mo_coeffs_alpha,
result.mo_energies_alpha,
result.mo_coeffs_beta,
result.mo_energies_beta,
n_alpha,
n_beta,
)
alpha_hf = 0.0
func_name = getattr(result, "functional", "")
if func_name:
from ._vibeqc_core import Functional
try:
alpha_hf = float(Functional(func_name, 2).hf_exchange_fraction)
except Exception:
alpha_hf = 0.0
return _compute_bipole_gradient(
system,
basis,
D_total,
W_k_list,
n_elec,
lattice_opts=lattice_opts,
alpha_hf=alpha_hf,
ewald_alpha=getattr(result, "ewald_alpha_bohr_inv", None),
)
def compute_bipole_gradient_fd(
system: PeriodicSystem,
basis_name: str,
kmesh,
options=None,
*,
step_bohr: float = 1e-3,
**bipole_kwargs,
) -> np.ndarray:
"""Central-difference BIPOLE RHF gradient via repeated SCF.
Runs 6N full BIPOLE SCF calculations at displaced geometries.
Exact but 6N× more expensive than the analytic path.
Parameters
----------
system : PeriodicSystem
Reference geometry.
basis_name : str
Basis set name (rebuilt per displaced geometry).
kmesh : BlochKMesh
k-point mesh.
options : PeriodicRHFOptions, optional
SCF options.
step_bohr : float
Half-step for central difference (default 1e-3 bohr).
**bipole_kwargs
Forwarded to ``run_pbc_bipole_rhf`` (ewald_precision, etc.).
Returns
-------
np.ndarray
``(n_atoms, 3)`` gradient in Ha/bohr.
"""
from copy import deepcopy
from ._vibeqc_core import Atom
from ._vibeqc_core import BasisSet as _BasisSet
from .pbc_bipole import run_pbc_bipole_rhf
n_atoms = len(system.unit_cell)
grad = np.zeros((n_atoms, 3), dtype=np.float64)
for a in range(n_atoms):
for cart in range(3):
e_plus = None
e_minus = None
for sign in (+1, -1):
delta = sign * step_bohr
new_atoms = []
for i, atom in enumerate(system.unit_cell):
xyz = list(atom.xyz)
if i == a:
xyz[cart] += delta
new_atoms.append(Atom(int(atom.Z), xyz))
sys_disp = PeriodicSystem(
system.dim,
np.asarray(system.lattice, dtype=np.float64),
new_atoms,
charge=system.charge,
multiplicity=system.multiplicity,
)
basis_disp = _BasisSet(
sys_disp.unit_cell_molecule(),
basis_name,
)
res = run_pbc_bipole_rhf(
sys_disp,
basis_disp,
kmesh,
options,
progress=False,
**bipole_kwargs,
)
if sign == +1:
e_plus = res.energy
else:
e_minus = res.energy
grad[a, cart] = (e_plus - e_minus) / (2.0 * step_bohr)
return grad