I am using CVODES to calculate adjoint sensitivities for a very basic equation with solution discontinuity at t=1. When integrating the forward solution, I integrate over each interval with CVodeF() in CV_ONE_STEP mode in a loop, calling CVodeReInit() to restart the integration at the discontinuity. This yields the correct forward solution. I then call CVodeB() for the backward integration for an adjoint sensitivity analysis.
My question regards forward integration restarts and how to handle them during the backward integration. At the beginning of my program, I call CVodeAdjInit() with Nd = 1, so I believe I am saving checkpoints at every integration step. However, when examining the checkpoints and also trying to recall the forward solution (with CVodeGetAdjY()) at and around t=1, I don't see the correct jump discontinuity at the correct time.
Are these restarts explicitly saved (in the checkpointing scheme or elsewhere)? Are there additional CVODES functions I should call to inform the backward integrator of these restarts? Any general guidance with using CVODES for an adjoint sensitivity analysis (in the presence of forward solution discontinuities) would be much appreciated.
I've included example code that illustrates this. We integrate dy/dt = -0.05*y from t = 0 to t = 2, with an impulse of 0.1 applied at t = 1. In this example, I am not doing anything with the adjoint state lambda - the main purpose is to illustrate how the forward solution y is recalled during the backward integration.
#include <stdio.h>
#include <cvodes/cvodes.h>
#include <nvector/nvector_serial.h> /* access to serial N_Vector */
#include <sunmatrix/sunmatrix_dense.h> /* access to dense SUNMatrix */
#include <sunlinsol/sunlinsol_dense.h> /* access to dense SUNLinearSolver */
/* Number of equations in system (1 for this example) */
#define NEQ 1
/* Accessor macros */
#define Ith(v, i) NV_Ith_S(v,i) /* i-th vector component, i=0..NEQ-1 */
#define IJth(A,i,j) SM_ELEMENT_D(A,i,j) /* IJth numbers rows,cols 0..NEQ-1 */
/* Decay rate*/
#define R 0.05;
static int rhs(realtype t, N_Vector y, N_Vector ydot, void *user_data);
static int Jac(realtype t, N_Vector y, N_Vector fy, SUNMatrix J,
void *user_data, N_Vector tmp1, N_Vector tmp2, N_Vector tmp3);
static int rhs_adj(realtype t, N_Vector y, N_Vector lam, N_Vector lamdot, void *user_data);
static int Jac_adj(realtype t,
N_Vector y, N_Vector lam, N_Vector fyB,
SUNMatrix JB, void *user_data,
N_Vector tmp1B, N_Vector tmp2B, N_Vector tmp3B);
int main() {
uint32_t maxord = 5; //BDF order
double abstol = 1e-8;
double reltol = 1e-8;
double abstol_adj = 1e-8;
double reltol_adj = 1e-8;
int adj_steps = 1; //number of integration steps between two consecutive checkpoints
int ncheck; //number of checkpoints
int indexB; //index for the backward problem
/* impulse at t=1*/
double impulse_time = 1.0;
double impulse_amount = 0.1;
/* integrate to t=2 */
double final_time = 2.0;
/* y = state */
N_Vector y = N_VNew_Serial(NEQ);
Ith(y, 0) = 1.0; //init condit.
/* lambda in adjoint equation. Needed for the backward integration, though will be ignored for this example */
N_Vector lam = N_VNew_Serial(NEQ);
Ith(lam, 0) = 0.0; //init condit.
/* initialize cvodes, set tolerances, etc. */
void *cvode_mem = CVodeCreate(CV_BDF);
CVodeSetMaxOrd(cvode_mem, maxord);
CVodeInit(cvode_mem, rhs, 0.0, y);
CVodeSStolerances(cvode_mem, reltol, abstol);
/* Create SUNMatrix and linear solver for the forward problem and set Jacobian fcn*/
SUNMatrix A = SUNDenseMatrix(NEQ, NEQ);
SUNLinearSolver LS = SUNLinSol_Dense(y, A);
CVodeSetLinearSolver(cvode_mem, LS, A);
CVodeSetJacFn(cvode_mem, Jac);
/* Inform the forward problem there will be an adjoint problem to solve as well */
CVodeAdjInit(cvode_mem, adj_steps, CV_HERMITE);
/* Initialization steps for adj. problem*/
CVodeCreateB(cvode_mem, CV_BDF, &indexB);
CVodeSetMaxOrdB(cvode_mem, indexB, maxord);
CVodeInitB(cvode_mem, indexB, rhs_adj, final_time, lam);
CVodeSStolerancesB(cvode_mem, indexB, reltol_adj, abstol_adj);
/* Create SUNMatrix and linear solver for the adj problem and attach */
SUNMatrix AB = SUNDenseMatrix(NEQ, NEQ);
SUNLinearSolver LSB = SUNLinSol_Dense(y, AB);
CVodeSetLinearSolverB(cvode_mem, indexB, LSB, AB);
CVodeSetJacFnB(cvode_mem, indexB, Jac_adj);
/* The forward integration */
realtype time; // updated by each integration step by CVodeF
double current_time = 0.0;
double goal_time = impulse_time; // we will first integrate to the time of the impulse
int impulse_applied = 0;
int init_needed = 0;
while (current_time < final_time) {
/* need to re-initialize cvodes after the jump */
if (init_needed) {
CVodeReInit(cvode_mem, current_time, y);
printf("Re-init after impulse\n");
init_needed = 0;
}
while (current_time < goal_time) {
/* main forward integration step */
CVodeF(cvode_mem, goal_time, y, &time, CV_ONE_STEP, &ncheck);
current_time = time;
printf("t = %10.8f, y = %10.8f | ncheck = %d\n", current_time, Ith(y, 0), ncheck);
}
/* apply impulse */
if (!impulse_applied && impulse_time <= current_time) {
current_time = impulse_time;
CVodeGetDky(cvode_mem, impulse_time, 0, y);
printf("****** Before impulse: t = %10.8f, y = %10.8f \n", current_time, Ith(y, 0));
Ith(y, 0) += impulse_amount;
printf("****** After impulse: t = %10.8f, y = %10.8f \n", current_time, Ith(y, 0));
init_needed = 1;
impulse_applied = 1;
goal_time = final_time;
}
}
/* Now, integrate backwards in time for the adjoint problem */
goal_time = 1.0;
printf("\nPerforming backward integration ...\n");
while (current_time > goal_time) {
/* main backward integration step */
CVodeB(cvode_mem, goal_time, CV_ONE_STEP);
/* need to call CVodeGetB to get the time of the last CVodeB step */
CVodeGetB(cvode_mem, indexB, &time, lam);
/* obtain interpolated forward solution at current time as well */
CVodeGetAdjY(cvode_mem, time, y);
printf("t = %10.8f, y = %10.8f \n", time, Ith(y, 0));
current_time = time;
}
printf("Around impulse: \n");
double times[5] = {1.002, 1.001, 1.0, 0.999, 0.998};
for (int i = 0; i < 5; i++){
CVodeGetAdjY(cvode_mem, times[i], y);
printf("t = %10.8f, y = %10.8f \n", times[i], Ith(y, 0));
}
/* cleanup */
N_VDestroy(y);
N_VDestroy(lam);
CVodeFree(&cvode_mem);
SUNLinSolFree(LS);
SUNLinSolFree(LSB);
SUNMatDestroy(A);
SUNMatDestroy(AB);
return 0;
}
static int rhs(realtype t, N_Vector y, N_Vector ydot, void *user_data) {
/* exp decay */
double r = R;
Ith(ydot, 0) = -r * Ith(y, 0);
return (0);
}
static int Jac(realtype t, N_Vector y, N_Vector fy, SUNMatrix J,
void *user_data, N_Vector tmp1, N_Vector tmp2, N_Vector tmp3) {
/* Jacobian of rhs */
double r = R;
IJth(J, 0, 0) = -r;
return (0);
}
static int rhs_adj(realtype t, N_Vector y, N_Vector lam, N_Vector lamdot, void *user_data) {
/* RHS of adjoint problem
* Note: the adjoint problem is lam' = -(J^*)lam - g_y^*, where J is the Jacobian matrix of the main problem.
* For this example, we take g = 0 everywhere */
double r = R;
Ith(lamdot, 0) = r * Ith(lam, 0);
return (0);
}
static int Jac_adj(realtype t,
N_Vector y, N_Vector lam, N_Vector fyB,
SUNMatrix JB, void *user_data,
N_Vector tmp1B, N_Vector tmp2B, N_Vector tmp3B) {
/* Jacobian of adjoint problem */
double r = R;
IJth(JB, 0, 0) = r;
return (0);
}
During the forward integration, we have the output (around t = 1)
...
t = 0.77864484, y = 0.96181582 | ncheck = 18
t = 0.94641606, y = 0.95378132 | ncheck = 19
t = 1.11418727, y = 0.94581393 | ncheck = 20
****** Before impulse: t = 1.00000000, y = 0.95122937
****** After impulse: t = 1.00000000, y = 1.05122937
Re-init after impulse
t = 1.00197548, y = 1.05112555 | ncheck = 21
t = 1.00395097, y = 1.05102173 | ncheck = 22
...
During the backward phase (here, y is obtained via CVodeGetAdjY),
...
t = 1.00934328, y = 1.05073841
t = 1.00395097, y = 1.05102173
t = 1.00197548, y = 1.05112555
t = 0.98222065, y = 0.95207536
The recalled y value at t = 1.00197548 is correct at that time (this is the first step after the impulse taken during the forward integration), but if I then query y at times near the impulse (again with CVodeGetAdjY):
Around impulse:
t = 1.00200000, y = 0.95113425
t = 1.00100000, y = 0.95118181
t = 1.00000000, y = 0.95122937
t = 0.99900000, y = 0.95127693
t = 0.99800000, y = 0.95132450
the impulse is not apparent. The recalled y at t = 1.0 is the pre-impulse value. It appears as though, even though CVodeReInit() is called immediately after the impulse during the forward integration, the post-impulse y is not "seen" during the backward integration. Moreover, if the backward integrator had required any steps between checkpoints, the interpolated y between 1.00197548 and t = 1.0 would likely be off.
In other words, my question is: After a re-init of the forward problem, is there a way to ensure that such a restart is saved and accessible from the checkpoint data?