Skip to content

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

Merged
Jammy2211 merged 1 commit into
mainfrom
claude/autoarray-mapper-zero-signal-nan-jcck8q
Sep 10, 2026
Merged

Jammy2211 merged 1 commit into
mainfrom
claude/autoarray-mapper-zero-signal-nan-jcck8q

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

Closes #548.

adaptive_pixel_signals_from normalised the pixel signals by their maximum with
xp.where(max_sig > 0, pixel_signals / max_sig, pixel_signals). The where guards the
selection 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_sig
exactly 0 and puts a 0/0 NaN in the unselected branch.

Both divisions now use a safe denominator, matching the idiom already used for
pixel_counts two lines above. The step-8 exponentiation 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.

The diagnosis in the issue needed sharpening, and it changes what this PR claims. The
two backends already agreed in the forward pass: jnp.where selects correctly, so a JAX
forward result was finite. What actually diverged is (a) NumPy's discarded
RuntimeWarning: invalid value encountered in divide and (b) the gradient — under
jax.grad the NaN escapes the discarded branch. Six of the eight new tests pass on the
unfixed 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: nan in the originating workspace run — something there
was differentiating. That residual is called out rather than papered over.

API Changes

None — internal changes only.

Two degenerate corners of adaptive_pixel_signals_from change deliberately, both away
from a wrong answer, and neither is a signature or symbol change:

  • 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.
    (The sibling image_mesh/abstract_weighted.py already takes np.abs(adapt_data), so
    treating 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_autoarray1438 passed, 74 skipped, 0 failed (with the
    [optional] extras installed; see the note below).
  • pytest test_autoarray/inversion/pixelization/mappers/ — 27 passed.
  • Eight new tests in test_mapper_util.py: NumPy value and finiteness, a
    RuntimeWarning-as-error assertion, finiteness parametrised over
    signal_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 existing requires_jax skipif
    convention (as in test_delaunay.py).
  • Verified red before the fix: the clean-warning and finite-gradient legs fail on
    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.

Site Verdict
inversion/mappers/mapper_util.py:84 the bug — fixed here
inversion/mesh/interpolator/delaunay.py:390 safe — divides by the literal 3.0
inversion/mesh/interpolator/sibson.py:826 safe — divides by the literal 3.0
inversion/mesh/image_mesh/abstract_weighted.py:72 safe — explicit if max_value <= 0.0 early return
dataset/preprocess.py ×3, dataset/imaging/dataset.py:507 NumPy-only preprocessing; divide by a user-supplied scalar, not on the JAX path
fit/fit_util.py:251, 452, 474 same bug class, live — filed separately

fit/fit_util.py:251 (chi_squared_map_with_mask_from) is on the JAX likelihood-gradient
path and was confirmed by direct reproduction: with a masked-out pixel carrying zero noise
the forward value is finite (2.0) while jax.grad returns [2., 1., nan]. Different
module, 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, eight inversion/inversion tests fail with
ModuleNotFoundError: No module named 'numba' (w-tilde pixelized reconstructions are
disabled without it). That is an environment artifact, not a defect — installing numba
takes the suite to a clean 1438 passed. Recorded here because it briefly looked like
pre-existing breakage on main and is not.

Generated by the PyAutoLabs agent workflow.


Generated by Claude Code

…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 Jammy2211 added the pending-release PR queued for the next release build label Sep 10, 2026 — with Claude
@Jammy2211
Jammy2211 merged commit 667deed into main Sep 10, 2026
3 checks passed
@Jammy2211
Jammy2211 deleted the claude/autoarray-mapper-zero-signal-nan-jcck8q branch September 10, 2026 03:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release PR queued for the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

2 participants