beautification of imex_sweeper test

Author: Daniel Ruprecht

tests/test_imexsweeper.py

@@ -14,7 +14,38 @@
 class TestImexSweeper(unittest.TestCase):
 
   #
+  # Some auxiliary functions which are not tests themselves
   #
+  def setupLevelStepProblem(self):
+    step = stepclass.step(params={})
+    L = lvl.level(problem_class=swfw_scalar, problem_params=self.pparams, dtype_u=mesh, dtype_f=rhs_imex_mesh, sweeper_class=imex, sweeper_params=self.swparams, level_params={}, hook_class=hookclass.hooks, id="imextest")
+    step.register_level(L)
+    step.status.dt   = 1.0
+    step.status.time = 0.0
+    u0 = step.levels[0].prob.u_exact(step.status.time)
+    step.init_step(u0)
+    nnodes  = step.levels[0].sweep.coll.num_nodes
+    level   = step.levels[0]
+    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
   #
   def setUp(self):
     self.pparams = {}
@@ -25,160 +56,129 @@ def setUp(self):
     self.swparams['collocation_class'] = collclass.CollGaussLobatto
     self.swparams['num_nodes'] = 2
 
+  # ***************
+  # **** TESTS ****
+  # ***************
+
+
   #
-  #
+  # Check that a level object can be instantiated
   #
   def test_caninstantiate(self):
     L = lvl.level(problem_class=swfw_scalar, problem_params=self.pparams, dtype_u=mesh, dtype_f=rhs_imex_mesh, sweeper_class=imex, sweeper_params=self.swparams, level_params={}, hook_class=hookclass.hooks, id="imextest")
+    assert isinstance(L.sweep, imex), "sweeper in generated level is not an object of type imex"
 
   #
-  #
+  # Check that a level object can be registered in a step object (needed as prerequiste to execute update_nodes
   #
   def test_canregisterlevel(self):
     step = stepclass.step(params={})
     L = lvl.level(problem_class=swfw_scalar, problem_params=self.pparams, dtype_u=mesh, dtype_f=rhs_imex_mesh, sweeper_class=imex, sweeper_params=self.swparams, level_params={}, hook_class=hookclass.hooks, id="imextest")
     step.register_level(L)
+    # At this point, it should not be possible to actually execute functions of the sweeper because the parameters set in setupLevelStepProblem are not yet initialised
+    with self.assertRaises(Exception):
+      step.sweep.predict()
+    with self.assertRaises(Exception):
+      step.sweep.update_nodes()
+    with self.assertRaises(Exception):
+      step.sweep.compute_end_point()
 
   #
-  #
+  # Check that the sweeper functions update_nodes and compute_end_point can be executed
   #
   def test_canrunsweep(self):
 
-    step = stepclass.step(params={})
-    L = lvl.level(problem_class=swfw_scalar, problem_params=self.pparams, dtype_u=mesh, dtype_f=rhs_imex_mesh, sweeper_class=imex, sweeper_params=self.swparams, level_params={}, hook_class=hookclass.hooks, id="imextest")
-    step.register_level(L)
-    step.status.dt   = 1.0
-    step.status.time = 0.0
-    nnodes  = step.levels[0].sweep.coll.num_nodes
-    u0 = step.levels[0].prob.u_exact(step.status.time)
-    step.init_step(u0)
-    step.levels[0].sweep.predict()
-    u0full = np.array([ step.levels[0].u[l].values.flatten() for l in range(1,nnodes+1) ])
-
-    step.levels[0].sweep.update_nodes()
-    assert step.levels[0].uend is None, "uend should be None previous to running compute_end_point"
-    step.levels[0].sweep.compute_end_point()
-    #print "Sweep: %s" % step.levels[0].uend.values
+    # After running setupLevelStepProblem, the functions predict, update_nodes and compute_end_point should run
+    step, level, problem, nnodes = self.setupLevelStepProblem()   
+    assert level.u[0] is not None, "After init_step, level.u[0] should no longer be of type None" 
+    assert level.u[1] is None, "Before predict, level.u[1] and following should be of type None"
+    level.sweep.predict()
+    # Should now be able to run update nodes
+    level.sweep.update_nodes()
+    assert level.uend is None, "uend should be None previous to running compute_end_point"
+    level.sweep.compute_end_point()
+    assert level.uend is not None, "uend still None after running compute_end_point"
 
   #
   # Make sure a sweep in matrix form is equal to a sweep in node-to-node form
   #
   def test_sweepequalmatrix(self):
 
-    step = stepclass.step(params={})
-    L = lvl.level(problem_class=swfw_scalar, problem_params=self.pparams, dtype_u=mesh, dtype_f=rhs_imex_mesh, sweeper_class=imex, sweeper_params=self.swparams, level_params={}, hook_class=hookclass.hooks, id="imextest")
-    step.register_level(L)
-    step.status.dt   = 1.0
-    step.status.time = 0.0
-    nnodes  = step.levels[0].sweep.coll.num_nodes
-    u0 = step.levels[0].prob.u_exact(step.status.time)
-    step.init_step(u0)
+    step, level, problem, nnodes = self.setupLevelStepProblem()
     step.levels[0].sweep.predict()
-    u0full = np.array([ step.levels[0].u[l].values.flatten() for l in range(1,nnodes+1) ])
+    u0full = np.array([ level.u[l].values.flatten() for l in range(1,nnodes+1) ])
 
     # Perform node-to-node SDC sweep
-    step.levels[0].sweep.update_nodes()
+    level.sweep.update_nodes()
+
+    LHS, RHS = self.setupSweeperMatrices(step, level, problem)
 
-    # Build SDC sweep matrix
-    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:]
-    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) )
     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) ])
     assert np.linalg.norm(unew - usweep, np.infty)<1e-14, "Single SDC sweeps in matrix and node-to-node formulation yield different results"        
 
   #
-  #
+  # Make sure the implemented update formula matches the matrix update formula
   #
   @unittest.skip("Needs fix of isse #52 before passing")
   def test_updateformula(self):
 
-    step = stepclass.step(params={})
-    L = lvl.level(problem_class=swfw_scalar, problem_params=self.pparams, dtype_u=mesh, dtype_f=rhs_imex_mesh, sweeper_class=imex, sweeper_params=self.swparams, level_params={}, hook_class=hookclass.hooks, id="imextest")
-    step.register_level(L)
-    step.status.dt   = 1.0
-    step.status.time = 0.0
-    nnodes  = step.levels[0].sweep.coll.num_nodes
-    u0 = step.levels[0].prob.u_exact(step.status.time)
-    step.init_step(u0)
-    step.levels[0].sweep.predict()
-    u0full = np.array([ step.levels[0].u[l].values.flatten() for l in range(1,nnodes+1) ])
-
-    # Build SDC sweep matrix
-    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:]
+    step, level, problem, nnodes = self.setupLevelStepProblem()
+    level.sweep.predict()
+    u0full = np.array([ level.u[l].values.flatten() for l in range(1,nnodes+1) ])
 
     # Perform update step in sweeper
-    step.levels[0].sweep.update_nodes()
-    ustages = np.array([ step.levels[0].u[l].values.flatten() for l in range(1,nnodes+1) ])
-
-    step.levels[0].sweep.compute_end_point()
-    uend_sweep = step.levels[0].uend.values
-    uend_mat   = u0.values + step.status.dt*step.levels[0].sweep.coll.weights.dot(ustages*(problem.lambda_s[0] + problem.lambda_f[0]))
+    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   = u0.values + 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
   #
   def test_collocationinvariant(self):
 
-    step = stepclass.step(params={})
-    L = lvl.level(problem_class=swfw_scalar, problem_params=self.pparams, dtype_u=mesh, dtype_f=rhs_imex_mesh, sweeper_class=imex, sweeper_params=self.swparams, level_params={}, hook_class=hookclass.hooks, id="imextest")
-    step.register_level(L)
-    step.status.dt   = 1.0
-    step.status.time = 0.0
-    nnodes  = step.levels[0].sweep.coll.num_nodes
-    u0 = step.levels[0].prob.u_exact(step.status.time)
-    step.init_step(u0)
-    step.levels[0].sweep.predict()
-    u0full = np.array([ step.levels[0].u[l].values.flatten() for l in range(1,nnodes+1) ])
-
-    # Build SDC sweep matrix
-    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:]
+    step, level, problem, nnodes = self.setupLevelStepProblem()
+    level.sweep.predict()
+    u0full = np.array([ level.u[l].values.flatten() for l in range(1,nnodes+1) ])
 
+    QE, QI, Q = self.setupQMatrices(level)
+
+    # Build collocation matrix
     Mcoll = np.eye(nnodes) - step.status.dt*Q*(problem.lambda_s[0] + problem.lambda_f[0])
+
+    # Solve collocation problem directly
     ucoll = np.linalg.inv(Mcoll).dot(u0full)
+    
+    # Put stages of collocation solution into level
     for l in range(0,nnodes):
-      step.levels[0].u[l+1].values = ucoll[l]
-      step.levels[0].f[l+1].impl.values = problem.lambda_f[0]*ucoll[l]
-      step.levels[0].f[l+1].expl.values = problem.lambda_s[0]*ucoll[l]
+      level.u[l+1].values = ucoll[l]
+      level.f[l+1].impl.values = problem.lambda_f[0]*ucoll[l]
+      level.f[l+1].expl.values = problem.lambda_s[0]*ucoll[l]
 
     # Perform node-to-node SDC sweep
-    step.levels[0].sweep.update_nodes()
+    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) )
+    # 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"
-    unew_sweep = np.array([ step.levels[0].u[l].values.flatten() for l in range(1,nnodes+1) ])
+    unew_sweep = np.array([ level.u[l].values.flatten() for l in range(1,nnodes+1) ])
     assert np.linalg.norm( unew_sweep - ucoll, np.infty )<1e-14, "Collocation solution not invariant under node-to-node sweep"
 
   #
   #
   #
   def test_canrunmatrixsweep(self):
-    step = stepclass.step(params={})
-    L = lvl.level(problem_class=swfw_scalar, problem_params=self.pparams, dtype_u=mesh, dtype_f=rhs_imex_mesh, sweeper_class=imex, sweeper_params=self.swparams, level_params={}, hook_class=hookclass.hooks, id="imextest")
-    step.register_level(L)
-    step.status.dt   = 1.0
-    step.status.time = 0.0
-    u0 = step.levels[0].prob.u_exact(step.status.time)
-    step.init_step(u0)
-    nnodes  = step.levels[0].sweep.coll.num_nodes
-    level   = step.levels[0]
-    problem = level.prob
+    step, level, problem, nnodes = self.setupLevelStepProblem()
+
     QE = level.sweep.QE[1:,1:]
     QI = level.sweep.QI[1:,1:]
     Q  = level.sweep.coll.Qmat[1:,1:]
@@ -204,3 +204,5 @@ def test_canrunmatrixsweep(self):
     #print ufull
     #uend   = u0.values + step.status.dt*level.sweep.coll.weights.dot( (problem.lambda_f[0]+problem.lambda_s[0])*ufull )
     #print "Matrix: %s" % uend
+
+