From aba85793673df6bd02976f21bb00a56c87f7d001 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 17 Aug 2026 09:15:14 -0400 Subject: [PATCH 1/2] feat: per-parameter step scaling for the multi-start gradient searches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ScalerPriorWidth` gives each parameter its own step scale, derived from its prior, applied as a change of variables so the rule steps in `phi = theta / s` while the objective is still evaluated at the physical `theta = s * phi`. Default `ScalerNone`, skipped rather than applied as a multiply by ones, so the default path and its compiled step are both unchanged. WHY THIS FAMILY. Every other search in the library already adapts per-parameter scale: the nested samplers propose in the unit cube, Emcee's stretch move is affine-invariant, Zeus self-tunes per direction, BlackJAXNUTS adapts a MASS MATRIX via window_adaptation, and L-BFGS accumulates curvature. MultiStart {Adam,Lion,ADABelief,Prodigy} is the sole exception, and not by oversight: the Adam rules normalise per-coordinate by GRADIENT magnitude rather than PARAMETER scale, so with m/sqrt(v) ~ 1 every coordinate moves ~lr in PHYSICAL units, and Prodigy adds only a single GLOBAL `d`. On the reference imaging/mge cell prior widths span 40x, so a step sensible for einstein_radius (width 8.0) is a wall-crossing for bulge.centre (width 0.2) — measured at 31.5% of lane-steps clipped (autolens_profiling#131). This is therefore the diagonal preconditioner the family lacks, the analogue of the NUTS mass matrix, not a novel hack. WHAT IT COSTS. NUTS LEARNS its preconditioner from the geometry during warmup; this one is static and prior-derived — cheap, no warmup, no new hazard, but a prior width is a PROXY for the posterior scale and is poor where the likelihood is far tighter than the prior. Documented in the module rather than hidden. THE RULE is the physical extent of a unit-cube-sized step: Uniform -> width; LogUniform -> sqrt(lo*hi)*log(hi/lo) at the median, NOT the physical width, which for LogUniform(1e-4, 1e4) is wrong by a factor of 543; Gaussian and TruncatedGaussian -> sigma, NOT the truncation width, which for ell_comps (sigma 0.3 in a [-1,1] box) overstates the scale six-fold; LogGaussian -> exp(mean)*sigma. Anything degenerate falls back to 1.0 and logs — the vector is a DIVISOR, so a zero would send the whole population to infinity. Normalised to a geometric mean of exactly 1, so only the RATIOS change and an A/B is not confounded with an effective learning-rate change. APPLIED AS A CHANGE OF VARIABLES, not `params += s * updates`. Prodigy estimates `d` from the distance actually travelled (params0, grad_sum), so rescaling its update externally would leave its own estimate inconsistent with the trajectory it believes it took. Composed as `fitness.call(phi * scale)`, JAX's chain rule yields the gradient w.r.t. phi and there is no manual gradient rescaling anywhere. Optimising the density OF phi instead would fold a Jacobian in and move the MAP — silently, with every counter looking healthy — so `test_scaler.py` pins the invariance directly. The map is linear, so an optimum that genuinely sits on a boundary stays finite; that is what killed the logit reparameterisation this replaces. `search_internal` stays entirely in PHYSICAL units (`params` multiplied on write, divided on resume). The scaler does not enter the search identifier, so a scaled and an unscaled arm share an output directory; a file whose units depended on an invisible knob would resume as a silently wrong population rather than as an error. `samples_via_internal_from` reads that array directly and has no scaler to consult. COMPLEMENTARY TO THE CLIPPER, not a replacement. Scaling makes reaching a wall rarer, never impossible, and where the likelihood genuinely prefers a value outside the prior the clipped lane on the bound is the correct MAP under the declared prior. `Clipper.project` takes an optional `scale` and clips against `bounds / s` in the caller's own coordinates, rather than round-tripping the whole batch through physical space every step. VISIBILITY. The derived vector is appended to the `model.info` FILE — not to the `Model.info` property, which has no access to the search and so cannot render a search-dependent block. `search.summary` gains a `Scaler = ...` line under the same discipline as the clipper's: `ScalerNone` emits nothing, so an existing run's summary is unchanged byte for byte. There is no per-step count to report; scaling's effect is read in the clipped lane-step rate. NOT offered on LBFGS, which takes a `clipper` but does not own its step loop — a knob that accepts a value and ignores it is worse than no knob. A test pins that. Verified end-to-end on the real JAX path (the library suite is NumPy-only by house rule, so it does not execute `_fit`): MultiStartProdigy on a 40x-spread model reaches the same optimum to ~1e-9 with scaling on and off, Prodigy's own `d` estimates differ (so the scaler is not an accidental no-op), the stored per-start parameters come back physical, and the clip count falls 2 -> 0. Note when reading logs: under a scaler the reported `d` is in SCALED units, because that is the space the rule steps in. Do not compare it across a scaled and an unscaled arm; compare the clip rate. Issue #1483. Co-Authored-By: Claude Opus 5 --- autofit/__init__.py | 3 + autofit/non_linear/clipper.py | 35 +- autofit/non_linear/paths/directory.py | 25 +- autofit/non_linear/scaler.py | 424 +++++++++++++++++ .../search/mle/multi_start_gradient/search.py | 173 ++++++- autofit/text/text_util.py | 54 ++- .../search/mle/test_multi_start_gradient.py | 35 +- test_autofit/non_linear/test_scaler.py | 436 ++++++++++++++++++ 8 files changed, 1142 insertions(+), 43 deletions(-) create mode 100644 autofit/non_linear/scaler.py create mode 100644 test_autofit/non_linear/test_scaler.py diff --git a/autofit/__init__.py b/autofit/__init__.py index 557ae50ab..02579a881 100644 --- a/autofit/__init__.py +++ b/autofit/__init__.py @@ -83,6 +83,9 @@ from .non_linear.clipper import AbstractClipper from .non_linear.clipper import ClipperNone from .non_linear.clipper import ClipperPriorBox +from .non_linear.scaler import AbstractScaler +from .non_linear.scaler import ScalerNone +from .non_linear.scaler import ScalerPriorWidth from .non_linear.initializer import InitializerBall from .non_linear.initializer import InitializerPrior from .non_linear.initializer import InitializerParamBounds diff --git a/autofit/non_linear/clipper.py b/autofit/non_linear/clipper.py index ddadca834..ceae70aae 100644 --- a/autofit/non_linear/clipper.py +++ b/autofit/non_linear/clipper.py @@ -98,20 +98,30 @@ def bounds_from_model(self, model) -> Tuple[np.ndarray, np.ndarray]: """ @abstractmethod - def project(self, vector, model, xp=np): + def project(self, vector, model, xp=np, scale=None): """ Project ``vector`` onto the prior support. Parameters ---------- vector - A physical parameter vector, either a single ``(n_params,)`` vector or - a batched ``(n_starts, n_params)`` array of them. Broadcasting handles - both, so no ``vmap`` is required of the caller. + A parameter vector, either a single ``(n_params,)`` vector or a + batched ``(n_starts, n_params)`` array of them. Broadcasting handles + both, so no ``vmap`` is required of the caller. Physical unless + ``scale`` is given, in which case it is in scaled coordinates. model The model whose priors define the support. xp The array module, ``numpy`` or ``jax.numpy``. + scale + The per-parameter step scale from a + :mod:`~autofit.non_linear.scaler`, when the caller is stepping in + ``phi = theta / scale`` rather than in physical parameters. The + **bounds** are divided by it, so the projection happens in the + caller's own coordinates and no round-trip through physical space is + needed. Scales are strictly positive, so dividing preserves the + ordering of each ``(lower, upper)`` pair and ``+/-inf`` stay + ``+/-inf``. Returns ------- @@ -134,7 +144,7 @@ def bounds_from_model(self, model) -> Tuple[np.ndarray, np.ndarray]: n = model.prior_count return np.full(n, -np.inf), np.full(n, np.inf) - def project(self, vector, model, xp=np): + def project(self, vector, model, xp=np, scale=None): return vector, xp.zeros_like(vector, dtype=bool) @@ -264,9 +274,22 @@ def bounds_from_model(self, model) -> Tuple[np.ndarray, np.ndarray]: return lower_inset, upper_inset - def project(self, vector, model, xp=np): + def project(self, vector, model, xp=np, scale=None): lower, upper = self.bounds_from_model(model) + if scale is not None: + # Divided AFTER the insets are applied, not before. The inset is + # relative to the box width, so scaling the raw limits first and + # insetting afterwards gives the identical box -- but the half-open + # `strict_epsilon` is ABSOLUTE, and dividing it by the scale would + # shrink the nudge for a large-scale coordinate until it no longer + # lands strictly inside a support that excludes its limit. Insetting + # in physical space and then mapping the finished bounds keeps the + # inset meaning what it says in the space the prior is declared in. + scale = np.asarray(scale, dtype=float) + lower = lower / scale + upper = upper / scale + dtype = getattr(vector, "dtype", None) lower = xp.asarray(lower, dtype=dtype) upper = xp.asarray(upper, dtype=dtype) diff --git a/autofit/non_linear/paths/directory.py b/autofit/non_linear/paths/directory.py index e4ede6874..47b551bf8 100644 --- a/autofit/non_linear/paths/directory.py +++ b/autofit/non_linear/paths/directory.py @@ -483,9 +483,32 @@ def _save_metadata(self, search_name): def _save_model_info(self, model): """ Save the model.info file, which summarises every parameter and prior. + + A search that applies a per-parameter step scale + (:mod:`autofit.non_linear.scaler`) appends its scale vector to the bottom + of the same file. It goes in the FILE rather than into ``model.info`` the + property because the scale is a property of the *search*, which the model + knows nothing about — a search-dependent block cannot be rendered by a + model-only property, and splitting it into a second file would leave the + canonical "what is this model" artefact silently omitting the fact that + the search does not step in the units the file lists. + + Guarded the same way ``_save_model_start_point``'s source is: a search + with no scaler, an older pickled search, or a scaler that declines to + describe itself all fall through to writing ``model.info`` alone. """ + info = model.info + + try: + scaler_info = self.search.scaler.info_from_model(model=model) + except (NotImplementedError, AttributeError): + scaler_info = "" + + if scaler_info: + info += f"\n\n{scaler_info}" + with open_(self.output_path / "model.info", "w+") as f: - f.write(model.info) + f.write(info) if should_output("model_graph") and hasattr(model, "graph_info"): with open_(self.output_path / "model.graph", "w+") as f: diff --git a/autofit/non_linear/scaler.py b/autofit/non_linear/scaler.py new file mode 100644 index 000000000..72c019374 --- /dev/null +++ b/autofit/non_linear/scaler.py @@ -0,0 +1,424 @@ +""" +Per-parameter step scaling for the gradient searches. + +The problem +----------- + +``MultiStartGradient`` steps in **physical** parameter space with a single global +step scale, while prior box widths across one model routinely span orders of +magnitude. On the reference ``imaging/mge`` lens model they span **40x**:: + + mass.einstein_radius UniformPrior(0, 8) width 8.00 + shear.gamma_1/2 UniformPrior(-0.3, 0.3) width 0.60 + bulge.centre_0/1 UniformPrior(-0.1, 0.1) width 0.20 + +A step that is sensible for ``einstein_radius`` is a wall-crossing for +``bulge.centre``. Lanes are not initialised at the edges -- ``_broad_starts`` +draws in the unit cube at ``(0.15, 0.85)``, the middle 70% of every prior -- they +*walk* there, and on that cell clipping fires on 31.5% of lane-steps +(autolens_profiling#131). + +Every other search family already handles this +--------------------------------------------- + +The multi-start gradient family is the **only** search in the library that steps +without adapting to per-parameter scale: + +=========================== ========================================= ============== +search step mechanism adapts scale? +=========================== ========================================= ============== +Nautilus / Dynesty proposals in the unit cube by construction +Emcee affine-invariant stretch move yes, provably +Zeus ensemble slice sampling (``tune=True``) yes, per direction +BlackJAXNUTS ``blackjax.window_adaptation`` yes -- step size AND mass matrix +LBFGS / BFGS quasi-Newton inverse-Hessian estimate partial -- starts from the identity +Drawer independent draws from the priors n/a, no stepping +MultiStart{Adam,...,Prodigy} fixed / global step in physical space **no** +=========================== ========================================= ============== + +The reason this family is the exception is specific, and it is not an oversight +that can be fixed by tuning: Adam-family rules normalise per-coordinate by +**gradient magnitude**, not **parameter scale**. Since ``m_hat / sqrt(v_hat) ~ 1`` +for a consistent gradient, every coordinate moves by roughly ``learning_rate`` in +*physical* units -- equalising precisely the wrong thing when the natural scales +differ 40x. Prodigy inherits that per-coordinate behaviour and adds a single +**global** estimate ``d``, so it cannot repair a per-coordinate problem either. + +A ``Scaler`` is therefore the direct analogue of NUTS's **mass matrix** and +L-BFGS's **inverse Hessian**: a diagonal preconditioner, an entirely established +idea that this one family happens to lack. + +Static, and what that costs +--------------------------- + +NUTS *learns* its preconditioner from the geometry during warmup. A ``Scaler`` +is **static** and derived from the priors, which is a deliberate trade: it costs +one array multiply, needs no warmup, and cannot itself introduce a failure mode. +The price is that **a prior width is a proxy for the posterior scale**, and where +the likelihood is far tighter than the prior the proxy is poor -- the coordinate +is then merely under-scaled rather than mis-scaled, but the benefit shrinks. A +diagonal learned from the observed gradient/step history is the natural successor +if this proves too blunt. + +Why a linear change of variables, and not the unit cube +------------------------------------------------------- + +Stepping in the unit cube would make the step automatically commensurate with +every box width, and it was rejected -- deliberately, and the objections stand: + +- a **logit** reparameterisation sends an optimum that genuinely sits on a + boundary to infinity, and the reference cell demonstrably has such optima + (6 of 16 lanes end pinned to a bound); +- the **inverse-CDF** transform for non-uniform priors has ``dtheta/du -> inf`` + at the cube faces, trading one numerical hazard for another; +- it invalidates every stored benchmark. + +Scaling by a constant diagonal buys the same normalisation and dodges all three: +the map is **linear**, so boundary optima stay finite and no face singularity +exists. + +The invariance that must not be broken +-------------------------------------- + +The search runs in ``phi = theta / s`` and evaluates the objective at +``theta = s * phi`` -- the **physical-space** posterior, at the transformed point. +The Jacobian of a constant diagonal map is constant, so it adds a constant to the +log-density and **cannot move the MAP**. + +Optimising the density **of ``phi``** instead would fold a non-constant Jacobian +into the objective and move the answer -- and it would fail *silently*, since +every counter and diagnostic would look healthy. That is why the objective is +composed as ``fitness(s * phi)`` rather than the prior being re-expressed, and +why ``test_scaler.py`` pins the invariance directly. + +The scale must also be applied as a **change of variables**, never as a post-hoc +``params += s * updates``: Prodigy estimates its step scale from the distance +actually travelled (``params0``, ``grad_sum``), so rescaling its update +externally would leave its own estimate inconsistent with the trajectory it +believes it took. Run the rule in ``phi`` and it needs no changes at all. + +Composition with ``Clipper`` +---------------------------- + +A ``Scaler`` and a :mod:`~autofit.non_linear.clipper` are **complementary, not +alternatives**. Scaling treats the cause -- it reduces how often a lane reaches a +wall. Clipping guarantees the invariant: scaling makes overshoot rarer, never +impossible (a large gradient, bad curvature, or a grown step scale can still +cross), and without clipping that failure is silent. And where the likelihood +genuinely prefers a value outside the prior, **the clipped lane sitting on the +bound is the correct MAP answer under the declared prior** -- only clipping can +express that. ``AbstractClipper.project`` accepts the scale so it can clip in +``phi`` against ``bounds / s``. +""" + +import logging +from abc import ABC, abstractmethod +from typing import List, NamedTuple, Tuple + +import numpy as np + +from autofit.mapper.prior.gaussian import GaussianPrior +from autofit.mapper.prior.log_gaussian import LogGaussianPrior +from autofit.mapper.prior.log_uniform import LogUniformPrior +from autofit.mapper.prior.truncated_gaussian import TruncatedGaussianPrior +from autofit.mapper.prior.uniform import UniformPrior + +logger = logging.getLogger(__name__) + + +class ScaleTerm(NamedTuple): + """ + One parameter's contribution to the scale vector, kept in a struct so that + :meth:`AbstractScaler.scale_from_model` and + :meth:`AbstractScaler.info_from_model` are computed from **one** source + rather than deriving the numbers twice and drifting apart -- the same + discipline ``ClipperPriorBox._limits_from_model`` exists for. + + Parameters + ---------- + basis + The name of the prior quantity the scale was read from (``"width"``, + ``"sigma"``, ``"median x log-ratio"``, ``"fallback"``). Reported to the + user, because *which* quantity was used is the part of this that is a + judgement call. + basis_value + The value of that quantity. + scale + The un-normalised scale. + """ + + basis: str + basis_value: float + scale: float + + +class AbstractScaler(ABC): + """ + A strategy supplying a per-parameter step scale for a search that steps in + physical parameter space. + + Subclasses are pluggable per-search, exactly like + :class:`~autofit.non_linear.clipper.AbstractClipper`, and the default is + :class:`ScalerNone` so behaviour is unchanged until a user opts in. + """ + + @abstractmethod + def scale_from_model(self, model) -> np.ndarray: + """ + The ``(n_params,)`` scale vector, in **physical parameter order** (that + is, ``model.priors_ordered_by_id`` -- the same order + :meth:`~autofit.non_linear.clipper.AbstractClipper.bounds_from_model` + returns, so the two compose without a reindex). + + Every entry is strictly positive and finite. That is a guarantee, not an + expectation: the vector is used as a divisor, so a zero would send the + whole population to infinity and a ``NaN`` would poison every lane. + """ + + def info_from_model(self, model) -> str: + """ + A human-readable summary of the scale vector, appended to the written + ``model.info`` file. + + The default is empty, which is what causes nothing to be appended for a + search that does no scaling. + """ + return "" + + +class ScalerNone(AbstractScaler): + """ + The no-op scaler, and **the default**. + + The scale is all ones, so a search configured with this is bit-identical to + one with no scaler concept at all. Searches should test for it and skip the + change of variables entirely rather than multiplying by ones, so that the + compiled step is unchanged too. + """ + + def scale_from_model(self, model) -> np.ndarray: + return np.ones(model.prior_count, dtype=float) + + +class ScalerPriorWidth(AbstractScaler): + """ + A static diagonal preconditioner derived from the priors. + + The rule, per prior, is **the physical extent of a unit-cube-sized step**: + + - ``UniformPrior(lo, hi)`` -> ``hi - lo``. ``dtheta/du`` is constant and is + exactly the width. + + - ``LogUniformPrior(lo, hi)`` -> ``sqrt(lo * hi) * log(hi / lo)``, the same + quantity evaluated at the prior median, since ``theta(u) = lo * (hi/lo)**u`` + gives ``dtheta/du = theta * log(hi / lo)``. The **physical width is not + usable here** and the error is not marginal: for ``LogUniform(1e-4, 1e4)`` + the width is ``1e4`` against a true local scale of ``18.4``, a factor of + 543. A log-spaced coordinate's natural step is multiplicative, and this is + the linear-space image of it. + + - ``GaussianPrior(mu, sigma)`` -> ``sigma``. The prior is unbounded, so no + width exists. + + - ``TruncatedGaussianPrior(mu, sigma, lo, hi)`` -> ``sigma``, **not** the + truncation width. The two differ by a lot in practice rather than in + principle: ``ell_comps`` is ``sigma=0.3`` inside a ``[-1, 1]`` box, so the + width overstates its scale six-fold. + + - ``LogGaussianPrior(mu, sigma)`` -> ``exp(mu) * sigma``, the Gaussian rule + carried through the exponential at the median (``mu`` and ``sigma`` are in + **natural log** space, so ``exp(mu)`` is the median and ``sigma`` its + fractional spread). + + - anything else, or a non-finite or non-positive result -> ``1.0``, logged. + + A deliberate inconsistency, recorded rather than hidden + ------------------------------------------------------ + + The Uniform rule uses the full ``dtheta/du``; the Gaussian rules use + ``sigma``, which is ``dtheta/du`` at the median divided by ``sqrt(2 * pi)``. + So a Gaussian coordinate is scaled about **2.5x smaller** relative to a + Uniform one than the strict unit-step principle alone would give. That is the + intended behaviour: ``sigma`` is the quantity a user actually reasons about + when they write a Gaussian prior, and matching the constant would inflate + every Gaussian coordinate's step by a factor nobody asked for. It is written + down here because an undocumented factor of 2.5 between two branches of one + rule is exactly the kind of thing that later reads as a bug. + + Normalisation + ------------- + + The vector is normalised so its **geometric mean is exactly 1**. That leaves + the *global* step magnitude untouched and alters only the *ratios*, so an A/B + against an unscaled run measures rescaling alone and is not confounded with + an effective learning-rate change. The geometric mean is the right centre for + a set of quantities that are compared multiplicatively; the arithmetic mean + would be dominated by the single widest prior. + """ + + def _terms_from_model(self, model) -> List[ScaleTerm]: + """ + One :class:`ScaleTerm` per prior, in ``model.priors_ordered_by_id`` order, + **before** normalisation. + + Priors are matched by type rather than by reading ``lower_limit`` / + ``upper_limit`` off every one of them (which is how + ``ClipperPriorBox._limits_from_model`` works). The two want different + things: the clipper wants *the support*, which every prior exposes + uniformly, whereas this wants *the scale*, which lives in a different + attribute per prior family and for ``TruncatedGaussianPrior`` is + specifically **not** the support. + """ + terms = [] + + for prior in model.priors_ordered_by_id: + term = self._term_for(prior) + + if not np.isfinite(term.scale) or term.scale <= 0.0: + logger.warning( + f"ScalerPriorWidth: {type(prior).__name__} produced a " + f"non-positive or non-finite scale ({term.scale}); falling " + f"back to 1.0 for this parameter. The step scale for it is " + f"therefore unchanged, not zero." + ) + term = ScaleTerm(basis="fallback", basis_value=term.scale, scale=1.0) + + terms.append(term) + + return terms + + @staticmethod + def _term_for(prior) -> ScaleTerm: + """The scale for a single prior, before the fallback guard.""" + if isinstance(prior, UniformPrior): + width = float(prior.upper_limit) - float(prior.lower_limit) + return ScaleTerm(basis="width", basis_value=width, scale=width) + + if isinstance(prior, LogUniformPrior): + lower = float(prior.lower_limit) + upper = float(prior.upper_limit) + # Guarded before the logs rather than after: `log` of a non-positive + # limit is `nan`/`-inf` and would reach the fallback as a warning + # about arithmetic rather than about the prior, plus a NumPy + # RuntimeWarning on the way. A LogUniformPrior with a non-positive + # limit is malformed, not merely awkward to scale. + if lower <= 0.0 or upper <= 0.0: + return ScaleTerm(basis="log-ratio", basis_value=np.nan, scale=np.nan) + log_ratio = np.log(upper / lower) + median = np.sqrt(lower * upper) + return ScaleTerm( + basis="median x log-ratio", + basis_value=float(log_ratio), + scale=float(median * log_ratio), + ) + + # TruncatedGaussianPrior is checked before GaussianPrior purely for + # readability -- they are sibling subclasses of `Prior`, not parent and + # child, so the order is not load-bearing today. It is written this way + # so that it does not silently become load-bearing if that ever changes. + if isinstance(prior, TruncatedGaussianPrior): + sigma = float(prior.sigma) + return ScaleTerm(basis="sigma", basis_value=sigma, scale=sigma) + + if isinstance(prior, GaussianPrior): + sigma = float(prior.sigma) + return ScaleTerm(basis="sigma", basis_value=sigma, scale=sigma) + + if isinstance(prior, LogGaussianPrior): + sigma = float(prior.sigma) + median = float(np.exp(prior.mean)) + return ScaleTerm( + basis="median x sigma", + basis_value=sigma, + scale=median * sigma, + ) + + logger.warning( + f"ScalerPriorWidth: no scale rule for {type(prior).__name__}; using " + f"1.0 for this parameter." + ) + return ScaleTerm(basis="fallback", basis_value=np.nan, scale=1.0) + + @staticmethod + def _normalised(scales: np.ndarray) -> np.ndarray: + """ + ``scales`` divided by their geometric mean, so the returned vector has a + geometric mean of exactly 1. + + Computed in log space (``exp(mean(log(s)))``) rather than as + ``prod(s) ** (1/n)``: the product of fifteen widths spanning 40x + underflows or overflows for no reason, whereas the sum of their logs does + not. Every entry is guaranteed positive and finite by + :meth:`_terms_from_model` before this is called. + """ + return scales / np.exp(np.mean(np.log(scales))) + + def scale_from_model(self, model) -> np.ndarray: + scales = np.array( + [term.scale for term in self._terms_from_model(model)], dtype=float + ) + + if scales.size == 0: + return scales + + return self._normalised(scales) + + def terms_and_scale_from_model(self, model) -> Tuple[List[ScaleTerm], np.ndarray]: + """ + The per-prior terms alongside the normalised scale vector, for callers + that want to report *why* each scale is what it is (the ``model.info`` + block) without recomputing it. + """ + terms = self._terms_from_model(model) + scales = np.array([term.scale for term in terms], dtype=float) + + if scales.size == 0: + return terms, scales + + return terms, self._normalised(scales) + + def info_from_model(self, model) -> str: + """ + The ``model.info`` block: one line per parameter giving its prior type, + the quantity the scale was read from, and the normalised scale. + + Formatted as aligned ``path ...`` lines, matching how the initializer + already reports start points into ``model.start`` + (:meth:`~autofit.non_linear.initializer.InitializerParamBounds.info_from_model`) + rather than as a ``TextFormatter`` tree, because this is a flat + per-parameter table and the tree adds nothing to it. + """ + terms, scales = self.terms_and_scale_from_model(model) + + if len(terms) == 0: + return "" + + rows = [] + for prior, term, scale in zip(model.priors_ordered_by_id, terms, scales): + rows.append( + ( + ".".join(str(path) for path in model.path_for_prior(prior)), + type(prior).__name__, + term.basis, + f"{term.basis_value:.4e}", + f"{scale:.4e}", + ) + ) + + widths = [max(len(row[column]) for row in rows) for column in range(4)] + + info = f"Per-Parameter Step Scaling ({type(self).__name__})\n\n" + for path, prior_name, basis, basis_value, scale in rows: + info += ( + f"{path:<{widths[0]}} {prior_name:<{widths[1]}} " + f"{basis:<{widths[2]}} {basis_value:>{widths[3]}} " + f"scale {scale}\n" + ) + + info += ( + f"\nScales are normalised so their geometric mean is 1.0, which " + f"changes only the RATIOS between parameters and leaves the global " + f"step magnitude unchanged.\n" + ) + + return info diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index 7cbaa3784..4b4443ab6 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -11,6 +11,7 @@ from autofit.non_linear.analysis import Analysis from autofit.non_linear.fitness import Fitness from autofit.non_linear.clipper import AbstractClipper, ClipperNone +from autofit.non_linear.scaler import AbstractScaler, ScalerNone from autofit.non_linear.initializer import AbstractInitializer from autofit.non_linear.samples.sample import Sample from autofit.non_linear.samples.samples import Samples @@ -48,6 +49,7 @@ def __init__( iterations_per_log: int = 10, initializer: Optional[AbstractInitializer] = None, clipper: Optional[AbstractClipper] = None, + scaler: Optional[AbstractScaler] = None, reset_momentum_on_clip: bool = False, iterations_per_full_update: int = None, iterations_per_quick_update: int = None, @@ -134,6 +136,31 @@ def __init__( enforcement on moves where the search converges and would shift every stored multi-start benchmark, so the default flip is deliberately a separate, re-baselined change. + scaler + A per-parameter step scale derived from the priors, applied as a + change of variables so the rule steps in ``phi = theta / scale`` while + the objective is still evaluated at the physical ``theta = scale * + phi`` (see :mod:`autofit.non_linear.scaler`). + + It exists because this rule steps in **physical** space with a single + global step scale, while prior widths across one model routinely span + orders of magnitude (40x on the reference lens cell), so a step that + is sensible for the widest parameter is a wall-crossing for the + narrowest. Every other search family already handles this — the + nested samplers propose in the unit cube, Emcee is affine-invariant, + NUTS adapts a mass matrix, L-BFGS accumulates curvature — and this + family is the sole exception, because the Adam rules normalise by + *gradient* magnitude rather than *parameter* scale and Prodigy's + estimate is a single global scalar. + + Default ``ScalerNone`` — a no-op that is skipped rather than applied + as a multiply by ones, so the default path and its compiled step are + both unchanged. + + Complementary to ``clipper``, not a replacement for it: scaling makes + reaching a wall rarer, never impossible, and where the likelihood + genuinely prefers a value outside the prior only clipping can express + the resulting MAP. resurrect Restart-on-death. When ``True``, any start whose objective goes non-finite is redrawn each step (fresh params from the start band + @@ -242,6 +269,12 @@ def __init__( self.start_upper_limit = start_upper_limit self.resurrect = resurrect self.seed = seed + # Resolved here rather than on ``AbstractMLE`` beside ``clipper``, + # deliberately. The clipper is shared because LBFGS genuinely uses it + # (it hands the bounds to scipy); a scaler has no meaning for a search + # that does not own its step loop, and hanging an inert knob off LBFGS + # would be a knob that silently does nothing. + self.scaler = scaler or ScalerNone() self.reset_momentum_on_clip = reset_momentum_on_clip self.convergence = ( convergence if convergence is not None else MultiStartGradientConvergence() @@ -574,7 +607,7 @@ def _warn_if_unbatched_exceeds_memory(self, model, analysis): return projected = fixed + self.n_starts * per_start - gb = 1024 ** 3 + gb = 1024**3 self.logger.warning( f"{type(self).__name__} is set to evaluate all " f"{self.n_starts} starts in one jax.vmap (batch_size=None). " @@ -678,10 +711,49 @@ def _fit( # entirely rather than applying an identity every step. has_clipper = not isinstance(self.clipper, ClipperNone) + # Per-parameter step scaling (see ``autofit.non_linear.scaler``). The + # step loop runs in ``phi = theta / scale``; everything OUTSIDE it — + # ``_broad_starts``, ``best_params``, ``search_internal`` — stays in + # physical parameters, so the only thing that ever sees ``phi`` is the + # optimizer. Same short-circuit as above: under the default ``ScalerNone`` + # ``scale`` is ``None`` and no multiply is traced at all, so the compiled + # step is byte-for-byte what it was. + has_scaler = not isinstance(self.scaler, ScalerNone) + scale = self.scaler.scale_from_model(model=model) if has_scaler else None + scale_jnp = jnp.asarray(scale) if has_scaler else None + + def _to_physical(vector): + return vector if scale_jnp is None else vector * scale_jnp + + # The objective the OPTIMIZER differentiates. Composed as + # ``fitness.call(phi * scale)``, so what is optimised remains the + # PHYSICAL-space posterior evaluated at the transformed point, and JAX's + # chain rule produces the gradient with respect to ``phi`` for free — + # there is no manual gradient rescaling anywhere in this file, which is + # the entire reason the scale is applied here rather than to the update. + # + # Optimising the density OF ``phi`` instead would fold a Jacobian into the + # objective and move the MAP, and it would do so SILENTLY: every counter, + # every diagnostic and the wall time would all look healthy. For a + # constant diagonal map the Jacobian is a constant, so composing this way + # provably cannot move the answer — ``test_scaler.py`` pins that. + # + # ``_value_and_grad`` itself is deliberately left as the physical + # objective: ``_broad_starts`` below jits it directly to filter candidate + # draws, and those draws are, and stay, physical. + _value_and_grad_stepped = ( + jax.value_and_grad(lambda phi: fitness.call(_to_physical(phi))) + if has_scaler + else _value_and_grad + ) + def _value_and_grad_finite(vector): - fom, grad = _value_and_grad(vector) + fom, grad = _value_and_grad_stepped(vector) violation = ( - model.model_constraint_from_vector(vector, xp=jnp) + # Evaluated at the PHYSICAL vector: a declared model constraint is + # a statement about the model's parameters, not about whatever + # coordinates the optimizer happens to be stepping in. + model.model_constraint_from_vector(_to_physical(vector), xp=jnp) if has_constraint else jnp.asarray(0.0) ) @@ -703,9 +775,7 @@ def batched_value_and_grad(params): for lo, hi, pad in _chunk_slices(params.shape[0], batch_size): chunk = params[lo:hi] if pad: - chunk = jnp.concatenate( - [chunk, jnp.tile(chunk[-1:], (pad, 1))] - ) + chunk = jnp.concatenate([chunk, jnp.tile(chunk[-1:], (pad, 1))]) foms, grads, grad_finite, violation = _vmapped(chunk) if pad: foms = foms[:-pad] @@ -732,7 +802,15 @@ def batched_value_and_grad(params): try: search_internal = self.paths.load_search_internal() + # ``search_internal["params"]`` is PHYSICAL, always — see the write + # site below for why — so it is mapped into the stepped coordinates + # here rather than being trusted to already be in them. That is what + # lets a run written without a scaler resume WITH one, and the other + # way round, instead of the file's meaning depending on a search knob + # that does not enter the identifier. params = jnp.asarray(search_internal["params"]) + if has_scaler: + params = params / scale_jnp opt_state = optax.tree_utils.tree_get(search_internal, "opt_state") best_params = np.asarray(search_internal["best_params"]) best_fom = float(search_internal["best_fom"]) @@ -792,17 +870,28 @@ def batched_value_and_grad(params): if not self.silence: self.logger.info(self._compile_message(batched=False)) + # Drawn PHYSICAL and mapped into the stepped coordinates afterwards, + # so the draw band, the finite-gradient filter and the RNG stream are + # all untouched by scaling — the starting POPULATION is identical + # with and without it, which is what makes an A/B measure the + # stepping and nothing else. params = self._broad_starts( model=model, value_and_grad_single=jax.jit(_value_and_grad), jnp=jnp, ) + best_params = np.asarray(params[0]) + + if has_scaler: + params = params / scale_jnp + # Per-start optimizer state: one independent state per start, so # learning-rate-free rules never share a global scalar estimate. + # Built AFTER the rescale, so a rule that seeds its state from the + # initial parameters (Prodigy anchors ``params0`` there) anchors it in + # the coordinates it will actually step in. opt_state = jax.vmap(optimizer.init)(params) - - best_params = np.asarray(params[0]) best_fom = np.inf fom_history = [] alive_history = [] @@ -887,7 +976,14 @@ def batched_value_and_grad(params): best_index = int(np.argmin(foms_np)) if foms_np[best_index] < best_fom: best_fom = float(foms_np[best_index]) + # Mapped back to PHYSICAL on capture, not on the way out. + # ``best_params`` is read by ``samples_via_internal_from`` and + # by the resume path, neither of which knows a scaler exists; + # converting at every read instead of once here is how one of + # them eventually gets missed. best_params = np.asarray(params[best_index]) + if has_scaler: + best_params = best_params * scale fom_history.append(best_fom) @@ -921,6 +1017,7 @@ def batched_value_and_grad(params): jax=jax, jnp=jnp, rng=resurrect_rng, + scale=scale, ) updates, opt_state = step_update(grads, opt_state, params, foms) @@ -935,9 +1032,16 @@ def batched_value_and_grad(params): # # ``project`` broadcasts the ``(n_params,)`` bounds against the # ``(n_starts, n_params)`` batch, so no vmap is needed here. + # + # ``scale`` is handed through rather than the parameters being + # round-tripped to physical and back: the bounds are a constant + # ``(n_params,)`` vector divided once, whereas the round trip + # would multiply and divide the whole ``(n_starts, n_params)`` + # batch every step and reintroduce floating-point drift into + # coordinates the clipper had just placed exactly on a bound. if has_clipper: params, clipped_mask = self.clipper.project( - vector=params, model=model, xp=jnp + vector=params, model=model, xp=jnp, scale=scale ) # Per-LANE, not per-coordinate: a lane clipped in three # parameters at once is one clipped lane-step, matching how @@ -976,6 +1080,14 @@ def batched_value_and_grad(params): # logging step: it is an (n_starts,) device->host copy, which # merely pulls forward the sync the next iteration's # ``np.asarray(foms)`` would force anyway. + # + # Under a scaler this ``d`` is in SCALED units, because that + # is the space the rule is stepping in. It is reported as the + # rule's own estimate rather than converted, since there is no + # single physical value to convert it to — one ``d`` now spans + # a whole vector of physical step sizes, which is the point of + # the feature. Do not compare ``d`` across a scaled and an + # unscaled arm; compare the clip rate instead. estim_lr = optax.tree_utils.tree_get(opt_state, "estim_lr") self.logger.info( @@ -998,8 +1110,20 @@ def batched_value_and_grad(params): elif total_steps >= self.n_steps: stop_reason = "max_steps" + # ``params`` is written back in PHYSICAL parameters even when the + # search stepped in scaled ones. The scaler does not enter the search + # identifier, so the same output directory can be written by a scaled + # run and read by an unscaled one; a file whose units depended on a + # knob that is invisible to the identifier would resume as a silently + # wrong population rather than as an error. ``samples_via_internal_from`` + # reads this array directly for the per-start parameters too, and it + # has no scaler to consult. The scale vector rides along so a reader + # can see what was used without inferring it. search_internal = { - "params": np.asarray(params), + "params": ( + np.asarray(params) * scale if has_scaler else np.asarray(params) + ), + "scale": scale, "opt_state": opt_state, "best_params": best_params, "best_fom": best_fom, @@ -1030,7 +1154,9 @@ def batched_value_and_grad(params): # this loop just built (carrying ``stop_reason``, ``converged`` and # ``fom_history``) is what ``_fit`` returns and what that final # update is computed from. - if not self._is_final_boundary(converged=converged, total_steps=total_steps): + if not self._is_final_boundary( + converged=converged, total_steps=total_steps + ): self.perform_update( model=model, analysis=analysis, @@ -1113,7 +1239,7 @@ def step_update(grads, opt_state, params, values): return optimizer, step_update def _reinit_dead_starts( - self, params, opt_state, dead_idx, model, optimizer, jax, jnp, rng + self, params, opt_state, dead_idx, model, optimizer, jax, jnp, rng, scale=None ): """ Redraw the dead starts (``dead_idx``) and reinitialise their per-start @@ -1125,6 +1251,14 @@ def _reinit_dead_starts( and merged into the live state pytree with a boolean mask (``jnp.where`` per leaf). ``np.asarray`` of a JAX array is read-only, so the redraw happens in an ``np.array`` copy. + + ``params`` is in whatever coordinates the caller is stepping in. When + ``scale`` is given those are the scaled ones, so the redraw — which is + necessarily physical, since ``vector_from_unit_vector`` returns physical + parameters — is divided by it before it is written back into the row. + Redrawing into the wrong coordinates would place a resurrected lane at a + point the prior never proposed, and it would do so only on lanes that had + already died, which is exactly where nobody looks. """ n = params.shape[0] @@ -1133,9 +1267,10 @@ def _reinit_dead_starts( unit_vector = rng.uniform( self.start_lower_limit, self.start_upper_limit, size=model.prior_count ) - params_np[k] = np.asarray( + redrawn = np.asarray( model.vector_from_unit_vector(unit_vector=list(unit_vector), xp=jnp) ) + params_np[k] = redrawn if scale is None else redrawn / scale params = jnp.asarray(params_np) fresh_state = jax.vmap(optimizer.init)(params) @@ -1338,10 +1473,16 @@ def samples_via_internal_from( # current process's share — the same reasoning as the ``.get`` # defaults above. "clipper": type(self.clipper).__name__, + # Per-parameter step scaling (PyAutoFit#1483), recorded by NAME for + # the same reasons as the clipper. It is recorded ALWAYS, including + # the ``ScalerNone`` default, because neither knob enters the search + # identifier: two arms differing only in the scaler write to the same + # output directory, and a result file that omitted the key when the + # feature was off would be indistinguishable from one written before + # the feature existed. + "scaler": type(self.scaler).__name__, "reset_momentum_on_clip": self.reset_momentum_on_clip, - "n_clipped_lane_steps": int( - search_internal.get("n_clipped_lane_steps", 0) - ), + "n_clipped_lane_steps": int(search_internal.get("n_clipped_lane_steps", 0)), # The seed this search's own draws used, so a result file says which # member of a multi-seed study produced it. ``None`` records the # default (historical fixed seeds) rather than being omitted, so a diff --git a/autofit/text/text_util.py b/autofit/text/text_util.py index d3e00adf8..010a14ffc 100644 --- a/autofit/text/text_util.py +++ b/autofit/text/text_util.py @@ -14,7 +14,9 @@ def padding(item, target=6): return f"{prefix}{string}" -def result_max_lh_info_from(max_log_likelihood_sample : List[float], max_log_likelihood : float, model) -> List[str]: +def result_max_lh_info_from( + max_log_likelihood_sample: List[float], max_log_likelihood: float, model +) -> List[str]: """ Output the maximum log likelihood model only, for quick reference. """ @@ -37,8 +39,8 @@ def result_max_lh_info_from(max_log_likelihood_sample : List[float], max_log_lik paths = [] for (_, prior), value in zip( - model.unique_path_prior_tuples, - max_log_likelihood_sample, + model.unique_path_prior_tuples, + max_log_likelihood_sample, ): for path in model.all_paths_for_prior(prior): paths.append((path, value)) @@ -150,9 +152,7 @@ def _clipper_summary_from(samples_info) -> [str]: line = [f"Clipper = {clipper}\n"] if "n_clipped_lane_steps" not in samples_info: - line.append( - "Clipped Lane-Steps = not measured (bounds enforced by scipy)\n" - ) + line.append("Clipped Lane-Steps = not measured (bounds enforced by scipy)\n") return line n_clipped = int(samples_info["n_clipped_lane_steps"]) @@ -167,6 +167,35 @@ def _clipper_summary_from(samples_info) -> [str]: return line +#: The scaler that rescales nothing. A run using it steps in physical parameters, +#: the historical behaviour, and reports no scaling line — see +#: ``_scaler_summary_from``. +_SCALER_NONE = "ScalerNone" + + +def _scaler_summary_from(samples_info) -> [str]: + """ + The ``search.summary`` line naming the per-parameter step scaler, or none. + + Same discipline as ``_clipper_summary_from``: ``ScalerNone`` and a search + predating the ``Scaler`` (no key at all) both emit **nothing**, so the default + path's summary is unchanged byte for byte. + + There is no count to report. Unlike clipping, scaling has no per-step event — + it is a fixed change of variables applied once at fit start, so "how often did + it fire" is not a question with an answer. Its effect is read *indirectly*, in + the clipped lane-step rate above: a scaler that is doing its job drives that + rate down. The vector itself is written to ``model.info``, which is where a + reader who wants the actual numbers should look. + """ + scaler = samples_info.get("scaler") + + if scaler is None or scaler == _SCALER_NONE: + return [] + + return [f"Scaler = {scaler}\n"] + + def search_summary_from_samples(samples) -> [str]: line = [f"Total Samples = {samples.total_samples}\n"] if hasattr(samples, "total_accepted_samples"): @@ -219,6 +248,7 @@ def search_summary_from_samples(samples) -> [str]: line.append(f"Gradient-NaN Lane-Step Rate = {n_grad_nan / lane_steps}\n") line += _clipper_summary_from(samples_info=samples_info) + line += _scaler_summary_from(samples_info=samples_info) if samples.time is not None: line.append(f"Time To Run = {dt.timedelta(seconds=float(samples.time))}\n") @@ -229,10 +259,10 @@ def search_summary_from_samples(samples) -> [str]: def search_summary_to_file( - samples, - log_likelihood_function_time, - filename, - visualization_time=None, + samples, + log_likelihood_function_time, + filename, + visualization_time=None, ): summary = search_summary_from_samples(samples=samples) summary.append( @@ -253,9 +283,7 @@ def search_summary_to_file( pass if visualization_time is not None: - summary.append( - f"Visualization Time (seconds) = {visualization_time}" - ) + summary.append(f"Visualization Time (seconds) = {visualization_time}") frm.output_list_of_strings_to_file(file=filename, list_of_strings=summary) diff --git a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py index 55b1eba79..46af6de3d 100644 --- a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py +++ b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py @@ -385,7 +385,9 @@ def test__reset_clipped_momentum__recurses_through_optax_chain_tuples(): # optax.adam's state arrives wrapped in a chain tuple, so the reset has to # descend through plain tuples as well as NamedTuples. mask = np.array([[True, False]]) - nested = (_AdamLikeState(count=np.array([1]), mu=np.ones((1, 2)), nu=np.ones((1, 2))),) + nested = ( + _AdamLikeState(count=np.array([1]), mu=np.ones((1, 2)), nu=np.ones((1, 2))), + ) out = af.MultiStartAdam._reset_clipped_momentum( opt_state=nested, clipped_mask=mask, jnp=np @@ -402,6 +404,25 @@ def test__dict_round_trip__reset_momentum_on_clip(): assert af.MultiStartAdam().reset_momentum_on_clip is False +def test__dict_round_trip__scaler(): + restored = from_dict(to_dict(af.MultiStartAdam(scaler=af.ScalerPriorWidth()))) + + assert isinstance(restored.scaler, af.ScalerPriorWidth) + # Default is the no-op, so the step loop skips the change of variables + # entirely rather than multiplying by ones and recompiling a different step. + assert isinstance(af.MultiStartAdam().scaler, af.ScalerNone) + + +def test__scaler_is_not_offered_by_searches_that_do_not_own_their_step_loop(): + """ + ``LBFGS`` takes a ``clipper`` (it hands the bounds to scipy) but must NOT take + a ``scaler``: it does not own its step loop, so the knob would silently do + nothing. A knob that accepts a value and ignores it is worse than no knob. + """ + assert "clipper" in inspect.signature(af.LBFGS.__init__).parameters + assert "scaler" not in inspect.signature(af.LBFGS.__init__).parameters + + def test__samples_via_internal_from(): model = af.Model(example.Gaussian) @@ -649,11 +670,11 @@ def samples_for(n_starts, total_steps): @pytest.mark.parametrize( "converged, total_steps, is_final", [ - (False, 50, False), # ordinary intermediate boundary -> update here + (False, 50, False), # ordinary intermediate boundary -> update here (False, 299, False), # still short of the ceiling - (False, 300, True), # ceiling reached - (False, 301, True), # overshoot is still terminal - (True, 50, True), # early convergence, well short of the ceiling + (False, 300, True), # ceiling reached + (False, 301, True), # overshoot is still terminal + (True, 50, True), # early convergence, well short of the ceiling (True, 300, True), ], ) @@ -680,7 +701,7 @@ def test__is_final_boundary(converged, total_steps, is_final): "restored, expected", [ ("converged", "converged"), # must survive: the loop guard reads it - ("max_steps", None), # stale — the run it described is over + ("max_steps", None), # stale — the run it described is over (None, None), ("some_future_reason", None), ], @@ -925,7 +946,7 @@ def test__samples_info__reports_a_cleared_stop_reason_as_unfinished(): # pure Python and lives here, while the JAX measurement it consumes # (`Analysis.batched_memory_bytes`) is exercised in autofit_workspace_test. -GB = 1024 ** 3 +GB = 1024**3 # The 2026-07-30 release failure: XLA reported 85,898,814,480 bytes for a # 48-start interferometer jvp, i.e. 1,789,558,635 bytes per start. diff --git a/test_autofit/non_linear/test_scaler.py b/test_autofit/non_linear/test_scaler.py new file mode 100644 index 000000000..f640ff82a --- /dev/null +++ b/test_autofit/non_linear/test_scaler.py @@ -0,0 +1,436 @@ +import numpy as np +import pytest + +import autofit as af +from autofit import example +from autofit.non_linear.scaler import ScalerNone, ScalerPriorWidth + +# Pure NumPy, deliberately. The scaler's contract is algebraic — what the scale +# vector IS, and that composing the objective through it cannot move the answer — +# and none of that needs JAX. The end-to-end JAX fit under a real scaler is +# validated downstream, matching the discipline stated in +# ``test_multi_start_gradient.py``. + + +def model_from(centre, normalization, sigma): + """ + A flat three-parameter model, one prior per argument. + + All three are required rather than defaulted: ``example.Gaussian`` always has + exactly three priors, so leaving one unset would silently mix a config-default + prior into the scale vector and make the assertions below depend on the + packaged config rather than on the rule under test. + """ + model = af.Model(example.Gaussian) + model.centre = centre + model.normalization = normalization + model.sigma = sigma + return model + + +def scale_of(prior): + """The rule applied to one prior, without the model plumbing in the way.""" + return ScalerPriorWidth._term_for(prior) + + +def test__scaler_none__is_all_ones(): + model = model_from( + centre=af.UniformPrior(lower_limit=0.0, upper_limit=8.0), + normalization=af.UniformPrior(lower_limit=-0.1, upper_limit=0.1), + sigma=af.GaussianPrior(mean=1.0, sigma=2.0), + ) + + scale = ScalerNone().scale_from_model(model=model) + + assert scale == pytest.approx(np.ones(3)) + + +def test__scaler_none__writes_no_model_info_block(): + """ + The empty string is what makes ``_save_model_info`` append nothing at all, so + a default run's ``model.info`` is byte-identical to one written before this + feature existed. + """ + model = model_from( + centre=af.UniformPrior(lower_limit=0.0, upper_limit=8.0), + normalization=af.UniformPrior(lower_limit=-0.3, upper_limit=0.3), + sigma=af.UniformPrior(lower_limit=-0.1, upper_limit=0.1), + ) + + assert ScalerNone().info_from_model(model=model) == "" + + +def test__uniform_prior__scale_is_the_box_width(): + model = model_from( + centre=af.UniformPrior(lower_limit=0.0, upper_limit=8.0), + normalization=af.UniformPrior(lower_limit=-0.3, upper_limit=0.3), + sigma=af.UniformPrior(lower_limit=-0.1, upper_limit=0.1), + ) + + terms = ScalerPriorWidth()._terms_from_model(model=model) + + assert [term.basis for term in terms] == ["width"] * 3 + assert [term.scale for term in terms] == pytest.approx([8.0, 0.6, 0.2]) + + +def test__gaussian_prior__scale_is_sigma_not_a_width(): + assert scale_of(af.GaussianPrior(mean=100.0, sigma=0.5)).basis == "sigma" + assert scale_of(af.GaussianPrior(mean=100.0, sigma=0.5)).scale == pytest.approx(0.5) + assert scale_of(af.GaussianPrior(mean=-3.0, sigma=2.0)).scale == pytest.approx(2.0) + + +def test__truncated_gaussian__scale_is_sigma_NOT_the_truncation_width(): + """ + The distinction this rule exists for. ``ell_comps`` is the real case: sigma + 0.3 inside a [-1, 1] box, so using the truncation width would overstate its + natural scale six-fold and under-step it by the same factor. + """ + term = scale_of( + af.TruncatedGaussianPrior( + mean=0.0, sigma=0.3, lower_limit=-1.0, upper_limit=1.0 + ) + ) + + assert term.basis == "sigma" + assert term.scale == pytest.approx(0.3) + assert term.scale != pytest.approx(2.0) # the truncation width + + +def test__log_uniform_prior__scale_is_median_times_log_ratio_NOT_the_width(): + """ + A log-spaced coordinate's natural step is multiplicative. ``dtheta/du`` for + ``theta(u) = lo * (hi/lo)**u`` is ``theta * log(hi/lo)``, which at the median + ``sqrt(lo*hi)`` is the value asserted here. The physical width is wrong by a + factor of 543 on this prior, which is why it is asserted against explicitly + rather than merely left unused. + """ + term = scale_of(af.LogUniformPrior(lower_limit=1.0e-4, upper_limit=1.0e4)) + + expected = np.sqrt(1.0e-4 * 1.0e4) * np.log(1.0e4 / 1.0e-4) + + assert term.basis == "median x log-ratio" + assert term.scale == pytest.approx(expected) + assert term.scale == pytest.approx(18.42, abs=0.01) + assert term.scale != pytest.approx(1.0e4) # the physical width + + +def test__log_uniform__reduces_to_the_uniform_rule_in_the_narrow_limit(): + """ + A sanity check on the unifying principle rather than on an implementation + detail: for a narrow log-uniform prior the multiplicative step and the linear + one agree, so the two branches of the rule meet rather than merely coexisting. + """ + lower, upper = 1.0, 1.0 + 1.0e-6 + + term = scale_of(af.LogUniformPrior(lower_limit=lower, upper_limit=upper)) + + assert term.scale == pytest.approx(upper - lower, rel=1.0e-5) + + +def test__log_gaussian_prior__scale_is_median_times_sigma(): + term = scale_of(af.LogGaussianPrior(mean=2.0, sigma=0.5)) + + assert term.basis == "median x sigma" + assert term.scale == pytest.approx(np.exp(2.0) * 0.5) + + +def test__geometric_mean_of_the_scale_is_exactly_one(): + """ + The normalisation is what keeps an A/B honest: it changes only the RATIOS, so + the global step magnitude is untouched and the comparison is not confounded + with an effective learning-rate change. + """ + model = model_from( + centre=af.UniformPrior(lower_limit=0.0, upper_limit=8.0), + normalization=af.UniformPrior(lower_limit=-0.3, upper_limit=0.3), + sigma=af.UniformPrior(lower_limit=-0.1, upper_limit=0.1), + ) + + scale = ScalerPriorWidth().scale_from_model(model=model) + + assert np.exp(np.mean(np.log(scale))) == pytest.approx(1.0) + + +def test__normalisation_preserves_every_ratio(): + model = model_from( + centre=af.UniformPrior(lower_limit=0.0, upper_limit=8.0), + normalization=af.UniformPrior(lower_limit=-0.3, upper_limit=0.3), + sigma=af.UniformPrior(lower_limit=-0.1, upper_limit=0.1), + ) + + scaler = ScalerPriorWidth() + raw = np.array([term.scale for term in scaler._terms_from_model(model=model)]) + scale = scaler.scale_from_model(model=model) + + assert scale / scale[0] == pytest.approx(raw / raw[0]) + + # And the 40x spread the feature exists for survives normalisation. + assert scale.max() / scale.min() == pytest.approx(40.0) + + +def test__scale_is_always_positive_and_finite(): + """ + A guarantee rather than an expectation: the vector is used as a DIVISOR, so a + zero sends the whole population to infinity and a NaN poisons every lane. A + degenerate prior must degrade to "unscaled", never to "broken". + """ + # ``UniformPrior`` refuses a zero width at CONSTRUCTION, so the degenerate + # case has to be reached by mutating one afterwards. That is the realistic + # route anyway -- the limits are plain attributes and nothing re-validates + # them -- and it is why the scaler cannot simply assume the priors it is + # handed are well formed. + degenerate = af.UniformPrior(lower_limit=1.0, upper_limit=2.0) + degenerate.upper_limit = 1.0 + + model = model_from( + centre=degenerate, + normalization=af.GaussianPrior(mean=0.0, sigma=0.0), + sigma=af.UniformPrior(lower_limit=0.0, upper_limit=4.0), + ) + + scale = ScalerPriorWidth().scale_from_model(model=model) + + assert np.all(np.isfinite(scale)) + assert np.all(scale > 0.0) + + # Degrades to "unscaled", not to "broken": the two bad coordinates fall back + # to 1.0 BEFORE normalisation, so they cannot drag the geometric mean either. + assert np.exp(np.mean(np.log(scale))) == pytest.approx(1.0) + + +def test__log_uniform_with_a_non_positive_limit__falls_back_rather_than_nan(): + prior = af.LogUniformPrior(lower_limit=1.0e-8, upper_limit=1.0) + prior.lower_limit = -1.0 + + term = scale_of(prior) + if np.isnan(term.scale): + # ``_term_for`` reports the malformed prior; ``_terms_from_model`` is + # where the fallback is applied, so go through it for the final value. + model = model_from( + centre=prior, + normalization=af.UniformPrior(lower_limit=0.0, upper_limit=1.0), + sigma=af.UniformPrior(lower_limit=0.0, upper_limit=1.0), + ) + term = ScalerPriorWidth()._terms_from_model(model=model)[0] + + assert term.basis == "fallback" + assert term.scale == 1.0 + + +def test__unknown_prior_type__falls_back_to_one(): + class _UnscalablePrior(af.UniformPrior): + pass + + term = ScalerPriorWidth._term_for( + _UnscalablePrior(lower_limit=0.0, upper_limit=1.0) + ) + + # A subclass of a KNOWN prior still resolves through isinstance, which is the + # desired behaviour — a user subclassing UniformPrior gets the width rule. + assert term.basis == "width" + + +def test__info_from_model__names_the_basis_used_for_each_prior(): + """ + The block is the user-facing half of this feature, and WHICH quantity each + scale came from is the part that is a judgement call — so it is reported, + not just the resulting number. + """ + model = model_from( + centre=af.UniformPrior(lower_limit=0.0, upper_limit=8.0), + normalization=af.TruncatedGaussianPrior( + mean=0.0, sigma=0.3, lower_limit=-1.0, upper_limit=1.0 + ), + sigma=af.UniformPrior(lower_limit=0.0, upper_limit=4.0), + ) + + info = ScalerPriorWidth().info_from_model(model=model) + + assert "Per-Parameter Step Scaling (ScalerPriorWidth)" in info + assert "centre" in info + assert "UniformPrior" in info + assert "TruncatedGaussianPrior" in info + assert "width" in info + assert "sigma" in info + assert "geometric mean is 1.0" in info + + # One line per parameter, so a 2-parameter model cannot silently report 1. + body = [line for line in info.splitlines() if "scale " in line] + assert len(body) == 3 + + +# --------------------------------------------------------------------------- +# The invariance the whole design turns on. +# --------------------------------------------------------------------------- + + +def _rosenbrock_like(theta): + """ + A deliberately ill-conditioned objective with a unique interior minimum at + ``(3.0, 0.05)`` — chosen so the two coordinates' natural scales differ by the + same order the real model's priors do. + """ + return 5.0 * (theta[0] - 3.0) ** 2 + 8000.0 * (theta[1] - 0.05) ** 2 + + +def _gradient(objective, theta, eps=1.0e-7): + grad = np.zeros_like(theta) + for i in range(theta.size): + step = np.zeros_like(theta) + step[i] = eps + grad[i] = (objective(theta + step) - objective(theta - step)) / (2.0 * eps) + return grad + + +def _descend(objective, start, learning_rate, n_steps): + theta = np.array(start, dtype=float) + for _ in range(n_steps): + theta = theta - learning_rate * _gradient(objective, theta) + return theta + + +def test__objective_composed_through_the_scale_is_the_SAME_objective(): + """ + The algebraic pin. Composing as ``f(s * phi)`` evaluates the physical-space + objective at the transformed point, so it is equal at every corresponding + pair — exactly, not approximately. + + Optimising the density OF ``phi`` instead would fold a Jacobian in and move + the MAP, and it would fail SILENTLY: every counter, every diagnostic and the + wall time would look healthy. This is the test that says which of the two was + implemented. + """ + scale = np.array([2.5, 0.04]) + + rng = np.random.default_rng(0) + for _ in range(50): + theta = rng.uniform(-10.0, 10.0, size=2) + phi = theta / scale + + # ``rel`` rather than exact: ``scale * (theta / scale)`` is not bitwise + # ``theta`` in floating point. The claim under test is that the objective + # is the SAME FUNCTION composed through the map, not that the round trip + # is lossless -- a Jacobian-folding implementation would differ by orders + # of magnitude here, not by an ulp. + assert _rosenbrock_like(scale * phi) == pytest.approx( + _rosenbrock_like(theta), rel=1.0e-12 + ) + + +def test__scaling_cannot_move_the_MAP(): + """ + The end-to-end statement of the same thing: descend on the raw objective, and + descend on the scaled one and map back, and the recovered argmin agrees. + + The scaled run uses the SMALLER learning rate that the unscaled one needs to + remain stable at all, so this is not the scaled run being flattered — it is + the answer being shown to be independent of the coordinates, which is the + property that makes the whole change safe. + """ + scale = np.array([2.5, 0.04]) + start = np.array([8.0, -4.0]) + + unscaled = _descend(_rosenbrock_like, start, learning_rate=1.0e-5, n_steps=200000) + + scaled_phi = _descend( + lambda phi: _rosenbrock_like(scale * phi), + start / scale, + learning_rate=1.0e-5, + n_steps=200000, + ) + scaled = scale * scaled_phi + + assert unscaled == pytest.approx(np.array([3.0, 0.05]), abs=1.0e-4) + assert scaled == pytest.approx(np.array([3.0, 0.05]), abs=1.0e-4) + assert scaled == pytest.approx(unscaled, abs=1.0e-4) + + +def test__scaling_reaches_the_minimum_a_shared_step_size_cannot(): + """ + Why the feature exists at all, stated as a test rather than as a claim. + + One shared step size cannot serve two coordinates whose curvatures differ by + orders of magnitude. It is set by the STIFF coordinate -- push it higher and + that one oscillates or diverges -- which leaves the soft coordinate crawling. + Here the stiff coordinate has converged to machine precision while the soft + one is still visibly short of its optimum after the same 5000 steps. + + This is the same failure the real search shows, in miniature: it is the + NARROW-prior coordinates that reach the walls while the wide ones are still + travelling. Rescaling equalises the two and one rate then serves both. + """ + scale = np.array([2.5, 0.04]) + start = np.array([8.0, -4.0]) + learning_rate = 1.0e-4 + + unscaled = _descend(_rosenbrock_like, start, learning_rate, n_steps=5000) + + scaled = scale * _descend( + lambda phi: _rosenbrock_like(scale * phi), + start / scale, + learning_rate, + n_steps=5000, + ) + + # The stiff coordinate is done; the soft one is not, and it is not close. + assert unscaled[1] == pytest.approx(0.05, abs=1.0e-6) + assert np.abs(unscaled[0] - 3.0) > 1.0e-2 + + # The same budget, the same single learning rate, both coordinates converged. + assert scaled == pytest.approx(np.array([3.0, 0.05]), abs=1.0e-4) + + +# --------------------------------------------------------------------------- +# Composition with the clipper. +# --------------------------------------------------------------------------- + + +def test__clipper_projects_in_scaled_coordinates_against_scaled_bounds(): + model = model_from( + centre=af.UniformPrior(lower_limit=0.0, upper_limit=8.0), + normalization=af.UniformPrior(lower_limit=-0.3, upper_limit=0.3), + sigma=af.UniformPrior(lower_limit=-0.1, upper_limit=0.1), + ) + + clipper = af.ClipperPriorBox() + scale = ScalerPriorWidth().scale_from_model(model=model) + + theta = np.array([9.0, -0.5, 0.4]) # all three outside their boxes + phi = theta / scale + + projected_physical, mask_physical = clipper.project(vector=theta, model=model) + projected_phi, mask_phi = clipper.project(vector=phi, model=model, scale=scale) + + assert (mask_phi == mask_physical).all() + assert scale * projected_phi == pytest.approx(projected_physical) + + +def test__clipper_with_no_scale__is_unchanged(): + model = model_from( + centre=af.UniformPrior(lower_limit=0.0, upper_limit=8.0), + normalization=af.UniformPrior(lower_limit=0.0, upper_limit=1.0), + sigma=af.UniformPrior(lower_limit=0.0, upper_limit=1.0), + ) + + clipper = af.ClipperPriorBox() + vector = np.array([9.0, 0.5, 0.5]) + + assert clipper.project(vector=vector, model=model)[0] == pytest.approx( + clipper.project(vector=vector, model=model, scale=None)[0] + ) + + +def test__clipper_none__accepts_a_scale_and_ignores_it(): + model = model_from( + centre=af.UniformPrior(lower_limit=0.0, upper_limit=8.0), + normalization=af.UniformPrior(lower_limit=0.0, upper_limit=1.0), + sigma=af.UniformPrior(lower_limit=0.0, upper_limit=1.0), + ) + + vector = np.array([9.0, 0.5, 0.5]) + projected, mask = af.ClipperNone().project( + vector=vector, model=model, scale=np.array([2.0, 1.0, 1.0]) + ) + + assert projected == pytest.approx(vector) + assert not mask.any() From 4f629ffd8a6f11c194b600a49702498d2dff57bf Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 17 Aug 2026 09:32:51 -0400 Subject: [PATCH 2/2] docs(scaler): record what one scaled unit does NOT mean, and the clipper-default policy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two clarifications from review, both about honesty rather than behaviour. SCALE DEFINITION. Mixing Uniform -> width with Gaussian -> sigma means one scaled unit does NOT carry a single consistent probabilistic interpretation across coordinates. That is fine for the ENGINEERING goal here (stop a step sized for a wide coordinate crossing a narrow one's wall, for which order-of-magnitude commensurability is enough) but it is not a claim to prior-standardised coordinates, and the module now says so. The unified definitions to reach for if that is ever wanted are named: the prior's actual standard deviation, a robust central quantile width, or the local inverse-CDF derivative dtheta/du at a reference point — the last being the continuous generalisation of what the uniform and log-uniform rules already do. None is implemented, deliberately: if the simple width/sigma version does not move the clip rate, a more principled one will not either, and the diagnosis is wrong. CLIPPER DEFAULT. Reframed from hygiene to constrained-optimizer semantics. A search advertising a posterior with hard prior support is solving a CONSTRAINED problem, so a state outside that support is infeasible, not merely poor. Projection does not alter the constrained optimum, makes the invariant explicit, prevents silent invalid trajectories, stops prior exits masking later pathologies, and is what makes the diagnostics interpretable. Hard-support enforcement is therefore the INTENDED default for the gradient/MLE searches; the reason it is not yet is empirical breadth, not mathematics — everything measured is MultiStartProdigy on one lens cell, and Adam/Lion/ADABelief are unmeasured while being MORE exposed (Adam steps ~lr in physical units in every coordinate, Lion exactly it). When it flips it must not be sold as a long-budget accuracy gain: at 16x3000 the answer does not move at all. The scaler and the clipper stay separate features regardless. Even at a clip rate driven from 31% to 0.1%, projection is retained as the last-line invariant. Issue #1483. Co-Authored-By: Claude Opus 5 --- autofit/non_linear/scaler.py | 75 ++++++++++++++++++++++++++++++++---- 1 file changed, 67 insertions(+), 8 deletions(-) diff --git a/autofit/non_linear/scaler.py b/autofit/non_linear/scaler.py index 72c019374..2f16def72 100644 --- a/autofit/non_linear/scaler.py +++ b/autofit/non_linear/scaler.py @@ -60,6 +60,32 @@ diagonal learned from the observed gradient/step history is the natural successor if this proves too blunt. +What "one scaled unit" does and does not mean +--------------------------------------------- + +:class:`ScalerPriorWidth` mixes ``width`` for uniform priors with ``sigma`` for +Gaussian ones, so **one scaled unit does not carry a single consistent +probabilistic interpretation** across coordinates. That is a real limitation and +it is stated rather than papered over. + +It is acceptable because the goal here is an **engineering** one -- stop a step +that suits a wide coordinate from crossing a narrow coordinate's wall -- and for +that, order-of-magnitude commensurability is what matters. It is *not* a claim to +prior-standardised coordinates. + +If genuinely prior-standardised coordinates are ever wanted, the unified +definitions to reach for, in rough order of sophistication, are: the prior's +actual **standard deviation**; a **robust central quantile width** (e.g. the +inter-quantile range between the 15.9% and 84.1% points, which degrades +gracefully for heavy-tailed and bounded priors alike); or the **local inverse-CDF +derivative** ``dtheta/du`` evaluated at a reference point, which is the +continuous generalisation of what the uniform and log-uniform rules already do. +Each is a strictly better-defined quantity than the mixture above. + +None of them is implemented, deliberately. The simple width/sigma version is the +one to measure first: if it does not move the clip rate, a more principled scale +definition will not either, and the diagnosis behind the whole feature is wrong. + Why a linear change of variables, and not the unit cube ------------------------------------------------------- @@ -101,14 +127,47 @@ ---------------------------- A ``Scaler`` and a :mod:`~autofit.non_linear.clipper` are **complementary, not -alternatives**. Scaling treats the cause -- it reduces how often a lane reaches a -wall. Clipping guarantees the invariant: scaling makes overshoot rarer, never -impossible (a large gradient, bad curvature, or a grown step scale can still -cross), and without clipping that failure is silent. And where the likelihood -genuinely prefers a value outside the prior, **the clipped lane sitting on the -bound is the correct MAP answer under the declared prior** -- only clipping can -express that. ``AbstractClipper.project`` accepts the scale so it can clip in -``phi`` against ``bounds / s``. +alternatives**, and this stays true however well scaling works. Even if scaling +drove the clip rate from 31% to 0.1%, projection would be retained as the +last-line invariant. + +The reason is not hygiene, it is **constrained-optimizer semantics**. A search +that advertises itself as optimising a posterior with hard prior support is +solving a *constrained* problem, and a state outside that support is not an +inconveniently poor candidate -- it is **infeasible**. Projection onto the box: + +- does not alter the constrained optimum; +- makes that invariant explicit rather than incidental; +- prevents silent invalid trajectories; +- prevents prior exits from masking later pathologies; +- makes the search diagnostics interpretable at all. + +Scaling treats the *cause* -- it reduces how often a lane reaches a wall -- and +can only ever make overshoot rarer, never impossible: a large gradient, bad +curvature, or a grown step scale can still cross. Without clipping that crossing +is silent. And where the likelihood genuinely prefers a value outside the prior, +**the clipped lane sitting on the bound is the correct MAP answer under the +declared prior** -- only clipping can express it. + +``AbstractClipper.project`` accepts the scale so it can clip in ``phi`` against +``bounds / s``. + +On the clipper's default +------------------------ + +Hard-support enforcement is the **intended** default for the gradient/MLE +searches on those semantics. It is not the default yet, and the hesitation is +empirical breadth rather than mathematics: everything measured so far is +``MultiStartProdigy`` on a single lens cell, and ``MultiStartAdam`` / +``MultiStartLion`` / ``MultiStartADABelief`` remain unmeasured -- which matters, +because the fixed-rate rules are *more* exposed to the underlying problem, not +less (Adam steps by roughly ``learning_rate`` in physical units in every +coordinate; Lion, being sign-based, by exactly it). + +When the default does flip, it must not be sold as a long-budget accuracy +improvement. It is not one: measured at 16x3000 the answer does not move at all, +and at 105 steps the finite-budget gain was 114 nats. The case is correctness of +the optimisation problem being solved. """ import logging