Skip to content

Keep Variance/InverseVariance uncertainty propagation in the array namespace - #975

Open
mwcraig wants to merge 2 commits into
astropy:mainfrom
mwcraig:fix-962-variance-wrappers
Open

Keep Variance/InverseVariance uncertainty propagation in the array namespace#975
mwcraig wants to merge 2 commits into
astropy:mainfrom
mwcraig:fix-962-variance-wrappers

Conversation

@mwcraig

@mwcraig mwcraig commented Aug 22, 2026

Copy link
Copy Markdown
Member

_VarianceUncertaintyWrapper and _InverseVarianceWrapper in ccdproc/_ccddata_wrapper_for_array_api.py were bare subclasses, so arithmetic on a CCDData carrying a VarianceUncertainty / InverseVariance went through astropy's default to_variance / from_variance hooks and the propagated uncertainty escaped to NumPy on JAX (and other non-NumPy backends).

This PR factors the four _propagate_add/subtract/multiply/divide overrides that _StdDevUncertaintyWrapper already had into a small _ArrayAPIPropagationMixin; each wrapper supplies namespace-aware hooks via _variance_hooks(xp):

  • StdDev: xp.square / xp.sqrt (unchanged behaviour)
  • Variance: identity / xp.asarray
  • InverseVariance: 1/x / 1/xp.asarray(x)

The two jax backend_xfail params on test_cosmicray_gain_correct (VarianceUncertainty and InverseVariance with gain_apply=True) are removed; under jax they go XPASS -> PASS.

Note: intermediates may still round-trip through NumPy inside astropy's _VariancePropagationMixin itself (it calls np.sqrt / np.abs and checks isinstance(correlation, np.ndarray)); that is upstream and tracked in #940. This PR only fixes the ccdproc-side wrapper so the result stays in the data's namespace.

Test matrix

backend pytest ccdproc
numpy 380 passed, 5 skipped
jax (JAX_ENABLE_X64=True) 368 passed, 10 skipped, 7 xpassed (same 7 XPASS as on main; 2 fewer xfails — the two removed here)
dask 370 passed, 15 skipped
dask + CCDPROC_ENFORCE_ESCAPE_BASELINE=1 370 passed, 15 skipped

ccdproc/tests/test_cosmicray.py under array-api-strict: 1 failed, 7 passed, 33 xfailed both before and after (pre-existing failure, unrelated).

Fixes #962

🤖 Generated with Claude Code

https://claude.ai/code/session_01DTN9DnPnLKK2u7knnMJ2gA

@codecov

codecov Bot commented Aug 22, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 95.00000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 96.65%. Comparing base (c2ce2e2) to head (60d9004).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
ccdproc/_ccddata_wrapper_for_array_api.py 95.00% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #975      +/-   ##
==========================================
+ Coverage   96.13%   96.65%   +0.51%     
==========================================
  Files           8        8              
  Lines        1655     1672      +17     
==========================================
+ Hits         1591     1616      +25     
+ Misses         64       56       -8     
Flag Coverage Δ
dask 95.80% <100.00%> (+0.58%) ⬆️
jax 95.92% <100.00%> (+0.58%) ⬆️
numpy 96.53% <95.00%> (+0.51%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@mwcraig
mwcraig marked this pull request as ready for review August 22, 2026 18:24
…mespace

_VarianceUncertaintyWrapper and _InverseVarianceWrapper were bare
subclasses, so their propagation went through astropy's default
to_variance/from_variance hooks and the result escaped to NumPy on JAX
and other non-NumPy backends. Factor the four _propagate_* overrides
from _StdDevUncertaintyWrapper into a shared _ArrayAPIPropagationMixin
whose subclasses supply the namespace-aware hooks, and use it for all
three wrapper classes.

Remove the two jax backend_xfail markers in test_cosmicray.py that
covered this; they now pass outright.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DTN9DnPnLKK2u7knnMJ2gA
@mwcraig
mwcraig force-pushed the fix-962-variance-wrappers branch from 088dc12 to 9549a1a Compare August 23, 2026 14:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Variance and InverseVariance uncertainty propagation escapes to numpy on jax (only the StdDev wrapper is array-API-aware)

1 participant