fix(inversion): zero-signal adapt image gives NaN on the JAX path

autoarray/inversion/mappers/mapper_util.py

@@ -78,13 +78,39 @@ def adaptive_pixel_signals_from(
     pixel_counts = pixel_counts[:pixels]
 
     # 7) Normalize
+    #
+    # Both divisions here use a *safe denominator* rather than a `where` around
+    # the quotient. `xp.where(cond, a / b, a)` guards only the selection: the
+    # division is still evaluated for every element, so a zero denominator
+    # produces a NaN that NumPy discards along with the unselected branch but
+    # that JAX propagates (the standard `where`-inside-`grad` trap). A
+    # zero-signal adapt image — one whose emission lands nowhere near the
+    # pixels being adapted — makes `max_sig` exactly 0, and the two backends
+    # then disagree on the same input: NumPy returns a finite likelihood with a
+    # `RuntimeWarning`, JAX returns NaN.
     pixel_counts = xp.where(pixel_counts > 0, pixel_counts, 1.0)
     pixel_signals = pixel_signals / pixel_counts
     max_sig = xp.max(pixel_signals)
-    pixel_signals = xp.where(max_sig > 0, pixel_signals / max_sig, pixel_signals)
+    max_sig = xp.where(max_sig > 0, max_sig, 1.0)
+    pixel_signals = pixel_signals / max_sig
 
     # 8) Exponentiate
-    return pixel_signals**signal_scale
+    #
+    # Same trap, one step further on. `0.0 ** signal_scale` is finite forwards
+    # but its derivative, `signal_scale * 0.0 ** (signal_scale - 1)`, is
+    # infinite for `signal_scale < 1` and NaN under `jax.grad`; and a negative
+    # signal raised to a fractional power is NaN on both backends. Zero-signal
+    # pixels are the common case (a pixel the adapt image does not reach), so
+    # exponentiate a safe base and select the result for those pixels
+    # afterwards, keeping the whole line finite and differentiable.
+    #
+    # The signals are documented above as varying between 0 and 1, so a
+    # non-positive signal contributes nothing: it is given 0, not the NaN (or,
+    # for an even integer `signal_scale`, the spurious positive weight) that a
+    # direct exponentiation would return.
+    zero_signal = pixel_signals <= 0.0
+    safe_signals = xp.where(zero_signal, 1.0, pixel_signals)
+    return xp.where(zero_signal, 0.0, safe_signals**signal_scale)
 
 
 def sparse_triplets_from(

test_autoarray/inversion/pixelization/mappers/test_mapper_util.py

@@ -1,3 +1,6 @@
+import importlib.util
+import warnings
+
 import numpy as np
 import pytest
 
@@ -438,3 +441,129 @@ def test_mapped_to_source_via_mapping_matrix_from():
 
 
 #    assert (mapped_to_source == np.array([3.5, 8.0, 3.5])).all()
+
+
+# ----------------------------------------------------------------------------
+# Zero-signal adapt image: NumPy/JAX parity (PyAutoArray bug, 2026-09-06)
+#
+# `adaptive_pixel_signals_from` normalises by the maximum pixel signal. When the
+# adapt image carries no signal where the pixels are -- e.g. an adapt image
+# built from the *unlensed* source profile, a compact blob sitting where the
+# Einstein ring is not -- every pixel signal is zero and so is the maximum.
+#
+# The normalisation used to read `xp.where(max_sig > 0, pixel_signals / max_sig,
+# pixel_signals)`, which guards the *selection* but still evaluates the 0/0
+# division. NumPy discarded the resulting NaN along with the unselected branch
+# (emitting only a `RuntimeWarning`) while JAX propagated it, so the two
+# backends returned different answers for the same input: a finite likelihood
+# on NumPy, NaN on JAX, with `fitness._vmap` collapsing to the resample
+# figure-of-merit.
+# ----------------------------------------------------------------------------
+
+# jax is an `[optional]` extra and is absent on the NumPy-only matrix env, so
+# the JAX parity test skips rather than fails there (same convention as
+# test_delaunay.py).
+requires_jax = pytest.mark.skipif(
+    importlib.util.find_spec("jax") is None,
+    reason="requires jax (installed via the [optional] extras; absent on the NumPy-only matrix env)",
+)
+
+
+def _zero_signal_kwargs():
+    """Three pixels, each mapped by one sub-pixel, over an all-zero adapt image."""
+    return dict(
+        pixels=3,
+        signal_scale=1.0,
+        pix_indexes_for_sub_slim_index=np.array([[0], [1], [2]]),
+        pix_size_for_sub_slim_index=np.ones(3, dtype="int"),
+        pixel_weights=np.ones((3, 1), dtype="int"),
+        slim_index_for_sub_slim_index=np.array([0, 1, 2]),
+        adapt_data=np.zeros(3),
+    )
+
+
+def test__adaptive_pixel_signals_from__zero_signal_adapt_data__is_finite():
+    pixel_signals = aa.util.mapper.adaptive_pixel_signals_from(**_zero_signal_kwargs())
+
+    assert np.isfinite(np.asarray(pixel_signals)).all()
+    assert np.asarray(pixel_signals) == pytest.approx(np.zeros(3), abs=1.0e-10)
+
+
+def test__adaptive_pixel_signals_from__zero_signal_adapt_data__no_divide_warning():
+    # The 0/0 that used to raise this warning is the same one that became a NaN
+    # on the JAX path, so the clean-warning assertion pins the fix at its source.
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", RuntimeWarning)
+
+        aa.util.mapper.adaptive_pixel_signals_from(**_zero_signal_kwargs())
+
+
+@pytest.mark.parametrize("signal_scale", [0.5, 1.0, 2.0])
+def test__adaptive_pixel_signals_from__zero_signal__finite_for_any_signal_scale(
+    signal_scale,
+):
+    # `0.0 ** signal_scale` is finite forwards for every positive exponent, but
+    # a fractional one made the *derivative* infinite; the guard has to hold for
+    # all three without changing the value.
+    kwargs = _zero_signal_kwargs()
+    kwargs["signal_scale"] = signal_scale
+
+    pixel_signals = aa.util.mapper.adaptive_pixel_signals_from(**kwargs)
+
+    assert np.isfinite(np.asarray(pixel_signals)).all()
+    assert np.asarray(pixel_signals) == pytest.approx(np.zeros(3), abs=1.0e-10)
+
+
+@requires_jax
+def test__adaptive_pixel_signals_from__zero_signal__jax_matches_numpy():
+    import jax.numpy as jnp
+
+    kwargs = _zero_signal_kwargs()
+
+    numpy_signals = np.asarray(aa.util.mapper.adaptive_pixel_signals_from(**kwargs))
+    jax_signals = np.asarray(
+        aa.util.mapper.adaptive_pixel_signals_from(**kwargs, xp=jnp)
+    )
+
+    assert np.isfinite(jax_signals).all()
+    assert jax_signals == pytest.approx(numpy_signals, abs=1.0e-10)
+
+
+@requires_jax
+def test__adaptive_pixel_signals_from__signal_present__jax_matches_numpy():
+    # The parity has to hold where the function was already well-defined too,
+    # otherwise the zero-signal guard could pass by breaking the ordinary path.
+    import jax.numpy as jnp
+
+    kwargs = _zero_signal_kwargs()
+    kwargs["adapt_data"] = np.array([2.0, 1.0, 1.0])
+
+    numpy_signals = np.asarray(aa.util.mapper.adaptive_pixel_signals_from(**kwargs))
+    jax_signals = np.asarray(
+        aa.util.mapper.adaptive_pixel_signals_from(**kwargs, xp=jnp)
+    )
+
+    assert numpy_signals == pytest.approx(np.array([1.0, 0.5, 0.5]), abs=1.0e-10)
+    assert jax_signals == pytest.approx(numpy_signals, abs=1.0e-10)
+
+
+@requires_jax
+def test__adaptive_pixel_signals_from__zero_signal__grad_is_finite():
+    # The NaN this task exists for reached the likelihood through `grad`/`vmap`,
+    # not through the forward pass alone, so the guard is pinned there as well.
+    import jax
+    import jax.numpy as jnp
+
+    kwargs = _zero_signal_kwargs()
+    adapt_data = kwargs.pop("adapt_data")
+
+    def total_signal(data):
+        return jnp.sum(
+            aa.util.mapper.adaptive_pixel_signals_from(
+                **kwargs, adapt_data=data, xp=jnp
+            )
+        )
+
+    gradient = jax.grad(total_signal)(jnp.asarray(adapt_data))
+
+    assert np.isfinite(np.asarray(gradient)).all()