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..2f16def72 --- /dev/null +++ b/autofit/non_linear/scaler.py @@ -0,0 +1,483 @@ +""" +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. + +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 +------------------------------------------------------- + +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**, 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 +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()