* Added `apply_over_axes` as requested in ticket #1480

2 files changed, 326 insertions, 288 deletions

diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
index 5804cf3c1..094478545 100644
--- a/numpy/ma/extras.py
+++ b/numpy/ma/extras.py
@@ -11,14 +11,14 @@ A collection of utilities for `numpy.ma`.
 __author__ = "Pierre GF Gerard-Marchant ($Author: jarrod.millman $)"
 __version__ = '1.0'
 __revision__ = "$Revision: 3473 $"
-__date__     = '$Date: 2007-10-29 17:18:13 +0200 (Mon, 29 Oct 2007) $'
+__date__ = '$Date: 2007-10-29 17:18:13 +0200 (Mon, 29 Oct 2007) $'
 
-__all__ = ['apply_along_axis', 'atleast_1d', 'atleast_2d', 'atleast_3d',
-           'average',
+__all__ = ['apply_along_axis', 'apply_over_axes', 'atleast_1d', 'atleast_2d',
+           'atleast_3d', 'average',
            'clump_masked', 'clump_unmasked', 'column_stack', 'compress_cols',
            'compress_rowcols', 'compress_rows', 'count_masked', 'corrcoef',
            'cov',
-           'diagflat', 'dot','dstack',
+           'diagflat', 'dot', 'dstack',
            'ediff1d',
            'flatnotmasked_contiguous', 'flatnotmasked_edges',
            'hsplit', 'hstack',
@@ -37,8 +37,8 @@ import itertools
 import warnings
 
 import core as ma
-from core import MaskedArray, MAError, add, array, asarray, concatenate, count,\
-    filled, getmask, getmaskarray, make_mask_descr, masked, masked_array,\
+from core import MaskedArray, MAError, add, array, asarray, concatenate, count, \
+    filled, getmask, getmaskarray, make_mask_descr, masked, masked_array, \
     mask_or, nomask, ones, sort, zeros
 #from core import *
 
@@ -271,7 +271,7 @@ class _fromnxfunction:
 
     def __call__(self, *args, **params):
         func = getattr(np, self.__name__)
-        if len(args)==1:
+        if len(args) == 1:
             x = args[0]
             if isinstance(x, ndarray):
                 _d = func(np.asarray(x), **params)
@@ -284,7 +284,7 @@ class _fromnxfunction:
         else:
             arrays = []
             args = list(args)
-            while len(args)>0 and issequence(args[0]):
+            while len(args) > 0 and issequence(args[0]):
                 arrays.append(args.pop(0))
             res = []
             for x in arrays:
@@ -317,8 +317,8 @@ def flatten_inplace(seq):
     """Flatten a sequence in place."""
     k = 0
     while (k != len(seq)):
-        while hasattr(seq[k],'__iter__'):
-            seq[k:(k+1)] = seq[k]
+        while hasattr(seq[k], '__iter__'):
+            seq[k:(k + 1)] = seq[k]
         k += 1
     return seq
 
@@ -333,12 +333,12 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
         axis += nd
     if (axis >= nd):
         raise ValueError("axis must be less than arr.ndim; axis=%d, rank=%d."
-            % (axis,nd))
-    ind = [0]*(nd-1)
-    i = np.zeros(nd,'O')
+            % (axis, nd))
+    ind = [0] * (nd - 1)
+    i = np.zeros(nd, 'O')
     indlist = range(nd)
     indlist.remove(axis)
-    i[axis] = slice(None,None)
+    i[axis] = slice(None, None)
     outshape = np.asarray(arr.shape).take(indlist)
     i.put(indlist, ind)
     j = i.copy()
@@ -364,8 +364,8 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
             # increment the index
             ind[-1] += 1
             n = -1
-            while (ind[n] >= outshape[n]) and (n > (1-nd)):
-                ind[n-1] += 1
+            while (ind[n] >= outshape[n]) and (n > (1 - nd)):
+                ind[n - 1] += 1
                 ind[n] = 0
                 n -= 1
             i.put(indlist, ind)
@@ -391,8 +391,8 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
             # increment the index
             ind[-1] += 1
             n = -1
-            while (ind[n] >= holdshape[n]) and (n > (1-nd)):
-                ind[n-1] += 1
+            while (ind[n] >= holdshape[n]) and (n > (1 - nd)):
+                ind[n - 1] += 1
                 ind[n] = 0
                 n -= 1
             i.put(indlist, ind)
@@ -411,6 +411,32 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
 apply_along_axis.__doc__ = np.apply_along_axis.__doc__
 
 
+def apply_over_axes(func, a, axes):
+    """
+    (This docstring will be overwritten)
+    """
+    val = np.asarray(a)
+    msk = getmaskarray(a)
+    N = a.ndim
+    if array(axes).ndim == 0:
+        axes = (axes,)
+    for axis in axes:
+        if axis < 0: axis = N + axis
+        args = (val, axis)
+        res = ma.array(func(*(val, axis)), mask=func(*(msk, axis)))
+        if res.ndim == val.ndim:
+            (val, msk) = (res._data, res._mask)
+        else:
+            res = ma.expand_dims(res, axis)
+            if res.ndim == val.ndim:
+                (val, msk) = (res._data, res._mask)
+            else:
+                raise ValueError("Function is not returning"\
+                                 " an array of correct shape")
+    return val
+apply_over_axes.__doc__ = np.apply_over_axes.__doc__
+
+
 def average(a, axis=None, weights=None, returned=False):
     """
     Return the weighted average of array over the given axis.
@@ -496,15 +522,15 @@ def average(a, axis=None, weights=None, returned=False):
                     wsh = (1,)
                 if wsh == ash:
                     w = np.array(w, float, copy=0)
-                    n = add.reduce(a*w, axis)
+                    n = add.reduce(a * w, axis)
                     d = add.reduce(w, axis)
                     del w
                 elif wsh == (ash[axis],):
                     ni = ash[axis]
-                    r = [None]*len(ash)
+                    r = [None] * len(ash)
                     r[axis] = slice(None, None, 1)
-                    w = eval ("w["+ repr(tuple(r)) + "] * ones(ash, float)")
-                    n = add.reduce(a*w, axis, dtype=float)
+                    w = eval ("w[" + repr(tuple(r)) + "] * ones(ash, float)")
+                    n = add.reduce(a * w, axis, dtype=float)
                     d = add.reduce(w, axis, dtype=float)
                     del w, r
                 else:
@@ -520,26 +546,26 @@ def average(a, axis=None, weights=None, returned=False):
                     wsh = (1,)
                 if wsh == ash:
                     w = array(w, dtype=float, mask=mask, copy=0)
-                    n = add.reduce(a*w, axis, dtype=float)
+                    n = add.reduce(a * w, axis, dtype=float)
                     d = add.reduce(w, axis, dtype=float)
                 elif wsh == (ash[axis],):
                     ni = ash[axis]
-                    r = [None]*len(ash)
+                    r = [None] * len(ash)
                     r[axis] = slice(None, None, 1)
-                    w = eval ("w["+ repr(tuple(r)) + \
+                    w = eval ("w[" + repr(tuple(r)) + \
                               "] * masked_array(ones(ash, float), mask)")
-                    n = add.reduce(a*w, axis, dtype=float)
+                    n = add.reduce(a * w, axis, dtype=float)
                     d = add.reduce(w, axis, dtype=float)
                 else:
                     raise ValueError, 'average: weights wrong shape.'
                 del w
     if n is masked or d is masked:
         return masked
-    result = n/d
+    result = n / d
     del n
 
     if isinstance(result, MaskedArray):
-        if ((axis is None) or (axis==0 and a.ndim == 1)) and \
+        if ((axis is None) or (axis == 0 and a.ndim == 1)) and \
            (result.mask is nomask):
             result = result._data
         if returned:
@@ -615,12 +641,12 @@ def median(a, axis=None, out=None, overwrite_input=False):
 
     """
     def _median1D(data):
-        counts = filled(count(data),0)
+        counts = filled(count(data), 0)
         (idx, rmd) = divmod(counts, 2)
         if rmd:
-            choice = slice(idx, idx+1)
+            choice = slice(idx, idx + 1)
         else:
-            choice = slice(idx-1, idx+1)
+            choice = slice(idx - 1, idx + 1)
         return data[choice].mean(0)
     #
     if overwrite_input:
@@ -710,7 +736,7 @@ def compress_rowcols(x, axis=None):
     if axis in [None, 1, -1]:
         for j in np.unique(masked[1]):
             idxc.remove(j)
-    return x._data[idxr][:,idxc]
+    return x._data[idxr][:, idxc]
 
 def compress_rows(a):
     """
@@ -827,7 +853,7 @@ def mask_rowcols(a, axis=None):
     if not axis:
         a[np.unique(maskedval[0])] = masked
     if axis in [None, 1, -1]:
-        a[:,np.unique(maskedval[1])] = masked
+        a[:, np.unique(maskedval[1])] = masked
     return a
 
 def mask_rows(a, axis=None):
@@ -921,7 +947,7 @@ def mask_cols(a, axis=None):
     return mask_rowcols(a, 1)
 
 
-def dot(a,b, strict=False):
+def dot(a, b, strict=False):
     """
     Return the dot product of two arrays.
 
@@ -1114,15 +1140,15 @@ def in1d(ar1, ar2, assume_unique=False):
         ar1, rev_idx = unique(ar1, return_inverse=True)
         ar2 = unique(ar2)
 
-    ar = ma.concatenate( (ar1, ar2) )
+    ar = ma.concatenate((ar1, ar2))
     # We need this to be a stable sort, so always use 'mergesort'
     # here. The values from the first array should always come before
     # the values from the second array.
     order = ar.argsort(kind='mergesort')
     sar = ar[order]
     equal_adj = (sar[1:] == sar[:-1])
-    flag = ma.concatenate( (equal_adj, [False] ) )
-    indx = order.argsort(kind='mergesort')[:len( ar1 )]
+    flag = ma.concatenate((equal_adj, [False]))
+    indx = order.argsort(kind='mergesort')[:len(ar1)]
 
     if assume_unique:
         return flag[indx]
@@ -1199,10 +1225,10 @@ def intersect1d_nu(ar1, ar2):
 def setmember1d(ar1, ar2):
     """ This function is deprecated. Use ma.in1d() instead."""
     ar1 = ma.asanyarray(ar1)
-    ar2 = ma.asanyarray( ar2 )
-    ar = ma.concatenate((ar1, ar2 ))
-    b1 = ma.zeros(ar1.shape, dtype = np.int8)
-    b2 = ma.ones(ar2.shape, dtype = np.int8)
+    ar2 = ma.asanyarray(ar2)
+    ar = ma.concatenate((ar1, ar2))
+    b1 = ma.zeros(ar1.shape, dtype=np.int8)
+    b2 = ma.ones(ar2.shape, dtype=np.int8)
     tt = ma.concatenate((b1, b2))
 
     # We need this to be a stable sort, so always use 'mergesort' here. The
@@ -1213,12 +1239,12 @@ def setmember1d(ar1, ar2):
     aux2 = tt[perm]
 #    flag = ediff1d( aux, 1 ) == 0
     flag = ma.concatenate((aux[1:] == aux[:-1], [False]))
-    ii = ma.where( flag * aux2 )[0]
-    aux = perm[ii+1]
-    perm[ii+1] = perm[ii]
+    ii = ma.where(flag * aux2)[0]
+    aux = perm[ii + 1]
+    perm[ii + 1] = perm[ii]
     perm[ii] = aux
     #
-    indx = perm.argsort(kind='mergesort')[:len( ar1 )]
+    indx = perm.argsort(kind='mergesort')[:len(ar1)]
     #
     return flag[indx]
 
@@ -1246,7 +1272,7 @@ def _covhelper(x, y=None, rowvar=True, allow_masked=True):
         rowvar = True
     # Make sure that rowvar is either 0 or 1
     rowvar = int(bool(rowvar))
-    axis = 1-rowvar
+    axis = 1 - rowvar
     if rowvar:
         tup = (slice(None), None)
     else:
@@ -1267,7 +1293,7 @@ def _covhelper(x, y=None, rowvar=True, allow_masked=True):
                     x.unshare_mask()
                     y.unshare_mask()
                     xmask = x._mask = y._mask = ymask = common_mask
-        x = ma.concatenate((x,y),axis)
+        x = ma.concatenate((x, y), axis)
         xnotmask = np.logical_not(np.concatenate((xmask, ymask), axis)).astype(int)
     x -= x.mean(axis=rowvar)[tup]
     return (x, xnotmask, rowvar)
@@ -1321,10 +1347,10 @@ def cov(x, y=None, rowvar=True, bias=False, allow_masked=True):
     """
     (x, xnotmask, rowvar) = _covhelper(x, y, rowvar, allow_masked)
     if not rowvar:
-        fact = np.dot(xnotmask.T, xnotmask)*1. - (1 - bool(bias))
+        fact = np.dot(xnotmask.T, xnotmask) * 1. - (1 - bool(bias))
         result = (dot(x.T, x.conj(), strict=False) / fact).squeeze()
     else:
-        fact = np.dot(xnotmask, xnotmask.T)*1. - (1 - bool(bias))
+        fact = np.dot(xnotmask, xnotmask.T) * 1. - (1 - bool(bias))
         result = (dot(x, x.T.conj(), strict=False) / fact).squeeze()
     return result
 
@@ -1369,10 +1395,10 @@ def corrcoef(x, y=None, rowvar=True, bias=False, allow_masked=True):
     (x, xnotmask, rowvar) = _covhelper(x, y, rowvar, allow_masked)
     # Compute the covariance matrix
     if not rowvar:
-        fact = np.dot(xnotmask.T, xnotmask)*1. - (1 - bool(bias))
+        fact = np.dot(xnotmask.T, xnotmask) * 1. - (1 - bool(bias))
         c = (dot(x.T, x.conj(), strict=False) / fact).squeeze()
     else:
-        fact = np.dot(xnotmask, xnotmask.T)*1. - (1 - bool(bias))
+        fact = np.dot(xnotmask, xnotmask.T) * 1. - (1 - bool(bias))
         c = (dot(x, x.T.conj(), strict=False) / fact).squeeze()
     # Check whether we have a scalar
     try:
@@ -1384,20 +1410,20 @@ def corrcoef(x, y=None, rowvar=True, bias=False, allow_masked=True):
         _denom = ma.sqrt(ma.multiply.outer(diag, diag))
     else:
         _denom = diagflat(diag)
-        n = x.shape[1-rowvar]
+        n = x.shape[1 - rowvar]
         if rowvar:
-            for i in range(n-1):
-                for j in range(i+1,n):
+            for i in range(n - 1):
+                for j in range(i + 1, n):
                     _x = mask_cols(vstack((x[i], x[j]))).var(axis=1,
-                                                             ddof=1-bias)
-                    _denom[i,j] = _denom[j,i] = ma.sqrt(ma.multiply.reduce(_x))
+                                                             ddof=1 - bias)
+                    _denom[i, j] = _denom[j, i] = ma.sqrt(ma.multiply.reduce(_x))
         else:
-            for i in range(n-1):
-                for j in range(i+1,n):
-                    _x = mask_cols(vstack((x[:,i], x[:,j]))).var(axis=1,
-                                                                 ddof=1-bias)
-                    _denom[i,j] = _denom[j,i] = ma.sqrt(ma.multiply.reduce(_x))
-    return c/_denom
+            for i in range(n - 1):
+                for j in range(i + 1, n):
+                    _x = mask_cols(vstack((x[:, i], x[:, j]))).var(axis=1,
+                                                                 ddof=1 - bias)
+                    _denom[i, j] = _denom[j, i] = ma.sqrt(ma.multiply.reduce(_x))
+    return c / _denom
 
 #####--------------------------------------------------------------------------
 #---- --- Concatenation helpers ---
@@ -1418,7 +1444,7 @@ class MAxisConcatenator(AxisConcatenator):
     def __init__(self, axis=0):
         AxisConcatenator.__init__(self, axis, matrix=False)
 
-    def __getitem__(self,key):
+    def __getitem__(self, key):
         if isinstance(key, str):
             raise MAError, "Unavailable for masked array."
         if type(key) is not tuple:
@@ -1466,7 +1492,7 @@ class MAxisConcatenator(AxisConcatenator):
         if final_dtypedescr is not None:
             for k in scalars:
                 objs[k] = objs[k].astype(final_dtypedescr)
-        res = concatenate(tuple(objs),axis=self.axis)
+        res = concatenate(tuple(objs), axis=self.axis)
         return self._retval(res)
 
 class mr_class(MAxisConcatenator):
@@ -1538,10 +1564,10 @@ def flatnotmasked_edges(a):
     """
     m = getmask(a)
     if m is nomask or not np.any(m):
-        return [0,-1]
+        return [0, -1]
     unmasked = np.flatnonzero(~m)
     if len(unmasked) > 0:
-        return unmasked[[0,-1]]
+        return unmasked[[0, -1]]
     else:
         return None
 
@@ -1591,9 +1617,9 @@ def notmasked_edges(a, axis=None):
     if axis is None or a.ndim == 1:
         return flatnotmasked_edges(a)
     m = getmaskarray(a)
-    idx = array(np.indices(a.shape), mask=np.asarray([m]*a.ndim))
+    idx = array(np.indices(a.shape), mask=np.asarray([m] * a.ndim))
     return [tuple([idx[i].min(axis).compressed() for i in range(a.ndim)]),
-            tuple([idx[i].max(axis).compressed() for i in range(a.ndim)]),]
+            tuple([idx[i].max(axis).compressed() for i in range(a.ndim)]), ]
 
 
 def flatnotmasked_contiguous(a):
@@ -1635,15 +1661,15 @@ def flatnotmasked_contiguous(a):
     """
     m = getmask(a)
     if m is nomask:
-        return (a.size, [0,-1])
+        return (a.size, [0, -1])
     unmasked = np.flatnonzero(~m)
     if len(unmasked) == 0:
         return None
     result = []
-    for (k, group) in itertools.groupby(enumerate(unmasked), lambda (i,x):i-x):
+    for (k, group) in itertools.groupby(enumerate(unmasked), lambda (i, x):i - x):
         tmp = np.array([g[1] for g in group], int)
 #        result.append((tmp.size, tuple(tmp[[0,-1]])))
-        result.append( slice(tmp[0], tmp[-1]) )
+        result.append(slice(tmp[0], tmp[-1]))
     result.sort()
     return result
 
@@ -1690,19 +1716,19 @@ def notmasked_contiguous(a, axis=None):
     a = asarray(a)
     nd = a.ndim
     if nd > 2:
-        raise NotImplementedError,"Currently limited to atmost 2D array."
+        raise NotImplementedError, "Currently limited to atmost 2D array."
     if axis is None or nd == 1:
         return flatnotmasked_contiguous(a)
     #
     result = []
     #
-    other = (axis+1)%2
+    other = (axis + 1) % 2
     idx = [0, 0]
     idx[axis] = slice(None, None)
     #
     for i in range(a.shape[other]):
         idx[other] = i
-        result.append( flatnotmasked_contiguous(a[idx]) )
+        result.append(flatnotmasked_contiguous(a[idx]))
     return result
 
 
@@ -1831,16 +1857,16 @@ def polyfit(x, y, deg, rcond=None, full=False):
         y = mask_rows(y)
         my = getmask(y)
         if my is not nomask:
-            m = mask_or(mx, my[:,0])
+            m = mask_or(mx, my[:, 0])
         else:
             m = mx
     else:
-        raise TypeError,"Expected a 1D or 2D array for y!"
+        raise TypeError, "Expected a 1D or 2D array for y!"
     if m is not nomask:
         x[m] = y[m] = masked
     # Set rcond
     if rcond is None :
-        rcond = len(x)*np.finfo(x.dtype).eps
+        rcond = len(x) * np.finfo(x.dtype).eps
     # Scale x to improve condition number
     scale = abs(x).max()
     if scale != 0 :
diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
index 9e2c48b50..d6cda7e70 100644
--- a/numpy/ma/tests/test_extras.py
+++ b/numpy/ma/tests/test_extras.py
@@ -9,7 +9,7 @@ Adapted from the original test_ma by Pierre Gerard-Marchant
 __author__ = "Pierre GF Gerard-Marchant ($Author: jarrod.millman $)"
 __version__ = '1.0'
 __revision__ = "$Revision: 3473 $"
-__date__     = '$Date: 2007-10-29 17:18:13 +0200 (Mon, 29 Oct 2007) $'
+__date__ = '$Date: 2007-10-29 17:18:13 +0200 (Mon, 29 Oct 2007) $'
 
 import numpy as np
 from numpy.testing import TestCase, run_module_suite
@@ -31,13 +31,13 @@ class TestGeneric(TestCase):
         test = masked_all((2,), dtype=dt)
         control = array([(0, 0), (0, 0)], mask=[(1, 1), (1, 1)], dtype=dt)
         assert_equal(test, control)
-        test = masked_all((2,2), dtype=dt)
+        test = masked_all((2, 2), dtype=dt)
         control = array([[(0, 0), (0, 0)], [(0, 0), (0, 0)]],
                         mask=[[(1, 1), (1, 1)], [(1, 1), (1, 1)]],
                         dtype=dt)
         assert_equal(test, control)
         # Nested dtype
-        dt = np.dtype([('a','f'), ('b', [('ba', 'f'), ('bb', 'f')])])
+        dt = np.dtype([('a', 'f'), ('b', [('ba', 'f'), ('bb', 'f')])])
         test = masked_all((2,), dtype=dt)
         control = array([(1, (1, 1)), (1, (1, 1))],
                          mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt)
@@ -46,7 +46,7 @@ class TestGeneric(TestCase):
         control = array([(1, (1, 1)), (1, (1, 1))],
                          mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt)
         assert_equal(test, control)
-        test = masked_all((1,1), dtype=dt)
+        test = masked_all((1, 1), dtype=dt)
         control = array([[(1, (1, 1))]], mask=[[(1, (1, 1))]], dtype=dt)
         assert_equal(test, control)
 
@@ -65,7 +65,7 @@ class TestGeneric(TestCase):
         control = array([(10, 10), (10, 10)], mask=[(1, 1), (1, 1)], dtype=dt)
         assert_equal(test, control)
         # Nested dtype
-        dt = np.dtype([('a','f'), ('b', [('ba', 'f'), ('bb', 'f')])])
+        dt = np.dtype([('a', 'f'), ('b', [('ba', 'f'), ('bb', 'f')])])
         control = array([(1, (1, 1)), (1, (1, 1))],
                         mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt)
         test = masked_all_like(control)
@@ -85,7 +85,7 @@ class TestGeneric(TestCase):
         a = masked_array(np.arange(10))
         a[[0, 1, 2, 6, 8, 9]] = masked
         test = clump_unmasked(a)
-        control = [slice(3, 6), slice(7, 8),]
+        control = [slice(3, 6), slice(7, 8), ]
         assert_equal(test, control)
 
 
@@ -94,7 +94,7 @@ class TestAverage(TestCase):
     "Several tests of average. Why so many ? Good point..."
     def test_testAverage1(self):
         "Test of average."
-        ott = array([0.,1.,2.,3.], mask=[True, False, False, False])
+        ott = array([0., 1., 2., 3.], mask=[True, False, False, False])
         assert_equal(2.0, average(ott, axis=0))
         assert_equal(2.0, average(ott, weights=[1., 1., 2., 1.]))
         result, wts = average(ott, weights=[1., 1., 2., 1.], returned=1)
@@ -104,28 +104,28 @@ class TestAverage(TestCase):
         assert_equal(average(ott, axis=0).mask, [True])
         ott = array([0., 1., 2., 3.], mask=[True, False, False, False])
         ott = ott.reshape(2, 2)
-        ott[:,1] = masked
+        ott[:, 1] = masked
         assert_equal(average(ott, axis=0), [2.0, 0.0])
         assert_equal(average(ott, axis=1).mask[0], [True])
-        assert_equal([2.,0.], average(ott, axis=0))
+        assert_equal([2., 0.], average(ott, axis=0))
         result, wts = average(ott, axis=0, returned=1)
         assert_equal(wts, [1., 0.])
 
     def test_testAverage2(self):
         "More tests of average."
-        w1 = [0,1,1,1,1,0]
-        w2 = [[0,1,1,1,1,0],[1,0,0,0,0,1]]
+        w1 = [0, 1, 1, 1, 1, 0]
+        w2 = [[0, 1, 1, 1, 1, 0], [1, 0, 0, 0, 0, 1]]
         x = arange(6, dtype=float_)
         assert_equal(average(x, axis=0), 2.5)
         assert_equal(average(x, axis=0, weights=w1), 2.5)
-        y = array([arange(6, dtype=float_), 2.0*arange(6)])
-        assert_equal(average(y, None), np.add.reduce(np.arange(6))*3./12.)
-        assert_equal(average(y, axis=0), np.arange(6) * 3./2.)
+        y = array([arange(6, dtype=float_), 2.0 * arange(6)])
+        assert_equal(average(y, None), np.add.reduce(np.arange(6)) * 3. / 12.)
+        assert_equal(average(y, axis=0), np.arange(6) * 3. / 2.)
         assert_equal(average(y, axis=1),
                      [average(x, axis=0), average(x, axis=0) * 2.0])
-        assert_equal(average(y, None, weights=w2), 20./6.)
+        assert_equal(average(y, None, weights=w2), 20. / 6.)
         assert_equal(average(y, axis=0, weights=w2),
-                     [0.,1.,2.,3.,4.,10.])
+                     [0., 1., 2., 3., 4., 10.])
         assert_equal(average(y, axis=1),
                      [average(x, axis=0), average(x, axis=0) * 2.0])
         m1 = zeros(6)
@@ -139,11 +139,11 @@ class TestAverage(TestCase):
         assert_equal(average(masked_array(x, m5), axis=0), 0.0)
         assert_equal(count(average(masked_array(x, m4), axis=0)), 0)
         z = masked_array(y, m3)
-        assert_equal(average(z, None), 20./6.)
-        assert_equal(average(z, axis=0), [0.,1.,99.,99.,4.0, 7.5])
+        assert_equal(average(z, None), 20. / 6.)
+        assert_equal(average(z, axis=0), [0., 1., 99., 99., 4.0, 7.5])
         assert_equal(average(z, axis=1), [2.5, 5.0])
-        assert_equal(average(z,axis=0, weights=w2),
-                     [0.,1., 99., 99., 4.0, 10.0])
+        assert_equal(average(z, axis=0, weights=w2),
+                     [0., 1., 99., 99., 4.0, 10.0])
 
     def test_testAverage3(self):
         "Yet more tests of average!"
@@ -159,13 +159,13 @@ class TestAverage(TestCase):
         r2, w2 = average(ones((2, 2, 3)), weights=ones((2, 2, 3)), returned=1)
         assert_equal(shape(w2), shape(r2))
         a2d = array([[1, 2], [0, 4]], float)
-        a2dm = masked_array(a2d, [[False, False],[True, False]])
+        a2dm = masked_array(a2d, [[False, False], [True, False]])
         a2da = average(a2d, axis=0)
         assert_equal(a2da, [0.5, 3.0])
         a2dma = average(a2dm, axis=0)
         assert_equal(a2dma, [1.0, 3.0])
         a2dma = average(a2dm, axis=None)
-        assert_equal(a2dma, 7./3.)
+        assert_equal(a2dma, 7. / 3.)
         a2dma = average(a2dm, axis=1)
         assert_equal(a2dma, [1.5, 4.0])
 
@@ -184,33 +184,33 @@ class TestConcatenator(TestCase):
 
     def test_1d(self):
         "Tests mr_ on 1D arrays."
-        assert_array_equal(mr_[1,2,3,4,5,6],array([1,2,3,4,5,6]))
+        assert_array_equal(mr_[1, 2, 3, 4, 5, 6], array([1, 2, 3, 4, 5, 6]))
         b = ones(5)
-        m = [1,0,0,0,0]
-        d = masked_array(b,mask=m)
-        c = mr_[d,0,0,d]
-        self.assertTrue(isinstance(c,MaskedArray) or isinstance(c,core.MaskedArray))
-        assert_array_equal(c,[1,1,1,1,1,0,0,1,1,1,1,1])
-        assert_array_equal(c.mask, mr_[m,0,0,m])
+        m = [1, 0, 0, 0, 0]
+        d = masked_array(b, mask=m)
+        c = mr_[d, 0, 0, d]
+        self.assertTrue(isinstance(c, MaskedArray) or isinstance(c, core.MaskedArray))
+        assert_array_equal(c, [1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1])
+        assert_array_equal(c.mask, mr_[m, 0, 0, m])
 
     def test_2d(self):
         "Tests mr_ on 2D arrays."
-        a_1 = rand(5,5)
-        a_2 = rand(5,5)
-        m_1 = np.round_(rand(5,5),0)
-        m_2 = np.round_(rand(5,5),0)
-        b_1 = masked_array(a_1,mask=m_1)
-        b_2 = masked_array(a_2,mask=m_2)
-        d = mr_['1',b_1,b_2]  # append columns
-        self.assertTrue(d.shape == (5,10))
-        assert_array_equal(d[:,:5],b_1)
-        assert_array_equal(d[:,5:],b_2)
-        assert_array_equal(d.mask, np.r_['1',m_1,m_2])
-        d = mr_[b_1,b_2]
-        self.assertTrue(d.shape == (10,5))
-        assert_array_equal(d[:5,:],b_1)
-        assert_array_equal(d[5:,:],b_2)
-        assert_array_equal(d.mask, np.r_[m_1,m_2])
+        a_1 = rand(5, 5)
+        a_2 = rand(5, 5)
+        m_1 = np.round_(rand(5, 5), 0)
+        m_2 = np.round_(rand(5, 5), 0)
+        b_1 = masked_array(a_1, mask=m_1)
+        b_2 = masked_array(a_2, mask=m_2)
+        d = mr_['1', b_1, b_2]  # append columns
+        self.assertTrue(d.shape == (5, 10))
+        assert_array_equal(d[:, :5], b_1)
+        assert_array_equal(d[:, 5:], b_2)
+        assert_array_equal(d.mask, np.r_['1', m_1, m_2])
+        d = mr_[b_1, b_2]
+        self.assertTrue(d.shape == (10, 5))
+        assert_array_equal(d[:5, :], b_1)
+        assert_array_equal(d[5:, :], b_2)
+        assert_array_equal(d.mask, np.r_[m_1, m_2])
 
 
 
@@ -256,26 +256,26 @@ class TestNotMasked(TestCase):
 
     def test_contiguous(self):
         "Tests notmasked_contiguous"
-        a = masked_array(np.arange(24).reshape(3,8),
-                         mask=[[0,0,0,0,1,1,1,1],
-                               [1,1,1,1,1,1,1,1],
-                               [0,0,0,0,0,0,1,0],])
+        a = masked_array(np.arange(24).reshape(3, 8),
+                         mask=[[0, 0, 0, 0, 1, 1, 1, 1],
+                               [1, 1, 1, 1, 1, 1, 1, 1],
+                               [0, 0, 0, 0, 0, 0, 1, 0], ])
         tmp = notmasked_contiguous(a, None)
-        assert_equal(tmp[-1], slice(23,23,None))
-        assert_equal(tmp[-2], slice(16,21,None))
-        assert_equal(tmp[-3], slice(0,3,None))
+        assert_equal(tmp[-1], slice(23, 23, None))
+        assert_equal(tmp[-2], slice(16, 21, None))
+        assert_equal(tmp[-3], slice(0, 3, None))
         #
         tmp = notmasked_contiguous(a, 0)
         self.assertTrue(len(tmp[-1]) == 1)
         self.assertTrue(tmp[-2] is None)
-        assert_equal(tmp[-3],tmp[-1])
+        assert_equal(tmp[-3], tmp[-1])
         self.assertTrue(len(tmp[0]) == 2)
         #
         tmp = notmasked_contiguous(a, 1)
-        assert_equal(tmp[0][-1], slice(0,3,None))
+        assert_equal(tmp[0][-1], slice(0, 3, None))
         self.assertTrue(tmp[1] is None)
-        assert_equal(tmp[2][-1], slice(7,7,None))
-        assert_equal(tmp[2][-2], slice(0,5,None))
+        assert_equal(tmp[2][-1], slice(7, 7, None))
+        assert_equal(tmp[2][-2], slice(0, 5, None))
 
 
 
@@ -283,114 +283,114 @@ class Test2DFunctions(TestCase):
     "Tests 2D functions"
     def test_compress2d(self):
         "Tests compress2d"
-        x = array(np.arange(9).reshape(3,3), mask=[[1,0,0],[0,0,0],[0,0,0]])
-        assert_equal(compress_rowcols(x), [[4,5],[7,8]] )
-        assert_equal(compress_rowcols(x,0), [[3,4,5],[6,7,8]] )
-        assert_equal(compress_rowcols(x,1), [[1,2],[4,5],[7,8]] )
-        x = array(x._data, mask=[[0,0,0],[0,1,0],[0,0,0]])
-        assert_equal(compress_rowcols(x), [[0,2],[6,8]] )
-        assert_equal(compress_rowcols(x,0), [[0,1,2],[6,7,8]] )
-        assert_equal(compress_rowcols(x,1), [[0,2],[3,5],[6,8]] )
-        x = array(x._data, mask=[[1,0,0],[0,1,0],[0,0,0]])
-        assert_equal(compress_rowcols(x), [[8]] )
-        assert_equal(compress_rowcols(x,0), [[6,7,8]] )
-        assert_equal(compress_rowcols(x,1,), [[2],[5],[8]] )
-        x = array(x._data, mask=[[1,0,0],[0,1,0],[0,0,1]])
-        assert_equal(compress_rowcols(x).size, 0 )
-        assert_equal(compress_rowcols(x,0).size, 0 )
-        assert_equal(compress_rowcols(x,1).size, 0 )
+        x = array(np.arange(9).reshape(3, 3), mask=[[1, 0, 0], [0, 0, 0], [0, 0, 0]])
+        assert_equal(compress_rowcols(x), [[4, 5], [7, 8]])
+        assert_equal(compress_rowcols(x, 0), [[3, 4, 5], [6, 7, 8]])
+        assert_equal(compress_rowcols(x, 1), [[1, 2], [4, 5], [7, 8]])
+        x = array(x._data, mask=[[0, 0, 0], [0, 1, 0], [0, 0, 0]])
+        assert_equal(compress_rowcols(x), [[0, 2], [6, 8]])
+        assert_equal(compress_rowcols(x, 0), [[0, 1, 2], [6, 7, 8]])
+        assert_equal(compress_rowcols(x, 1), [[0, 2], [3, 5], [6, 8]])
+        x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 0]])
+        assert_equal(compress_rowcols(x), [[8]])
+        assert_equal(compress_rowcols(x, 0), [[6, 7, 8]])
+        assert_equal(compress_rowcols(x, 1,), [[2], [5], [8]])
+        x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 1]])
+        assert_equal(compress_rowcols(x).size, 0)
+        assert_equal(compress_rowcols(x, 0).size, 0)
+        assert_equal(compress_rowcols(x, 1).size, 0)
     #
     def test_mask_rowcols(self):
         "Tests mask_rowcols."
-        x = array(np.arange(9).reshape(3,3), mask=[[1,0,0],[0,0,0],[0,0,0]])
-        assert_equal(mask_rowcols(x).mask, [[1,1,1],[1,0,0],[1,0,0]] )
-        assert_equal(mask_rowcols(x,0).mask, [[1,1,1],[0,0,0],[0,0,0]] )
-        assert_equal(mask_rowcols(x,1).mask, [[1,0,0],[1,0,0],[1,0,0]] )
-        x = array(x._data, mask=[[0,0,0],[0,1,0],[0,0,0]])
-        assert_equal(mask_rowcols(x).mask, [[0,1,0],[1,1,1],[0,1,0]] )
-        assert_equal(mask_rowcols(x,0).mask, [[0,0,0],[1,1,1],[0,0,0]] )
-        assert_equal(mask_rowcols(x,1).mask, [[0,1,0],[0,1,0],[0,1,0]] )
-        x = array(x._data, mask=[[1,0,0],[0,1,0],[0,0,0]])
-        assert_equal(mask_rowcols(x).mask, [[1,1,1],[1,1,1],[1,1,0]] )
-        assert_equal(mask_rowcols(x,0).mask, [[1,1,1],[1,1,1],[0,0,0]] )
-        assert_equal(mask_rowcols(x,1,).mask, [[1,1,0],[1,1,0],[1,1,0]] )
-        x = array(x._data, mask=[[1,0,0],[0,1,0],[0,0,1]])
+        x = array(np.arange(9).reshape(3, 3), mask=[[1, 0, 0], [0, 0, 0], [0, 0, 0]])
+        assert_equal(mask_rowcols(x).mask, [[1, 1, 1], [1, 0, 0], [1, 0, 0]])
+        assert_equal(mask_rowcols(x, 0).mask, [[1, 1, 1], [0, 0, 0], [0, 0, 0]])
+        assert_equal(mask_rowcols(x, 1).mask, [[1, 0, 0], [1, 0, 0], [1, 0, 0]])
+        x = array(x._data, mask=[[0, 0, 0], [0, 1, 0], [0, 0, 0]])
+        assert_equal(mask_rowcols(x).mask, [[0, 1, 0], [1, 1, 1], [0, 1, 0]])
+        assert_equal(mask_rowcols(x, 0).mask, [[0, 0, 0], [1, 1, 1], [0, 0, 0]])
+        assert_equal(mask_rowcols(x, 1).mask, [[0, 1, 0], [0, 1, 0], [0, 1, 0]])
+        x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 0]])
+        assert_equal(mask_rowcols(x).mask, [[1, 1, 1], [1, 1, 1], [1, 1, 0]])
+        assert_equal(mask_rowcols(x, 0).mask, [[1, 1, 1], [1, 1, 1], [0, 0, 0]])
+        assert_equal(mask_rowcols(x, 1,).mask, [[1, 1, 0], [1, 1, 0], [1, 1, 0]])
+        x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 1]])
         self.assertTrue(mask_rowcols(x).all() is masked)
-        self.assertTrue(mask_rowcols(x,0).all() is masked)
-        self.assertTrue(mask_rowcols(x,1).all() is masked)
+        self.assertTrue(mask_rowcols(x, 0).all() is masked)
+        self.assertTrue(mask_rowcols(x, 1).all() is masked)
         self.assertTrue(mask_rowcols(x).mask.all())
-        self.assertTrue(mask_rowcols(x,0).mask.all())
-        self.assertTrue(mask_rowcols(x,1).mask.all())
+        self.assertTrue(mask_rowcols(x, 0).mask.all())
+        self.assertTrue(mask_rowcols(x, 1).mask.all())
     #
     def test_dot(self):
         "Tests dot product"
-        n = np.arange(1,7)
-        #
-        m = [1,0,0,0,0,0]
-        a = masked_array(n, mask=m).reshape(2,3)
-        b = masked_array(n, mask=m).reshape(3,2)
-        c = dot(a,b,True)
-        assert_equal(c.mask, [[1,1],[1,0]])
-        c = dot(b,a,True)
-        assert_equal(c.mask, [[1,1,1],[1,0,0],[1,0,0]])
-        c = dot(a,b,False)
+        n = np.arange(1, 7)
+        #
+        m = [1, 0, 0, 0, 0, 0]
+        a = masked_array(n, mask=m).reshape(2, 3)
+        b = masked_array(n, mask=m).reshape(3, 2)
+        c = dot(a, b, True)
+        assert_equal(c.mask, [[1, 1], [1, 0]])
+        c = dot(b, a, True)
+        assert_equal(c.mask, [[1, 1, 1], [1, 0, 0], [1, 0, 0]])
+        c = dot(a, b, False)
         assert_equal(c, np.dot(a.filled(0), b.filled(0)))
-        c = dot(b,a,False)
+        c = dot(b, a, False)
         assert_equal(c, np.dot(b.filled(0), a.filled(0)))
         #
-        m = [0,0,0,0,0,1]
-        a = masked_array(n, mask=m).reshape(2,3)
-        b = masked_array(n, mask=m).reshape(3,2)
-        c = dot(a,b,True)
-        assert_equal(c.mask,[[0,1],[1,1]])
-        c = dot(b,a,True)
-        assert_equal(c.mask, [[0,0,1],[0,0,1],[1,1,1]])
-        c = dot(a,b,False)
+        m = [0, 0, 0, 0, 0, 1]
+        a = masked_array(n, mask=m).reshape(2, 3)
+        b = masked_array(n, mask=m).reshape(3, 2)
+        c = dot(a, b, True)
+        assert_equal(c.mask, [[0, 1], [1, 1]])
+        c = dot(b, a, True)
+        assert_equal(c.mask, [[0, 0, 1], [0, 0, 1], [1, 1, 1]])
+        c = dot(a, b, False)
         assert_equal(c, np.dot(a.filled(0), b.filled(0)))
-        assert_equal(c, dot(a,b))
-        c = dot(b,a,False)
+        assert_equal(c, dot(a, b))
+        c = dot(b, a, False)
         assert_equal(c, np.dot(b.filled(0), a.filled(0)))
         #
-        m = [0,0,0,0,0,0]
-        a = masked_array(n, mask=m).reshape(2,3)
-        b = masked_array(n, mask=m).reshape(3,2)
-        c = dot(a,b)
-        assert_equal(c.mask,nomask)
-        c = dot(b,a)
-        assert_equal(c.mask,nomask)
-        #
-        a = masked_array(n, mask=[1,0,0,0,0,0]).reshape(2,3)
-        b = masked_array(n, mask=[0,0,0,0,0,0]).reshape(3,2)
-        c = dot(a,b,True)
-        assert_equal(c.mask,[[1,1],[0,0]])
-        c = dot(a,b,False)
-        assert_equal(c, np.dot(a.filled(0),b.filled(0)))
-        c = dot(b,a,True)
-        assert_equal(c.mask,[[1,0,0],[1,0,0],[1,0,0]])
-        c = dot(b,a,False)
-        assert_equal(c, np.dot(b.filled(0),a.filled(0)))
-        #
-        a = masked_array(n, mask=[0,0,0,0,0,1]).reshape(2,3)
-        b = masked_array(n, mask=[0,0,0,0,0,0]).reshape(3,2)
-        c = dot(a,b,True)
-        assert_equal(c.mask,[[0,0],[1,1]])
-        c = dot(a,b)
-        assert_equal(c, np.dot(a.filled(0),b.filled(0)))
-        c = dot(b,a,True)
-        assert_equal(c.mask,[[0,0,1],[0,0,1],[0,0,1]])
-        c = dot(b,a,False)
+        m = [0, 0, 0, 0, 0, 0]
+        a = masked_array(n, mask=m).reshape(2, 3)
+        b = masked_array(n, mask=m).reshape(3, 2)
+        c = dot(a, b)
+        assert_equal(c.mask, nomask)
+        c = dot(b, a)
+        assert_equal(c.mask, nomask)
+        #
+        a = masked_array(n, mask=[1, 0, 0, 0, 0, 0]).reshape(2, 3)
+        b = masked_array(n, mask=[0, 0, 0, 0, 0, 0]).reshape(3, 2)
+        c = dot(a, b, True)
+        assert_equal(c.mask, [[1, 1], [0, 0]])
+        c = dot(a, b, False)
+        assert_equal(c, np.dot(a.filled(0), b.filled(0)))
+        c = dot(b, a, True)
+        assert_equal(c.mask, [[1, 0, 0], [1, 0, 0], [1, 0, 0]])
+        c = dot(b, a, False)
+        assert_equal(c, np.dot(b.filled(0), a.filled(0)))
+        #
+        a = masked_array(n, mask=[0, 0, 0, 0, 0, 1]).reshape(2, 3)
+        b = masked_array(n, mask=[0, 0, 0, 0, 0, 0]).reshape(3, 2)
+        c = dot(a, b, True)
+        assert_equal(c.mask, [[0, 0], [1, 1]])
+        c = dot(a, b)
+        assert_equal(c, np.dot(a.filled(0), b.filled(0)))
+        c = dot(b, a, True)
+        assert_equal(c.mask, [[0, 0, 1], [0, 0, 1], [0, 0, 1]])
+        c = dot(b, a, False)
         assert_equal(c, np.dot(b.filled(0), a.filled(0)))
         #
-        a = masked_array(n, mask=[0,0,0,0,0,1]).reshape(2,3)
-        b = masked_array(n, mask=[0,0,1,0,0,0]).reshape(3,2)
-        c = dot(a,b,True)
-        assert_equal(c.mask,[[1,0],[1,1]])
-        c = dot(a,b,False)
-        assert_equal(c, np.dot(a.filled(0),b.filled(0)))
-        c = dot(b,a,True)
-        assert_equal(c.mask,[[0,0,1],[1,1,1],[0,0,1]])
-        c = dot(b,a,False)
-        assert_equal(c, np.dot(b.filled(0),a.filled(0)))
+        a = masked_array(n, mask=[0, 0, 0, 0, 0, 1]).reshape(2, 3)
+        b = masked_array(n, mask=[0, 0, 1, 0, 0, 0]).reshape(3, 2)
+        c = dot(a, b, True)
+        assert_equal(c.mask, [[1, 0], [1, 1]])
+        c = dot(a, b, False)
+        assert_equal(c, np.dot(a.filled(0), b.filled(0)))
+        c = dot(b, a, True)
+        assert_equal(c.mask, [[0, 0, 1], [1, 1, 1], [0, 0, 1]])
+        c = dot(b, a, False)
+        assert_equal(c, np.dot(b.filled(0), a.filled(0)))
 
 
 
@@ -398,51 +398,63 @@ class TestApplyAlongAxis(TestCase):
     #
     "Tests 2D functions"
     def test_3d(self):
-        a = arange(12.).reshape(2,2,3)
+        a = arange(12.).reshape(2, 2, 3)
         def myfunc(b):
             return b[1]
-        xa = apply_along_axis(myfunc,2,a)
-        assert_equal(xa,[[1,4],[7,10]])
+        xa = apply_along_axis(myfunc, 2, a)
+        assert_equal(xa, [[1, 4], [7, 10]])
+
 
 
+class TestApplyOverAxes(TestCase):
+    "Tests apply_over_axes"
+    def test_basic(self):
+        a = arange(24).reshape(2, 3, 4)
+        test = apply_over_axes(np.sum, a, [0, 2])
+        ctrl = np.array([[[ 60], [ 92], [124]]])
+        assert_equal(test, ctrl)
+        a[(a % 2).astype(np.bool)] = masked
+        test = apply_over_axes(np.sum, a, [0, 2])
+        ctrl = np.array([[[ 30], [ 44], [60]]])
+
 
 class TestMedian(TestCase):
     #
     def test_2d(self):
         "Tests median w/ 2D"
-        (n,p) = (101,30)
-        x = masked_array(np.linspace(-1.,1.,n),)
+        (n, p) = (101, 30)
+        x = masked_array(np.linspace(-1., 1., n),)
         x[:10] = x[-10:] = masked
-        z = masked_array(np.empty((n,p), dtype=float))
-        z[:,0] = x[:]
+        z = masked_array(np.empty((n, p), dtype=float))
+        z[:, 0] = x[:]
         idx = np.arange(len(x))
-        for i in range(1,p):
+        for i in range(1, p):
             np.random.shuffle(idx)
-            z[:,i] = x[idx]
-        assert_equal(median(z[:,0]), 0)
+            z[:, i] = x[idx]
+        assert_equal(median(z[:, 0]), 0)
         assert_equal(median(z), 0)
         assert_equal(median(z, axis=0), np.zeros(p))
         assert_equal(median(z.T, axis=1), np.zeros(p))
     #
     def test_2d_waxis(self):
         "Tests median w/ 2D arrays and different axis."
-        x = masked_array(np.arange(30).reshape(10,3))
+        x = masked_array(np.arange(30).reshape(10, 3))
         x[:3] = x[-3:] = masked
         assert_equal(median(x), 14.5)
-        assert_equal(median(x, axis=0), [13.5,14.5,15.5])
-        assert_equal(median(x,axis=1), [0,0,0,10,13,16,19,0,0,0])
-        assert_equal(median(x,axis=1).mask, [1,1,1,0,0,0,0,1,1,1])
+        assert_equal(median(x, axis=0), [13.5, 14.5, 15.5])
+        assert_equal(median(x, axis=1), [0, 0, 0, 10, 13, 16, 19, 0, 0, 0])
+        assert_equal(median(x, axis=1).mask, [1, 1, 1, 0, 0, 0, 0, 1, 1, 1])
     #
     def test_3d(self):
         "Tests median w/ 3D"
-        x = np.ma.arange(24).reshape(3,4,2)
-        x[x%3==0] = masked
-        assert_equal(median(x,0), [[12,9],[6,15],[12,9],[18,15]])
-        x.shape = (4,3,2)
-        assert_equal(median(x,0),[[99,10],[11,99],[13,14]])
-        x = np.ma.arange(24).reshape(4,3,2)
-        x[x%5==0] = masked
-        assert_equal(median(x,0), [[12,10],[8,9],[16,17]])
+        x = np.ma.arange(24).reshape(3, 4, 2)
+        x[x % 3 == 0] = masked
+        assert_equal(median(x, 0), [[12, 9], [6, 15], [12, 9], [18, 15]])
+        x.shape = (4, 3, 2)
+        assert_equal(median(x, 0), [[99, 10], [11, 99], [13, 14]])
+        x = np.ma.arange(24).reshape(4, 3, 2)
+        x[x % 5 == 0] = masked
+        assert_equal(median(x, 0), [[12, 10], [8, 9], [16, 17]])
 
 
 
@@ -461,7 +473,7 @@ class TestCov(TestCase):
 
     def test_2d_wo_missing(self):
         "Test cov on 1 2D variable w/o missing values"
-        x = self.data.reshape(3,4)
+        x = self.data.reshape(3, 4)
         assert_almost_equal(np.cov(x), cov(x))
         assert_almost_equal(np.cov(x, rowvar=False), cov(x, rowvar=False))
         assert_almost_equal(np.cov(x, rowvar=False, bias=True),
@@ -495,19 +507,19 @@ class TestCov(TestCase):
         "Test cov on 2D variable w/ missing value"
         x = self.data
         x[-1] = masked
-        x = x.reshape(3,4)
+        x = x.reshape(3, 4)
         valid = np.logical_not(getmaskarray(x)).astype(int)
         frac = np.dot(valid, valid.T)
-        xf = (x - x.mean(1)[:,None]).filled(0)
-        assert_almost_equal(cov(x), np.cov(xf) * (x.shape[1]-1) / (frac - 1.))
+        xf = (x - x.mean(1)[:, None]).filled(0)
+        assert_almost_equal(cov(x), np.cov(xf) * (x.shape[1] - 1) / (frac - 1.))
         assert_almost_equal(cov(x, bias=True),
                             np.cov(xf, bias=True) * x.shape[1] / frac)
         frac = np.dot(valid.T, valid)
         xf = (x - x.mean(0)).filled(0)
         assert_almost_equal(cov(x, rowvar=False),
-                            np.cov(xf, rowvar=False) * (x.shape[0]-1)/(frac - 1.))
+                            np.cov(xf, rowvar=False) * (x.shape[0] - 1) / (frac - 1.))
         assert_almost_equal(cov(x, rowvar=False, bias=True),
-                            np.cov(xf, rowvar=False, bias=True) * x.shape[0]/frac)
+                            np.cov(xf, rowvar=False, bias=True) * x.shape[0] / frac)
 
 
 
@@ -527,7 +539,7 @@ class TestCorrcoef(TestCase):
 
     def test_2d_wo_missing(self):
         "Test corrcoef on 1 2D variable w/o missing values"
-        x = self.data.reshape(3,4)
+        x = self.data.reshape(3, 4)
         assert_almost_equal(np.corrcoef(x), corrcoef(x))
         assert_almost_equal(np.corrcoef(x, rowvar=False),
                             corrcoef(x, rowvar=False))
@@ -562,11 +574,11 @@ class TestCorrcoef(TestCase):
         "Test corrcoef on 2D variable w/ missing value"
         x = self.data
         x[-1] = masked
-        x = x.reshape(3,4)
+        x = x.reshape(3, 4)
 
         test = corrcoef(x)
         control = np.corrcoef(x)
-        assert_almost_equal(test[:-1,:-1], control[:-1,:-1])
+        assert_almost_equal(test[:-1, :-1], control[:-1, :-1])
 
 
 
@@ -576,27 +588,27 @@ class TestPolynomial(TestCase):
         "Tests polyfit"
         # On ndarrays
         x = np.random.rand(10)
-        y = np.random.rand(20).reshape(-1,2)
-        assert_almost_equal(polyfit(x,y,3),np.polyfit(x,y,3))
+        y = np.random.rand(20).reshape(-1, 2)
+        assert_almost_equal(polyfit(x, y, 3), np.polyfit(x, y, 3))
         # ON 1D maskedarrays
         x = x.view(MaskedArray)
         x[0] = masked
         y = y.view(MaskedArray)
-        y[0,0] = y[-1,-1] = masked
+        y[0, 0] = y[-1, -1] = masked
         #
-        (C,R,K,S,D) = polyfit(x,y[:,0],3,full=True)
-        (c,r,k,s,d) = np.polyfit(x[1:], y[1:,0].compressed(), 3, full=True)
-        for (a,a_) in zip((C,R,K,S,D),(c,r,k,s,d)):
+        (C, R, K, S, D) = polyfit(x, y[:, 0], 3, full=True)
+        (c, r, k, s, d) = np.polyfit(x[1:], y[1:, 0].compressed(), 3, full=True)
+        for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)):
             assert_almost_equal(a, a_)
         #
-        (C,R,K,S,D) = polyfit(x,y[:,-1],3,full=True)
-        (c,r,k,s,d) = np.polyfit(x[1:-1], y[1:-1,-1], 3, full=True)
-        for (a,a_) in zip((C,R,K,S,D),(c,r,k,s,d)):
+        (C, R, K, S, D) = polyfit(x, y[:, -1], 3, full=True)
+        (c, r, k, s, d) = np.polyfit(x[1:-1], y[1:-1, -1], 3, full=True)
+        for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)):
             assert_almost_equal(a, a_)
         #
-        (C,R,K,S,D) = polyfit(x,y,3,full=True)
-        (c,r,k,s,d) = np.polyfit(x[1:-1], y[1:-1,:], 3, full=True)
-        for (a,a_) in zip((C,R,K,S,D),(c,r,k,s,d)):
+        (C, R, K, S, D) = polyfit(x, y, 3, full=True)
+        (c, r, k, s, d) = np.polyfit(x[1:-1], y[1:-1, :], 3, full=True)
+        for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)):
             assert_almost_equal(a, a_)
 
 
@@ -632,7 +644,7 @@ class TestArraySetOps(TestCase):
         "Test all masked"
         data = masked_array([1, 1, 1], mask=True)
         test = unique(data, return_index=True, return_inverse=True)
-        assert_equal(test[0], masked_array([1,], mask=[True]))
+        assert_equal(test[0], masked_array([1, ], mask=[True]))
         assert_equal(test[1], [0])
         assert_equal(test[2], [0, 0, 0])
         #
@@ -642,10 +654,10 @@ class TestArraySetOps(TestCase):
         assert_equal(test[0], masked_array(masked))
         assert_equal(test[1], [0])
         assert_equal(test[2], [0])
-    
+
     def test_ediff1d(self):
         "Tests mediff1d"
-        x = masked_array(np.arange(5), mask=[1,0,0,0,1])
+        x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1])
         control = array([1, 1, 1, 4], mask=[1, 0, 0, 1])
         test = ediff1d(x)
         assert_equal(test, control)
@@ -654,14 +666,14 @@ class TestArraySetOps(TestCase):
     #
     def test_ediff1d_tobegin(self):
         "Test ediff1d w/ to_begin"
-        x = masked_array(np.arange(5), mask=[1,0,0,0,1])
+        x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1])
         test = ediff1d(x, to_begin=masked)
         control = array([0, 1, 1, 1, 4], mask=[1, 1, 0, 0, 1])
         assert_equal(test, control)
         assert_equal(test.data, control.data)
         assert_equal(test.mask, control.mask)
         #
-        test = ediff1d(x, to_begin=[1,2,3])
+        test = ediff1d(x, to_begin=[1, 2, 3])
         control = array([1, 2, 3, 1, 1, 1, 4], mask=[0, 0, 0, 1, 0, 0, 1])
         assert_equal(test, control)
         assert_equal(test.data, control.data)
@@ -669,14 +681,14 @@ class TestArraySetOps(TestCase):
     #
     def test_ediff1d_toend(self):
         "Test ediff1d w/ to_end"
-        x = masked_array(np.arange(5), mask=[1,0,0,0,1])
+        x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1])
         test = ediff1d(x, to_end=masked)
         control = array([1, 1, 1, 4, 0], mask=[1, 0, 0, 1, 1])
         assert_equal(test, control)
         assert_equal(test.data, control.data)
         assert_equal(test.mask, control.mask)
         #
-        test = ediff1d(x, to_end=[1,2,3])
+        test = ediff1d(x, to_end=[1, 2, 3])
         control = array([1, 1, 1, 4, 1, 2, 3], mask=[1, 0, 0, 1, 0, 0, 0])
         assert_equal(test, control)
         assert_equal(test.data, control.data)
@@ -684,14 +696,14 @@ class TestArraySetOps(TestCase):
     #
     def test_ediff1d_tobegin_toend(self):
         "Test ediff1d w/ to_begin and to_end"
-        x = masked_array(np.arange(5), mask=[1,0,0,0,1])
+        x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1])
         test = ediff1d(x, to_end=masked, to_begin=masked)
         control = array([0, 1, 1, 1, 4, 0], mask=[1, 1, 0, 0, 1, 1])
         assert_equal(test, control)
         assert_equal(test.data, control.data)
         assert_equal(test.mask, control.mask)
         #
-        test = ediff1d(x, to_end=[1,2,3], to_begin=masked)
+        test = ediff1d(x, to_end=[1, 2, 3], to_begin=masked)
         control = array([0, 1, 1, 1, 4, 1, 2, 3], mask=[1, 1, 0, 0, 1, 0, 0, 0])
         assert_equal(test, control)
         assert_equal(test.data, control.data)
@@ -735,8 +747,8 @@ class TestArraySetOps(TestCase):
         test = setxor1d(a, b)
         assert_equal(test, array([3, 4, 7, -1], mask=[0, 0, 0, 1]))
         #
-        a = array( [1, 2, 3] )
-        b = array( [6, 5, 4] )
+        a = array([1, 2, 3])
+        b = array([6, 5, 4])
         test = setxor1d(a, b)
         assert(isinstance(test, MaskedArray))
         assert_equal(test, [1, 2, 3, 4, 5, 6])
@@ -747,10 +759,10 @@ class TestArraySetOps(TestCase):
         assert(isinstance(test, MaskedArray))
         assert_equal(test, [1, 2, 3, 4, 5, 6])
         #
-        assert_array_equal([], setxor1d([],[]))
+        assert_array_equal([], setxor1d([], []))
 
 
-    def test_in1d( self ):
+    def test_in1d(self):
         "Test in1d"
         a = array([1, 2, 5, 7, -1], mask=[0, 0, 0, 0, 1])
         b = array([1, 2, 3, 4, 5, -1], mask=[0, 0, 0, 0, 0, 1])
@@ -762,10 +774,10 @@ class TestArraySetOps(TestCase):
         test = in1d(a, b)
         assert_equal(test, [True, True, False, True, True])
         #
-        assert_array_equal([], in1d([],[]))
+        assert_array_equal([], in1d([], []))
 
 
-    def test_union1d( self ):
+    def test_union1d(self):
         "Test union1d"
         a = array([1, 2, 5, 7, 5, -1], mask=[0, 0, 0, 0, 0, 1])
         b = array([1, 2, 3, 4, 5, -1], mask=[0, 0, 0, 0, 0, 1])
@@ -773,10 +785,10 @@ class TestArraySetOps(TestCase):
         control = array([1, 2, 3, 4, 5, 7, -1], mask=[0, 0, 0, 0, 0, 0, 1])
         assert_equal(test, control)
         #
-        assert_array_equal([], union1d([],[]))
+        assert_array_equal([], union1d([], []))
 
 
-    def test_setdiff1d( self ):
+    def test_setdiff1d(self):
         "Test setdiff1d"
         a = array([6, 5, 4, 7, 7, 1, 2, 1], mask=[0, 0, 0, 0, 0, 0, 0, 1])
         b = array([2, 4, 3, 3, 2, 1, 5])
@@ -790,9 +802,9 @@ class TestArraySetOps(TestCase):
 
     def test_setdiff1d_char_array(self):
         "Test setdiff1d_charray"
-        a = np.array(['a','b','c'])
-        b = np.array(['a','b','s'])
-        assert_array_equal(setdiff1d(a,b), np.array(['c']))
+        a = np.array(['a', 'b', 'c'])
+        b = np.array(['a', 'b', 's'])
+        assert_array_equal(setdiff1d(a, b), np.array(['c']))
 
 
 class TestShapeBase(TestCase):