diff --git a/autoarray/inversion/mappers/mapper_util.py b/autoarray/inversion/mappers/mapper_util.py index 4cbd2c346..629d8dfac 100644 --- a/autoarray/inversion/mappers/mapper_util.py +++ b/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( diff --git a/test_autoarray/inversion/pixelization/mappers/test_mapper_util.py b/test_autoarray/inversion/pixelization/mappers/test_mapper_util.py index 787aa47b9..a3c550312 100644 --- a/test_autoarray/inversion/pixelization/mappers/test_mapper_util.py +++ b/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()