feat: project ell_comps onto its disk with an opt-in joint clipper (#1537)

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

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.

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.

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(