Add a spec-only nanmedian fallback for Combiner.median_combine - #978
Add a spec-only nanmedian fallback for Combiner.median_combine#978mwcraig wants to merge 3 commits into
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #978 +/- ##
==========================================
+ Coverage 96.13% 96.32% +0.18%
==========================================
Files 8 9 +1
Lines 1655 1685 +30
==========================================
+ Hits 1591 1623 +32
+ Misses 64 62 -2
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Add ccdproc/_nanmedian.py, a NaN-aware median written purely in terms of the array API standard (sort, count non-NaN, gather the two middle elements), and use it from _default_median when the selected namespace has no nanmedian instead of raising RuntimeError. Works on numpy, jax, dask and array-api-strict. Fix the docs, which claimed a bottleneck -> numpy fallback that did not exist, and drop the numpy-median and pytest.skip workarounds in test_combiner.py now that a median is available on every backend. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DTN9DnPnLKK2u7knnMJ2gA
7c96d93 to
c32570b
Compare
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01RQMJZUaaxfqGDk41GLSaFK
|
Reviewed for correctness and test size. The algorithm is right: I compared the fallback to Two things worth changing before merge (inline): the NaN-ordering assumption in Not attributable to this PR: the 16 remaining strict-job failures in |
|
|
||
|
|
||
| @pytest.mark.parametrize("length", [1, 2, 3, 4, 5, 6]) | ||
| def test_nanmedian_odd_and_even_lengths(length): |
There was a problem hiding this comment.
Add a short description of the intent of each test.
There was a problem hiding this comment.
Agreed. If the parametrized rewrite in the other thread on this file is adopted, the per-case intent moves into the trailing comments on each parametrize entry (# odd/even lengths, # all-NaN column, ...) plus a one-line docstring on each of the two tests; otherwise I will add a docstring to each test_* function.
— Written by Claude at @mwcraig's direction.
mwcraig
left a comment
mwcraig
left a comment
There was a problem hiding this comment.
Inline details for the summary comment above.
| axis = axis % ndim | ||
|
|
||
| device = array_api_compat.device(x) | ||
| s = xp.sort(x, axis=axis) |
There was a problem hiding this comment.
The docstring (line 21) says NaNs sort to the end "per the array API specification", but the standard leaves NaN ordering in sort implementation-defined. It works today only because strict/jax/dask all mirror numpy's [R, nan] order. A one-line hardening removes the dependency entirely:
| s = xp.sort(x, axis=axis) | |
| s = xp.sort( | |
| xp.where(xp.isnan(x), xp.asarray(xp.inf, dtype=x.dtype, device=device), x), | |
| axis=axis, | |
| ) |
Every position < n is then a real value no matter where the backend puts NaNs (real +inf values still land at the end, and only positions < n are ever read). If you'd rather not, please at least fix the docstring claim.
There was a problem hiding this comment.
Confirmed: the array API sort spec lists NaN ordering as implementation-defined, so the docstring claim is wrong as written and the current code only works because numpy/strict/jax/dask all happen to put NaN last. I will take the xp.where(isnan, inf, x) hardening — it costs one extra where on an array we already compute isnan for, and real +inf entries are still counted in n and land after every finite value, so positions < n stay correct. Docstring will be updated to say NaNs are replaced by +inf before sorting.
— Written by Claude at @mwcraig's direction.
| s = xp.sort(x, axis=axis) | ||
|
|
||
| # Number of non-NaN values along the axis, kept broadcastable. | ||
| n = xp.sum(xp.astype(~xp.isnan(x), xp.int64), axis=axis, keepdims=True) |
There was a problem hiding this comment.
Requesting int64 explicitly trips jax's x64-truncation UserWarning when JAX_ENABLE_X64 is unset, which filterwarnings = ["error"] turns into a failure — 12/15 of the new tests fail in that configuration. CI only passes because tox.ini sets JAX_ENABLE_X64=True for the jax job. The count never needs 64 bits:
| n = xp.sum(xp.astype(~xp.isnan(x), xp.int64), axis=axis, keepdims=True) | |
| n = xp.sum(xp.astype(~xp.isnan(x), xp.int32), axis=axis, keepdims=True) |
(The float64 promotion on line 56 has the same fragility for integer input, but that path is much rarer; fine to leave.)
There was a problem hiding this comment.
Verified locally: with JAX_ENABLE_X64 unset and CCDPROC_ARRAY_LIBRARY=jax, test_nanmedian.py gives 12 failed / 3 passed, all from the x64-truncation UserWarning being promoted to an error. int32 is plenty for a count along one axis, so I will switch to xp.int32. I will leave the float64 promotion for integer input as you suggest (it is also the numpy nanmedian result dtype, so matching it there is deliberate).
— Written by Claude at @mwcraig's direction.
There was a problem hiding this comment.
Took the int32 change in 64f5d91, and I ended up taking the float64 one too rather than leaving it.
The reason: once the count was int32, the promotion on line 56 was the only thing left failing without x64, so leaving it meant test_nanmedian.py still could not run clean in that configuration.
CCDPROC_ARRAY_LIBRARY=jax, JAX_ENABLE_X64 unset |
test_nanmedian.py |
|---|---|
| before either fix | 12 failed, 3 passed |
int32 only |
1 failed, 16 passed — test_nanmedian_matches_numpy[data12-0], the integer-input case |
| both | 17 passed |
With JAX_ENABLE_X64=True it is 17 passed either way, so CI would not have caught the difference.
The promotion now asks the namespace instead of hardcoding:
info = xp.__array_namespace_info__()
x = xp.astype(x, info.default_dtypes(device=device)["real floating"])I checked __array_namespace_info__().default_dtypes() is available on all four backends we test: numpy → float64, array-api-strict → array_api_strict.float64, dask → float64, jax → float32 without x64 and float64 with. This required moving the promotion below the device = array_api_compat.device(x) line, since default_dtypes takes the device.
The trade-off is real and worth stating: on jax without x64, integer input now yields float32 rather than matching numpy.nanmedian's float64. That seems clearly better than requesting a dtype the backend does not have — and on every backend that has float64 the result is unchanged. The docstring now says "promoted to the namespace's default real floating dtype" rather than naming float64.
Happy to revert this half if you would rather keep the numpy dtype match and accept that the file needs JAX_ENABLE_X64=True.
— Written by Claude at @mwcraig's direction.
| "nanmedian fallback supports only a single integer axis." | ||
| ) | ||
|
|
||
| xp = xp or array_api_compat.array_namespace(x) |
There was a problem hiding this comment.
Nit: if xp is None: is the idiom used elsewhere in the package.
There was a problem hiding this comment.
Will change to if xp is None: to match combiner.py (_ccddata_wrapper_for_array_api.py uses the or form, but the explicit check is the better idiom anyway since a namespace object should never be tested for truthiness).
— Written by Claude at @mwcraig's direction.
| assert default.keywords == {"xp": xp} | ||
| return | ||
| else: | ||
| pytest.skip(f"{xp.__name__} has no {function_name}") |
There was a problem hiding this comment.
This branch is only reachable for nanmean/nansum/nanstd on a namespace that lacks them — and on the strict job those don't get here, they raise RuntimeError inside _default_average etc. (3 failures currently). The skip is misleading; I'd drop it and let assert default is expected fail honestly.
There was a problem hiding this comment.
Agreed. _default_average/_default_sum/_default_std raise RuntimeError from inside the except AttributeError branch when the namespace lacks the function, so on a namespace without them the test fails before reaching this elif/else — the skip never actually fires. I will drop it and let assert default is expected fail directly. (Those strict-job failures are the pre-existing gap for nanmean/nansum/nanstd, out of scope for this PR which only adds the nanmedian fallback.)
— Written by Claude at @mwcraig's direction.
| else: | ||
| elif hasattr(xp, function_name): | ||
| expected = getattr(xp, function_name) | ||
| elif function_name == "nanmedian": |
There was a problem hiding this comment.
This elif and test_default_median_falls_back_without_nanmedian below test the same thing. The fake-module test runs deterministically on every job (including the coverage job codecov measures), so I'd keep that one and drop this 8-line branch.
There was a problem hiding this comment.
Agreed — test_default_median_falls_back_without_nanmedian exercises exactly this path with a stand-in namespace on every job, and every real backend in CI provides nanmedian so this branch never runs there. I will remove it.
— Written by Claude at @mwcraig's direction.
There was a problem hiding this comment.
Implemented in 64f5d91 — but I kept this branch, because the premise turned out not to hold.
array-api-strict is a CI backend (tox.ini:13, strict: CCDPROC_ARRAY_LIBRARY = array-api-strict) and it has no nanmedian, so this branch does run there — it is the only end-to-end check that a real namespace lacking nanmedian gets the fallback bound to it. The types.ModuleType stand-in in test_default_median_falls_back_without_nanmedian covers the same code path, but it cannot catch a regression in how a real namespace is probed, since it never exercises hasattr against an actual array namespace.
I removed the branch first and measured it. Isolating the change (strict, test_combiner.py):
| state | result |
|---|---|
| HEAD, before the edit | 68 failed, 23 passed, 1 xfailed |
| with this branch removed | 69 failed, 22 passed, 1 xfailed |
new _nanmedian.py, HEAD test_combiner.py |
68 failed, 23 passed, 1 xfailed |
The third row confirms the flip comes from this edit, not from the implementation changes. The failure mode is also worse than a clean assertion: getattr(xp, "nanmedian") raises AttributeError before the assert is reached.
The pytest.skip from the sibling thread is gone, so what remains is the elif hasattr(...) plus an else that asserts the fallback binding, with a comment noting that nanmean/nansum/nanstd raise RuntimeError inside default_func above and so never reach it.
Unrelated but worth knowing, since it affects numbers quoted across the #971 batch: the strict counts are not stable run-to-run. pytest-randomly reorders the suite and some strict test_combiner.py tests are order-dependent — I saw 75 failed / 12 passed and 68 failed / 23 passed from the same tree. With -p no:randomly it is consistently 68/23. The "75 → 74" figures in this PR's description were measured under a random order and are not comparable; the honest statement is that strict is unchanged by this PR.
— Written by Claude at @mwcraig's direction.
|
|
||
|
|
||
| @pytest.mark.parametrize("length", [1, 2, 3, 4, 5, 6]) | ||
| def test_nanmedian_odd_and_even_lengths(length): |
There was a problem hiding this comment.
The eight _check wrappers and two error tests can collapse into two parametrized tests — 85 → 43 lines, 15 → 17 cases, same assertions (the all-NaN column goes through the same numpy comparison with equal_nan=True instead of hand-rolled asserts). I ran this version against this branch on numpy, array-api-strict, dask and jax (x64): 17 passed on each.
Proposed test_nanmedian.py
# Licensed under a 3-clause BSD style license - see LICENSE.rst
import array_api_extra as xpx
import numpy as np
import pytest
from ccdproc._nanmedian import nanmedian
from ccdproc.conftest import testing_array_device as xp_device
from ccdproc.conftest import testing_array_library as xp
_rng = np.random.default_rng(906)
_some_nan = _rng.normal(size=(4, 3))
_some_nan[[0, 1, 2, 3], [1, 2, 0, 2]] = np.nan
@pytest.mark.filterwarnings("ignore:All-NaN slice encountered:RuntimeWarning")
@pytest.mark.parametrize(
("data", "axis"),
[
*[(_rng.normal(size=(n, 7)), 0) for n in range(1, 7)], # odd/even lengths
(np.array([3.0, 1.0, 2.0, 4.0]), 0), # 1-D
(_rng.normal(size=(5, 4, 3)), 0), # 3-D
(_some_nan, 0),
(_some_nan, 1),
(_some_nan, -1),
(np.array([[1.0, np.nan], [2.0, np.nan], [3.0, np.nan]]), 0), # all-NaN column
(np.array([[1, 4], [2, 3], [5, 6], [4, 1]]), 0), # integer input
],
)
def test_nanmedian_matches_numpy(data, axis):
result = nanmedian(xp.asarray(data, device=xp_device), axis=axis)
expected = xp.asarray(np.nanmedian(data, axis=axis), device=xp_device)
assert result.shape == expected.shape
assert xp.isdtype(result.dtype, "real floating")
assert xp.all(xpx.isclose(result, expected, equal_nan=True))
@pytest.mark.parametrize(
("axis", "error"),
[(None, NotImplementedError), ((0, 1), NotImplementedError), (2, ValueError), (-3, ValueError)],
)
def test_nanmedian_bad_axis(axis, error):
with pytest.raises(error):
nanmedian(xp.asarray(np.ones((2, 2)), device=xp_device), axis=axis)There was a problem hiding this comment.
Happy to adopt this. It keeps the same assertions (the all-NaN column now goes through isclose(..., equal_nan=True) against np.nanmedian, which is the same check expressed uniformly) and picks up two extra error cases. Two small notes on the proposed version: the module-level filterwarnings marker becomes a per-test decorator, which is fine since only the matching test can emit the all-NaN warning; and the error-case parametrize line exceeds the line length, so I will wrap it for ruff.
— Written by Claude at @mwcraig's direction.
Review feedback on astropy#978: - The array API leaves NaN ordering in `sort` implementation-defined, so relying on NaNs sorting last was not portable. Replace NaNs with `+inf` before sorting; genuine `+inf` entries compare equal to the sentinels, so the positions below `n` are unaffected. Fix the docstring, which claimed the old ordering came from the specification. - Count non-NaN entries as `int32` rather than `int64`, and promote integer input to the namespace's default real dtype rather than hardcoding `float64`. Both tripped jax's x64-truncation `UserWarning` when `JAX_ENABLE_X64` is unset, which `filterwarnings = ["error"]` turns into a failure: `test_nanmedian.py` went 12 failed / 3 passed in that configuration, and now passes in full both with and without x64. - Use the `if xp is None:` idiom for the namespace default. - Collapse `test_nanmedian.py` into two parametrized tests: 85 -> 50 lines, 15 -> 17 cases, same assertions. - Drop the `pytest.skip` for namespaces lacking nanmean/nansum/nanstd in `test_bottleneck_defaults_respect_array_namespace`; those raise `RuntimeError` inside the default_func call above, so the skip never fired and only obscured the real failure. The nanmedian branch of that test is kept: array-api-strict is a CI backend and does not provide `nanmedian`, so it is the only end-to-end check that a real namespace without it gets the fallback bound.
1 participant
Adds
ccdproc/_nanmedian.py, a NaN-aware median written only in terms of array-API-standard functions (replace NaNs with+inf, sort the axis, count non-NaN entries, gather and average the two middle elements; all-NaN slices give NaN; integer input is promoted to the namespace's default real floating dtype).Combiner._default_mediannow returns this fallback (bound to the namespace viafunctools.partial) when the namespace has nonanmedian, instead of raisingRuntimeError. The fallback is O(n log n) along the combination axis; nativenanmedian/ bottleneck are still preferred when available.Also:
docs/array_api.rst: replace the claim of a bottleneck -> numpy fallback (which did not exist) with the real chain (bottleneck for numpy ->xp.nanmedian-> spec-only fallback).test_combiner.py: drop the numpy-median/.compute()workarounds and thepytest.skipfor namespaces withoutmedian;test_bottleneck_defaults_respect_array_namespacenow checks the fallback is returned whennanmedianis absent.ccdproc/tests/test_nanmedian.py, two parametrized tests covering odd/even lengths, 1-D/2-D/3-D, partial and all-NaN columns, integer input, axis=1/-1, and rejected axes.Review follow-ups
sortis implementation-defined. The first version relied on NaNs sorting last, which the standard does not guarantee — it only worked because numpy/strict/jax/dask all happen to agree. NaNs are now replaced by+infbefore sorting, so the firstnpositions hold the non-NaN values whatever the backend does. Genuine+infentries compare equal to the sentinels and are still counted inn, so the result is unchanged for them. Verified against a namespace whosesortdeliberately places NaNs first: the old code returned[1., 2., nan]where the answer is[2., 3., 3.]; the new code is correct.int32rather thanint64, and integer input is promoted viaxp.__array_namespace_info__().default_dtypes(device=device)["real floating"]rather thanxp.float64. Both tripped jax's x64-truncationUserWarning, whichfilterwarnings = ["error"]turns into a failure: withJAX_ENABLE_X64unset,test_nanmedian.pywas 12 failed / 3 passed and is now 17 passed. On jax without x64 an integer input therefore yieldsfloat32rather than numpy'sfloat64; every backend that hasfloat64is unchanged.Test matrix
pytest ccdprocJAX_ENABLE_X64=True),pytest ccdprocJAX_ENABLE_X64unset,test_nanmedian.pyCCDPROC_LOG_ARRAY_ESCAPES=1 CCDPROC_ENFORCE_ESCAPE_BASELINE=1)test_nanmedian.pytest_combiner.pyThe 7 jax XPASS are identical on upstream/main (astropy/jax
DeprecationWarningmarkers, unrelated).The strict
test_combiner.pyfailures are pre-existing and tracked in #971; they hit theCombinerconstructor ("Nested Arrays are not allowed", being fixed in #976) or other known strict bugs, not the median. An earlier revision of this description quoted "75 failed / 12 passed -> 74 failed / 13 passed"; those numbers were measured underpytest-randomly's random order, which some stricttest_combiner.pytests are sensitive to, and are not reproducible. Measured with-p no:randomlythe count is a stable 68/23 both before and after this PR.Fixes #906
🤖 Generated with Claude Code
https://claude.ai/code/session_01DTN9DnPnLKK2u7knnMJ2gA