text: '''
  Created on 21.03.2016
  
  @author: Michael Kraus (michael.kraus@ipp.mpg.de)
  '''
  
  import sys, petsc4py
  from importlib import import_module
  
  from petsc4py import PETSc
  
  import h5py
  import numpy as np
  
  import argparse, datetime, time
  import pstats, cProfile
  
  from rmhd.config.config import Config
  
  from rmhd.solvers.common.PETScDerivatives import PETScDerivatives
  from rmhd.solvers.linear.PETScPoissonCFD2 import PETScPoisson
  
  
  class rmhd2d(object):
      '''
      PETSc/Python Reduced MHD Solver in 2D.
      '''
          
      def __init__(self, cfgfile, mode = "none"):
          '''
          Constructor
          '''
          
          petsc4py.init(sys.argv)
  
          if PETSc.COMM_WORLD.getRank() == 0:
              print("")
              print("Reduced MHD 2D")
              print("==============")
              print("")
  
          # solver mode
          self.mode = mode
          
          # set run id to timestamp
          self.run_id = datetime.datetime.fromtimestamp(time.time()).strftime("%y%m%d%H%M%S")
          
          if PETSc.COMM_WORLD.getRank() == 0:
              print("  Config: %s" % cfgfile)
          
          # load run config file
          self.cfg = Config(cfgfile)
          
          # timestep setup
          self.ht    = self.cfg['grid']['ht']              # timestep size
          self.nt    = self.cfg['grid']['nt']              # number of timesteps
          self.nsave = self.cfg['io']['nsave']             # save only every nsave'th timestep
          
          # grid setup
          self.nx = self.cfg['grid']['nx']                 # number of points in x
          self.ny = self.cfg['grid']['ny']                 # number of points in y
          
          self.Lx = self.cfg['grid']['Lx']                 # spatial domain in x
          x1      = self.cfg['grid']['x1']                 # 
          x2      = self.cfg['grid']['x2']                 # 
          
          self.Ly = self.cfg['grid']['Ly']                 # spatial domain in y
          y1      = self.cfg['grid']['y1']                 # 
          y2      = self.cfg['grid']['y2']                 # 
          
          self.hx = self.cfg['grid']['hx']                 # gridstep size in x
          self.hy = self.cfg['grid']['hy']                 # gridstep size in y
          
          # create time vector
          self.time = PETSc.Vec().createMPI(1, comm=PETSc.COMM_WORLD)
          self.time.setName('t')
          
          # electron skin depth
          self.de = self.cfg['initial_data']['skin_depth']
          
          # double bracket dissipation
          self.nu = self.cfg['initial_data']['dissipation']
          
          # set global tolerance
          self.tolerance = self.cfg['solver']['petsc_snes_atol'] * self.nx * self.ny
          
          # direct solver package
          self.solver_package = self.cfg['solver']['lu_solver_package']
          
          # set some PETSc solver options
          OptDB = PETSc.Options()
          
          OptDB.setValue('ksp_rtol',   self.cfg['solver']['petsc_ksp_rtol'])
          OptDB.setValue('ksp_atol',   self.cfg['solver']['petsc_ksp_atol'])
          OptDB.setValue('ksp_max_it', self.cfg['solver']['petsc_ksp_max_iter'])
          
          OptDB.setValue('pc_type', 'hypre')
          OptDB.setValue('pc_hypre_type', 'boomeramg')
          
          
          # create DA with single dof
          self.da1 = PETSc.DA().create(dim=2, dof=1,
                                      sizes=[self.nx, self.ny],
                                      proc_sizes=[PETSc.DECIDE, PETSc.DECIDE],
                                      boundary_type=('periodic', 'periodic'),
                                      stencil_width=1,
                                      stencil_type='box')
          
          # create DA (dof = 4 for A, J, P, O)
          self.da4 = PETSc.DA().create(dim=2, dof=4,
                                       sizes=[self.nx, self.ny],
                                       proc_sizes=[PETSc.DECIDE, PETSc.DECIDE],
                                       boundary_type=('periodic', 'periodic'),
                                       stencil_width=1,
                                       stencil_type='box')
          
          # create DA for x grid
          self.dax = PETSc.DA().create(dim=1, dof=1,
                                      sizes=[self.nx],
                                      proc_sizes=[PETSc.DECIDE],
                                      boundary_type=('periodic'))
          
          # create DA for y grid
          self.day = PETSc.DA().create(dim=1, dof=1,
                                      sizes=[self.ny],
                                      proc_sizes=[PETSc.DECIDE],
                                      boundary_type=('periodic'))
          
          
          # initialise grid
          self.da1.setUniformCoordinates(xmin=x1, xmax=x2,
                                         ymin=y1, ymax=y2)
          
          self.da4.setUniformCoordinates(xmin=x1, xmax=x2,
                                         ymin=y1, ymax=y2)
          
          self.dax.setUniformCoordinates(xmin=x1, xmax=x2)
          
          self.day.setUniformCoordinates(xmin=y1, xmax=y2)
          
          
          # create solution and RHS vector
          self.dx = self.da4.createGlobalVec()
          self.dy = self.da4.createGlobalVec()
          self.x  = self.da4.createGlobalVec()
          self.b  = self.da4.createGlobalVec()
          self.f  = self.da4.createGlobalVec()
          self.Pb = self.da1.createGlobalVec()
  
          self.FA = self.da1.createGlobalVec()
          self.FJ = self.da1.createGlobalVec()
          self.FP = self.da1.createGlobalVec()
          self.FO = self.da1.createGlobalVec()
          
          # create initial guess vectors
          self.igFA1= self.da1.createGlobalVec()
          self.igFA2= self.da1.createGlobalVec()
          self.igFO1= self.da1.createGlobalVec()
          self.igFO2= self.da1.createGlobalVec()
          
          #  nullspace vectors
          self.x0 = self.da4.createGlobalVec()
          self.P0 = self.da1.createGlobalVec()
          
          # create vectors for magnetic and velocity field
          self.A  = self.da1.createGlobalVec()        # magnetic vector potential A
          self.J  = self.da1.createGlobalVec()        # current density           J
          self.P  = self.da1.createGlobalVec()        # streaming function        psi
          self.O  = self.da1.createGlobalVec()        # vorticity                 omega
  
          self.Bx = self.da1.createGlobalVec()
          self.By = self.da1.createGlobalVec()
          self.Vx = self.da1.createGlobalVec()
          self.Vy = self.da1.createGlobalVec()
          
          # set variable names
          self.A.setName('A')
          self.J.setName('J')
          self.P.setName('P')
          self.O.setName('O')
          
          self.Bx.setName('Bx')
          self.By.setName('By')
          self.Vx.setName('Vx')
          self.Vy.setName('Vy')
          
          
          # initialise nullspace
          self.x0.set(0.)
          x0_arr = self.da4.getVecArray(self.x0)[...]
          
          x0_arr[:,:,2] = 1.
          self.x0.assemble()
          self.x0.normalize()
          
          self.solver_nullspace  = PETSc.NullSpace().create(constant=False, vectors=(self.x0,))
          self.poisson_nullspace = PETSc.NullSpace().create(constant=True)
          
          # initialise Poisson matrix
          self.Pm = self.da1.createMat()
          self.Pm.setOption(PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR, False)
          self.Pm.setUp()
          self.Pm.setNullSpace(self.poisson_nullspace)
          
          # create Poisson solver object
          self.petsc_poisson  = PETScPoisson(self.da1, self.nx, self.ny, self.hx, self.hy)
          
          # setup linear Poisson solver
          self.poisson_ksp = PETSc.KSP().create()
          self.poisson_ksp.setFromOptions()
          self.poisson_ksp.setOperators(self.Pm)
          self.poisson_ksp.setTolerances(rtol=self.cfg['solver']['poisson_ksp_rtol'],
                                         atol=self.cfg['solver']['poisson_ksp_atol'],
                                         max_it=self.cfg['solver']['poisson_ksp_max_iter'])
          self.poisson_ksp.setType('cg')
          self.poisson_ksp.getPC().setType('hypre')
          self.poisson_ksp.setUp()
          
          self.petsc_poisson.formMat(self.Pm)
          
          # create derivatives object
          self.derivatives = PETScDerivatives(self.da1, self.nx, self.ny, self.ht, self.hx, self.hy)
          
          # read initial data
          if self.cfg["io"]["hdf5_input"] != None and self.cfg["io"]["hdf5_input"] != "":
              if self.cfg["initial_data"]["python"] != None and self.cfg["initial_data"]["python"] != "":
                  if PETSc.COMM_WORLD.getRank() == 0:
                      print("WARNING: Both io.hdf5_input and initial_data.python are set!")
                      print("         Reading initial data from HDF5 file.")
              
              self.read_initial_data_from_hdf5()
          else:
              self.read_initial_data_from_python()
          
          # copy initial data vectors to x
          self.copy_x_from_da1_to_da4()
          
          # create HDF5 output file and write parameters
          hdf5_filename = self.cfg['io']['hdf5_output']
          last_dot      = hdf5_filename.rfind('.')
          hdf5_filename = hdf5_filename[:last_dot] + "." + str(self.run_id) + hdf5_filename[last_dot:]
          
          if PETSc.COMM_WORLD.getRank() == 0:
              print("  Output: %s" % hdf5_filename)
          
          hdf5out = h5py.File(hdf5_filename, "w", driver="mpio", comm=PETSc.COMM_WORLD.tompi4py())
          
          hdf5out.attrs["run_id"] = self.run_id
          
          for cfg_group in self.cfg:
              for cfg_item in self.cfg[cfg_group]:
                  if self.cfg[cfg_group][cfg_item] != None:
                      value = self.cfg[cfg_group][cfg_item]
                  else:
                      value = ""
                      
                  hdf5out.attrs[cfg_group + "." + cfg_item] = value
          
          hdf5out.attrs["solver.solver_mode"] = self.mode
          
          if self.cfg["initial_data"]["python"] != None and self.cfg["initial_data"]["python"] != "":
              python_input = open("examples/" + self.cfg['initial_data']['python'] + ".py", 'r')
              python_file = python_input.read()
              python_input.close()
          else:
              python_file = ""
              
          hdf5out.attrs["initial_data.python_file"] = python_file
          hdf5out.close()
          
          # create HDF5 viewer
          self.hdf5_viewer = PETSc.ViewerHDF5().create(hdf5_filename,
                                            mode=PETSc.Viewer.Mode.APPEND,
                                            comm=PETSc.COMM_WORLD)
          
          self.hdf5_viewer.pushGroup("/")
          
          
          # write grid data to hdf5 file
          coords_x = self.dax.getCoordinates()
          coords_y = self.day.getCoordinates()
          
          coords_x.setName('x')
          coords_y.setName('y')
  
          self.hdf5_viewer(coords_x)
          self.hdf5_viewer(coords_y)
          
          
          # write initial data to hdf5 file
          self.save_to_hdf5(0)
          
          # output some more information
          if PETSc.COMM_WORLD.getRank() == 0:
              print("")
              print("  nt = %i" % (self.nt))
              print("  nx = %i" % (self.nx))
              print("  ny = %i" % (self.ny))
              print("")
              print("  ht = %f" % (self.ht))
              print("  hx = %f" % (self.hx))
              print("  hy = %f" % (self.hy))
              print("")
              print("  PETSc SNES rtol = %e" % self.cfg['solver']['petsc_snes_rtol'])
              print("             atol = %e" % self.cfg['solver']['petsc_snes_atol'])
              print("             stol = %e" % self.cfg['solver']['petsc_snes_stol'])
              print("         max iter = %i" % self.cfg['solver']['petsc_snes_max_iter'])
              print("")
              print("  PETSc KSP  rtol = %e" % self.cfg['solver']['petsc_ksp_rtol'])
              print("             atol = %e" % self.cfg['solver']['petsc_ksp_atol'])
              print("         max iter = %i" % self.cfg['solver']['petsc_ksp_max_iter'])
              print("")
  
  
      def __del__(self):
          self.poisson_ksp.destroy()
          self.Pm.destroy()
          self.hdf5_viewer.destroy()
  
  
      def run(self):
          raise NotImplementedError
  
      
      def read_initial_data_from_python(self):
          python_module = "examples." + self.cfg['initial_data']['python']
          
          if PETSc.COMM_WORLD.getRank() == 0:
              print("  Input:  %s" % python_module)
              print("")
          
          # get whole grid
          xGrid, yGrid = self.get_coordinate_vectors_from_das()
          
          # set initial data
          (xs, xe), (ys, ye) = self.da1.getRanges()
          
          init_data = import_module(python_module)
          
          if hasattr(init_data, 'magnetic_A'):
              if PETSc.COMM_WORLD.getRank() == 0:
                  print("  Computing magnetic potential from initial data.")
              
              A_arr = self.da1.getVecArray(self.A)
              
              for i in range(xs, xe):
                  for j in range(ys, ye):
                      A_arr[i,j] = init_data.magnetic_A(xGrid[i], yGrid[j], self.Lx, self.Ly)
          
              # apply Fourier filtering to magnetic potential 
              self.fourier_filter_magnetic_potential()
              
              # compute current density
              self.derivatives.laplace_vec(self.A, self.J, -1.)
          
          elif hasattr(init_data, 'magnetic_J'):
              if PETSc.COMM_WORLD.getRank() == 0:
                  print("  Computing current density from initial data.")
              
              J_arr = self.da1.getVecArray(self.J)
              
              for i in range(xs, xe):
                  for j in range(ys, ye):
                      J_arr[i,j] = init_data.magnetic_J(xGrid[i], yGrid[j], self.Lx, self.Ly)
              
              # solve for consistent initial magnetic potential
              self.A.set(0.)
              self.petsc_poisson.formRHS(self.J, self.Pb)
              self.poisson_nullspace.remove(self.Pb)
              self.poisson_ksp.solve(self.Pb, self.A)
          
              self.derivatives.laplace_vec(self.A, self.J, -1.)
              
          else:
              if PETSc.COMM_WORLD.getRank() == 0:
                  print("  WARNING: Neither magnetic potential nor current density given in initial data.")
          
          
          if hasattr(init_data, 'velocity_P'):
              if PETSc.COMM_WORLD.getRank() == 0:
                  print("  Computing streaming function from initial data.")
              
              P_arr = self.da1.getVecArray(self.P)
              
              for i in range(xs, xe):
                  for j in range(ys, ye):
                      P_arr[i,j] = init_data.velocity_P(xGrid[i], yGrid[j], self.Lx, self.Ly)
              
              # remmove nullspace
              self.poisson_nullspace.remove(self.P)
              
              # compute vorticity
              self.derivatives.laplace_vec(self.P, self.O, -1.)
          
          elif hasattr(init_data, 'velocity_O'):
              if PETSc.COMM_WORLD.getRank() == 0:
                  print("  Computing vorticity from initial data.")
              
              O_arr = self.da1.getVecArray(self.O)
              
              for i in range(xs, xe):
                  for j in range(ys, ye):
                      O_arr[i,j] = init_data.velocity_O(xGrid[i], yGrid[j], self.Lx, self.Ly)
              
              self.poisson_nullspace.remove(self.O)
              
              # solve for consistent initial streaming function
              self.P.set(0.)
              self.petsc_poisson.formRHS(self.O, self.Pb)
              self.poisson_nullspace.remove(self.Pb)
              self.poisson_ksp.solve(self.Pb, self.P)
          
          else:
              if PETSc.COMM_WORLD.getRank() == 0:
                  print("  WARNING: Neither streaming function nor vorticity given in initial data.")
          
          if PETSc.COMM_WORLD.getRank() == 0:
              print("")
          
      
      def read_initial_data_from_hdf5(self):
          hdf5_filename = self.cfg["io"]["hdf5_input"]
          
          if PETSc.COMM_WORLD.getRank() == 0:
              print("  Input:  %s" % hdf5_filename)
              
          hdf5in = h5py.File(hdf5_filename, "r", driver="mpio", comm=PETSc.COMM_WORLD.tompi4py())
          
          assert self.nx == hdf5in.attrs["grid.nx"]
          assert self.ny == hdf5in.attrs["grid.ny"]
          assert self.hx == hdf5in.attrs["grid.hx"]
          assert self.hy == hdf5in.attrs["grid.hy"]
          assert self.Lx == hdf5in.attrs["grid.Lx"]
          assert self.Ly == hdf5in.attrs["grid.Ly"]
          
          assert self.de == hdf5in.attrs["initial_data.skin_depth"]
          
          timestep = len(hdf5in["t"][...].flatten()) - 1
          
          hdf5in.close()
          
          hdf5_viewer = PETSc.ViewerHDF5().create(hdf5_filename,
                                            mode=PETSc.Viewer.Mode.READ,
                                            comm=PETSc.COMM_WORLD)
          
          hdf5_viewer.setTimestep(timestep)
          
          self.A.load(hdf5_viewer)
          self.J.load(hdf5_viewer)
          self.P.load(hdf5_viewer)
          self.O.load(hdf5_viewer)
          
          hdf5_viewer.destroy()
      
      
      def copy_x_from_da4_to_da1(self):
          x_arr = self.da4.getVecArray(self.x)
          self.da1.getVecArray(self.A)[:,:] = x_arr[:,:,0]
          self.da1.getVecArray(self.J)[:,:] = x_arr[:,:,1]
          self.da1.getVecArray(self.P)[:,:] = x_arr[:,:,2]
          self.da1.getVecArray(self.O)[:,:] = x_arr[:,:,3]
  
  
      def copy_x_from_da1_to_da4(self):
          x_arr = self.da4.getVecArray(self.x)
          x_arr[:,:,0] = self.da1.getVecArray(self.A)[:,:]
          x_arr[:,:,1] = self.da1.getVecArray(self.J)[:,:]
          x_arr[:,:,2] = self.da1.getVecArray(self.P)[:,:]
          x_arr[:,:,3] = self.da1.getVecArray(self.O)[:,:]
      
      
      def get_coordinate_vectors_from_das(self):
          # get coordinate vectors
          coords_x = self.dax.getCoordinates()
          coords_y = self.day.getCoordinates()
           
          # save x coordinate arrays
          scatter, xVec = PETSc.Scatter.toAll(coords_x)
          
          scatter.begin(coords_x, xVec, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)
          scatter.end  (coords_x, xVec, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)
            
          xGrid = xVec.getValues(range(self.nx)).copy()
            
          scatter.destroy()
          xVec.destroy()
            
          # save y coordinate arrays
          scatter, yVec = PETSc.Scatter.toAll(coords_y)
          
          scatter.begin(coords_y, yVec, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)
          scatter.end  (coords_y, yVec, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)
                      
          yGrid = yVec.getValues(range(self.ny)).copy()
            
          scatter.destroy()
          yVec.destroy()
          
          return xGrid, yGrid
      
      
      def fourier_filter_magnetic_potential(self):
          # Fourier Filtering
          self.nfourier = self.cfg['initial_data']['nfourier']
            
          if self.nfourier >= 0:
              (xs, xe), (ys, ye) = self.da1.getRanges()
              
              # obtain whole A vector everywhere
              scatter, Aglobal = PETSc.Scatter.toAll(self.A)
              
              scatter.begin(self.A, Aglobal, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)
              scatter.end  (self.A, Aglobal, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)
              
              petsc_indices = self.da1.getAO().app2petsc(np.arange(self.nx*self.ny, dtype=np.int32))
              
              Ainit = Aglobal.getValues(petsc_indices).copy().reshape((self.ny, self.nx))
              
              scatter.destroy()
              Aglobal.destroy()
              
              # compute FFT, cut, compute inverse FFT
              from scipy.fftpack import rfft, irfft
              
              Afft = rfft(Ainit, axis=1)
              
  #             Afft[:,0] = 0.
              Afft[:,self.nfourier+1:] = 0.
              
              self.da1.getVecArray(self.A)[:,:] = irfft(Afft).T[xs:xe, ys:ye] 
      
      
      def calculate_initial_guess(self, initial=False):
          
          # set some variables for hermite extrapolation
          t0 = 0.
          t1 = 1.
          t  = 2.
          
          a0 = 2./(t0-t1)
          a1 = 2./(t1-t0)
          
          b0 = 1./(t0-t1)**2
          b1 = 1./(t1-t0)**2
          
          d0 = 1./(t-t0)
          d1 = 1./(t-t1)
          
          e0 = d0*b0
          e1 = d1*b1
          
          self.hermite_x0 = e0*(d0-a0)
          self.hermite_x1 = e1*(d1-a1)
          
          self.hermite_f0 = e0*self.ht
          self.hermite_f1 = e1*self.ht
          
          self.hermite_den = 1. / (self.hermite_x0 + self.hermite_x1)
          
          
          self.petsc_solver.Ap.copy(self.A)
          self.petsc_solver.Op.copy(self.O)
          
          if initial:
              self.derivatives.arakawa_vec(self.petsc_solver.Ap, self.petsc_solver.Pp, self.igFA1)
              self.derivatives.arakawa_vec(self.petsc_solver.Op, self.petsc_solver.Pp, self.igFO1)
              self.derivatives.arakawa_vec(self.petsc_solver.Ap, self.petsc_solver.Jp, self.igFO2)
              
              self.A.axpy(0.5*self.ht, self.igFA1)
              self.O.axpy(0.5*self.ht, self.igFO1)
              self.O.axpy(0.5*self.ht, self.igFO2)
              
              if self.de != 0.:
                  self.derivatives.arakawa_vec(self.petsc_solver.Jp, self.petsc_solver.Pp, self.igFA2)
                  self.igFA2.scale(self.de**2)
                  self.A.axpy(0.5*self.ht, self.igFA2)
              
              self.derivatives.laplace_vec(self.A, self.J, -1.)
              
              self.P.set(0.)
              self.O.copy(self.Pb)
              self.poisson_nullspace.remove(self.Pb)
              self.poisson_ksp.solve(self.Pb, self.P)
              
              self.derivatives.arakawa_vec(self.A, self.P, self.igFA1)
              self.derivatives.arakawa_vec(self.O, self.P, self.igFO1)
              self.derivatives.arakawa_vec(self.A, self.J, self.igFO2)
              
              self.petsc_solver.Ap.copy(self.A)
              self.petsc_solver.Op.copy(self.O)
              
              self.A.axpy(self.ht, self.igFA1)
              self.O.axpy(self.ht, self.igFO1)
              self.O.axpy(self.ht, self.igFO2)
              
              if self.de != 0.:
                  self.derivatives.arakawa_vec(self.J, self.P, self.igFA2)
                  self.igFA2.scale(self.de**2)
                  self.A.axpy(self.ht, self.igFA2)
  #                 self.A.axpy(-self.de**2, self.petsc_solver.Jp)
  #                 self.A.axpy(+self.de**2, self.J)
  
              self.derivatives.laplace_vec(self.A, self.J, -1.)
              
              self.P.set(0.)
              self.O.copy(self.Pb)
              self.poisson_nullspace.remove(self.Pb)
              self.poisson_ksp.solve(self.Pb, self.P)
              
              self.derivatives.arakawa_vec(self.petsc_solver.Ap, self.petsc_solver.Pp, self.igFA1)
              self.derivatives.arakawa_vec(self.petsc_solver.Op, self.petsc_solver.Pp, self.igFO1)
              self.derivatives.arakawa_vec(self.petsc_solver.Ap, self.petsc_solver.Jp, self.igFO2)
              
              if self.de != 0.:
                  self.derivatives.arakawa_vec(self.petsc_solver.Jp, self.petsc_solver.Pp, self.igFA2)
                  self.igFA2.scale(self.de**2)
              
          else:
              self.A.set(0.)
              self.O.set(0.)
      
              self.A.axpy(self.hermite_x0, self.petsc_solver.Ah)
              self.A.axpy(self.hermite_f0, self.igFA1)
              
              if self.de != 0.:
  #                 self.A.axpy(+self.de**2, self.petsc_solver.Jh)
                  self.A.axpy(self.hermite_f0, self.igFA2)
                  
              self.O.axpy(self.hermite_x0, self.petsc_solver.Oh)
              self.O.axpy(self.hermite_f0, self.igFO1)
              self.O.axpy(self.hermite_f0, self.igFO2)
              
              self.derivatives.arakawa_vec(self.petsc_solver.Ap, self.petsc_solver.Pp, self.igFA1)
              self.derivatives.arakawa_vec(self.petsc_solver.Op, self.petsc_solver.Pp, self.igFO1)
              self.derivatives.arakawa_vec(self.petsc_solver.Ap, self.petsc_solver.Jp, self.igFO2)
              
              if self.de != 0.:
                  self.derivatives.arakawa_vec(self.petsc_solver.Jp, self.petsc_solver.Pp, self.igFA2)
                  self.igFA2.scale(self.de**2)
              
              self.A.axpy(self.hermite_x1, self.petsc_solver.Ap)
              self.A.axpy(self.hermite_f1, self.igFA1)
              
              if self.de != 0.:
  #                 self.A.axpy(-self.de**2, self.petsc_solver.Jp)
                  self.A.axpy(self.hermite_f1, self.igFA2)
              
              self.O.axpy(self.hermite_x1, self.petsc_solver.Op)
              self.O.axpy(self.hermite_f1, self.igFO1)
              self.O.axpy(self.hermite_f1, self.igFO2)
              
              self.A.scale(self.hermite_den)
              self.O.scale(self.hermite_den)
          
          self.derivatives.laplace_vec(self.A, self.J, -1.)
          
          self.P.set(0.)
          self.O.copy(self.Pb)
          self.poisson_nullspace.remove(self.Pb)
          self.poisson_ksp.solve(self.Pb, self.P)
          
          
      def save_to_hdf5(self, timestep):
          if timestep % self.nsave == 0:
              if PETSc.COMM_WORLD.getRank() == 0:
                  self.time.setValue(0, self.ht*timestep)
              
              # copy solution to A, J, psi, omega vectors
              self.copy_x_from_da4_to_da1()
              
              # calculate B and V field
              self.derivatives.dy(self.A, self.Bx, +1.)
              self.derivatives.dx(self.A, self.By, -1.)
              self.derivatives.dy(self.P, self.Vx, +1.)
              self.derivatives.dx(self.P, self.Vy, -1.)
              
              # save timestep
              self.hdf5_viewer.setTimestep(timestep // self.nsave)
              self.hdf5_viewer(self.time)
              
              self.hdf5_viewer(self.A)
              self.hdf5_viewer(self.J)
              self.hdf5_viewer(self.P)
              self.hdf5_viewer(self.O)
              
              self.hdf5_viewer(self.Bx)
              self.hdf5_viewer(self.By)
              self.hdf5_viewer(self.Vx)
              self.hdf5_viewer(self.Vy)
          
  
      def check_jacobian(self):
          
          (xs, xe), (ys, ye) = self.da1.getRanges()
          
          eps = 1.E-7
          
          # calculate initial guess
  #        self.calculate_initial_guess()
          
          # update previous iteration
          self.petsc_solver.update_previous(self.x)
          
          # calculate jacobian
          self.petsc_solver.formMat(self.Jac)
          
          # create working vectors
          Jx  = self.da4.createGlobalVec()
          dJ  = self.da4.createGlobalVec()
          ex  = self.da4.createGlobalVec()
          dx  = self.da4.createGlobalVec()
          dF  = self.da4.createGlobalVec()
          Fxm = self.da4.createGlobalVec()
          Fxp = self.da4.createGlobalVec()
          
          
  #         sx = -2
  #         sx = -1
          sx =  0
  #         sx = +1
  #         sx = +2
  
  #         sy = -1
          sy =  0
  #         sy = +1
          
          nfield=4
          
          for ifield in range(0, nfield):
              for ix in range(xs, xe):
                  for iy in range(ys, ye):
                      for tfield in range(0, nfield):
                          
                          # compute ex
                          ex_arr = self.da4.getVecArray(ex)
                          ex_arr[:] = 0.
                          ex_arr[(ix+sx) % self.nx, (iy+sy) % self.ny, ifield] = 1.
                          
                          
                          # compute J.e
                          self.Jac.function(ex, dJ)
                          
                          dJ_arr = self.da4.getVecArray(dJ)
                          Jx_arr = self.da4.getVecArray(Jx)
                          Jx_arr[ix, iy, tfield] = dJ_arr[ix, iy, tfield]
                          
                          
                          # compute F(x - eps ex)
                          self.x.copy(dx)
                          dx_arr = self.da4.getVecArray(dx)
                          dx_arr[(ix+sx) % self.nx, (iy+sy) % self.ny, ifield] -= eps
                          
                          self.petsc_solver.function(dx, Fxm)
                          
                          
                          # compute F(x + eps ex)
                          self.x.copy(dx)
                          dx_arr = self.da4.getVecArray(dx)
                          dx_arr[(ix+sx) % self.nx, (iy+sy) % self.ny, ifield] += eps
                          
                          self.petsc_solver.function(dx, Fxp)
                          
                          
                          # compute dF = [F(x + eps ex) - F(x - eps ex)] / (2 eps)
                          Fxm_arr = self.da4.getVecArray(Fxm)
                          Fxp_arr = self.da4.getVecArray(Fxp)
                          dF_arr  = self.da4.getVecArray(dF)
                          
                          dF_arr[ix, iy, tfield] = ( Fxp_arr[ix, iy, tfield] - Fxm_arr[ix, iy, tfield] ) / (2. * eps)
                          
              
              diff = np.zeros(nfield)
              
              for tfield in range(0,nfield):
  #                print()
  #                print("Fields: (%5i, %5i)" % (ifield, tfield))
  #                print()
                  
                  Jx_arr = self.da4.getVecArray(Jx)[...][:, :, tfield]
                  dF_arr = self.da4.getVecArray(dF)[...][:, :, tfield]
                  
                  
  #                 print("Jacobian:")
  #                 print(Jx_arr)
  #                 print()
  #                 
  #                 print("[F(x+dx) - F(x-dx)] / [2 eps]:")
  #                 print(dF_arr)
  #                 print()
  #                
  #                print("Difference:")
  #                print(Jx_arr - dF_arr)
  #                print()
                  
                  
  #                if ifield == 3 and tfield == 2:
  #                    print("Jacobian:")
  #                    print(Jx_arr)
  #                    print()
  #                    
  #                    print("[F(x+dx) - F(x-dx)] / [2 eps]:")
  #                    print(dF_arr)
  #                    print()
                  
                  
                  diff[tfield] = (Jx_arr - dF_arr).max()
              
              print()
          
              for tfield in range(0,nfield):
                  print("max(difference)[field=%i, equation=%i] = %16.8E" % ( ifield, tfield, diff[tfield] ))
              
              print()
      
  
  
  
  
  
  
  if __name__ == '__main__':
  
      parser = argparse.ArgumentParser(description='PETSc Reduced MHD Solver in 2D')
      parser.add_argument('-c', '--config', metavar='<cfg_file>', type=str, required=True,
                          help='Configuration File')
      parser.add_argument('-m', '--mode', metavar='[asm, lu, ppc, split]', type=str, required=True,
                          help='Solver Mode')
      parser.add_argument('-p', '--profiler', action='store_true', required=False,
                          help='Activate Profiler')
      parser.add_argument('-j', '--jacobian', action='store_true', required=False,
                          help='Check Jacobian')
  #     
      args = parser.parse_args()
      
      runfile = args.config
      mode    = args.mode
      
      if mode == "ppc":
          # physics based preconditioner
          from run_rmhd2d_ppc import rmhd2d_ppc
          petscvp = rmhd2d_ppc(runfile)
      elif mode == "split":
          # split solver with physics based preconditioner
          from run_rmhd2d_split import rmhd2d_split
          petscvp = rmhd2d_split(runfile)
      elif mode == "asm":
          # additive schwarz preconditioner
          from run_rmhd2d_asm import rmhd2d_asm
          petscvp = rmhd2d_asm(runfile)
      else:
          # direct solver (lu decomposition)
          from run_rmhd2d_lu import rmhd2d_lu
          petscvp = rmhd2d_lu(runfile)
  
      if args.profiler:
          cProfile.runctx("petscvp.run()", globals(), locals(), "profile.prof")
          
          if PETSc.COMM_WORLD.getRank() == 0:
              s = pstats.Stats("profile.prof")
              s.strip_dirs().sort_stats("time").print_stats()
      elif args.jacobian:
          petscvp.check_jacobian()
      else:
          petscvp.run()
  