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