Precond: stationary not as real inverse

integrators/_details/interfaces/tpsPreconditioner.py

@@ -159,6 +159,7 @@ class TPSPreconditioner(_preconditionerInterface._PreconditionerInterface):
     import scipy.sparse.linalg
     import scipy.sparse as sp
     import numpy as np
+    import pyamg
 
     if(self.parameters.alphaPreconditioner == None):
       A = self._A_stat[:self._A_stat.shape[0]//3, :self._A_stat.shape[0]//3]
@@ -166,40 +167,72 @@ class TPSPreconditioner(_preconditionerInterface._PreconditionerInterface):
       A = self._A_stat_precond[:self._A_stat_precond.shape[0]//3 \
         , :self._A_stat_precond.shape[0]//3]
 
-    Ainv = np.zeros(A.shape, dtype=float)
-
-    e = np.zeros(A.shape[0], dtype=float)
-    e[0] = 1.0
-
-    print("Inverting matrix for stationary preconditioner")
-    sol, succ = scipy.sparse.linalg.cg(A, e \
-      , tol=self._solvetol, maxiter=4000)
-    idx = (abs(sol) > self._solvetol)
-    Ainv[idx,0] = sol[idx]
-    for j in range(1, A.shape[0]):
-      e[j-1] = 0.0
-      e[j] = 1.0
-      sol, succ = scipy.sparse.linalg.cg(A, e, tol=self._solvetol, maxiter=4000)
-      idx = (abs(sol) > self._solvetol)
-      Ainv[idx, j] = sol[idx]
-      if (j % 100 == 0):
-        print("Finished ", j, "out of ", A.shape[0] \
-          , "cols of stationary preconditioner")
-
-    N = A.shape[0]
-
-    cooAinv = sp.coo_matrix(Ainv)
-    self._preStationary = sp.csr_matrix(( \
-      np.append(cooAinv.data, cooAinv.data), ( \
-      np.append(2*cooAinv.row, 2*cooAinv.row+1), \
-      np.append(2*cooAinv.col, 2*cooAinv.col+1) )))
+    # trying AMG, SPLU, CG preconditioner, CG as default. 
+    # way faster than true inverse, especially for large NV
+    opt = "cg" # "ml", "splu"
+    if opt == "ml":
+      ml = pyamg.ruge_stuben_solver(A)
+      pre = ml.aspreconditioner()
       
-    self._preStationary.eliminate_zeros()
-    
-    stationary = lambda x: self._preStationary.dot(x)
+      Pre = lambda x: np.vstack( \
+        (pre * x[0::2], pre * x[1::2])).reshape((-1,),order='F')
+    elif opt == "splu":
+      cgs = pyamg.coarse_grid_solver("splu")
+      Pre = lambda x: np.vstack( \
+        (cgs(A, x[0::2]), cgs(A, x[1::2]))).reshape((-1,),order='F')
+    elif opt == "cg":
+      metaPre = 1.0 / A.diagonal()
+      MetaPre = lambda x: metaPre * x
+      P = scipy.sparse.linalg.LinearOperator(A.shape, matvec=MetaPre)
+      Pre = lambda x: np.vstack((
+        scipy.sparse.linalg.cg(A, x[0::2], tol=self._solvetol/10, M=P)[0], \
+        scipy.sparse.linalg.cg(A, x[1::2], tol=self._solvetol/10, M=P)[0])) \
+        .reshape((-1,),order='F')
 
     self._preconditioner = scipy.sparse.linalg.LinearOperator( \
-      self._preStationary.shape, matvec=stationary)
+      (A.shape[0]*2, A.shape[0]*2), matvec=Pre)
+
+
+#    Ainv = np.zeros(A.shape, dtype=float)
+#
+#    e = np.zeros(A.shape[0], dtype=float)
+#    e[0] = 1.0
+#
+#    print("Inverting matrix for stationary preconditioner")
+#    sol, succ = scipy.sparse.linalg.cg(A, e \
+#      , tol=self._solvetol, maxiter=4000)
+#    idx = (abs(sol) > self._solvetol)
+#    Ainv[idx,0] = sol[idx]
+#    for j in range(1, A.shape[0]):
+#      e[j-1] = 0.0
+#      e[j] = 1.0
+#      sol, succ = scipy.sparse.linalg.cg(A, e, tol=self._solvetol, maxiter=4000)
+#      idx = (abs(sol) > self._solvetol)
+#      Ainv[idx, j] = sol[idx]
+#      if (j % 100 == 0):
+#        print("Finished ", j, "out of ", A.shape[0] \
+#          , "cols of stationary preconditioner")
+#
+#    N = A.shape[0]
+#
+#    csrAinv = sp.csr_matrix(Ainv)
+##    cooAinv = sp.coo_matrix(Ainv)
+##    self._preStationary = sp.csr_matrix(( \
+##      np.append(cooAinv.data, cooAinv.data), ( \
+##      np.append(2*cooAinv.row, 2*cooAinv.row+1), \
+##      np.append(2*cooAinv.col, 2*cooAinv.col+1) )))
+#      
+##    self._preStationary.eliminate_zeros()
+#    
+###    stationary = lambda x: np.concatenate(( \
+###      csrAinv.dot(x[:x.size//2]), csrAinv.dot(x[x.size//2:])))
+#    # TODO THIS solves the above problem, saves memory also. TODO adapt to solving stuff
+#    stationary = lambda x: np.vstack((csrAinv.dot(x[0::2]),csrAinv.dot(x[1::2]))).reshape((-1,),order='F')
+##    stationary = lambda x: self._preStationary.dot(x)
+#
+#    self._preconditioner = scipy.sparse.linalg.LinearOperator( \
+#      (2*csrAinv.shape[0], 2*csrAinv.shape[1]), matvec=stationary)
+##      self._preStationary.shape, matvec=stationary)
 
 
 #------------------------------------------------------------------------------#