Matrix functions in sweeper

examples/fwsw/plot_stability.py

@@ -18,8 +18,8 @@
     N_s = 100
     N_f = 400
 
-    lam_s_max = 12.0
-    lam_f_max = 24.0
+    lam_s_max = 3.0
+    lam_f_max = 12.0
     lambda_s = 1j*np.linspace(0.0, lam_s_max, N_s)
     lambda_f = 1j*np.linspace(0.0, lam_f_max, N_f)
 
@@ -31,7 +31,7 @@
     swparams = {}
     swparams['collocation_class'] = collclass.CollGaussLegendre
     swparams['num_nodes'] = 3
-    K = 0
+    K = 3
     do_coll_update = True
 
     #
@@ -47,10 +47,6 @@
     nnodes  = step.levels[0].sweep.coll.num_nodes
     level   = step.levels[0]
     problem = level.prob
-
-    QE = level.sweep.QE[1:,1:]
-    QI = level.sweep.QI[1:,1:]
-    Q  = level.sweep.coll.Qmat[1:,1:]
 
     stab = np.zeros((N_f, N_s), dtype='complex')
 
@@ -59,15 +55,9 @@
         lambda_fast = lambda_f[j]
         lambda_slow = lambda_s[i]
         if K is not 0:
-
-          LHS = np.eye(nnodes) - step.status.dt*( lambda_fast*QI + lambda_slow*QE )
-          RHS = step.status.dt*( (lambda_fast+lambda_slow)*Q - (lambda_fast*QI + lambda_slow*QE) )
-
-          Pinv = np.linalg.inv(LHS)
-          Mat_sweep = np.linalg.matrix_power(Pinv.dot(RHS), K)
-          for k in range(0,K):
-            Mat_sweep = Mat_sweep + np.linalg.matrix_power(Pinv.dot(RHS),k).dot(Pinv)
-
+          lambdas = [ lambda_fast, lambda_slow ]
+          LHS, RHS = level.sweep.get_scalar_problems_sweeper_mats( lambdas = lambdas )
+          Mat_sweep = level.sweep.get_scalar_problems_manysweep_mat( nsweeps = K, lambdas = lambdas )
         else:
           # Compute stability function of collocation solution
           Mat_sweep = np.linalg.inv(np.eye(nnodes)-step.status.dt*(lambda_fast + lambda_slow)*Q)

pySDC/sweeper_classes/imex_1st_order.py

@@ -129,7 +129,6 @@ def update_nodes(self):
 
         return None
 
-
     def compute_end_point(self):
         """
         Compute u at the right point of the interval
@@ -155,3 +154,45 @@ def compute_end_point(self):
               L.uend += L.tau[-1]
 
         return None
+
+    def get_sweeper_mats(self):
+      """
+      Returns the three matrices Q, QI, QE which define the sweeper. The first row and column, corresponding to the left starting value, are removed to correspond to the notation
+      introduced in Ruprecht & Speck, ``Spectral deferred corrections with fast-wave slow-wave splitting'', 2016
+      """
+      QE = self.QE[1:,1:]
+      QI = self.QI[1:,1:]
+      Q  = self.coll.Qmat[1:,1:]
+      return QE, QI, Q
+
+    def get_scalar_problems_sweeper_mats(self, lambdas=[None, None]):
+      """
+      For a scalar problem, an IMEX-SDC sweep can be written in matrix formulation. This function returns the corresponding matrices. 
+      See Ruprecht & Speck, ``Spectral deferred corrections with fast-wave slow-wave splitting'', 2016 for the derivation.
+
+      The first entry in lambdas is lambda_fast, the second is lambda_slow.
+      """
+      QE, QI, Q = self.get_sweeper_mats()
+      if lambdas==[None,None]:
+        pass
+        # should use lambdas from attached problem and make sure it is a scalar IMEX 
+        raise NotImplementedError("At the moment, the values for lambda have to be provided")
+      else:
+        lambda_fast = lambdas[0]
+        lambda_slow = lambdas[1]
+      nnodes = self.coll.num_nodes
+      dt = self.level.dt
+      LHS = np.eye(nnodes) - dt*( lambda_fast*QI + lambda_slow*QE )
+      RHS = dt*( (lambda_fast+lambda_slow)*Q - (lambda_fast*QI + lambda_slow*QE) )
+      return LHS, RHS
+
+    def get_scalar_problems_manysweep_mat(self, nsweeps, lambdas=[None,None]):
+      """
+      For a scalar problem, K sweeps of IMEX-SDC can be written in matrix form.
+      """
+      LHS, RHS = self.get_scalar_problems_sweeper_mats(lambdas = lambdas)
+      Pinv = np.linalg.inv(LHS)
+      Mat_sweep = np.linalg.matrix_power(Pinv.dot(RHS), nsweeps)
+      for k in range(0,nsweeps):
+        Mat_sweep += np.linalg.matrix_power(Pinv.dot(RHS), k).dot(Pinv)
+      return Mat_sweep

tests/test_imexsweeper.py

@@ -32,21 +32,6 @@ def setupLevelStepProblem(self):
     problem = level.prob
     return step, level, problem, nnodes
 
-  def setupQMatrices(self, level):
-    QE = level.sweep.QE[1:,1:]
-    QI = level.sweep.QI[1:,1:]
-    Q  = level.sweep.coll.Qmat[1:,1:]
-    return QE, QI, Q
-
-  def setupSweeperMatrices(self, step, level, problem):
-    nnodes  = step.levels[0].sweep.coll.num_nodes
-    # Build SDC sweep matrix
-    QE, QI, Q = self.setupQMatrices(level)
-    dt = step.status.dt
-    LHS = np.eye(nnodes) - step.status.dt*( problem.lambda_f[0]*QI + problem.lambda_s[0]*QE )
-    RHS = step.status.dt*( (problem.lambda_f[0]+problem.lambda_s[0])*Q - (problem.lambda_f[0]*QI + problem.lambda_s[0]*QE) )
-    return LHS, RHS
-
   #
   # General setUp function used by all tests
   #
@@ -118,7 +103,8 @@ def test_sweepequalmatrix(self):
       # Perform node-to-node SDC sweep
       level.sweep.update_nodes()
 
-      LHS, RHS = self.setupSweeperMatrices(step, level, problem)
+      lambdas = [ problem.lambda_f[0] , problem.lambda_s[0] ]
+      LHS, RHS = level.sweep.get_scalar_problems_sweeper_mats( lambdas = lambdas )
 
       unew = np.linalg.inv(LHS).dot( u0full + RHS.dot(u0full) )
       usweep = np.array([ level.u[l].values.flatten() for l in range(1,nnodes+1) ])
@@ -137,14 +123,14 @@ def test_updateformula(self):
       # Perform update step in sweeper
       level.sweep.update_nodes()
       ustages = np.array([ level.u[l].values.flatten() for l in range(1,nnodes+1) ])
-
       # Compute end value through provided function
       level.sweep.compute_end_point()
       uend_sweep = level.uend.values
       # Compute end value from matrix formulation
       uend_mat   = self.pparams['u0'] + step.status.dt*level.sweep.coll.weights.dot(ustages*(problem.lambda_s[0] + problem.lambda_f[0]))
       assert np.linalg.norm(uend_sweep - uend_mat, np.infty)<1e-14, "Update formula in sweeper gives different result than matrix update formula"
 
+
   #
   # Compute the exact collocation solution by matrix inversion and make sure it is a fixed point
   #
@@ -155,7 +141,7 @@ def test_collocationinvariant(self):
       level.sweep.predict()
       u0full = np.array([ level.u[l].values.flatten() for l in range(1,nnodes+1) ])
 
-      QE, QI, Q = self.setupQMatrices(level)
+      QE, QI, Q = level.sweep.get_sweeper_mats()
 
       # Build collocation matrix
       Mcoll = np.eye(nnodes) - step.status.dt*Q*(problem.lambda_s[0] + problem.lambda_f[0])
@@ -172,9 +158,9 @@ def test_collocationinvariant(self):
       # Perform node-to-node SDC sweep
       level.sweep.update_nodes()
 
-      # Build matrices for matrix formulation of sweep
-      LHS = np.eye(nnodes) - step.status.dt*( problem.lambda_f[0]*QI + problem.lambda_s[0]*QE )
-      RHS = step.status.dt*( (problem.lambda_f[0]+problem.lambda_s[0])*Q - (problem.lambda_f[0]*QI + problem.lambda_s[0]*QE) )
+      lambdas = [ problem.lambda_f[0] , problem.lambda_s[0] ]
+      LHS, RHS = level.sweep.get_scalar_problems_sweeper_mats( lambdas = lambdas )
+
       # Make sure both matrix and node-to-node sweep leave collocation unaltered
       unew = np.linalg.inv(LHS).dot( u0full + RHS.dot(ucoll) )
       assert np.linalg.norm( unew - ucoll, np.infty )<1e-14, "Collocation solution not invariant under matrix SDC sweep"
@@ -198,18 +184,16 @@ def test_manysweepsequalmatrix(self):
         level.sweep.update_nodes()
       usweep = np.array([ level.u[l].values.flatten() for l in range(1,nnodes+1) ])
 
-      LHS, RHS = self.setupSweeperMatrices(step, level, problem)
+      lambdas = [ problem.lambda_f[0] , problem.lambda_s[0] ]
+      LHS, RHS = level.sweep.get_scalar_problems_sweeper_mats( lambdas = lambdas )
+
       unew = u0full
       for i in range(0,K):
         unew = np.linalg.inv(LHS).dot( u0full + RHS.dot(unew) )
 
       assert np.linalg.norm(unew - usweep, np.infty)<1e-14, "Doing multiple node-to-node sweeps yields different result than same number of matrix-form sweeps"   
 
-      # Build single matrix representing K sweeps    
-      Pinv = np.linalg.inv(LHS)
-      Mat_sweep = np.linalg.matrix_power(Pinv.dot(RHS), K)
-      for i in range(0,K):
-        Mat_sweep = Mat_sweep + np.linalg.matrix_power(Pinv.dot(RHS),i).dot(Pinv)
+      Mat_sweep = level.sweep.get_scalar_problems_manysweep_mat( nsweeps = K, lambdas = lambdas )
       usweep_onematrix = Mat_sweep.dot(u0full)
       assert np.linalg.norm( usweep_onematrix - usweep, np.infty )<1e-14, "Single-matrix multiple sweep formulation yields different result than multiple sweeps in node-to-node or matrix form form"
 
@@ -231,12 +215,11 @@ def test_manysweepupdate(self):
       level.sweep.compute_end_point()
       uend_sweep = level.uend.values
 
-      LHS, RHS = self.setupSweeperMatrices(step, level, problem)
+      lambdas = [ problem.lambda_f[0] , problem.lambda_s[0] ]
+      LHS, RHS = level.sweep.get_scalar_problems_sweeper_mats( lambdas = lambdas )
+
       # Build single matrix representing K sweeps    
-      Pinv = np.linalg.inv(LHS)
-      Mat_sweep = np.linalg.matrix_power(Pinv.dot(RHS), K)
-      for i in range(0,K):
-        Mat_sweep = Mat_sweep + np.linalg.matrix_power(Pinv.dot(RHS),i).dot(Pinv)
+      Mat_sweep = level.sweep.get_scalar_problems_manysweep_mat( nsweeps = K, lambdas = lambdas )
       # Now build update function
       update = 1.0 + (problem.lambda_s[0] + problem.lambda_f[0])*level.sweep.coll.weights.dot(Mat_sweep.dot(np.ones(nnodes)))
       # Multiply u0 by value of update function to get end value directly