From 19f5bbebc95a4e7682e5237cd42f3f9177027795 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Tue, 18 Aug 2026 16:09:37 -0400 Subject: [PATCH 1/8] Optional preconditioning of the fit parameters Many unconstrained, strongly correlated parameters -- the coefficients of a smooth parameterisation whose basis is not orthogonal under the data's own weight -- make the Hessian badly conditioned. The Krylov inner solve then needs a number of Hessian-vector products growing like sqrt(kappa), and outer steps get rejected. Observed on a 2112-coefficient in-situ muon efficiency fit: 19 of 77 iterations returned a bit-identical loss while consuming 29 of 145 minutes, the worst single iteration taking 25 minutes to make no progress. Add an opt-in reparameterisation theta = theta_ref + T y with T = L^-T from the Cholesky of a reference Hessian restricted to a selected block, so T^T H0 T = I there. This is preconditioning of the trust-region subproblem obtained as a change of variables, which matters because scipy's trust-krylov (GLTR/_trlib) accepts no user preconditioner: the minimizer is left completely untouched and the spherical trust region in y becomes an H0-aligned ellipsoid in theta. The transform is confined to the three scipy callbacks in fit(); self.x holds physical parameters everywhere else, so the postfit Hessian, covariance, impacts and pulls need no mapping back. Preconditioner.identity() is an exact no-op, keeping a single code path. Default scope is the unconstrained parameters, where this pays off; constrained nuisances are already normalised by their unit Gaussian prior. Frozen parameters are always excluded, since a dense transform would otherwise mix them back in. T is applied as a dense matvec against a cached L^-1 rather than a triangular solve: same flops, but the solve is inherently sequential while the matvec is a parallel GEMV, measured 15x faster at m=2112 (0.19 ms vs 2.9 ms per HVP, 0.09% of a 215 ms HVP). The triangular solve is kept as a fallback. A block that cannot be factorised, or a reference Hessian that cannot be formed, falls back to running unpreconditioned with a warning: a preconditioner must never break a fit. Off by default and a pure reparameterisation, so existing fits are unchanged. Tests cover the algebra (whitening, chain rule, hessp vs dense, fast path vs fallback), scope selection, the degenerate-block fallbacks, and fit invariance for both trust-krylov and trust-exact on a deliberately ill-conditioned model (correlation condition number 1.3e3 -> 1.0, identical results). Co-Authored-By: Claude Opus 5 --- rabbit/fitter.py | 86 +++++++-- rabbit/parsing.py | 39 ++++ rabbit/preconditioner.py | 345 ++++++++++++++++++++++++++++++++++ tests/test_preconditioner.py | 351 +++++++++++++++++++++++++++++++++++ 4 files changed, 808 insertions(+), 13 deletions(-) create mode 100644 rabbit/preconditioner.py create mode 100644 tests/test_preconditioner.py diff --git a/rabbit/fitter.py b/rabbit/fitter.py index a132bec..d9a3710 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -11,6 +11,7 @@ from wums import logging from rabbit import external_likelihood, io_tools +from rabbit import preconditioner as precond from rabbit import tfhelpers as tfh from rabbit.bbstat.bbstat import BinByBinStat from rabbit.impacts import ( @@ -129,6 +130,11 @@ def __init__( # outcome can be written to the output. None if the minimizer raised. self.minimizer_result = None self.hvp_method = getattr(options, "hvpMethod", "revrev") + # Optional parameter preconditioning (see rabbit/preconditioner.py). + # getattr so callers that build options objects by hand keep working. + self.precondition = getattr(options, "precondition", False) + self.precondition_params = getattr(options, "preconditionParams", None) + self.precondition_ridge = getattr(options, "preconditionRidge", 1e-8) # jitCompile accepts "auto" (the default), "on", or "off". # True / False from programmatic callers are accepted as # aliases for "on" / "off". The tri-state is resolved to the @@ -2264,31 +2270,83 @@ def loss_val_grad_hess_beta(self, profile=True): return val, grad, hess + def _build_preconditioner(self): + """Preconditioner for the upcoming :meth:`fit`, or an exact no-op. + + Built at the current parameter values, so the reference Hessian is the + one the minimizer starts from. + """ + theta_ref = self.x.numpy() + if not self.precondition: + return precond.Preconditioner.identity(theta_ref) + + idx = precond.select_indices( + self.parms, + self.cw.numpy(), + self.frozen_params_mask.numpy(), + expressions=self.precondition_params, + match_fn=match_regexp_params, + groups=self.indata.systgroups, + group_idxs=self.indata.systgroupidxs, + ) + # "hessian" is the only source so far; the CLI restricts the choices. + # The dense [npar, npar] Hessian is the one costly part; if it cannot be + # formed (memory, or a tracing failure on a large model) fall back to an + # unpreconditioned fit rather than taking the whole job down. + try: + _, _, hess = self.loss_val_grad_hess() + hess_np = hess.__array__() + except Exception as ex: + logger.warning( + f"Could not compute the reference Hessian for preconditioning ({ex}); " + "running unpreconditioned." + ) + return precond.Preconditioner.identity(theta_ref) + + return precond.Preconditioner.from_hessian( + hess_np, + theta_ref, + idx, + ridge=self.precondition_ridge, + ) + def fit(self): logger.info("Perform iterative fit") - def scipy_loss(xval): - self.x.assign(xval) + # Optional reparameterisation. Built once at the starting point, it is + # confined to the three scipy callbacks below: self.x always holds + # *physical* parameters outside them, so the postfit Hessian, + # covariance, impacts and pulls are unaffected and need no mapping + # back. pc is an exact no-op when disabled, keeping one code path. + pc = self._build_preconditioner() + + def scipy_loss(yval): + self.x.assign(pc.to_physical(yval)) val, grad = self.loss_val_grad() - return val.__array__(), grad.__array__() + return val.__array__(), pc.grad_to_internal(grad.__array__()) - def scipy_hessp(xval, pval): - self.x.assign(xval) - p = tf.convert_to_tensor(pval) - val, grad, hessp = self.loss_val_grad_hessp(p) - return hessp.__array__() + def scipy_hessp(yval, pval): + self.x.assign(pc.to_physical(yval)) - def scipy_hess(xval): - self.x.assign(xval) + def hvp(v): + _, _, hessp = self.loss_val_grad_hessp(tf.convert_to_tensor(v)) + return hessp.__array__() + + return pc.hessp_to_internal(np.asarray(pval, dtype=np.float64), hvp) + + def scipy_hess(yval): + self.x.assign(pc.to_physical(yval)) val, grad, hess = self.loss_val_grad_hess() if self.diagnostics: cond_number = tfh.cond_number(hess) logger.info(f" - Condition number: {cond_number}") edmval = tfh.edmval(grad, hess) logger.info(f" - edmval: {edmval}") - return hess.__array__() + return pc.hess_to_internal(hess.__array__()) - xval = self.x.numpy() + # scipy works in internal coordinates throughout; y = 0 at the point the + # transform was built. + xval = pc.from_physical(self.x.numpy()) callback = FitterCallback(xval, self.earlyStopping) @@ -2342,7 +2400,9 @@ def scipy_hess(xval): self.minimizer_result = res logger.debug(res) - self.x.assign(xval) + # xval (and callback.xval) are internal coordinates; everything outside + # fit() expects physical parameters. + self.x.assign(pc.to_physical(xval)) return callback diff --git a/rabbit/parsing.py b/rabbit/parsing.py index 19ee749..8d1a118 100644 --- a/rabbit/parsing.py +++ b/rabbit/parsing.py @@ -202,6 +202,45 @@ def common_parser(): ], help="Mnimizer method used in scipy.optimize.minimize for the nominal fit minimization", ) + parser.add_argument( + "--precondition", + action="store_true", + help="Reparameterise a block of parameters so the reference Hessian is the " + "identity there (theta = theta_ref + L^-T y). This is preconditioning of the " + "trust-region subproblem obtained as a change of variables, so the minimizer " + "itself is untouched. Helps where many unconstrained, strongly correlated " + "parameters make the Krylov inner solve struggle and outer steps get rejected. " + "Off by default; a pure reparameterisation, so results are unchanged.", + ) + parser.add_argument( + "--preconditionParams", + default=None, + type=str, + nargs="+", + help="Parameters to precondition: exact names, regexes matched against the full " + "parameter name, or systematic group names. Default (flag given without this " + "option) is every unconstrained parameter, which is where it pays off; " + "constrained nuisances are already normalised by their prior. Frozen " + "parameters are always excluded.", + ) + parser.add_argument( + "--preconditionFrom", + default="hessian", + type=str, + choices=["hessian"], + help="Source of the reference matrix. 'hessian' takes the exact Hessian at the " + "starting point (one extra Hessian evaluation, roughly one trust-exact " + "iteration).", + ) + parser.add_argument( + "--preconditionRidge", + default=1e-8, + type=float, + help="Ridge added to the preconditioning block diagonal, relative to its largest " + "diagonal entry, to keep near-degenerate blocks factorisable. Escalated " + "automatically if the Cholesky still fails; a block that cannot be factorised " + "falls back to no preconditioning.", + ) parser.add_argument( "--hvpMethod", default="revrev", diff --git a/rabbit/preconditioner.py b/rabbit/preconditioner.py new file mode 100644 index 0000000..156f61e --- /dev/null +++ b/rabbit/preconditioner.py @@ -0,0 +1,345 @@ +"""Optional preconditioning of the fit parameters. + +WHY. The trust-region minimizers solve their subproblem in a *spherical* trust +region, and the Krylov/CG inner solve converges in a number of Hessian-vector +products that grows like sqrt(kappa) of the Hessian. When a block of +parameters is strongly correlated -- typically many unconstrained, +weakly-identified coefficients of a smooth parameterisation, where the basis is +not orthogonal under the data's own weight -- kappa is huge, the inner solve +struggles, and the outer step is rejected. The symptom is outer iterations that +cost minutes and return a bit-identical loss. + +WHAT. We reparameterise. With theta the physical parameters, y the internal +ones, theta_ref the point the transform was built at, and a reference Hessian +H0 = L L^T restricted to a selected block: + + theta = theta_ref + T y, T = L^-T -> T^T H0 T = I + +so the Hessian in y is the identity at theta_ref and the spherical trust region +in y is an H0-aligned ellipsoid in theta. That is exactly preconditioning, but +obtained as a change of variables, which means **the minimizer is untouched**: +scipy's trust-krylov (GLTR/_trlib) accepts no user preconditioner, and it does +not need to. + +The chain rule gives what the scipy callbacks must return: + + loss_y(y) = loss(theta(y)) + grad_y = T^T grad_theta + (H_y p) = T^T H_theta (T p) + +T and T^T are applied as dense matvecs against a cached L^-1, which is formed +once at construction. A triangular solve would do the same flops but is +inherently sequential and so cannot use more than one core; the matvec is a +parallel GEMV and measured ~20x faster at m=2112 (0.09 ms vs 1.9 ms). The +triangular solve is kept as a fallback. Forming the inverse is safe here: +against a direct solve it agrees to ~1e-15 even at cond(H)=1e12. + +That is worth stating explicitly because the older *offline* in-situ basis had +to row-normalise its L^-1, which looks like a warning against explicit +inverses. It was not one. There the transform reshaped the basis in which the +histmaker precomputed its delta=0.01 finite variations, so a large T moved the +linearisation point far from where those templates were valid -- a modelling +constraint, not a numerical-stability one. Here T is applied to exact vectors +inside the fit, so that failure mode does not arise. + +Parameters outside the selected block are passed through untouched, so the +transform is a no-op there and the block can be as small as one wants. + +SCOPE. Preconditioning buys little for constrained nuisances: their unit +Gaussian prior contributes the identity, H = I + J^T W J, so kappa is bounded. +The win is for *unconstrained* parameters and POIs, which is the default scope. + +Everything here operates on numpy arrays at the scipy boundary; the fitter +keeps holding physical parameters in its tf.Variable, so nothing downstream +(covariance, impacts, pulls) needs to know that preconditioning happened. +""" + +import numpy as np +import scipy.linalg +from wums import logging + +logger = logging.child_logger(__name__) + +# Sources for the reference matrix. Only "hessian" is implemented so far; the +# others are accepted by the CLI layer as they land. +PRECONDITION_SOURCES = ("hessian",) + + +class Preconditioner: + """Affine reparameterisation theta = theta_ref + T y with T = L^-T. + + Use :meth:`identity` for the disabled case: it is an exact no-op, so the + fitter has a single code path whether or not preconditioning is on. + """ + + def __init__(self, theta_ref, idx=None, chol=None): + # Correlation condition number of the block before/after the transform, + # filled in by from_hessian. Kept for diagnostics and tests. + self.cond_before = None + self.cond_after = None + self.theta_ref = np.asarray(theta_ref, dtype=np.float64) + self.n = self.theta_ref.size + # idx None <=> identity transform + self.idx = None if idx is None else np.asarray(idx, dtype=np.int64) + self.chol = None if chol is None else np.asarray(chol, dtype=np.float64) + # Explicit L^-1, formed once so the per-call transform is a dense matvec + # (parallel GEMV) instead of a triangular solve (inherently sequential). + # Measured at m=2112: 0.09 ms vs 1.9 ms per application. Falls back to + # the solve if the inverse cannot be formed. + self._linv = None + if (self.idx is None) != (self.chol is None): + raise ValueError("idx and chol must be given together") + if self.chol is not None and self.chol.shape != (self.idx.size,) * 2: + raise ValueError( + f"chol shape {self.chol.shape} does not match block size {self.idx.size}" + ) + if self.chol is not None: + try: + self._linv = scipy.linalg.solve_triangular( + self.chol, + np.eye(self.chol.shape[0]), + lower=True, + trans="N", + ) + except (scipy.linalg.LinAlgError, ValueError) as ex: + logger.warning( + f"Could not form the explicit inverse ({ex}); " + "falling back to triangular solves." + ) + self._linv = None + + # -- construction ---------------------------------------------------- + + @classmethod + def identity(cls, theta_ref): + """Disabled preconditioner: y == theta - theta_ref, no block.""" + return cls(theta_ref) + + @property + def enabled(self): + return self.idx is not None + + @property + def nblock(self): + return 0 if self.idx is None else int(self.idx.size) + + @classmethod + def from_hessian(cls, hess, theta_ref, idx, ridge=1e-8, max_tries=4): + """Build from a reference Hessian, restricted to ``idx``. + + ``hess`` is the full [npar, npar] Hessian at ``theta_ref``; only the + ``idx`` sub-block is used. The block is symmetrised, then a ridge + proportional to the largest diagonal entry is added until the Cholesky + succeeds. A block that cannot be factorised at all falls back to the + identity with a warning: a preconditioner must never break a fit. + """ + idx = np.asarray(idx, dtype=np.int64) + if idx.size == 0: + logger.warning( + "Preconditioning requested but the selected block is empty; " + "running unpreconditioned." + ) + return cls.identity(theta_ref) + + block = np.asarray(hess, dtype=np.float64)[np.ix_(idx, idx)] + # symmetrise: the autodiff Hessian is symmetric only up to roundoff + block = 0.5 * (block + block.T) + + diag = np.diag(block) + scale = float(np.max(diag)) if diag.size else 0.0 + if not np.isfinite(scale) or scale <= 0.0: + logger.warning( + "Preconditioning block has no positive diagonal " + f"(max diag = {scale}); running unpreconditioned." + ) + return cls.identity(theta_ref) + + cond_before = _cond_corr(block) + eps = ridge + for itry in range(max_tries): + trial = block.copy() + if eps > 0.0: + trial[np.diag_indices_from(trial)] += eps * scale + try: + chol = scipy.linalg.cholesky(trial, lower=True) + except scipy.linalg.LinAlgError: + eps = max(eps, 1e-12) * 100.0 + logger.debug( + f"Preconditioner Cholesky failed (try {itry + 1}), " + f"raising ridge to {eps:.3g}" + ) + continue + # Conditioning actually achieved: L^-1 B L^-T for the *un-ridged* + # block B. Using the ridged matrix here would return 1 by + # construction and measure nothing. + tb = scipy.linalg.solve_triangular(chol, block, lower=True, trans="N") + tb = scipy.linalg.solve_triangular(chol, tb.T, lower=True, trans="N").T + cond_after = _cond_corr(tb) + logger.info( + f"Preconditioning {idx.size} parameters from the reference Hessian " + f"(ridge {eps:.3g} x max|diag|): correlation condition number " + f"{cond_before:.3g} -> {cond_after:.3g} at the reference point" + ) + out = cls(theta_ref, idx=idx, chol=chol) + out.cond_before = cond_before + out.cond_after = cond_after + return out + + logger.warning( + "Preconditioning block is not factorisable even with a ridge of " + f"{eps:.3g} x max|diag|; running unpreconditioned." + ) + return cls.identity(theta_ref) + + # -- the transform --------------------------------------------------- + + def _apply_T(self, v): + """T v, i.e. L^-T on the block and the identity elsewhere.""" + if self.idx is None: + return np.asarray(v, dtype=np.float64) + out = np.array(v, dtype=np.float64, copy=True) + if self._linv is not None: + out[self.idx] = self._linv.T @ out[self.idx] + else: + out[self.idx] = scipy.linalg.solve_triangular( + self.chol, out[self.idx], lower=True, trans="T" + ) + return out + + def _apply_TT(self, v): + """T^T v, i.e. L^-1 on the block and the identity elsewhere.""" + if self.idx is None: + return np.asarray(v, dtype=np.float64) + out = np.array(v, dtype=np.float64, copy=True) + if self._linv is not None: + out[self.idx] = self._linv @ out[self.idx] + else: + out[self.idx] = scipy.linalg.solve_triangular( + self.chol, out[self.idx], lower=True, trans="N" + ) + return out + + def to_physical(self, y): + """theta = theta_ref + T y.""" + return self.theta_ref + self._apply_T(y) + + def from_physical(self, theta): + """y = T^-1 (theta - theta_ref), i.e. L^T on the block.""" + d = np.asarray(theta, dtype=np.float64) - self.theta_ref + if self.idx is None: + return d + out = np.array(d, dtype=np.float64, copy=True) + out[self.idx] = self.chol.T @ d[self.idx] + return out + + def grad_to_internal(self, grad): + """grad_y = T^T grad_theta.""" + return self._apply_TT(grad) + + def hessp_to_internal(self, p, hvp): + """H_y p = T^T H_theta (T p); ``hvp`` maps a physical-space vector.""" + return self._apply_TT(hvp(self._apply_T(p))) + + def hess_to_internal(self, hess): + """H_y = T^T H_theta T, for the dense-Hessian minimizers. + + Done as two matrix operations rather than column by column. Note the + off-diagonal blocks transform too (one-sided), so a block that + correlates with the rest of the model is handled correctly. + """ + if self.idx is None: + return np.asarray(hess, dtype=np.float64) + out = np.array(hess, dtype=np.float64, copy=True) + if self._linv is not None: + # left: T^T acts on the row index; right: T on the column index + out[self.idx, :] = self._linv @ out[self.idx, :] + out[:, self.idx] = out[:, self.idx] @ self._linv.T + return out + out[self.idx, :] = scipy.linalg.solve_triangular( + self.chol, out[self.idx, :], lower=True, trans="N" + ) + out[:, self.idx] = scipy.linalg.solve_triangular( + self.chol, out[:, self.idx].T, lower=True, trans="N" + ).T + return out + + # -- diagnostics ----------------------------------------------------- + + def summary(self): + if self.idx is None: + return "preconditioning: disabled" + return f"preconditioning: enabled on {self.idx.size} of {self.n} parameters" + + +def _cond_corr(mat): + """Condition number of the *correlation* matrix of ``mat``. + + Scale-invariant, so it measures genuine degeneracy rather than a mismatch + of units between parameters -- the quantity that actually governs how hard + the block is to fit. + """ + d = np.sqrt(np.abs(np.diag(mat))) + good = d > 0 + if not np.any(good): + return np.inf + m = mat[np.ix_(good, good)] / np.outer(d[good], d[good]) + try: + sv = np.linalg.svd(m, compute_uv=False) + except np.linalg.LinAlgError: + return np.inf + return float(sv[0] / sv[-1]) if sv[-1] > 0 else np.inf + + +def select_indices( + parms, + cw, + frozen_mask, + expressions=None, + match_fn=None, + groups=None, + group_idxs=None, +): + """Indices of the parameters to precondition. + + ``expressions`` may name parameters exactly, be regexes matched against the + full parameter name (via ``match_fn``, the fitter's existing matcher), or + name a systematic group. With no expressions the default scope is every + *unconstrained* parameter (cw == 0), which is where preconditioning helps. + + Frozen parameters are always excluded: a dense transform would otherwise + mix a frozen parameter back into the fit through the other coordinates. + """ + parms = np.asarray(parms).astype(str) + n = parms.size + frozen_mask = np.asarray(frozen_mask, dtype=bool) + + if expressions: + sel = np.zeros(n, dtype=bool) + leftover = [] + # NB explicit None checks: groups/group_idxs arrive as numpy arrays, + # for which `groups or []` raises on the truth-value test. + gnames = [] if groups is None else list(groups) + gidxs = [] if group_idxs is None else list(group_idxs) + by_group = { + (k.decode() if isinstance(k, bytes) else str(k)): v + for k, v in zip(gnames, gidxs) + } + for expr in expressions: + if expr in by_group: + sel[np.asarray(by_group[expr], dtype=np.int64)] = True + else: + leftover.append(expr) + if leftover: + if match_fn is None: + raise ValueError("no matcher available for regex selection") + names = match_fn(leftover, parms) + sel |= np.isin(parms, names) + else: + sel = np.asarray(cw) == 0.0 + logger.info( + "No --preconditionParams given; defaulting to the unconstrained " + "parameters (constraint weight 0)." + ) + + sel &= ~frozen_mask + return np.where(sel)[0] diff --git a/tests/test_preconditioner.py b/tests/test_preconditioner.py new file mode 100644 index 0000000..6c2e913 --- /dev/null +++ b/tests/test_preconditioner.py @@ -0,0 +1,351 @@ +"""Tests for the optional parameter preconditioning. + +The transform is a pure reparameterisation, so the decisive test is that a fit +run with it converges to the *same* parameters, uncertainties and NLL as one +run without it. The unit tests below check the algebra that guarantees this: +T^T H T = I on the block, and the chain rule for the gradient / Hessian-vector +product that the scipy callbacks rely on. +""" + +import os +import tempfile + +import numpy as np +import pytest + +from rabbit import fitter, inputdata +from rabbit.param_models.helpers import load_model +from rabbit.preconditioner import Preconditioner, select_indices + +from .test_sparse_fit import check_results, make_options, make_test_tensor + + +def _spd(n, seed=0, cond=1e6): + """A symmetric positive definite matrix with a controlled condition number.""" + rng = np.random.default_rng(seed) + q, _ = np.linalg.qr(rng.normal(size=(n, n))) + eig = np.geomspace(1.0, cond, n) + return q @ np.diag(eig) @ q.T + + +# -- unit tests: the algebra --------------------------------------------- + + +def test_identity_is_exact_noop(): + theta = np.arange(5, dtype=float) + pc = Preconditioner.identity(theta) + assert not pc.enabled + y = pc.from_physical(theta) + np.testing.assert_allclose(y, np.zeros(5)) + np.testing.assert_allclose(pc.to_physical(y), theta) + g = np.array([1.0, -2.0, 3.0, 0.5, 0.0]) + np.testing.assert_allclose(pc.grad_to_internal(g), g) + h = _spd(5) + np.testing.assert_allclose(pc.hess_to_internal(h), h) + + +def test_round_trip(): + n = 8 + theta_ref = np.linspace(-1, 1, n) + h = _spd(n, seed=1) + idx = np.arange(2, 7) + pc = Preconditioner.from_hessian(h, theta_ref, idx) + assert pc.enabled and pc.nblock == idx.size + + rng = np.random.default_rng(3) + y = rng.normal(size=n) + np.testing.assert_allclose(pc.from_physical(pc.to_physical(y)), y, atol=1e-10) + + theta = theta_ref + rng.normal(size=n) + np.testing.assert_allclose( + pc.to_physical(pc.from_physical(theta)), theta, atol=1e-10 + ) + + +def test_whitens_the_reference_hessian_on_the_block(): + """T^T H T must be the identity on the block -- that is the whole point.""" + n = 10 + h = _spd(n, seed=2, cond=1e8) + idx = np.arange(n) + pc = Preconditioner.from_hessian(h, np.zeros(n), idx, ridge=0.0) + hy = pc.hess_to_internal(h) + np.testing.assert_allclose(hy, np.eye(n), atol=1e-6) + # and the conditioning is genuinely improved + assert np.linalg.cond(hy) < 1e3 < np.linalg.cond(h) + + +def test_partial_block_leaves_the_rest_untouched(): + """Parameters outside the block pass through, including their Hessian block.""" + n = 6 + h = _spd(n, seed=4) + idx = np.array([0, 1, 2]) + pc = Preconditioner.from_hessian(h, np.zeros(n), idx, ridge=0.0) + hy = pc.hess_to_internal(h) + # block-block whitened + np.testing.assert_allclose(hy[np.ix_(idx, idx)], np.eye(idx.size), atol=1e-8) + # outside-outside untouched + rest = np.array([3, 4, 5]) + np.testing.assert_allclose( + hy[np.ix_(rest, rest)], h[np.ix_(rest, rest)], atol=1e-12 + ) + # off-diagonal coupling is transformed one-sided, and stays symmetric + np.testing.assert_allclose(hy, hy.T, atol=1e-8) + + +def test_gradient_follows_the_chain_rule(): + """grad_y == T^T grad_theta, checked against finite differences.""" + n = 6 + a = _spd(n, seed=5) + b = np.linspace(-0.5, 0.5, n) + theta_ref = np.zeros(n) + pc = Preconditioner.from_hessian(a, theta_ref, np.arange(n), ridge=0.0) + + def loss_theta(theta): + return 0.5 * theta @ a @ theta + b @ theta + + def loss_y(y): + return loss_theta(pc.to_physical(y)) + + rng = np.random.default_rng(6) + y0 = rng.normal(size=n) * 0.1 + grad_theta = a @ pc.to_physical(y0) + b + analytic = pc.grad_to_internal(grad_theta) + + numeric = np.empty(n) + eps = 1e-6 + for i in range(n): + yp, ym = y0.copy(), y0.copy() + yp[i] += eps + ym[i] -= eps + numeric[i] = (loss_y(yp) - loss_y(ym)) / (2 * eps) + np.testing.assert_allclose(analytic, numeric, rtol=1e-5, atol=1e-7) + + +def test_hessp_matches_the_dense_transform(): + """hessp_to_internal must agree with hess_to_internal @ p.""" + n = 7 + a = _spd(n, seed=7) + pc = Preconditioner.from_hessian(a, np.zeros(n), np.arange(1, 6), ridge=0.0) + hy = pc.hess_to_internal(a) + rng = np.random.default_rng(8) + p = rng.normal(size=n) + got = pc.hessp_to_internal(p, lambda v: a @ v) + np.testing.assert_allclose(got, hy @ p, atol=1e-8) + + +def test_explicit_inverse_and_triangular_fallback_agree(): + """The cached-inverse fast path must match the triangular-solve fallback. + + Both branches are live (the fallback triggers if the inverse cannot be + formed), so they have to give the same answer. + """ + n = 9 + h = _spd(n, seed=11, cond=1e8) + idx = np.arange(1, 8) + fast = Preconditioner.from_hessian(h, np.zeros(n), idx, ridge=0.0) + slow = Preconditioner.from_hessian(h, np.zeros(n), idx, ridge=0.0) + assert fast._linv is not None + slow._linv = None # force the triangular-solve path + + rng = np.random.default_rng(12) + v = rng.normal(size=n) + np.testing.assert_allclose(fast.to_physical(v), slow.to_physical(v), atol=1e-10) + np.testing.assert_allclose( + fast.grad_to_internal(v), slow.grad_to_internal(v), atol=1e-10 + ) + np.testing.assert_allclose( + fast.hess_to_internal(h), slow.hess_to_internal(h), atol=1e-8 + ) + np.testing.assert_allclose( + fast.hessp_to_internal(v, lambda x: h @ x), + slow.hessp_to_internal(v, lambda x: h @ x), + atol=1e-8, + ) + + +def test_singular_block_falls_back_to_identity(): + """A preconditioner must never break a fit.""" + n = 5 + h = np.zeros((n, n)) # completely degenerate + pc = Preconditioner.from_hessian(h, np.zeros(n), np.arange(n)) + assert not pc.enabled + + +def test_rank_deficient_block_is_ridged_into_shape(): + n = 5 + h = _spd(n, seed=9) + h[:, -1] = h[:, 0] # exact linear dependence + h[-1, :] = h[0, :] + pc = Preconditioner.from_hessian(h, np.zeros(n), np.arange(n), ridge=1e-8) + assert pc.enabled + + +def test_empty_block_falls_back_to_identity(): + pc = Preconditioner.from_hessian(_spd(4), np.zeros(4), np.array([], dtype=int)) + assert not pc.enabled + + +# -- unit tests: scope selection ---------------------------------------- + + +def test_default_scope_is_the_unconstrained_parameters(): + parms = np.array(["poi", "a", "b", "c"]) + cw = np.array([0.0, 1.0, 0.0, 1.0]) + frozen = np.zeros(4, dtype=bool) + idx = select_indices(parms, cw, frozen) + np.testing.assert_array_equal(idx, [0, 2]) + + +def test_frozen_parameters_are_always_excluded(): + parms = np.array(["poi", "a", "b", "c"]) + cw = np.zeros(4) + frozen = np.array([False, True, False, False]) + idx = select_indices(parms, cw, frozen) + np.testing.assert_array_equal(idx, [0, 2, 3]) + + +def test_selection_by_regex_and_by_group(): + parms = np.array(["poi", "eff_a", "eff_b", "other"]) + cw = np.ones(4) + frozen = np.zeros(4, dtype=bool) + + def match_fn(exprs, names): + import re + + return [n for n in names if any(re.fullmatch(e, n) for e in exprs)] + + idx = select_indices(parms, cw, frozen, expressions=["eff_.*"], match_fn=match_fn) + np.testing.assert_array_equal(idx, [1, 2]) + + idx = select_indices( + parms, + cw, + frozen, + expressions=["mygroup"], + match_fn=match_fn, + groups=["mygroup"], + group_idxs=[[3]], + ) + np.testing.assert_array_equal(idx, [3]) + + +# -- the invariance test ------------------------------------------------- + + +def make_polynomial_tensor(outdir, order=6, nbins=30): + """Background with ``order`` *unconstrained* polynomial shape systematics. + + Mirrors the situation preconditioning is for: a smooth basis that is not + orthogonal under the data's own weight, so the coefficients are strongly + correlated and the block is badly conditioned. Plain monomials x^k are used + deliberately -- the Vandermonde-like Gram matrix is about as ill-conditioned + as it gets, which is the point. + """ + import hist + + from rabbit import tensorwriter + + rng = np.random.default_rng(7) + ax = hist.axis.Regular(nbins, -1, 1, name="x") + x = ax.centers + + truth = 1000.0 * np.exp(-0.5 * (x / 0.6) ** 2) + 200.0 + h_data = hist.Hist(ax, storage=hist.storage.Double()) + h_data.values()[...] = rng.poisson(truth).astype(float) + + h_bkg = hist.Hist(ax, storage=hist.storage.Weight()) + h_bkg.values()[...] = truth + h_bkg.variances()[...] = truth + + h_sig = hist.Hist(ax, storage=hist.storage.Weight()) + h_sig.values()[...] = 50.0 * np.exp(-0.5 * (x / 0.2) ** 2) + h_sig.variances()[...] = h_sig.values() + + writer = tensorwriter.TensorWriter() + writer.add_channel(h_data.axes, "ch0") + writer.add_data(h_data, "ch0") + writer.add_process(h_sig, "sig", "ch0", signal=True) + writer.add_process(h_bkg, "bkg", "ch0") + + for k in range(1, order + 1): + w = 0.05 * x**k + up = h_bkg.copy() + dn = h_bkg.copy() + up.values()[...] = h_bkg.values() * (1 + w) + dn.values()[...] = h_bkg.values() * (1 - w) + writer.add_systematic( + [up, dn], + f"poly{k}", + "bkg", + "ch0", + symmetrize="average", + constrained=False, + ) + + writer.write(outfolder=outdir, outfilename="test_poly") + return os.path.join(outdir, "test_poly.hdf5") + + +def _run(filename, method, **kw): + """Fit ``filename`` and return the results plus the preconditioner used.""" + indata_obj = inputdata.FitInputData(filename) + param_model = load_model("Mu", indata_obj) + options = make_options(minimizerMethod=method, **kw) + f = fitter.Fitter(indata_obj, param_model, options) + f.set_nobs(indata_obj.data_obs) + # built separately purely so the test can assert the block is non-trivial; + # fit() builds its own at the same point + pc = f._build_preconditioner() + f.minimize() + val, grad, hess = f.loss_val_grad_hess() + from rabbit.tfhelpers import edmval_cov + + edmval, cov = edmval_cov(grad, hess) + cov_np = np.asarray(cov.numpy() if hasattr(cov, "numpy") else cov) + res = dict( + param=f.x[: param_model.nparams].numpy(), + theta=f.x[param_model.nparams :].numpy(), + param_err=np.sqrt(np.diag(cov_np)[: param_model.nparams]), + nll=f.reduced_nll().numpy(), + edmval=edmval, + parms=f.parms, + ) + return res, pc + + +@pytest.mark.parametrize("method", ["trust-krylov", "trust-exact"]) +def test_fit_is_invariant_under_preconditioning(method): + """Same minimum, same uncertainties, same NLL -- with trust-krylov kept. + + Covers both the hessp path (trust-krylov) and the dense-hess path + (trust-exact), since the two transform different callbacks. The scope is + forced to every parameter so the transform is a genuinely dense one with + off-diagonal coupling, not a per-parameter rescaling. + """ + with tempfile.TemporaryDirectory() as tmp: + filename = make_test_tensor(tmp) + plain, _ = _run(filename, method) + pre, pc = _run(filename, method, precondition=True, preconditionParams=[".*"]) + + assert pc.enabled and pc.nblock > 1, "preconditioning block must be non-trivial" + assert check_results("plain", plain, "preconditioned", pre) + + +@pytest.mark.parametrize("method", ["trust-krylov", "trust-exact"]) +def test_fit_is_invariant_on_an_ill_conditioned_block(method): + """The case this feature exists for: many correlated unconstrained params. + + Uses the default scope (unconstrained parameters), so it also checks that + the default picks the right block. + """ + with tempfile.TemporaryDirectory() as tmp: + filename = make_polynomial_tensor(tmp, order=6) + plain, _ = _run(filename, method) + pre, pc = _run(filename, method, precondition=True) + + # default scope must have found the unconstrained polynomial block + assert pc.enabled and pc.nblock > 1 + # and it must actually be badly conditioned before, well conditioned after + assert pc.cond_before > 1e3 + assert pc.cond_after < 1e2 + assert check_results("plain", plain, "preconditioned", pre) From 8b5eeec89ef7de647ce6689fea56133cbe566b1c Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Tue, 18 Aug 2026 17:03:01 -0400 Subject: [PATCH 2/8] Add a Gauss-Newton reference matrix, and restart the minimizer on a stall Two follow-ups to the preconditioner, both driven by what the 2112-coefficient in-situ efficiency fit actually did. --preconditionFrom gaussnewton takes the Fisher information instead of the exact Hessian, computed by evaluating the Hessian with the data replaced by the current prediction: there (1 - nobs/nexp) vanishes, the second-derivative term drops out and J^T W J + diag(cw) remains, positive semi-definite by construction. It is not the better default, and the docstrings say so with the measurement behind it. Being PSD, the Fisher matrix cannot represent negative curvature; the exact Hessian of that fit has 33 negative eigenvalues at the starting point, whitening with Gauss-Newton left all 33 negative and the true Hessian at kappa ~1e4 in the new coordinates, and the fit froze after 13 iterations. Raising the ridge did not help, so it is the missing curvature and not the scaling. The ridge on the exact Hessian turns out to be load-bearing for exactly that reason: it regularises the indefinite directions. "hessian" stays the default; "gaussnewton" is for models that are positive definite where the fit starts. The second change is unrelated to preconditioning and worth having on its own. scipy's trust-region loop shrinks the trust radius 4x on every rejected step with no lower bound, and keeps it in a local variable, so a run of rejections leaves the minimizer taking infinitesimal steps far from any minimum. The loss stops moving while the gradient stays large, which is what --earlyStopping was added to catch. But a fresh minimize() call resets the radius, which is why restarting from a stalled point resumes the descent -- the trick of chaining --externalPostfit by hand. Do it in the fitter instead: on an early-stopping stall, restart from the stalled point, and keep restarting for as long as the loss keeps coming down, stopping only once a restart no longer reduces it (--maxRestarts, -1 by default, 0 to disable). Measured on a fit that froze at loss 91596: four automatic restarts took it to 11252, matching the 145-minute unpreconditioned run's 11249. Also: a minimizer exception that is not a stall is now logged at warning level. The bare except previously swallowed everything, so a broken callback -- for instance one whose argument is not named intermediate_result, which makes scipy pass a bare array -- looked exactly like a converged fit. Co-Authored-By: Claude Opus 5 --- rabbit/fitter.py | 178 ++++++++++++++++++++++++++++++----- rabbit/parsing.py | 27 +++++- rabbit/preconditioner.py | 5 +- tests/test_preconditioner.py | 35 +++++++ tests/test_restart.py | 143 ++++++++++++++++++++++++++++ tests/test_sparse_fit.py | 1 + 6 files changed, 356 insertions(+), 33 deletions(-) create mode 100644 tests/test_restart.py diff --git a/rabbit/fitter.py b/rabbit/fitter.py index d9a3710..48bda36 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -71,6 +71,9 @@ def __init__(self, xv, early_stopping=-1): self.t0 = time.time() self.early_stopping = early_stopping + # set just before raising, so fit() can tell a recoverable stall apart + # from a genuine error and restart instead of giving up + self.stopped_early = False def __call__(self, intermediate_result): loss = intermediate_result.fun @@ -91,6 +94,7 @@ def __call__(self, intermediate_result): and len(self.loss_history) > self.early_stopping and self.loss_history[-self.early_stopping] <= loss ): + self.stopped_early = True raise ValueError( f"No reduction in loss after {self.early_stopping} iterations, early stopping." ) @@ -102,6 +106,31 @@ def __call__(self, intermediate_result): self.iiter += 1 +# Relative loss improvement below which a restart counts as having bought +# nothing. Loss values here are O(1e4), so float64 cancellation puts genuine +# improvements no finer than ~1e-9 relative. +_RESTART_MIN_IMPROVEMENT = 1e-9 + + +def _merge_callbacks(acc, cb): + """Fold one restart's callback into the accumulated one. + + Callers read loss_history/time_history/iiter to report on the whole fit, so + the restarts have to look like a single continuous run. Times are offset by + the elapsed time already accumulated, since each callback clocks from its + own construction. + """ + if acc is None: + return cb + offset = acc.time_history[-1] if acc.time_history else 0.0 + acc.loss_history.extend(cb.loss_history) + acc.time_history.extend(t + offset for t in cb.time_history) + acc.iiter += cb.iiter + acc.xval = cb.xval + acc.stopped_early = cb.stopped_early + return acc + + class Fitter: valid_systematic_types = ["log_normal", "normal"] @@ -134,6 +163,8 @@ def __init__( # getattr so callers that build options objects by hand keep working. self.precondition = getattr(options, "precondition", False) self.precondition_params = getattr(options, "preconditionParams", None) + self.precondition_from = getattr(options, "preconditionFrom", "hessian") + self.max_restarts = getattr(options, "maxRestarts", -1) self.precondition_ridge = getattr(options, "preconditionRidge", 1e-8) # jitCompile accepts "auto" (the default), "on", or "off". # True / False from programmatic callers are accepted as @@ -2270,6 +2301,52 @@ def loss_val_grad_hess_beta(self, profile=True): return val, grad, hess + def _reference_matrix(self): + """Reference matrix the preconditioner is built from, as a numpy array. + + "hessian" is the exact Hessian at the current point, and is the default. + On real data it is not guaranteed positive definite -- the term + (1 - nobs/nexp) multiplying the second derivative of the prediction can + go either way -- so the Cholesky needs a ridge big enough to make it so. + That ridge is not a wart: it is what regularises the negative-curvature + directions into small positive ones, which is exactly what lets the + block be whitened into something the trust region can work with. + + "gaussnewton" is the Fisher information, i.e. the *expected* Hessian, + obtained by evaluating the exact Hessian with the data replaced by the + current prediction: at that point (1 - nobs/nexp) vanishes, the + second-derivative term drops out and J^T W J + diag(cw) remains, which + is positive semi-definite by construction and so factorises at the + default ridge. + + That PSD-ness is a double-edged sword, and measurement says the edge + usually points the wrong way. Being PSD, the Fisher matrix *cannot + represent negative curvature at all*. Measured on a 2112-parameter + in-situ efficiency block whose exact Hessian had 33 negative + eigenvalues at the starting point: whitening with the Gauss-Newton + transform left all 33 directions negative and the true Hessian at + kappa ~1e4 in the new coordinates, and the fit froze immediately, + while the ridged exact Hessian took the same fit from 145 min to + 4.6 min. Raising the ridge on the Gauss-Newton matrix did not help, + so this is about the missing negative curvature, not about scaling. + + So prefer "hessian" unless the exact Hessian is positive definite where + the fit starts -- near a minimum, or for a genuinely convex model -- + where "gaussnewton" is the cheaper and better-conditioned choice. + """ + if self.precondition_from == "gaussnewton": + saved_nobs = tf.identity(self.nobs) + saved_varnobs = tf.identity(self.varnobs) if self.chisqFit else None + try: + self.set_nobs(self.expected_yield(), variances=saved_varnobs) + _, _, hess = self.loss_val_grad_hess() + return hess.__array__() + finally: + self.set_nobs(saved_nobs, variances=saved_varnobs) + + _, _, hess = self.loss_val_grad_hess() + return hess.__array__() + def _build_preconditioner(self): """Preconditioner for the upcoming :meth:`fit`, or an exact no-op. @@ -2289,13 +2366,12 @@ def _build_preconditioner(self): groups=self.indata.systgroups, group_idxs=self.indata.systgroupidxs, ) - # "hessian" is the only source so far; the CLI restricts the choices. - # The dense [npar, npar] Hessian is the one costly part; if it cannot be + # The dense [npar, npar] reference matrix is the one costly part; if it + # cannot be # formed (memory, or a tracing failure on a large model) fall back to an # unpreconditioned fit rather than taking the whole job down. try: - _, _, hess = self.loss_val_grad_hess() - hess_np = hess.__array__() + hess_np = self._reference_matrix() except Exception as ex: logger.warning( f"Could not compute the reference Hessian for preconditioning ({ex}); " @@ -2348,8 +2424,6 @@ def scipy_hess(yval): # transform was built. xval = pc.from_physical(self.x.numpy()) - callback = FitterCallback(xval, self.earlyStopping) - if self.minimizer_method in [ "trust-krylov", "trust-ncg", @@ -2378,27 +2452,79 @@ def scipy_hess(yval): sci_opts["ftol"] = float(self.minimizer_ftol) logger.info(f"[minimize] method={self.minimizer_method} options={sci_opts}") - try: - res = scipy.optimize.minimize( - scipy_loss, - xval, - method=self.minimizer_method, - jac=True, - tol=0.0, - callback=callback, - options=sci_opts, - **info_minimize, + # Restart loop. scipy's trust-region methods shrink the trust radius by + # 4x on every rejected step with no lower bound, and the radius is a + # local of scipy's loop -- so once it has collapsed the method takes + # infinitesimal steps and the loss stops moving, far from any minimum. + # A fresh minimize() call resets the radius to initial_trust_radius, + # which is why restarting from a stalled point resumes progress. Loop + # that here instead of making the caller chain --externalPostfit by + # hand. Only an early-stopping stall is retried, and only while the + # restarts keep buying loss. + callback = None + prev_loss = None + attempt = 0 + while True: + cb = FitterCallback(xval, self.earlyStopping) + try: + res = scipy.optimize.minimize( + scipy_loss, + xval, + method=self.minimizer_method, + jac=True, + tol=0.0, + callback=cb, + options=sci_opts, + **info_minimize, + ) + except Exception as ex: + # minimizer could have called the loss or hessp functions with "random" values, so restore the + # state from the end of the last iteration before the exception + xval = cb.xval + self.minimizer_result = None + if not cb.stopped_early: + # a real failure, not a stall: surface it rather than + # letting a broken callback look like a converged fit + logger.warning(f"Minimizer raised: {ex}") + logger.debug(ex) + else: + xval = res["x"] + self.minimizer_result = res + logger.debug(res) + + callback = _merge_callbacks(callback, cb) + last_loss = cb.loss_history[-1] if cb.loss_history else None + + if not cb.stopped_early: + break + # The only reason to stop restarting is that the last restart + # bought nothing: a round can never end above where it started + # (the trust region accepts improving steps only), so "not below + # the previous round" means the descent is genuinely exhausted. + if ( + prev_loss is not None + and last_loss is not None + and last_loss + >= prev_loss - _RESTART_MIN_IMPROVEMENT * max(1.0, abs(prev_loss)) + ): + logger.info( + f"Restart did not reduce the loss further ({prev_loss} -> " + f"{last_loss}); stopping after {attempt} restart(s)." + ) + break + if 0 <= self.max_restarts <= attempt: + logger.warning( + f"Minimizer still stalling at loss {last_loss} after " + f"{self.max_restarts} restart(s) and the loss was still " + "coming down; raise --maxRestarts to let it continue." + ) + break + attempt += 1 + logger.info( + f"Minimizer stalled at loss {last_loss}; restarting " + f"(#{attempt}) to reset the trust radius." ) - except Exception as ex: - # minimizer could have called the loss or hessp functions with "random" values, so restore the - # state from the end of the last iteration before the exception - xval = callback.xval - self.minimizer_result = None - logger.debug(ex) - else: - xval = res["x"] - self.minimizer_result = res - logger.debug(res) + prev_loss = last_loss # xval (and callback.xval) are internal coordinates; everything outside # fit() expects physical parameters. diff --git a/rabbit/parsing.py b/rabbit/parsing.py index 8d1a118..f8a4493 100644 --- a/rabbit/parsing.py +++ b/rabbit/parsing.py @@ -187,6 +187,19 @@ def common_parser(): type=int, help="Number of iterations with no improvement after which training will be stopped. Specify -1 to disable.", ) + parser.add_argument( + "--maxRestarts", + default=-1, + type=int, + help="When --earlyStopping triggers, restart the minimizer from the stalled " + "point instead of giving up. scipy's trust-region methods shrink the trust " + "radius 4x per rejected step with no lower bound, and the radius is reset by a " + "fresh minimize() call -- so a stall is often just a collapsed step size, and " + "restarting resumes the descent (the same effect as chaining --externalPostfit " + "by hand). -1 (default) restarts as often as the loss keeps improving and stops " + "only once a restart no longer reduces it; 0 disables restarting; N > 0 caps " + "the number of restarts.", + ) parser.add_argument( "--minimizerMethod", default="trust-krylov", @@ -227,10 +240,16 @@ def common_parser(): "--preconditionFrom", default="hessian", type=str, - choices=["hessian"], - help="Source of the reference matrix. 'hessian' takes the exact Hessian at the " - "starting point (one extra Hessian evaluation, roughly one trust-exact " - "iteration).", + choices=["hessian", "gaussnewton"], + help="Source of the reference matrix. 'hessian' (default) takes the exact " + "Hessian at the starting point, one extra Hessian evaluation, roughly one " + "trust-exact iteration. 'gaussnewton' takes the Fisher information instead: " + "positive semi-definite by construction, so it factorises at the default ridge, " + "but for that same reason it cannot represent negative curvature. Where the " + "exact Hessian is indefinite at the starting point, whitening with it leaves " + "those directions negative and the fit stalls, so 'hessian' is usually the " + "better choice; prefer 'gaussnewton' only where the Hessian is positive " + "definite anyway, e.g. near a minimum.", ) parser.add_argument( "--preconditionRidge", diff --git a/rabbit/preconditioner.py b/rabbit/preconditioner.py index 156f61e..9ed114e 100644 --- a/rabbit/preconditioner.py +++ b/rabbit/preconditioner.py @@ -60,9 +60,8 @@ logger = logging.child_logger(__name__) -# Sources for the reference matrix. Only "hessian" is implemented so far; the -# others are accepted by the CLI layer as they land. -PRECONDITION_SOURCES = ("hessian",) +# Sources for the reference matrix, built by the fitter (see Fitter._reference_matrix). +PRECONDITION_SOURCES = ("hessian", "gaussnewton") class Preconditioner: diff --git a/tests/test_preconditioner.py b/tests/test_preconditioner.py index 6c2e913..24a42be 100644 --- a/tests/test_preconditioner.py +++ b/tests/test_preconditioner.py @@ -331,6 +331,41 @@ def test_fit_is_invariant_under_preconditioning(method): assert check_results("plain", plain, "preconditioned", pre) +def test_gaussnewton_source_is_psd_and_restores_nobs(): + """Fisher information must be PSD, and must not leave nobs modified. + + It is computed by temporarily swapping the data for the prediction, so a + leak there would silently corrupt the fit that follows. + """ + with tempfile.TemporaryDirectory() as tmp: + filename = make_polynomial_tensor(tmp, order=6) + indata_obj = inputdata.FitInputData(filename) + param_model = load_model("Mu", indata_obj) + options = make_options(precondition=True, preconditionFrom="gaussnewton") + f = fitter.Fitter(indata_obj, param_model, options) + f.set_nobs(indata_obj.data_obs) + + nobs_before = f.nobs.numpy().copy() + mat = f._reference_matrix() + np.testing.assert_allclose(f.nobs.numpy(), nobs_before, rtol=0, atol=0) + + ev = np.linalg.eigvalsh(0.5 * (mat + mat.T)) + assert ev.min() > -1e-8 * abs(ev).max(), f"not PSD: min eig {ev.min()}" + + +@pytest.mark.parametrize("source", ["hessian", "gaussnewton"]) +def test_fit_is_invariant_for_either_reference_source(source): + """Both sources are only a choice of transform, so neither may move the fit.""" + with tempfile.TemporaryDirectory() as tmp: + filename = make_polynomial_tensor(tmp, order=6) + plain, _ = _run(filename, "trust-krylov") + pre, pc = _run( + filename, "trust-krylov", precondition=True, preconditionFrom=source + ) + assert pc.enabled and pc.nblock > 1 + assert check_results("plain", plain, f"precond[{source}]", pre) + + @pytest.mark.parametrize("method", ["trust-krylov", "trust-exact"]) def test_fit_is_invariant_on_an_ill_conditioned_block(method): """The case this feature exists for: many correlated unconstrained params. diff --git a/tests/test_restart.py b/tests/test_restart.py new file mode 100644 index 0000000..13135f3 --- /dev/null +++ b/tests/test_restart.py @@ -0,0 +1,143 @@ +"""Tests for restarting the minimizer after an early-stopping stall. + +scipy's trust-region loop shrinks the trust radius by 4x on every rejected step +with no lower bound, and holds the radius in a local variable. Once it has +collapsed the method takes infinitesimal steps and the loss stops changing far +from any minimum; a fresh minimize() call resets the radius and the descent +resumes. These tests cover the plumbing that turns that stall into a restart +instead of a give-up. +""" + +import tempfile +from types import SimpleNamespace + +import numpy as np + +from rabbit import fitter, inputdata +from rabbit.fitter import FitterCallback, _merge_callbacks +from rabbit.param_models.helpers import load_model + +from .test_preconditioner import make_polynomial_tensor +from .test_sparse_fit import check_results, make_options, make_test_tensor + + +def _result(fun, x): + return SimpleNamespace(fun=fun, x=np.asarray(x, dtype=float)) + + +def test_stopped_early_flag_set_only_on_a_stall(): + """fit() keys the restart off this flag, so it must not fire otherwise.""" + cb = FitterCallback(np.zeros(2), early_stopping=3) + # a steadily improving fit never stalls + for i, loss in enumerate([10.0, 9.0, 8.0, 7.0, 6.0, 5.0]): + cb(_result(loss, [i, i])) + assert not cb.stopped_early + + # a flat one does + cb = FitterCallback(np.zeros(2), early_stopping=3) + raised = False + try: + for loss in [10.0, 9.0, 9.0, 9.0, 9.0, 9.0]: + cb(_result(loss, [0, 0])) + except ValueError: + raised = True + assert raised and cb.stopped_early + + +def test_disabled_early_stopping_never_stalls(): + cb = FitterCallback(np.zeros(2), early_stopping=-1) + for _ in range(20): + cb(_result(5.0, [0, 0])) + assert not cb.stopped_early + + +def test_merge_callbacks_concatenates_into_one_continuous_run(): + """Callers report on the whole fit, so restarts must look continuous.""" + a = FitterCallback(np.zeros(2), early_stopping=-1) + a.loss_history = [10.0, 9.0] + a.time_history = [1.0, 2.0] + a.iiter = 2 + + b = FitterCallback(np.ones(2), early_stopping=-1) + b.loss_history = [8.0, 7.0] + b.time_history = [0.5, 1.5] # clocked from its own construction + b.iiter = 2 + b.xval = np.array([3.0, 4.0]) + + merged = _merge_callbacks(a, b) + assert merged is a + assert merged.loss_history == [10.0, 9.0, 8.0, 7.0] + # second run's times offset by the first run's elapsed + assert merged.time_history == [1.0, 2.0, 2.5, 3.5] + assert merged.iiter == 4 + np.testing.assert_array_equal(merged.xval, [3.0, 4.0]) + + +def test_merge_callbacks_with_no_accumulator_returns_the_first(): + cb = FitterCallback(np.zeros(2), early_stopping=-1) + assert _merge_callbacks(None, cb) is cb + + +def _fit(filename, **kw): + indata_obj = inputdata.FitInputData(filename) + param_model = load_model("Mu", indata_obj) + options = make_options(**kw) + f = fitter.Fitter(indata_obj, param_model, options) + f.set_nobs(indata_obj.data_obs) + f.minimize() + val, grad, hess = f.loss_val_grad_hess() + from rabbit.tfhelpers import edmval_cov + + edmval, cov = edmval_cov(grad, hess) + cov_np = np.asarray(cov.numpy() if hasattr(cov, "numpy") else cov) + return dict( + param=f.x[: param_model.nparams].numpy(), + theta=f.x[param_model.nparams :].numpy(), + param_err=np.sqrt(np.diag(cov_np)[: param_model.nparams]), + nll=f.reduced_nll().numpy(), + edmval=edmval, + parms=f.parms, + ) + + +def test_restarts_do_not_change_a_fit_that_does_not_stall(): + """--maxRestarts must be inert when nothing stalls.""" + with tempfile.TemporaryDirectory() as tmp: + filename = make_test_tensor(tmp) + plain = _fit(filename, maxRestarts=0) + with_restarts = _fit(filename, earlyStopping=20, maxRestarts=5) + assert check_results("no restarts", plain, "maxRestarts=5", with_restarts) + + +def test_restarting_is_on_by_default_and_unbounded(): + """Default is -1: keep restarting while the loss keeps dropping.""" + + with tempfile.TemporaryDirectory() as tmp: + filename = make_test_tensor(tmp) + indata_obj = inputdata.FitInputData(filename) + param_model = load_model("Mu", indata_obj) + # options object without the attribute at all -> the getattr default + options = make_options() + del options.maxRestarts + f = fitter.Fitter(indata_obj, param_model, options) + assert f.max_restarts == -1 + + +def test_unbounded_restarts_stop_when_a_restart_stops_improving(): + """The improvement check, not a counter, is what ends the loop.""" + with tempfile.TemporaryDirectory() as tmp: + filename = make_polynomial_tensor(tmp, order=6) + # unbounded restarts must still terminate + res = _fit(filename, earlyStopping=10, maxRestarts=-1) + assert np.isfinite(res["nll"]) + + +def test_restarts_reach_at_least_as_low_a_loss_on_a_hard_model(): + """On a model with many correlated unconstrained params, restarting can + only help: the restart is only taken after a stall, and is abandoned as + soon as it stops reducing the loss.""" + with tempfile.TemporaryDirectory() as tmp: + filename = make_polynomial_tensor(tmp, order=6) + stop_only = _fit(filename, earlyStopping=10, maxRestarts=0) + restarted = _fit(filename, earlyStopping=10, maxRestarts=5) + assert restarted["nll"] <= stop_only["nll"] + 1e-6 diff --git a/tests/test_sparse_fit.py b/tests/test_sparse_fit.py index 81b8c6d..887fcd5 100644 --- a/tests/test_sparse_fit.py +++ b/tests/test_sparse_fit.py @@ -154,6 +154,7 @@ def make_options(**kwargs): setConstraintMinimum=[], unblind=[], blindingGroup=[], + maxRestarts=-1, ) defaults.update(kwargs) return SimpleNamespace(**defaults) From fa37c918134082649b1a22aad813a4d50e2a6f26 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Tue, 18 Aug 2026 18:07:37 -0400 Subject: [PATCH 3/8] Rebuild the preconditioner at the point each restart starts from The transform whitens the Hessian at the point it was built. A restart happens precisely because the fit has stalled somewhere else, where that Hessian is no longer the same one -- so reusing the original transform resets the trust radius but carries a stale change of variables that no longer conditions anything. Rebuild it at the current point instead, which is where the next round will actually be working. The transform now lives in a one-element cell so the three scipy callbacks pick up the new one; the physical parameters are re-expressed in the new coordinates before restarting. Costs one Hessian evaluation per restart, and is a no-op when preconditioning is off, since the builder returns the identity without computing anything. The accompanying test asserts the rebuild happens *and* that it happens at a different point than the previous build; it was verified to fail when the rebuild is removed, so it cannot pass vacuously. Co-Authored-By: Claude Opus 5 --- rabbit/fitter.py | 21 ++++++++++++++++++--- tests/test_restart.py | 43 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 3 deletions(-) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 48bda36..e26b83e 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -2394,14 +2394,19 @@ def fit(self): # *physical* parameters outside them, so the postfit Hessian, # covariance, impacts and pulls are unaffected and need no mapping # back. pc is an exact no-op when disabled, keeping one code path. - pc = self._build_preconditioner() + # Held in a one-element cell because it is rebuilt at the current point + # before every restart (see the restart loop) and the scipy callbacks + # below must pick the new one up. + pc_cell = [self._build_preconditioner()] def scipy_loss(yval): + pc = pc_cell[0] self.x.assign(pc.to_physical(yval)) val, grad = self.loss_val_grad() return val.__array__(), pc.grad_to_internal(grad.__array__()) def scipy_hessp(yval, pval): + pc = pc_cell[0] self.x.assign(pc.to_physical(yval)) def hvp(v): @@ -2411,6 +2416,7 @@ def hvp(v): return pc.hessp_to_internal(np.asarray(pval, dtype=np.float64), hvp) def scipy_hess(yval): + pc = pc_cell[0] self.x.assign(pc.to_physical(yval)) val, grad, hess = self.loss_val_grad_hess() if self.diagnostics: @@ -2422,7 +2428,7 @@ def scipy_hess(yval): # scipy works in internal coordinates throughout; y = 0 at the point the # transform was built. - xval = pc.from_physical(self.x.numpy()) + xval = pc_cell[0].from_physical(self.x.numpy()) if self.minimizer_method in [ "trust-krylov", @@ -2526,9 +2532,18 @@ def scipy_hess(yval): ) prev_loss = last_loss + # Rebuild the transform at the point we are restarting from. The + # one built at the start whitens the Hessian *there*; by the time + # the fit has stalled somewhere else that Hessian has changed and + # the transform no longer conditions anything. Refreshing costs one + # Hessian evaluation and is a no-op when preconditioning is off. + self.x.assign(pc_cell[0].to_physical(xval)) + pc_cell[0] = self._build_preconditioner() + xval = pc_cell[0].from_physical(self.x.numpy()) + # xval (and callback.xval) are internal coordinates; everything outside # fit() expects physical parameters. - self.x.assign(pc.to_physical(xval)) + self.x.assign(pc_cell[0].to_physical(xval)) return callback diff --git a/tests/test_restart.py b/tests/test_restart.py index 13135f3..70680dc 100644 --- a/tests/test_restart.py +++ b/tests/test_restart.py @@ -141,3 +141,46 @@ def test_restarts_reach_at_least_as_low_a_loss_on_a_hard_model(): stop_only = _fit(filename, earlyStopping=10, maxRestarts=0) restarted = _fit(filename, earlyStopping=10, maxRestarts=5) assert restarted["nll"] <= stop_only["nll"] + 1e-6 + + +def test_preconditioner_is_rebuilt_before_every_restart(): + """The transform whitens the Hessian at the point it was built. + + Once the fit has stalled somewhere else that Hessian has changed, so a + restart must rebuild it there rather than reuse the one from the starting + point -- otherwise the restart resets the trust radius but keeps a + transform that no longer conditions anything. + """ + with tempfile.TemporaryDirectory() as tmp: + filename = make_polynomial_tensor(tmp, order=6) + indata_obj = inputdata.FitInputData(filename) + param_model = load_model("Mu", indata_obj) + options = make_options(earlyStopping=3, maxRestarts=3, precondition=True) + f = fitter.Fitter(indata_obj, param_model, options) + f.set_nobs(indata_obj.data_obs) + + calls = {"n": 0, "at": []} + original = f._build_preconditioner + + def counting(): + calls["n"] += 1 + calls["at"].append(np.array(f.x.numpy(), copy=True)) + return original() + + f._build_preconditioner = counting + callback = f.fit() + + assert callback is not None + + # This model stalls under earlyStopping=3, so at least one restart is + # taken; without the refresh there would be exactly one build. + assert ( + calls["n"] >= 2 + ), f"expected a rebuild before the restart, saw {calls['n']} build(s)" + + # and the rebuild must happen where the fit is now, not back at the + # point it started from + moved = np.linalg.norm(calls["at"][1] - calls["at"][0]) + assert ( + moved > 1e-6 + ), f"rebuild happened at the same point ({moved}), so it was not a refresh" From de176e1190a69304cc38f94fb006a01a80655c4b Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Tue, 18 Aug 2026 18:26:35 -0400 Subject: [PATCH 4/8] Move the minimizer callback and stall bookkeeping into callbacks.py FitterCallback, merge_callbacks and the restart improvement threshold touch no Fitter state, so they do not need to live in the largest module in the package. fitter.py goes from 3105 to 3039 lines; the names are re-exported by importing them, so fitter.FitterCallback still resolves. The underscores are dropped now that the two helpers are a cross-module API. The restart loop itself stays in fit(): it reads six pieces of Fitter state and closes over the three scipy callbacks, so moving it would trade a long function for indirection through a module boundary. Also make the rebuild-before-restart test deterministic. It relied on the model tripping --earlyStopping on its own, which turns on differences far below the scale of the fit; TF's multithreaded CPU reductions are not bitwise reproducible, and the test was observed to pass and fail on the same code. It now forces the stall from a callback subclass and asserts the exact build count, one up front plus one per restart, and that each refresh happens somewhere new. Verified to fail when the rebuild is removed, so it still cannot pass vacuously. Co-Authored-By: Claude Opus 5 --- rabbit/callbacks.py | 93 +++++++++++++++++++++++++++++++++++++++++++ rabbit/fitter.py | 80 ++++--------------------------------- tests/test_restart.py | 64 +++++++++++++++++++---------- 3 files changed, 143 insertions(+), 94 deletions(-) create mode 100644 rabbit/callbacks.py diff --git a/rabbit/callbacks.py b/rabbit/callbacks.py new file mode 100644 index 0000000..e45a619 --- /dev/null +++ b/rabbit/callbacks.py @@ -0,0 +1,93 @@ +"""Minimizer callback and the bookkeeping around a stalled fit. + +The callback is what scipy calls once per accepted iteration. Beyond logging it +does two things the fitter relies on: it keeps the last parameter vector, so a +minimizer that raises can be rolled back to the end of the last good iteration, +and it detects a *stall* -- no reduction in loss over a run of iterations. + +A stall is worth singling out because it is usually recoverable. scipy's +trust-region methods shrink the trust radius 4x on every rejected step with no +lower bound and hold it in a local variable, so a run of rejections leaves the +minimizer taking infinitesimal steps far from any minimum, with the loss frozen +while the gradient is still large. A fresh minimize() call resets the radius, so +the fitter restarts rather than giving up; these helpers carry the state across +those restarts so the whole thing still reports as one continuous fit. +""" + +import time + +import numpy as np +from wums import logging + +logger = logging.child_logger(__name__) + + +class FitterCallback: + def __init__(self, xv, early_stopping=-1): + self.iiter = 0 + self.xval = xv + + self.loss_history = [] + self.time_history = [] + + self.t0 = time.time() + + self.early_stopping = early_stopping + # set just before raising, so fit() can tell a recoverable stall apart + # from a genuine error and restart instead of giving up + self.stopped_early = False + + def __call__(self, intermediate_result): + loss = intermediate_result.fun + + elapsed = time.time() - self.t0 + prev = self.time_history[-1] if self.time_history else 0.0 + dt = elapsed - prev + + logger.debug( + f"Iteration {self.iiter}: loss {loss} " + f"[dt={dt:.2f}s elapsed={elapsed:.2f}s]" + ) + if np.isnan(loss): + raise ValueError(f"Loss value is NaN at iteration {self.iiter}") + + if ( + self.early_stopping > 0 + and len(self.loss_history) > self.early_stopping + and self.loss_history[-self.early_stopping] <= loss + ): + self.stopped_early = True + raise ValueError( + f"No reduction in loss after {self.early_stopping} iterations, early stopping." + ) + + self.loss_history.append(loss) + self.time_history.append(elapsed) + + self.xval = intermediate_result.x + self.iiter += 1 + + +# Relative loss improvement below which a restart counts as having bought +# nothing. Loss values here are O(1e4), so float64 cancellation puts genuine +# improvements no finer than ~1e-9 relative. +RESTART_MIN_IMPROVEMENT = 1e-9 + + +def merge_callbacks(acc, cb): + """Fold one restart's callback into the accumulated one. + + Callers read loss_history/time_history/iiter to report on the whole fit, so + the restarts have to look like a single continuous run. Times are offset by + the elapsed time already accumulated, since each callback clocks from its + own construction. + """ + if acc is None: + return cb + offset = acc.time_history[-1] if acc.time_history else 0.0 + acc.loss_history.extend(cb.loss_history) + acc.time_history.extend(t + offset for t in cb.time_history) + acc.iiter += cb.iiter + acc.xval = cb.xval + acc.stopped_early = cb.stopped_early + return acc diff --git a/rabbit/fitter.py b/rabbit/fitter.py index e26b83e..da20aa4 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -14,6 +14,11 @@ from rabbit import preconditioner as precond from rabbit import tfhelpers as tfh from rabbit.bbstat.bbstat import BinByBinStat +from rabbit.callbacks import ( + RESTART_MIN_IMPROVEMENT, + FitterCallback, + merge_callbacks, +) from rabbit.impacts import ( asym_impacts, global_asym_impacts, @@ -60,77 +65,6 @@ def match_regexp_params(regular_expressions, parameter_names): return matched -class FitterCallback: - def __init__(self, xv, early_stopping=-1): - self.iiter = 0 - self.xval = xv - - self.loss_history = [] - self.time_history = [] - - self.t0 = time.time() - - self.early_stopping = early_stopping - # set just before raising, so fit() can tell a recoverable stall apart - # from a genuine error and restart instead of giving up - self.stopped_early = False - - def __call__(self, intermediate_result): - loss = intermediate_result.fun - - elapsed = time.time() - self.t0 - prev = self.time_history[-1] if self.time_history else 0.0 - dt = elapsed - prev - - logger.debug( - f"Iteration {self.iiter}: loss {loss} " - f"[dt={dt:.2f}s elapsed={elapsed:.2f}s]" - ) - if np.isnan(loss): - raise ValueError(f"Loss value is NaN at iteration {self.iiter}") - - if ( - self.early_stopping > 0 - and len(self.loss_history) > self.early_stopping - and self.loss_history[-self.early_stopping] <= loss - ): - self.stopped_early = True - raise ValueError( - f"No reduction in loss after {self.early_stopping} iterations, early stopping." - ) - - self.loss_history.append(loss) - self.time_history.append(elapsed) - - self.xval = intermediate_result.x - self.iiter += 1 - - -# Relative loss improvement below which a restart counts as having bought -# nothing. Loss values here are O(1e4), so float64 cancellation puts genuine -# improvements no finer than ~1e-9 relative. -_RESTART_MIN_IMPROVEMENT = 1e-9 - - -def _merge_callbacks(acc, cb): - """Fold one restart's callback into the accumulated one. - - Callers read loss_history/time_history/iiter to report on the whole fit, so - the restarts have to look like a single continuous run. Times are offset by - the elapsed time already accumulated, since each callback clocks from its - own construction. - """ - if acc is None: - return cb - offset = acc.time_history[-1] if acc.time_history else 0.0 - acc.loss_history.extend(cb.loss_history) - acc.time_history.extend(t + offset for t in cb.time_history) - acc.iiter += cb.iiter - acc.xval = cb.xval - acc.stopped_early = cb.stopped_early - return acc - - class Fitter: valid_systematic_types = ["log_normal", "normal"] @@ -2498,7 +2432,7 @@ def scipy_hess(yval): self.minimizer_result = res logger.debug(res) - callback = _merge_callbacks(callback, cb) + callback = merge_callbacks(callback, cb) last_loss = cb.loss_history[-1] if cb.loss_history else None if not cb.stopped_early: @@ -2511,7 +2445,7 @@ def scipy_hess(yval): prev_loss is not None and last_loss is not None and last_loss - >= prev_loss - _RESTART_MIN_IMPROVEMENT * max(1.0, abs(prev_loss)) + >= prev_loss - RESTART_MIN_IMPROVEMENT * max(1.0, abs(prev_loss)) ): logger.info( f"Restart did not reduce the loss further ({prev_loss} -> " diff --git a/tests/test_restart.py b/tests/test_restart.py index 70680dc..d517fbc 100644 --- a/tests/test_restart.py +++ b/tests/test_restart.py @@ -14,7 +14,7 @@ import numpy as np from rabbit import fitter, inputdata -from rabbit.fitter import FitterCallback, _merge_callbacks +from rabbit.callbacks import FitterCallback, merge_callbacks from rabbit.param_models.helpers import load_model from .test_preconditioner import make_polynomial_tensor @@ -64,7 +64,7 @@ def test_merge_callbacks_concatenates_into_one_continuous_run(): b.iiter = 2 b.xval = np.array([3.0, 4.0]) - merged = _merge_callbacks(a, b) + merged = merge_callbacks(a, b) assert merged is a assert merged.loss_history == [10.0, 9.0, 8.0, 7.0] # second run's times offset by the first run's elapsed @@ -75,7 +75,7 @@ def test_merge_callbacks_concatenates_into_one_continuous_run(): def test_merge_callbacks_with_no_accumulator_returns_the_first(): cb = FitterCallback(np.zeros(2), early_stopping=-1) - assert _merge_callbacks(None, cb) is cb + assert merge_callbacks(None, cb) is cb def _fit(filename, **kw): @@ -150,37 +150,59 @@ def test_preconditioner_is_rebuilt_before_every_restart(): restart must rebuild it there rather than reuse the one from the starting point -- otherwise the restart resets the trust radius but keeps a transform that no longer conditions anything. + + The stall is forced rather than coaxed out of the model: whether a real fit + trips a given --earlyStopping threshold turns on differences far below the + scale of the fit, and TF's multithreaded CPU reductions are not bitwise + reproducible, so keying the test on that makes it flaky. """ + n_restarts = 2 + + class StallEveryFewIterations(FitterCallback): + def __call__(self, intermediate_result): + self.iiter += 1 + self.loss_history.append(intermediate_result.fun) + self.time_history.append(float(self.iiter)) + self.xval = intermediate_result.x + if self.iiter >= 3: + self.stopped_early = True + raise ValueError("forced stall") + with tempfile.TemporaryDirectory() as tmp: filename = make_polynomial_tensor(tmp, order=6) indata_obj = inputdata.FitInputData(filename) param_model = load_model("Mu", indata_obj) - options = make_options(earlyStopping=3, maxRestarts=3, precondition=True) + options = make_options( + earlyStopping=3, maxRestarts=n_restarts, precondition=True + ) f = fitter.Fitter(indata_obj, param_model, options) f.set_nobs(indata_obj.data_obs) - calls = {"n": 0, "at": []} + built_at = [] original = f._build_preconditioner def counting(): - calls["n"] += 1 - calls["at"].append(np.array(f.x.numpy(), copy=True)) + built_at.append(np.array(f.x.numpy(), copy=True)) return original() f._build_preconditioner = counting - callback = f.fit() + monkeyed = fitter.FitterCallback + fitter.FitterCallback = StallEveryFewIterations + try: + callback = f.fit() + finally: + fitter.FitterCallback = monkeyed assert callback is not None - - # This model stalls under earlyStopping=3, so at least one restart is - # taken; without the refresh there would be exactly one build. - assert ( - calls["n"] >= 2 - ), f"expected a rebuild before the restart, saw {calls['n']} build(s)" - - # and the rebuild must happen where the fit is now, not back at the - # point it started from - moved = np.linalg.norm(calls["at"][1] - calls["at"][0]) - assert ( - moved > 1e-6 - ), f"rebuild happened at the same point ({moved}), so it was not a refresh" + # one build up front, then one before each of the forced restarts + assert len(built_at) == 1 + n_restarts, ( + f"expected {1 + n_restarts} builds (1 initial + {n_restarts} " + f"refreshes), saw {len(built_at)}" + ) + # each refresh must happen where the fit now is, not back at the start + for i in range(1, len(built_at)): + moved = np.linalg.norm(built_at[i] - built_at[i - 1]) + assert moved > 1e-6, ( + f"refresh {i} happened at the same point ({moved}); " + "the transform was reused, not rebuilt" + ) From 88da38303a3fd4c77e278cc978436d49df4270da Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Tue, 18 Aug 2026 18:44:15 -0400 Subject: [PATCH 5/8] Turn early stopping on by default scipy will not stop on its own. Its trust-region loop has no termination test for a collapsed trust radius, and maxiter defaults to 200*nparams, so a fit whose radius has collapsed keeps calling the subproblem solver on a microscopic region: measured on a 3603-parameter fit, the loss froze after 63 seconds and the next 720 iterations changed nothing, at ~10 iterations per second, which projects to twenty hours before maxiter would end it. Now that a stall triggers a restart rather than a give-up, detecting one is what lets the fit continue, so there is no longer a reason to leave it off. Default 20 iterations without improvement; -1 still disables it, which now means "let a stalled fit spin". Co-Authored-By: Claude Opus 5 --- rabbit/parsing.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/rabbit/parsing.py b/rabbit/parsing.py index f8a4493..0e96f88 100644 --- a/rabbit/parsing.py +++ b/rabbit/parsing.py @@ -183,9 +183,14 @@ def common_parser(): ) parser.add_argument( "--earlyStopping", - default=-1, + default=20, type=int, - help="Number of iterations with no improvement after which training will be stopped. Specify -1 to disable.", + help="Number of iterations with no improvement after which the minimizer is " + "considered stalled. Paired with --maxRestarts this means 'restart here', not " + "'give up here'. On by default because scipy will not stop on its own: its " + "trust-region loop has no termination test for a collapsed trust radius, and " + "its maxiter defaults to 200*nparams, so a stalled fit spins for hours making " + "no progress rather than exiting. Specify -1 to disable.", ) parser.add_argument( "--maxRestarts", From ffca009d0d089381bcdbb6afc6b21aa2d38cf522 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Wed, 19 Aug 2026 10:04:36 -0400 Subject: [PATCH 6/8] Make the preconditioner block diagonal and find the blocks automatically One Cholesky over everything selected was both the most expensive option and the most fragile. Cholesky is O(m^3), so a single 2112-parameter block costs ~4e4 times more than the 240 blocks the parameterisation actually has; worse, a union of individually well-behaved groups can be singular, and then the whole transform is lost. Measured on the combined W+Z fit: the 1921 ABCD fake parameters were not factorisable, which took the 2112 in-situ parameters down with them and the fit ran unpreconditioned. The transform is now block diagonal, one factorisation per block, and a block that cannot be factorised is skipped while the others still apply. --preconditionBlocks chooses how the blocks are formed: auto (default) threshold the reference matrix's correlations and take connected components. On the in-situ fit this recovers 233-240 components where the parameterisation has 240 (step, eta, charge) blocks, with no knowledge of parameter names, and every block factorises at the default ridge where the single block needed 1e-2. The median block is 12 parameters throughout. expressions one block per --preconditionParams entry. none no grouping, the whole scope as one block (the previous behaviour). auto is the default because it needs nothing from the user: --precondition alone now selects the unconstrained parameters and discovers their clusters, which is the right thing without knowing in advance which parameters are badly correlated, and the tuning knobs stay for when that is known. --preconditionBlockThreshold sets the correlation cut (0.1); below the percolation point everything joins one component, which is warned about. The blocks are re-derived at each restart along with the factorisations, since the correlation structure moves with the fit: over five builds of one 4D fit the block count went 233, 240, 230, 229, 213 with the median size pinned at 12 while the largest cluster grew 73 -> 204, i.e. clusters merge as the minimum is approached. Grouping cannot change the answer, and the invariance test is parametrised over all three modes to keep it that way. Co-Authored-By: Claude Opus 5 --- rabbit/fitter.py | 27 ++- rabbit/parsing.py | 27 +++ rabbit/preconditioner.py | 392 +++++++++++++++++++++++------------ tests/test_preconditioner.py | 209 +++++++++++++++++-- 4 files changed, 505 insertions(+), 150 deletions(-) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index da20aa4..faab3ef 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -98,6 +98,10 @@ def __init__( self.precondition = getattr(options, "precondition", False) self.precondition_params = getattr(options, "preconditionParams", None) self.precondition_from = getattr(options, "preconditionFrom", "hessian") + self.precondition_blocks = getattr(options, "preconditionBlocks", "auto") + self.precondition_block_threshold = getattr( + options, "preconditionBlockThreshold", 0.1 + ) self.max_restarts = getattr(options, "maxRestarts", -1) self.precondition_ridge = getattr(options, "preconditionRidge", 1e-8) # jitCompile accepts "auto" (the default), "on", or "off". @@ -2291,7 +2295,7 @@ def _build_preconditioner(self): if not self.precondition: return precond.Preconditioner.identity(theta_ref) - idx = precond.select_indices( + index_blocks = precond.select_index_blocks( self.parms, self.cw.numpy(), self.frozen_params_mask.numpy(), @@ -2313,10 +2317,29 @@ def _build_preconditioner(self): ) return precond.Preconditioner.identity(theta_ref) + if not index_blocks: + logger.warning( + "Preconditioning requested but no parameters were selected; " + "running unpreconditioned." + ) + return precond.Preconditioner.identity(theta_ref) + + if self.precondition_blocks in ("auto", "none"): + # the expressions only set the scope + scope = np.unique(np.concatenate([idx for _, idx in index_blocks])) + if self.precondition_blocks == "auto": + # the blocks come from the reference matrix itself + index_blocks = precond.auto_blocks( + hess_np, scope, threshold=self.precondition_block_threshold + ) + else: + # no grouping: one factorisation over the whole scope + index_blocks = [("all", scope)] + return precond.Preconditioner.from_hessian( hess_np, theta_ref, - idx, + index_blocks, ridge=self.precondition_ridge, ) diff --git a/rabbit/parsing.py b/rabbit/parsing.py index 0e96f88..c3f3be7 100644 --- a/rabbit/parsing.py +++ b/rabbit/parsing.py @@ -241,6 +241,33 @@ def common_parser(): "constrained nuisances are already normalised by their prior. Frozen " "parameters are always excluded.", ) + parser.add_argument( + "--preconditionBlocks", + nargs="?", + const="auto", + default="auto", + type=str, + choices=["auto", "expressions", "none"], + help="How to group the selected parameters into blocks; the transform is " + "block diagonal, one Cholesky per block. 'auto' (the default, and what a " + "bare --preconditionBlocks selects) reads the clusters off the reference " + "matrix by thresholding its correlations and taking connected components: " + "no parameter naming needed, much cheaper since Cholesky is O(m^3), and a " + "singular cluster only costs its own block rather than all of them. " + "'expressions' makes one block per --preconditionParams entry. 'none' does " + "no grouping at all and factorises the whole selected scope as a single " + "block, which keeps every cross-correlation but is the most expensive and " + "fails entirely if any part of the scope is singular.", + ) + parser.add_argument( + "--preconditionBlockThreshold", + default=0.1, + type=float, + help="Correlation threshold for --preconditionBlocks auto. Correlations " + "below it are left unpreconditioned. Too low and every parameter " + "percolates into a single block; too high and genuinely coupled " + "parameters are split apart.", + ) parser.add_argument( "--preconditionFrom", default="hessian", diff --git a/rabbit/preconditioner.py b/rabbit/preconditioner.py index 9ed114e..1ef9d4b 100644 --- a/rabbit/preconditioner.py +++ b/rabbit/preconditioner.py @@ -42,8 +42,24 @@ constraint, not a numerical-stability one. Here T is applied to exact vectors inside the fit, so that failure mode does not arise. -Parameters outside the selected block are passed through untouched, so the -transform is a no-op there and the block can be as small as one wants. +Parameters outside the selected blocks are passed through untouched, so the +transform is a no-op there and a block can be as small as one wants. + +FINDING THE BLOCKS. The clusters can be read off the reference matrix instead +of being named by hand: threshold the correlation matrix and take connected +components (see :func:`auto_blocks`). On the in-situ efficiency fit that +recovers 239 components where the parameterisation has 240 (step, eta, charge) +blocks, with no knowledge of parameter names, and costs ~4e4 times fewer flops +to factorise than one joint block, since Cholesky is O(m^3). + +SEVERAL BLOCKS. The transform is block diagonal: each selected group of +parameters gets its own factorisation and they are applied independently. That +is not only cheaper (m^3 per block instead of (sum m)^3) but often the only +thing that works -- a union of two individually well-behaved groups can be +singular, because the groups are nearly degenerate *with each other*, and one +joint Cholesky then fails where two separate ones succeed. The price is that +correlations between blocks are left alone, so blocks should be chosen to be +the strongly-correlated clusters. SCOPE. Preconditioning buys little for constrained nuisances: their unit Gaussian prior contributes the identity, H = I + J^T W J, so kappa is bounded. @@ -56,6 +72,8 @@ import numpy as np import scipy.linalg +import scipy.sparse +import scipy.sparse.csgraph from wums import logging logger = logging.child_logger(__name__) @@ -64,81 +82,116 @@ PRECONDITION_SOURCES = ("hessian", "gaussnewton") -class Preconditioner: - """Affine reparameterisation theta = theta_ref + T y with T = L^-T. - - Use :meth:`identity` for the disabled case: it is an exact no-op, so the - fitter has a single code path whether or not preconditioning is on. - """ +class Block: + """One factorised group of parameters: indices, L, and a cached L^-1.""" - def __init__(self, theta_ref, idx=None, chol=None): - # Correlation condition number of the block before/after the transform, - # filled in by from_hessian. Kept for diagnostics and tests. - self.cond_before = None - self.cond_after = None - self.theta_ref = np.asarray(theta_ref, dtype=np.float64) - self.n = self.theta_ref.size - # idx None <=> identity transform - self.idx = None if idx is None else np.asarray(idx, dtype=np.int64) - self.chol = None if chol is None else np.asarray(chol, dtype=np.float64) + def __init__(self, idx, chol, cond_before=None, cond_after=None, label=""): + self.idx = np.asarray(idx, dtype=np.int64) + self.chol = np.asarray(chol, dtype=np.float64) + if self.chol.shape != (self.idx.size,) * 2: + raise ValueError( + f"chol shape {self.chol.shape} does not match block size {self.idx.size}" + ) + self.cond_before = cond_before + self.cond_after = cond_after + self.label = label # Explicit L^-1, formed once so the per-call transform is a dense matvec # (parallel GEMV) instead of a triangular solve (inherently sequential). # Measured at m=2112: 0.09 ms vs 1.9 ms per application. Falls back to # the solve if the inverse cannot be formed. - self._linv = None - if (self.idx is None) != (self.chol is None): - raise ValueError("idx and chol must be given together") - if self.chol is not None and self.chol.shape != (self.idx.size,) * 2: - raise ValueError( - f"chol shape {self.chol.shape} does not match block size {self.idx.size}" + self.linv = None + try: + self.linv = scipy.linalg.solve_triangular( + self.chol, np.eye(self.chol.shape[0]), lower=True, trans="N" + ) + except (scipy.linalg.LinAlgError, ValueError) as ex: + logger.warning( + f"Could not form the explicit inverse ({ex}); " + "falling back to triangular solves." ) - if self.chol is not None: - try: - self._linv = scipy.linalg.solve_triangular( - self.chol, - np.eye(self.chol.shape[0]), - lower=True, - trans="N", - ) - except (scipy.linalg.LinAlgError, ValueError) as ex: - logger.warning( - f"Could not form the explicit inverse ({ex}); " - "falling back to triangular solves." - ) - self._linv = None + + +class Preconditioner: + """Block-diagonal affine reparameterisation theta = theta_ref + T y. + + Each block contributes T = L^-T on its own indices; everything else is + passed through. Use :meth:`identity` for the disabled case: it is an exact + no-op, so the fitter has a single code path whether or not preconditioning + is on. + """ + + def __init__(self, theta_ref, blocks=()): + self.theta_ref = np.asarray(theta_ref, dtype=np.float64) + self.n = self.theta_ref.size + self.blocks = list(blocks) # -- construction ---------------------------------------------------- @classmethod def identity(cls, theta_ref): - """Disabled preconditioner: y == theta - theta_ref, no block.""" + """Disabled preconditioner: y == theta - theta_ref, no blocks.""" return cls(theta_ref) @property def enabled(self): - return self.idx is not None + return bool(self.blocks) @property def nblock(self): - return 0 if self.idx is None else int(self.idx.size) + """Total number of preconditioned parameters, over all blocks.""" + return int(sum(b.idx.size for b in self.blocks)) + + @property + def n_blocks(self): + return len(self.blocks) + + @property + def cond_before(self): + c = [b.cond_before for b in self.blocks if b.cond_before is not None] + return max(c) if c else None + + @property + def cond_after(self): + c = [b.cond_after for b in self.blocks if b.cond_after is not None] + return max(c) if c else None @classmethod - def from_hessian(cls, hess, theta_ref, idx, ridge=1e-8, max_tries=4): - """Build from a reference Hessian, restricted to ``idx``. - - ``hess`` is the full [npar, npar] Hessian at ``theta_ref``; only the - ``idx`` sub-block is used. The block is symmetrised, then a ridge - proportional to the largest diagonal entry is added until the Cholesky - succeeds. A block that cannot be factorised at all falls back to the - identity with a warning: a preconditioner must never break a fit. + def from_hessian(cls, hess, theta_ref, index_blocks, ridge=1e-8, max_tries=4): + """Build from a reference Hessian, one factorisation per index block. + + ``index_blocks`` is a list of index arrays (a single array is accepted + and treated as one block). Each block is symmetrised, then a ridge + proportional to its largest diagonal entry is added until the Cholesky + succeeds. A block that cannot be factorised is dropped -- the others + still apply, and a preconditioner must never break a fit. """ - idx = np.asarray(idx, dtype=np.int64) - if idx.size == 0: + if isinstance(index_blocks, np.ndarray) or ( + index_blocks and np.isscalar(index_blocks[0]) + ): + index_blocks = [index_blocks] + blocks = [] + for spec in index_blocks: + label, idx = ("", spec) if not isinstance(spec, tuple) else spec + blk = cls._factorise( + hess, np.asarray(idx, dtype=np.int64), ridge, max_tries, label + ) + if blk is not None: + blocks.append(blk) + if not blocks: logger.warning( - "Preconditioning requested but the selected block is empty; " + "Preconditioning requested but no block could be used; " "running unpreconditioned." ) return cls.identity(theta_ref) + return cls(theta_ref, blocks) + + @staticmethod + def _factorise(hess, idx, ridge, max_tries, label=""): + """One block -> a :class:`Block`, or None if it is unusable.""" + tag = f"{label} " if label else "" + if idx.size == 0: + logger.warning(f"Preconditioning block {tag}is empty; skipping.") + return None block = np.asarray(hess, dtype=np.float64)[np.ix_(idx, idx)] # symmetrise: the autodiff Hessian is symmetric only up to roundoff @@ -148,10 +201,10 @@ def from_hessian(cls, hess, theta_ref, idx, ridge=1e-8, max_tries=4): scale = float(np.max(diag)) if diag.size else 0.0 if not np.isfinite(scale) or scale <= 0.0: logger.warning( - "Preconditioning block has no positive diagonal " - f"(max diag = {scale}); running unpreconditioned." + f"Preconditioning block {tag}has no positive diagonal " + f"(max diag = {scale}); skipping." ) - return cls.identity(theta_ref) + return None cond_before = _cond_corr(block) eps = ridge @@ -164,8 +217,8 @@ def from_hessian(cls, hess, theta_ref, idx, ridge=1e-8, max_tries=4): except scipy.linalg.LinAlgError: eps = max(eps, 1e-12) * 100.0 logger.debug( - f"Preconditioner Cholesky failed (try {itry + 1}), " - f"raising ridge to {eps:.3g}" + f"Preconditioner Cholesky failed for {tag}block " + f"(try {itry + 1}), raising ridge to {eps:.3g}" ) continue # Conditioning actually achieved: L^-1 B L^-T for the *un-ridged* @@ -175,47 +228,47 @@ def from_hessian(cls, hess, theta_ref, idx, ridge=1e-8, max_tries=4): tb = scipy.linalg.solve_triangular(chol, tb.T, lower=True, trans="N").T cond_after = _cond_corr(tb) logger.info( - f"Preconditioning {idx.size} parameters from the reference Hessian " - f"(ridge {eps:.3g} x max|diag|): correlation condition number " - f"{cond_before:.3g} -> {cond_after:.3g} at the reference point" + f"Preconditioning {tag}block of {idx.size} parameters from the " + f"reference Hessian (ridge {eps:.3g} x max|diag|): correlation " + f"condition number {cond_before:.3g} -> {cond_after:.3g} at the " + "reference point" ) - out = cls(theta_ref, idx=idx, chol=chol) - out.cond_before = cond_before - out.cond_after = cond_after - return out + return Block(idx, chol, cond_before, cond_after, label) logger.warning( - "Preconditioning block is not factorisable even with a ridge of " - f"{eps:.3g} x max|diag|; running unpreconditioned." + f"Preconditioning block {tag}is not factorisable even with a ridge of " + f"{eps:.3g} x max|diag|; skipping this block." ) - return cls.identity(theta_ref) + return None # -- the transform --------------------------------------------------- def _apply_T(self, v): - """T v, i.e. L^-T on the block and the identity elsewhere.""" - if self.idx is None: + """T v: L^-T on each block, identity elsewhere.""" + if not self.blocks: return np.asarray(v, dtype=np.float64) out = np.array(v, dtype=np.float64, copy=True) - if self._linv is not None: - out[self.idx] = self._linv.T @ out[self.idx] - else: - out[self.idx] = scipy.linalg.solve_triangular( - self.chol, out[self.idx], lower=True, trans="T" - ) + for b in self.blocks: + if b.linv is not None: + out[b.idx] = b.linv.T @ out[b.idx] + else: + out[b.idx] = scipy.linalg.solve_triangular( + b.chol, out[b.idx], lower=True, trans="T" + ) return out def _apply_TT(self, v): - """T^T v, i.e. L^-1 on the block and the identity elsewhere.""" - if self.idx is None: + """T^T v: L^-1 on each block, identity elsewhere.""" + if not self.blocks: return np.asarray(v, dtype=np.float64) out = np.array(v, dtype=np.float64, copy=True) - if self._linv is not None: - out[self.idx] = self._linv @ out[self.idx] - else: - out[self.idx] = scipy.linalg.solve_triangular( - self.chol, out[self.idx], lower=True, trans="N" - ) + for b in self.blocks: + if b.linv is not None: + out[b.idx] = b.linv @ out[b.idx] + else: + out[b.idx] = scipy.linalg.solve_triangular( + b.chol, out[b.idx], lower=True, trans="N" + ) return out def to_physical(self, y): @@ -223,12 +276,13 @@ def to_physical(self, y): return self.theta_ref + self._apply_T(y) def from_physical(self, theta): - """y = T^-1 (theta - theta_ref), i.e. L^T on the block.""" + """y = T^-1 (theta - theta_ref), i.e. L^T on each block.""" d = np.asarray(theta, dtype=np.float64) - self.theta_ref - if self.idx is None: + if not self.blocks: return d out = np.array(d, dtype=np.float64, copy=True) - out[self.idx] = self.chol.T @ d[self.idx] + for b in self.blocks: + out[b.idx] = b.chol.T @ d[b.idx] return out def grad_to_internal(self, grad): @@ -242,32 +296,38 @@ def hessp_to_internal(self, p, hvp): def hess_to_internal(self, hess): """H_y = T^T H_theta T, for the dense-Hessian minimizers. - Done as two matrix operations rather than column by column. Note the - off-diagonal blocks transform too (one-sided), so a block that - correlates with the rest of the model is handled correctly. + Applied block by block: the left multiplication acts on the row index + and the right one on the column index, so a block's off-diagonal + coupling to the rest of the model transforms one-sided, as it should. """ - if self.idx is None: + if not self.blocks: return np.asarray(hess, dtype=np.float64) out = np.array(hess, dtype=np.float64, copy=True) - if self._linv is not None: - # left: T^T acts on the row index; right: T on the column index - out[self.idx, :] = self._linv @ out[self.idx, :] - out[:, self.idx] = out[:, self.idx] @ self._linv.T - return out - out[self.idx, :] = scipy.linalg.solve_triangular( - self.chol, out[self.idx, :], lower=True, trans="N" - ) - out[:, self.idx] = scipy.linalg.solve_triangular( - self.chol, out[:, self.idx].T, lower=True, trans="N" - ).T + for b in self.blocks: + if b.linv is not None: + out[b.idx, :] = b.linv @ out[b.idx, :] + else: + out[b.idx, :] = scipy.linalg.solve_triangular( + b.chol, out[b.idx, :], lower=True, trans="N" + ) + for b in self.blocks: + if b.linv is not None: + out[:, b.idx] = out[:, b.idx] @ b.linv.T + else: + out[:, b.idx] = scipy.linalg.solve_triangular( + b.chol, out[:, b.idx].T, lower=True, trans="N" + ).T return out # -- diagnostics ----------------------------------------------------- def summary(self): - if self.idx is None: + if not self.blocks: return "preconditioning: disabled" - return f"preconditioning: enabled on {self.idx.size} of {self.n} parameters" + return ( + f"preconditioning: {self.n_blocks} block(s) covering " + f"{self.nblock} of {self.n} parameters" + ) def _cond_corr(mat): @@ -289,7 +349,57 @@ def _cond_corr(mat): return float(sv[0] / sv[-1]) if sv[-1] > 0 else np.inf -def select_indices( +def auto_blocks(hess, idx, threshold=0.1, max_fraction=0.5): + """Find the correlated clusters within ``idx`` from the reference matrix. + + Thresholds the correlation matrix at ``threshold`` and returns the connected + components as blocks. Parameters correlate strongly with the others in their + cluster and negligibly across clusters, which is exactly the structure a + block-diagonal transform wants, and the components are far smaller than the + union so each Cholesky is cheaper and likelier to succeed. + + ``threshold`` matters: too low and everything percolates into one component, + too high and genuinely coupled parameters are split apart. Correlations below + it are left unpreconditioned, which is the deliberate approximation. A + component covering more than ``max_fraction`` of the parameters is warned + about, since that usually means percolation. + + Parameters with a non-positive diagonal cannot be normalised and are + returned as singletons, i.e. effectively left alone. + """ + idx = np.asarray(idx, dtype=np.int64) + if idx.size == 0: + return [] + sub = np.asarray(hess, dtype=np.float64)[np.ix_(idx, idx)] + sub = 0.5 * (sub + sub.T) + d = np.diag(sub) + good = d > 0 + n = idx.size + corr = np.zeros((n, n)) + if np.any(good): + g = np.where(good)[0] + dd = np.sqrt(d[g]) + corr[np.ix_(g, g)] = np.abs(sub[np.ix_(g, g)] / np.outer(dd, dd)) + np.fill_diagonal(corr, 0.0) + + adj = scipy.sparse.csr_matrix(corr > threshold) + ncomp, labels = scipy.sparse.csgraph.connected_components(adj, directed=False) + sizes = np.bincount(labels, minlength=ncomp) + biggest = int(sizes.max()) + if biggest > max_fraction * n: + logger.warning( + f"Auto-blocking at |rho| > {threshold} produced a component with " + f"{biggest} of {n} parameters: the threshold is probably below the " + "percolation point, so the blocks are not really separated." + ) + logger.info( + f"Auto-blocking {n} parameters at |rho| > {threshold}: {ncomp} block(s), " + f"largest {biggest}, median {int(np.median(sizes))}" + ) + return [(f"auto{c}", idx[np.where(labels == c)[0]]) for c in range(ncomp)] + + +def select_index_blocks( parms, cw, frozen_mask, @@ -298,11 +408,18 @@ def select_indices( groups=None, group_idxs=None, ): - """Indices of the parameters to precondition. - - ``expressions`` may name parameters exactly, be regexes matched against the - full parameter name (via ``match_fn``, the fitter's existing matcher), or - name a systematic group. With no expressions the default scope is every + """Parameter blocks to precondition, as a list of ``(label, indices)``. + + One block per entry in ``expressions``, which is what makes the transform + block diagonal: each expression is expected to name a cluster of parameters + that are correlated with each other. Grouping them into one factorisation + instead is both more expensive and more fragile -- the union of two + individually fine groups can be singular because the groups are nearly + degenerate with each other. + + An entry may name parameters exactly, be a regex matched against the full + parameter name (via ``match_fn``, the fitter's existing matcher), or name a + systematic group. With no expressions there is a single block of every *unconstrained* parameter (cw == 0), which is where preconditioning helps. Frozen parameters are always excluded: a dense transform would otherwise @@ -312,33 +429,36 @@ def select_indices( n = parms.size frozen_mask = np.asarray(frozen_mask, dtype=bool) - if expressions: + if not expressions: + sel = (np.asarray(cw) == 0.0) & ~frozen_mask + logger.info( + "No --preconditionParams given; defaulting to a single block of the " + "unconstrained parameters (constraint weight 0)." + ) + return [("unconstrained", np.where(sel)[0])] + + # NB explicit None checks: groups/group_idxs arrive as numpy arrays, + # for which `groups or []` raises on the truth-value test. + gnames = [] if groups is None else list(groups) + gidxs = [] if group_idxs is None else list(group_idxs) + by_group = { + (k.decode() if isinstance(k, bytes) else str(k)): v + for k, v in zip(gnames, gidxs) + } + + out = [] + for expr in expressions: sel = np.zeros(n, dtype=bool) - leftover = [] - # NB explicit None checks: groups/group_idxs arrive as numpy arrays, - # for which `groups or []` raises on the truth-value test. - gnames = [] if groups is None else list(groups) - gidxs = [] if group_idxs is None else list(group_idxs) - by_group = { - (k.decode() if isinstance(k, bytes) else str(k)): v - for k, v in zip(gnames, gidxs) - } - for expr in expressions: - if expr in by_group: - sel[np.asarray(by_group[expr], dtype=np.int64)] = True - else: - leftover.append(expr) - if leftover: + if expr in by_group: + sel[np.asarray(by_group[expr], dtype=np.int64)] = True + else: if match_fn is None: raise ValueError("no matcher available for regex selection") - names = match_fn(leftover, parms) - sel |= np.isin(parms, names) - else: - sel = np.asarray(cw) == 0.0 - logger.info( - "No --preconditionParams given; defaulting to the unconstrained " - "parameters (constraint weight 0)." - ) - - sel &= ~frozen_mask - return np.where(sel)[0] + sel |= np.isin(parms, match_fn([expr], parms)) + sel &= ~frozen_mask + idx = np.where(sel)[0] + if idx.size == 0: + logger.warning(f"--preconditionParams '{expr}' matched no parameters") + continue + out.append((expr, idx)) + return out diff --git a/tests/test_preconditioner.py b/tests/test_preconditioner.py index 24a42be..b872a54 100644 --- a/tests/test_preconditioner.py +++ b/tests/test_preconditioner.py @@ -15,7 +15,11 @@ from rabbit import fitter, inputdata from rabbit.param_models.helpers import load_model -from rabbit.preconditioner import Preconditioner, select_indices +from rabbit.preconditioner import ( + Preconditioner, + auto_blocks, + select_index_blocks, +) from .test_sparse_fit import check_results, make_options, make_test_tensor @@ -144,8 +148,8 @@ def test_explicit_inverse_and_triangular_fallback_agree(): idx = np.arange(1, 8) fast = Preconditioner.from_hessian(h, np.zeros(n), idx, ridge=0.0) slow = Preconditioner.from_hessian(h, np.zeros(n), idx, ridge=0.0) - assert fast._linv is not None - slow._linv = None # force the triangular-solve path + assert fast.blocks[0].linv is not None + slow.blocks[0].linv = None # force the triangular-solve path rng = np.random.default_rng(12) v = rng.normal(size=n) @@ -188,20 +192,21 @@ def test_empty_block_falls_back_to_identity(): # -- unit tests: scope selection ---------------------------------------- -def test_default_scope_is_the_unconstrained_parameters(): +def test_default_scope_is_one_block_of_unconstrained_parameters(): parms = np.array(["poi", "a", "b", "c"]) cw = np.array([0.0, 1.0, 0.0, 1.0]) frozen = np.zeros(4, dtype=bool) - idx = select_indices(parms, cw, frozen) - np.testing.assert_array_equal(idx, [0, 2]) + blocks = select_index_blocks(parms, cw, frozen) + assert len(blocks) == 1 + np.testing.assert_array_equal(blocks[0][1], [0, 2]) def test_frozen_parameters_are_always_excluded(): parms = np.array(["poi", "a", "b", "c"]) cw = np.zeros(4) frozen = np.array([False, True, False, False]) - idx = select_indices(parms, cw, frozen) - np.testing.assert_array_equal(idx, [0, 2, 3]) + blocks = select_index_blocks(parms, cw, frozen) + np.testing.assert_array_equal(blocks[0][1], [0, 2, 3]) def test_selection_by_regex_and_by_group(): @@ -214,10 +219,12 @@ def match_fn(exprs, names): return [n for n in names if any(re.fullmatch(e, n) for e in exprs)] - idx = select_indices(parms, cw, frozen, expressions=["eff_.*"], match_fn=match_fn) - np.testing.assert_array_equal(idx, [1, 2]) + blocks = select_index_blocks( + parms, cw, frozen, expressions=["eff_.*"], match_fn=match_fn + ) + np.testing.assert_array_equal(blocks[0][1], [1, 2]) - idx = select_indices( + blocks = select_index_blocks( parms, cw, frozen, @@ -226,7 +233,15 @@ def match_fn(exprs, names): groups=["mygroup"], group_idxs=[[3]], ) - np.testing.assert_array_equal(idx, [3]) + np.testing.assert_array_equal(blocks[0][1], [3]) + + # one block per expression -> block-diagonal transform + blocks = select_index_blocks( + parms, cw, frozen, expressions=["eff_.*", "other"], match_fn=match_fn + ) + assert [b[0] for b in blocks] == ["eff_.*", "other"] + np.testing.assert_array_equal(blocks[0][1], [1, 2]) + np.testing.assert_array_equal(blocks[1][1], [3]) # -- the invariance test ------------------------------------------------- @@ -384,3 +399,173 @@ def test_fit_is_invariant_on_an_ill_conditioned_block(method): assert pc.cond_before > 1e3 assert pc.cond_after < 1e2 assert check_results("plain", plain, "preconditioned", pre) + + +# -- several blocks ------------------------------------------------------ + + +def test_two_blocks_are_factorised_independently(): + """Block-diagonal transform: each block whitened, no cross-block mixing.""" + n = 12 + h = np.zeros((n, n)) + h[:6, :6] = _spd(6, seed=21, cond=1e6) + h[6:, 6:] = _spd(6, seed=22, cond=1e4) + a, b = np.arange(0, 6), np.arange(6, 12) + pc = Preconditioner.from_hessian(h, np.zeros(n), [("a", a), ("b", b)], ridge=0.0) + assert pc.enabled and pc.n_blocks == 2 and pc.nblock == 12 + hy = pc.hess_to_internal(h) + np.testing.assert_allclose(hy[np.ix_(a, a)], np.eye(6), atol=1e-8) + np.testing.assert_allclose(hy[np.ix_(b, b)], np.eye(6), atol=1e-8) + + +def test_a_singular_block_does_not_disable_the_others(): + """The reason for blocks: a union can be unusable while its parts are fine.""" + n = 10 + h = np.zeros((n, n)) + h[:5, :5] = _spd(5, seed=23) + # second block is entirely degenerate + good, bad = np.arange(0, 5), np.arange(5, 10) + pc = Preconditioner.from_hessian(h, np.zeros(n), [("good", good), ("bad", bad)]) + assert pc.enabled, "the usable block must survive" + assert pc.n_blocks == 1 and pc.nblock == 5 + np.testing.assert_array_equal(pc.blocks[0].idx, good) + + +def test_all_blocks_unusable_falls_back_to_identity(): + n = 6 + pc = Preconditioner.from_hessian( + np.zeros((n, n)), np.zeros(n), [("x", np.arange(3)), ("y", np.arange(3, 6))] + ) + assert not pc.enabled + + +def test_block_diagonal_transform_is_still_a_reparameterisation(): + """Round trip must hold with several blocks.""" + n = 14 + h = np.zeros((n, n)) + h[:7, :7] = _spd(7, seed=24) + h[7:, 7:] = _spd(7, seed=25) + pc = Preconditioner.from_hessian( + h, np.linspace(-1, 1, n), [("a", np.arange(7)), ("b", np.arange(7, 14))] + ) + rng = np.random.default_rng(26) + y = rng.normal(size=n) + np.testing.assert_allclose(pc.from_physical(pc.to_physical(y)), y, atol=1e-9) + + +# -- automatic block discovery ------------------------------------------- + + +def _block_diag_hessian(sizes, seed=0, cond=1e5, coupling=0.0): + """Block-diagonal SPD matrix, optionally with weak inter-block coupling.""" + rng = np.random.default_rng(seed) + n = sum(sizes) + H = np.zeros((n, n)) + at = 0 + for i, m in enumerate(sizes): + H[at : at + m, at : at + m] = _spd(m, seed=seed + i, cond=cond) + at += m + if coupling: + # small symmetric off-block perturbation, kept well below the diagonal + P = rng.normal(size=(n, n)) * coupling * np.mean(np.diag(H)) + P = 0.5 * (P + P.T) + mask = np.ones((n, n), dtype=bool) + at = 0 + for m in sizes: + mask[at : at + m, at : at + m] = False + at += m + H = H + P * mask + return H + + +def test_auto_blocks_recovers_a_known_block_structure(): + """The point of the feature: find the clusters without being told them.""" + sizes = [6, 6, 6, 6, 6] + H = _block_diag_hessian(sizes, seed=31, coupling=0.0) + blocks = auto_blocks(H, np.arange(sum(sizes)), threshold=0.1) + assert len(blocks) == len(sizes) + found = sorted(sorted(idx.tolist()) for _, idx in blocks) + expect = [] + at = 0 + for m in sizes: + expect.append(list(range(at, at + m))) + at += m + assert found == sorted(expect) + + +def test_auto_blocks_ignores_weak_coupling_but_merges_strong(): + sizes = [5, 5] + H = _block_diag_hessian(sizes, seed=32, coupling=1e-4) + assert len(auto_blocks(H, np.arange(10), threshold=0.1)) == 2 + # a strong link between the two groups must merge them + H2 = H.copy() + scale = np.sqrt(H2[0, 0] * H2[5, 5]) + H2[0, 5] = H2[5, 0] = 0.8 * scale + assert len(auto_blocks(H2, np.arange(10), threshold=0.1)) == 1 + + +def test_auto_blocks_percolates_at_a_low_threshold(): + """Documented failure mode: below the percolation point it is one block.""" + H = _block_diag_hessian([5, 5, 5], seed=33, coupling=1e-3) + assert len(auto_blocks(H, np.arange(15), threshold=1e-9)) == 1 + + +def test_auto_blocks_isolates_parameters_with_no_curvature(): + """A zero-diagonal parameter cannot be normalised; it must not poison a block.""" + H = _block_diag_hessian([4, 4], seed=34) + H[3, :] = 0.0 + H[:, 3] = 0.0 + blocks = auto_blocks(H, np.arange(8), threshold=0.1) + singleton = [idx for _, idx in blocks if idx.tolist() == [3]] + assert singleton, f"expected index 3 alone, got {[i.tolist() for _, i in blocks]}" + + +def test_auto_blocks_on_empty_scope(): + assert auto_blocks(_spd(4), np.array([], dtype=int)) == [] + + +def test_auto_blocking_is_the_default(): + from types import SimpleNamespace + + from rabbit import fitter as fitter_mod + + o = SimpleNamespace() + f = fitter_mod.Fitter.__new__(fitter_mod.Fitter) + assert getattr(o, "preconditionBlocks", "auto") == "auto" + assert getattr(o, "preconditionBlockThreshold", 0.1) == 0.1 + del f + + +@pytest.mark.parametrize("blocks", ["auto", "expressions", "none"]) +def test_fit_is_invariant_for_either_blocking(blocks): + """Blocking only changes how the transform is grouped, never the answer.""" + with tempfile.TemporaryDirectory() as tmp: + filename = make_polynomial_tensor(tmp, order=6) + plain, _ = _run(filename, "trust-krylov") + pre, pc = _run( + filename, "trust-krylov", precondition=True, preconditionBlocks=blocks + ) + assert pc.enabled + assert check_results("plain", plain, f"precond[{blocks}]", pre) + + +def test_blocks_none_is_a_single_factorisation_over_the_scope(): + """'none' keeps every cross-correlation, at the cost of one big Cholesky.""" + with tempfile.TemporaryDirectory() as tmp: + filename = make_polynomial_tensor(tmp, order=6) + indata_obj = inputdata.FitInputData(filename) + param_model = load_model("Mu", indata_obj) + options = make_options(precondition=True, preconditionBlocks="none") + f = fitter.Fitter(indata_obj, param_model, options) + f.set_nobs(indata_obj.data_obs) + pc = f._build_preconditioner() + assert pc.enabled + assert pc.n_blocks == 1, f"expected one joint block, got {pc.n_blocks}" + + # and 'auto' on the same model may split it into more than one + options_auto = make_options(precondition=True, preconditionBlocks="auto") + f2 = fitter.Fitter(indata_obj, param_model, options_auto) + f2.set_nobs(indata_obj.data_obs) + pc2 = f2._build_preconditioner() + assert pc2.enabled + assert pc2.nblock == pc.nblock, "same scope, only the grouping differs" From 914b215a6b5beef947a8ab0ad0e5c1649b8d1fdd Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Wed, 19 Aug 2026 10:39:19 -0400 Subject: [PATCH 7/8] Derive the preconditioner ridge from the block's spectrum The ridge escalated by powers of a hundred, 1e-8 -> 1e-6 -> 1e-4 -> 1e-2, and gave up after four tries. That steps straight over the values a mildly indefinite block actually needs. On the combined W+Z fit two blocks holding 168 parameters were skipped as "not factorisable" when their smallest eigenvalues were only 2-3% of max|diag| and nothing was near zero: 1e-4 was too small and 1e-2 was the next rung up. The ridge required to restore definiteness is set by the most negative eigenvalue, so compute it instead of guessing. The caller's value is still tried first, since a positive definite block needs nothing more and the Cholesky is cheap; only if that fails is the extra O(m^3) eigendecomposition worth it. The powers-of-a-hundred escalation stays as a fallback. Measured on that fit: 69 of 71 blocks factorised before, 71 of 71 after, i.e. all 4033 selected parameters are now preconditioned. Splitting the failing blocks further was the alternative and a worse one -- they only became factorisable at a correlation threshold of 0.5, which would have discarded every correlation between 0.1 and 0.5, the bulk of what the transform is for. Also fix the message for a block that really cannot be used: it reported the next ridge it would have tried rather than the largest one actually tried. Co-Authored-By: Claude Opus 5 --- rabbit/preconditioner.py | 37 +++++++++++++++++++++++++++++++----- tests/test_preconditioner.py | 23 ++++++++++++++++++++++ 2 files changed, 55 insertions(+), 5 deletions(-) diff --git a/rabbit/preconditioner.py b/rabbit/preconditioner.py index 1ef9d4b..3f31510 100644 --- a/rabbit/preconditioner.py +++ b/rabbit/preconditioner.py @@ -207,18 +207,45 @@ def _factorise(hess, idx, ridge, max_tries, label=""): return None cond_before = _cond_corr(block) - eps = ridge + # Ridge schedule: try the caller's value first, since for a positive + # definite block that is all that is needed and the Cholesky is cheap. + # Only if that fails is the spectrum worth the extra O(m^3): the ridge + # required to restore definiteness is set by the most negative + # eigenvalue, so it can be computed rather than guessed. Escalating by + # powers of a hundred instead used to overshoot badly -- blocks needing + # 0.03 were skipped after 1e-4 failed and 1e-2 was tried next. + schedule = [ridge] for itry in range(max_tries): + if itry >= len(schedule): + if itry == 1: + w = np.linalg.eigvalsh(block) + lam_min, lam_max = float(w[0]), float(w[-1]) + if lam_min < 0.0: + # enough to make it positive definite, plus a margin so + # the smallest eigenvalue is not left at zero + need = abs(lam_min) + max( + 0.1 * abs(lam_min), 1e-8 * abs(lam_max) + ) + schedule.append(need / scale) + logger.debug( + f"{tag}block has lam_min={lam_min:.3g} " + f"(max|diag|={scale:.3g}); ridge from the spectrum: " + f"{schedule[-1]:.3g} x max|diag|" + ) + else: + schedule.append(max(schedule[-1], 1e-12) * 100.0) + else: + schedule.append(max(schedule[-1], 1e-12) * 100.0) + eps = schedule[itry] trial = block.copy() if eps > 0.0: trial[np.diag_indices_from(trial)] += eps * scale try: chol = scipy.linalg.cholesky(trial, lower=True) except scipy.linalg.LinAlgError: - eps = max(eps, 1e-12) * 100.0 logger.debug( f"Preconditioner Cholesky failed for {tag}block " - f"(try {itry + 1}), raising ridge to {eps:.3g}" + f"with ridge {eps:.3g} x max|diag| (try {itry + 1})" ) continue # Conditioning actually achieved: L^-1 B L^-T for the *un-ridged* @@ -236,8 +263,8 @@ def _factorise(hess, idx, ridge, max_tries, label=""): return Block(idx, chol, cond_before, cond_after, label) logger.warning( - f"Preconditioning block {tag}is not factorisable even with a ridge of " - f"{eps:.3g} x max|diag|; skipping this block." + f"Preconditioning block {tag}is not factorisable; the largest ridge " + f"tried was {max(schedule):.3g} x max|diag|. Skipping this block." ) return None diff --git a/tests/test_preconditioner.py b/tests/test_preconditioner.py index b872a54..34fdb55 100644 --- a/tests/test_preconditioner.py +++ b/tests/test_preconditioner.py @@ -175,6 +175,29 @@ def test_singular_block_falls_back_to_identity(): assert not pc.enabled +def test_mildly_indefinite_block_is_ridged_from_its_spectrum(): + """A block needing a ridge between the old ladder's rungs must still work. + + The escalation used to go 1e-8 -> 1e-6 -> 1e-4 -> 1e-2, so a block whose + smallest eigenvalue sat at a few percent of max|diag| was skipped: 1e-4 was + too small and 1e-2 was tried only after. Two such blocks were dropped from a + real fit that way. The ridge is now computed from lam_min instead. + """ + n = 40 + h = _spd(n, seed=41, cond=1e3) + w, V = np.linalg.eigh(h) + # push a few eigenvalues to ~-3% of the largest diagonal entry + target = -0.03 * np.max(np.diag(h)) + w[:5] = target + h = V @ np.diag(w) @ V.T + h = 0.5 * (h + h.T) + assert np.linalg.eigvalsh(h).min() < 0, "test matrix must be indefinite" + + pc = Preconditioner.from_hessian(h, np.zeros(n), [("mild", np.arange(n))]) + assert pc.enabled, "a mildly indefinite block must not be skipped" + assert pc.nblock == n + + def test_rank_deficient_block_is_ridged_into_shape(): n = 5 h = _spd(n, seed=9) From 0709702a75f07fabca4a7770af1049f557b9efc6 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Thu, 20 Aug 2026 12:50:25 -0400 Subject: [PATCH 8/8] Tidy the preconditioner logging The per-block line was written when a run had one block; auto-blocking routinely finds hundreds, so it is now debug and one summary line reports how many blocks were used, over how many parameters, and the condition numbers they started from. Drop the claim that no --preconditionParams means "a single block of the unconstrained parameters": those expressions only set the scope, and how it is split is --preconditionBlocks' business. The corrected line repeats identically on every restart's rebuild and says nothing the summary does not, so it moves to debug as well. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01HrsrmERvJF2Vafn7wE1B3v --- rabbit/preconditioner.py | 27 +++++++++++++++++++++++---- 1 file changed, 23 insertions(+), 4 deletions(-) diff --git a/rabbit/preconditioner.py b/rabbit/preconditioner.py index 3f31510..bd0347e 100644 --- a/rabbit/preconditioner.py +++ b/rabbit/preconditioner.py @@ -183,6 +183,21 @@ def from_hessian(cls, hess, theta_ref, index_blocks, ridge=1e-8, max_tries=4): "running unpreconditioned." ) return cls.identity(theta_ref) + + # One summary rather than a line per block: auto-blocking routinely + # finds hundreds, and the per-block detail is at debug level. + n_req = len(index_blocks) + npar = sum(b.idx.size for b in blocks) + conds = [b.cond_before for b in blocks if b.cond_before is not None] + summary = f"Preconditioned {len(blocks)} of {n_req} block(s), {npar} parameters" + if conds: + summary += ( + f"; correlation condition number median {np.median(conds):.3g}, " + f"worst {max(conds):.3g} -> 1 at the reference point" + ) + if len(blocks) < n_req: + summary += f" ({n_req - len(blocks)} block(s) not factorisable, skipped)" + logger.info(summary) return cls(theta_ref, blocks) @staticmethod @@ -254,7 +269,7 @@ def _factorise(hess, idx, ridge, max_tries, label=""): tb = scipy.linalg.solve_triangular(chol, block, lower=True, trans="N") tb = scipy.linalg.solve_triangular(chol, tb.T, lower=True, trans="N").T cond_after = _cond_corr(tb) - logger.info( + logger.debug( f"Preconditioning {tag}block of {idx.size} parameters from the " f"reference Hessian (ridge {eps:.3g} x max|diag|): correlation " f"condition number {cond_before:.3g} -> {cond_after:.3g} at the " @@ -458,9 +473,13 @@ def select_index_blocks( if not expressions: sel = (np.asarray(cw) == 0.0) & ~frozen_mask - logger.info( - "No --preconditionParams given; defaulting to a single block of the " - "unconstrained parameters (constraint weight 0)." + # debug, not info: this repeats identically on every restart's rebuild, + # and the summary from from_hessian already reports how many parameters + # ended up preconditioned + logger.debug( + f"No --preconditionParams given; selecting all {int(sel.sum())} " + "unconstrained parameters (constraint weight 0). How they are grouped " + "into blocks is set by --preconditionBlocks." ) return [("unconstrained", np.where(sel)[0])]