fix: compose ClipperPriorBoxJoint with a bijector/scaler instead of refusing it

autofit/non_linear/bijector.py

@@ -230,7 +230,8 @@ def _check_resolved(self):
             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."
+                "/ identity_scales -- the per-coordinate kind arrays are not "
+                "yet resolved."
             )
 
     @property
@@ -239,6 +240,34 @@ def kinds(self) -> List[str]:
         self._check_resolved()
         return [_KIND_NAMES[int(code)] for code in self._kind_code]
 
+    @property
+    def identity_scales(self) -> List[Optional[float]]:
+        """
+        The per-coordinate **linear** scale, or ``None`` where there is not one.
+
+        Entry ``i`` is ``self._scale[i]`` when coordinate ``i`` is
+        ``identity``-kind — the one case in which ``forward`` is exactly
+        ``theta / scale``, an affine map — and ``None`` for ``log`` / ``logit``,
+        where ``scale`` is not read at all and no single linear factor describes
+        the coordinate.
+
+        It exists so a consumer can ask "is this coordinate merely rescaled,
+        and by how much?" without reaching into ``_kind_code`` / ``_scale``.
+        :class:`~autofit.non_linear.clipper.ClipperPriorBoxJoint` is the caller:
+        a ball is a statement about physical coordinates, and it stays a ball
+        under a *common* linear rescale of both its members (radius ``R`` maps
+        to ``R / s``) but not otherwise — so it needs the scale where there is
+        one and an explicit "there is not" where there is not. ``None`` rather
+        than ``1.0`` for the non-identity kinds, deliberately: a ``log``
+        coordinate is not a coordinate scaled by one, and reporting it as such
+        would let a caller compose with it silently and wrongly.
+        """
+        self._check_resolved()
+        return [
+            float(scale) if int(code) == _IDENTITY else None
+            for code, scale in zip(self._kind_code, self._scale)
+        ]
+
     def forward(self, theta, xp=np):
         """
         The physical-to-transformed map, ``theta -> phi``, vectorised over the

autofit/non_linear/clipper.py

@@ -97,6 +97,11 @@
 :meth:`ClipperPriorBox._inset_from_model`); the plain, no-bijector
 ``bounds_from_model`` path is unaffected — every existing caller keeps its
 current (physical) inset exactly.
+
+A **ball** composes with the same maps under a narrower condition, resolved per
+pair rather than per map: it survives a common linear rescale of both its
+members and nothing else. See :class:`ClipperPriorBoxJoint` ("Composition with
+a scaler/bijector") and :meth:`ClipperPriorBoxJoint.pairs_in_stepped_coordinates`.
 """
 
 import logging
@@ -497,6 +502,43 @@ class ClipperPriorBoxJoint(ClipperPriorBox):
     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.
+
+    Composition with a scaler/bijector
+    ---------------------------------
+
+    A ball is a statement about **physical** coordinates, so a search stepping
+    in some other coordinates can only be given a ball projection if the ball
+    is still a ball there. That is a question about the *pair*, not about the
+    map as a whole, and it has a clean answer:
+
+    - both members ``identity``-kind under a **common** linear scale ``s``: the
+      step coordinates are ``(theta_i / s, theta_j / s)``, so
+      ``theta_i**2 + theta_j**2 <= R**2`` is exactly
+      ``phi_i**2 + phi_j**2 <= (R / s)**2``. A disk of radius ``R`` in physical
+      coordinates **is** a disk of radius ``R / s`` in the stepped ones, and
+      the radial shrink below — which is scale-covariant, being a pure
+      multiply — projects onto it correctly. The pair is accepted, with its
+      radius divided;
+    - anything else — ``log``, ``logit``, or two ``identity`` members with
+      *different* scales — is not a disk. Unequal scales give an **ellipse**,
+      whose nearest-point projection is not a radial shrink at all (it is the
+      root of a quartic), and the non-linear kinds give a region with no closed
+      form. Projecting a circle onto those coordinates would enforce a
+      different, silently wrong constraint, so :meth:`project` raises instead,
+      naming the offending index pair and the kinds it found.
+
+    This is deliberately resolved **per pair** rather than per map. A model
+    reparameterised through ``log`` on, say, an ``einstein_radius`` while its
+    ``ell_comps`` stay linear is the common case, and refusing it wholesale is
+    what left those lanes at the box corner ``|e| = 1.414`` — inside every prior
+    box, outside the disk, and in the flat region the clipper exists to escape.
+
+    The projection is **not** round-tripped through the bijector (map back to
+    physical, project, map forward). That would be correct for every kind, but
+    it would also drag the ``logit`` epsilon clamps across coordinates the ball
+    has nothing to do with, breaking the bit-identity the interior-point path
+    promises, saturating gradients at the clamped edges, and adding traced ops
+    to every step. Dividing the radius costs nothing and is exact.
     """
 
     # Divisor floor for the radial shrink. Small enough to be irrelevant for any
@@ -505,6 +547,93 @@ class ClipperPriorBoxJoint(ClipperPriorBox):
     # this), and large enough that `r / tiny` cannot overflow float32.
     _TINY = 1.0e-30
 
+    def pairs_in_stepped_coordinates(self, pairs, scale=None, bijector=None):
+        """
+        The declared ball pairs, restated in whatever coordinates the search is
+        actually stepping in — or a ``ValueError`` naming the pair that cannot
+        be restated.
+
+        Parameters
+        ----------
+        pairs
+            ``(index_0, index_1, radius)`` triples as
+            :meth:`~autofit.mapper.prior_model.abstract.AbstractPriorModel.
+            ball_constraint_index_pairs` returns them, in **physical**
+            coordinates.
+        scale
+            The per-parameter linear scale a ``scaler`` is applying, or
+            ``None``.
+        bijector
+            A **resolved** bijector (``from_model`` already called) the search
+            is stepping through, or ``None``. Mutually exclusive with ``scale``.
+
+        Returns
+        -------
+        The same triples with each ``radius`` divided by that pair's common
+        linear scale — unchanged when neither a scaler nor a bijector is in
+        play, since the common scale is then ``1.0``.
+
+        Raises
+        ------
+        ValueError
+            When a pair is not a disk in the stepped coordinates. See the class
+            docstring's "Composition with a scaler/bijector" for which pairs
+            those are and why the projection cannot be defined for them.
+
+        This is a **model-resolution-time** check, separable from
+        :meth:`project` on purpose: ``MultiStartGradient`` calls it once, as
+        soon as it has resolved its bijector against the model and before it
+        compiles a step, so a combination that cannot work dies there rather
+        than a minute into a multi-hour fit.
+        """
+        if scale is None and bijector is None:
+            return list(pairs)
+
+        if scale is not None:
+            # A raw scaler is linear by construction, so every coordinate is
+            # `identity`-kind; only the equality of the two scales is in doubt.
+            scales = [float(value) for value in np.asarray(scale, dtype=float)]
+            kinds = ["identity"] * len(scales)
+        else:
+            scales = bijector.identity_scales
+            kinds = bijector.kinds
+
+        resolved = []
+
+        for index_0, index_1, radius in pairs:
+            scale_0 = scales[index_0]
+            scale_1 = scales[index_1]
+
+            common = (
+                scale_0 is not None
+                and scale_1 is not None
+                and scale_0 == scale_1
+                and scale_0 > 0.0
+                and np.isfinite(scale_0)
+            )
+
+            if not common:
+                raise ValueError(
+                    f"{type(self).__name__} cannot project onto the ball "
+                    f"declared on parameter indices ({index_0}, {index_1}) "
+                    "while the search steps in scaled or transformed "
+                    f"coordinates: that pair maps as ({kinds[index_0]}, "
+                    f"{kinds[index_1]}) with linear scales ({scale_0}, "
+                    f"{scale_1}), which is not a disk. A ball survives a "
+                    "COMMON linear scale `s` (radius `R` becomes `R / s`) and "
+                    "is composed with; unequal scales give an ellipse and a "
+                    "log/logit coordinate a region with no closed form, and "
+                    "projecting a circle onto either would enforce a "
+                    "different, silently wrong constraint. Leave this pair "
+                    "identity-mapped under one common scale (e.g. a "
+                    "`BijectorPerPath` carrying the non-linear kinds on the "
+                    "OTHER paths), or use ClipperPriorBox instead."
+                )
+
+            resolved.append((index_0, index_1, radius / scale_0))
+
+        return resolved
+
     def project(self, vector, model, xp=np, scale=None, bijector=None):
         projected, clipped_mask = super().project(
             vector=vector,
@@ -519,20 +648,15 @@ def project(self, vector, model, xp=np, scale=None, bijector=None):
         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."
-            )
+        # Restated in the coordinates `projected` is already in (the box above
+        # returns the STEPPED vector on both the scaler and the bijector path),
+        # so the radial shrink below is unchanged -- it just shrinks onto a
+        # radius that has been divided by the pair's common scale.
+        pairs = self.pairs_in_stepped_coordinates(
+            pairs=pairs,
+            scale=scale,
+            bijector=bijector,
+        )
 
         # One multiplicative factor per parameter, defaulting to 1.0, so the
         # whole projection is a single elementwise multiply and every

autofit/non_linear/search/mle/multi_start_gradient/search.py

@@ -160,6 +160,17 @@ def __init__(
             enforcement on moves where the search converges and would shift every
             stored multi-start benchmark, so the default flip is deliberately a
             separate, re-baselined change.
+
+            ``ClipperPriorBoxJoint`` additionally projects onto the balls the
+            model's classes declare (the ``ell_comps`` disk being the canonical
+            one). It composes with ``scaler`` / ``bijector`` **per ball pair**,
+            not wholesale: a pair both of whose members the map merely rescales
+            by one common linear factor ``s`` is still a disk in the stepped
+            coordinates, of radius ``R / s``, and is projected onto normally; a
+            pair a ``log`` / ``logit`` kind or two unequal scales touch is not a
+            disk at all and raises. That check runs once against the model at
+            the start of ``_fit`` — not at construction, where there is no model
+            to ask — and always before the first likelihood evaluation.
         scaler
             A per-parameter step scale derived from the priors, applied as a
             change of variables so the rule steps in ``phi = theta / scale`` while
@@ -207,6 +218,13 @@ def __init__(
             Default ``BijectorNone`` -- a no-op, skipped entirely rather than
             applied as an identity, so the default path and its compiled step
             are both unchanged.
+
+            Composable with ``ClipperPriorBoxJoint`` where the map leaves each
+            declared ball pair identity-mapped under one common scale — so a
+            model wanting ``log`` on (say) an ``einstein_radius`` while its
+            ``ell_comps`` step linearly gets both the reparameterisation and
+            the disk. See ``clipper`` above for the condition and where it is
+            checked.
         resurrect
             Restart-on-death. When ``True``, any start whose objective goes
             non-finite is redrawn each step (fresh params from the start band +
@@ -362,26 +380,16 @@ 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."
-            )
+        # A `ClipperPriorBoxJoint` alongside a scaler/bijector is NOT rejected
+        # here, deliberately, even though the pair above is. Whether the two
+        # compose is a question about the MODEL -- which coordinate pairs carry
+        # a ball, and how the map treats each of them -- and no model exists at
+        # construction. A ball survives a common linear rescale of both its
+        # members and nothing else, so the same search+clipper+bijector triple
+        # is perfectly well defined on one model and undefinable on the next.
+        # The check is made once in `_fit`, as soon as the bijector has been
+        # resolved against the model and before any step is compiled, so a bad
+        # combination still dies before the first likelihood evaluation.
 
         self.reset_momentum_on_clip = reset_momentum_on_clip
         self.record_lane_nan_history = bool(record_lane_nan_history)
@@ -884,6 +892,27 @@ def _fit(
         if has_bijector:
             self.bijector.from_model(model=model)
 
+        # The joint clipper's ball is declared in PHYSICAL coordinates, and it
+        # is still a ball in the stepped ones only where the map rescales both
+        # its members by one common linear factor (see
+        # `ClipperPriorBoxJoint.pairs_in_stepped_coordinates`). Resolved HERE:
+        # this is the first moment both halves are known -- the bijector was
+        # resolved against the model one line above, and the model is what says
+        # which pairs carry a ball at all -- and it is still before `_vmapped`
+        # is built, so an unresolvable combination dies without a single
+        # likelihood evaluation rather than a minute into a multi-hour fit.
+        # Its return value is discarded: the step loop calls `clipper.project`,
+        # which re-derives the same restated pairs itself. What is wanted here
+        # is only the raise.
+        if isinstance(self.clipper, ClipperPriorBoxJoint) and (
+            has_scaler or has_bijector
+        ):
+            self.clipper.pairs_in_stepped_coordinates(
+                pairs=model.ball_constraint_index_pairs(),
+                scale=scale if has_scaler else None,
+                bijector=self.bijector if has_bijector else None,
+            )
+
         def _to_physical(vector):
             if scale_jnp is not None:
                 return vector * scale_jnp