Merge pull request #66 from danielru/add/rk_imex

pancetta · pancetta · commit d2bc3d65fcec · 2016-01-25T09:47:01.000+01:00
Add/rk imex
diff --git a/examples/acoustic_1d_imex/plot_dispersion.py b/examples/acoustic_1d_imex/plot_dispersion.py
@@ -5,7 +5,7 @@
 
 from pySDC.datatype_classes.complex_mesh import mesh, rhs_imex_mesh
 from pySDC.sweeper_classes.imex_1st_order import imex_1st_order as imex
-from standard_integrators import dirk
+from standard_integrators import dirk, rk_imex
 # for simplicity, import the scalar problem to generate Q matrices
 from examples.fwsw.ProblemClass import swfw_scalar 
 import numpy as np
@@ -41,8 +41,8 @@ def findomega(stab_fh):
     pparams['u0'] = 1.0
     swparams = {}
     swparams['collocation_class'] = collclass.CollGaussLegendre
-    swparams['num_nodes'] = 2
-    K = 2
+    swparams['num_nodes'] = 3
+    K = 4
     dirk_order = K
     
     c_speed = 1.0
@@ -68,8 +68,8 @@ def findomega(stab_fh):
     Nsamples = 30
     k_vec = np.linspace(0, np.pi, Nsamples+1, endpoint=False)
     k_vec = k_vec[1:]
-    phase = np.zeros((2,Nsamples))
-    amp_factor = np.zeros((2,Nsamples))
+    phase = np.zeros((3,Nsamples))
+    amp_factor = np.zeros((3,Nsamples))
     for i in range(0,np.size(k_vec)):
       Cs = -1j*k_vec[i]*np.array([[0.0, c_speed],[c_speed, 0.0]], dtype='complex')
       Uadv = -1j*k_vec[i]*np.array([[U_speed, 0.0], [0.0, U_speed]], dtype='complex')
@@ -85,8 +85,8 @@ def findomega(stab_fh):
       # ---> The update formula for this case need verification!!
       update = step.status.dt*np.kron( level.sweep.coll.weights, Uadv+Cs )
       
-      y1 = np.array([1,0])
-      y2 = np.array([0,1])
+      y1 = np.array([1,0], dtype='complex')
+      y2 = np.array([0,1], dtype='complex')
       e1 = np.kron( np.ones(nnodes), y1 )
       stab_fh_1      = y1 + update.dot( Mat_sweep.dot(e1) )
       e2 = np.kron( np.ones(nnodes), y2 )
@@ -104,21 +104,31 @@ def findomega(stab_fh):
       stab_fh2 = dirkts.timestep(y2, 1.0)
       stab_dirk = np.column_stack((stab_fh1, stab_fh2))
 
+      rkimex = rk_imex(M_fast = Cs, M_slow = Uadv, order = K)
+      stab_fh1 = rkimex.timestep(y1, 1.0)
+      stab_fh2 = rkimex.timestep(y2, 1.0)
+      stab_rk_imex = np.column_stack((stab_fh1, stab_fh2))
+
       sol_sdc = findomega(stab_sdc)
       sol_dirk = findomega(stab_dirk)
+      sol_rk_imex = findomega(stab_rk_imex)
       
       # Now solve for discrete phase 
       phase[0,i]      = sol_sdc.real/k_vec[i]
       amp_factor[0,i] = np.exp(sol_sdc.imag)
       phase[1,i]      = sol_dirk.real/k_vec[i]
       amp_factor[1,i] = np.exp(sol_dirk.imag)
+      phase[2,i]      = sol_rk_imex.real/k_vec[i]
+      amp_factor[2,i] = np.exp(sol_rk_imex.imag)
+
     ###
     rcParams['figure.figsize'] = 1.5, 1.5
     fs = 8
     fig  = plt.figure()
     plt.plot(k_vec, (U_speed+c_speed)+np.zeros(np.size(k_vec)), '--', color='k', linewidth=1.5, label='Exact')
-    plt.plot(k_vec, phase[0,:], '-', color='b', linewidth=1.5, label='SDC('+str(K)+')')
     plt.plot(k_vec, phase[1,:], '-', color='g', linewidth=1.5, label='DIRK('+str(dirkts.order)+')')
+    plt.plot(k_vec, phase[2,:], '-', color='r', linewidth=1.5, label='RK-IMEX('+str(rkimex.order)+')')
+    plt.plot(k_vec, phase[0,:], '-', color='b', linewidth=1.5, label='SDC('+str(K)+')')
     plt.xlabel('Wave number', fontsize=fs, labelpad=0.25)
     plt.ylabel('Phase speed', fontsize=fs, labelpad=0.5)
     plt.xlim([k_vec[0], k_vec[-1:]])
@@ -133,8 +143,9 @@ def findomega(stab_fh):
 
     fig  = plt.figure()
     plt.plot(k_vec, 1.0+np.zeros(np.size(k_vec)), '--', color='k', linewidth=1.5, label='Exact')
-    plt.plot(k_vec, amp_factor[0,:], '-', color='b', linewidth=1.5, label='SDC('+str(K)+')')
     plt.plot(k_vec, amp_factor[1,:], '-', color='g', linewidth=1.5, label='DIRK('+str(dirkts.order)+')')
+    plt.plot(k_vec, amp_factor[2,:], '-', color='r', linewidth=1.5, label='RK-IMEX('+str(rkimex.order)+')')
+    plt.plot(k_vec, amp_factor[0,:], '-', color='b', linewidth=1.5, label='SDC('+str(K)+')')
     plt.xlabel('Wave number', fontsize=fs, labelpad=0.25)
     plt.ylabel('Amplification factor', fontsize=fs, labelpad=0.5)
     fig.gca().tick_params(axis='both', labelsize=fs)
diff --git a/examples/acoustic_1d_imex/runmultiscale.py b/examples/acoustic_1d_imex/runmultiscale.py
@@ -12,7 +12,7 @@
 from pySDC import Log
 from pySDC.Stats import grep_stats, sort_stats
 
-from standard_integrators import bdf2, dirk, trapezoidal
+from standard_integrators import bdf2, dirk, trapezoidal, rk_imex
 
 from matplotlib import pyplot as plt
 from pylab import rcParams
@@ -32,7 +32,7 @@
     lparams['restol'] = 1E-10
 
     sparams = {}
-    sparams['maxiter'] = 4
+    sparams['maxiter'] = 2
 
     # setup parameters "in time"
     t0   = 0.0
@@ -59,7 +59,7 @@
     description['dtype_f']           = rhs_imex_mesh
     description['collocation_class'] = collclass.CollGaussLegendre
     # Number of nodes
-    description['num_nodes']         = 3
+    description['num_nodes']         = 2
     description['sweeper_class']     = imex_1st_order
     description['level_params']      = lparams
     description['hook_class']        = plot_solution
@@ -77,14 +77,16 @@
     uend,stats = mp.run_pfasst(MS,u0=uinit,t0=t0,dt=dt,Tend=Tend)
 
     # instantiate standard integrators to be run for comparison
-    trap  = trapezoidal( P.A+P.Dx, 0.5 )
+    trap  = trapezoidal( (P.A+P.Dx).astype('complex'), 0.5 )
     bdf2  = bdf2( P.A+P.Dx)
-    dirk  = dirk( P.A+P.Dx, sparams['maxiter'])
-
+    dirk  = dirk( (P.A+P.Dx).astype('complex'), sparams['maxiter'])
+    rkimex = rk_imex(P.A.astype('complex'), P.Dx.astype('complex'), sparams['maxiter'])
+    
     y0_tp = np.concatenate( (uinit.values[0,:], uinit.values[1,:]) )
     y0_bdf  = y0_tp
-    y0_dirk = y0_tp
-      
+    y0_dirk = y0_tp.astype('complex')
+    y0_imex = y0_tp.astype('complex')
+    
     # Perform 154 time steps with standard integrators
     for i in range(0,154):  
 
@@ -101,15 +103,20 @@
       # DIRK scheme
       ynew_dirk = dirk.timestep(y0_dirk, dt)
 
+      # IMEX scheme
+      ynew_imex = rkimex.timestep(y0_imex, dt)
+      
       y0_tp   = ynew_tp
       ym1_bdf = y0_bdf
       y0_bdf  = ynew_bdf
       y0_dirk = ynew_dirk
-
+      y0_imex = ynew_imex
+      
     # Finished running standard integrators
       unew_tp, pnew_tp     = np.split(ynew_tp, 2)
       unew_bdf, pnew_bdf   = np.split(ynew_bdf, 2)
       unew_dirk, pnew_dirk = np.split(ynew_dirk, 2)
+      unew_imex, pnew_imex = np.split(ynew_imex, 2)
 
     rcParams['figure.figsize'] = 5, 2.5
     fig = plt.figure()
@@ -119,8 +126,13 @@
     x_0     = 0.75
     x_1     = 0.25
 
-    plt.plot(P.mesh, pnew_tp,  '-', color='c', label='Trapezoidal')
-    plt.plot(P.mesh, uend.values[1,:], '-', color='b', label='SDC('+str(sparams['maxiter'])+')')
+    print ('Maximum pressure in SDC: %5.3e' % np.linalg.norm(uend.values[1,:], np.inf))
+    print ('Maximum pressure in DIRK: %5.3e' % np.linalg.norm(pnew_dirk, np.inf))
+    print ('Maximum pressure in RK-IMEX: %5.3e' % np.linalg.norm(pnew_imex, np.inf))
+
+    #plt.plot(P.mesh, pnew_tp,  '-', color='c', label='Trapezoidal')
+    plt.plot(P.mesh, pnew_imex,  '-', color='c', label='IMEX('+str(rkimex.order)+')')
+    plt.plot(P.mesh, uend.values[1,:], '--', color='b', label='SDC('+str(sparams['maxiter'])+')')
     plt.plot(P.mesh, pnew_bdf, '-', color='r', label='BDF-2')
     plt.plot(P.mesh, pnew_dirk, color='g', label='DIRK('+str(dirk.order)+')')
     #plt.plot(P.mesh, uex.values[1,:],  '+', color='r', label='p (exact)')
diff --git a/examples/acoustic_1d_imex/standard_integrators.py b/examples/acoustic_1d_imex/standard_integrators.py
@@ -3,6 +3,101 @@
 import scipy.sparse.linalg as LA
 import scipy.sparse as sp
 
+#
+# Runge-Kutta IMEX methods of order 1 to 3
+#
+class rk_imex:
+  
+  def __init__(self, M_fast, M_slow, order):
+    assert np.shape(M_fast)[0]==np.shape(M_fast)[1], "A_fast must be square"
+    assert np.shape(M_slow)[0]==np.shape(M_slow)[1], "A_slow must be square"
+    assert np.shape(M_fast)[0]==np.shape(M_slow)[0], "A_fast and A_slow must be of the same size"
+
+    assert order in [1,2,3,4], "Order must be 1, 2, 3 or 4"
+    self.order = order
+
+    if self.order==1:
+      self.A      = np.array([[0,0],[0,1]])
+      self.A_hat  = np.array([[0,0],[1,0]])
+      self.b      = np.array([0,1])
+      self.b_hat  = np.array([1,0])
+      self.nstages = 2
+
+    elif self.order==2:
+      self.A      = np.array([[0,0],[0,0.5]])
+      self.A_hat  = np.array([[0,0],[0.5,0]])
+      self.b      = np.array([0,1])
+      self.b_hat  = np.array([0,1])
+      self.nstages = 2
+
+    elif self.order==3:
+      # parameter from Pareschi and Russo, J. Sci. Comp. 2005
+      alpha = 0.24169426078821
+      beta  = 0.06042356519705
+      eta   = 0.12915286960590
+      self.A_hat   = np.array([ [0,0,0,0], [0,0,0,0], [0,1.0,0,0], [0, 1.0/4.0, 1.0/4.0, 0] ])
+      self.A       = np.array([ [alpha, 0, 0, 0], [-alpha, alpha, 0, 0], [0, 1.0-alpha, alpha, 0], [beta, eta, 0.5-beta-eta-alpha, alpha] ])
+      self.b_hat   = np.array([0, 1.0/6.0, 1.0/6.0, 2.0/3.0])
+      self.b       = self.b_hat
+      self.nstages = 4
+
+    elif self.order==4:
+
+      self.A_hat = np.array([ [0,0,0,0,0,0],
+                    [1./2,0,0,0,0,0],
+                    [13861./62500.,6889./62500.,0,0,0,0],
+                    [-116923316275./2393684061468.,-2731218467317./15368042101831.,9408046702089./11113171139209.,0,0,0],
+                    [-451086348788./2902428689909.,-2682348792572./7519795681897.,12662868775082./11960479115383.,3355817975965./11060851509271.,0,0],
+                 [647845179188./3216320057751.,73281519250./8382639484533.,552539513391./3454668386233.,3354512671639./8306763924573.,4040./17871.,0]])
+      self.A = np.array([[0,0,0,0,0,0],
+                  [1./4,1./4,0,0,0,0],
+                  [8611./62500.,-1743./31250.,1./4,0,0,0],
+                  [5012029./34652500.,-654441./2922500.,174375./388108.,1./4,0,0],
+                  [15267082809./155376265600.,-71443401./120774400.,730878875./902184768.,2285395./8070912.,1./4,0],
+                  [82889./524892.,0,15625./83664.,69875./102672.,-2260./8211,1./4]])
+      self.b = np.array([82889./524892.,0,15625./83664.,69875./102672.,-2260./8211,1./4])
+      self.b_hat = np.array([4586570599./29645900160.,0,178811875./945068544.,814220225./1159782912.,-3700637./11593932.,61727./225920.])
+      self.nstages = 6
+    
+    self.M_fast = sp.csc_matrix(M_fast)
+    self.M_slow = sp.csc_matrix(M_slow)
+    self.ndof   = np.shape(M_fast)[0]
+
+    self.stages = np.zeros((self.nstages, self.ndof), dtype='complex')
+
+  def timestep(self, u0, dt):
+
+    # Solve for stages
+    for i in range(0,self.nstages):
+
+      # Construct RHS
+      rhs = np.copy(u0)
+      for j in range(0,i):
+        rhs += dt*self.A_hat[i,j]*(self.f_slow(self.stages[j,:])) + dt*self.A[i,j]*(self.f_fast(self.stages[j,:]))
+
+      # Solve for stage i
+      if self.A[i,i] == 0:
+        # Avoid call to spsolve with identity matrix
+        self.stages[i,:] = np.copy(rhs)
+      else:
+        self.stages[i,:] = self.f_fast_solve( rhs, dt*self.A[i,i] )
+    
+    # Update 
+    for i in range(0,self.nstages):
+      u0 += dt*self.b_hat[i]*(self.f_slow(self.stages[i,:])) + dt*self.b[i]*(self.f_fast(self.stages[i,:]))
+
+    return u0
+
+  def f_slow(self, u):
+    return self.M_slow.dot(u)
+
+  def f_fast(self, u):
+    return self.M_fast.dot(u)
+
+  def f_fast_solve(self, rhs, alpha):
+    L = sp.eye(self.ndof) - alpha*self.M_fast
+    return LA.spsolve(L, rhs)
+
 #
 # Trapezoidal rule
 #
diff --git a/examples/boussinesq_2d_imex/plotgmrescounter.py b/examples/boussinesq_2d_imex/plotgmrescounter.py
@@ -9,21 +9,23 @@
 from unflatten import unflatten
 
 if __name__ == "__main__":
-  xx = np.load('xaxis.npy')
-  uend = np.load('sdc.npy')
+  xx    = np.load('xaxis.npy')
+  uend  = np.load('sdc.npy')
   udirk = np.load('dirk.npy')
+  uimex = np.load('rkimex.npy')
   uref  = np.load('uref.npy')
 
   print "Estimated discretisation error of DIRK: %5.3e" % ( np.linalg.norm(udirk.flatten() - uref.flatten(), np.inf)/np.linalg.norm(uref.flatten(),np.inf) )
   print "Estimated discretisation error of SDC:  %5.3e" % ( np.linalg.norm(uend.flatten() - uref.flatten(), np.inf)/np.linalg.norm(uref.flatten(),np.inf) )
+  print "Estimated discretisation error of RK-IMEX:  %5.3e" % ( np.linalg.norm(uimex.flatten() - uref.flatten(), np.inf)/np.linalg.norm(uref.flatten(),np.inf) )
 
   fs = 8
   rcParams['figure.figsize'] = 5.0, 2.5
   fig = plt.figure()
 
   plt.plot(xx[:,5], udirk[2,:,5], '--', color='r', markersize=fs-2, label='DIRK', dashes=(3,3))
   plt.plot(xx[:,5], uend[2,:,5], '-', color='b', label='SDC')
-  #plt.plot(xx[:,5], udirk2[2,:,5], '--', color='r', markersize=fs-2, label='DIRK(2)', dashes=(3,3))
+  plt.plot(xx[:,5], uimex[2,:,5], '--', color='g', markersize=fs-2, label='RK-IMEX', dashes=(3,3))
   #plt.plot(xx[:,5], utrap[2,:,5], '--', color='k', markersize=fs-2, label='Trap', dashes=(3,3))
   plt.legend(loc='lower left', fontsize=fs, prop={'size':fs})
   plt.yticks(fontsize=fs)
diff --git a/examples/boussinesq_2d_imex/rungmrescounter.py b/examples/boussinesq_2d_imex/rungmrescounter.py
@@ -21,7 +21,7 @@
 
 from unflatten import unflatten
 
-from standard_integrators import dirk, bdf2, trapezoidal
+from standard_integrators import dirk, bdf2, trapezoidal, rk_imex
 
 if __name__ == "__main__":
 
@@ -47,13 +47,13 @@
     # setup parameters "in time"
     t0     = 0
     Tend   = 3000
-    Nsteps =  100
+    Nsteps =  500
     dt = Tend/float(Nsteps)
 
     # This comes as read-in for the problem class
     pparams = {}
-    pparams['nvars']    = [(4,300,20)]
-    #pparams['nvars']    = [(4,150,10)]
+    pparams['nvars']    = [(4,450,30)]
+    #pparams['nvars']    = [(4,300,30)]
     pparams['u_adv']    = 0.02
     pparams['c_s']      = 0.3
     pparams['Nfreq']    = 0.01
@@ -63,7 +63,7 @@
     pparams['order_upw'] = [5]
     pparams['gmres_maxiter'] = [500]
     pparams['gmres_restart'] = [10]
-    pparams['gmres_tol']     = [1e-6]
+    pparams['gmres_tol']     = [1e-9]
 
     # This comes as read-in for the transfer operations
     tparams = {}
@@ -102,27 +102,43 @@
     for i in range(0,Nsteps):
       udirk = dirkp.timestep(udirk, dt)  
 
+    Pref = P
+    # For reference solution, increase GMRES tolerance
+    Pref.gmes_tol = 1e-6
     dirkref = dirk(P, 4)
     uref    = u0
     dt_ref  = dt/10.0
     for i in range(0,10*Nsteps):
       uref = dirkref.timestep(uref, dt_ref)
+
+    rkimex = rk_imex(P, dirk_order)
+    uimex  = u0
+    dt_imex = dt
+    for i in range(0,Nsteps):
+      uimex = rkimex.timestep(uimex, dt_imex)
   
     # call main function to get things done...
     uend,stats = mp.run_pfasst(MS,u0=uinit,t0=t0,dt=dt,Tend=Tend)
     udirk = unflatten(udirk, 4, P.N[0], P.N[1])
+    uimex = unflatten(uimex, 4, P.N[0], P.N[1])
     uref  = unflatten(uref,  4, P.N[0], P.N[1])
 
     np.save('xaxis', P.xx)
     np.save('sdc', uend.values)
     np.save('dirk', udirk)
+    np.save('rkimex', uimex)
     np.save('uref', uref)
     
-    print " #### Logging report for DIRK-%1i #### " % dirk_order
+    print " #### Logging report for DIRK-%1i #### " % dirkp.order
     print "Number of calls to implicit solver: %5i" % dirkp.logger.solver_calls
     print "Total number of GMRES iterations: %5i" % dirkp.logger.iterations
     print "Average number of iterations per call: %6.3f" % (float(dirkp.logger.iterations)/float(dirkp.logger.solver_calls))
-
+    print " "
+    print " #### Logging report for RK-IMEX-%1i #### " % rkimex.order
+    print "Number of calls to implicit solver: %5i" % rkimex.logger.solver_calls
+    print "Total number of GMRES iterations: %5i" % rkimex.logger.iterations
+    print "Average number of iterations per call: %6.3f" % (float(rkimex.logger.iterations)/float(rkimex.logger.solver_calls))
+    print " "
     print " #### Logging report for SDC-(%1i,%1i) #### " % (swparams['num_nodes'], sparams['maxiter'])
     print "Number of calls to implicit solver: %5i" % P.logger.solver_calls
     print "Total number of GMRES iterations: %5i" % P.logger.iterations
diff --git a/examples/boussinesq_2d_imex/standard_integrators.py b/examples/boussinesq_2d_imex/standard_integrators.py
