<commit_before>from pySDC import Level as lvl
from pySDC import Hooks as hookclass
from pySDC import CollocationClasses as collclass
from pySDC import Step as stepclass

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 examples.SWFW.ProblemClass import swfw_scalar 

from nose.tools import *
import unittest
import numpy as np

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 = {}
    self.pparams['lambda_s'] = np.array([-0.1*1j], dtype='complex')
    self.pparams['lambda_f'] = np.array([-1.0*1j], dtype='complex')
    self.pparams['u0'] = np.random.rand()
    self.swparams = {}
    self.swparams['collocation_class'] = collclass.CollGaussLobatto
    self.swparams['num_nodes'] = 2+np.random.randint(5)

  # ***************
  # **** 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):

    # 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, level, problem, nnodes = self.setupLevelStepProblem()
    step.levels[0].sweep.predict()
    u0full = np.array([ level.u[l].values.flatten() for l in range(1,nnodes+1) ])

    # Perform node-to-node SDC sweep
    level.sweep.update_nodes()

    LHS, RHS = self.setupSweeperMatrices(step, level, problem)

    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 issue #52 before passing")
  def test_updateformula(self):

    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
    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
  #<commit_msg>small fix in use of u0 -- skipped tests are passed for Gauss-Legendre nodes<commit_after>from pySDC import Level as lvl
from pySDC import Hooks as hookclass
from pySDC import CollocationClasses as collclass
from pySDC import Step as stepclass

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 examples.SWFW.ProblemClass import swfw_scalar 

from nose.tools import *
import unittest
import numpy as np

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 = {}
    self.pparams['lambda_s'] = np.array([-0.1*1j], dtype='complex')
    self.pparams['lambda_f'] = np.array([-1.0*1j], dtype='complex')
    self.pparams['u0'] = np.random.rand()
    self.swparams = {}
    self.swparams['collocation_class'] = collclass.CollGaussLobatto
    self.swparams['num_nodes'] = 2+np.random.randint(5)

  # ***************
  # **** 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):

    # 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, level, problem, nnodes = self.setupLevelStepProblem()
    step.levels[0].sweep.predict()
    u0full = np.array([ level.u[l].values.flatten() for l in range(1,nnodes+1) ])

    # Perform node-to-node SDC sweep
    level.sweep.update_nodes()

    LHS, RHS = self.setupSweeperMatrices(step, level, problem)

    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 issue #52 before passing")
  def test_updateformula(self):

    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
    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
  #
{'fraction_non_alphanumeric': 0.11421298957221447, 'fraction_numerical': 0.010099351342474751, 'mean_word_length': 4.0644490644490645, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4007500', 'n_tokens_mistral': 4316, 'n_tokens_neox': 4137, 'n_words': 1038}