Skip to content

Commit 22858d4

Browse files
committed
Ensure bottleneck result types are consistent with numpy
1 parent a723956 commit 22858d4

1 file changed

Lines changed: 15 additions & 33 deletions

File tree

‎astropy/stats/sigma_clipping.py‎

Lines changed: 15 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
# Licensed under a 3-clause BSD style license - see LICENSE.rst
22

3+
import functools
34
import warnings
45

56
import numpy as np
@@ -49,40 +50,29 @@ def _move_tuple_axes_first(array, axis):
4950
return array_new
5051

5152

52-
def _nanmean(array, axis=None):
53-
"""Bottleneck nanmean function that handle tuple axis."""
54-
if isinstance(axis, tuple):
55-
array = _move_tuple_axes_first(array, axis=axis)
56-
axis = 0
57-
58-
if isinstance(array, Quantity):
59-
return array.__array_wrap__(bottleneck.nanmean(array, axis=axis))
60-
else:
61-
return bottleneck.nanmean(array, axis=axis)
53+
def _apply_bottleneck(function, array, axis=None, **kwargs):
54+
"""Wrap bottleneck function to handle tuple axis.
6255
63-
64-
def _nanmedian(array, axis=None):
65-
"""Bottleneck nanmedian function that handle tuple axis."""
56+
Also takes care to ensure the output is of the expected type,
57+
i.e., a quantity, numpy array, or numpy scalar.
58+
"""
6659
if isinstance(axis, tuple):
6760
array = _move_tuple_axes_first(array, axis=axis)
6861
axis = 0
6962

63+
result = function(array, axis=axis, **kwargs)
7064
if isinstance(array, Quantity):
71-
return array.__array_wrap__(bottleneck.nanmedian(array, axis=axis))
65+
return array.__array_wrap__(result)
66+
elif isinstance(result, float):
67+
# For compatibility with numpy, always return a numpy scalar.
68+
return np.float64(result)
7269
else:
73-
return bottleneck.nanmedian(array, axis=axis)
70+
return result
7471

7572

76-
def _nanstd(array, axis=None, ddof=0):
77-
"""Bottleneck nanstd function that handle tuple axis."""
78-
if isinstance(axis, tuple):
79-
array = _move_tuple_axes_first(array, axis=axis)
80-
axis = 0
81-
82-
if isinstance(array, Quantity):
83-
return array.__array_wrap__(bottleneck.nanstd(array, axis=axis, ddof=ddof))
84-
else:
85-
return bottleneck.nanstd(array, axis=axis, ddof=ddof)
73+
_nanmean = functools.partial(_apply_bottleneck, bottleneck.nanmean)
74+
_nanmedian = functools.partial(_apply_bottleneck, bottleneck.nanmedian)
75+
_nanstd = functools.partial(_apply_bottleneck, bottleneck.nanstd)
8676

8777

8878
def _nanmadstd(array, axis=None):
@@ -316,15 +306,7 @@ def _compute_bounds(self, data, axis=None):
316306
with warnings.catch_warnings():
317307
warnings.simplefilter("ignore", category=RuntimeWarning)
318308
cen = self._cenfunc_parsed(data, axis=axis)
319-
if np.isscalar(cen) and not isinstance(cen, np.float64):
320-
# change bottleneck float scalar to np.float64
321-
cen = np.float64(cen)
322-
323309
std = self._stdfunc_parsed(data, axis=axis)
324-
if np.isscalar(std) and not isinstance(std, np.float64):
325-
# change bottleneck float scalar to np.float64
326-
std = np.float64(std)
327-
328310
self._min_value = cen - (std * self.sigma_lower)
329311
self._max_value = cen + (std * self.sigma_upper)
330312

0 commit comments

Comments
 (0)