fix(inversion): zero-signal adapt image gives NaN on the JAX path - #549
Merged
Merged
Conversation
…path `adaptive_pixel_signals_from` normalised by the maximum pixel signal with `xp.where(max_sig > 0, pixel_signals / max_sig, pixel_signals)`. The `where` guards the *selection* but the division is still evaluated, so a zero-signal adapt image — one whose emission lands nowhere near the pixels being adapted — made `max_sig` exactly 0 and produced a 0/0 NaN in the unselected branch. Forwards the two backends agreed (both `where` calls select the zeros), but NumPy emitted a `RuntimeWarning: invalid value encountered in divide` and, under `jax.grad`, the NaN propagated out of the discarded branch and into the likelihood — the standard `where`-inside-`grad` trap. Both divisions now use a safe denominator instead, matching the idiom already used for `pixel_counts` two lines above. The exponentiation in step 8 carries the same trap one step on: `0.0 ** signal_scale` is finite forwards but its derivative is infinite for `signal_scale < 1`, so it now exponentiates a safe base and selects zero for non-positive signals afterwards. Behaviour is unchanged wherever it was previously well-defined. Two degenerate corners change deliberately, both away from a wrong answer: - `signal_scale == 0` with a zero signal returned 1.0 (the `0 ** 0` convention) and now returns 0.0. - A negative pixel signal returned NaN for a fractional `signal_scale`, or a spurious positive weight for an even integer one, and now returns 0.0 — the signals are documented as varying between 0 and 1, so a non-positive signal contributes nothing. Tests cover the NumPy value and clean-warning case, finiteness across `signal_scale`, NumPy/JAX parity with and without signal, and a finite `jax.grad`. Verified red before the fix on the warning and gradient legs. Sweep for the same `where(cond, a / b, ...)` pattern (AST pass over the whole package, not grep): the two other hits in `autoarray.inversion` divide by the literal 3.0 and are safe, and `mesh/image_mesh/abstract_weighted.py` already carries an explicit non-positive-maximum early return. Three live hits in `autoarray/fit/fit_util.py` are the same class and are filed separately. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01J1ZtfVTWRCc5za6SmU2apo
Jammy2211
deleted the
claude/autoarray-mapper-zero-signal-nan-jcck8q
branch
September 10, 2026 03:24
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Closes #548.
adaptive_pixel_signals_fromnormalised the pixel signals by their maximum withxp.where(max_sig > 0, pixel_signals / max_sig, pixel_signals). Thewhereguards theselection but the division is still evaluated for every element, so a zero-signal adapt
image — one whose emission lands nowhere near the pixels being adapted — makes
max_sigexactly
0and puts a0/0NaN in the unselected branch.Both divisions now use a safe denominator, matching the idiom already used for
pixel_countstwo lines above. The step-8 exponentiation carries the same trap one stepon —
0.0 ** signal_scaleis finite forwards but its derivative is infinite forsignal_scale < 1— so it now exponentiates a safe base and selects zero for non-positivesignals afterwards.
The diagnosis in the issue needed sharpening, and it changes what this PR claims. The
two backends already agreed in the forward pass:
jnp.whereselects correctly, so a JAXforward result was finite. What actually diverged is (a) NumPy's discarded
RuntimeWarning: invalid value encountered in divideand (b) the gradient — underjax.gradthe NaN escapes the discarded branch. Six of the eight new tests pass on theunfixed code; only the clean-warning and finite-gradient legs go red. So this fixes a real
NaN on the gradient path, but it does not by itself explain the reported
JAX(no jit) fit.log_likelihood: nanin the originating workspace run — something therewas differentiating. That residual is called out rather than papered over.
API Changes
None — internal changes only.
Two degenerate corners of
adaptive_pixel_signals_fromchange deliberately, both awayfrom a wrong answer, and neither is a signature or symbol change:
signal_scale == 0with a zero signal returned1.0(the0 ** 0convention) and nowreturns
0.0.signal_scale, or a spuriouspositive weight for an even integer one, and now returns
0.0— the signals aredocumented as varying between 0 and 1, so a non-positive signal contributes nothing.
(The sibling
image_mesh/abstract_weighted.pyalready takesnp.abs(adapt_data), sotreating negative adapt signal as carrying no positive weight is consistent with the
surrounding code.)
Behaviour is otherwise identical wherever it was previously well-defined.
Test Plan
pytest test_autoarray— 1438 passed, 74 skipped, 0 failed (with the[optional]extras installed; see the note below).pytest test_autoarray/inversion/pixelization/mappers/— 27 passed.test_mapper_util.py: NumPy value and finiteness, aRuntimeWarning-as-error assertion, finiteness parametrised oversignal_scale in {0.5, 1.0, 2.0}, NumPy/JAX parity both with and without signal,and a finite
jax.grad. JAX legs carry the repo's existingrequires_jaxskipifconvention (as in
test_delaunay.py).the unfixed code and pass after.
Pattern sweep — issue #548 ask (3)
Done as an AST pass over the whole package, not grep, so multi-line calls could not
slip past: every
where(cond, x, y)with a division anywhere inside either branch.inversion/mappers/mapper_util.py:84inversion/mesh/interpolator/delaunay.py:3903.0inversion/mesh/interpolator/sibson.py:8263.0inversion/mesh/image_mesh/abstract_weighted.py:72if max_value <= 0.0early returndataset/preprocess.py×3,dataset/imaging/dataset.py:507fit/fit_util.py:251, 452, 474fit/fit_util.py:251(chi_squared_map_with_mask_from) is on the JAX likelihood-gradientpath and was confirmed by direct reproduction: with a masked-out pixel carrying zero noise
the forward value is finite (
2.0) whilejax.gradreturns[2., 1., nan]. Differentmodule, own test surface, so it is filed as its own PyAutoMind prompt rather than widening
this PR.
Note on running the suite locally
Without the
[optional]extras, eightinversion/inversiontests fail withModuleNotFoundError: No module named 'numba'(w-tilde pixelized reconstructions aredisabled without it). That is an environment artifact, not a defect — installing
numbatakes the suite to a clean 1438 passed. Recorded here because it briefly looked like
pre-existing breakage on
mainand is not.Generated by the PyAutoLabs agent workflow.
Generated by Claude Code