8000 Add SigmaClippedStats class by larrybradley · Pull Request #17221 · astropy/astropy · GitHub
[go: up one dir, main page]

Skip to content
18 changes: 17 additions & 1 deletion astropy/stats/nanfunctions.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,11 @@ def _apply_bottleneck(

result = function(array, axis=axis, **kwargs)
if isinstance(array, Quantity):
return array.__array_wrap__(result)
if function is bottleneck.nanvar:
result = array._result_as_quantity(result, array.unit**2, None)
else:
result = array.__array_wrap__(result)
return result
elif isinstance(result, float):
# For compatibility with numpy, always return a numpy scalar.
return np.float64(result)
Expand All @@ -86,16 +90,22 @@ def _apply_bottleneck(

bn_funcs = dict(
nansum=functools.partial(_apply_bottleneck, bottleneck.nansum),
nanmin=functools.partial(_apply_bottleneck, bottleneck.nanmin),
nanmax=functools.partial(_apply_bottleneck, bottleneck.nanmax),
nanmean=functools.partial(_apply_bottleneck, bottleneck.nanmean),
nanmedian=functools.partial(_apply_bottleneck, bottleneck.nanmedian),
nanstd=functools.partial(_apply_bottleneck, bottleneck.nanstd),
nanvar=functools.partial(_apply_bottleneck, bottleneck.nanvar),
)

np_funcs = dict(
nansum=np.nansum,
nanmin=np.nanmin,
nanmax=np.nanmax,
nanmean=np.nanmean,
nanmedian=np.nanmedian,
nanstd=np.nanstd,
nanvar=np.nanvar,
)

def _dtype_dispatch(func_name):
Expand All @@ -115,15 +125,21 @@ def wrapped(*args, **kwargs):
return wrapped

nansum = _dtype_dispatch("nansum")
nanmin = _dtype_dispatch("nanmin")
nanmax = _dtype_dispatch("nanmax")
nanmean = _dtype_dispatch("nanmean")
nanmedian = _dtype_dispatch("nanmedian")
nanstd = _dtype_dispatch("nanstd")
nanvar = _dtype_dispatch("nanvar")

else:
nansum = np.nansum
nanmin = np.nanmin
nanmax = np.nanmax
nanmean = np.nanmean
nanmedian = np.nanmedian
nanstd = np.nanstd
nanvar = np.nanvar


def nanmadstd(
Expand Down
Loading
0