From 84e8b8f56ce684ee6b2d5d5ea344b93825bf7c7c Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 24 Aug 2026 17:57:29 -0400 Subject: [PATCH 1/4] Add AbstractBijector: per-parameter log/logit change of variables Extends the linear diagonal AbstractScaler to a per-coordinate bijection (identity / log / logit), vectorised over (n_starts, n_params) and built once per model via from_model() so jax.jit traces a single jnp.where selection tree. BijectorNone is the no-op default (byte-identical: x/1.0 == x exactly); BijectorAuto picks log for LogUniform/LogGaussian priors; BijectorLogit is an explicitly secondary arm (the scaler's unit-cube objections about boundary optima still apply); BijectorPerPath lets a caller declare kinds by model path; BijectorDiagonal adapts an existing AbstractScaler into this framework without changing scaler= behaviour. log_det_jacobian is exposed for a future phi-space sampler (PyAutoFit #1521/#1522) and is never needed for MAP: composing the objective through a bijection relabels points without changing its value set, so no Jacobian may be added without moving the MAP. Compose AbstractClipper.project with a bijector, clipping in transformed space against bijector.bounds_forward(lower_inset, upper_inset) -- valid because every kind is monotone increasing. Fixes a real bug found while wiring this up: the existing physical-relative margin (margin * (upper - lower)) is wrong for a log-kind coordinate by orders of magnitude -- for LogUniform(1e-6, 1e6) the default margin produces a ~1.0 PHYSICAL inset, fencing off virtually the entire support. Log-kind coordinates now get a log-space margin (margin * log(upper/lower)) via ClipperPriorBox._inset_from_model(kinds=...); the no-bijector bounds_from_model path is unaffected. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01EDABYoH6giHXhFJUks8yd6 --- autofit/non_linear/bijector.py | 608 +++++++++++++++++++++++++++++++++ autofit/non_linear/clipper.py | 164 +++++++-- 2 files changed, 743 insertions(+), 29 deletions(-) create mode 100644 autofit/non_linear/bijector.py diff --git a/autofit/non_linear/bijector.py b/autofit/non_linear/bijector.py new file mode 100644 index 000000000..7654a8a99 --- /dev/null +++ b/autofit/non_linear/bijector.py @@ -0,0 +1,608 @@ +""" +Per-parameter change of variables for the gradient searches. + +Extends :mod:`autofit.non_linear.scaler` from a linear diagonal preconditioner +(``phi = theta / s``) to a per-coordinate **bijection**: an ``AbstractBijector`` +may map a coordinate through ``log`` (multiplicative steps for a log-spread +parameter) or a bounded ``logit`` (a coordinate that is genuinely walled on both +sides), while every other coordinate stays the linear ``identity`` map -- +optionally diagonally scaled, which is how a plain :class:`~autofit.non_linear. +scaler.AbstractScaler` is expressed here too (:class:`BijectorDiagonal`). + +Read :mod:`autofit.non_linear.scaler` first. Everything below assumes its +"why not the unit cube" objections (a logit reparameterisation sends a +boundary optimum to infinity; an inverse-CDF transform has a face +singularity) and its invariance argument (a constant-Jacobian change of +variables cannot move the MAP) as background, and extends both. + +Why per-coordinate bijections are safe where the unit cube was rejected +------------------------------------------------------------------------- + +The scaler's objections to the unit cube are about *every coordinate at +once*, unconditionally. This module does not revisit that decision -- it +answers a narrower question: is there any coordinate for which a *specific*, +individually-chosen non-linear map is a strict improvement, applied only +there? + +- **log**: only ever offered for :class:`~autofit.mapper.prior.log_uniform. + LogUniformPrior` (positive support) and :class:`~autofit.mapper.prior. + log_gaussian.LogGaussianPrior` (support ``(0, inf)``, already a lognormal -- + ``log`` recovers exactly the underlying Gaussian variable). Both have a + *positive, open* domain with no face to send to infinity in the direction + the search actually steps: the lower "wall" is the point at negative + infinity in ``phi``, so a step that approaches it is receding, not + accelerating toward a singularity. This is precisely the one case the + scaler's log-ratio scale rule already targets (``ScalerPriorWidth``'s + ``LogUniformPrior`` branch) -- the same coordinate, one step further: not + just scaled multiplicatively, but *stepped* multiplicatively. + +- **logit**: the scaler's specific objection stands and is not undone here -- + a genuinely boundary-pinned optimum is sent to infinity in ``phi``, exactly + as it warned. ``BijectorLogit`` is offered as a **secondary, opt-in** arm + for a model known *not* to have boundary optima (a mixed sweep of + ``BijectorAuto`` vs ``BijectorLogit`` is how that would be established), not + a default. + +- **identity**: unconditionally safe -- it is the map the whole rest of the + library already uses. + +The equivalence argument (why no Jacobian for MAP) +---------------------------------------------------- + +Let ``g = inverse`` be one coordinate's bijection, a strictly monotone map of +that coordinate's prior support onto ``phi``-space (every kind here is +monotone increasing). The search minimises + + F(phi) = fom(g(phi)) + +where ``fom`` is the physical-space objective (``-2 * log_posterior``). Because +``g`` is a bijection, composing with it **relabels points without changing the +objective's value set** -- the same value of ``fom`` is attained at +``theta = g(phi)`` as at ``phi`` under ``F``. So + + argmin F = g^-1(argmin fom) + +exactly, and **no Jacobian is added** to ``F``. Adding one would move the +answer: the pushforward density's mode is *not*, in general, +``g^-1(mode of theta)`` -- for ``log`` specifically it shifts by exactly the +``-log(lambda)`` term that ``LogUniformPrior.log_prior_from_value`` already +carries as its own (dropped) normalisation, i.e. folding a Jacobian into ``F`` +here would silently double up a term the prior density already accounts for +under change of variables. This is the same "compose the objective through the +map, never re-express the density" discipline +:mod:`autofit.non_linear.scaler` uses, generalised from a constant Jacobian +(which cannot move the MAP either way) to a non-constant one (which can, and +must therefore never be added for a MAP search). + +Why :meth:`AbstractBijector.log_det_jacobian` exists at all, then +-------------------------------------------------------------------- + +A *sampler* -- as opposed to a MAP optimizer -- is a different consumer with a +different contract: it wants **samples of the physical posterior**, not its +argmax. Drawing (or stepping) in ``phi`` and pushing the resulting samples +through ``g`` changes their density unless the Jacobian is folded in +(``p_theta(theta) = p_phi(phi) * |d phi/d theta|``, the ordinary change-of- +variables rule). ``log_det_jacobian`` is exposed for exactly that consumer -- +the multi-chain gradient sampler tracked as PyAutoFit#1521/#1522 -- and is +**never called by MultiStart**, which is a MAP search and must not use it (see +above). Its absence from the ``MultiStartGradient`` wiring in +:mod:`autofit.non_linear.search.mle.multi_start_gradient.search` is therefore +not an oversight to "complete" later; it is the correctness condition. + +Vectorisation and the double-where pattern +--------------------------------------------- + +Every method is vectorised over ``(n_starts, n_params)`` in +``model.priors_ordered_by_id`` order, matching +:meth:`~autofit.non_linear.scaler.AbstractScaler.scale_from_model` and +:meth:`~autofit.non_linear.clipper.AbstractClipper.bounds_from_model` so the +three compose without a reindex. The per-coordinate *kind* (identity / log / +logit) is resolved once, in :meth:`AbstractBijector.from_model`, into a +**static** ``numpy`` integer array baked onto the instance -- never re-derived +inside ``forward`` / ``inverse`` / ``log_det_jacobian`` themselves, so that +under ``jax.jit`` those methods trace to a single fixed program (one +``jnp.where`` selection tree per call) rather than branching in Python on +every trace. + +Selecting a branch with ``jnp.where`` still *evaluates every branch* -- that is +how ``jnp.where`` computes its gradient. A coordinate outside a given branch's +domain (e.g. ``log`` of a coordinate that is not log-kind, which may be zero or +negative) would then feed an invalid value into ``log`` / division, and even +though the outer ``where`` discards the *result*, ``0 * NaN = NaN`` in IEEE 754 +means the discarded branch's ``NaN`` can still poison the gradient through the +multiply-by-zero-mask that ``jnp.where``'s VJP performs. The fix used +throughout this module is the same "double where" safe-surrogate pattern +``LogUniformPrior.log_prior_from_value`` uses +(``autofit/mapper/prior/log_uniform.py``, ``in_bounds`` branch): substitute an +arbitrary in-domain value (e.g. ``1.0`` before a ``log``, ``0.0`` before an +``exp``) *inside* the risky call, via an inner ``where``, before the outer +``where`` picks the real branch. The discarded branch then evaluates to a +finite, safe number rather than ``NaN``/``-inf``, and the multiply-by-zero-mask +gradient trick is safe. +""" + +import logging +from abc import ABC, abstractmethod +from typing import Dict, List, Optional, Tuple + +import numpy as np + +from autofit.mapper.prior.log_gaussian import LogGaussianPrior +from autofit.mapper.prior.log_uniform import LogUniformPrior +from autofit.non_linear.scaler import AbstractScaler + +logger = logging.getLogger(__name__) + +# Static per-coordinate kind codes. Plain ints (not an Enum) because these are +# baked into `numpy` arrays that get compared with `==` inside `jnp.where` +# chains -- an Enum would work too, but would force every comparison through +# `.value` for no benefit. +_IDENTITY = 0 +_LOG = 1 +_LOGIT = 2 + +_KIND_NAMES = {_IDENTITY: "identity", _LOG: "log", _LOGIT: "logit"} +_KIND_CODES = {name: code for code, name in _KIND_NAMES.items()} + + +def _stable_sigmoid(x, xp): + """ + ``1 / (1 + exp(-x))``, computed so neither branch ever overflows ``exp``. + + For ``x >= 0`` uses ``1 / (1 + exp(-x))`` (``-x <= 0``, so ``exp`` cannot + overflow); for ``x < 0`` uses the algebraically identical + ``exp(x) / (1 + exp(x))`` (``x < 0``, so ``exp`` cannot overflow there + either). Finite and safe for every finite ``x``, so no additional + double-where guard is needed around it. + """ + positive = x >= 0 + z = xp.exp(-xp.abs(x)) + return xp.where(positive, 1.0 / (1.0 + z), z / (1.0 + z)) + + +def _softplus(x, xp): + """``log(1 + exp(x))``, computed in the standard overflow-safe form.""" + return xp.maximum(x, 0.0) + xp.log1p(xp.exp(-xp.abs(x))) + + +class AbstractBijector(ABC): + """ + A strategy supplying a per-parameter change of variables for a search that + steps in physical parameter space. + + Subclasses are pluggable per-search, exactly like + :class:`~autofit.non_linear.scaler.AbstractScaler` and + :class:`~autofit.non_linear.clipper.AbstractClipper`, and the default is + :class:`BijectorNone` so behaviour is unchanged until a user opts in. + + Unlike ``AbstractScaler`` (whose ``scale_from_model`` returns a fresh array + every call and holds no state), a bijector's per-coordinate *kind* cannot + be an argument of ``forward`` / ``inverse`` / ``log_det_jacobian`` without + breaking ``jax.jit`` tracing (see the module docstring), so it is resolved + **once**, against one model, and cached on the instance by + :meth:`from_model`. Construct one bijector per search/model, the same + discipline ``AbstractScaler`` follows in practice even though its own + statelessness does not strictly require it. + """ + + def __init__(self): + self._kind_code: Optional[np.ndarray] = None + self._scale: Optional[np.ndarray] = None + self._lo: Optional[np.ndarray] = None + self._hi: Optional[np.ndarray] = None + self._logit_eps = 1.0e-12 + + @abstractmethod + def _kind_scale_bounds_from_model( + self, model + ) -> Tuple[List[int], List[float], List[float], List[float]]: + """ + Subclass hook: the per-coordinate ``(kind_code, scale, lower, upper)`` + in ``model.priors_ordered_by_id`` order. + + ``scale`` is used only by ``identity``-kind coordinates (``phi = + theta / scale``); ``lower`` / ``upper`` only by ``logit``-kind ones + (the box the logit is taken over). Both are still returned for every + coordinate -- unused entries are never read, but ``from_model`` builds + one dense array per field rather than a ragged one keyed by kind. + """ + + def from_model(self, model) -> "AbstractBijector": + """ + Resolve and cache this bijector's per-coordinate kind/scale/bound + arrays against ``model``, and return ``self``. + + Must be called once, before ``forward`` / ``inverse`` / + ``log_det_jacobian`` / ``kinds``, and again for a different model -- + it overwrites the previous cache rather than accumulating state, the + same "one strategy object per search" discipline + :class:`~autofit.non_linear.scaler.AbstractScaler` follows. + """ + kind_codes, scale, lower, upper = self._kind_scale_bounds_from_model(model) + self._kind_code = np.asarray(kind_codes, dtype=np.int8) + self._scale = np.asarray(scale, dtype=float) + self._lo = np.asarray(lower, dtype=float) + self._hi = np.asarray(upper, dtype=float) + return self + + def _check_resolved(self): + if self._kind_code is None: + raise RuntimeError( + f"{type(self).__name__}.from_model(model) must be called before " + "forward / inverse / log_det_jacobian / bounds_forward / kinds " + "-- the per-coordinate kind arrays are not yet resolved." + ) + + @property + def kinds(self) -> List[str]: + """The per-coordinate kind, in ``model.priors_ordered_by_id`` order.""" + self._check_resolved() + return [_KIND_NAMES[int(code)] for code in self._kind_code] + + def forward(self, theta, xp=np): + """ + The physical-to-transformed map, ``theta -> phi``, vectorised over the + last axis (``(n_params,)`` or ``(n_starts, n_params)``). + """ + self._check_resolved() + + is_id = self._kind_code == _IDENTITY + is_log = self._kind_code == _LOG + is_logit = self._kind_code == _LOGIT + + scale = self._scale + lo = self._lo + hi = self._hi + + phi_id = theta / scale + + # Safe surrogate `1.0` inside `log` for every non-log coordinate (which + # may be zero, negative, or anything else); an extra floor guards a + # genuinely log-kind coordinate that has drifted to exactly zero. + theta_safe = xp.where(is_log, theta, 1.0) + theta_safe = xp.maximum(theta_safe, 1.0e-300) + phi_log = xp.log(theta_safe) + + # Safe surrogate `0.5` for `u` on every non-logit coordinate, clamped + # away from the exact 0/1 edges where `logit` diverges. + denom = xp.where(is_logit, hi - lo, 1.0) + u = xp.where(is_logit, (theta - lo) / denom, 0.5) + u = xp.clip(u, self._logit_eps, 1.0 - self._logit_eps) + phi_logit = xp.log(u) - xp.log(1.0 - u) + + return xp.where(is_id, phi_id, xp.where(is_log, phi_log, phi_logit)) + + def inverse(self, phi, xp=np): + """ + The transformed-to-physical map, ``phi -> theta`` -- the search's + ``g`` (see the module docstring's equivalence argument). Vectorised the + same way as :meth:`forward`. + """ + self._check_resolved() + + is_id = self._kind_code == _IDENTITY + is_log = self._kind_code == _LOG + is_logit = self._kind_code == _LOGIT + + scale = self._scale + lo = self._lo + hi = self._hi + + theta_id = phi * scale + + # Safe surrogate `0.0` inside `exp` for every non-log coordinate: `phi` + # there is in an unrelated space (identity/logit) and may be large. + phi_log_safe = xp.where(is_log, phi, 0.0) + theta_log = xp.exp(phi_log_safe) + + phi_logit_safe = xp.where(is_logit, phi, 0.0) + sigmoid = _stable_sigmoid(phi_logit_safe, xp) + theta_logit = lo + (hi - lo) * sigmoid + + return xp.where(is_id, theta_id, xp.where(is_log, theta_log, theta_logit)) + + def log_det_jacobian(self, phi, xp=np): + """ + ``log |d theta / d phi|``, summed over parameters, per row -- shape + ``()`` for a single ``(n_params,)`` vector or ``(n_starts,)`` for a + batch. **Never used to form a MAP objective** (see the module + docstring); this is for a sampler that must push ``phi``-space samples + through :meth:`inverse` into physical-space ones via the ordinary + change-of-variables rule. + + Closed forms, per kind (``theta = inverse(phi)``): + + - identity (scale ``s``): ``theta = s * phi``, so + ``d theta / d phi = s`` and the contribution is ``log(s)``. + - log: ``theta = exp(phi)``, so ``d theta / d phi = exp(phi) = theta`` + and the contribution is ``log(theta) = phi`` exactly. + - logit (box ``[lo, hi]``): ``theta = lo + (hi - lo) * sigmoid(phi)``, + so ``d theta / d phi = (hi - lo) * sigmoid(phi) * (1 - sigmoid(phi))`` + and the contribution is + ``log(hi - lo) - softplus(-phi) - softplus(phi)`` (the standard + stable log-sigmoid / log-one-minus-sigmoid pair). + """ + self._check_resolved() + + is_id = self._kind_code == _IDENTITY + is_log = self._kind_code == _LOG + is_logit = self._kind_code == _LOGIT + + scale = self._scale + width = self._hi - self._lo + + contrib_id = xp.log(scale) + contrib_log = phi + contrib_logit = ( + xp.log(xp.where(is_logit, width, 1.0)) + - _softplus(-phi, xp) + - _softplus(phi, xp) + ) + + per_coord = xp.where( + is_id, contrib_id, xp.where(is_log, contrib_log, contrib_logit) + ) + return xp.sum(per_coord, axis=-1) + + def bounds_forward( + self, lower: np.ndarray, upper: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray]: + """ + Map a physical ``(lower, upper)`` box into transformed coordinates, + elementwise, for :class:`~autofit.non_linear.clipper.AbstractClipper` + to clip against. + + Every kind here is strictly monotone **increasing**, so this commutes + with clipping: ``clip(forward(theta), forward(lo), forward(hi)) == + forward(clip(theta, lo, hi))``. Pure ``numpy`` -- the clipper's bounds + are a constant vector, never a traced value. + """ + self._check_resolved() + lower = np.asarray(lower, dtype=float) + upper = np.asarray(upper, dtype=float) + return self.forward(lower, xp=np), self.forward(upper, xp=np) + + def info_from_model(self, model) -> str: + """ + A human-readable summary of the resolved kinds, appended to the + written ``model.info`` file. Resolves against ``model`` first if this + bijector has not been used yet, matching + :meth:`~autofit.non_linear.scaler.AbstractScaler.info_from_model`'s + contract of being safe to call standalone. + """ + if self._kind_code is None: + self.from_model(model) + + priors = model.priors_ordered_by_id + if len(priors) == 0: + return "" + + rows = [] + for prior, kind in zip(priors, self.kinds): + path = model.path_for_prior(prior) + path_str = ".".join(str(p) for p in path) if path is not None else "?" + rows.append((path_str, type(prior).__name__, kind)) + + widths = [max(len(row[column]) for row in rows) for column in range(2)] + + info = f"Per-Parameter Bijector ({type(self).__name__})\n\n" + for path_str, prior_name, kind in rows: + info += f"{path_str:<{widths[0]}} {prior_name:<{widths[1]}} kind {kind}\n" + + return info + + +class BijectorNone(AbstractBijector): + """ + The no-op bijector, and **the default**. + + Every coordinate is ``identity`` with ``scale = 1.0``, so ``phi == theta`` + exactly (``x / 1.0 == x`` and ``x * 1.0 == x`` bitwise in IEEE 754) and + ``log_det_jacobian`` is exactly ``0.0``. Searches should test for + ``isinstance(bijector, BijectorNone)`` and skip the change of variables + entirely rather than applying it as a no-op, the same short-circuit + :class:`~autofit.non_linear.scaler.ScalerNone` relies on, so the compiled + step is unchanged too. + """ + + def _kind_scale_bounds_from_model(self, model): + n = model.prior_count + return ( + [_IDENTITY] * n, + [1.0] * n, + [0.0] * n, + [1.0] * n, + ) + + def info_from_model(self, model) -> str: + return "" + + +class BijectorAuto(AbstractBijector): + """ + ``log`` for every eligible coordinate, ``identity`` (unscaled) elsewhere. + + Eligible means :class:`~autofit.mapper.prior.log_uniform.LogUniformPrior` + with a strictly positive ``lower_limit`` (the only case the prior's own + constructor allows in the first place, so the fallback below is a + defence-in-depth guard rather than a reachable path through normal + construction), or :class:`~autofit.mapper.prior.log_gaussian. + LogGaussianPrior` (support ``(0, inf)``, unconditionally eligible). A + malformed candidate falls back to ``identity`` with a logged warning, + mirroring ``ScalerPriorWidth``'s fallback for the same prior type. + """ + + @staticmethod + def _kind_for(prior) -> Tuple[int, float, float]: + if isinstance(prior, LogUniformPrior): + lower = float(prior.lower_limit) + if lower > 0.0: + return _LOG, lower, float(prior.upper_limit) + logger.warning( + f"BijectorAuto: LogUniformPrior with a non-positive lower_limit " + f"({lower}); falling back to identity for this parameter." + ) + return _IDENTITY, 0.0, 1.0 + + if isinstance(prior, LogGaussianPrior): + return _LOG, 0.0, np.inf + + return _IDENTITY, 0.0, 1.0 + + def _kind_scale_bounds_from_model(self, model): + kind_codes, lo, hi = [], [], [] + for prior in model.priors_ordered_by_id: + code, low, high = self._kind_for(prior) + kind_codes.append(code) + lo.append(low) + hi.append(high) + return kind_codes, [1.0] * len(kind_codes), lo, hi + + +class BijectorLogit(AbstractBijector): + """ + ``logit`` on the prior box for every two-sided finite coordinate, + ``identity`` (unscaled) elsewhere. + + **Secondary arm.** :mod:`autofit.non_linear.scaler`'s objection to the + unit cube stands unchanged: a genuinely boundary-pinned optimum is sent to + infinity in ``phi``, and the reference cell has such optima. This class + exists so that claim can be measured directly (a ``BijectorAuto`` vs + ``BijectorLogit`` A/B), not as a recommended default. + + :class:`~autofit.mapper.prior.log_gaussian.LogGaussianPrior`'s support is + ``(0, inf)`` -- half-open, not two-sided -- so it always falls back to + ``identity`` here (with a logged warning), the same declared-support + correction :class:`~autofit.non_linear.clipper.ClipperPriorBox` applies. + """ + + @staticmethod + def _kind_for(prior) -> Tuple[int, float, float]: + low = float(prior.lower_limit) + high = float(prior.upper_limit) + if isinstance(prior, LogGaussianPrior): + low = 0.0 + + if np.isfinite(low) and np.isfinite(high) and high > low: + return _LOGIT, low, high + + logger.warning( + f"BijectorLogit: {type(prior).__name__} has no two-sided finite " + f"bound (lower={low}, upper={high}); falling back to identity for " + "this parameter." + ) + return _IDENTITY, 0.0, 1.0 + + def _kind_scale_bounds_from_model(self, model): + kind_codes, lo, hi = [], [], [] + for prior in model.priors_ordered_by_id: + code, low, high = self._kind_for(prior) + kind_codes.append(code) + lo.append(low) + hi.append(high) + return kind_codes, [1.0] * len(kind_codes), lo, hi + + +class BijectorPerPath(AbstractBijector): + """ + An explicit, user-declared ``{path: kind}`` mapping, resolved against a + model's own path helper (:meth:`~autofit.mapper.prior_model.abstract. + AbstractPriorModel.path_for_prior`). + + Parameters + ---------- + kind_by_path + Maps a dotted parameter path (as rendered by + ``".".join(str(p) for p in model.path_for_prior(prior))`` -- the same + format :meth:`~autofit.non_linear.scaler.ScalerPriorWidth. + info_from_model` reports) to one of ``"identity"``, ``"log"``, + ``"logit"``. A path not present in the mapping defaults to + ``"identity"``. A requested kind that is not eligible for that + prior (``"log"`` on anything but a positive-lower ``LogUniformPrior`` + or a ``LogGaussianPrior``; ``"logit"`` on a coordinate without a + two-sided finite bound) falls back to ``"identity"`` with a logged + warning rather than silently doing something else or raising -- + matching every other fallback in this module. + """ + + def __init__(self, kind_by_path: Dict[str, str]): + super().__init__() + self.kind_by_path = dict(kind_by_path) + + @staticmethod + def _resolve_one(requested: str, prior, path_str: str) -> Tuple[int, float, float]: + if requested == "log": + if isinstance(prior, LogUniformPrior) and float(prior.lower_limit) > 0.0: + return _LOG, float(prior.lower_limit), float(prior.upper_limit) + if isinstance(prior, LogGaussianPrior): + return _LOG, 0.0, np.inf + logger.warning( + f"BijectorPerPath: path {path_str!r} requested 'log' but its " + f"prior ({type(prior).__name__}) is not eligible (needs a " + "LogUniformPrior with a positive lower_limit, or a " + "LogGaussianPrior); falling back to identity." + ) + return _IDENTITY, 0.0, 1.0 + + if requested == "logit": + low = float(prior.lower_limit) + high = float(prior.upper_limit) + if isinstance(prior, LogGaussianPrior): + low = 0.0 + if np.isfinite(low) and np.isfinite(high) and high > low: + return _LOGIT, low, high + logger.warning( + f"BijectorPerPath: path {path_str!r} requested 'logit' but its " + f"prior ({type(prior).__name__}) has no two-sided finite bound " + f"(lower={low}, upper={high}); falling back to identity." + ) + return _IDENTITY, 0.0, 1.0 + + if requested != "identity": + logger.warning( + f"BijectorPerPath: path {path_str!r} requested unknown kind " + f"{requested!r}; falling back to identity." + ) + return _IDENTITY, 0.0, 1.0 + + def _kind_scale_bounds_from_model(self, model): + kind_codes, lo, hi = [], [], [] + for prior in model.priors_ordered_by_id: + path = model.path_for_prior(prior) + path_str = ".".join(str(p) for p in path) if path is not None else None + requested = self.kind_by_path.get(path_str, "identity") + code, low, high = self._resolve_one(requested, prior, path_str) + kind_codes.append(code) + lo.append(low) + hi.append(high) + return kind_codes, [1.0] * len(kind_codes), lo, hi + + +class BijectorDiagonal(AbstractBijector): + """ + Adapter expressing an :class:`~autofit.non_linear.scaler.AbstractScaler` + as a bijector: every coordinate is ``identity``, diagonally scaled by the + wrapped scaler's ``scale_from_model``. + + This subsumes :class:`~autofit.non_linear.scaler.ScalerPriorWidth` + exactly -- ``forward`` / ``inverse`` here reduce to precisely the + ``phi = theta / scale`` / ``theta = scale * phi`` maps + ``MultiStartGradient`` already applies under ``scaler=``. It exists so a + diagonal preconditioner can be composed inside this framework (e.g. + stacked with a per-path ``log`` on top, once a caller wants both) without + duplicating the scale rule, and so ``scaler=`` keeps working completely + unchanged: it is not wired into the ``scaler=`` code path itself, which + stays exactly as it was. + """ + + def __init__(self, scaler: AbstractScaler): + super().__init__() + self.scaler = scaler + + def _kind_scale_bounds_from_model(self, model): + scale = self.scaler.scale_from_model(model=model) + n = len(scale) + return [_IDENTITY] * n, list(scale), [0.0] * n, [1.0] * n + + def info_from_model(self, model) -> str: + return self.scaler.info_from_model(model) diff --git a/autofit/non_linear/clipper.py b/autofit/non_linear/clipper.py index eb11ea97b..2503d85e1 100644 --- a/autofit/non_linear/clipper.py +++ b/autofit/non_linear/clipper.py @@ -62,11 +62,40 @@ on **what kind of bound it is**, never on unguarded ``upper - lower`` arithmetic. That distinction is load-bearing rather than fussy — see :class:`ClipperPriorBox` for the three cases and why collapsing them breaks. + +Composition with a :mod:`~autofit.non_linear.bijector` +-------------------------------------------------------- + +``AbstractClipper.project`` also accepts a ``bijector``, alongside (never +together with) ``scale``, so a search stepping in ``phi = bijector.forward +(theta)`` can clip in that same space. Every bijector kind is elementwise +**strictly monotone increasing**, so forward and clip commute: clipping +``forward(theta)`` against ``forward(lower_inset)`` / ``forward(upper_inset)`` +gives exactly ``forward(clip(theta, lower_inset, upper_inset))``. + +**The physical inset is wrong for a ``log``-kind coordinate, and the error is +not marginal.** The two-sided inset above is a *relative-to-physical-width* +margin, ``lower + margin * (upper - lower)``. For +``LogUniformPrior(1e-6, 1e6)`` with the default ``margin=1e-6`` that is +``1e-6 + 1e-6 * (1e6 - 1e-6) ~ 1e-6 + 1.0 ~ 1.0`` — a PHYSICAL inset of order +1, on a prior whose entire point is that it is uniform across twelve decades. +Clipping (or mapping through a ``log`` bijector) against that bound fences off +every ``lambda < 1``, silently, for a margin whose entire purpose is to be +negligible. The bug is in the inset arithmetic, not in the bijector: a +``log``-kind coordinate's natural margin is a fraction of its **log-ratio**, +``margin * log(upper / lower)``, mapped back through ``exp`` — the same +"physical width is not usable here" correction +:class:`~autofit.non_linear.scaler.ScalerPriorWidth` already makes for its own +``LogUniformPrior`` scale rule. ``ClipperPriorBox`` applies this whenever a +``bijector`` reports a coordinate's kind as ``"log"`` (see +:meth:`ClipperPriorBox._inset_from_model`); the plain, no-bijector +``bounds_from_model`` path is unaffected — every existing caller keeps its +current (physical) inset exactly. """ import logging from abc import ABC, abstractmethod -from typing import Tuple +from typing import Optional, Tuple import numpy as np @@ -107,7 +136,7 @@ def bounds_from_model(self, model) -> Tuple[np.ndarray, np.ndarray]: """ @abstractmethod - def project(self, vector, model, xp=np, scale=None): + def project(self, vector, model, xp=np, scale=None, bijector=None): """ Project ``vector`` onto the prior support. @@ -117,7 +146,8 @@ def project(self, vector, model, xp=np, scale=None): 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. + ``scale`` or ``bijector`` is given, in which case it is in that + change of variables' coordinates. model The model whose priors define the support. xp @@ -130,7 +160,15 @@ def project(self, vector, model, xp=np, scale=None): 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``. + ``+/-inf``. Mutually exclusive with ``bijector``. + bijector + A resolved :class:`~autofit.non_linear.bijector.AbstractBijector` + (``.from_model`` already called), when the caller is stepping in + ``phi = bijector.forward(theta)``. The inset bounds are computed in + physical space (kind-aware — see the module docstring) and then + mapped through ``bijector.bounds_forward``, which commutes with + clipping because every bijector kind is monotone increasing. + Mutually exclusive with ``scale``. Returns ------- @@ -153,7 +191,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, scale=None): + def project(self, vector, model, xp=np, scale=None, bijector=None): return vector, xp.zeros_like(vector, dtype=bool) @@ -256,21 +294,65 @@ def _limits_from_model(self, model): np.array(upper_strict, dtype=bool), ) - def bounds_from_model(self, model) -> Tuple[np.ndarray, np.ndarray]: + def _inset_from_model( + self, model, kinds: Optional[list] = None + ) -> Tuple[np.ndarray, np.ndarray]: + """ + The inset ``(lower, upper)`` box, in physical parameter order, aware of + an optional per-coordinate ``kinds`` list (from + :attr:`~autofit.non_linear.bijector.AbstractBijector.kinds`). + + ``kinds=None`` (the default, and what :meth:`bounds_from_model` passes) + reproduces the original physical-relative-margin rule for every + two-sided finite coordinate, unchanged. Where ``kinds[i] == "log"`` the + margin is instead taken as a fraction of the LOG-ratio and mapped back + through ``exp`` — see the module docstring for why the physical margin + is wrong there (it is not a small correction: it can be O(1) against a + prior spanning many decades). + """ lower, upper, lower_strict, upper_strict = self._limits_from_model(model) two_sided = np.isfinite(lower) & np.isfinite(upper) - # The width is evaluated ONLY where both bounds are finite. `np.where` - # alone would not be enough -- it evaluates both branches, so `inf - -inf` - # would still be computed and still be NaN. The finite substitution has to - # happen before the subtraction, not after it. - safe_lower = np.where(two_sided, lower, 0.0) - safe_upper = np.where(two_sided, upper, 0.0) + is_log = np.zeros(len(lower), dtype=bool) + if kinds is not None: + is_log = np.array([kind == "log" for kind in kinds], dtype=bool) + # A "log" kind is only ever assigned (by the bijector classes) to a + # coordinate with a strictly positive, two-sided-finite physical + # bound; this additional guard is defence-in-depth so a stray or + # malformed `kinds` entry can never reach `np.log` of a + # non-positive or non-finite value below. + is_log = is_log & two_sided & (lower > 0.0) + + # PHYSICAL (linear-space) margin, for every two-sided finite coordinate + # that is NOT log-kind -- identical to the original rule. The width is + # evaluated ONLY where both bounds are finite and the coordinate is a + # linear target. `np.where` alone would not be enough -- it evaluates + # both branches, so `inf - -inf` would still be computed and still be + # NaN. The finite substitution has to happen before the subtraction, + # not after it. + linear_target = two_sided & ~is_log + safe_lower = np.where(linear_target, lower, 0.0) + safe_upper = np.where(linear_target, upper, 0.0) relative = self.margin * (safe_upper - safe_lower) - lower_inset = lower + np.where(two_sided, relative, 0.0) - upper_inset = upper - np.where(two_sided, relative, 0.0) + lower_inset = lower + np.where(linear_target, relative, 0.0) + upper_inset = upper - np.where(linear_target, relative, 0.0) + + # LOG-space margin, for every log-kind coordinate: a fraction of the + # log-ratio, mapped back through `exp`. `safe_log_lower` / + # `safe_log_upper` substitute `1.0` for every non-log-kind coordinate + # so `np.log` is never evaluated outside its domain, the same + # before-the-call substitution `bounds_from_model` above already uses. + safe_log_lower = np.where(is_log, lower, 1.0) + safe_log_upper = np.where(is_log, upper, 1.0) + log_relative = self.margin * np.log(safe_log_upper / safe_log_lower) + + log_lower_inset = np.log(safe_log_lower) + log_relative + log_upper_inset = np.log(safe_log_upper) - log_relative + + lower_inset = np.where(is_log, np.exp(log_lower_inset), lower_inset) + upper_inset = np.where(is_log, np.exp(log_upper_inset), upper_inset) # Half-open bounds get an absolute nudge, since `relative` is identically # zero for them. @@ -283,21 +365,45 @@ def bounds_from_model(self, model) -> Tuple[np.ndarray, np.ndarray]: return lower_inset, upper_inset - 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 + def bounds_from_model(self, model) -> Tuple[np.ndarray, np.ndarray]: + return self._inset_from_model(model, kinds=None) + + def project(self, vector, model, xp=np, scale=None, bijector=None): + if scale is not None and bijector is not None: + raise ValueError( + f"{type(self).__name__}.project received both `scale` and " + "`bijector` -- they are two different changes of variables for " + "the same step, and a caller must pick exactly one (a search " + "already enforces this at construction; see " + "AbstractMultiStartGradient.__init__)." + ) + + if bijector is not None: + # Inset in PHYSICAL space, kind-aware, then mapped through the + # bijector -- never the other way round. Every kind is monotone + # increasing, so this commutes with clipping (see the module + # docstring), and computing the inset in physical space is what + # lets the log-kind correction above be expressed in terms even a + # non-bijector-aware caller (`bounds_from_model`) already + # understands. + lower, upper = self._inset_from_model(model, kinds=bijector.kinds) + lower, upper = bijector.bounds_forward(lower, upper) + else: + 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) From af82c2441de959d4ee5cb15777f0fa1f73e702af Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 24 Aug 2026 18:00:26 -0400 Subject: [PATCH 2/4] Wire AbstractBijector into MultiStartGradient, beside scaler Adds bijector: Optional[AbstractBijector] = None alongside scaler, mutually exclusive (raises ValueError at construction if both are non-default). Wherever the step loop maps between physical and stepped coordinates for the scaler (resume load, fresh-start seeding, best/ lane-best capture, dead-lane redraw, the write-back), a parallel bijector.forward/inverse branch is added; search_internal["params"] stays PHYSICAL either way so resume is safe across a scaler/bijector change. search_internal also gains "bijector": bijector.kinds, and samples_info records the bijector's class name and (from search_internal, not live state) its resolved kinds -- mirroring how "scaler"/"clipper" are already reported. AbstractClipper.project's new `bijector=` param is threaded through the clip call and the dead-lane redraw. Everything is gated behind the same has_scaler/has_bijector-style short-circuit already used for the scaler/clipper/constraint checks, so the default (BijectorNone) step loop and its jax.jit trace are unchanged. Also adds two independently flag-gated, off-by-default diagnostics (neither touches the compiled step when unused): record_lane_nan_history (per-step per-lane value/grad-NaN bits, numpy.packbits'd to (n_steps, ceil(n_starts/8))) and trace_param_indices (a physical per-step trace for a chosen subset of parameters, (n_steps, n_starts, k)). Both resume via the same `.get`-default discipline as the existing lifetime counters. Exports the Bijector* classes from autofit/__init__.py beside the Scaler* ones. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01EDABYoH6giHXhFJUks8yd6 --- autofit/__init__.py | 6 + .../search/mle/multi_start_gradient/search.py | 248 ++++++++++++++++-- 2 files changed, 234 insertions(+), 20 deletions(-) diff --git a/autofit/__init__.py b/autofit/__init__.py index 943824a7f..6f674f8ee 100644 --- a/autofit/__init__.py +++ b/autofit/__init__.py @@ -83,6 +83,12 @@ from .non_linear.scaler import AbstractScaler from .non_linear.scaler import ScalerNone from .non_linear.scaler import ScalerPriorWidth +from .non_linear.bijector import AbstractBijector +from .non_linear.bijector import BijectorNone +from .non_linear.bijector import BijectorAuto +from .non_linear.bijector import BijectorLogit +from .non_linear.bijector import BijectorPerPath +from .non_linear.bijector import BijectorDiagonal from .non_linear.initializer import InitializerBall from .non_linear.initializer import InitializerPrior from .non_linear.initializer import InitializerParamBounds 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 129014d98..bbf34432c 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -2,7 +2,7 @@ import inspect import pickle -from typing import Optional, TYPE_CHECKING +from typing import Optional, Sequence, TYPE_CHECKING import numpy as np @@ -11,6 +11,7 @@ from autofit.non_linear.search.mle.abstract_mle import AbstractMLE from autofit.non_linear.analysis import Analysis from autofit.non_linear.fitness import Fitness +from autofit.non_linear.bijector import AbstractBijector, BijectorNone from autofit.non_linear.clipper import AbstractClipper, ClipperNone from autofit.non_linear.scaler import AbstractScaler, ScalerNone from autofit.non_linear.initializer import AbstractInitializer @@ -62,7 +63,10 @@ def __init__( initializer: Optional[AbstractInitializer] = None, clipper: Optional[AbstractClipper] = None, scaler: Optional[AbstractScaler] = None, + bijector: Optional[AbstractBijector] = None, reset_momentum_on_clip: bool = False, + record_lane_nan_history: bool = False, + trace_param_indices: Optional[Sequence[int]] = None, iterations_per_full_update: int = None, iterations_per_quick_update: int = None, silence: bool = False, @@ -177,6 +181,28 @@ def __init__( reaching a wall rarer, never impossible, and where the likelihood genuinely prefers a value outside the prior only clipping can express the resulting MAP. + bijector + A per-parameter change of variables (see + :mod:`autofit.non_linear.bijector`), applied instead of ``scaler``: + the rule steps in ``phi = bijector.forward(theta)`` while the + objective is still evaluated at the physical + ``theta = bijector.inverse(phi)``. Where ``scaler`` only ever + rescales linearly, a bijector may additionally reparameterise a + coordinate through ``log`` (for a coordinate whose natural step is + multiplicative, e.g. a ``LogUniformPrior``) or ``logit`` (secondary; + see :class:`~autofit.non_linear.bijector.BijectorLogit`). Mutually + exclusive with ``scaler`` -- passing both non-default raises + ``ValueError`` at construction, since they are two different + changes of variables for the same step. + + No Jacobian is added to the objective for the same reason none is + added under ``scaler``: composing through a bijection relabels + points without changing the objective's value set, so the MAP is + unaffected (see the module docstring's equivalence argument). + + Default ``BijectorNone`` -- a no-op, skipped entirely rather than + applied as an identity, so the default path and its compiled step + are both unchanged. resurrect Restart-on-death. When ``True``, any start whose objective goes non-finite is redrawn each step (fresh params from the start band + @@ -226,6 +252,27 @@ def __init__( this flag would then be discarding useful state. It is a knob for the case where pinning is an artefact of momentum rather than a statement about the data. + record_lane_nan_history + Record ``lane_value_nan_history`` / ``lane_grad_nan_history`` -- + one bit per lane per step, ``bitpacked`` to ``(n_steps, + ceil(n_starts / 8))`` via ``numpy.packbits`` -- alongside the + existing lifetime totals. The totals (``n_value_nan_lane_steps``, + ``n_grad_nan_lane_steps``) answer "how much"; this answers "when, + and which lane" at the cost of ``O(n_steps * n_starts)`` bits of + ``search_internal``, which is why it defaults ``False``. Uses the + same short-circuit style as ``has_scaler`` / ``has_clipper`` above, + so a default run's step loop is unchanged. + trace_param_indices + Optional physical-parameter indices (into + ``model.priors_ordered_by_id``) to record a full per-step trace + for, as ``trace_history`` -- shape ``(n_steps, n_starts, k)``, + PHYSICAL units. ``None`` (default) records nothing, at zero cost to + the step loop, the same short-circuit as + ``record_lane_nan_history``. Intended for post-hoc basin/step + diagnostics on a small, deliberately chosen set of coordinates -- + recording every parameter for every step is the + ``(n_steps, n_starts, n_params)`` array this knob exists to avoid + paying for by default. convergence Auto-convergence (early-stopping) settings. When ``check_for_convergence`` is ``True`` (the default) the search stops @@ -291,7 +338,31 @@ def __init__( # 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.bijector = bijector or BijectorNone() + + # Rejected at construction, not at fit time: `scaler` and `bijector` + # are two different changes of variables for the same step (linear + # rescale vs. a general per-coordinate bijection), and letting both be + # non-default would mean silently picking one over the other rather + # than surfacing the conflict to the caller who set both. + if not isinstance(self.scaler, ScalerNone) and not isinstance( + self.bijector, BijectorNone + ): + raise ValueError( + f"{type(self).__name__} received both a non-default `scaler` " + f"({type(self.scaler).__name__}) and a non-default `bijector` " + f"({type(self.bijector).__name__}). They are two different " + "changes of variables for the same step -- pass exactly one. " + "(A plain diagonal scale can be expressed as a bijector via " + "`autofit.BijectorDiagonal(scaler)` if both are wanted at " + "once.)" + ) + self.reset_momentum_on_clip = reset_momentum_on_clip + self.record_lane_nan_history = bool(record_lane_nan_history) + self.trace_param_indices = ( + list(trace_param_indices) if trace_param_indices is not None else None + ) self.convergence = ( convergence if convergence is not None else MultiStartGradientConvergence() ) @@ -776,8 +847,24 @@ def _fit( scale = self.scaler.scale_from_model(model=model) if has_scaler else None scale_jnp = jnp.asarray(scale) if has_scaler else None + # Per-parameter change of variables (see ``autofit.non_linear.bijector``). + # Mutually exclusive with the scaler (enforced at construction), so at + # most one of ``has_scaler`` / ``has_bijector`` is ever true. Resolved + # ONCE here -- ``from_model`` bakes the per-coordinate kind arrays onto + # ``self.bijector`` -- so ``forward`` / ``inverse`` below need no model + # argument and trace to a fixed program under ``jax.jit``. Same + # short-circuit as ``has_scaler``: under the default ``BijectorNone`` + # nothing here is even resolved, and the compiled step is unchanged. + has_bijector = not isinstance(self.bijector, BijectorNone) + if has_bijector: + self.bijector.from_model(model=model) + def _to_physical(vector): - return vector if scale_jnp is None else vector * scale_jnp + if scale_jnp is not None: + return vector * scale_jnp + if has_bijector: + return self.bijector.inverse(vector, xp=jnp) + return vector # The objective the OPTIMIZER differentiates. Composed as # ``fitness.call(phi * scale)``, so what is optimised remains the @@ -797,7 +884,7 @@ def _to_physical(vector): # 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 + if (has_scaler or has_bijector) else _value_and_grad ) @@ -865,6 +952,8 @@ def batched_value_and_grad(params): params = jnp.asarray(search_internal["params"]) if has_scaler: params = params / scale_jnp + elif has_bijector: + params = self.bijector.forward(params, xp=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"]) @@ -917,6 +1006,26 @@ def batched_value_and_grad(params): dtype=int, ) + # ``.get`` defaults, same doctrine as above: a ``search_internal`` + # written before these histories existed (or by a run with the + # recording flags off) has none of these keys. Restored as a list + # of per-step rows so the step loop can keep appending to them; + # re-stacked into an array at the next write-back. + _lane_value_nan_history = search_internal.get("lane_value_nan_history") + lane_value_nan_rows = ( + list(_lane_value_nan_history) + if _lane_value_nan_history is not None + else [] + ) + _lane_grad_nan_history = search_internal.get("lane_grad_nan_history") + lane_grad_nan_rows = ( + list(_lane_grad_nan_history) + if _lane_grad_nan_history is not None + else [] + ) + _trace_history = search_internal.get("trace_history") + trace_rows = list(_trace_history) if _trace_history is not None else [] + self.logger.info( "Resuming MultiStartGradient search (previous samples found)." ) @@ -978,6 +1087,8 @@ def batched_value_and_grad(params): if has_scaler: params = params / scale_jnp + elif has_bijector: + params = self.bijector.forward(params, xp=jnp) # Per-start optimizer state: one independent state per start, so # learning-rate-free rules never share a global scalar estimate. @@ -994,6 +1105,9 @@ def batched_value_and_grad(params): n_grad_nan_lane_steps = 0 n_constrained_lane_steps = 0 n_clipped_lane_steps = 0 + lane_value_nan_rows = [] + lane_grad_nan_rows = [] + trace_rows = [] stop_reason = None self.logger.info( @@ -1077,6 +1191,10 @@ def batched_value_and_grad(params): best_params = np.asarray(params[best_index]) if has_scaler: best_params = best_params * scale + elif has_bijector: + best_params = np.asarray( + self.bijector.inverse(best_params, xp=np) + ) # Per-lane best capture, PHYSICAL for the same reason as # ``best_params`` just above. The step index is the @@ -1090,13 +1208,39 @@ def batched_value_and_grad(params): gather_physical=lambda idx: ( np.asarray(params[idx]) * scale if has_scaler - else np.asarray(params[idx]) + else ( + np.asarray( + self.bijector.inverse(np.asarray(params[idx]), xp=np) + ) + if has_bijector + else np.asarray(params[idx]) + ) ), step=total_steps, ) fom_history.append(best_fom) + # Optional per-step lane-death and trace diagnostics, both off + # by default (``record_lane_nan_history`` / + # ``trace_param_indices``) and both skipped entirely rather + # than computed and discarded, the same short-circuit style as + # ``has_scaler`` / ``has_clipper``. Packed to one bit per lane + # (``numpy.packbits``) rather than kept as bool arrays: at + # ``n_starts=48`` and a multi-thousand-step run this is an 8x + # reduction in ``search_internal`` size for a feature whose + # whole point is being cheap to turn on. + if self.record_lane_nan_history: + lane_value_nan_rows.append(np.packbits(~alive)) + lane_grad_nan_rows.append( + np.packbits(alive & ~np.asarray(grad_finite).astype(bool)) + ) + if self.trace_param_indices: + physical_row = np.asarray(_to_physical(params)) + trace_rows.append( + np.asarray(physical_row)[:, self.trace_param_indices] + ) + # The size of the living population at this step. Recorded as a # history because the cumulative lane counters above are # survival INTEGRALS: a dead lane keeps adding to them every @@ -1128,6 +1272,7 @@ def batched_value_and_grad(params): jnp=jnp, rng=resurrect_rng, scale=scale, + bijector=self.bijector if has_bijector else None, ) # A resurrected slot is a NEW start: its per-lane record # resets rather than conflating two starts' basins in one @@ -1161,7 +1306,11 @@ def batched_value_and_grad(params): # 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, scale=scale + vector=params, + model=model, + xp=jnp, + scale=scale, + bijector=self.bijector if has_bijector else None, ) # Per-LANE, not per-coordinate: a lane clipped in three # parameters at once is one clipped lane-step, matching how @@ -1231,19 +1380,27 @@ def batched_value_and_grad(params): 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`` + # search stepped in scaled/transformed ones. Neither the scaler nor + # the bijector enters the search identifier, so the same output + # directory can be written by a scaled/transformed run and read by + # an untransformed 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. + # has no scaler/bijector to consult. The scale vector / bijector + # kinds ride along so a reader can see what was used without + # inferring it. + if has_scaler: + params_physical = np.asarray(params) * scale + elif has_bijector: + params_physical = np.asarray(self.bijector.inverse(params, xp=np)) + else: + params_physical = np.asarray(params) + search_internal = { - "params": ( - np.asarray(params) * scale if has_scaler else np.asarray(params) - ), + "params": params_physical, "scale": scale, + "bijector": self.bijector.kinds if has_bijector else None, "opt_state": opt_state, "best_params": best_params, "best_fom": best_fom, @@ -1261,6 +1418,30 @@ def batched_value_and_grad(params): "n_constrained_lane_steps": n_constrained_lane_steps, "n_clipped_lane_steps": n_clipped_lane_steps, "stop_reason": stop_reason, + # Optional per-step diagnostics (off by default -- see + # ``record_lane_nan_history`` / ``trace_param_indices``). + # ``None`` when off, rather than an empty array, so a reader can + # tell "not recorded" from "recorded, zero steps so far". + "lane_value_nan_history": ( + np.stack(lane_value_nan_rows) + if self.record_lane_nan_history and lane_value_nan_rows + else None + ), + "lane_grad_nan_history": ( + np.stack(lane_grad_nan_rows) + if self.record_lane_nan_history and lane_grad_nan_rows + else None + ), + "trace_history": ( + np.stack(trace_rows) + if self.trace_param_indices and trace_rows + else None + ), + "trace_param_indices": ( + list(self.trace_param_indices) + if self.trace_param_indices + else None + ), } self.paths.save_search_internal(obj=search_internal) @@ -1364,7 +1545,17 @@ 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, scale=None + self, + params, + opt_state, + dead_idx, + model, + optimizer, + jax, + jnp, + rng, + scale=None, + bijector=None, ): """ Redraw the dead starts (``dead_idx``) and reinitialise their per-start @@ -1381,9 +1572,12 @@ def _reinit_dead_starts( ``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. + When ``bijector`` is given instead, the redraw is mapped through its + ``forward`` the same way. ``scale`` and ``bijector`` are mutually + exclusive, enforced by the caller's constructor. 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] @@ -1395,7 +1589,12 @@ def _reinit_dead_starts( 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 + if scale is not None: + params_np[k] = redrawn / scale + elif bijector is not None: + params_np[k] = np.asarray(bijector.forward(redrawn, xp=np)) + else: + params_np[k] = redrawn params = jnp.asarray(params_np) fresh_state = jax.vmap(optimizer.init)(params) @@ -1606,6 +1805,15 @@ def samples_via_internal_from( # feature was off would be indistinguishable from one written before # the feature existed. "scaler": type(self.scaler).__name__, + # Per-parameter bijector (see ``autofit.non_linear.bijector``), + # recorded by NAME for the same reasons as the scaler, plus the + # resolved per-coordinate KINDS -- read from ``search_internal`` + # (persisted at write-back) rather than from ``self.bijector`` + # directly, so this is correct even when called on a + # ``search_internal`` loaded in a process that never ran ``_fit`` + # (and so never called ``self.bijector.from_model``). + "bijector": type(self.bijector).__name__, + "bijector_kinds": search_internal.get("bijector"), "reset_momentum_on_clip": self.reset_momentum_on_clip, "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 From 863739030d09d363d577af28f4ebfc1b5e081d1f Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 24 Aug 2026 18:08:10 -0400 Subject: [PATCH 3/4] Add tests for AbstractBijector and its clipper/search composition test_bijector.py (numpy-only, mirroring test_scaler.py's discipline): round trip and monotonicity per kind, the log-kind closed-form log-det-Jacobian checked against a finite difference of inverse() directly, kind-selection/fallback coverage for BijectorAuto/ BijectorLogit/BijectorPerPath (including a malformed LogUniformPrior), BijectorDiagonal reducing exactly to ScalerPriorWidth's map, BijectorNone's byte-identical no-op, the objective-composed-through-the- bijector equivalence pin and its end-to-end "cannot move the MAP" companion (L-BFGS-B in raw theta space vs. unconstrained BFGS in phi space -- the raw run needs a box because (log(theta)-2)^2 is undefined for theta <= 0 and unconstrained steps stray there, which is itself the motivating case for the log bijector), and clip-commutes-with-forward. test_clipper.py: scale+bijector together raises, ClipperNone ignores a bijector, the LogUniform log-space inset fix pinned directly against the bug (a LogUniform(1e-6, 1e6) physical-relative inset lands near 1.0, fencing off virtually the whole support; the log-space inset stays within 0.1% of the true bound), project(..., bijector=...) matching a direct physical clip mapped through, and the no-bijector bounds_from_model path proven byte-unchanged. test_multi_start_gradient.py: bijector dict round-trip, the scaler+bijector construction-time raise, bijector not offered by LBFGS, and samples_info recording the bijector's class name and (read from search_internal, not live resolved state) its kinds. Full suite green: test_autofit/non_linear (673 passed, 2 skipped) and the whole test_autofit suite (2167 passed, 3 skipped). ruff clean on every touched file except autofit/__init__.py, where the new Bijector* re-export lines repeat that file's pre-existing F401/E402 baseline (258 -> 270 errors, +12 for the 12 new import lines) rather than adding a new category of lint issue. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01EDABYoH6giHXhFJUks8yd6 --- .../search/mle/test_multi_start_gradient.py | 63 +++ test_autofit/non_linear/test_bijector.py | 497 ++++++++++++++++++ test_autofit/non_linear/test_clipper.py | 107 ++++ 3 files changed, 667 insertions(+) create mode 100644 test_autofit/non_linear/test_bijector.py 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 e1ad58b0f..52db42151 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 @@ -423,6 +423,33 @@ def test__scaler_is_not_offered_by_searches_that_do_not_own_their_step_loop(): assert "scaler" not in inspect.signature(af.LBFGS.__init__).parameters +def test__dict_round_trip__bijector(): + restored = from_dict(to_dict(af.MultiStartAdam(bijector=af.BijectorAuto()))) + + assert isinstance(restored.bijector, af.BijectorAuto) + # Default is the no-op, so the step loop skips the change of variables + # entirely rather than tracing an identity `jnp.where` selection. + assert isinstance(af.MultiStartAdam().bijector, af.BijectorNone) + + +def test__scaler_and_bijector_together__raises_at_construction(): + """ + They are two different changes of variables for the same step; letting both + be non-default would mean silently picking one over the other rather than + surfacing the conflict. + """ + with pytest.raises(ValueError): + af.MultiStartAdam(scaler=af.ScalerPriorWidth(), bijector=af.BijectorAuto()) + + # Either alone is fine. + af.MultiStartAdam(scaler=af.ScalerPriorWidth()) + af.MultiStartAdam(bijector=af.BijectorAuto()) + + +def test__bijector_is_not_offered_by_searches_that_do_not_own_their_step_loop(): + assert "bijector" not in inspect.signature(af.LBFGS.__init__).parameters + + def test__samples_via_internal_from(): model = af.Model(example.Gaussian) @@ -489,6 +516,42 @@ def test__samples_via_internal_from(): assert samples.samples_info["fom_history"] == pytest.approx([-4.0, -8.0, best_fom]) assert all(isinstance(x, float) for x in samples.samples_info["fom_history"]) + # Scaler/bijector are recorded ALWAYS (including their no-op defaults), + # since neither enters the search identifier. + assert samples.samples_info["scaler"] == "ScalerNone" + assert samples.samples_info["bijector"] == "BijectorNone" + assert samples.samples_info["bijector_kinds"] is None + + +def test__samples_info__bijector_kinds_are_read_from_search_internal_not_live_state(): + """ + ``bijector_kinds`` must come from the persisted ``search_internal["bijector"]`` + rather than from ``self.bijector.kinds`` directly -- the latter would raise + (unresolved) or read stale state in a process that loaded a + ``search_internal`` without ever calling ``_fit``. + """ + model = af.Model(example.Gaussian) + best_params = np.asarray(model.vector_from_unit_vector([0.5] * model.prior_count)) + per_start_params = np.stack([best_params]) + + search = af.MultiStartAdam(n_starts=1, n_steps=1, bijector=af.BijectorAuto()) + + samples = search.samples_via_internal_from( + model=model, + search_internal={ + "params": per_start_params, + "best_params": best_params, + "best_fom": -2.0, + "total_steps": 1, + "n_resurrections": 0, + "fom_history": np.asarray([-2.0]), + "bijector": ["identity", "identity", "identity"], + }, + ) + + assert samples.samples_info["bijector"] == "BijectorAuto" + assert samples.samples_info["bijector_kinds"] == ["identity", "identity", "identity"] + def test__samples_info__stop_reason_max_steps_and_legacy_search_internal(): model = af.Model(example.Gaussian) diff --git a/test_autofit/non_linear/test_bijector.py b/test_autofit/non_linear/test_bijector.py new file mode 100644 index 000000000..9517758c5 --- /dev/null +++ b/test_autofit/non_linear/test_bijector.py @@ -0,0 +1,497 @@ +import numpy as np +import pytest + +import autofit as af +from autofit import example +from autofit.non_linear.bijector import ( + BijectorAuto, + BijectorDiagonal, + BijectorLogit, + BijectorNone, + BijectorPerPath, +) +from autofit.non_linear.clipper import ClipperPriorBox +from autofit.non_linear.scaler import ScalerNone, ScalerPriorWidth + +# Pure NumPy, deliberately -- the same discipline `test_scaler.py` uses. The +# bijector's contract is algebraic (round trip, monotonicity, the closed-form +# log-det, and that composing the objective through it cannot move the MAP) +# and none of that needs JAX. + + +def model_from(centre, normalization, sigma): + model = af.Model(example.Gaussian) + model.centre = centre + model.normalization = normalization + model.sigma = sigma + return model + + +def log_uniform_model(): + return model_from( + centre=af.UniformPrior(lower_limit=0.0, upper_limit=8.0), + normalization=af.LogUniformPrior(lower_limit=1.0e-6, upper_limit=1.0e6), + sigma=af.GaussianPrior(mean=1.0, sigma=2.0), + ) + + +# --------------------------------------------------------------------------- +# BijectorNone +# --------------------------------------------------------------------------- + + +def test__bijector_none__is_identity_for_every_coordinate(): + model = log_uniform_model() + bijector = BijectorNone().from_model(model) + + assert bijector.kinds == ["identity", "identity", "identity"] + + +def test__bijector_none__forward_and_inverse_are_bit_identical_to_the_input(): + """ + `x / 1.0 == x` and `x * 1.0 == x` exactly in IEEE 754, so this is not merely + numerically close -- it is the byte-identical no-op the module docstring + promises, matching `ScalerNone`'s guarantee. + """ + model = log_uniform_model() + bijector = BijectorNone().from_model(model) + + theta = np.array([3.0, 100.0, -1.5]) + + assert (bijector.forward(theta) == theta).all() + assert (bijector.inverse(theta) == theta).all() + + +def test__bijector_none__log_det_jacobian_is_exactly_zero(): + model = log_uniform_model() + bijector = BijectorNone().from_model(model) + + phi = np.array([3.0, 100.0, -1.5]) + + assert bijector.log_det_jacobian(phi) == 0.0 + + +def test__bijector_none__writes_no_model_info_block(): + model = log_uniform_model() + + assert BijectorNone().info_from_model(model) == "" + + +def test__bijector_none__is_the_default_on_multi_start_gradient(): + assert isinstance(af.MultiStartAdam().bijector, BijectorNone) + + +# --------------------------------------------------------------------------- +# BijectorAuto: kind selection +# --------------------------------------------------------------------------- + + +def test__bijector_auto__picks_log_for_log_uniform_and_log_gaussian__identity_elsewhere(): + model = af.Model(example.Gaussian) + model.centre = af.UniformPrior(lower_limit=0.0, upper_limit=8.0) + model.normalization = af.LogUniformPrior(lower_limit=1.0e-6, upper_limit=1.0e6) + model.sigma = af.LogGaussianPrior(mean=0.0, sigma=1.0) + + bijector = BijectorAuto().from_model(model) + + assert bijector.kinds == ["identity", "log", "log"] + + +def test__bijector_auto__log_uniform_with_a_non_positive_limit__falls_back_to_identity( + caplog, +): + """ + Regression / defence-in-depth: `LogUniformPrior.__init__` already raises on a + non-positive `lower_limit`, so this path is not reachable through normal + construction -- it is guarded anyway, the same "malformed prior, not merely + awkward to scale" discipline `ScalerPriorWidth._term_for` uses. + """ + + class _MalformedLogUniform(af.LogUniformPrior): + def __init__(self): + # Bypass the constructor's own validation. + super().__init__(lower_limit=1.0, upper_limit=2.0) + self.lower_limit = 0.0 + + model = af.Model(example.Gaussian) + model.centre = af.UniformPrior(lower_limit=0.0, upper_limit=1.0) + model.normalization = _MalformedLogUniform() + model.sigma = af.UniformPrior(lower_limit=0.0, upper_limit=1.0) + + bijector = BijectorAuto().from_model(model) + + assert bijector.kinds[1] == "identity" + assert "non-positive lower_limit" in caplog.text + + +# --------------------------------------------------------------------------- +# Round trip, monotonicity, closed-form log-det -- per kind. +# --------------------------------------------------------------------------- + + +class TestIdentityKind: + def test__round_trip(self): + model = model_from( + centre=af.UniformPrior(lower_limit=-10.0, upper_limit=10.0), + normalization=af.UniformPrior(lower_limit=-10.0, upper_limit=10.0), + sigma=af.UniformPrior(lower_limit=-10.0, upper_limit=10.0), + ) + bijector = BijectorAuto().from_model(model) + + theta = np.array([1.0, -2.5, 7.3]) + assert bijector.inverse(bijector.forward(theta)) == pytest.approx(theta) + + def test__log_det_jacobian_is_zero_for_unscaled_identity(self): + model = model_from( + centre=af.UniformPrior(lower_limit=-10.0, upper_limit=10.0), + normalization=af.UniformPrior(lower_limit=-10.0, upper_limit=10.0), + sigma=af.UniformPrior(lower_limit=-10.0, upper_limit=10.0), + ) + bijector = BijectorAuto().from_model(model) + + assert bijector.log_det_jacobian(np.array([1.0, -2.5, 7.3])) == pytest.approx( + 0.0 + ) + + +class TestLogKind: + def test__round_trip(self): + model = log_uniform_model() + bijector = BijectorAuto().from_model(model) + + theta = np.array([3.0, 250.0, -1.5]) + assert bijector.inverse(bijector.forward(theta)) == pytest.approx(theta) + + def test__forward_is_strictly_monotone_increasing(self): + model = log_uniform_model() + bijector = BijectorAuto().from_model(model) + + thetas = np.geomspace(1.0e-6, 1.0e6, 25) + phis = np.array( + [bijector.forward(np.array([1.0, t, 0.0]))[1] for t in thetas] + ) + + assert (np.diff(phis) > 0.0).all() + + def test__log_det_jacobian_matches_the_closed_form(self): + """ + `theta = exp(phi)` for a log-kind coordinate, so + `d theta/d phi = exp(phi) = theta`, i.e. the per-coordinate contribution + to `log_det_jacobian` is exactly `phi`, checked here against a finite + difference of `inverse` directly (rather than re-deriving the same + algebra `log_det_jacobian` uses, which would not be an independent + check). + """ + model = log_uniform_model() + bijector = BijectorAuto().from_model(model) + + phi = np.array([1.0, 2.5, 0.5]) # centre/sigma identity, normalization log + eps = 1.0e-6 + d_theta = ( + bijector.inverse(phi + np.array([0.0, eps, 0.0])) + - bijector.inverse(phi - np.array([0.0, eps, 0.0])) + ) / (2.0 * eps) + numeric_log_det = np.log(np.abs(d_theta[1])) + + # Only one coordinate is log-kind here; isolate its own contribution by + # zeroing the (known, `log(1.0) == 0`) identity contributions of the + # unscaled centre/sigma coordinates. + analytic = bijector.log_det_jacobian(phi) + + assert analytic == pytest.approx(numeric_log_det, abs=1.0e-4) + + +class TestLogitKind: + def test__round_trip(self): + 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.TruncatedGaussianPrior( + mean=0.0, sigma=1.0, lower_limit=-1.0, upper_limit=1.0 + ), + ) + bijector = BijectorLogit().from_model(model) + + assert bijector.kinds == ["logit", "logit", "logit"] + + theta = np.array([3.0, 0.1, -0.4]) + assert bijector.inverse(bijector.forward(theta)) == pytest.approx(theta) + + def test__forward_is_strictly_monotone_increasing(self): + 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=-1.0, upper_limit=1.0), + ) + bijector = BijectorLogit().from_model(model) + + thetas = np.linspace(0.0001, 7.9999, 25) + phis = np.array( + [bijector.forward(np.array([t, 0.0, 0.0]))[0] for t in thetas] + ) + + assert (np.diff(phis) > 0.0).all() + + def test__log_det_jacobian_matches_a_finite_difference(self): + 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=-1.0, upper_limit=1.0), + ) + bijector = BijectorLogit().from_model(model) + + phi = np.array([0.3, -0.7, 1.2]) + eps = 1.0e-6 + d_theta_0 = ( + bijector.inverse(phi + np.array([eps, 0.0, 0.0]))[0] + - bijector.inverse(phi - np.array([eps, 0.0, 0.0]))[0] + ) / (2.0 * eps) + d_theta_1 = ( + bijector.inverse(phi + np.array([0.0, eps, 0.0]))[1] + - bijector.inverse(phi - np.array([0.0, eps, 0.0]))[1] + ) / (2.0 * eps) + d_theta_2 = ( + bijector.inverse(phi + np.array([0.0, 0.0, eps]))[2] + - bijector.inverse(phi - np.array([0.0, 0.0, eps]))[2] + ) / (2.0 * eps) + + numeric_total = ( + np.log(np.abs(d_theta_0)) + + np.log(np.abs(d_theta_1)) + + np.log(np.abs(d_theta_2)) + ) + + assert bijector.log_det_jacobian(phi) == pytest.approx( + numeric_total, abs=1.0e-4 + ) + + def test__log_gaussian_has_no_two_sided_bound__falls_back_to_identity(self, caplog): + 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.LogGaussianPrior(mean=0.0, sigma=1.0), + ) + bijector = BijectorLogit().from_model(model) + + assert bijector.kinds[2] == "identity" + assert "no two-sided finite bound" in caplog.text + + +# --------------------------------------------------------------------------- +# BijectorPerPath +# --------------------------------------------------------------------------- + + +class TestBijectorPerPath: + def test__resolves_kinds_by_dotted_path(self): + model = log_uniform_model() + + bijector = BijectorPerPath({"normalization": "log"}).from_model(model) + + assert bijector.kinds == ["identity", "log", "identity"] + + def test__unrequested_paths_default_to_identity(self): + model = log_uniform_model() + + bijector = BijectorPerPath({}).from_model(model) + + assert bijector.kinds == ["identity", "identity", "identity"] + + def test__requesting_log_on_an_ineligible_prior_falls_back_to_identity(self, caplog): + model = log_uniform_model() + + bijector = BijectorPerPath({"centre": "log"}).from_model(model) + + assert bijector.kinds[0] == "identity" + assert "not eligible" in caplog.text + + def test__unknown_kind_falls_back_to_identity(self, caplog): + model = log_uniform_model() + + bijector = BijectorPerPath({"centre": "banana"}).from_model(model) + + assert bijector.kinds[0] == "identity" + assert "unknown kind" in caplog.text + + +# --------------------------------------------------------------------------- +# BijectorDiagonal: subsumes ScalerPriorWidth exactly. +# --------------------------------------------------------------------------- + + +class TestBijectorDiagonal: + def test__reduces_to_the_scaler_map_exactly(self): + 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() + scale = scaler.scale_from_model(model) + + bijector = BijectorDiagonal(scaler).from_model(model) + + theta = np.array([3.0, 0.1, -0.05]) + + assert bijector.forward(theta) == pytest.approx(theta / scale) + assert bijector.inverse(theta / scale) == pytest.approx(theta) + assert bijector.kinds == ["identity", "identity", "identity"] + + def test__info_from_model_delegates_to_the_wrapped_scaler(self): + 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() + + assert BijectorDiagonal(scaler).info_from_model(model) == ( + scaler.info_from_model(model) + ) + + def test__scaler_none__is_the_identity_bijector_with_unit_scale(self): + 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), + ) + bijector = BijectorDiagonal(ScalerNone()).from_model(model) + + theta = np.array([3.0, 0.1, -0.05]) + assert bijector.forward(theta) == pytest.approx(theta) + + +# --------------------------------------------------------------------------- +# The equivalence argument: composing the objective through the bijector +# cannot move the MAP. Cloned from `test_scaler.py`'s `test__objective_ +# composed_through_the_scale_is_the_SAME_objective` / `test__scaling_cannot_ +# move_the_MAP`, with one LOG coordinate added. +# --------------------------------------------------------------------------- + + +def _log_coordinate_objective(theta): + """ + A three-parameter objective matching `log_uniform_model()`'s prior order: + `theta[0]` (`UniformPrior`, identity-kind) is minimised at `3.0`, + `theta[1]` (`LogUniformPrior`, log-kind) has a log-shaped optimum at + `exp(2.0) ~ 7.39`, and `theta[2]` (`GaussianPrior`, identity-kind) is + minimised at `0.05` -- so the identity and log branches are both + exercised at once. + """ + return ( + (3.0 - theta[0]) ** 2 + + (np.log(theta[1]) - 2.0) ** 2 + + 100.0 * (0.05 - theta[2]) ** 2 + ) + + +def test__objective_composed_through_the_bijector_is_the_SAME_objective(): + """ + The algebraic pin. Composing as `f(bijector.inverse(phi))` evaluates the + physical-space objective at the transformed point, so it is equal at every + corresponding pair -- exactly, not approximately. Folding a Jacobian in + instead would move the MAP and do so SILENTLY (see the module docstring's + equivalence argument); this is the test that says which of the two was + implemented. + """ + model = log_uniform_model() + bijector = BijectorAuto().from_model(model) + + rng = np.random.default_rng(0) + for _ in range(50): + theta = rng.uniform( + np.array([-10.0, 0.5, -10.0]), np.array([10.0, 500.0, 10.0]) + ) + phi = bijector.forward(theta) + + assert _log_coordinate_objective( + bijector.inverse(phi) + ) == pytest.approx(_log_coordinate_objective(theta), rel=1.0e-9) + + +def test__bijector_cannot_move_the_MAP(): + """ + The end-to-end statement: minimise the raw objective, and minimise the + bijector-composed one and map back, and the recovered argmin agrees -- + even though the third coordinate is stepping multiplicatively (in `log`) + while the first two step linearly (`identity`). + + ``scipy.optimize.minimize`` rather than a hand-rolled fixed-step descent, + for a reason that is itself part of the point: the raw-space objective is + ``(log(theta) - 2)^2`` in the second coordinate, undefined for + ``theta <= 0``, and unconstrained BFGS's own trial steps stray there and + diverge -- exactly the failure a real physical-space search guards + against with a ``Clipper`` (see :mod:`autofit.non_linear.clipper`), and + exactly what stepping in ``phi = log(theta)`` (unconstrained, valid on + the whole real line) sidesteps entirely. The raw-space run is therefore + given the box a real search would also need (``L-BFGS-B``, + ``theta[1] > 0``); the ``phi``-space run needs no such box and uses plain + ``BFGS``. + """ + from scipy import optimize + + model = log_uniform_model() + bijector = BijectorAuto().from_model(model) + + start = np.array([8.0, 500.0, 0.5]) + expected = np.array([3.0, np.exp(2.0), 0.05]) + + unscaled = optimize.minimize( + _log_coordinate_objective, + start, + method="L-BFGS-B", + bounds=[(None, None), (1.0e-8, None), (None, None)], + ).x + + phi_start = bijector.forward(start) + transformed_phi = optimize.minimize( + lambda phi: _log_coordinate_objective(bijector.inverse(phi)), + phi_start, + method="BFGS", + ).x + transformed = bijector.inverse(transformed_phi) + + assert unscaled == pytest.approx(expected, abs=1.0e-3) + assert transformed == pytest.approx(expected, abs=1.0e-3) + assert transformed == pytest.approx(unscaled, abs=1.0e-3) + + +# --------------------------------------------------------------------------- +# Composition with the clipper: clipping commutes with `forward`. +# --------------------------------------------------------------------------- + + +def test__clip_commutes_with_forward(): + """ + Every bijector kind is monotone increasing, so clipping `theta` against + physical bounds and then mapping through `forward` must give exactly the + same result as mapping `theta` through `forward` first and then clipping + against the forward-mapped bounds. This is the property + `ClipperPriorBox.project(..., bijector=...)` relies on. + """ + model = log_uniform_model() + bijector = BijectorAuto().from_model(model) + clipper = ClipperPriorBox() + + lower, upper = clipper._inset_from_model(model, kinds=bijector.kinds) + + theta = np.array([9.0, 1.0e-8, 100.0]) # centre and normalization out of box + + clip_then_forward = bijector.forward(np.clip(theta, lower, upper)) + + lower_t, upper_t = bijector.bounds_forward(lower, upper) + forward_then_clip = np.clip(bijector.forward(theta), lower_t, upper_t) + + assert clip_then_forward == pytest.approx(forward_then_clip) + + +def test__bijector_none__forward_and_bounds_forward_are_the_identity(): + model = log_uniform_model() + bijector = BijectorNone().from_model(model) + clipper = ClipperPriorBox() + + lower, upper = clipper.bounds_from_model(model) + lower_t, upper_t = bijector.bounds_forward(lower, upper) + + assert lower_t == pytest.approx(lower) + assert upper_t == pytest.approx(upper) diff --git a/test_autofit/non_linear/test_clipper.py b/test_autofit/non_linear/test_clipper.py index 850936036..0e5a89c2e 100644 --- a/test_autofit/non_linear/test_clipper.py +++ b/test_autofit/non_linear/test_clipper.py @@ -3,6 +3,7 @@ import autofit as af from autofit import exc +from autofit.non_linear.bijector import BijectorAuto from autofit.non_linear.clipper import ClipperNone, ClipperPriorBox @@ -388,3 +389,109 @@ def test__non_bound_supporting_method_raises_rather_than_being_ignored(self): with pytest.raises(exc.SearchException): search._bounds_from(model=model) + + +class TestBijectorComposition: + """ + ``AbstractClipper.project(..., bijector=...)`` -- see the module docstring's + "Composition with a bijector" section and + ``autofit.non_linear.bijector``'s equivalence argument. + """ + + def test__scale_and_bijector_together__raises(self): + model = _model(alpha=af.UniformPrior(lower_limit=0.0, upper_limit=1.0)) + bijector = BijectorAuto().from_model(model) + + with pytest.raises(ValueError): + ClipperPriorBox().project( + vector=np.array([0.5]), + model=model, + scale=np.array([1.0]), + bijector=bijector, + ) + + def test__clipper_none__accepts_a_bijector_and_ignores_it(self): + model = _model(alpha=af.UniformPrior(lower_limit=0.0, upper_limit=1.0)) + bijector = BijectorAuto().from_model(model) + + vector = np.array([5.0]) + projected, mask = ClipperNone().project( + vector=vector, model=model, bijector=bijector + ) + + assert projected is vector + assert not mask.any() + + def test__log_uniform_inset__is_taken_in_log_space_not_physical_space(self): + """ + The bug found while wiring this up (see the ``clipper`` module + docstring). The physical-relative inset (`margin * (upper - lower)`) + is `~1e-6 + 1e-6 * (1e6 - 1e-6) ~ 1.0` for `LogUniform(1e-6, 1e6)` -- + fencing off virtually the entire support. The log-space inset must + instead land close to the ORIGINAL bound, off by a factor of + `exp(margin * log(upper / lower))`, not by an additive ~1.0. + """ + model = _model(alpha=af.LogUniformPrior(lower_limit=1.0e-6, upper_limit=1.0e6)) + bijector = BijectorAuto().from_model(model) + assert bijector.kinds == ["log"] + + clipper = ClipperPriorBox(margin=1.0e-6) + + physical_lower, physical_upper = clipper.bounds_from_model(model) + log_aware_lower, log_aware_upper = clipper._inset_from_model( + model, kinds=bijector.kinds + ) + + # The un-aware inset is the bug: it is nowhere near the original 1e-6. + assert physical_lower[0] > 0.5 + + # The kind-aware inset stays close (multiplicatively) to the original + # bound, not displaced by an O(1) physical amount. + assert log_aware_lower[0] == pytest.approx(1.0e-6, rel=1.0e-3) + assert log_aware_upper[0] == pytest.approx(1.0e6, rel=1.0e-3) + assert log_aware_lower[0] > 1.0e-6 + assert log_aware_upper[0] < 1.0e6 + + def test__project_with_bijector__matches_a_direct_physical_clip_mapped_through( + self, + ): + """ + ``clipper.project(..., bijector=...)`` clips in `phi`-space against + `bijector.bounds_forward` of the (kind-aware) inset box. Because every + bijector kind is monotone increasing this must recover EXACTLY the + same physical point as clipping directly against that same inset box + and never routing through `phi` at all. + """ + model = _model( + alpha=af.UniformPrior(lower_limit=0.0, upper_limit=8.0), + beta=af.LogUniformPrior(lower_limit=1.0e-6, upper_limit=1.0e6), + gamma=af.GaussianPrior(mean=1.0, sigma=2.0), + ) + clipper = ClipperPriorBox() + bijector = BijectorAuto().from_model(model) + + theta = np.array([9.0, 1.0e-8, 0.4]) # alpha and beta both out of box + + lower, upper = clipper._inset_from_model(model, kinds=bijector.kinds) + direct_physical_clip = np.clip(theta, lower, upper) + + phi = bijector.forward(theta) + projected_phi, mask = clipper.project(vector=phi, model=model, bijector=bijector) + recovered = bijector.inverse(projected_phi) + + assert recovered == pytest.approx(direct_physical_clip) + assert list(mask) == [True, True, False] + + def test__no_bijector__bounds_from_model_is_unaffected_by_the_log_kind_fix(self): + """ + The plain (no-bijector) `bounds_from_model` path -- used directly by + LBFGS's `_bounds_from` and by any existing caller -- must be + byte-for-byte what it was before: only `project(..., bijector=...)` + is kind-aware. + """ + model = _model(alpha=af.LogUniformPrior(lower_limit=1.0e-6, upper_limit=1.0e6)) + clipper = ClipperPriorBox(margin=1.0e-6) + + lower, upper = clipper.bounds_from_model(model) + + assert lower[0] == pytest.approx(1.0e-6 + 1.0e-6 * (1.0e6 - 1.0e-6)) From bc6c7a8989d53196c4654a5de2a9a39666ebaef1 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 24 Aug 2026 18:08:38 -0400 Subject: [PATCH 4/4] Note Prodigy's global d is unmodified by scaler or bijector Extends the existing estim_lr/d progress-line comment (already noting d is in SCALED units under a scaler, not physical ones) to cover the bijector case explicitly: under a mixed identity/log bijector, d is a single global scalar spanning coordinates that are not even the same KIND of unit across parameters. Prodigy's own estimation rule (one global d from whole-tree norms) is unmodified by either knob -- this is existing, documented Prodigy behaviour, not something this change touches or attempts to fix. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01EDABYoH6giHXhFJUks8yd6 --- .../search/mle/multi_start_gradient/search.py | 21 ++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) 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 bbf34432c..5a157b2a1 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -1350,13 +1350,20 @@ def batched_value_and_grad(params): # 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. + # Under a scaler OR a bijector this ``d`` is in whatever + # coordinates the rule is stepping in (SCALED, or the + # bijector's ``phi``), not physical units — and under a + # bijector mixing ``identity`` and ``log`` coordinates it is + # a single global scalar spanning units that are not even + # the same KIND across coordinates (a linear step size for + # one, a multiplicative log-step for another). It is + # reported as the rule's own estimate rather than converted, + # since there is no single physical value to convert it to + # either way — this is Prodigy's own design (one global + # ``d`` estimated from whole-tree norms), unmodified by + # either knob and not something this feature attempts to + # fix. Do not compare ``d`` across arms that step in + # different coordinates; compare the clip rate instead. estim_lr = optax.tree_utils.tree_get(opt_state, "estim_lr") self.logger.info(