[petsc-dev] Discussion about time-dependent optimization moved from PR

[Barry Smith]

> On Oct 17, 2017, at 10:41 AM, Jed Brown <jed at jedbrown.org> wrote:
> 
> Barry Smith <bsmith at mcs.anl.gov> writes:
> 
>>>>> My
>>>>> preference is for TSCreateAdjointTS() which the user can customize to
>>>>> use arbitrary time integration methods (thus becoming a
>>>>> "continuous-in-time" adjoint) and to use different spatial
>>>>> discretizations (for consistency with the adjoint PDE).
>>>> 
>>>>  I understand, I think TSCreateAdjointTS() is not the right model. I think that the correct model is, for example
>>>> 
>>>> PetscSensiCreate()
>>>> PetscSensiSetObjective.... or whatever etc etc
>>>> PetscSensiGetIntegrator(,&TS)
>>>> PetscSensiGetAdjointIntegrator(,&TS).
>>> 
>>> How would PetscSensiGetAdjointIntegrator be implemented?  Since it needs
>>> to be able to create a discrete adjoint (which depends on the forward TS
>>> implementation), the control flow *must* go through TS.  What does that
>>> interface look like?
>> 
>>   I am not sure what you mean. 
>> 
>>   PetscSensi will contain the information about objective function etc as well as a pointer to the trajectory history and (of course) a pointer to the original TS. All the  auxiliary vectors etc related to sensitivities. 
>> 
>>   The adjoint TS that is returned has the calls need to get access to the forward solution that it needs etc from the PetscSensi object (which likely will call the forward TS for missing timestep values etc. But the calls go through the PetscSensi object from the adjoint TS, they don't go directly to the forward TS.)
> 
> To be concrete with TSRK, PetscSensi does not have access to the Butcher
> table for the forward method because that only exists in
> src/ts/impls/explicit/rk/rk.c.  Since the discrete adjoint depends on
> this detail of the forward method, there MUST be some code somewhere
> that depends on the forward Butcher table.  Since that data ONLY EXISTS
> in rk.c, whatever it is that PetscSensi does, there must be some code in
> rk.c to define the adjoint method.

   Of course, that would be incorporated in what is an essentially a TSRKADJOINT implementation that does the adjoint integration. Note that (1) this is simpler than the regular TSRK because it only needs to handle linear with given time steps (2) it also computes along sensitivity information with calls to TSAdjointComputeDRDPFunction() etc

   Here is what the code looks like now. Just needs minor refactoring to make it match my proposal.

static PetscErrorCode TSAdjointStep_RK(TS ts)
{
  TS_RK           *rk   = (TS_RK*)ts->data;
  RKTableau        tab  = rk->tableau;
  const PetscInt   s    = tab->s;
  const PetscReal *A = tab->A,*b = tab->b,*c = tab->c;
  PetscScalar     *w    = rk->work;
  Vec             *Y    = rk->Y,*VecDeltaLam = rk->VecDeltaLam,*VecDeltaMu = rk->VecDeltaMu,*VecSensiTemp = rk->VecSensiTemp;
  PetscInt         i,j,nadj;
  PetscReal        t = ts->ptime;
  PetscErrorCode   ierr;
  PetscReal        h = ts->time_step;
  Mat              J,Jp;

  PetscFunctionBegin;
  rk->status = TS_STEP_INCOMPLETE;
  for (i=s-1; i>=0; i--) {
    rk->stage_time = t + h*(1.0-c[i]);
    for (nadj=0; nadj<ts->numcost; nadj++) {
      ierr = VecCopy(ts->vecs_sensi[nadj],VecSensiTemp[nadj]);CHKERRQ(ierr);
      ierr = VecScale(VecSensiTemp[nadj],-h*b[i]);CHKERRQ(ierr);
      for (j=i+1; j<s; j++) {
        ierr = VecAXPY(VecSensiTemp[nadj],-h*A[j*s+i],VecDeltaLam[nadj*s+j]);CHKERRQ(ierr);
      }
    }
    /* Stage values of lambda */
    ierr = TSGetRHSJacobian(ts,&J,&Jp,NULL,NULL);CHKERRQ(ierr);
    ierr = TSComputeRHSJacobian(ts,rk->stage_time,Y[i],J,Jp);CHKERRQ(ierr);
    if (ts->vec_costintegral) {
      ierr = TSAdjointComputeDRDYFunction(ts,rk->stage_time,Y[i],ts->vecs_drdy);CHKERRQ(ierr);
    }
    for (nadj=0; nadj<ts->numcost; nadj++) {
      ierr = MatMultTranspose(J,VecSensiTemp[nadj],VecDeltaLam[nadj*s+i]);CHKERRQ(ierr);
      if (ts->vec_costintegral) {
        ierr = VecAXPY(VecDeltaLam[nadj*s+i],-h*b[i],ts->vecs_drdy[nadj]);CHKERRQ(ierr);
      }
    }

    /* Stage values of mu */
    if(ts->vecs_sensip) {
      ierr = TSAdjointComputeRHSJacobian(ts,rk->stage_time,Y[i],ts->Jacp);CHKERRQ(ierr);
      if (ts->vec_costintegral) {
        ierr = TSAdjointComputeDRDPFunction(ts,rk->stage_time,Y[i],ts->vecs_drdp);CHKERRQ(ierr);
      }

      for (nadj=0; nadj<ts->numcost; nadj++) {
        ierr = MatMultTranspose(ts->Jacp,VecSensiTemp[nadj],VecDeltaMu[nadj*s+i]);CHKERRQ(ierr);
        if (ts->vec_costintegral) {
          ierr = VecAXPY(VecDeltaMu[nadj*s+i],-h*b[i],ts->vecs_drdp[nadj]);CHKERRQ(ierr);
        }
      }
    }
  }

  for (j=0; j<s; j++) w[j] = 1.0;
  for (nadj=0; nadj<ts->numcost; nadj++) {
    ierr = VecMAXPY(ts->vecs_sensi[nadj],s,w,&VecDeltaLam[nadj*s]);CHKERRQ(ierr);
    if(ts->vecs_sensip) {
      ierr = VecMAXPY(ts->vecs_sensip[nadj],s,w,&VecDeltaMu[nadj*s]);CHKERRQ(ierr);
    }
  }
  rk->status = TS_STEP_COMPLETE;
  PetscFunctionReturn(0);
}

> 
> What is that code and what is the call stack to reach it?

Now you ask, how does the PetscSensi know to use that discrete adjoint integrator and return the right TS (instead of an empty one requiring the user to set the type)? This is simple, each TS can return 
its discrete adjoint integrator when requested, sure mechanically it is like a hidden TSGetAdjointTS() but from the public API point of view it is much clearer to come from the PetscSensi. You tell the PetscSensi all the information you need about the derivatives, you don't tell the TS. You are also right that not EVERYTHING related to sensitivities is only in the PetscSensi code since the TSRKADJOINT object knows about sensitivities that it carries around (and is in rk.c) but from the user point of view only the PetscSensi knows about sensitivities.

   Of course I could still be missing stuff

   Barry


> 
>>   Both proposed models shove one heck of a lot of stuff into TS; I think it would be a better design and better user API if we don't do that; we reserve the TS for being able to integrate something (and that is about all). Maybe it is not possible but I don't see why it is not possible.
>> 
>>   Since I don't know the details I asked for the people who do know the details (Hong and Stefano) to try to develop such an API (or clearly articular why it is impossible.)