fix(dm): make tol_grad a sufficient stop for the CG path

The CG loop required energy-plateau AND grad<=tol_grad (inherited from
jax_hf), so soft-mode runs sat below tol_grad for hundreds of iterations
and hit max_iter with converged=False.  With tol_grad > 0 the gradient
test is now the sole stopping criterion, mirroring solve_rtr: tested at
the freshly evaluated point before stepping, so the returned state is
the one that passed and the histories end with the qualifying (E, grad)
pair (this also removes the one-iteration staleness of the old top-of-
loop test and the max_iter off-by-one).  tol_grad == 0 keeps the
windowed-energy criterion unchanged.

Also documents the optimizer-internal history semantics — hist_E is the
pre-step HF energy without the -T*S entropy term; the CG finaliser takes
one extra diagonalise-and-reoccupy step so energy != hist_E[-1] for runs
stopped away from stationarity, while Newton finalises in place — and
adds TestGradientStop regression tests (note: a kernel without band
hybridization has exactly zero orbital gradient, so the tests seed off
the h-eigenbasis).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: skilledwolf

cpp/include/cpp_hf/solver_dm.hpp

@@ -22,6 +22,13 @@ namespace cpp_hf {
 struct SolverConfig {
     std::size_t max_iter = 200;
     f32 tol_E = 1.0e-7;
+    // Gradient stopping criterion.  When tol_grad > 0 it is the SOLE and
+    // SUFFICIENT stop for both optimizers: the run ends (converged = true) as
+    // soon as the w-normalised Frobenius norm of the orbital gradient — the
+    // exact quantity logged in hist_grad — drops to tol_grad, and the energy
+    // criterion below is inactive.  tol_grad == 0 (default) stops on the
+    // windowed-energy criterion instead (CG path only; the Newton path then
+    // falls back to an internal 1e-6 gradient tolerance).
     f32 tol_grad = 0.0;
     // Energy convergence is judged over a sliding window: stop when the energy
     // has improved by less than tol_E over the last plateau_window iterations.
@@ -71,6 +78,15 @@ struct DMResult {
     f32 energy = 0.0;
     std::size_t n_iter = 0;
     bool converged = false;
+    // Optimizer-internal per-iteration logs (one entry per outer iteration):
+    // hist_E is the HF total energy E[P] (w2d-weighted, refP-referenced, no
+    // -T*S entropy term) at each iteration's pre-step density; hist_grad is
+    // the w-normalised Frobenius norm of the orbital gradient at the same
+    // point — the exact quantity the tol_grad stop tests.  NOTE `energy` is
+    // recomputed after finalisation, which for the CG path takes one extra
+    // diagonalise-and-reoccupy step — so energy != hist_E.back() for a CG run
+    // that stopped away from stationarity (the Newton path finalises in place
+    // and satisfies energy == hist_E.back() exactly).
     std::vector<f32> hist_E;
     std::vector<f32> hist_grad;
 };
@@ -470,19 +486,25 @@ inline DMResult solve_dm(const HFKernel& K, const c64* P0, f32 n_e,
     f32 E = 0.0;
     f32 mu = out.mu;
     std::size_t k = 0;
+    bool converged = false;
 
     while (k < cfg.max_iter) {
-        bool energy_not_converged;
-        if (cfg.plateau_window > 0 && k > cfg.plateau_window) {
-            // Improvement over the last plateau_window recorded energies.
-            const f32 e_now = out.hist_E[k - 1];
-            const f32 e_past = out.hist_E[k - 1 - cfg.plateau_window];
-            energy_not_converged = (e_past - e_now) > cfg.tol_E;
-        } else {
-            energy_not_converged = (dE > cfg.tol_E);
+        // Energy-plateau stop — governs only when the gradient criterion is
+        // disabled (tol_grad == 0).  With tol_grad > 0 the gradient test
+        // below, taken at the freshly evaluated point BEFORE stepping, is the
+        // sole stopping criterion (mirrors solve_rtr).
+        if (cfg.tol_grad <= 0.0) {
+            bool energy_not_converged;
+            if (cfg.plateau_window > 0 && k > cfg.plateau_window) {
+                // Improvement over the last plateau_window recorded energies.
+                const f32 e_now = out.hist_E[k - 1];
+                const f32 e_past = out.hist_E[k - 1 - cfg.plateau_window];
+                energy_not_converged = (e_past - e_now) > cfg.tol_E;
+            } else {
+                energy_not_converged = (dE > cfg.tol_E);
+            }
+            if (!energy_not_converged) { converged = true; break; }
         }
-        bool grad_not_converged = (cfg.tol_grad > 0.0) && (grad_norm > cfg.tol_grad);
-        if (!(energy_not_converged || grad_not_converged)) break;
 
         // 1. Build Fock at projected hermitised density
         density_from_Qp(Q.data(), p.data(), P_cur.data(), nk, nb);
@@ -510,6 +532,17 @@ inline DMResult solve_dm(const HFKernel& K, const c64* P0, f32 n_e,
         }
         grad_norm = norm_matrix(G_Q.data(), w_norm.data(), nk, nb);
 
+        // Gradient stop (tol_grad > 0): sufficient criterion, tested at the
+        // point just evaluated so the returned state is the one that passed
+        // and the histories end with the qualifying (E, grad) pair.
+        if (cfg.tol_grad > 0.0 && grad_norm <= cfg.tol_grad) {
+            out.hist_E[k] = E;
+            out.hist_grad[k] = grad_norm;
+            ++k;
+            converged = true;
+            break;
+        }
+
         // 3. Precondition
         double eps_sq_sum = 0.0;
         for (std::size_t i = 0; i < nk * nb; ++i)
@@ -646,7 +679,7 @@ inline DMResult solve_dm(const HFKernel& K, const c64* P0, f32 n_e,
     out.mu = mu;
     out.energy = E_fin;
     out.n_iter = k;
-    out.converged = (k < cfg.max_iter);
+    out.converged = converged;
     out.hist_E.resize(k);
     out.hist_grad.resize(k);
     return out;

src/cpp_hf/solver.py

@@ -24,6 +24,14 @@ class SolverConfig:
     whose Fock matrices remain block diagonal in the original orbital basis.
     The solver falls back to the full dense path if a generated Fock matrix has
     off-block entries above numerical tolerance.
+
+    Stopping criteria: when ``tol_grad > 0`` it is the sole and *sufficient*
+    stop for both optimizers — the run ends (``converged=True``) as soon as
+    the w-normalised Frobenius norm of the orbital gradient (the quantity
+    logged in ``history["grad_norm"]``) drops to ``tol_grad``.  With
+    ``tol_grad == 0`` (default) the CG path stops on the windowed-energy
+    criterion (``tol_E`` / ``plateau_window``) and the Newton path falls back
+    to an internal 1e-6 gradient tolerance.
     """
 
     max_iter: int = 200
@@ -65,6 +73,21 @@ class SolverConfig:
 
 
 class SolveResult(NamedTuple):
+    """Result of :func:`solve_direct_minimization`.
+
+    ``history`` holds optimizer-internal per-iteration logs: ``"E"`` is the
+    HF total energy E[P] (w2d-weighted, refP-referenced, **without** the
+    -T*S entropy term of the free energy the optimizer minimises) at each
+    iteration's pre-step density, and ``"grad_norm"`` is the w-normalised
+    Frobenius norm of the orbital gradient at the same point — exactly the
+    quantity the ``tol_grad`` stop tests.  ``energy`` is recomputed after
+    finalisation: the CG path finalises with one extra diagonalise-and-
+    reoccupy step, so ``energy != history["E"][-1]`` for a CG run that
+    stopped away from stationarity; the Newton path finalises in place and
+    satisfies ``energy == history["E"][-1]`` exactly.  Do not compare
+    ``history["E"]`` tails across optimizers — compare ``energy``.
+    """
+
     Q: np.ndarray | None
     p: np.ndarray
     mu: np.ndarray

tests/test_solver.py

@@ -111,6 +111,61 @@ def test_4x4_grid_converges(self):
         assert bool(result.converged)
 
 
+class TestGradientStop:
+    """tol_grad > 0 is a sufficient stop for the CG path (mirrors newton).
+
+    Needs a kernel with band hybridization: without off-diagonal h the
+    orbital gradient is exactly zero from iteration 0 (all relaxation lives
+    in the occupation channel, which tol_grad intentionally excludes).
+    """
+
+    @staticmethod
+    def _hybridized_kernel(nk=2, T=0.1):
+        kernel_args = dict(
+            weights=np.ones((nk, nk), dtype=np.float32),
+            coulomb_q=np.full((nk, nk, 1, 1), 0.5, dtype=np.complex64),
+            T=T,
+        )
+        h = np.zeros((nk, nk, 2, 2), dtype=np.complex64)
+        h[..., 0, 0] = -0.5
+        h[..., 1, 1] = 0.5
+        h[..., 0, 1] = h[..., 1, 0] = 0.2
+        return HartreeFockKernel(hamiltonian=h, **kernel_args)
+
+    def _solve(self, config):
+        kernel = self._hybridized_kernel()
+        P0 = np.zeros_like(kernel.h)
+        # Seed off the h-eigenbasis so the orbital gradient starts nonzero.
+        P0[..., 0, 0] = 0.6
+        P0[..., 1, 1] = 0.15
+        P0[..., 0, 1] = P0[..., 1, 0] = 0.25
+        return solve(kernel, P0, 3.0, config=config)
+
+    def test_cg_stops_at_first_crossing_and_flags_converged(self):
+        tol = 1e-4
+        result = self._solve(SolverConfig(max_iter=300, tol_grad=tol))
+        assert bool(result.converged)
+        assert 1 < int(result.n_iter) < 300
+        hG = np.asarray(result.history["grad_norm"])
+        assert len(hG) == int(result.n_iter)
+        # Histories end with the qualifying point, and it is the first
+        # crossing — the run must not grind on after reaching tolerance.
+        assert hG[-1] <= tol
+        assert np.all(hG[:-1] > tol)
+
+    def test_cg_gradient_stop_matches_energy_stop_solution(self):
+        by_grad = self._solve(SolverConfig(max_iter=300, tol_grad=1e-7))
+        by_energy = self._solve(SolverConfig(max_iter=300, tol_E=1e-10))
+        assert bool(by_grad.converged) and bool(by_energy.converged)
+        np.testing.assert_allclose(
+            float(by_grad.energy), float(by_energy.energy), atol=1e-6)
+
+    def test_cg_unreachable_tolerance_flags_unconverged(self):
+        result = self._solve(SolverConfig(max_iter=3, tol_grad=1e-14))
+        assert not bool(result.converged)
+        assert int(result.n_iter) == 3
+
+
 class TestSolveResult:
     def test_result_is_named_tuple(self):
         kernel = _make_two_band_kernel()