source: branches/source_numpy_conversion/anuga/utilities/test_cg_solve.py @ 6768

Last change on this file since 6768 was 5902, checked in by rwilson, 16 years ago

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