From 96dad5191db10c04aba318786aefa619b58dc977 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 28 Aug 2026 10:00:39 -0400 Subject: [PATCH] feat: project ell_comps onto its disk with an opt-in joint clipper (#1537) `ell_comps = (e1, e2)` is physical only inside the unit disk, but its priors are two independent boxes -- a square whose corners describe no ellipse. 21.5% of the declared prior area is non-physical, and 20.1% of recorded MultiStart lane best points end there (autolens_profiling#182); 0 of the 246 lanes that reach the target basin do. No per-coordinate bound can see it: (0.8, 0.8) is inside both boxes and outside the disk. Adds the structural half of the constraint machinery, beside the existing `__model_constraint__` violation *measure*: - `__model_ball_constraints__`, a class-declared `((path, radius), ...)` naming the tuple prior confined to a disk. Duck-typed like `__model_constraint__`, so a profile library states its own geometry without inheriting from PyAutoFit. - `AbstractPriorModel.ball_constraint_index_pairs()` resolves declarations to `(index_0, index_1, radius)` triples into the physical parameter vector. Static geometry, so it is cached under an underscore-prefixed `__dict__` key (the `parameterization` convention the pytree/ModelInstance paths skip). - `af.ClipperPriorBoxJoint`, an OPT-IN `ClipperPriorBox` subclass that clips the box and then radially shrinks each declared pair onto its ball. Jittable: the factor is a `where`, the radius is compared squared, and the `sqrt` argument is substituted before the `sqrt` (the double-`where` idiom) so `grad` stays finite at the origin. Both members of a moved pair are masked, which is what lets MultiStartGradient zero the outward momentum in both coordinates. The default clipper is untouched: `ClipperPriorBox.__identifier_fields__` is pinned to `("margin", "strict_epsilon")` -- exactly what the identifier's argspec fallback inferred -- so existing search identifiers and output directories are byte-identical, and a subclass adding a constructor argument cannot silently re-key stored results. Unsupported combinations raise rather than degrading silently: `AbstractBFGS` refuses a joint clipper on a model that declares a ball (a ball is not a `scipy.optimize.Bounds`), and `AbstractMultiStartGradient` refuses one alongside a non-default scaler or bijector at construction (neither change of variables maps a disk to a disk). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_011joies4k5TdRqezPUK8YET --- autofit/__init__.py | 1 + autofit/mapper/prior_model/abstract.py | 112 +++++++ autofit/mapper/prior_model/constraint.py | 90 ++++++ autofit/mapper/prior_model/prior_model.py | 20 ++ autofit/non_linear/clipper.py | 172 +++++++++++ autofit/non_linear/search/mle/bfgs/search.py | 34 +- .../search/mle/multi_start_gradient/search.py | 27 +- .../mapper/model/test_model_constraint.py | 150 +++++++++ .../search/mle/test_multi_start_gradient.py | 125 +++++++- test_autofit/non_linear/test_clipper.py | 291 +++++++++++++++++- 10 files changed, 1018 insertions(+), 4 deletions(-) diff --git a/autofit/__init__.py b/autofit/__init__.py index 6f674f8ee..6053b1444 100644 --- a/autofit/__init__.py +++ b/autofit/__init__.py @@ -80,6 +80,7 @@ from .non_linear.clipper import AbstractClipper from .non_linear.clipper import ClipperNone from .non_linear.clipper import ClipperPriorBox +from .non_linear.clipper import ClipperPriorBoxJoint from .non_linear.scaler import AbstractScaler from .non_linear.scaler import ScalerNone from .non_linear.scaler import ScalerPriorWidth diff --git a/autofit/mapper/prior_model/abstract.py b/autofit/mapper/prior_model/abstract.py index 3dec7e519..3ec277036 100644 --- a/autofit/mapper/prior_model/abstract.py +++ b/autofit/mapper/prior_model/abstract.py @@ -855,6 +855,118 @@ def constrained_model_tuples(self): tuples = [(("",), self)] + tuples return tuples + def ball_constraint_index_pairs(self): + """ + The ``(index_0, index_1, radius)`` triples describing every pair of this + model's free parameters confined to a disk by a class-declared + ``__model_ball_constraints__``. + + The indices are into the physical parameter vector — the same ordering + ``instance_from_vector`` and + :meth:`~autofit.non_linear.clipper.AbstractClipper.project` use — so a + consumer needs nothing but the vector and this list to project onto the + ball. See :mod:`autofit.mapper.prior_model.constraint` for why the + declaration is structural rather than a second violation measure. + + This is **static geometry**: it depends only on how the model is + composed, never on the parameter values, so it is resolved once and + cached (following :attr:`parameterization`'s convention of an + underscore-prefixed ``__dict__`` key, which the pytree-flattening and + ``ModelInstance`` construction paths both skip). A search resolves it + outside its step loop and traces only the arithmetic. + + A declared ball is **skipped**, rather than raising, in two cases, both + of which are ordinary model composition rather than a mistake: + + - the declaring component has no such attribute at all, because every + coordinate was fixed to an instance (a spherical profile pins + ``ell_comps`` to ``(0.0, 0.0)``, so there is nothing to project); + - fewer than two of the coordinates are free priors, because the user + fixed one of them. A ball with a coordinate held constant is an + interval on the remainder, which is a different projection and not one + this pair-shaped list can express. + + Returns + ------- + A tuple of ``(index_0, index_1, radius)`` triples, sorted and + de-duplicated (two components sharing the same linked priors describe one + ball, not two). Empty when nothing in the model declares one. + """ + cached = self.__dict__.get("_ball_constraint_index_pairs_cache") + if cached is not None: + return cached + + from autofit.mapper.prior.tuple_prior import TuplePrior + from autofit.mapper.prior_model.constraint import ball_constraints_for + from autofit.mapper.prior_model.prior_model import Model + + models = [ + model + for _, model in self.attribute_tuples_with_type( + Model, ignore_children=False + ) + if model.has_ball_constraints + ] + if isinstance(self, Model) and self.has_ball_constraints: + models = [self] + models + + pairs = [] + + if models: + # Keyed on object identity, matching how the model holds its priors: + # `priors_ordered_by_id` and a component's `prior_tuples` return the + # same objects, so identity is exact and needs no equality semantics + # from `Prior`. + index_of = { + id(prior): index + for index, prior in enumerate(self.priors_ordered_by_id) + } + + for model in models: + for path, radius in ball_constraints_for(model.cls): + obj = model + for name in path: + obj = getattr(obj, name, None) + if obj is None: + break + + if not isinstance(obj, TuplePrior): + logger.debug( + f"{model.cls.__name__} declares a ball constraint on " + f"{path}, which is not a TuplePrior on this model " + f"(every coordinate is fixed); it is not projected." + ) + continue + + priors = [ + prior + for _, prior in sorted(obj.prior_tuples, key=lambda t: t[0]) + ] + + if len(priors) != 2: + logger.debug( + f"{model.cls.__name__} declares a ball constraint on " + f"{path}, which has {len(priors)} free priors rather " + f"than 2 on this model; it is not projected." + ) + continue + + indices = [index_of.get(id(prior)) for prior in priors] + + if None in indices: + logger.debug( + f"{model.cls.__name__} declares a ball constraint on " + f"{path}, whose priors are not in this model's own " + f"prior vector; it is not projected." + ) + continue + + pairs.append((indices[0], indices[1], radius)) + + pairs = tuple(sorted(set(pairs))) + self.__dict__["_ball_constraint_index_pairs_cache"] = pairs + return pairs + def model_constraint_from_vector(self, vector, xp=np): """ The largest constraint violation any component reports for this vector. diff --git a/autofit/mapper/prior_model/constraint.py b/autofit/mapper/prior_model/constraint.py index 1898fc187..a38ebb221 100644 --- a/autofit/mapper/prior_model/constraint.py +++ b/autofit/mapper/prior_model/constraint.py @@ -38,11 +38,48 @@ def __model_constraint__(self, xp=np): diagnostic counters that consume this today, but a magnitude is what a penalty term would need later, and it carries a usable gradient back into the valid region. Returning a bool would work for counting and then have to be redesigned. + +Declaring a ball +---------------- + +``__model_constraint__`` *measures* a violation; it cannot say how to fix one. +A search that wants to keep its parameters inside the valid region needs the +constraint's **structure**, not just its magnitude, and for the commonest case — +a pair of parameters confined to a disk, ``e0**2 + e1**2 < r**2`` — the structure +is fully described by the coordinates involved and the radius. + +``__model_ball_constraints__`` declares exactly that, as a class attribute (not a +method: it is static geometry, resolved once from the model rather than evaluated +per step): + +.. code-block:: python + + class EllProfile: + __model_ball_constraints__ = ((("ell_comps",), 0.999),) + +Each entry is a ``(path, radius)`` pair. ``path`` is a tuple of attribute names +navigating from the declaring component to the +:class:`~autofit.mapper.prior.tuple_prior.TuplePrior` whose components are the +ball's coordinates; ``radius`` is the radius the coordinates are confined within. + +This is deliberately **not** a second statement of the same fact as +``__model_constraint__``, but the projectable form of it. The measure is what a +counter or a penalty term reads; the ball is what +:class:`~autofit.non_linear.clipper.ClipperPriorBoxJoint` projects onto. A class +may declare either, or both, independently. + +The radius a class declares should be the threshold **its own maths** needs, not +the boundary of formal validity. For ``ell_comps`` those differ: the geometry is +undefined at magnitude ``1``, but the conversion to an axis ratio saturates at +``0.999``, and the annulus between them is a dead-gradient region a search can +sit in while still passing every validity check. Projecting onto the clamp keeps +lanes out of it; projecting onto ``1 - epsilon`` would park them in it. """ import numpy as np MODEL_CONSTRAINT = "__model_constraint__" +MODEL_BALL_CONSTRAINT = "__model_ball_constraints__" def declares_model_constraint(cls) -> bool: @@ -55,6 +92,59 @@ def declares_model_constraint(cls) -> bool: return callable(getattr(cls, MODEL_CONSTRAINT, None)) +def declares_ball_constraints(cls) -> bool: + """ + Whether ``cls`` declares one or more ball constraints. + + Duck-typed on the presence of a non-empty ``__model_ball_constraints__``, + exactly like :func:`declares_model_constraint`, so a profile library + describes its own geometry without inheriting from PyAutoFit. + """ + return bool(getattr(cls, MODEL_BALL_CONSTRAINT, None)) + + +def ball_constraints_for(cls) -> tuple: + """ + The ``((path, radius), ...)`` ball constraints ``cls`` declares, normalised. + + Each ``path`` is returned as a tuple of attribute names and each ``radius`` + as a float, so callers never have to defend against a class having written + a list where a tuple was expected. + + Parameters + ---------- + cls + A class, which may or may not declare ``__model_ball_constraints__``. + + Returns + ------- + An empty tuple when nothing is declared. + """ + declared = getattr(cls, MODEL_BALL_CONSTRAINT, None) + if not declared: + return () + + normalised = [] + for entry in declared: + try: + path, radius = entry + except (TypeError, ValueError) as e: + raise AssertionError( + f"{cls.__name__}.{MODEL_BALL_CONSTRAINT} entries must be " + f"(path, radius) pairs; got {entry!r}" + ) from e + + if isinstance(path, str): + raise AssertionError( + f"{cls.__name__}.{MODEL_BALL_CONSTRAINT} paths must be a tuple " + f"of attribute names, not a bare string; got {path!r}" + ) + + normalised.append((tuple(path), float(radius))) + + return tuple(normalised) + + def violation_for_instance(instance, xp=np): """ The non-negative violation measure ``instance`` reports for itself. diff --git a/autofit/mapper/prior_model/prior_model.py b/autofit/mapper/prior_model/prior_model.py index 759187c85..377e922a5 100644 --- a/autofit/mapper/prior_model/prior_model.py +++ b/autofit/mapper/prior_model/prior_model.py @@ -17,6 +17,9 @@ from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.mapper.prior_model.constraint import ( MODEL_CONSTRAINT, + MODEL_BALL_CONSTRAINT, + ball_constraints_for, + declares_ball_constraints, declares_model_constraint, ) from autofit.mapper.prior_model.util import gather_namespaces @@ -46,6 +49,17 @@ def has_model_constraint(self) -> bool: """ return declares_model_constraint(self.cls) + @property + def has_ball_constraints(self) -> bool: + """ + Whether this component's class declares one or more ball constraints. + + Resolved from ``cls`` for the same reason as + :attr:`has_model_constraint`: a model rebuilt or deserialised without + going through ``__init__`` must still report correctly. + """ + return declares_ball_constraints(self.cls) + def __str__(self): prior_string = ", ".join(map(str, self.prior_tuples)) return f"{self.name} {prior_string}" @@ -232,6 +246,12 @@ def __init__( f"{vars(cls)[MODEL_CONSTRAINT]!r}" ) + # Same reasoning for the ball declaration: a malformed `(path, radius)` + # entry fails here, where the model is built, rather than when a search + # first asks the model for its index pairs. + if MODEL_BALL_CONSTRAINT in vars(cls): + ball_constraints_for(cls) + # try: # # noinspection PyTypeChecker # register_pytree_node( diff --git a/autofit/non_linear/clipper.py b/autofit/non_linear/clipper.py index 752fefb5c..170155e8e 100644 --- a/autofit/non_linear/clipper.py +++ b/autofit/non_linear/clipper.py @@ -45,6 +45,12 @@ Both read one private ``_limits_from_model`` so the declarative and imperative views can never drift apart. +That split is also what bounds :class:`ClipperPriorBoxJoint`, the box-plus-ball +subclass: a ball has an imperative projection but no declarative ``Bounds``, so +it is available to the first consumer and explicitly refused by the second (see +:meth:`~autofit.non_linear.search.mle.bfgs.search.AbstractBFGS._bounds_from`) +rather than quietly degraded to its box. + ``project`` returns **which coordinates it clipped**, not just the new vector. That mask is what lets a caller zero optimiser momentum along clipped directions: projecting the parameters while the accumulated optimiser state keeps pushing @@ -240,6 +246,15 @@ class ClipperPriorBox(AbstractClipper): dead. """ + # Pinned explicitly rather than left to the identifier's fallback, which + # infers the fields from `__init__`'s argspec and therefore produces the + # identical `{margin, strict_epsilon}` today. The pin is what keeps it + # identical: a subclass that takes a further constructor argument would + # otherwise silently re-key every stored `ClipperPriorBox` result whose + # search shares the identifier machinery. Any subclass adding a parameter + # that changes where a lane can sit must extend this tuple deliberately. + __identifier_fields__ = ("margin", "strict_epsilon") + def __init__(self, margin: float = 1.0e-6, strict_epsilon: float = 1.0e-12): self.margin = float(margin) self.strict_epsilon = float(strict_epsilon) @@ -405,3 +420,160 @@ def project(self, vector, model, xp=np, scale=None, bijector=None): projected = xp.clip(vector, lower, upper) return projected, projected != vector + + +class ClipperPriorBoxJoint(ClipperPriorBox): + """ + The prior box, and then the class-declared **balls** inside it. + + A box prior on each of a pair of coordinates is a square, and where the model + means a disk the corners of that square are not merely unlikely, they are + **non-physical**. The canonical case is `ell_comps`: two independent + ``[-1, 1]`` priors whose valid region is ``e0**2 + e1**2 < 1``, so ``1 - pi/4`` + — 21.5% — of the declared prior area describes no ellipse at all. Nothing in + a box clipper can see that, because no per-coordinate bound can: ``(0.8, 0.8)`` + is inside both boxes and outside the disk. + + That region is not hypothetical. 20.1% of MultiStart lane best points across + the recorded `autolens_profiling` campaign end at ``|e| >= 1``, and 0 of the + 246 lanes that reach the target basin do — ending outside the disk is a + property of *failed* lanes, so the 20% is wasted budget rather than a wasted + answer (autolens_profiling#182). The reason the lanes are never pulled back + is the same reason :class:`ClipperPriorBox` exists at all: the conversion to + an axis ratio *saturates* past the clamp, so the objective is flat out there + and the gradient has nothing to say. + + Which is why this clipper **projects** rather than penalises: the projection + does not need a gradient to exist, and it moves the lane in one step from a + region where no gradient exists to the nearest point where one does. + + Where the geometry comes from + ----------------------------- + + Nothing here knows what ``ell_comps`` is. The radius and the coordinate pair + are read off the model, which reads them off the declaring class's + ``__model_ball_constraints__`` (see + :mod:`autofit.mapper.prior_model.constraint`) — so PyAutoGalaxy states its own + geometry and PyAutoFit projects onto whatever geometry it is told about. + + The order of the two projections + -------------------------------- + + The box is applied first and the balls second, never the reverse. The ball is + a *subset* of the box for any radius the box contains, so projecting onto the + ball last is what leaves the point inside both; projecting onto the box last + could push a point that was on the ball's surface back off it. It is not a + true joint projection onto the intersection — that would require iterating — + but for the case this exists for (a ball inscribed in, or well inside, its + box) the two-step projection lands in the intersection in one pass. + + Jittability + ----------- + + The radial shrink is written to survive ``jit``, ``vmap`` and ``grad``: + + - the factor is built with ``xp.where`` rather than a Python branch, so the + pair's radius is never a concrete condition; + - the radius is compared **squared**, so the ``sqrt`` is only ever needed on + the branch that is actually taken; + - the ``sqrt`` argument is itself passed through a ``where`` that substitutes + ``1.0`` on the untaken branch — the "double where" idiom. A lane sitting + exactly at the origin (where a *spherical* profile's linked components sit, + and where many priors start) would otherwise evaluate ``sqrt(0)``, whose + derivative is infinite; ``jax.grad`` computes **both** branches of a + ``where`` and multiplies the untaken one by zero, and ``0 * inf`` is + ``NaN``. The forward value is unaffected either way; only the gradient is, + and only silently. Substituting before the ``sqrt`` is the only thing that + keeps it finite; + - the divisor is additionally floored at ``tiny``, so a degenerate + ``radius=0`` declaration cannot divide by zero; + - the factors are assembled as a full multiplicative vector and applied with + a single multiply, rather than with ``.at[...].set(...)`` scatter updates, + so the traced program is a fixed handful of ops independent of how many + balls the model declares and works unchanged under ``numpy``. + + ``clipped_mask`` is set on **both** members of a projected pair, even though + a shrink of a factor ``1.0 - 1e-16`` moves one of them imperceptibly. The + mask's consumer is ``MultiStartGradient``'s momentum reset: a lane pushed out + of the disk carries outward momentum in *both* coordinates, and zeroing only + one of them leaves the pair spiralling back out on the next step. + """ + + # Divisor floor for the radial shrink. Small enough to be irrelevant for any + # coordinate the projection actually moves (a pair is only shrunk when + # `r > radius`, and every declared radius is orders of magnitude larger than + # this), and large enough that `r / tiny` cannot overflow float32. + _TINY = 1.0e-30 + + def project(self, vector, model, xp=np, scale=None, bijector=None): + projected, clipped_mask = super().project( + vector=vector, + model=model, + xp=xp, + scale=scale, + bijector=bijector, + ) + + pairs = model.ball_constraint_index_pairs() + + if not pairs: + return projected, clipped_mask + + if scale is not None or bijector is not None: + # A ball is a statement about PHYSICAL coordinates, and neither + # change of variables preserves it: a per-parameter `scale` maps the + # disk to an ellipse, and a `bijector` maps it to something with no + # closed form at all. Projecting the scaled coordinates onto a circle + # would enforce a different, silently wrong constraint -- so this + # says so rather than doing it. + raise ValueError( + f"{type(self).__name__} cannot project onto a ball while the " + "search steps in scaled or transformed coordinates: the ball is " + "a statement about physical parameters, and neither a scaler nor " + "a bijector maps a disk to a disk. Use ClipperPriorBox with the " + "scaler/bijector, or drop them to project onto the ball." + ) + + # One multiplicative factor per parameter, defaulting to 1.0, so the + # whole projection is a single elementwise multiply and every + # non-ball coordinate passes through exactly (`x * 1.0 == x`). + factor = xp.ones_like(projected) + shrunk = xp.zeros_like(projected, dtype=bool) + + for index_0, index_1, radius in pairs: + value_0 = projected[..., index_0] + value_1 = projected[..., index_1] + + r_squared = value_0**2 + value_1**2 + outside = r_squared > radius**2 + + # See the class docstring: the substitution has to happen BEFORE the + # `sqrt`, not after it, or `grad` differentiates `sqrt` at zero on + # the branch it then discards and the NaN survives the discard. + r = xp.sqrt(xp.where(outside, r_squared, xp.ones_like(r_squared))) + + pair_factor = xp.where( + outside, + radius / xp.maximum(r, self._TINY), + xp.ones_like(r), + ) + + # Built by broadcasting a one-hot selector rather than by scattering + # into `factor`, so the same code runs under numpy and under a JAX + # trace (where arrays are immutable) with no `.at` and no + # index-dependent control flow. + for index in (index_0, index_1): + selector = xp.asarray( + np.arange(projected.shape[-1]) == index, + ) + factor = xp.where(selector, pair_factor[..., None], factor) + shrunk = shrunk | selector + + projected_ball = projected * factor + + # `shrunk` marks the coordinates a ball COULD move; `factor != 1.0` marks + # the lanes it actually did. Both members of a moved pair are marked -- + # see the class docstring on the momentum reset. + moved = shrunk & (factor != 1.0) + + return projected_ball, clipped_mask | moved diff --git a/autofit/non_linear/search/mle/bfgs/search.py b/autofit/non_linear/search/mle/bfgs/search.py index e0471b079..705a63d6a 100644 --- a/autofit/non_linear/search/mle/bfgs/search.py +++ b/autofit/non_linear/search/mle/bfgs/search.py @@ -8,7 +8,11 @@ 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.clipper import AbstractClipper, ClipperNone +from autofit.non_linear.clipper import ( + AbstractClipper, + ClipperNone, + ClipperPriorBoxJoint, +) from autofit.non_linear.initializer import AbstractInitializer from autofit.non_linear.samples.sample import Sample from autofit.non_linear.samples.samples import Samples @@ -126,6 +130,20 @@ def _bounds_from(self, model): constant and returns a wrong fit with no error and no warning; at every other dimensionality it raises. Building an explicit ``Bounds`` is what makes the intent unambiguous. + + A :class:`~autofit.non_linear.clipper.ClipperPriorBoxJoint` on a model that + actually declares a ball is rejected rather than degraded to its box. A + ball is not expressible as a ``scipy.optimize.Bounds`` at all -- no + per-coordinate interval can exclude the corners of a square -- so silently + handing scipy the box alone would give the caller an + unconstrained-in-the-corners fit from a clipper they chose precisely to + constrain them, which is the same class of silent wrong answer the + ``_BOUND_SUPPORTING_METHODS`` check below exists to prevent. + + The refusal is keyed on the *model*, not on the clipper's type, because + the joint clipper is a strict no-op on a model whose classes declare no + geometry (see :class:`ClipperPriorBoxJoint`); refusing that case too would + stop the clipper being configured once for a whole pipeline. """ # Imported lazily, as everywhere else in autofit -- no module in the # package pulls scipy in at import time. @@ -134,6 +152,20 @@ def _bounds_from(self, model): if isinstance(self.clipper, ClipperNone): return None + if ( + isinstance(self.clipper, ClipperPriorBoxJoint) + and model.ball_constraint_index_pairs() + ): + raise exc.SearchException( + f"A {type(self.clipper).__name__} was passed to " + f"{type(self).__name__}, which enforces its bounds through " + "scipy, together with a model declaring a ball constraint. A " + "ball cannot be expressed as a `scipy.optimize.Bounds` -- only " + "the box could be passed on, silently dropping the ball this " + "clipper exists for. Use ClipperPriorBox here, or a search that " + "projects its own steps (e.g. af.MultiStartAdam)." + ) + if self.method not in self._BOUND_SUPPORTING_METHODS: raise exc.SearchException( f"A {type(self.clipper).__name__} was passed to a search using " 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 5a157b2a1..e732c46f7 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -12,7 +12,11 @@ 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.clipper import ( + AbstractClipper, + ClipperNone, + ClipperPriorBoxJoint, +) from autofit.non_linear.scaler import AbstractScaler, ScalerNone from autofit.non_linear.initializer import AbstractInitializer from autofit.non_linear.samples.sample import Sample @@ -358,6 +362,27 @@ def __init__( "once.)" ) + # Same reasoning, one rung along: a `ClipperPriorBoxJoint` projects onto + # a ball declared in PHYSICAL coordinates, and neither change of + # variables maps a disk to a disk (a diagonal scale gives an ellipse, a + # bijector gives something with no closed form). Surfaced here rather + # than at the first step, so a multi-hour fit does not die a minute in + # on a configuration that was wrong before it started. + if isinstance(self.clipper, ClipperPriorBoxJoint) and not ( + isinstance(self.scaler, ScalerNone) + and isinstance(self.bijector, BijectorNone) + ): + raise ValueError( + f"{type(self).__name__} received a " + f"{type(self.clipper).__name__} together with a non-default " + "`scaler` or `bijector`. The joint clipper projects onto a ball " + "declared in physical parameters, and neither change of " + "variables maps a disk to a disk -- projecting in the stepped " + "coordinates would enforce a different, silently wrong " + "constraint. Use ClipperPriorBox with the scaler/bijector, or " + "drop them to project onto the ball." + ) + self.reset_momentum_on_clip = reset_momentum_on_clip self.record_lane_nan_history = bool(record_lane_nan_history) self.trace_param_indices = ( diff --git a/test_autofit/mapper/model/test_model_constraint.py b/test_autofit/mapper/model/test_model_constraint.py index cb62a0f4e..c422789ec 100644 --- a/test_autofit/mapper/model/test_model_constraint.py +++ b/test_autofit/mapper/model/test_model_constraint.py @@ -1,9 +1,15 @@ +from typing import Tuple + import numpy as np import pytest import autofit as af +from autofit.mapper.prior.tuple_prior import TuplePrior from autofit.mapper.prior_model.constraint import ( + MODEL_BALL_CONSTRAINT, MODEL_CONSTRAINT, + ball_constraints_for, + declares_ball_constraints, declares_model_constraint, violation_for_instance, ) @@ -169,3 +175,147 @@ def test_buckets_sum_to_at_most_the_lane_count(self): ) assert (n_value_nan, n_grad_nan, n_constrained) == (1, 1, 1) assert n_value_nan + n_grad_nan + n_constrained <= len(foms) + + +class Elliptical: + """Stands in for `EllProfile`: a TUPLE parameter with a declared disk. + + `Saturating` above declares the same geometry as a violation *measure* on two + flat scalars. This one declares it structurally, on a real `TuplePrior`, + which is the only shape `ball_constraint_index_pairs` can resolve — the + declaration names a path to a tuple, not a pair of loose parameter names. + """ + + __model_ball_constraints__ = ((("ell_comps",), 0.999),) + + def __init__( + self, + ell_comps: Tuple[float, float] = (0.0, 0.0), + intensity: float = 1.0, + ): + self.ell_comps = ell_comps + self.intensity = intensity + + +def elliptical_model(free=2): + """An `Elliptical` model with `free` of its two components as priors. + + `free=0` is the spherical case: every component fixed, so the model has no + `ell_comps` tuple prior at all. + """ + model = af.Model(Elliptical) + model.intensity = af.UniformPrior(lower_limit=0.0, upper_limit=10.0) + + if free == 0: + return model + + tuple_prior = TuplePrior() + tuple_prior.ell_comps_0 = af.UniformPrior(lower_limit=-1.0, upper_limit=1.0) + if free == 2: + tuple_prior.ell_comps_1 = af.UniformPrior(lower_limit=-1.0, upper_limit=1.0) + else: + tuple_prior.ell_comps_1 = 0.5 + model.ell_comps = tuple_prior + + return model + + +class TestBallDeclaration: + def test_duck_typed_detection(self): + assert declares_ball_constraints(Elliptical) + assert not declares_ball_constraints(Unconstrained) + assert not declares_ball_constraints(Saturating) + + def test_model_exposes_the_declaration(self): + assert af.Model(Elliptical).has_ball_constraints + assert not af.Model(Unconstrained).has_ball_constraints + + def test_the_two_declarations_are_independent(self): + """A class may declare a measure, a ball, or both. `Elliptical` declares + only the ball and `Saturating` only the measure, and neither is inferred + from the other.""" + assert not declares_model_constraint(Elliptical) + assert not declares_ball_constraints(Saturating) + + def test_entries_are_normalised(self): + class Listy: + __model_ball_constraints__ = [[["ell_comps"], 1]] + + assert ball_constraints_for(Listy) == ((("ell_comps",), 1.0),) + + def test_undeclared_is_empty(self): + assert ball_constraints_for(Unconstrained) == () + + def test_malformed_declaration_fails_at_composition(self): + class Malformed: + def __init__(self, centre=0.0): + self.centre = centre + + setattr(Malformed, MODEL_BALL_CONSTRAINT, ("ell_comps", 0.999)) + + with pytest.raises(AssertionError): + af.Model(Malformed) + + def test_a_bare_string_path_fails_rather_than_iterating_its_letters(self): + class Stringy: + __model_ball_constraints__ = (("ell_comps", 0.999),) + + with pytest.raises(AssertionError): + ball_constraints_for(Stringy) + + +class TestBallConstraintIndexPairs: + def test_indices_point_at_the_declared_tuple(self): + model = elliptical_model() + + ((index_0, index_1, radius),) = model.ball_constraint_index_pairs() + + names = [tuple_.name for tuple_ in model.prior_tuples_ordered_by_id] + assert names[index_0] == "ell_comps_0" + assert names[index_1] == "ell_comps_1" + assert radius == pytest.approx(0.999) + + def test_indices_are_into_the_whole_collections_vector(self): + """The offset matters: a nested component's pair must be indexed against + the vector the search actually steps, not against its own sub-model.""" + model = af.Collection(a=unconstrained_model(), b=elliptical_model()) + + ((index_0, index_1, _),) = model.ball_constraint_index_pairs() + + names = [tuple_.name for tuple_ in model.prior_tuples_ordered_by_id] + assert names[index_0] == "ell_comps_0" + assert names[index_1] == "ell_comps_1" + assert index_0 >= unconstrained_model().prior_count + + def test_one_pair_per_declaring_component(self): + model = af.Collection(a=elliptical_model(), b=elliptical_model()) + + assert len(model.ball_constraint_index_pairs()) == 2 + + def test_linked_components_describe_one_ball_not_two(self): + a = elliptical_model() + b = elliptical_model() + b.ell_comps = a.ell_comps + + assert len(af.Collection(a=a, b=b).ball_constraint_index_pairs()) == 1 + + def test_undeclared_model_has_none(self): + assert ( + af.Collection(a=unconstrained_model()).ball_constraint_index_pairs() == () + ) + + def test_every_component_fixed_is_skipped(self): + """The spherical case: `ell_comps` is pinned to an instance, so there is + no pair of free parameters to project and nothing to do.""" + assert elliptical_model(free=0).ball_constraint_index_pairs() == () + + def test_one_component_fixed_is_skipped(self): + """A ball with a coordinate held constant is an interval on the + remainder, which this pair-shaped list cannot express.""" + assert elliptical_model(free=1).ball_constraint_index_pairs() == () + + def test_is_cached(self): + model = elliptical_model() + assert ( + model.ball_constraint_index_pairs() is model.ball_constraint_index_pairs() + ) 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 52db42151..1f8625653 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 @@ -1,11 +1,12 @@ import inspect -from typing import NamedTuple +from typing import NamedTuple, Tuple import numpy as np import pytest import autofit as af from autofit import example +from autofit.mapper.prior.tuple_prior import TuplePrior from autofit.non_linear.search import abstract_search from autofit.non_linear.search.mle.multi_start_gradient.search import ( _chunk_slices, @@ -1227,3 +1228,125 @@ def test__fit_wires_the_per_lane_best_seams(): assert "lane_best_params[dead_idx] = np.nan" in source assert "lane_best_foms[dead_idx] = np.inf" in source assert "lane_best_steps[dead_idx] = total_steps" in source + + +# --------------------------------------------------------------------------- +# Joint ball clipping (PyAutoFit#1537) +# +# Kept NumPy, like the rest of this file: the projection and the momentum reset +# are both `xp`-generic, so the step loop's contract is exercised by calling +# them exactly as the loop does. The end-to-end JAX fit lives in +# autofit_workspace_test. +# --------------------------------------------------------------------------- + +BALL_RADIUS = 0.999 + + +class EllipticalLane: + """Stands in for an elliptical profile: a tuple parameter confined to a disk.""" + + __model_ball_constraints__ = ((("ell_comps",), BALL_RADIUS),) + + def __init__( + self, + ell_comps: Tuple[float, float] = (0.0, 0.0), + intensity: float = 1.0, + ): + self.ell_comps = ell_comps + self.intensity = intensity + + +def _lane_model(): + """Parameter vector ``(ell_comps_0, ell_comps_1, intensity)``. + + The `ell_comps` priors span `[-2, 2]`, deliberately wider than the disk, so + a lane at `|e| = 1.4` is inside every prior box and outside the disk — the + region the box clipper cannot see. + """ + model = af.Model(EllipticalLane) + + tuple_prior = TuplePrior() + tuple_prior.ell_comps_0 = af.UniformPrior(lower_limit=-2.0, upper_limit=2.0) + tuple_prior.ell_comps_1 = af.UniformPrior(lower_limit=-2.0, upper_limit=2.0) + model.ell_comps = tuple_prior + model.intensity = af.UniformPrior(lower_limit=0.0, upper_limit=10.0) + + return model + + +def test__a_lane_seeded_outside_the_disk_is_inside_it_after_one_step(): + """`|e| = 1.4` is what 20.1% of recorded MultiStart lane best points look + like (autolens_profiling#182): inside both prior boxes, outside the disk, and + invisible to `ClipperPriorBox`. One application of the search's own clipper — + the same call the step loop makes — has to end it inside.""" + search = af.MultiStartAdam(clipper=af.ClipperPriorBoxJoint(margin=0.0)) + model = _lane_model() + + # (n_starts, n_params), as the step loop holds it. Lane 0 is out of the disk, + # lane 1 well inside it. + params = np.array( + [ + [1.4 / np.sqrt(2.0), 1.4 / np.sqrt(2.0), 5.0], + [0.1, -0.2, 5.0], + ] + ) + + assert np.hypot(params[0, 0], params[0, 1]) == pytest.approx(1.4) + + params, clipped_mask = search.clipper.project( + vector=params, model=model, xp=np, scale=None, bijector=None + ) + + assert np.hypot(params[0, 0], params[0, 1]) <= BALL_RADIUS + assert np.hypot(params[1, 0], params[1, 1]) == pytest.approx(np.hypot(0.1, 0.2)) + + assert clipped_mask[0].tolist() == [True, True, False] + assert not clipped_mask[1].any() + + +def test__the_clipped_lanes_momentum_is_reset_in_both_components(): + """The projection alone is not enough: a lane keeps the velocity that pushed + it out of the disk, so without the reset it is projected back onto the same + surface every step. Both members of the pair have to be zeroed — zeroing one + leaves the pair spiralling straight back out.""" + + class MockAdamState(NamedTuple): + mu: np.ndarray + nu: np.ndarray + + search = af.MultiStartAdam(clipper=af.ClipperPriorBoxJoint(margin=0.0)) + model = _lane_model() + + params = np.array( + [ + [1.4 / np.sqrt(2.0), 1.4 / np.sqrt(2.0), 5.0], + [0.1, -0.2, 5.0], + ] + ) + + _, clipped_mask = search.clipper.project( + vector=params, model=model, xp=np, scale=None, bijector=None + ) + + opt_state = MockAdamState(mu=np.ones_like(params), nu=np.ones_like(params)) + + reset = search._reset_clipped_momentum( + opt_state=opt_state, clipped_mask=clipped_mask, jnp=np + ) + + assert reset.mu[0].tolist() == [0.0, 0.0, 1.0] + assert reset.mu[1].tolist() == [1.0, 1.0, 1.0] + assert reset.nu[0].tolist() == [0.0, 0.0, 1.0] + + +def test__the_default_clipper_leaves_a_lane_outside_the_disk_where_it_is(): + """The counterfactual the change exists for: neither `ClipperNone` nor + `ClipperPriorBox` can move this lane, because no per-coordinate bound + excludes the corners of a square.""" + model = _lane_model() + + params = np.array([[1.4 / np.sqrt(2.0), 1.4 / np.sqrt(2.0), 5.0]]) + + for clipper in (af.ClipperNone(), af.ClipperPriorBox(margin=0.0)): + projected, _ = clipper.project(vector=params, model=model, xp=np) + assert np.hypot(projected[0, 0], projected[0, 1]) == pytest.approx(1.4) diff --git a/test_autofit/non_linear/test_clipper.py b/test_autofit/non_linear/test_clipper.py index 61d9133ed..9cc5305ea 100644 --- a/test_autofit/non_linear/test_clipper.py +++ b/test_autofit/non_linear/test_clipper.py @@ -1,10 +1,18 @@ +from typing import Tuple + import numpy as np import pytest import autofit as af from autofit import exc +from autofit.mapper.identifier import Identifier +from autofit.mapper.prior.tuple_prior import TuplePrior from autofit.non_linear.bijector import BijectorAuto -from autofit.non_linear.clipper import ClipperNone, ClipperPriorBox +from autofit.non_linear.clipper import ( + ClipperNone, + ClipperPriorBox, + ClipperPriorBoxJoint, +) class Widget: @@ -503,3 +511,284 @@ def test__no_bijector__bounds_from_model_is_unaffected_by_the_log_kind_fix(self) lower, upper = clipper.bounds_from_model(model) assert lower[0] == pytest.approx(1.0e-6 + 1.0e-6 * (1.0e6 - 1.0e-6)) + + +# The radius `Ellipse` below declares. Deliberately not 1.0: the class it stands +# in for (`autogalaxy.profiles.geometry_profiles.EllProfile`) declares the +# ellipticity CLAMP, past which its own conversion saturates, rather than the +# boundary of formal validity. +BALL_RADIUS = 0.999 + + +class Ellipse: + """A component declaring a ball on its `ell_comps` tuple prior. + + Stands in for an elliptical profile without importing one: the clipper knows + nothing about ellipticity, only about `__model_ball_constraints__`. + """ + + __model_ball_constraints__ = ((("ell_comps",), BALL_RADIUS),) + + def __init__( + self, + ell_comps: Tuple[float, float] = (0.0, 0.0), + intensity: float = 1.0, + ): + self.ell_comps = ell_comps + self.intensity = intensity + + +def _ball_model(lower_limit=-1.0, upper_limit=1.0): + """An `Ellipse` model whose parameter vector is + ``(ell_comps_0, ell_comps_1, intensity)``. + + Priors are set explicitly because the fixture class has no entry in the + priors config, and the tuple prior is built by hand because a class with no + config gets its tuple default as an instance rather than as priors. + """ + model = af.Model(Ellipse) + + tuple_prior = TuplePrior() + tuple_prior.ell_comps_0 = af.UniformPrior( + lower_limit=lower_limit, upper_limit=upper_limit + ) + tuple_prior.ell_comps_1 = af.UniformPrior( + lower_limit=lower_limit, upper_limit=upper_limit + ) + model.ell_comps = tuple_prior + model.intensity = af.UniformPrior(lower_limit=0.0, upper_limit=10.0) + + return model + + +class TestJointBall: + def test__a_corner_of_the_box_is_projected_onto_the_disk(self): + """`(-1, -1)` is inside BOTH prior boxes and a factor `sqrt(2)` outside + the disk -- the region no per-coordinate bound can see, and the whole + reason this clipper exists.""" + model = _ball_model() + + projected, _ = ClipperPriorBoxJoint(margin=0.0).project( + vector=np.array([-1.0, -1.0, 5.0]), model=model + ) + + assert np.hypot(projected[0], projected[1]) == pytest.approx(BALL_RADIUS) + + def test__the_projection_preserves_the_angle(self): + """A radial shrink, not a per-coordinate clip: the lane keeps the + ellipse orientation it had found and loses only the magnitude the + geometry cannot support.""" + # A wider box than the disk, so the point under test is outside the ball + # and inside the box -- otherwise the box clip lands first and there is + # no angle left for the ball to preserve. + model = _ball_model(lower_limit=-2.0, upper_limit=2.0) + + projected, _ = ClipperPriorBoxJoint(margin=0.0).project( + vector=np.array([-1.2, 1.6, 5.0]), model=model + ) + + assert projected[0] / projected[1] == pytest.approx(-0.6 / 0.8) + assert projected[0] == pytest.approx(-0.6 * BALL_RADIUS) + assert projected[1] == pytest.approx(0.8 * BALL_RADIUS) + + def test__an_interior_point_is_bit_identical(self): + model = _ball_model() + vector = np.array([0.1, 0.2, 5.0]) + + projected, mask = ClipperPriorBoxJoint(margin=0.0).project( + vector=vector, model=model + ) + + assert (projected == vector).all() + assert not mask.any() + + def test__the_origin_is_not_nan(self): + """`r = 0` is where a spherical profile's linked components sit and where + many priors start. The `sqrt` and the division must both survive it.""" + model = _ball_model() + + projected, mask = ClipperPriorBoxJoint(margin=0.0).project( + vector=np.array([0.0, 0.0, 5.0]), model=model + ) + + assert np.isfinite(projected).all() + assert projected[0] == 0.0 + assert projected[1] == 0.0 + assert not mask.any() + + def test__the_mask_is_set_on_both_members_and_only_them(self): + """The mask's consumer is the momentum reset. A lane pushed out of the + disk carries outward momentum in BOTH coordinates; zeroing one leaves the + pair spiralling straight back out.""" + model = _ball_model() + + _, mask = ClipperPriorBoxJoint(margin=0.0).project( + vector=np.array([-1.0, -1.0, 5.0]), model=model + ) + + assert mask[0] + assert mask[1] + assert not mask[2] + + def test__batched_input_projects_per_lane(self): + model = _ball_model() + + vector = np.array( + [ + [-1.0, -1.0, 5.0], + [0.1, 0.2, 5.0], + [0.0, 0.0, 1.0], + ] + ) + + projected, mask = ClipperPriorBoxJoint(margin=0.0).project( + vector=vector, model=model + ) + + assert np.hypot(projected[0, 0], projected[0, 1]) == pytest.approx(BALL_RADIUS) + assert (projected[1] == vector[1]).all() + assert (projected[2] == vector[2]).all() + + assert mask[0].tolist() == [True, True, False] + assert not mask[1].any() + assert not mask[2].any() + + def test__the_box_is_applied_before_the_ball(self): + """A point outside BOTH is clipped onto the box and then shrunk onto the + ball, landing inside both. Applying them the other way round would push a + point off the ball it had just been placed on.""" + model = _ball_model() + + projected, mask = ClipperPriorBoxJoint(margin=0.0).project( + vector=np.array([5.0, 0.0, 5.0]), model=model + ) + + assert projected[0] == pytest.approx(BALL_RADIUS) + assert projected[1] == 0.0 + assert mask[0] + assert mask[1] + + def test__a_model_declaring_no_ball_is_the_plain_box_clipper(self): + """The subclass is opt-in on BOTH sides: nothing happens to a model whose + classes declare no geometry, so it can be configured globally.""" + model = _model( + alpha=af.UniformPrior(lower_limit=-1.0, upper_limit=1.0), + beta=af.UniformPrior(lower_limit=-1.0, upper_limit=1.0), + ) + vector = np.array([5.0, -5.0]) + + box, box_mask = ClipperPriorBox(margin=0.0).project(vector=vector, model=model) + joint, joint_mask = ClipperPriorBoxJoint(margin=0.0).project( + vector=vector, model=model + ) + + assert (box == joint).all() + assert (box_mask == joint_mask).all() + + def test__the_prior_box_is_still_enforced(self): + """Inheriting the box is not decoration: a coordinate outside its prior + and inside the disk is still clipped.""" + model = _ball_model(lower_limit=-0.5, upper_limit=0.5) + + projected, mask = ClipperPriorBoxJoint(margin=0.0).project( + vector=np.array([0.9, 0.0, 5.0]), model=model + ) + + assert projected[0] == pytest.approx(0.5) + assert mask[0] + + def test__a_scaler_or_bijector_is_refused_rather_than_applied(self): + """A ball is a statement about physical parameters; a diagonal scale maps + the disk to an ellipse and a bijector to something with no closed form. + Projecting in the stepped coordinates would enforce a different, silently + wrong constraint.""" + model = _ball_model() + clipper = ClipperPriorBoxJoint(margin=0.0) + + with pytest.raises(ValueError): + clipper.project( + vector=np.array([-1.0, -1.0, 5.0]), + model=model, + scale=np.array([1.0, 1.0, 1.0]), + ) + + with pytest.raises(ValueError): + clipper.project( + vector=np.array([-1.0, -1.0, 5.0]), + model=model, + bijector=BijectorAuto().from_model(model=model), + ) + + +class TestJointBallIdentifier: + def test__the_base_clippers_identifier_fields_are_pinned(self): + """Pinned so that a subclass taking a further constructor argument cannot + silently re-key every stored `ClipperPriorBox` result. The pinned tuple + reproduces exactly what the identifier's argspec fallback inferred.""" + assert ClipperPriorBox.__identifier_fields__ == ("margin", "strict_epsilon") + + def test__the_joint_clipper_is_a_different_identifier(self): + """It changes where a lane can sit and therefore the result, so two runs + differing only in clipper must not share an output directory.""" + model = _ball_model() + + box = Identifier([af.MultiStartAdam(clipper=ClipperPriorBox()), model, None]) + joint = Identifier( + [af.MultiStartAdam(clipper=ClipperPriorBoxJoint()), model, None] + ) + + assert str(box) != str(joint) + + +# The `jit`/`grad`/`vmap` behaviour of the radial shrink -- in particular that the +# "double where" keeps `grad` finite at the origin, where `sqrt(0)` has an infinite +# derivative -- is deliberately NOT tested here: unit tests stay numpy-only because +# JAX is an optional dependency. It is exercised by the traced path in +# `autolens_workspace_test`, and the numpy tests above pin the same arithmetic. + + +class TestJointBallSearchWiring: + def test__lbfgs_refuses_it_rather_than_dropping_the_ball(self): + """`LBFGS` enforces bounds through scipy, which has no ball. Passing the + box alone would give an unconstrained-in-the-corners fit from a clipper + chosen precisely to constrain them.""" + search = af.LBFGS(clipper=ClipperPriorBoxJoint()) + + with pytest.raises(exc.SearchException): + search._bounds_from(model=_ball_model()) + + def test__lbfgs_accepts_it_on_a_model_that_declares_no_ball(self): + """The refusal is keyed on the model, not the clipper's type: on a model + with no declared geometry the joint clipper is the plain box clipper, so a + pipeline may configure it once without breaking its scipy searches.""" + model = _model( + alpha=af.UniformPrior(lower_limit=-1.0, upper_limit=1.0), + beta=af.UniformPrior(lower_limit=-1.0, upper_limit=1.0), + ) + search = af.LBFGS(clipper=ClipperPriorBoxJoint()) + + bounds = search._bounds_from(model=model) + + assert bounds.lb == pytest.approx(search.clipper.bounds_from_model(model)[0]) + + def test__a_scaler_is_refused_at_construction_not_at_the_first_step(self): + """A multi-hour fit must not die a minute in on a configuration that was + wrong before it started.""" + with pytest.raises(ValueError): + af.MultiStartAdam( + clipper=ClipperPriorBoxJoint(), scaler=af.ScalerPriorWidth() + ) + + def test__a_bijector_is_refused_at_construction_too(self): + with pytest.raises(ValueError): + af.MultiStartAdam(clipper=ClipperPriorBoxJoint(), bijector=BijectorAuto()) + + def test__the_plain_box_clipper_still_composes_with_a_scaler(self): + """The refusal is scoped to the joint clipper: nothing about the existing + box-plus-scaler combination changes.""" + search = af.MultiStartAdam( + clipper=ClipperPriorBox(), scaler=af.ScalerPriorWidth() + ) + + assert isinstance(search.clipper, ClipperPriorBox) + assert isinstance(search.scaler, af.ScalerPriorWidth)