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 a132bec..faab3ef 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -11,8 +11,14 @@ 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.callbacks import ( + RESTART_MIN_IMPROVEMENT, + FitterCallback, + merge_callbacks, +) from rabbit.impacts import ( asym_impacts, global_asym_impacts, @@ -59,48 +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 - - 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 - ): - 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 - - class Fitter: valid_systematic_types = ["log_normal", "normal"] @@ -129,6 +93,17 @@ 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_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". # True / False from programmatic callers are accepted as # aliases for "on" / "off". The tri-state is resolved to the @@ -2264,33 +2239,153 @@ 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. + + 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) + + index_blocks = precond.select_index_blocks( + 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, + ) + # 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_np = self._reference_matrix() + except Exception as ex: + logger.warning( + f"Could not compute the reference Hessian for preconditioning ({ex}); " + "running unpreconditioned." + ) + 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, + index_blocks, + 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. + # 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__(), 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): + pc = pc_cell[0] + 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): + pc = pc_cell[0] + 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__() - - xval = self.x.numpy() + return pc.hess_to_internal(hess.__array__()) - callback = FitterCallback(xval, self.earlyStopping) + # scipy works in internal coordinates throughout; y = 0 at the point the + # transform was built. + xval = pc_cell[0].from_physical(self.x.numpy()) if self.minimizer_method in [ "trust-krylov", @@ -2320,29 +2415,92 @@ def scipy_hess(xval): 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) - - self.x.assign(xval) + 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_cell[0].to_physical(xval)) return callback diff --git a/rabbit/parsing.py b/rabbit/parsing.py index 19ee749..c3f3be7 100644 --- a/rabbit/parsing.py +++ b/rabbit/parsing.py @@ -183,9 +183,27 @@ def common_parser(): ) parser.add_argument( "--earlyStopping", + default=20, + type=int, + 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", default=-1, type=int, - help="Number of iterations with no improvement after which training will be stopped. Specify -1 to disable.", + 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", @@ -202,6 +220,78 @@ 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( + "--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", + type=str, + 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", + 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..bd0347e --- /dev/null +++ b/rabbit/preconditioner.py @@ -0,0 +1,510 @@ +"""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 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. +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 +import scipy.sparse +import scipy.sparse.csgraph +from wums import logging + +logger = logging.child_logger(__name__) + +# Sources for the reference matrix, built by the fitter (see Fitter._reference_matrix). +PRECONDITION_SOURCES = ("hessian", "gaussnewton") + + +class Block: + """One factorised group of parameters: indices, L, and a cached L^-1.""" + + 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 + 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." + ) + + +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 blocks.""" + return cls(theta_ref) + + @property + def enabled(self): + return bool(self.blocks) + + @property + def nblock(self): + """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, 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. + """ + 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 no block could be used; " + "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 + 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 + 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( + f"Preconditioning block {tag}has no positive diagonal " + f"(max diag = {scale}); skipping." + ) + return None + + cond_before = _cond_corr(block) + # 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: + logger.debug( + f"Preconditioner Cholesky failed for {tag}block " + f"with ridge {eps:.3g} x max|diag| (try {itry + 1})" + ) + 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.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 " + "reference point" + ) + return Block(idx, chol, cond_before, cond_after, label) + + logger.warning( + f"Preconditioning block {tag}is not factorisable; the largest ridge " + f"tried was {max(schedule):.3g} x max|diag|. Skipping this block." + ) + return None + + # -- the transform --------------------------------------------------- + + def _apply_T(self, v): + """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) + 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: 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) + 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): + """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 each block.""" + d = np.asarray(theta, dtype=np.float64) - self.theta_ref + if not self.blocks: + return d + out = np.array(d, dtype=np.float64, copy=True) + for b in self.blocks: + out[b.idx] = b.chol.T @ d[b.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. + + 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 not self.blocks: + return np.asarray(hess, dtype=np.float64) + out = np.array(hess, dtype=np.float64, copy=True) + 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 not self.blocks: + return "preconditioning: disabled" + return ( + f"preconditioning: {self.n_blocks} block(s) covering " + f"{self.nblock} 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 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, + expressions=None, + match_fn=None, + groups=None, + group_idxs=None, +): + """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 + 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 not expressions: + sel = (np.asarray(cw) == 0.0) & ~frozen_mask + # 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])] + + # 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) + 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") + 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 new file mode 100644 index 0000000..34fdb55 --- /dev/null +++ b/tests/test_preconditioner.py @@ -0,0 +1,594 @@ +"""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, + auto_blocks, + select_index_blocks, +) + +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.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) + 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_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) + 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_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) + 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]) + 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(): + 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)] + + blocks = select_index_blocks( + parms, cw, frozen, expressions=["eff_.*"], match_fn=match_fn + ) + np.testing.assert_array_equal(blocks[0][1], [1, 2]) + + blocks = select_index_blocks( + parms, + cw, + frozen, + expressions=["mygroup"], + match_fn=match_fn, + groups=["mygroup"], + group_idxs=[[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 ------------------------------------------------- + + +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) + + +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. + + 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) + + +# -- 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" diff --git a/tests/test_restart.py b/tests/test_restart.py new file mode 100644 index 0000000..d517fbc --- /dev/null +++ b/tests/test_restart.py @@ -0,0 +1,208 @@ +"""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.callbacks 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 + + +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. + + 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=n_restarts, precondition=True + ) + f = fitter.Fitter(indata_obj, param_model, options) + f.set_nobs(indata_obj.data_obs) + + built_at = [] + original = f._build_preconditioner + + def counting(): + built_at.append(np.array(f.x.numpy(), copy=True)) + return original() + + f._build_preconditioner = counting + monkeyed = fitter.FitterCallback + fitter.FitterCallback = StallEveryFewIterations + try: + callback = f.fit() + finally: + fitter.FitterCallback = monkeyed + + assert callback is not None + # 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" + ) 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)