1from numpy import array, kron, diag 2from numpy.testing import assert_, assert_equal 3 4from scipy.sparse import spfuncs 5from scipy.sparse import csr_matrix, csc_matrix, bsr_matrix 6from scipy.sparse._sparsetools import (csr_scale_rows, csr_scale_columns, 7 bsr_scale_rows, bsr_scale_columns) 8from scipy.sparse.sputils import matrix 9 10 11class TestSparseFunctions: 12 def test_scale_rows_and_cols(self): 13 D = matrix([[1,0,0,2,3], 14 [0,4,0,5,0], 15 [0,0,6,7,0]]) 16 17 #TODO expose through function 18 S = csr_matrix(D) 19 v = array([1,2,3]) 20 csr_scale_rows(3,5,S.indptr,S.indices,S.data,v) 21 assert_equal(S.todense(), diag(v)*D) 22 23 S = csr_matrix(D) 24 v = array([1,2,3,4,5]) 25 csr_scale_columns(3,5,S.indptr,S.indices,S.data,v) 26 assert_equal(S.todense(), D@diag(v)) 27 28 # blocks 29 E = kron(D,[[1,2],[3,4]]) 30 S = bsr_matrix(E,blocksize=(2,2)) 31 v = array([1,2,3,4,5,6]) 32 bsr_scale_rows(3,5,2,2,S.indptr,S.indices,S.data,v) 33 assert_equal(S.todense(), diag(v)@E) 34 35 S = bsr_matrix(E,blocksize=(2,2)) 36 v = array([1,2,3,4,5,6,7,8,9,10]) 37 bsr_scale_columns(3,5,2,2,S.indptr,S.indices,S.data,v) 38 assert_equal(S.todense(), E@diag(v)) 39 40 E = kron(D,[[1,2,3],[4,5,6]]) 41 S = bsr_matrix(E,blocksize=(2,3)) 42 v = array([1,2,3,4,5,6]) 43 bsr_scale_rows(3,5,2,3,S.indptr,S.indices,S.data,v) 44 assert_equal(S.todense(), diag(v)@E) 45 46 S = bsr_matrix(E,blocksize=(2,3)) 47 v = array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]) 48 bsr_scale_columns(3,5,2,3,S.indptr,S.indices,S.data,v) 49 assert_equal(S.todense(), E@diag(v)) 50 51 def test_estimate_blocksize(self): 52 mats = [] 53 mats.append([[0,1],[1,0]]) 54 mats.append([[1,1,0],[0,0,1],[1,0,1]]) 55 mats.append([[0],[0],[1]]) 56 mats = [array(x) for x in mats] 57 58 blks = [] 59 blks.append([[1]]) 60 blks.append([[1,1],[1,1]]) 61 blks.append([[1,1],[0,1]]) 62 blks.append([[1,1,0],[1,0,1],[1,1,1]]) 63 blks = [array(x) for x in blks] 64 65 for A in mats: 66 for B in blks: 67 X = kron(A,B) 68 r,c = spfuncs.estimate_blocksize(X) 69 assert_(r >= B.shape[0]) 70 assert_(c >= B.shape[1]) 71 72 def test_count_blocks(self): 73 def gold(A,bs): 74 R,C = bs 75 I,J = A.nonzero() 76 return len(set(zip(I//R,J//C))) 77 78 mats = [] 79 mats.append([[0]]) 80 mats.append([[1]]) 81 mats.append([[1,0]]) 82 mats.append([[1,1]]) 83 mats.append([[0,1],[1,0]]) 84 mats.append([[1,1,0],[0,0,1],[1,0,1]]) 85 mats.append([[0],[0],[1]]) 86 87 for A in mats: 88 for B in mats: 89 X = kron(A,B) 90 Y = csr_matrix(X) 91 for R in range(1,6): 92 for C in range(1,6): 93 assert_equal(spfuncs.count_blocks(Y, (R, C)), gold(X, (R, C))) 94 95 X = kron([[1,1,0],[0,0,1],[1,0,1]],[[1,1]]) 96 Y = csc_matrix(X) 97 assert_equal(spfuncs.count_blocks(X, (1, 2)), gold(X, (1, 2))) 98 assert_equal(spfuncs.count_blocks(Y, (1, 2)), gold(X, (1, 2))) 99