cleanup, unified notation, shortened impl

integrators/_details/interfaces/tpsPreconditioner.py

@@ -126,24 +126,19 @@ class TPSPreconditioner(_preconditionerInterface._PreconditionerInterface):
     if opt == "ml":
       ml = pyamg.ruge_stuben_solver(A)
       pre = ml.aspreconditioner()
-      
-      self._prePractical = lambda y: np.concatenate(( \
-        (pre * y[:y.size//3], pre * y[y.size//3:2*y.size//3], pre * y[2*y.size//3:])))
+      pre1D = lambda z: pre * z
     elif opt == "splu":
       cgs = pyamg.coarse_grid_solver("splu")
-      self._prePractical = lambda y: np.concatenate(( \
-        (cgs(A, y[:y.size//3]), \
-        cgs(A, y[y.size//3:2*y.size//3]), \
-        cgs(A, y[2*y.size//3:]))))
+      pre1D = lambda z: cgs(A, z)
     elif opt == "cg":
       metaPre = 1.0 / A.diagonal()
       metaP = lambda x: metaPre*x
       P = scipy.sparse.linalg.LinearOperator(A.shape, matvec=metaP)
-      self._prePractical = lambda y: np.concatenate(( \
-        scipy.sparse.linalg.cg(A, y[:y.size//3], tol=self._solvetol/10, M=P)[0], \
-        scipy.sparse.linalg.cg(A, y[y.size//3:2*y.size//3], tol=self._solvetol/10, M=P)[0], \
-        scipy.sparse.linalg.cg(A, y[2*y.size//3:], tol=self._solvetol/10, M=P)[0]))
+      pre1D = lambda z: scipy.sparse.linalg.cg(A, z, tol=self._solvetol/10, M=P)[0]
   
+    self._prePractical = lambda y: np.concatenate((
+      pre1D(y[:y.size//3]), pre1D(y[y.size//3:2*y.size//3]), pre1D(y[2*y.size//3:])))
+
     practical2D = lambda x: \
       self._Q.dot(self._prePractical(self._Q.transpose().dot(x)))
 
@@ -174,24 +169,21 @@ class TPSPreconditioner(_preconditionerInterface._PreconditionerInterface):
     if opt == "ml":
       ml = pyamg.ruge_stuben_solver(A)
       pre = ml.aspreconditioner()
-      
-      Pre = lambda x: np.vstack( \
-        (pre * x[0::2], pre * x[1::2])).reshape((-1,),order='F')
+      pre1D = lambda z: pre * z
     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')
+      pre1D = lambda z: cgs(A, z)
     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')
+      pre1D = lambda z: scipy.sparse.linalg.cg(A, z, tol=self._solvetol/10, M=P)[0]
 
+    stationary2D = lambda x: \
+      np.vstack((pre1D(x[0::2]), pre1D(x[1::2]))).reshape((-1,),order='F')
+    
     self._preconditioner = scipy.sparse.linalg.LinearOperator( \
-      (A.shape[0]*2, A.shape[0]*2), matvec=Pre)
+      (A.shape[0]*2, A.shape[0]*2), matvec=stationary2D)
 
 
 #------------------------------------------------------------------------------#