import warnings from numpy.testing import * import numpy.lib from numpy.lib import * from numpy.core import * from numpy import matrix, asmatrix class TestAny(TestCase): def test_basic(self): y1 = [0,0,1,0] y2 = [0,0,0,0] y3 = [1,0,1,0] assert(any(y1)) assert(any(y3)) assert(not any(y2)) def test_nd(self): y1 = [[0,0,0],[0,1,0],[1,1,0]] assert(any(y1)) assert_array_equal(sometrue(y1,axis=0),[1,1,0]) assert_array_equal(sometrue(y1,axis=1),[0,1,1]) class TestAll(TestCase): def test_basic(self): y1 = [0,1,1,0] y2 = [0,0,0,0] y3 = [1,1,1,1] assert(not all(y1)) assert(all(y3)) assert(not all(y2)) assert(all(~array(y2))) def test_nd(self): y1 = [[0,0,1],[0,1,1],[1,1,1]] assert(not all(y1)) assert_array_equal(alltrue(y1,axis=0),[0,0,1]) assert_array_equal(alltrue(y1,axis=1),[0,0,1]) class TestAverage(TestCase): def test_basic(self): y1 = array([1,2,3]) assert(average(y1,axis=0) == 2.) y2 = array([1.,2.,3.]) assert(average(y2,axis=0) == 2.) y3 = [0.,0.,0.] assert(average(y3,axis=0) == 0.) y4 = ones((4,4)) y4[0,1] = 0 y4[1,0] = 2 assert_almost_equal(y4.mean(0), average(y4, 0)) assert_almost_equal(y4.mean(1), average(y4, 1)) y5 = rand(5,5) assert_almost_equal(y5.mean(0), average(y5, 0)) assert_almost_equal(y5.mean(1), average(y5, 1)) y6 = matrix(rand(5,5)) assert_array_equal(y6.mean(0), average(y6,0)) def test_weights(self): y = arange(10) w = arange(10) assert_almost_equal(average(y, weights=w), (arange(10)**2).sum()*1./arange(10).sum()) y1 = array([[1,2,3],[4,5,6]]) w0 = [1,2] actual = average(y1,weights=w0,axis=0) desired = array([3.,4.,5.]) assert_almost_equal(actual, desired) w1 = [0,0,1] desired = array([3., 6.]) assert_almost_equal(average(y1, weights=w1, axis=1), desired) # This should raise an error. Can we test for that ? # assert_equal(average(y1, weights=w1), 9./2.) # 2D Case w2 = [[0,0,1],[0,0,2]] desired = array([3., 6.]) assert_array_equal(average(y1, weights=w2, axis=1), desired) assert_equal(average(y1, weights=w2), 5.) def test_returned(self): y = array([[1,2,3],[4,5,6]]) # No weights avg, scl = average(y, returned=True) assert_equal(scl, 6.) avg, scl = average(y, 0, returned=True) assert_array_equal(scl, array([2.,2.,2.])) avg, scl = average(y, 1, returned=True) assert_array_equal(scl, array([3.,3.])) # With weights w0 = [1,2] avg, scl = average(y, weights=w0, axis=0, returned=True) assert_array_equal(scl, array([3., 3., 3.])) w1 = [1,2,3] avg, scl = average(y, weights=w1, axis=1, returned=True) assert_array_equal(scl, array([6., 6.])) w2 = [[0,0,1],[1,2,3]] avg, scl = average(y, weights=w2, axis=1, returned=True) assert_array_equal(scl, array([1.,6.])) class TestSelect(TestCase): def _select(self,cond,values,default=0): output = [] for m in range(len(cond)): output += [V[m] for V,C in zip(values,cond) if C[m]] or [default] return output def test_basic(self): choices = [array([1,2,3]), array([4,5,6]), array([7,8,9])] conditions = [array([0,0,0]), array([0,1,0]), array([0,0,1])] assert_array_equal(select(conditions,choices,default=15), self._select(conditions,choices,default=15)) assert_equal(len(choices),3) assert_equal(len(conditions),3) class TestInsert(TestCase): def test_basic(self): a = [1,2,3] assert_equal(insert(a,0,1), [1,1,2,3]) assert_equal(insert(a,3,1), [1,2,3,1]) assert_equal(insert(a,[1,1,1],[1,2,3]), [1,1,2,3,2,3]) class TestAmax(TestCase): def test_basic(self): a = [3,4,5,10,-3,-5,6.0] assert_equal(amax(a),10.0) b = [[3,6.0, 9.0], [4,10.0,5.0], [8,3.0,2.0]] assert_equal(amax(b,axis=0),[8.0,10.0,9.0]) assert_equal(amax(b,axis=1),[9.0,10.0,8.0]) class TestAmin(TestCase): def test_basic(self): a = [3,4,5,10,-3,-5,6.0] assert_equal(amin(a),-5.0) b = [[3,6.0, 9.0], [4,10.0,5.0], [8,3.0,2.0]] assert_equal(amin(b,axis=0),[3.0,3.0,2.0]) assert_equal(amin(b,axis=1),[3.0,4.0,2.0]) class TestPtp(TestCase): def test_basic(self): a = [3,4,5,10,-3,-5,6.0] assert_equal(ptp(a,axis=0),15.0) b = [[3,6.0, 9.0], [4,10.0,5.0], [8,3.0,2.0]] assert_equal(ptp(b,axis=0),[5.0,7.0,7.0]) assert_equal(ptp(b,axis=-1),[6.0,6.0,6.0]) class TestCumsum(TestCase): def test_basic(self): ba = [1,2,10,11,6,5,4] ba2 = [[1,2,3,4],[5,6,7,9],[10,3,4,5]] for ctype in [int8,uint8,int16,uint16,int32,uint32, float32,float64,complex64,complex128]: a = array(ba,ctype) a2 = array(ba2,ctype) assert_array_equal(cumsum(a,axis=0), array([1,3,13,24,30,35,39],ctype)) assert_array_equal(cumsum(a2,axis=0), array([[1,2,3,4],[6,8,10,13], [16,11,14,18]],ctype)) assert_array_equal(cumsum(a2,axis=1), array([[1,3,6,10], [5,11,18,27], [10,13,17,22]],ctype)) class TestProd(TestCase): def test_basic(self): ba = [1,2,10,11,6,5,4] ba2 = [[1,2,3,4],[5,6,7,9],[10,3,4,5]] for ctype in [int16,uint16,int32,uint32, float32,float64,complex64,complex128]: a = array(ba,ctype) a2 = array(ba2,ctype) if ctype in ['1', 'b']: self.failUnlessRaises(ArithmeticError, prod, a) self.failUnlessRaises(ArithmeticError, prod, a2, 1) self.failUnlessRaises(ArithmeticError, prod, a) else: assert_equal(prod(a,axis=0),26400) assert_array_equal(prod(a2,axis=0), array([50,36,84,180],ctype)) assert_array_equal(prod(a2,axis=-1),array([24, 1890, 600],ctype)) class TestCumprod(TestCase): def test_basic(self): ba = [1,2,10,11,6,5,4] ba2 = [[1,2,3,4],[5,6,7,9],[10,3,4,5]] for ctype in [int16,uint16,int32,uint32, float32,float64,complex64,complex128]: a = array(ba,ctype) a2 = array(ba2,ctype) if ctype in ['1', 'b']: self.failUnlessRaises(ArithmeticError, cumprod, a) self.failUnlessRaises(ArithmeticError, cumprod, a2, 1) self.failUnlessRaises(ArithmeticError, cumprod, a) else: assert_array_equal(cumprod(a,axis=-1), array([1, 2, 20, 220, 1320, 6600, 26400],ctype)) assert_array_equal(cumprod(a2,axis=0), array([[ 1, 2, 3, 4], [ 5, 12, 21, 36], [50, 36, 84, 180]],ctype)) assert_array_equal(cumprod(a2,axis=-1), array([[ 1, 2, 6, 24], [ 5, 30, 210, 1890], [10, 30, 120, 600]],ctype)) class TestDiff(TestCase): def test_basic(self): x = [1,4,6,7,12] out = array([3,2,1,5]) out2 = array([-1,-1,4]) out3 = array([0,5]) assert_array_equal(diff(x),out) assert_array_equal(diff(x,n=2),out2) assert_array_equal(diff(x,n=3),out3) def test_nd(self): x = 20*rand(10,20,30) out1 = x[:,:,1:] - x[:,:,:-1] out2 = out1[:,:,1:] - out1[:,:,:-1] out3 = x[1:,:,:] - x[:-1,:,:] out4 = out3[1:,:,:] - out3[:-1,:,:] assert_array_equal(diff(x),out1) assert_array_equal(diff(x,n=2),out2) assert_array_equal(diff(x,axis=0),out3) assert_array_equal(diff(x,n=2,axis=0),out4) class TestGradient(TestCase): def test_basic(self): x = array([[1,1],[3,4]]) dx = [array([[2.,3.],[2.,3.]]), array([[0.,0.],[1.,1.]])] assert_array_equal(gradient(x), dx) def test_badargs(self): # for 2D array, gradient can take 0,1, or 2 extra args x = array([[1,1],[3,4]]) assert_raises(SyntaxError, gradient, x, array([1.,1.]), array([1.,1.]), array([1.,1.])) class TestAngle(TestCase): def test_basic(self): x = [1+3j,sqrt(2)/2.0+1j*sqrt(2)/2,1,1j,-1,-1j,1-3j,-1+3j] y = angle(x) yo = [arctan(3.0/1.0),arctan(1.0),0,pi/2,pi,-pi/2.0, -arctan(3.0/1.0),pi-arctan(3.0/1.0)] z = angle(x,deg=1) zo = array(yo)*180/pi assert_array_almost_equal(y,yo,11) assert_array_almost_equal(z,zo,11) class TestTrimZeros(TestCase): """ only testing for integer splits. """ def test_basic(self): a= array([0,0,1,2,3,4,0]) res = trim_zeros(a) assert_array_equal(res,array([1,2,3,4])) def test_leading_skip(self): a= array([0,0,1,0,2,3,4,0]) res = trim_zeros(a) assert_array_equal(res,array([1,0,2,3,4])) def test_trailing_skip(self): a= array([0,0,1,0,2,3,0,4,0]) res = trim_zeros(a) assert_array_equal(res,array([1,0,2,3,0,4])) class TestExtins(TestCase): def test_basic(self): a = array([1,3,2,1,2,3,3]) b = extract(a>1,a) assert_array_equal(b,[3,2,2,3,3]) def test_place(self): a = array([1,4,3,2,5,8,7]) place(a,[0,1,0,1,0,1,0],[2,4,6]) assert_array_equal(a,[1,2,3,4,5,6,7]) def test_both(self): a = rand(10) mask = a > 0.5 ac = a.copy() c = extract(mask, a) place(a,mask,0) place(a,mask,c) assert_array_equal(a,ac) class TestVectorize(TestCase): def test_simple(self): def addsubtract(a,b): if a > b: return a - b else: return a + b f = vectorize(addsubtract) r = f([0,3,6,9],[1,3,5,7]) assert_array_equal(r,[1,6,1,2]) def test_scalar(self): def addsubtract(a,b): if a > b: return a - b else: return a + b f = vectorize(addsubtract) r = f([0,3,6,9],5) assert_array_equal(r,[5,8,1,4]) def test_large(self): x = linspace(-3,2,10000) f = vectorize(lambda x: x) y = f(x) assert_array_equal(y, x) class TestDigitize(TestCase): def test_forward(self): x = arange(-6,5) bins = arange(-5,5) assert_array_equal(digitize(x,bins),arange(11)) def test_reverse(self): x = arange(5,-6,-1) bins = arange(5,-5,-1) assert_array_equal(digitize(x,bins),arange(11)) def test_random(self): x = rand(10) bin = linspace(x.min(), x.max(), 10) assert all(digitize(x,bin) != 0) class TestUnwrap(TestCase): def test_simple(self): #check that unwrap removes jumps greather that 2*pi assert_array_equal(unwrap([1,1+2*pi]),[1,1]) #check that unwrap maintans continuity assert(all(diff(unwrap(rand(10)*100))3, [4, 0]) assert y.ndim == 0 assert y == 0 def compare_results(res,desired): for i in range(len(desired)): assert_array_equal(res[i],desired[i]) if __name__ == "__main__": run_module_suite()