Changeset 2841

Timestamp:

May 10, 2006, 4:08:08 PM (18 years ago)

Author:

duncan

Message:

cg_solve now works if b is an array

Location:

inundation/utilities

Files:

• cg_solve.py (modified) (4 diffs)
• test_cg_solve.py (modified) (3 diffs)

inundation/utilities/cg_solve.py

@@ -3 +3 @@
 class ConvergenceError(exceptions.Exception): pass
 
+from Numeric import dot, array, Float, zeros
+   
 import logging, logging.config
 logger = logging.getLogger('cg_solve')
@@ -13 +15 @@
 
 def conjugate_gradient(A,b,x0=None,imax=10000,tol=1.0e-8,iprint=0):
+    """
+    Try to solve linear equation Ax = b using
+    conjugate gradient method
+
+    If b is an array, solve it as if it was a set of vectors, solving each
+    vector.
+    """
+    
+    if x0 is None:
+        x0 = zeros(b.shape, typecode=Float)
+    else:
+        x0 = array(x0, typecode=Float)
+
+    b  = array(b, typecode=Float)
+    if len(b.shape) != 1 :
+       
+        for i in range(b.shape[1]):
+            x0[:,i] = _conjugate_gradient(A, b[:,i], x0[:,i],
+                                          imax, tol, iprint)
+    else:
+        x0 = _conjugate_gradient(A, b, x0, imax, tol, iprint)
+
+    return x0
+    
+def _conjugate_gradient(A,b,x0=None,imax=10000,tol=1.0e-8,iprint=0):
    """
    Try to solve linear equation Ax = b using
@@ -29 +56 @@
    """
 
-   from Numeric import dot, array, Float, zeros
 
    b  = array(b, typecode=Float)
@@ -75 +101 @@
 
    return x
+

inundation/utilities/test_cg_solve.py

@@ -8 +8 @@
 from Numeric import dot, allclose, array, transpose, arange, ones, Float
 from utilities.cg_solve import *
+from utilities.cg_solve import _conjugate_gradient
 from utilities.sparse import Sparse, Sparse_CSR
 
@@ -178 +179 @@
 
         try:
-            x = conjugate_gradient(A,xe,xe,iprint=0)
+            x = _conjugate_gradient(A,xe,xe,iprint=0)
         except VectorShapeError:
             pass
@@ -185 +186 @@
             raise TestError, msg
 
+
+    def test_sparse_solve_matrix(self):
+        """Solve Small Sparse Matrix"""
+
+        A = [[2.0, -1.0, 0.0, 0.0 ],
+             [-1.0, 2.0, -1.0, 0.0],
+             [0.0, -1.0, 2.0, -1.0],
+             [0.0,0.0, -1.0, 2.0]]
+
+        A = Sparse(A)
+
+        xe = [[0.0, 0.0],[1.0, 1.0],[2.0 ,2.0],[3.0, 3.0]]
+        b = A*xe
+        x = [[0.0, 0.0],[0.0, 0.0],[0.0 ,0.0],[0.0, 0.0]]
+        x = conjugate_gradient(A,b,x,iprint=0)
+
+        assert allclose(x,xe)
 
 #-------------------------------------------------------------