source: anuga_core/source/anuga/utilities/test_cg_solve.py @ 3512

Last change on this file since 3512 was 2841, checked in by duncan, 19 years ago

cg_solve now works if b is an array

File size: 5.2 KB
Line 
1#!/usr/bin/env python
2
3import exceptions
4class TestError(exceptions.Exception): pass
5import unittest
6
7
8from Numeric import dot, allclose, array, transpose, arange, ones, Float
9from utilities.cg_solve import *
10from utilities.cg_solve import _conjugate_gradient
11from utilities.sparse import Sparse, Sparse_CSR
12
13
14class Test_CG_Solve(unittest.TestCase):
15
16    def test_sparse_solve(self):
17        """Solve Small Sparse Matrix"""
18
19        A = [[2.0, -1.0, 0.0, 0.0 ],
20             [-1.0, 2.0, -1.0, 0.0],
21             [0.0, -1.0, 2.0, -1.0],
22             [0.0,0.0, -1.0, 2.0]]
23
24        A = Sparse(A)
25
26        xe = [0.0, 1.0, 2.0, 3.0]
27        b  = A*xe
28        x =  [0.0, 0.0, 0.0, 0.0]
29
30        x = conjugate_gradient(A,b,x,iprint=0)
31
32        assert allclose(x,xe)
33
34    def test_max_iter(self):
35        """Test max iteration Small Sparse Matrix"""
36
37        A = [[2.0, -1.0, 0.0, 0.0 ],
38             [-1.0, 2.0, -1.0, 0.0],
39             [0.0, -1.0, 2.0, -1.0],
40             [0.0,0.0, -1.0, 2.0]]
41
42        A = Sparse(A)
43
44        xe = [0.0, 1.0, 2.0, 3.0]
45        b  = A*xe
46        x =  [0.0, 0.0, 0.0, 0.0]
47
48        try:
49            x = conjugate_gradient(A,b,x,iprint=0,imax=2)
50        except ConvergenceError:
51            pass
52        else:
53            msg = 'Should have raised exception'
54            raise TestError, msg
55
56
57    def test_solve_large(self):
58        """Standard 1d laplacian """
59
60        n = 50
61        A = Sparse(n,n)
62
63        for i in arange(0,n):
64            A[i,i] = 1.0
65            if i > 0 :
66                A[i,i-1] = -0.5
67            if i < n-1 :
68                A[i,i+1] = -0.5
69
70        xe = ones( (n,), Float)
71
72        b  = A*xe
73        x = conjugate_gradient(A,b,b,tol=1.0e-5,iprint=1)
74
75        assert allclose(x,xe)
76
77    def test_solve_large_2d(self):
78        """Standard 2d laplacian"""
79
80        n = 20
81        m = 10
82
83        A = Sparse(m*n, m*n)
84
85        for i in arange(0,n):
86            for j in arange(0,m):
87                I = j+m*i
88                A[I,I] = 4.0
89                if i > 0  :
90                    A[I,I-m] = -1.0
91                if i < n-1 :
92                    A[I,I+m] = -1.0
93                if j > 0  :
94                    A[I,I-1] = -1.0
95                if j < m-1 :
96                    A[I,I+1] = -1.0
97
98        xe = ones( (n*m,), Float)
99
100        b  = A*xe
101        x = conjugate_gradient(A,b,b,iprint=0)
102
103        assert allclose(x,xe)
104
105    def test_solve_large_2d_csr_matrix(self):
106        """Standard 2d laplacian with csr format
107        """
108
109        n = 100
110        m = 100
111
112        A = Sparse(m*n, m*n)
113
114        for i in arange(0,n):
115            for j in arange(0,m):
116                I = j+m*i
117                A[I,I] = 4.0
118                if i > 0  :
119                    A[I,I-m] = -1.0
120                if i < n-1 :
121                    A[I,I+m] = -1.0
122                if j > 0  :
123                    A[I,I-1] = -1.0
124                if j < m-1 :
125                    A[I,I+1] = -1.0
126
127        xe = ones( (n*m,), Float)
128
129        # Convert to csr format
130        #print 'start covert'
131        A = Sparse_CSR(A)
132        #print 'finish covert'
133        b = A*xe
134        x = conjugate_gradient(A,b,b,iprint=20)
135
136        assert allclose(x,xe)
137
138
139    def test_solve_large_2d_with_default_guess(self):
140        """Standard 2d laplacian using default first guess"""
141
142        n = 20
143        m = 10
144
145        A = Sparse(m*n, m*n)
146
147        for i in arange(0,n):
148            for j in arange(0,m):
149                I = j+m*i
150                A[I,I] = 4.0
151                if i > 0  :
152                    A[I,I-m] = -1.0
153                if i < n-1 :
154                    A[I,I+m] = -1.0
155                if j > 0  :
156                    A[I,I-1] = -1.0
157                if j < m-1 :
158                    A[I,I+1] = -1.0
159
160        xe = ones( (n*m,), Float)
161
162        b  = A*xe
163        x = conjugate_gradient(A,b)
164
165        assert allclose(x,xe)
166
167
168    def test_vector_shape_error(self):
169        """Raise VectorShapeError"""
170
171        A = [[2.0, -1.0, 0.0, 0.0 ],
172             [-1.0, 2.0, -1.0, 0.0],
173             [0.0, -1.0, 2.0, -1.0],
174             [0.0,0.0, -1.0, 2.0]]
175
176        A = Sparse(A)
177
178        xe = [[0.0,2.0], [1.0,3.0], [2.0,4.0], [3.0,2.0]]
179
180        try:
181            x = _conjugate_gradient(A,xe,xe,iprint=0)
182        except VectorShapeError:
183            pass
184        else:
185            msg = 'Should have raised exception'
186            raise TestError, msg
187
188
189    def test_sparse_solve_matrix(self):
190        """Solve Small Sparse Matrix"""
191
192        A = [[2.0, -1.0, 0.0, 0.0 ],
193             [-1.0, 2.0, -1.0, 0.0],
194             [0.0, -1.0, 2.0, -1.0],
195             [0.0,0.0, -1.0, 2.0]]
196
197        A = Sparse(A)
198
199        xe = [[0.0, 0.0],[1.0, 1.0],[2.0 ,2.0],[3.0, 3.0]]
200        b = A*xe
201        x = [[0.0, 0.0],[0.0, 0.0],[0.0 ,0.0],[0.0, 0.0]]
202        x = conjugate_gradient(A,b,x,iprint=0)
203
204        assert allclose(x,xe)
205
206#-------------------------------------------------------------
207if __name__ == "__main__":
208     suite = unittest.makeSuite(Test_CG_Solve,'test')
209     #runner = unittest.TextTestRunner(verbosity=2)
210     runner = unittest.TextTestRunner(verbosity=1)
211     runner.run(suite)
212
Note: See TracBrowser for help on using the repository browser.