Source code for vibeqc.periodic_gradient

"""Phase G1a — analytic Γ-only periodic RHF / RKS atomic gradient.

Closed-shell extension of :func:`vibeqc.compute_gradient` (molecular)
to a periodic system. Assembles the four contributions to the per-
cell gradient:

  F_A = ∂E_nn^pc/∂R_A
      + Σ_g  tr(D(g) · ∂h_core(g)/∂R_A)              (1-e Pulay)
      - Σ_g  tr(W(g) · ∂S(g)/∂R_A)                   (overlap-Lagrangian)
      + Σ_g Σ_h Σ_h'  Γ(g, h, h') · ∂(μν0|λh σh')/∂R_A  (2-e Pulay)

**Validation status (G1a series).**

Molecular limit (cell list reduces to a single cell, ``cutoff_bohr <
inter-image distance``):

  - All four contributions match the molecular gradient bit-for-bit
    (≤ 1e-7 Ha/bohr); validated against
    :func:`vibeqc.compute_gradient` on H₂ in 20-Å cubic box.
  - 27× speedup vs the FD reference at this regime.

True periodic (cells overlap; cross-cell density / ERIs significant):

  - **Nuclear-rep + 1-e Pulay + overlap-Lagrangian**: matches FD to
    Newton's-3rd-law precision (1e-13). These three are the bulk of
    the gradient on systems with negligible α_HF·K (pure DFT — LDA,
    PBE, BLYP, …).
  - **2-e J piece** (exchange_scale = 0): matches FD to N3-precision
    (~1e-13). This is what pure-DFT periodic gradients use.
  - **2-e K piece** (α_HF ≠ 0; HF or hybrid DFT): currently shows a
    residual ~5e-3 Ha/bohr disagreement with FD on small periodic
    cells (1D H chain, a = 2 Å, STO-3G). The structural bug is
    isolated to the K-style integral derivative-buffer routing in
    ``eri_lattice_gradient_contribution``; molecular limit passes
    bit-for-bit. **Tracked as a v0.6.x patch.**

So today this module produces:

  - ✅ Exact analytic gradients for pure-DFT periodic SCF (LDA, PBE,
    BLYP, … with α_HF = 0).
  - ⚠️  Approximate gradients for HF or hybrid DFT periodic SCF —
    the J piece is exact but K has a known bug. Fall back to
    :func:`compute_gradient_periodic_rhf_fd` (slow but FD-correct)
    for production HF / hybrid DFT periodic geometry optimization
    until the K bug is closed.
"""

from __future__ import annotations

from typing import Optional

import numpy as np

from ._vibeqc_core import (
    BasisSet,
    LatticeMatrixSet,
    LatticeSumOptions,
    PeriodicRHFResult,
    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,
    two_electron_gradient_contribution,
)


__all__ = ["compute_gradient_periodic_rhf_gamma"]


def _fold_gamma_real(M_set: LatticeMatrixSet) -> np.ndarray:
    """Σ_g M(g) — the Γ-point Bloch fold (k = 0 → all cells equally
    weighted). For a single-cell mesh this returns the g=0 block;
    for a true multi-cell list, sums them."""
    nbf = int(M_set.nbf)
    out = np.zeros((nbf, nbf), dtype=np.float64)
    for blk in M_set.blocks:
        out += np.asarray(blk, dtype=np.float64)
    return out


def _gamma_density_lattice_set(
    template: LatticeMatrixSet, D: np.ndarray
) -> LatticeMatrixSet:
    """Build a Γ-only lattice-resolved density on ``template``'s cell
    list with the home-cell block set to ``D`` and image blocks zeroed.

    Convention. The dedicated Γ Ewald/RHF driver uses the direct-space
    molecular-limit convention: only the unit-cell density is known from
    a one-point mesh, so image-density blocks are zero. The analytic
    gradient must use the same density convention as the SCF energy and
    finite-difference reference.

    Multi-k generalisation (G1c): replace this with the proper
    ``real_space_density_from_kpoints`` call so D(g) carries its
    correct k-dependent Bloch phase.
    """
    D_arr = np.asarray(D, dtype=np.float64)
    zero = np.zeros_like(D_arr)
    for c, cell in enumerate(template.cells):
        idx = tuple(int(v) for v in np.asarray(cell.index).reshape(3))
        template.set_block(c, D_arr if idx == (0, 0, 0) else zero)
    return template


[docs] def compute_gradient_periodic_rhf_gamma( system: PeriodicSystem, basis: BasisSet, result: PeriodicRHFResult, *, lattice_opts: Optional[LatticeSumOptions] = None, alpha_hf: float = 1.0, ) -> np.ndarray: """Analytic Γ-only periodic RHF atomic gradient. Parameters ---------- system, basis Periodic system and AO basis. result Converged :class:`PeriodicRHFResult` from :func:`vibeqc.run_rhf_periodic` (must have ``converged=True``). lattice_opts :class:`LatticeSumOptions` controlling the lattice-sum cutoffs. If ``None``, defaults from :class:`LatticeSumOptions()`. For the gradient to match the SCF energy gradient, these cutoffs **must** match the values used during the SCF (i.e. the same ``opts.lattice_opts`` you passed to ``run_rhf_periodic``). Returns ------- np.ndarray ``(n_atoms, 3)`` gradient in Ha/bohr. Notes ----- **Current scope (G1a-1)**: 1-electron Pulay + nuclear-repulsion + overlap-Lagrangian terms via the new lattice-summed C++ primitives. The 2-electron Pulay term falls back to the molecular code path on the Γ-folded total density. This is exact in the molecular limit (single cell, AO overlap between cells negligible) and approximate for truly periodic systems where cross-cell ERIs contribute. **G1a-2** will replace the 2-e fallback with the full lattice-summed periodic ERI gradient. """ if not result.converged: raise ValueError( "compute_gradient_periodic_rhf_gamma: PeriodicRHFResult " "is not converged." ) if lattice_opts is None: lattice_opts = LatticeSumOptions() # Fetch the Γ-folded RHF reference data. D_gamma = np.asarray(result.density, dtype=np.float64) C = np.asarray(result.mo_coeffs, dtype=np.float64) eps = np.asarray(result.mo_energies, dtype=np.float64) F_result = np.asarray(getattr(result, "fock", np.empty((0, 0))), dtype=np.float64) n_elec = system.n_electrons() if n_elec % 2 != 0: raise ValueError( "compute_gradient_periodic_rhf_gamma: open-shell RHF not " "supported (closed-shell only). Use the periodic UHF " "gradient driver (G1d, post-G1a) for open-shell systems." ) nocc = n_elec // 2 # Energy-weighted density: W = 2 Σ_i ε_i C_μi C_νi. Some periodic # result paths historically stored pre-Fock orbital energies, so when # the converged Fock is available we recover ε_i = C_i^T F C_i. if F_result.shape == (C.shape[0], C.shape[0]): eps_for_w = np.einsum("ui,uv,vi->i", C, F_result, C, optimize=True) else: eps_for_w = eps W_gamma = 2.0 * (C[:, :nocc] * eps_for_w[:nocc][None, :]) @ C[:, :nocc].T # Wrap D_gamma and W_gamma into LatticeMatrixSets for the C++ side. # The lattice cell list comes from a fresh compute_overlap_lattice # call so the cell layout matches what the SCF used internally. D_set = compute_overlap_lattice(basis, system, lattice_opts) W_set = compute_overlap_lattice(basis, system, lattice_opts) _gamma_density_lattice_set(D_set, D_gamma) _gamma_density_lattice_set(W_set, W_gamma) # Sum the four contributions. grad = np.zeros((len(system.unit_cell), 3), dtype=np.float64) 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_set, lattice_opts)) grad += np.asarray( nuclear_lattice_gradient_contribution( basis, system, D_set, lattice_opts)) home_cell_only = all( tuple(int(v) for v in np.asarray(cell.index).reshape(3)) == (0, 0, 0) for cell in D_set.cells ) if home_cell_only: grad += np.asarray( two_electron_gradient_contribution( basis, system.unit_cell_molecule(), D_gamma, float(alpha_hf), ) ) else: # 2-electron Pulay — full lattice-summed periodic ERI gradient. # alpha_hf=1 → plain HF (J − ½ K). alpha_hf=0 → pure DFT (J only). # Hybrids inherit the K-piece bug (G1a-2 patch). grad += np.asarray( eri_lattice_gradient_contribution( basis, system, D_set, lattice_opts, float(alpha_hf))) return grad