Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions CHANGES.rst
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,11 @@ New Features
classified differently from astropy (NumPy data still use astropy);
``'median'``/``'mean'``/``'std'``/``'mad_std'`` use the namespace's
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

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed — verified the non-breaking claim independently (_nanfuncs first appears in unreleased #986; latest release tag is 2.5.1, which has no _nanfuncs module), so it's a doc-completeness fix, not a compat note. Will append the clause to the #1006 entry.

Written by Claude at @mwcraig's direction.

forms; the None/tuple axis handling formerly in ``_mad_fallback`` moved
into the shared ``_nanfuncs._setup``. [#1006]

Other Changes and Additions
^^^^^^^^^^^^^^^^^^^^^^^^^^^
Expand Down
185 changes: 122 additions & 63 deletions ccdproc/_nanfuncs.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,10 @@

import array_api_compat

# Host-side axis normalisation for tuple axes: operates on python ints
# only, never on array data, so it does not tie the fallbacks to numpy.
from numpy.lib.array_utils import normalize_axis_tuple

__all__ = ["median", "nanmad", "nanmean", "nanmedian", "nanstd", "nansum"]


Expand Down Expand Up @@ -58,64 +62,107 @@ def _promote_to_real(x, xp, device):

def _setup(x, axis, xp):
"""
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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This goes in notes

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

``None`` flattens ``x`` so the caller reduces over everything; a tuple
or list moves the listed axes to the end and merges them into one, so
the caller's single-axis reduction reduces over all of them at once.
Either way the caller only ever sees a single non-negative integer
axis.

Parameters
----------
x : array
Input array.
axis : int
Axis along which the caller will reduce. Booleans, ``None`` and
tuples of axes are rejected; anything else goes through
`operator.index`, so numpy integer scalars are accepted.
axis : int, tuple of int, list of int or None
Axis or axes along which the caller will reduce. Booleans are
rejected -- bool subclasses int, so ``axis=True`` would silently
mean axis 1 -- while numpy integer scalars are accepted. Negative
values count from the last axis.
xp : array namespace or None
Namespace to use. ``None`` resolves it from ``x``.

Returns
-------
x : array
The input, promoted if necessary to the namespace's default real
floating dtype.
floating dtype, flattened when ``axis`` is ``None``, and with the
listed axes moved to the end and merged into one when ``axis`` is
a tuple or list.
axis : int
The axis, normalised to a non-negative integer.
The single axis of the returned ``x`` to reduce, normalised to a
non-negative integer.
xp : array namespace
The resolved namespace.
device : device
The device ``x`` lives on.
restore : callable
Maps an array shaped like the returned ``x`` back to the layout of
the input ``x``; the identity for a single integer ``axis``.
Reductions remove the reduced axis and never need it;
``ccdproc.combiner._sigma_clip_mask`` keeps the full shape and
uses it to hand its mask back in the caller's layout.

Raises
------
NotImplementedError
If ``axis`` is not a single integer.
TypeError
If ``axis``, or an entry of a tuple/list ``axis``, is a bool or
not an integer.
ValueError
If ``axis`` is out of bounds for ``x``.
If ``axis``, or an entry of a tuple/list ``axis``, is out of
bounds for ``x``, or a tuple/list names an axis more than once
(including via a negative alias).
"""
if xp is None:
xp = array_api_compat.array_namespace(x)
device = array_api_compat.device(x)
x = _promote_to_real(x, xp, device)
ndim = x.ndim

if axis is None:
shape = x.shape

def restore(a):
return xp.reshape(a, shape)

return xp.reshape(x, (-1,)), 0, xp, device, restore

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 isinstance(axis, tuple | list):
# normalize_axis_tuple would silently treat True as 1.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Explain in more detail what is going on in this block -- no comments in the code, just explain in the reply

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Step by step:

  1. The bool guard. normalize_axis_tuple converts entries with operator.index, and a Python True passes that (operator.index(True) == 1), so axis=(0, True) would silently reduce axes (0, 1). The guard turns that into a TypeError before the conversion runs.
  2. normalize_axis_tuple(axis, ndim) does the host-side validation and normalisation in one call: each entry through operator.index, negatives wrapped mod ndim (so -1 on a 3-D array becomes 2), AxisError on out-of-bounds, ValueError on duplicates — including duplicates via a negative alias, e.g. (0, -3) on 3-D. Out comes a tuple of unique non-negative ints.
  3. kept is the axes not being reduced, in their original order; order = kept + axes is then a permutation of range(ndim) that moves every reduced axis to the end. permute_dims(x, order) rearranges the array into that layout, and the reshape collapses 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.
  4. 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.
  5. inverse is the permutation that undoes order (inverse[ax] says where original axis ax ended up), and restore maps a full-shape array in the permuted-merged layout back to the caller's layout: un-merge (reshape to permuted_shape), then un-permute (permute_dims with inverse). The reductions drop the axis and never use it; combiner._sigma_clip_mask computes its mask at full shape in the working layout and uses restore to hand it back in the shape and axis order of the input.

Written by Claude at @mwcraig's direction.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add brief comments along these lines to the code

if any(isinstance(ax, bool) for ax in axis):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

raise TypeError("axis entries must be integers, not bool")
axes = normalize_axis_tuple(axis, ndim)
# Move the reduced axes to the end and merge them into one, so that
# a single-axis reduction reduces over all of them at once.
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,))

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

inverse = tuple(order.index(ax) for ax in range(ndim))

def restore(a):
return xp.permute_dims(xp.reshape(a, permuted_shape), inverse)

return x, len(kept), xp, device, restore

# bool subclasses int -- axis=True would silently mean axis 1 -- so it is
# rejected explicitly, while operator.index accepts the numpy integer
# scalars that isinstance(axis, int) would refuse.
if axis is None or isinstance(axis, bool):
raise NotImplementedError(
"NaN-aware reduction fallbacks support only a single integer axis."
)
if isinstance(axis, bool):
raise TypeError("axis must be an integer, not bool")
try:
axis = operator.index(axis)
except TypeError:
raise NotImplementedError(
"NaN-aware reduction fallbacks support only a single integer axis."
raise TypeError(
f"axis must be an integer, a tuple or list of integers, or None, "
f"got {axis!r}"
) from None

if xp is None:
xp = array_api_compat.array_namespace(x)

ndim = x.ndim
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

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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: a

Together with the axis is None fold, _setup's body goes from ~43 to ~34 lines with proven-identical behavior.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.



def _sum_and_count(x, axis, xp, device, *, keepdims):
Expand Down Expand Up @@ -197,19 +244,20 @@ def nansum(x, /, *, axis=0, xp=None):
x : array
Input array. Integer and boolean inputs are promoted to the
namespace's default real floating dtype.
axis : int, optional
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

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 deco

with 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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Implement it

every axis; a tuple or list sums over all the listed axes at once.
xp : array namespace, optional
Namespace to use. Defaults to ``array_api_compat.array_namespace(x)``.

Returns
-------
array
Sum of ``x`` along ``axis``, with that axis removed. Slices that are
entirely NaN sum to zero, matching `numpy.nansum`.
Sum of ``x`` along ``axis``, with the reduced axes removed (0-d
when ``axis`` is ``None``). Slices that are entirely NaN sum to
zero, matching `numpy.nansum`.
"""
x, axis, xp, device = _setup(x, axis, xp)
x, axis, xp, device, _ = _setup(x, axis, xp)
total, _ = _sum_and_count(x, axis, xp, device, keepdims=False)
return total

Expand All @@ -223,21 +271,22 @@ def nanmean(x, /, *, axis=0, xp=None):
x : array
Input array. Integer and boolean inputs are promoted to the
namespace's default real floating dtype.
axis : int, optional
Axis along which to average. 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 average. Default is 0. ``None``
averages over every axis; a tuple or list over all the listed axes.
xp : array namespace, optional
Namespace to use. Defaults to ``array_api_compat.array_namespace(x)``.

Returns
-------
array
Mean of ``x`` along ``axis``, with that axis removed. Slices that
are entirely NaN yield NaN silently, matching ``bottleneck.nanmean``
Mean of ``x`` along ``axis``, with the reduced axes removed (0-d
when ``axis`` is ``None``). Slices that are entirely NaN yield
NaN silently, matching ``bottleneck.nanmean``
(the numpy-backend default); `numpy.nanmean` warns here, but a fully
masked pixel is a routine input for the combiner, not an anomaly.
"""
x, axis, xp, device = _setup(x, axis, xp)
x, axis, xp, device, _ = _setup(x, axis, xp)
total, count = _sum_and_count(x, axis, xp, device, keepdims=False)
return _safe_divide(total, count, xp, device)

Expand All @@ -255,16 +304,18 @@ def nanstd(x, /, *, axis=0, xp=None):
x : array
Input array. Integer and boolean inputs are promoted to the
namespace's default real floating dtype.
axis : int, optional
Axis along which to compute the deviation. 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 compute the deviation. Default is 0.
``None`` reduces over every axis; a tuple or list over all the
listed axes.
xp : array namespace, optional
Namespace to use. Defaults to ``array_api_compat.array_namespace(x)``.

Returns
-------
array
Standard deviation of ``x`` along ``axis``, with that axis removed.
Standard deviation of ``x`` along ``axis``, with the reduced axes
removed (0-d when ``axis`` is ``None``).
Slices that are entirely NaN yield NaN silently, matching
``bottleneck.nanstd`` (the numpy-backend default); `numpy.nanstd`
warns here, but a fully masked pixel is a routine input for the
Expand All @@ -279,7 +330,7 @@ def nanstd(x, /, *, axis=0, xp=None):
single-pass form suffers when the values are large relative to their
spread, which is not unusual for CCD counts.
"""
x, axis, xp, device = _setup(x, axis, xp)
x, axis, xp, device, _ = _setup(x, axis, xp)

isnan = xp.isnan(x)
zero = xp.asarray(0, dtype=x.dtype, device=device)
Expand Down Expand Up @@ -314,18 +365,20 @@ def nanmedian(x, /, *, axis=0, xp=None):
x : array
Input array. Integer and boolean inputs are promoted to the
namespace's default real floating dtype.
axis : int, optional
Axis along which to compute the median. Default is 0. Booleans,
``None`` and tuples of axes are not supported; numpy integer
scalars are accepted.
axis : int, tuple of int, list of int or None, optional
Axis or axes along which to compute the median. Default is 0.
``None`` reduces over every axis and a tuple or list over all the
listed axes; booleans are rejected, numpy integer scalars are
accepted.
xp : array namespace, optional
Namespace to use. Defaults to ``array_api_compat.array_namespace(x)``.

Returns
-------
array
Median of ``x`` along ``axis``, with that axis removed. Slices that
are entirely NaN yield NaN silently, matching
Median of ``x`` along ``axis``, with the reduced axes removed (0-d
when ``axis`` is ``None``). Slices that are entirely NaN yield NaN
silently, matching
``bottleneck.nanmedian`` (the numpy-backend default);
`numpy.nanmedian` warns here, but a fully masked pixel is a routine
input for the combiner, not an anomaly.
Expand All @@ -337,7 +390,7 @@ def nanmedian(x, /, *, axis=0, xp=None):
or ``bottleneck.nanmedian``. Prefer a native ``nanmedian`` when the
namespace offers one.
"""
x, axis, xp, device = _setup(x, axis, xp)
x, axis, xp, device, _ = _setup(x, axis, xp)
ndim = x.ndim

# Replacing NaNs with +inf keeps them past every real value regardless of
Expand Down Expand Up @@ -385,18 +438,20 @@ def median(x, /, *, axis=0, xp=None):
x : array
Input array. Integer and boolean inputs are promoted to the
namespace's default real floating dtype.
axis : int, optional
Axis along which to compute the median. Default is 0. Booleans,
``None`` and tuples of axes are not supported; numpy integer
scalars are accepted.
axis : int, tuple of int, list of int or None, optional
Axis or axes along which to compute the median. Default is 0.
``None`` reduces over every axis and a tuple or list over all the
listed axes; booleans are rejected, numpy integer scalars are
accepted.
xp : array namespace, optional
Namespace to use. Defaults to ``array_api_compat.array_namespace(x)``.

Returns
-------
array
Median of ``x`` along ``axis``, with that axis removed. Slices that
contain any NaN yield NaN, matching `numpy.median`; this is the
Median of ``x`` along ``axis``, with the reduced axes removed (0-d
when ``axis`` is ``None``). Slices that contain any NaN yield NaN,
matching `numpy.median`; this is the
difference from `nanmedian`, which ignores NaNs entirely.

Notes
Expand All @@ -409,7 +464,7 @@ def median(x, /, *, axis=0, xp=None):
with a final `where` over whether any NaN is present along ``axis``,
since `nanmedian` alone would silently drop NaNs instead.
"""
x, axis, xp, device = _setup(x, axis, xp)
x, axis, xp, device, _ = _setup(x, axis, xp)
nan = xp.asarray(xp.nan, dtype=x.dtype, device=device)
return xp.where(xp.any(xp.isnan(x), axis=axis), nan, nanmedian(x, axis=axis, xp=xp))

Expand All @@ -423,26 +478,30 @@ def nanmad(x, /, *, axis=0, xp=None, median=None):
x : array
Input array. Integer and boolean inputs are promoted to the
namespace's default real floating dtype.
axis : int, optional
Axis along which to compute the deviation. Default is 0. Booleans,
``None`` and tuples of axes are not supported; numpy integer
scalars are accepted.
axis : int, tuple of int, list of int or None, optional
Axis or axes along which to compute the deviation. Default is 0.
``None`` reduces over every axis and a tuple or list over all the
listed axes; booleans are rejected, numpy integer scalars are
accepted.
xp : array namespace, optional
Namespace to use. Defaults to ``array_api_compat.array_namespace(x)``.
median : callable, optional
Reduction used for both medians, called as ``median(x, axis=axis)``.
Reduction used for both medians, called as ``median(x, axis=axis)``,
always with a single integer ``axis``: a ``None`` or tuple/list
``axis`` has already been flattened or merged away by `_setup`.
Default is `nanmedian`. A keyword rather than a module-level tier
(as `ccdproc.combiner._default_median` provides) so this module has
no dependency on `ccdproc.combiner`.

Returns
-------
array
``median(|x - median(x)|)`` along ``axis``, with that axis removed.
``median(|x - median(x)|)`` along ``axis``, with the reduced axes
removed (0-d when ``axis`` is ``None``).
Unscaled: multiply by ``1.482602218505602`` for an estimate of the
standard deviation, as `astropy.stats.mad_std` does.
"""
x, axis, xp, device = _setup(x, axis, xp)
x, axis, xp, device, _ = _setup(x, axis, xp)
if median is None:
median = partial(nanmedian, xp=xp)
center = xp.expand_dims(median(x, axis=axis), axis=axis)
Expand Down
Loading