Actual source code: cn.c

      1: #define PETSCTS_DLL

      3: /*
      4:        Code for Timestepping with implicit Crank-Nicholson method.
      5:     THIS IS NOT YET COMPLETE -- DO NOT USE!!
      6: */
     7:  #include src/ts/tsimpl.h

      9: typedef struct {
     10:   Vec  update;      /* work vector where new solution is formed */
     11:   Vec  func;        /* work vector where F(t[i],u[i]) is stored */
     12:   Vec  rhs;         /* work vector for RHS; vec_sol/dt */
     13: } TS_CN;

     15: /*------------------------------------------------------------------------------*/
     18: /*
     19:    TSComputeRHSFunctionEuler - Evaluates the right-hand-side function.

     21:    Note: If the user did not provide a function but merely a matrix,
     22:    this routine applies the matrix.
     23: */
     24: PetscErrorCode TSComputeRHSFunctionEuler(TS ts,PetscReal t,Vec x,Vec y)
     25: {
     27:   PetscScalar    neg_two = -2.0,neg_mdt = -1.0/ts->time_step;

     34:   if (ts->ops->rhsfunction) {
     35:     PetscStackPush("TS user right-hand-side function");
     36:     (*ts->ops->rhsfunction)(ts,t,x,y,ts->funP);
     37:     PetscStackPop;
     38:     return(0);
     39:   }

     41:   if (ts->ops->rhsmatrix) { /* assemble matrix for this timestep */
     42:     MatStructure flg;
     43:     PetscStackPush("TS user right-hand-side matrix function");
     44:     (*ts->ops->rhsmatrix)(ts,t,&ts->A,&ts->B,&flg,ts->jacP);
     45:     PetscStackPop;
     46:   }
     47:   MatMult(ts->A,x,y);
     48:   /* shift: y = y -2*x */
     49:   VecAXPY(y,neg_two,x);
     50:   /* scale: y = y -2*x */
     51:   VecScale(y,neg_mdt);

     53:   /* apply user-provided boundary conditions (only needed if these are time dependent) */
     54:   TSComputeRHSBoundaryConditions(ts,t,y);

     56:   return(0);
     57: }

     59: /*
     60:     Version for linear PDE where RHS does not depend on time. Has built a
     61:   single matrix that is to be used for all timesteps.
     62: */
     65: static PetscErrorCode TSStep_CN_Linear_Constant_Matrix(TS ts,PetscInt *steps,PetscReal *ptime)
     66: {
     67:   TS_CN          *cn = (TS_CN*)ts->data;
     68:   Vec            sol = ts->vec_sol,update = cn->update;
     69:   Vec            rhs = cn->rhs;
     71:   PetscInt       i,max_steps = ts->max_steps,its;
     72:   PetscScalar    dt = ts->time_step,two = 2.0;
     73:   KSP            ksp;

     76:   TSGetKSP(ts,&ksp);
     77:   *steps = -ts->steps;
     78:   TSMonitor(ts,ts->steps,ts->ptime,sol);

     80:   /* set initial guess to be previous solution */
     81:   VecCopy(sol,update);

     83:   for (i=0; i<max_steps; i++) {
     84:     ts->ptime += ts->time_step;
     85:     if (ts->ptime > ts->max_time) break;

     87:     /* phase 1 - explicit step */
     88:     TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
     89:     VecAXPBY(sol,dt,two,update);

     91:     /* phase 2 - implicit step */
     92:     VecCopy(sol,rhs);
     93:     /* apply user-provided boundary conditions (only needed if they are time dependent) */
     94:     TSComputeRHSBoundaryConditions(ts,ts->ptime,rhs);

     96:     KSPSolve(ts->ksp,rhs,update);
     97:     KSPGetIterationNumber(ksp,&its);
     98:     ts->linear_its += PetscAbsInt(its);
     99:     VecCopy(update,sol);
    100:     ts->steps++;
    101:     TSMonitor(ts,ts->steps,ts->ptime,sol);
    102:   }

    104:   *steps += ts->steps;
    105:   *ptime  = ts->ptime;
    106:   return(0);
    107: }
    108: /*
    109:       Version where matrix depends on time
    110: */
    113: static PetscErrorCode TSStep_CN_Linear_Variable_Matrix(TS ts,PetscInt *steps,PetscReal *ptime)
    114: {
    115:   TS_CN          *cn = (TS_CN*)ts->data;
    116:   Vec            sol = ts->vec_sol,update = cn->update,rhs = cn->rhs;
    118:   PetscInt       i,max_steps = ts->max_steps,its;
    119:   PetscScalar    dt = ts->time_step,two = 2.0;
    120:   MatStructure   str;
    121:   KSP            ksp;

    124:   TSGetKSP(ts,&ksp);
    125:   *steps = -ts->steps;
    126:   TSMonitor(ts,ts->steps,ts->ptime,sol);

    128:   /* set initial guess to be previous solution */
    129:   VecCopy(sol,update);

    131:   for (i=0; i<max_steps; i++) {
    132:     ts->ptime += ts->time_step;
    133:     if (ts->ptime > ts->max_time) break;
    134:     /*
    135:         evaluate matrix function
    136:     */
    137:     (*ts->ops->rhsmatrix)(ts,ts->ptime,&ts->A,&ts->B,&str,ts->jacP);
    138:     TSScaleShiftMatrices(ts,ts->A,ts->B,str);

    140:     /* phase 1 - explicit step */
    141:     TSComputeRHSFunctionEuler(ts,ts->ptime,sol,update);
    142:     VecAXPBY(sol,dt,two,update);

    144:     /* phase 2 - implicit step */
    145:     VecCopy(sol,rhs);

    147:     /* apply user-provided boundary conditions (only needed if they are time dependent) */
    148:     TSComputeRHSBoundaryConditions(ts,ts->ptime,rhs);

    150:     KSPSetOperators(ts->ksp,ts->A,ts->B,str);
    151:     KSPSolve(ts->ksp,rhs,update);
    152:     KSPGetIterationNumber(ksp,&its);
    153:     ts->linear_its += PetscAbsInt(its);
    154:     VecCopy(update,sol);
    155:     ts->steps++;
    156:     TSMonitor(ts,ts->steps,ts->ptime,sol);
    157:   }

    159:   *steps += ts->steps;
    160:   *ptime  = ts->ptime;
    161:   return(0);
    162: }
    163: /*
    164:     Version for nonlinear PDE.
    165: */
    168: static PetscErrorCode TSStep_CN_Nonlinear(TS ts,PetscInt *steps,PetscReal *ptime)
    169: {
    170:   Vec            sol = ts->vec_sol;
    172:   PetscInt       i,max_steps = ts->max_steps,its,lits;
    173:   TS_CN          *cn = (TS_CN*)ts->data;
    174:
    176:   *steps = -ts->steps;
    177:   TSMonitor(ts,ts->steps,ts->ptime,sol);

    179:   for (i=0; i<max_steps; i++) {
    180:     ts->ptime += ts->time_step;
    181:     if (ts->ptime > ts->max_time) break;
    182:     VecCopy(sol,cn->update);
    183:     SNESSolve(ts->snes,PETSC_NULL,cn->update);
    184:     SNESGetIterationNumber(ts->snes,&its);
    185:     SNESGetNumberLinearIterations(ts->snes,&lits);
    186:     ts->nonlinear_its += its; ts->linear_its += lits;
    187:     VecCopy(cn->update,sol);
    188:     ts->steps++;
    189:     TSMonitor(ts,ts->steps,ts->ptime,sol);
    190:   }

    192:   *steps += ts->steps;
    193:   *ptime  = ts->ptime;
    194:   return(0);
    195: }

    197: /*------------------------------------------------------------*/
    200: static PetscErrorCode TSDestroy_CN(TS ts)
    201: {
    202:   TS_CN          *cn = (TS_CN*)ts->data;

    206:   if (cn->update) {VecDestroy(cn->update);}
    207:   if (cn->func) {VecDestroy(cn->func);}
    208:   if (cn->rhs) {VecDestroy(cn->rhs);}
    209:   PetscFree(cn);
    210:   return(0);
    211: }

    213: /*
    214:     This defines the nonlinear equation that is to be solved with SNES

    216:               U^{n+1} - dt*F(U^{n+1}) - U^{n}
    217: */
    220: PetscErrorCode TSCnFunction(SNES snes,Vec x,Vec y,void *ctx)
    221: {
    222:   TS             ts = (TS) ctx;
    223:   PetscScalar    mdt = 1.0/ts->time_step,*unp1,*un,*Funp1;
    225:   PetscInt       i,n;

    228:   /* apply user provided function */
    229:   TSComputeRHSFunction(ts,ts->ptime,x,y);
    230:   /* (u^{n+1} - U^{n})/dt - F(u^{n+1}) */
    231:   VecGetArray(ts->vec_sol,&un);
    232:   VecGetArray(x,&unp1);
    233:   VecGetArray(y,&Funp1);
    234:   VecGetLocalSize(x,&n);

    236:   for (i=0; i<n; i++) {
    237:     Funp1[i] = mdt*(unp1[i] - un[i]) - Funp1[i];
    238:   }
    239:   VecRestoreArray(ts->vec_sol,&un);
    240:   VecRestoreArray(x,&unp1);
    241:   VecRestoreArray(y,&Funp1);
    242:   return(0);
    243: }

    245: /*
    246:    This constructs the Jacobian needed for SNES

    248:              J = I/dt - J_{F}   where J_{F} is the given Jacobian of F.
    249: */
    252: PetscErrorCode TSCnJacobian(SNES snes,Vec x,Mat *AA,Mat *BB,MatStructure *str,void *ctx)
    253: {
    254:   TS             ts = (TS) ctx;

    258:   /* construct user's Jacobian */
    259:   TSComputeRHSJacobian(ts,ts->ptime,x,AA,BB,str);

    261:   /* shift and scale Jacobian */
    262:   TSScaleShiftMatrices(ts,*AA,*BB,*str);
    263:   return(0);
    264: }

    266: /* ------------------------------------------------------------*/
    269: static PetscErrorCode TSSetUp_CN_Linear_Constant_Matrix(TS ts)
    270: {
    271:   TS_CN          *cn = (TS_CN*)ts->data;

    275:   VecDuplicate(ts->vec_sol,&cn->update);
    276:   VecDuplicate(ts->vec_sol,&cn->rhs);
    277:
    278:   /* build linear system to be solved */
    279:   TSScaleShiftMatrices(ts,ts->A,ts->B,SAME_NONZERO_PATTERN);
    280:   KSPSetOperators(ts->ksp,ts->A,ts->B,SAME_NONZERO_PATTERN);
    281:   return(0);
    282: }

    286: static PetscErrorCode TSSetUp_CN_Linear_Variable_Matrix(TS ts)
    287: {
    288:   TS_CN          *cn = (TS_CN*)ts->data;

    292:   VecDuplicate(ts->vec_sol,&cn->update);
    293:   VecDuplicate(ts->vec_sol,&cn->rhs);
    294:   return(0);
    295: }

    299: static PetscErrorCode TSSetUp_CN_Nonlinear(TS ts)
    300: {
    301:   TS_CN          *cn = (TS_CN*)ts->data;

    305:   VecDuplicate(ts->vec_sol,&cn->update);
    306:   VecDuplicate(ts->vec_sol,&cn->func);
    307:   SNESSetFunction(ts->snes,cn->func,TSCnFunction,ts);
    308:   SNESSetJacobian(ts->snes,ts->A,ts->B,TSCnJacobian,ts);
    309:   return(0);
    310: }
    311: /*------------------------------------------------------------*/

    315: static PetscErrorCode TSSetFromOptions_CN_Linear(TS ts)
    316: {

    320:   KSPSetFromOptions(ts->ksp);
    321:   return(0);
    322: }

    326: static PetscErrorCode TSSetFromOptions_CN_Nonlinear(TS ts)
    327: {

    331:   SNESSetFromOptions(ts->snes);
    332:   return(0);
    333: }

    337: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
    338: {
    340:   return(0);
    341: }

    343: /* ------------------------------------------------------------ */
    344: /*MC
    345:       TS_CN - ODE solver using the implicit Crank-Nicholson method

    347:   Level: beginner

    349: .seealso:  TSCreate(), TS, TSSetType()

    351: M*/
    355: PetscErrorCode PETSCTS_DLLEXPORT TSCreate_CN(TS ts)
    356: {
    357:   TS_CN          *cn;
    359:   KSP            ksp;

    362:   ts->ops->destroy         = TSDestroy_CN;
    363:   ts->ops->view            = TSView_CN;

    365:   if (ts->problem_type == TS_LINEAR) {
    366:     if (!ts->A) {
    367:       SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"Must set rhs matrix for linear problem");
    368:     }
    369:     if (!ts->ops->rhsmatrix) {
    370:       ts->ops->setup  = TSSetUp_CN_Linear_Constant_Matrix;
    371:       ts->ops->step   = TSStep_CN_Linear_Constant_Matrix;
    372:     } else {
    373:       ts->ops->setup  = TSSetUp_CN_Linear_Variable_Matrix;
    374:       ts->ops->step   = TSStep_CN_Linear_Variable_Matrix;
    375:     }
    376:     ts->ops->setfromoptions  = TSSetFromOptions_CN_Linear;
    377:     KSPCreate(ts->comm,&ts->ksp);
    378:     TSGetKSP(ts,&ksp);
    379:     KSPSetInitialGuessNonzero(ksp,PETSC_TRUE);
    380:   } else if (ts->problem_type == TS_NONLINEAR) {
    381:     ts->ops->setup           = TSSetUp_CN_Nonlinear;
    382:     ts->ops->step            = TSStep_CN_Nonlinear;
    383:     ts->ops->setfromoptions  = TSSetFromOptions_CN_Nonlinear;
    384:     SNESCreate(ts->comm,&ts->snes);
    385:   } else SETERRQ(PETSC_ERR_ARG_OUTOFRANGE,"No such problem");

    387:   PetscNew(TS_CN,&cn);
    388:   PetscLogObjectMemory(ts,sizeof(TS_CN));
    389:   ts->data = (void*)cn;

    391:   return(0);
    392: }