Merge pull request numpy#7382 from ahaldane/tidy_average_median

MAINT: cleanup np.average

Author: charris

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

numpy/lib/tests/test_function_base.py

@@ -167,6 +167,29 @@ def test_returned(self):
         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]),