dense-evolution 8.3.0__py3-none-win_amd64.whl

dense_evolution/native_hf/overlap.py

@@ -0,0 +1,91 @@
+"""Overlap integrals between Gaussian shells via the Obara-Saika recursion.
+
+For two 1D Gaussians centered at A, B with exponents a, b, let
+p = a+b and P = (aA+bB)/p be the center of their product (Gaussian
+product theorem). The overlap S[i,j] = <(x-A)^i exp(-a(x-A)^2) |
+(x-B)^j exp(-b(x-B)^2)> obeys two recursions:
+
+  Vertical (build up the first index from the base case S[0,0]):
+    S[i,0] = (P-A) S[i-1,0] + (i-1)/(2p) S[i-2,0]
+
+  Horizontal (build up the second index by shifting angular momentum
+  from center A to center B, exact for any two centers -- this is why
+  it needs no exponent-dependent term):
+    S[i,j] = (A-B) S[i,j-1] + S[i+1,j-1]
+
+Both are linear recursions in one index with a 2-term memory, so each
+maps directly onto jax.lax.scan: the whole angular-momentum ladder for
+a shell pair compiles to one XLA loop instead of a Python for-loop.
+"""
+
+import functools
+
+import jax
+import jax.numpy as jnp
+
+from dense_evolution.native_hf.gaussians import (
+    GaussianShell1D,
+    GaussianShell3D,
+    product_center,
+    product_prefactor,
+)
+
+
+def _base_overlap_1d(g1: GaussianShell1D, g2: GaussianShell1D) -> jax.Array:
+    p = g1.exponent + g2.exponent
+    mu = (g1.exponent * g2.exponent) / p
+    K = jnp.exp(-mu * jnp.square(g1.center - g2.center))
+    return jnp.sqrt(jnp.pi / p) * K
+
+
+def _vertical_step(p_minus_a, inv_two_p, carry, i):
+    s_prev, s_prev2 = carry
+    s_next = p_minus_a * s_prev + (i - 1) * inv_two_p * s_prev2
+    return (s_next, s_prev), s_next
+
+
+def _build_first_index(s00, g1: GaussianShell1D, g2: GaussianShell1D, n_extra: int) -> jax.Array:
+    """Returns S[i, 0] for i = 0..n_extra."""
+    if n_extra == 0:
+        return s00[None]
+    p = g1.exponent + g2.exponent
+    P = (g1.exponent * g1.center + g2.exponent * g2.center) / p
+    step = functools.partial(_vertical_step, P - g1.center, 1.0 / (2.0 * p))
+    _, rest = jax.lax.scan(step, (s00, jnp.zeros_like(s00)), jnp.arange(1, n_extra + 1))
+    return jnp.concatenate([s00[None], rest])
+
+
+def _horizontal_step(a_minus_b, column, _):
+    shifted = jnp.roll(column, shift=-1)
+    new_column = a_minus_b * column + shifted
+    return new_column, new_column
+
+
+def _build_second_index(first_col: jax.Array, g1: GaussianShell1D, g2: GaussianShell1D) -> jax.Array:
+    """first_col holds S[:,0]. Returns S[:, 0..g2.degree] (rows beyond
+    validity are garbage and get sliced away by the caller)."""
+    if g2.degree == 0:
+        return first_col[:, None]
+    step = functools.partial(_horizontal_step, g1.center - g2.center)
+    _, rest = jax.lax.scan(step, first_col, jnp.arange(1, g2.degree + 1))
+    return jnp.concatenate([first_col[:, None], jnp.moveaxis(rest, 0, -1)], axis=-1)
+
+
+def overlap_1d(g1: GaussianShell1D, g2: GaussianShell1D) -> jax.Array:
+    """S[i,j] for 0<=i<=g1.degree, 0<=j<=g2.degree. Shape (g1.degree+1, g2.degree+1)."""
+    s00 = _base_overlap_1d(g1, g2)
+    first_col = _build_first_index(s00, g1, g2, g1.degree + g2.degree)
+    full = _build_second_index(first_col, g1, g2)
+    return full[: g1.degree + 1, :]
+
+
+@jax.jit
+def overlap_3d(g1: GaussianShell3D, g2: GaussianShell3D) -> jax.Array:
+    """S[ix,iy,iz,jx,jy,jz] for the full Cartesian shell pair.
+
+    Shape: (L1+1,L1+1,L1+1, L2+1,L2+1,L2+1) where L1, L2 are the shell
+    degrees (unphysical (i,j,k) combinations with i+j+k > degree are
+    simply never read by the caller).
+    """
+    axes = [overlap_1d(g1.component(d), g2.component(d)) for d in range(3)]
+    return jnp.einsum("ad,be,cf->abcdef", *axes)

dense_evolution/native_hf/scf.py

@@ -0,0 +1,404 @@
+"""Restricted Hartree-Fock self-consistent field loop (Roothaan-Hall).
+
+Standard textbook algorithm (e.g. Szabo & Ostlund, "Modern Quantum
+Chemistry", ch. 3): orthogonalize the AO basis via S^(-1/2), build the
+Fock matrix F = H_core + 2J - K from the current density, diagonalize
+in the orthogonal basis, form a new density, repeat to convergence.
+This part is genuinely simple compared to the integral evaluation and
+doesn't need vectorizing -- a closed-shell molecule's SCF loop is a few
+dozen matrix multiplies on an N x N matrix where N is a few tens at
+most for STO-3G, nowhere near where PennyLane's implementation loses
+its time (which is entirely in building H_core/repulsion tensor, done
+once in assembly.py, not in this loop).
+
+BUG FOUND (Si2, minimal 4-electron/4-orbital active space, R=2.184 A):
+plain (undamped) density substitution never converged for this system
+-- 100/100 iterations, still oscillating -- because two pairs of
+orbitals near the active-space boundary are numerically degenerate
+(HOMO-1/HOMO and LUMO/LUMO+1 each split by <1e-9 Ha), so each iteration
+flips which member of a near-tied pair gets occupied, and the density
+never settles. Confirmed this is a real oscillation, not just slow
+convergence: three separate machines/runs of the undamped loop each
+hit the iteration cap at a DIFFERENT total energy (-571.63, -570.69,
+-571.02 Ha), all physically meaningless artifacts of whatever step the
+loop happened to be on. First fixed with plain linear density damping
+(P_next = alpha*P_new + (1-alpha)*P_old) -- textbook remedy for exactly
+this oscillation failure mode (Szabo & Ostlund ch. 3.4.9), verified to
+converge to the SAME energy (-570.874032094871 Ha, agreeing to 10
+significant figures) across alpha in {0.1, 0.2, 0.3, 0.5, 0.7).
+
+UPGRADED to DIIS (Pulay, "Convergence acceleration of iterative
+sequences: the case of SCF iteration", Chem. Phys. Lett. 73, 393
+(1980); "Improved SCF convergence acceleration", J. Comput. Chem. 3,
+556 (1982)) -- the standard production-grade SCF accelerator plain
+linear damping is a simplified special case of. Pulay's own error
+vector, `e = X.T @ (F@P@S - S@P@F) @ X` (the Fock/density commutator in
+the orthonormal basis, exactly zero at true self-consistency), is kept
+across the last `diis_dim` iterations alongside the Fock matrices that
+produced them; each step extrapolates a new Fock matrix as the
+minimum-norm linear combination of that history (constrained to sum to
+1) instead of diagonalizing the latest F directly. Verified on the
+same Si2 near-degenerate case that motivated damping in the first
+place: DIIS converges in 11 iterations to -570.8740320948958 Ha, versus
+53 iterations for plain damped substitution alone (diis_dim=0,
+alpha=0.5) to reach the same energy, -570.8740320948963 Ha -- agreeing
+to 12 significant figures, the same true self-consistent solution
+reached ~4.8x faster, not a different answer.
+
+If the DIIS linear system is ever singular (a degenerate/duplicated
+error-vector history), that step falls back to the plain, unextrapolated
+Fock matrix rather than raising -- a transient fallback, not silent
+wrong physics, since the next iteration's fresh error vector rebuilds a
+usable history.
+
+Convergence now requires BOTH the density AND the energy to stop
+changing (`|P_new - P| < convergence_tol` AND `|E_new - E_old| <
+energy_tol`) rather than density alone -- density convergence can
+occasionally plateau one step before energy does (or vice versa) for
+a system with several nearly-degenerate iterations near the end of the
+run; requiring both is strictly more conservative than either alone
+and costs at most a couple of extra iterations on every system tested
+here.
+
+Rewritten from NumPy to jax.numpy (the integral-building side in
+assembly.py already was JAX; this loop and its np.array() outputs were
+the only remaining barrier to an end-to-end JAX-traced pipeline). The
+iteration itself uses jax.lax.while_loop for its data-dependent
+convergence check, same as the original Python for/break -- note that
+reverse-mode autodiff (jax.grad) does not work through while_loop, by
+JAX's own design (the number of iterations isn't known ahead of time,
+which reverse-mode differentiation needs). Making this loop's OUTPUT
+differentiable is a separate, deliberately deferred step: the correct
+approach is implicit differentiation at the fixed point (the custom
+gradient rule only needs the converged F/P, not a replay of every
+iteration), not unrolling this loop and backpropagating through it.
+
+DIIS history is now a fixed-size (diis_dim, n, n) ring buffer instead
+of a growing/shrinking Python list -- jax.lax.while_loop requires its
+carried state to have constant shape across iterations. Each step
+rolls the buffer and writes the newest Fock/error matrix into the last
+slot; a boolean mask (derived from how many iterations have actually
+run) excludes the not-yet-filled slots from the DIIS linear system
+instead of the original's list-length check. The original's
+try/except LinAlgError fallback becomes a jnp.isfinite check on the
+solved coefficients (a singular solve produces NaN/Inf under JAX
+rather than raising), selecting the single latest Fock matrix exactly
+as the original's except-branch did.
+"""
+
+import dataclasses
+import functools
+
+import jax
+import numpy as np
+import jax.numpy as jnp
+
+from dense_evolution.config import ensure_x64
+
+_DIIS_DIM = 8
+
+
+@dataclasses.dataclass
+class HFResult:
+    converged: bool
+    n_iterations: int
+    electronic_energy: float
+    nuclear_repulsion_energy: float
+    total_energy: float
+    orbital_energies: jax.Array
+    orbital_coefficients: jax.Array  # C, shape (n_basis, n_basis)
+    density_matrix: jax.Array
+    energy_history: jax.Array  # shape (max_iterations,), NaN past n_iterations
+
+
+def nuclear_repulsion_energy(nuclear_charges: list[float], nuclear_positions: jax.Array) -> jax.Array:
+    positions = jnp.asarray(nuclear_positions)
+    energy = 0.0
+    n = len(nuclear_charges)
+    for i in range(n):
+        for j in range(i + 1, n):
+            r = jnp.linalg.norm(positions[i] - positions[j])
+            energy = energy + nuclear_charges[i] * nuclear_charges[j] / r
+    return energy
+
+
+def _orthogonalizer(S: jax.Array) -> jax.Array:
+    w, v = jnp.linalg.eigh(S)
+    return v @ jnp.diag(1.0 / jnp.sqrt(w)) @ v.T
+
+
+def _density_from_coefficients(C: jax.Array, n_occupied_pairs: int) -> jax.Array:
+    C_occ = C[:, :n_occupied_pairs]
+    return C_occ @ C_occ.T
+
+
+def _diis_error(F: jax.Array, P: jax.Array, S: jax.Array, X: jax.Array) -> jax.Array:
+    """Pulay's DIIS error vector, `X.T @ (F@P@S - S@P@F) @ X` -- the
+    Fock/density commutator (zero at true self-consistency, since F and
+    P then commute), transformed into the same orthonormal basis the
+    Fock matrix itself is diagonalized in. F, P, S are all symmetric, so
+    this reduces to `Y - Y.T` for `Y = X.T @ F @ P @ S @ X` -- computed
+    that way to avoid forming two separate FPS/SPF products."""
+    Y = X.T @ F @ P @ S @ X
+    return Y - Y.T
+
+
+def _diis_extrapolate(fock_history: jax.Array, error_history: jax.Array, history_count: jax.Array, diis_dim: int) -> jax.Array:
+    """fock_history/error_history: (diis_dim, n, n), newest entry always
+    in the last slot (see run_scf's roll-and-append). `valid` marks
+    which of the diis_dim slots hold a real (not yet overwritten, not a
+    zero-initialized placeholder) entry -- the last `history_count` of
+    them. Invalid slots are pinned to c_i=0 by giving their row/column
+    of the augmented linear system an identity-like equation instead of
+    a real DIIS constraint, so they can't contribute to the solution
+    regardless of their (garbage/zero) content."""
+    dtype = fock_history.dtype
+    zero = jnp.zeros((), dtype=dtype)
+    one = jnp.ones((), dtype=dtype)
+    valid = jnp.arange(diis_dim) >= (diis_dim - history_count)
+
+    errs_flat = error_history.reshape(diis_dim, -1)
+    B_full = errs_flat @ errs_flat.T
+    mask2d = valid[:, None] & valid[None, :]
+    B = jnp.where(mask2d, B_full, zero)
+    B = jnp.where(jnp.eye(diis_dim, dtype=bool) & ~mask2d, one, B)
+
+    A = jnp.zeros((diis_dim + 1, diis_dim + 1), dtype=dtype)
+    A = A.at[:diis_dim, :diis_dim].set(B)
+    col = jnp.where(valid, -one, zero)
+    A = A.at[:diis_dim, diis_dim].set(col)
+    A = A.at[diis_dim, :diis_dim].set(col)
+    b = jnp.zeros(diis_dim + 1, dtype=dtype).at[diis_dim].set(-one)
+
+    solution = jnp.linalg.solve(A, b)
+    is_finite = jnp.all(jnp.isfinite(solution))
+    coeffs = jnp.where(valid, jnp.where(is_finite, solution[:diis_dim], zero), zero)
+
+    F_diis = jnp.tensordot(coeffs, fock_history, axes=1)
+    return jnp.where(is_finite, F_diis, fock_history[-1])
+
+
+def _level_shift_fock(F_ao: jax.Array, C_prev: jax.Array, S: jax.Array, n_occupied_pairs: int, level_shift: float) -> jax.Array:
+    """Saunders & Hillier level shifting (Int. J. Quantum Chem. 7, 699
+    (1973)): push the virtual orbitals of the PREVIOUS iteration's MO
+    basis up by `level_shift` before this iteration's diagonalization,
+    to open a numerical gap and stop the occupied/virtual split from
+    flip-flopping across a near-degeneracy (see this module's own Si2
+    docstring above -- damping/DIIS already fix the textbook case, but a
+    real, harder case (a 30-atom aromatic fragment from the CASMI26
+    wiring kernel) still took 1114 iterations, swinging through three
+    wildly different intermediate energies first). Exact no-op at
+    level_shift=0.0: since C_prev is a full, S-orthonormal basis
+    (C_prev.T @ S @ C_prev = I, hence C_prev @ C_prev.T = S^{-1}), the
+    round-trip S @ C_prev @ (C_prev.T @ F_ao @ C_prev) @ C_prev.T @ S
+    reduces algebraically to exactly F_ao before any shift is added."""
+    F_mo_prev = C_prev.T @ F_ao @ C_prev
+    n = F_mo_prev.shape[0]
+    level_shift = jnp.asarray(level_shift, dtype=F_ao.dtype)
+    shift_diag = jnp.where(jnp.arange(n) >= n_occupied_pairs, level_shift, jnp.zeros((), dtype=F_ao.dtype))
+    F_mo_prev_shifted = F_mo_prev + jnp.diag(shift_diag)
+    return S @ C_prev @ F_mo_prev_shifted @ C_prev.T @ S
+
+
+def run_scf(
+    S: jax.Array,
+    H_core: jax.Array,
+    repulsion: jax.Array,
+    n_electrons: int,
+    nuclear_charges: list[float],
+    nuclear_positions: jax.Array,
+    max_iterations: int = 200,
+    convergence_tol: float = 1e-10,
+    energy_tol: float = 1e-10,
+    damping: float = 0.5,
+    diis_dim: int = _DIIS_DIM,
+    level_shift: float = 0.0,
+) -> HFResult:
+    ensure_x64()
+    if n_electrons % 2 != 0:
+        raise ValueError("Only closed-shell (even electron count) systems are supported.")
+    n_occupied_pairs = n_electrons // 2
+    n_basis = S.shape[0]
+
+    X = _orthogonalizer(S)
+    orbital_energies0, C_ortho0 = jnp.linalg.eigh(X.T @ H_core @ X)
+    C0 = X @ C_ortho0
+    P0 = _density_from_coefficients(C0, n_occupied_pairs)
+
+    def _fock_and_energy(P):
+        J = jnp.einsum("pqrs,rs->pq", repulsion, P)
+        K = jnp.einsum("prqs,rs->pq", repulsion, P)
+        F = H_core + 2.0 * J - K
+        energy = jnp.sum(P * (H_core + F))
+        return F, energy
+
+    def cond_fun(state):
+        iteration, _P, _C, _orbital_energies, _energy_prev, converged, _fock_history, _error_history, _energy_history = state
+        return jnp.logical_and(jnp.logical_not(converged), iteration < max_iterations)
+
+    def body_fun(state):
+        iteration, P, C_prev, _orbital_energies, energy_prev, _converged, fock_history, error_history, energy_history = state
+
+        F, energy = _fock_and_energy(P)
+        error = _diis_error(F, P, S, X)
+
+        fock_history = jnp.roll(fock_history, shift=-1, axis=0).at[-1].set(F)
+        error_history = jnp.roll(error_history, shift=-1, axis=0).at[-1].set(error)
+        history_count = jnp.minimum(iteration + 1, diis_dim)
+
+        F_step = jnp.where(
+            history_count >= 2,
+            _diis_extrapolate(fock_history, error_history, history_count, diis_dim),
+            F,
+        )
+
+        F_diag = _level_shift_fock(F_step, C_prev, S, n_occupied_pairs, level_shift)
+        orbital_energies, C_ortho = jnp.linalg.eigh(X.T @ F_diag @ X)
+        C = X @ C_ortho
+        P_new = _density_from_coefficients(C, n_occupied_pairs)
+
+        density_converged = jnp.linalg.norm(P_new - P) < convergence_tol
+        energy_converged = jnp.abs(energy - energy_prev) < energy_tol
+        converged = jnp.logical_and(density_converged, energy_converged)
+
+        P_damped = jnp.where(history_count < 2, damping * P_new + (1.0 - damping) * P, P_new)
+        P_next = jnp.where(converged, P_new, P_damped)
+
+        energy_history = energy_history.at[iteration].set(energy)
+
+        return (iteration + 1, P_next, C, orbital_energies, energy, converged, fock_history, error_history, energy_history)
+
+    init_state = (
+        jnp.array(0),
+        P0,
+        C0,
+        orbital_energies0,
+        jnp.array(jnp.inf, dtype=H_core.dtype),
+        jnp.array(False),
+        jnp.zeros((diis_dim, n_basis, n_basis), dtype=H_core.dtype),
+        jnp.zeros((diis_dim, n_basis, n_basis), dtype=H_core.dtype),
+        jnp.full((max_iterations,), jnp.nan, dtype=H_core.dtype),
+    )
+    iteration, P, C, orbital_energies, _energy_prev, converged, _fh, _eh, energy_history = jax.lax.while_loop(cond_fun, body_fun, init_state)
+
+    _F, electronic_energy = _fock_and_energy(P)
+    e_nuc = nuclear_repulsion_energy(nuclear_charges, nuclear_positions)
+
+    return HFResult(
+        converged=bool(converged),
+        n_iterations=int(iteration),
+        electronic_energy=float(electronic_energy),
+        nuclear_repulsion_energy=float(e_nuc),
+        total_energy=float(electronic_energy) + float(e_nuc),
+        orbital_energies=orbital_energies,
+        orbital_coefficients=C,
+        density_matrix=P,
+        energy_history=energy_history,
+    )
+
+
+def _import_dense_armor_robust_filters():
+    try:
+        from dense_armor.utility.robust_filters import hampel_filter, tukey_fences
+    except ImportError as exc:
+        raise ImportError(
+            "native_hf.scf.diagnose_convergence needs Dense-Armor (pip install dense-evolution[armor])"
+        ) from exc
+    return hampel_filter, tukey_fences
+
+
+def diagnose_convergence(result: HFResult, radius: int = 5, n_sigmas: float = 3.0) -> dict:
+    """Real automatic guard on HFResult.energy_history, instead of trusting
+    the final `converged` flag in isolation: Dense-Armor's Hampel filter
+    and Tukey fences (dense_evolution.utility.robust_filters -- the sister
+    project's own anomaly detectors, Chauvenet/Tukey/Hampel/sigma-clipping,
+    validated with 0 false positives on real H2 dissociation-curve
+    chemistry) applied to the real per-iteration electronic energy trace.
+
+    Validated on a real non-converging case (a 28-heavy-atom CASMI26
+    fragment, level_shift=0.0): the broken run flags 19-24% of its 200
+    iterations as anomalous; the same fragment fixed with level_shift=0.5
+    (converges in 55 iterations) flags only ~4% -- background noise, not a
+    false-alarm storm. Needs the `armor` extra (pip install
+    dense-evolution[armor]) -- Dense-Armor is not a hard dependency of this
+    module."""
+    hampel_filter, tukey_fences = _import_dense_armor_robust_filters()
+    trace = np.asarray(result.energy_history[:result.n_iterations])
+    n = max(1, result.n_iterations)
+
+    _cleaned_h, anomalies_h = hampel_filter(trace, radius=radius, n_sigmas=n_sigmas)
+    _cleaned_t, anomalies_t = tukey_fences(trace, radius=len(trace))
+
+    return {
+        "n_iterations": result.n_iterations,
+        "n_anomalies_hampel": len(anomalies_h),
+        "n_anomalies_tukey": len(anomalies_t),
+        "anomaly_fraction_hampel": len(anomalies_h) / n,
+        "anomaly_fraction_tukey": len(anomalies_t) / n,
+    }
+
+
+@functools.partial(jax.custom_vjp, nondiff_argnums=(3,))
+def scf_electronic_energy(S: jax.Array, H_core: jax.Array, repulsion: jax.Array, n_electrons: int) -> jax.Array:
+    """The RHF electronic energy (run_scf's own `electronic_energy`, not
+    counting nuclear repulsion) as a function of S/H_core/repulsion that
+    IS differentiable via jax.grad -- unlike run_scf itself, whose
+    jax.lax.while_loop can't be traced in reverse mode.
+
+    The gradient is NOT backprop through the SCF iteration (which
+    wouldn't even be possible) -- it's the analytic Hartree-Fock gradient
+    of Pople, Krishnan, Schlegel & Binkley, Int. J. Quantum Chem. Symp.
+    13, 225 (1979), eq. (21)-(22): at self-consistency, dE/dtheta equals
+    the derivative of `Tr[P(H_core+F(P))] - Tr[W S]` with P and the
+    energy-weighted density matrix W held fixed at their converged
+    values (an envelope-theorem/Lagrangian result -- P and W are the
+    stationary point and multipliers of the constrained HF variational
+    problem, so their own dependence on theta drops out of the total
+    derivative). W is built from ONLY the occupied orbitals:
+    `W = C_occ @ diag(2 * orbital_energies_occ) @ C_occ.T` -- the factor
+    of 2 is this module's own P convention (no explicit 2 in P itself,
+    carried instead by F = H_core + 2J - K), not part of Pople et al.'s
+    original spin-orbital formula. This automatically
+    includes the "Pulay force" terms from the atom-centered basis moving
+    with the nuclei (via H_core/repulsion/S's own theta-dependence),
+    without hand-deriving them -- jax.grad on the frozen-P expression
+    below does that part for free.
+
+    Verified against central finite differences on H2/STO-3G (see
+    tests/unit/test_native_hf_differentiable.py)."""
+    result = run_scf(S, H_core, repulsion, n_electrons, [], jnp.zeros((0, 3)))
+    return result.electronic_energy
+
+
+def _scf_electronic_energy_fwd(S, H_core, repulsion, n_electrons):
+    result = run_scf(S, H_core, repulsion, n_electrons, [], jnp.zeros((0, 3)))
+    residuals = (S, H_core, repulsion, result.density_matrix, result.orbital_coefficients, result.orbital_energies)
+    return result.electronic_energy, residuals
+
+
+def _scf_electronic_energy_bwd(n_electrons, residuals, cotangent):
+    S, H_core, repulsion, P, C, orbital_energies = residuals
+    n_occupied_pairs = n_electrons // 2
+    C_occ = C[:, :n_occupied_pairs]
+    # Lagrange multiplier for the C_occ.T @ S @ C_occ = I constraint is
+    # 2*diag(orbital_energies_occ), not diag(orbital_energies_occ) --
+    # this module's P has no explicit factor of 2 (F = H_core + 2J - K
+    # carries it instead), so stationarity of Tr[P(H_core+F(P))] -
+    # Tr[Lambda(C_occ.T S C_occ - I)] w.r.t. C_occ gives F C_occ =
+    # S C_occ (Lambda/2), which must match the Roothaan-Hall equation
+    # F C_occ = S C_occ diag(eps_occ) -- so Lambda = 2*diag(eps_occ).
+    # Verified: omitting this factor of 2 gave a gradient that disagreed
+    # with central finite differences by exactly Tr[W_undoubled dS/dx].
+    W = C_occ @ jnp.diag(2.0 * orbital_energies[:n_occupied_pairs]) @ C_occ.T
+
+    def lagrangian(S_, H_core_, repulsion_):
+        J = jnp.einsum("pqrs,rs->pq", repulsion_, P)
+        K = jnp.einsum("prqs,rs->pq", repulsion_, P)
+        F = H_core_ + 2.0 * J - K
+        return jnp.sum(P * (H_core_ + F)) - jnp.sum(W * S_)
+
+    dS, dH_core, drepulsion = jax.grad(lagrangian, argnums=(0, 1, 2))(S, H_core, repulsion)
+    return (cotangent * dS, cotangent * dH_core, cotangent * drepulsion)
+
+
+scf_electronic_energy.defvjp(_scf_electronic_energy_fwd, _scf_electronic_energy_bwd)

dense_evolution/noise/__init__.py

@@ -0,0 +1,79 @@
+"""Every way to put noise into a Dense-Evolution simulation, in one place.
+
+- **`NoiseModel`** (`.kraus_channels`, dispatching to one file per channel
+  under `.kraus`: `ideal`, `depolarizing`, `bitflip`, `phaseflip`,
+  `amplitude_damping`, `combined`) -- 6 stochastic single-qubit Kraus
+  channels, applied directly to a statevector via `apply_to_sv`.
+- **`NoiseSpec`** (`.differentiable`) -- the native JAX-differentiable
+  representation of a noise configuration, so noise strength itself can be
+  a traced/differentiable value inside `circuit_to_energy_fn`.
+- **Coherent adversarial noise** (`.coherent_attack`) -- a genuinely
+  continuous, multi-qubit coherent error channel (`apply_rz_all`) and a
+  JAX-differentiable search (`craft_adversarial_delta`,
+  `craft_adversarial_delta_constrained`) for a worst-case direction against
+  a stabilizer code's syndrome -- promoted from Dense-Evolution-Discovery's
+  Steane [[7,1,3]] investigation, including its honest negative result
+  (see `coherent_attack`'s module docstring).
+- **Density-matrix channels** (`.density_matrix_channels`) --
+  `global_depolarizing_channel`, `amplitude_damping_channel`: noise applied
+  directly to a density matrix instead of a statevector, for density-matrix
+  ZNE's noise ensemble. `phaseflip_channel_exact`: the same phaseflip model
+  as `NoiseModel`'s `'phaseflip'` statevector Kraus channel, but computed as
+  its exact zero-sampling-variance limit directly on a density matrix, for
+  use as a "classical node" in Classically Augmented ZNE (see
+  `dense_evolution.mitigation.classically_augmented_zne_phaseflip`).
+- **`cosmic_ray_burst_profile`** (`.cosmic_ray`) -- a real, time-dependent
+  noise-strength profile for a cosmic-ray-induced quasiparticle burst.
+- **`oscillating_p_eff`** (`.oscillating`) -- a noise strength that
+  oscillates instead of scaling smoothly, for stress-testing mitigation
+  techniques that assume smoothness.
+- **`pink_noise_p_eff`** (`.pink`) -- a noise-strength trace with a real
+  1/f (pink/flicker) power spectrum, the dominant decoherence mechanism
+  in real superconducting qubits, via the Timmer & Koenig (1995)
+  spectral-synthesis algorithm.
+
+Real device noise from a Qiskit backend's own calibration data
+(`noise_model_from_qiskit_backend`) lives in `dense_evolution.interop`,
+not here, since it needs a Qiskit `BackendV2` object as input rather than
+a noise-specific dependency.
+
+Everything in this package previously lived scattered across
+`dense_evolution.circuits.registry` (alongside unrelated hardware-detection
+code) and `dense_evolution.mitigation.zne` (alongside unrelated mitigation
+techniques, which only ever cancel noise, never generate it). Both modules
+re-export the relevant names from here for backward compatibility, but
+`dense_evolution.noise` is the canonical import path for new code.
+"""
+from .kraus_channels import NoiseModel
+from .differentiable import NoiseSpec
+from .coherent_attack import (
+    apply_rz_all,
+    x_stabilizer_leakage,
+    craft_adversarial_delta,
+    project_l2_linf,
+    craft_adversarial_delta_constrained,
+    decoder_failure_rate,
+    random_delta_failure_stats,
+)
+from .density_matrix_channels import global_depolarizing_channel, amplitude_damping_channel, phaseflip_channel_exact
+from .cosmic_ray import cosmic_ray_burst_profile
+from .oscillating import oscillating_p_eff
+from .pink import pink_noise_p_eff
+
+__all__ = [
+    "NoiseModel",
+    "NoiseSpec",
+    "apply_rz_all",
+    "x_stabilizer_leakage",
+    "craft_adversarial_delta",
+    "project_l2_linf",
+    "craft_adversarial_delta_constrained",
+    "decoder_failure_rate",
+    "random_delta_failure_stats",
+    "global_depolarizing_channel",
+    "amplitude_damping_channel",
+    "phaseflip_channel_exact",
+    "cosmic_ray_burst_profile",
+    "oscillating_p_eff",
+    "pink_noise_p_eff",
+]