Solving an Eigen::Matrix system type RHS with PETSc implicit timestepping + JFNK

teststiffjfnk.cpp

#include <iostream>
#include <math.h>
#include <Eigen/Dense>
#include <petscts.h>

typedef struct {
  //  Vec         solution; /* global exact solution vector */
  Eigen::MatrixXd soln,rhs;
  const double y0 = 1.0e-3; // epsilon in original ODE
} AppCtx;

// not the correct way to do this, look into raw pointer interface to prevent copying
void PETScVecToEigenMat(Vec& pvec,unsigned int nrow,unsigned int ncol,Eigen::MatrixXd& mat){
  mat.resize(nrow,ncol);
  int pidx = 0;
  for (int i=0; i!=nrow; ++i){
    for (int j=0; j!=ncol; ++j){
      PetscInt ix[] = {pidx};
      PetscScalar y[] = {0};
      VecGetValues(pvec,1,ix,y);
      mat(i,j) = y[0];
      pidx++;
    }
  }
}

// not the correct way to do this, look into raw pointer interface to prevent copying
void EigenMatToPETScVec(const Eigen::MatrixXd& in,Vec& out){
  unsigned int size = in.size();
  const double* data = in.data(); // raw pointer to eigen storage
  for (unsigned int i=0; i!=size; ++i){
    VecSetValue(out,i,data[i],INSERT_VALUES);
  }
  VecAssemblyBegin(out);
  VecAssemblyEnd(out);
}

PetscErrorCode Monitor(TS ts,PetscInt step,PetscReal time,Vec u,void *ctx){
  PetscInt ix[] = {0};
  PetscScalar y[] = {0};
  VecGetValues(u,1,ix,y);
  PetscPrintf(PETSC_COMM_SELF,"%f,%f;\n",time,y[0]); // petsc doesn't like stdout here..
}


// example from http://www.scholarpedia.org/article/Stiff_systems
// y' = y^2 * (1-y)
void StiffRHS(const Eigen::MatrixXd& y,const double& t,Eigen::MatrixXd& dydt){
  dydt.resize(1,1);
  dydt(0,0) = pow(y(0,0),2.0) * (1.0-y(0,0));
}

// interface the function petsc needs to our desired interface
PetscErrorCode PetscRHSFunctionInterface(TS ts,PetscReal time,Vec X,Vec F,void *ctx){
  AppCtx *appctx = (AppCtx*)ctx;
  PETScVecToEigenMat(X,1,1,appctx->soln);  // copy petsc vector to internal solution vector
  StiffRHS(appctx->soln,time,appctx->rhs); // call our RHS
  EigenMatToPETScVec(appctx->rhs,F);       // copy RHS evaluation to petsc vector
}

PetscErrorCode PetscIFunctionInterface(TS ts,PetscReal time,Vec X,Vec Xdot,Vec F,void *ctx){
  AppCtx *appctx = (AppCtx*)ctx;
  PETScVecToEigenMat(X,1,1,appctx->soln);  // copy petsc vector to internal solution vector
  StiffRHS(appctx->soln,time,appctx->rhs); // call our RHS
  EigenMatToPETScVec(appctx->rhs,F);    // copy RHS evaluation to petsc vector
  VecAYPX(F,-1.0,Xdot);  // F = Xdot - RHS  
}


int main(){
  AppCtx ctx;
  ctx.soln.resize(1,1);
  ctx.rhs.resize(1,1);
  double       tfinal = 2.0/ctx.y0;
  unsigned int nsteps = 50;
  double           dt = tfinal/double(nsteps);
  TS               ts;                     /* timestepping context */

  PetscInitialize(0,0,0,0);

  // create solution vector and RHS vector
  Vec PSolnVec,PRHSVec;
  VecCreateSeq(PETSC_COMM_SELF,1,&PSolnVec);
  VecDuplicate(PSolnVec,&PRHSVec);

  // create timestepper context
  TSCreate(PETSC_COMM_SELF,&ts);
  TSSetSolution(ts,PSolnVec); // tell the timestepper context where to compute solutions
  TSSetType(ts,TSSSP); // sets stepping to RK-SSP
  TSSSPSetNumStages(ts,4);
  TSSSPSetType(ts,TSSSPRKS3); // sets SSP to rk3 with 4 stages
  TSSetInitialTimeStep(ts,0,dt);
  TSSetDuration(ts,nsteps*2,tfinal);
  
  TSSetRHSFunction(ts,PRHSVec,PetscRHSFunctionInterface,&ctx); // interface to RHS

  TSMonitorSet(ts,Monitor,&ctx,NULL); // user-defined monitor routine

  ctx.soln(0,0) = ctx.y0; 
  EigenMatToPETScVec(ctx.soln,PSolnVec); // initial condition into solution

  TSSolve(ts,PSolnVec); // run it. should explode with explicit stepping.

  TSDestroy(&ts);

  /**********************/
  std::cout << "*** end explicit, begin implicit ***\n";
  /**********************/


  TS its; // new time-stepping context for implicit stepping
  TSCreate(PETSC_COMM_SELF,&its);
  TSSetProblemType(its,TS_NONLINEAR);
  TSSetType(its,TSBEULER); // backward euler
  Vec Residual; // vector to hold the residual for implicit scheme
  VecDuplicate(PSolnVec,&Residual); // size it
  TSSetIFunction(its,Residual,PetscIFunctionInterface,&ctx); // implicit function for solution
  TSSetInitialTimeStep(its,0,dt);
  TSSetDuration(its,nsteps*2,tfinal);

  ctx.soln(0,0) = ctx.y0; 
  EigenMatToPETScVec(ctx.soln,PSolnVec); // reset initial condition into solution
  TSSetSolution(its,PSolnVec);

  // Set up matrix-free jacobian
  SNES snes; // non-linear algebraic solver context
  
  TSGetSNES(its,&snes);
  Mat JacMatFree=NULL;
  MatCreateSNESMF(snes,&JacMatFree); // tell non-linear solver this is matrix-free
  SNESSetJacobian(snes,JacMatFree,NULL,SNESComputeJacobianDefault,&ctx);
  
  KSP ksp; // linear solver context
  PC pc;   // preconditioner context
  SNESGetKSP(snes,&ksp);
  KSPSetType(ksp,"gmres"); // use gmres
  KSPGetPC(ksp,&pc);
  PCSetType(pc,PCNONE); // no preconditioner

  TSSetMaxSNESFailures(its,-1); // -1 unlimited, probably a bad idea in general.

  TSMonitorSet(its,Monitor,&ctx,NULL); // user-defined monitor routine

  TSSolve(its,PSolnVec);

}