diff options
Diffstat (limited to 'numpy/lib')
| -rw-r--r-- | numpy/lib/function_base.py | 21 | ||||
| -rw-r--r-- | numpy/lib/tests/test_function_base.py | 23 |
2 files changed, 36 insertions, 8 deletions
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index 5c1654fc3..efa51173a 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -898,15 +898,19 @@ def average(a, axis=None, weights=None, returned=False): TypeError: Axis must be specified when shapes of a and weights differ. """ - if not isinstance(a, np.matrix): - a = np.asarray(a) + a = np.asanyarray(a) if weights is None: avg = a.mean(axis) scl = avg.dtype.type(a.size/avg.size) else: - a = a + 0.0 - wgt = np.asarray(weights) + wgt = np.asanyarray(weights) + + if issubclass(a.dtype.type, (np.integer, np.bool_)): + result_dtype = np.result_type(a.dtype, wgt.dtype, 'f8') + else: + result_dtype = np.result_type(a.dtype, wgt.dtype) + # Sanity checks if a.shape != wgt.shape: if axis is None: @@ -921,17 +925,18 @@ def average(a, axis=None, weights=None, returned=False): "Length of weights not compatible with specified axis.") # setup wgt to broadcast along axis - wgt = np.array(wgt, copy=0, ndmin=a.ndim).swapaxes(-1, axis) + wgt = np.broadcast_to(wgt, (a.ndim-1)*(1,) + wgt.shape) + wgt = wgt.swapaxes(-1, axis) - scl = wgt.sum(axis=axis, dtype=np.result_type(a.dtype, wgt.dtype)) + scl = wgt.sum(axis=axis, dtype=result_dtype) if (scl == 0.0).any(): raise ZeroDivisionError( "Weights sum to zero, can't be normalized") - avg = np.multiply(a, wgt).sum(axis)/scl + avg = np.multiply(a, wgt, dtype=result_dtype).sum(axis)/scl if returned: - scl = np.multiply(avg, 0) + scl + scl = np.broadcast_to(scl, avg.shape) return avg, scl else: return avg diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index 235b7f2fe..943544dd5 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -167,6 +167,29 @@ class TestAverage(TestCase): avg, scl = average(y, weights=w2, axis=1, returned=True) assert_array_equal(scl, np.array([1., 6.])) + def test_subclasses(self): + class subclass(np.ndarray): + pass + a = np.array([[1,2],[3,4]]).view(subclass) + w = np.array([[1,2],[3,4]]).view(subclass) + + assert_equal(type(np.average(a, weights=w)), subclass) + + # also test matrices + a = np.matrix([[1,2],[3,4]]) + w = np.matrix([[1,2],[3,4]]) + + r = np.average(a, axis=0, weights=w) + assert_equal(type(r), np.matrix) + assert_equal(r, [[2.5, 10.0/3]]) + + def test_upcasting(self): + types = [('i4', 'i4', 'f8'), ('i4', 'f4', 'f8'), ('f4', 'i4', 'f8'), + ('f4', 'f4', 'f4'), ('f4', 'f8', 'f8')] + for at, wt, rt in types: + a = np.array([[1,2],[3,4]], dtype=at) + w = np.array([[1,2],[3,4]], dtype=wt) + assert_equal(np.average(a, weights=w).dtype, np.dtype(rt)) class TestSelect(TestCase): choices = [np.array([1, 2, 3]), |