Hoist None/tuple axis handling into _nanfuncs._setup - #1006
Conversation
Move the flatten-for-None and permute-and-merge steps from core._mad_fallback into the shared _nanfuncs._setup, as agreed in the review of astropy#1000. _setup now accepts a single integer, None, or a tuple/list of axes, and returns a restore callable that maps a full-shape array back to the caller's layout (the identity for a single integer axis). This gives every reduction fallback in _nanfuncs (nansum, nanmean, nanstd, nanmedian, median, nanmad) tuple-axis and axis=None support for free, shrinks _mad_fallback to the median tier plus a delegation, and lifts the single-integer-axis restriction of Combiner.sigma_clipping's array-API path, which now accepts the same axis forms astropy.stats.sigma_clip does (the restore callable hands the clip mask back in the data's own shape). Bare-bool and non-integer axes now raise TypeError instead of NotImplementedError, matching the TypeError a bool entry in a tuple already raised. Closes astropy#1004 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LVxnTLrWKStBdPwcmDNhxA
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1006 +/- ##
==========================================
+ Coverage 97.96% 97.97% +0.01%
==========================================
Files 9 9
Lines 1912 1927 +15
==========================================
+ Hits 1873 1888 +15
Misses 39 39
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:
|
| which :func:`~astropy.stats.sigma_clip` accepts, raise | ||
| `NotImplementedError` there. ``masked`` and ``return_bounds`` | ||
| are never accepted, on any array namespace -- this method | ||
| honoured for every array namespace; like |
There was a problem hiding this comment.
American spelling, please
There was a problem hiding this comment.
Agreed — and this diff has more of them than this line: honoured here, Normalise/normalised/axis normalisation in _nanfuncs.py (module comment and _setup docstring), and normalised/honour in test_combiner.py comments. Will switch all of them to American spelling on the branch (honored, normalize/normalized/normalization).
— Written by Claude at @mwcraig's direction.
| Validate ``axis``, resolve the namespace and device, promote to float. | ||
| Normalise ``axis``, resolve the namespace and device, promote to float. | ||
|
|
||
| ``axis`` may be a single integer, ``None`` or a tuple/list of integers. |
There was a problem hiding this comment.
Agreed — the paragraph explaining how the axis forms are handled (flatten for None, permute-and-merge for a tuple/list) is extended discussion, not signature summary, and numpydoc puts that in Notes. Will move it into a Notes section of _setup's docstring, keeping the one-line summary up top.
— Written by Claude at @mwcraig's direction.
| Axis along which to sum. Default is 0. ``None`` and tuples of axes | ||
| are not supported. | ||
| axis : int, tuple of int, list of int or None, optional | ||
| Axis or axes along which to sum. Default is 0. ``None`` sums over |
There was a problem hiding this comment.
Can this be pulled out some how since it is repeated so many times? Or maybe factor out most of the docstring with function-specific substitutions?
There was a problem hiding this comment.
Yes — the x/axis/xp Parameters blocks of the six public functions are identical up to the verb ("to sum" / "over which to average" / ...) and the default, so they can be generated. Concrete shape (the scipy doccer pattern, host-side only):
_COMMON_PARAMS = """x : array
Input array. Integer and boolean inputs are promoted to the
namespace's default real floating dtype.
axis : int, tuple of int, list of int or None, optional
Axis or axes along which {verb}. Default is 0. ``None`` {verb_none}
over every axis; a tuple or list {verb_plural} over all the listed
axes at once.
xp : array namespace, optional
Namespace to use. Defaults to ``array_api_compat.array_namespace(x)``."""
def _fill_doc(**subs):
def deco(func):
if func.__doc__: # python -OO strips docstrings
func.__doc__ = func.__doc__.format(params=_COMMON_PARAMS.format(**subs))
return func
return decowith each function keeping its own summary line and Returns block inline — those differ meaningfully (all-NaN slices sum to zero for nansum but give NaN for nanmean/nanstd, median propagates NaN, ...), and templating them would obscure exactly the part a reader needs.
Trade-off to sign off on: the raw source (and editors that read source rather than __doc__) shows {params} at the definition; Sphinx/numpydoc and help() render the filled version. If that's acceptable, happy to implement; if not, my honest alternative is leaving the six short blocks as-is — they're repetitive but each is self-contained.
— Written by Claude at @mwcraig's direction.
| return xp.reshape(x, (-1,)), 0, xp, device, restore | ||
|
|
||
| if isinstance(axis, tuple | list): | ||
| # normalize_axis_tuple would silently treat True as 1. |
There was a problem hiding this comment.
Explain in more detail what is going on in this block -- no comments in the code, just explain in the reply
There was a problem hiding this comment.
Step by step:
- The bool guard.
normalize_axis_tupleconverts entries withoperator.index, and a PythonTruepasses that (operator.index(True) == 1), soaxis=(0, True)would silently reduce axes(0, 1). The guard turns that into aTypeErrorbefore the conversion runs. normalize_axis_tuple(axis, ndim)does the host-side validation and normalisation in one call: each entry throughoperator.index, negatives wrapped modndim(so-1on a 3-D array becomes2),AxisErroron out-of-bounds,ValueErroron duplicates — including duplicates via a negative alias, e.g.(0, -3)on 3-D. Out comes a tuple of unique non-negative ints.keptis the axes not being reduced, in their original order;order = kept + axesis then a permutation ofrange(ndim)that moves every reduced axis to the end.permute_dims(x, order)rearranges the array into that layout, and thereshapecollapses the trailing reduced axes into one axis whose length is the product of their sizes. Every reduction in this module (sum, mean, std, median, mad) is insensitive to the order of elements within the reduced set, so how the merge interleaves them is irrelevant to the result.- The caller then runs its ordinary single-axis reduction over axis
len(kept)— the merged trailing axis — which reduces over all the requested axes at once. inverseis the permutation that undoesorder(inverse[ax]says where original axisaxended up), andrestoremaps a full-shape array in the permuted-merged layout back to the caller's layout: un-merge (reshapetopermuted_shape), then un-permute (permute_dimswithinverse). The reductions drop the axis and never use it;combiner._sigma_clip_maskcomputes its mask at full shape in the working layout and usesrestoreto hand it back in the shape and axis order of the input.
— Written by Claude at @mwcraig's direction.
There was a problem hiding this comment.
Add brief comments along these lines to the code
There was a problem hiding this comment.
Pull request overview
Hoists multi-axis handling into the shared reduction setup and extends array-API sigma clipping accordingly.
Changes:
- Supports
None, tuple, list, and negative axes in reduction fallbacks. - Restores sigma-clipping masks to the original layout.
- Expands tests and documentation for new axis forms.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
CHANGES.rst |
Documents expanded axis support. |
ccdproc/_nanfuncs.py |
Centralizes axis normalization and layout restoration. |
ccdproc/combiner.py |
Enables multi-axis array-API sigma clipping. |
ccdproc/core.py |
Delegates MAD axis handling to _nanfuncs. |
ccdproc/tests/test_nanfuncs.py |
Tests supported and invalid axis forms. |
ccdproc/tests/test_combiner.py |
Tests multi-axis clipping and mask restoration. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| kept = tuple(ax for ax in range(ndim) if ax not in axes) | ||
| order = kept + axes | ||
| permuted_shape = tuple(x.shape[ax] for ax in order) | ||
| x = xp.reshape(xp.permute_dims(x, order), permuted_shape[: len(kept)] + (-1,)) |
There was a problem hiding this comment.
Confirmed — this is the same defect as #1006 (comment), and your framing of the mechanism is right: with total size 0 the -1 is genuinely ambiguous to reshape (any merged length satisfies 0 × n = 0), even though the intended value — the product of the reduced axes' sizes — is well-defined. Reproduced on this branch: nansum(np.zeros((0, 3, 4)), axis=(1, 2)) → ValueError: cannot reshape array of size 0 into shape (0,newaxis) where np.nansum returns shape (0,). Will fix by reshaping with the explicit product instead of -1, and add an empty-array regression row.
— Written by Claude at @mwcraig's direction.
mwcraig
left a comment
mwcraig
left a comment
There was a problem hiding this comment.
Adversarial self-review (multi-agent, via Claude Code): every finding below was verified by running code against this branch — differential harnesses vs numpy/astropy on numpy + array-api-strict + dask, plus full runs of the changed test files (all pass). Two confirmed edge-case bugs and a handful of simplifications, all as inline comments; two findings fall outside the diff:
test_ccdproc.py:450-466 — test_mad_fallback_rejects_bad_axis_tuple now duplicates test_bad_axis. Post-hoist, both grids exercise the identical _setup lines, which test_nanfuncs.py::test_bad_axis already covers 5-fold. It could shrink to a single case proving _mad_fallback propagates _setup's error, with its match= strings moved into test_bad_axis (which currently asserts only exception types).
Unrelated hunks: core.py carries five string-concatenation cleanups (~lines 1014, 1728, 2062, 2750, 2757) with no connection to the axis work.
Probed and clean: restore round-trips exactly for arbitrary tuple orders on 4-D; all six reductions match numpy/astropy for None/tuple/negative/all-NaN cases; _sigma_clip_mask masks are element-identical to astropy's across axis forms; no in-place mutation, boolean indexing, or data-dependent shapes (backend-safe); duplicate/out-of-range axes raise correctly; no NotImplementedError catchers existed anywhere, and the TypeError change is asserted in test_bad_axis. Simplifications considered and rejected: inlining _mad_fallback (still owns the tiered med closure and is a test seam), routing single-int axes through the tuple branch (permute/copy on the hot combiner path), and passing tuple axes straight to xp.sum (forks _setup's contract while median/mad still need the merge).
|
|
||
| if isinstance(axis, tuple | list): | ||
| # normalize_axis_tuple would silently treat True as 1. | ||
| if any(isinstance(ax, bool) for ax in axis): |
There was a problem hiding this comment.
Bug (found independently by two reviewers): np.bool_ slips through this guard. np.True_ is not a subclass of Python bool, so isinstance(ax, bool) is False for it and normalize_axis_tuple then silently treats it as axis 1 — nansum(x, axis=(0, np.True_)) reduces axes (0, 1), which is exactly the hazard the comment above says this rejects. numpy itself raises TypeError here. Same hole in the scalar guard at line 152 (operator.index(np.True_) returns 1, currently with a DeprecationWarning that numpy will turn into an error — and the test suite's filterwarnings = ["error", ...] means a test_bad_axis row for this fails on the warning, not the intended TypeError). The scalar hole is pre-existing; the tuple-entry one is new here. Fix in both places: isinstance(ax, bool | np.bool_) — numpy is already imported host-side. This also falsifies _setup's docstring claim that booleans are rejected while numpy scalars are accepted.
There was a problem hiding this comment.
Confirmed, with one version nuance worth recording. On the oldest supported numpy (2.0.x — pyproject.toml floor is numpy>=2.0): normalize_axis_tuple((0, np.True_), 2) returns (0, 1) with only a DeprecationWarning — exactly the silent wrong-axes reduction described (and under the suite's warnings-as-errors config, the wrong failure mode for a test row). On current numpy (verified 2.5.2) that deprecation has expired: operator.index(np.True_) now raises TypeError, so both the tuple entry and the scalar already fail there, just with numpy's message instead of ours.
So the fix is still exactly right for the supported range: isinstance(ax, bool | np.bool_) in both guards gives uniform behaviour and our clearer message on every numpy from 2.0 up, and makes the docstring's "booleans are rejected" claim true again. Will apply, plus test_bad_axis rows for np.True_ as a scalar and as a tuple entry (safe on all supported numpys once the guard is explicit).
— Written by Claude at @mwcraig's direction.
| kept = tuple(ax for ax in range(ndim) if ax not in axes) | ||
| order = kept + axes | ||
| permuted_shape = tuple(x.shape[ax] for ax in order) | ||
| x = xp.reshape(xp.permute_dims(x, order), permuted_shape[: len(kept)] + (-1,)) |
There was a problem hiding this comment.
Bug: the -1 merge reshape crashes when a kept dimension has size 0, where numpy returns an empty result. nansum(np.zeros((0, 3, 4)), axis=(1, 2)) → ValueError: cannot reshape array of size 0 into shape (0,newaxis) on every backend, while np.nansum gives shape (0,). Inherited from the old _mad_fallback merge, but the hoist spreads it to all six reductions plus sigma_clipping/mad_std. One-line fix: reshape with the explicit product instead of -1, e.g. permuted_shape[: len(kept)] + (math.prod(permuted_shape[len(kept):]),). Rare input in practice, but it's a genuine numpy divergence in the newly advertised feature.
There was a problem hiding this comment.
Confirmed — reproduced verbatim on this branch: nansum(np.zeros((0, 3, 4)), axis=(1, 2)) → ValueError: cannot reshape array of size 0 into shape (0,newaxis) on numpy, where np.nansum returns shape (0,). (Copilot found the same bug independently in its comment above.) Will fix with the explicit product as suggested — math.prod needs a new import math, everything else is one line — and add the empty-array regression row. Bonus: the explicit product also makes axis=() on an empty array work (prod(()) == 1), where -1 would hit the same ambiguity.
— Written by Claude at @mwcraig's direction.
| if axis is None: | ||
| shape = x.shape | ||
|
|
||
| def restore(a): | ||
| return xp.reshape(a, shape) | ||
|
|
||
| return xp.reshape(x, (-1,)), 0, xp, device, restore |
There was a problem hiding this comment.
Simplification: this whole branch folds into the tuple branch. Replacing it with
if axis is None:
axis = tuple(range(ndim))produces identical results: kept=() makes the permutation the identity, the merge reshape becomes reshape(x, (-1,)), the returned axis is len(kept) == 0, and restore is permute_dims(reshape(a, shape), identity). Verified with a differential harness (all six public functions + _sigma_clip_mask, numpy and array-api-strict, 11 axis forms including 0-d input): every value, shape, mask, and error identical. Only cost is a no-op identity permute_dims on this path (a view on numpy). Removes one of _setup's three restore definitions and an early return.
There was a problem hiding this comment.
Agreed — the mechanism checks out from the code: axis=None → tuple(range(ndim)) gives kept=(), so the permutation is the identity, the merge reshape degenerates to reshape(x, (-1,)) (or the full-size product once the explicit-product fix from the other thread lands — the two compose cleanly), the returned axis is len(kept) == 0, and restore is the same reshape the dedicated branch builds, behind a no-op permute_dims. Your differential harness covering the 0-d case settles the rest. Will apply, dropping the dedicated None branch and its restore.
— Written by Claude at @mwcraig's direction.
| if not -ndim <= axis < ndim: | ||
| raise ValueError(f"axis {axis} is out of bounds for array of dimension {ndim}") | ||
| axis = axis % ndim | ||
|
|
||
| device = array_api_compat.device(x) | ||
| x = _promote_to_real(x, xp, device) | ||
|
|
||
| return x, axis, xp, device | ||
| return x, axis % ndim, xp, device, lambda a: a |
There was a problem hiding this comment.
Simplification: this hand-rolled bounds check is normalize_axis_tuple(axis, ndim)[0]. The AxisError it raises carries a character-for-character identical message and subclasses ValueError, so test_bad_axis still passes unchanged. It also removes an existing asymmetry: the tuple branch already raises AxisError for out-of-bounds entries while this branch raises plain ValueError. (Keep the bool guard and the operator.index try/except above — their message is better.)
return x, normalize_axis_tuple(axis, ndim)[0], xp, device, lambda a: aTogether with the axis is None fold, _setup's body goes from ~43 to ~34 lines with proven-identical behavior.
There was a problem hiding this comment.
Verified both claims: normalize_axis_tuple(2, 2) raises AxisError('axis 2 is out of bounds for array of dimension 2') — character-identical to the hand-rolled message — and AxisError subclasses ValueError, so test_bad_axis passes unchanged. The asymmetry point is real too: the tuple branch already surfaces AxisError for out-of-bounds entries while this branch raises bare ValueError, and this change removes that. Will apply, keeping the bool guard and the operator.index try/except above for their better TypeError message, together with the axis is None fold.
— Written by Claude at @mwcraig's direction.
| ref_axis = axis | ||
| if isinstance(axis, tuple): | ||
| ref_axis = tuple(ax % np_data.ndim for ax in axis) |
There was a problem hiding this comment.
Test simplification: this normalization is duplicated. _sigma_clip_reference already computes the same ax % ndim normalization for its bounds shape (line ~1761). Normalizing the tuple once inside _sigma_clip_reference before calling sigma_clip (astropy's bottleneck path can't take negative tuple entries anyway) lets this ref_axis block and its comment disappear, and the helper under test still receives the tuple exactly as written.
There was a problem hiding this comment.
Agreed — _sigma_clip_reference already computes the same ax % ndim normalisation for its bounds shape, so normalising the tuple once at the top of the helper (before its sigma_clip call, whose bottleneck dispatch can't take negative tuple entries) serves both uses, and this ref_axis block and its comment disappear while _sigma_clip_mask still receives the axis exactly as the parametrize row writes it. Will apply.
— Written by Claude at @mwcraig's direction.
| are never accepted, on any array namespace -- this method | ||
| honoured for every array namespace; like | ||
| :func:`~astropy.stats.sigma_clip`, ``axis`` may be a single | ||
| integer, ``None`` or a tuple of axes on every one. ``masked`` |
There was a problem hiding this comment.
Nit: "a tuple of axes" — _setup and the _nanfuncs docstrings say "tuple or list", and a list is accepted here too (there's a test for it). One word for consistency.
There was a problem hiding this comment.
Agreed. Double-checked that "or list" is true on both dispatch paths before changing the doc: astropy.stats.sigma_clip accepts axis=[1, 2] (verified), so the numpy path is fine, and _setup handles a list off-numpy (test_nanfuncs.py::test_list_axis_matches_tuple). Will make it "a tuple or list of axes".
— Written by Claude at @mwcraig's direction.
| NaN-aware reductions or ccdproc's fallbacks. [#1001] | ||
| - ``Combiner.sigma_clipping`` outside NumPy now accepts ``axis=None`` and a | ||
| tuple of axes, as ``astropy.stats.sigma_clip`` does on the NumPy path, and | ||
| the reduction fallbacks in ``ccdproc._nanfuncs`` gained the same axis |
There was a problem hiding this comment.
The NotImplementedError→TypeError change for bad axis isn't mentioned. Verified non-breaking (_nanfuncs landed in unreleased #986, latest tag 2.5.1), so strictly optional — but one clause here, e.g. "an invalid axis now raises TypeError rather than NotImplementedError", makes the entry airtight.
There was a problem hiding this comment.
|
Replies to the two outside-diff findings from the self-review (each inline finding has an answer in its own thread — every one verified and accepted; the two edge-case bugs reproduce exactly as described):
Unrelated Planned change set from all threads, for one push: the two bug fixes ( — Written by Claude at @mwcraig's direction. |
…document Bug fixes, both found in review: - Reject np.bool_ axes explicitly, as scalars and as tuple/list entries: on the oldest supported numpy (2.0) operator.index still converts np.True_ to 1 with only a DeprecationWarning, so axis=(0, np.True_) silently reduced axes (0, 1) there. Newer numpy raises TypeError itself; the guard makes the behaviour and message uniform. - Spell out the merged-axis length in _setup's tuple merge instead of reshaping with -1, which is ambiguous (and raises) when a kept axis has size 0; numpy's reductions return an empty result there. Found independently by two reviewers. Simplifications from review, all behaviour-preserving: - Fold the axis=None branch into the tuple branch (None is tuple(range(ndim))). - Replace the hand-rolled scalar bounds check with normalize_axis_tuple, whose AxisError is a ValueError with an identical message; the tuple branch already surfaced AxisError. - Normalize tuple axes once inside test_combiner's _sigma_clip_reference instead of at its call site. - Shrink test_ccdproc's bad-axis grid to a single delegation check; the messages are now pinned in test_nanfuncs::test_bad_axis. Docs and tests: move _setup's axis-handling discussion into a Notes section, add brief comments to the permute-and-merge block, explain the list-equals-tuple test, document that sigma_clipping takes a list of axes too, note the NotImplementedError->TypeError change in CHANGES, switch the diff's British spellings to American, and drop the unrelated string-concatenation cleanups from core.py. New differential rows cover axis=(), numpy ints in a tuple, and the two empty-array shapes (where numpy.median itself trips over the same -1 merge; the test substitutes np.nanmedian as the reference there). Full runs of the three affected files: numpy 697, dask 701, jax full suite 968, array-api-strict 689 + 12 pre-existing xfails; ruff and black clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LVxnTLrWKStBdPwcmDNhxA
Requested in review on astropy#1006: the six public reductions repeated the same three Parameters entries with incidental wording drift. They are now written once in _COMMON_PARAMS and filled into each docstring's {params} placeholder by the _fill_doc decorator (the scipy doccer pattern), with only the axis action phrase substituted per function. Function-specific behaviour -- what an all-NaN slice yields, NaN propagation -- stays inline in each Returns section, and nanmad keeps its extra ``median`` entry after the template. The filled docstrings render identically to the hand-written ones (help(), Sphinx, and numpydoc all read __doc__, which is complete after import); the trade-off is that raw source shows the {params} placeholder. python -OO strips docstrings, so the decorator fills only when __doc__ exists. The unified axis entry also carries the booleans-rejected/numpy-scalars-accepted note everywhere, which previously appeared in only three of the six. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LVxnTLrWKStBdPwcmDNhxA
|
The change set promised across the review threads has landed in two commits:
One find from adding the empty-array rows: Local runs after the changes: numpy 964 (full suite), dask 701, jax 968 (full suite), array-api-strict 689 + 12 pre-existing xfails; ruff, black, and pre-commit clean. — Written by Claude at @mwcraig's direction. |
Follow-up agreed in the #1000 review; closes #1004.
Moves the flatten-for-
Noneand permute-and-merge steps fromcore._mad_fallbackinto the shared_nanfuncs._setup, which now accepts a single integer,None, or a tuple/list of axes and returns arestorecallable mapping a full-shape array back to the caller's layout (identity for a single integer axis).What this buys:
_nanfuncs(nansum,nanmean,nanstd,nanmedian,median,nanmad) gains tuple-axis andaxis=Nonesupport for free;_mad_fallbackshrinks to the median tier plus a delegation tonanmad;Combiner.sigma_clipping's array-API path lifts its single-integer-axis restriction and now accepts the same axis formsastropy.stats.sigma_clipdoes, with the clip mask handed back in the data's own shape viarestore.Behavior note: a bare-bool or non-integer
axisnow raisesTypeErrorinstead ofNotImplementedError, matching theTypeErrora bool entry in a tuple already raised.Tests: the differential grids in
test_nanfuncs.py/test_combiner.pygainedNone, tuple (including negative entries), and single-entry-tuple cases checked against numpy/astropy, a list-equals-tuple equivalence test, expanded bad-axis cases, and Combiner-level axis-form tests. Full suites pass locally: numpy 943, dask 676/676, jax 676/676, array-api-strict 664 passed + 12 pre-existing xfails on the three affected test files.🤖 Generated with Claude Code
https://claude.ai/code/session_01LVxnTLrWKStBdPwcmDNhxA