diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/ma/extras.py | 180 | ||||
| -rw-r--r-- | numpy/ma/tests/test_extras.py | 434 |
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): |