1 #include <petscdmda.h>
2 #include <petscts.h>
3 #include <adolc/interfaces.h>
4 #include "contexts.cxx"
5 
6 /*
7    REQUIRES configuration of PETSc with option --download-adolc.
8 
9    For documentation on ADOL-C, see
10      $PETSC_ARCH/externalpackages/ADOL-C-2.6.0/ADOL-C/doc/adolc-manual.pdf
11 */
12 
13 /*
14   ADOL-C implementation for Jacobian vector product, using the forward mode of AD.
15   Intended to overload MatMult in matrix-free methods where implicit timestepping
16   has been used.
17 
18   For an implicit Jacobian we may use the rule that
19      G = M*xdot + f(x)    ==>     dG/dx = a*M + df/dx,
20   where a = d(xdot)/dx is a constant. Evaluated at x0 and acting upon a vector x1:
21      (dG/dx)(x0) * x1 = (a*df/d(xdot)(x0) + df/dx)(x0) * x1.
22 
23   Input parameters:
24   A_shell - Jacobian matrix of MatShell type
25   X       - vector to be multiplied by A_shell
26 
27   Output parameters:
28   Y       - product of A_shell and X
29 */
PetscAdolcIJacobianVectorProduct(Mat A_shell,Vec X,Vec Y)30 PetscErrorCode PetscAdolcIJacobianVectorProduct(Mat A_shell,Vec X,Vec Y)
31 {
32   AdolcMatCtx        *mctx;
33   PetscErrorCode    ierr;
34   PetscInt          m,n,i,j,k = 0,d;
35   const PetscScalar *x0;
36   PetscScalar       *action,*x1;
37   Vec               localX1;
38   DM                da;
39   DMDALocalInfo     info;
40 
41   PetscFunctionBegin;
42 
43   /* Get matrix-free context info */
44   ierr = MatShellGetContext(A_shell,(void**)&mctx);CHKERRQ(ierr);
45   m = mctx->m;
46   n = mctx->n;
47 
48   /* Get local input vectors and extract data, x0 and x1*/
49   ierr = TSGetDM(mctx->ts,&da);CHKERRQ(ierr);
50   ierr = DMDAGetLocalInfo(da,&info);CHKERRQ(ierr);
51   ierr = DMGetLocalVector(da,&localX1);CHKERRQ(ierr);
52   ierr = DMGlobalToLocalBegin(da,X,INSERT_VALUES,localX1);CHKERRQ(ierr);
53   ierr = DMGlobalToLocalEnd(da,X,INSERT_VALUES,localX1);CHKERRQ(ierr);
54 
55   ierr = VecGetArrayRead(mctx->localX0,&x0);CHKERRQ(ierr);
56   ierr = VecGetArray(localX1,&x1);CHKERRQ(ierr);
57 
58   /* dF/dx part */
59   ierr = PetscMalloc1(m,&action);CHKERRQ(ierr);
60   ierr = PetscLogEventBegin(mctx->event1,0,0,0,0);CHKERRQ(ierr);
61   fos_forward(mctx->tag1,m,n,0,x0,x1,NULL,action);
62   for (j=info.gys; j<info.gys+info.gym; j++) {
63     for (i=info.gxs; i<info.gxs+info.gxm; i++) {
64       for (d=0; d<2; d++) {
65         if ((i >= info.xs) && (i < info.xs+info.xm) && (j >= info.ys) && (j < info.ys+info.ym)) {
66           ierr = VecSetValuesLocal(Y,1,&k,&action[k],INSERT_VALUES);CHKERRQ(ierr);
67         }
68         k++;
69       }
70     }
71   }
72   ierr = PetscLogEventEnd(mctx->event1,0,0,0,0);CHKERRQ(ierr);
73   k = 0;
74   ierr = VecAssemblyBegin(Y);CHKERRQ(ierr); /* Note: Need to assemble between separate calls */
75   ierr = VecAssemblyEnd(Y);CHKERRQ(ierr);   /*       to INSERT_VALUES and ADD_VALUES         */
76 
77   /* a * dF/d(xdot) part */
78   ierr = PetscLogEventBegin(mctx->event2,0,0,0,0);CHKERRQ(ierr);
79   fos_forward(mctx->tag2,m,n,0,x0,x1,NULL,action);
80   for (j=info.gys; j<info.gys+info.gym; j++) {
81     for (i=info.gxs; i<info.gxs+info.gxm; i++) {
82       for (d=0; d<2; d++) {
83         if ((i >= info.xs) && (i < info.xs+info.xm) && (j >= info.ys) && (j < info.ys+info.ym)) {
84           action[k] *= mctx->shift;
85           ierr = VecSetValuesLocal(Y,1,&k,&action[k],ADD_VALUES);CHKERRQ(ierr);
86         }
87         k++;
88       }
89     }
90   }
91   ierr = PetscLogEventEnd(mctx->event2,0,0,0,0);CHKERRQ(ierr);
92   ierr = VecAssemblyBegin(Y);CHKERRQ(ierr);
93   ierr = VecAssemblyEnd(Y);CHKERRQ(ierr);
94   ierr = PetscFree(action);CHKERRQ(ierr);
95 
96   /* Restore local vector */
97   ierr = VecRestoreArray(localX1,&x1);CHKERRQ(ierr);
98   ierr = VecRestoreArrayRead(mctx->localX0,&x0);CHKERRQ(ierr);
99   ierr = DMRestoreLocalVector(da,&localX1);CHKERRQ(ierr);
100   PetscFunctionReturn(0);
101 }
102 
103 /*
104   ADOL-C implementation for Jacobian vector product, using the forward mode of AD.
105   Intended to overload MatMult in matrix-free methods where implicit timestepping
106   has been applied to a problem of the form
107                              du/dt = F(u).
108 
109   Input parameters:
110   A_shell - Jacobian matrix of MatShell type
111   X       - vector to be multiplied by A_shell
112 
113   Output parameters:
114   Y       - product of A_shell and X
115 */
PetscAdolcIJacobianVectorProductIDMass(Mat A_shell,Vec X,Vec Y)116 PetscErrorCode PetscAdolcIJacobianVectorProductIDMass(Mat A_shell,Vec X,Vec Y)
117 {
118   AdolcMatCtx       *mctx;
119   PetscErrorCode    ierr;
120   PetscInt          m,n,i,j,k = 0,d;
121   const PetscScalar *x0;
122   PetscScalar       *action,*x1;
123   Vec               localX1;
124   DM                da;
125   DMDALocalInfo     info;
126 
127   PetscFunctionBegin;
128 
129   /* Get matrix-free context info */
130   ierr = MatShellGetContext(A_shell,(void**)&mctx);CHKERRQ(ierr);
131   m = mctx->m;
132   n = mctx->n;
133 
134   /* Get local input vectors and extract data, x0 and x1*/
135   ierr = TSGetDM(mctx->ts,&da);CHKERRQ(ierr);
136   ierr = DMDAGetLocalInfo(da,&info);CHKERRQ(ierr);
137   ierr = DMGetLocalVector(da,&localX1);CHKERRQ(ierr);
138   ierr = DMGlobalToLocalBegin(da,X,INSERT_VALUES,localX1);CHKERRQ(ierr);
139   ierr = DMGlobalToLocalEnd(da,X,INSERT_VALUES,localX1);CHKERRQ(ierr);
140 
141   ierr = VecGetArrayRead(mctx->localX0,&x0);CHKERRQ(ierr);
142   ierr = VecGetArray(localX1,&x1);CHKERRQ(ierr);
143 
144   /* dF/dx part */
145   ierr = PetscMalloc1(m,&action);CHKERRQ(ierr);
146   ierr = PetscLogEventBegin(mctx->event1,0,0,0,0);CHKERRQ(ierr);
147   fos_forward(mctx->tag1,m,n,0,x0,x1,NULL,action);
148   for (j=info.gys; j<info.gys+info.gym; j++) {
149     for (i=info.gxs; i<info.gxs+info.gxm; i++) {
150       for (d=0; d<2; d++) {
151         if ((i >= info.xs) && (i < info.xs+info.xm) && (j >= info.ys) && (j < info.ys+info.ym)) {
152           ierr = VecSetValuesLocal(Y,1,&k,&action[k],INSERT_VALUES);CHKERRQ(ierr);
153         }
154         k++;
155       }
156     }
157   }
158   ierr = PetscLogEventEnd(mctx->event1,0,0,0,0);CHKERRQ(ierr);
159   k = 0;
160   ierr = VecAssemblyBegin(Y);CHKERRQ(ierr); /* Note: Need to assemble between separate calls */
161   ierr = VecAssemblyEnd(Y);CHKERRQ(ierr);   /*       to INSERT_VALUES and ADD_VALUES         */
162   ierr = PetscFree(action);CHKERRQ(ierr);
163 
164   /* Restore local vector */
165   ierr = VecRestoreArray(localX1,&x1);CHKERRQ(ierr);
166   ierr = VecRestoreArrayRead(mctx->localX0,&x0);CHKERRQ(ierr);
167   ierr = DMRestoreLocalVector(da,&localX1);CHKERRQ(ierr);
168 
169   /* a * dF/d(xdot) part */
170   ierr = PetscLogEventBegin(mctx->event2,0,0,0,0);CHKERRQ(ierr);
171   ierr = VecAXPY(Y,mctx->shift,X);
172   ierr = PetscLogEventEnd(mctx->event2,0,0,0,0);CHKERRQ(ierr);
173   PetscFunctionReturn(0);
174 }
175 
176 /*
177   ADOL-C implementation for Jacobian transpose vector product, using the reverse mode of AD.
178   Intended to overload MatMultTranspose in matrix-free methods where implicit timestepping
179   has been used.
180 
181   Input parameters:
182   A_shell - Jacobian matrix of MatShell type
183   Y       - vector to be multiplied by A_shell transpose
184 
185   Output parameters:
186   X       - product of A_shell transpose and X
187 */
PetscAdolcIJacobianTransposeVectorProduct(Mat A_shell,Vec Y,Vec X)188 PetscErrorCode PetscAdolcIJacobianTransposeVectorProduct(Mat A_shell,Vec Y,Vec X)
189 {
190   AdolcMatCtx       *mctx;
191   PetscErrorCode    ierr;
192   PetscInt          m,n,i,j,k = 0,d;
193   const PetscScalar *x;
194   PetscScalar       *action,*y;
195   Vec               localY;
196   DM                da;
197   DMDALocalInfo     info;
198 
199   PetscFunctionBegin;
200 
201   /* Get matrix-free context info */
202   ierr = MatShellGetContext(A_shell,(void**)&mctx);CHKERRQ(ierr);
203   m = mctx->m;
204   n = mctx->n;
205 
206   /* Get local input vectors and extract data, x0 and x1*/
207   ierr = TSGetDM(mctx->ts,&da);CHKERRQ(ierr);
208   ierr = DMDAGetLocalInfo(da,&info);CHKERRQ(ierr);
209   ierr = DMGetLocalVector(da,&localY);CHKERRQ(ierr);
210   ierr = DMGlobalToLocalBegin(da,Y,INSERT_VALUES,localY);CHKERRQ(ierr);
211   ierr = DMGlobalToLocalEnd(da,Y,INSERT_VALUES,localY);CHKERRQ(ierr);
212   ierr = VecGetArrayRead(mctx->localX0,&x);CHKERRQ(ierr);
213   ierr = VecGetArray(localY,&y);CHKERRQ(ierr);
214 
215   /* dF/dx part */
216   ierr = PetscMalloc1(n,&action);CHKERRQ(ierr);
217   ierr = PetscLogEventBegin(mctx->event3,0,0,0,0);CHKERRQ(ierr);
218   if (!mctx->flg)
219     zos_forward(mctx->tag1,m,n,1,x,NULL);
220   fos_reverse(mctx->tag1,m,n,y,action);
221   for (j=info.gys; j<info.gys+info.gym; j++) {
222     for (i=info.gxs; i<info.gxs+info.gxm; i++) {
223       for (d=0; d<2; d++) {
224         if ((i >= info.xs) && (i < info.xs+info.xm) && (j >= info.ys) && (j < info.ys+info.ym)) {
225           ierr = VecSetValuesLocal(X,1,&k,&action[k],INSERT_VALUES);CHKERRQ(ierr);
226         }
227         k++;
228       }
229     }
230   }
231   ierr = PetscLogEventEnd(mctx->event3,0,0,0,0);CHKERRQ(ierr);
232   k = 0;
233   ierr = VecAssemblyBegin(X);CHKERRQ(ierr); /* Note: Need to assemble between separate calls */
234   ierr = VecAssemblyEnd(X);CHKERRQ(ierr);   /*       to INSERT_VALUES and ADD_VALUES         */
235 
236   /* a * dF/d(xdot) part */
237   ierr = PetscLogEventBegin(mctx->event4,0,0,0,0);CHKERRQ(ierr);
238   if (!mctx->flg) {
239     zos_forward(mctx->tag2,m,n,1,x,NULL);
240     mctx->flg = PETSC_TRUE;
241   }
242   fos_reverse(mctx->tag2,m,n,y,action);
243   for (j=info.gys; j<info.gys+info.gym; j++) {
244     for (i=info.gxs; i<info.gxs+info.gxm; i++) {
245       for (d=0; d<2; d++) {
246         if ((i >= info.xs) && (i < info.xs+info.xm) && (j >= info.ys) && (j < info.ys+info.ym)) {
247           action[k] *= mctx->shift;
248           ierr = VecSetValuesLocal(X,1,&k,&action[k],ADD_VALUES);CHKERRQ(ierr);
249         }
250         k++;
251       }
252     }
253   }
254   ierr = PetscLogEventEnd(mctx->event4,0,0,0,0);CHKERRQ(ierr);
255   ierr = VecAssemblyBegin(X);CHKERRQ(ierr);
256   ierr = VecAssemblyEnd(X);CHKERRQ(ierr);
257   ierr = PetscFree(action);CHKERRQ(ierr);
258 
259   /* Restore local vector */
260   ierr = VecRestoreArray(localY,&y);CHKERRQ(ierr);
261   ierr = VecRestoreArrayRead(mctx->localX0,&x);CHKERRQ(ierr);
262   ierr = DMRestoreLocalVector(da,&localY);CHKERRQ(ierr);
263   PetscFunctionReturn(0);
264 }
265 
266 /*
267   ADOL-C implementation for Jacobian transpose vector product, using the reverse mode of AD.
268   Intended to overload MatMultTranspose in matrix-free methods where implicit timestepping
269   has been applied to a problem of the form
270                           du/dt = F(u).
271 
272   Input parameters:
273   A_shell - Jacobian matrix of MatShell type
274   Y       - vector to be multiplied by A_shell transpose
275 
276   Output parameters:
277   X       - product of A_shell transpose and X
278 */
PetscAdolcIJacobianTransposeVectorProductIDMass(Mat A_shell,Vec Y,Vec X)279 PetscErrorCode PetscAdolcIJacobianTransposeVectorProductIDMass(Mat A_shell,Vec Y,Vec X)
280 {
281   AdolcMatCtx       *mctx;
282   PetscErrorCode    ierr;
283   PetscInt          m,n,i,j,k = 0,d;
284   const PetscScalar *x;
285   PetscScalar       *action,*y;
286   Vec               localY;
287   DM                da;
288   DMDALocalInfo     info;
289 
290   PetscFunctionBegin;
291 
292   /* Get matrix-free context info */
293   ierr = MatShellGetContext(A_shell,(void**)&mctx);CHKERRQ(ierr);
294   m = mctx->m;
295   n = mctx->n;
296 
297   /* Get local input vectors and extract data, x0 and x1*/
298   ierr = TSGetDM(mctx->ts,&da);CHKERRQ(ierr);
299   ierr = DMDAGetLocalInfo(da,&info);CHKERRQ(ierr);
300   ierr = DMGetLocalVector(da,&localY);CHKERRQ(ierr);
301   ierr = DMGlobalToLocalBegin(da,Y,INSERT_VALUES,localY);CHKERRQ(ierr);
302   ierr = DMGlobalToLocalEnd(da,Y,INSERT_VALUES,localY);CHKERRQ(ierr);
303   ierr = VecGetArrayRead(mctx->localX0,&x);CHKERRQ(ierr);
304   ierr = VecGetArray(localY,&y);CHKERRQ(ierr);
305 
306   /* dF/dx part */
307   ierr = PetscMalloc1(n,&action);CHKERRQ(ierr);
308   ierr = PetscLogEventBegin(mctx->event3,0,0,0,0);CHKERRQ(ierr);
309   if (!mctx->flg) zos_forward(mctx->tag1,m,n,1,x,NULL);
310   fos_reverse(mctx->tag1,m,n,y,action);
311   for (j=info.gys; j<info.gys+info.gym; j++) {
312     for (i=info.gxs; i<info.gxs+info.gxm; i++) {
313       for (d=0; d<2; d++) {
314         if ((i >= info.xs) && (i < info.xs+info.xm) && (j >= info.ys) && (j < info.ys+info.ym)) {
315           ierr = VecSetValuesLocal(X,1,&k,&action[k],INSERT_VALUES);CHKERRQ(ierr);
316         }
317         k++;
318       }
319     }
320   }
321   ierr = PetscLogEventEnd(mctx->event3,0,0,0,0);CHKERRQ(ierr);
322   k = 0;
323   ierr = VecAssemblyBegin(X);CHKERRQ(ierr); /* Note: Need to assemble between separate calls */
324   ierr = VecAssemblyEnd(X);CHKERRQ(ierr);   /*       to INSERT_VALUES and ADD_VALUES         */
325   ierr = PetscFree(action);CHKERRQ(ierr);
326 
327   /* Restore local vector */
328   ierr = VecRestoreArray(localY,&y);CHKERRQ(ierr);
329   ierr = VecRestoreArrayRead(mctx->localX0,&x);CHKERRQ(ierr);
330   ierr = DMRestoreLocalVector(da,&localY);CHKERRQ(ierr);
331 
332   /* a * dF/d(xdot) part */
333   ierr = PetscLogEventBegin(mctx->event4,0,0,0,0);CHKERRQ(ierr);
334   ierr = VecAXPY(X,mctx->shift,Y);CHKERRQ(ierr);
335   ierr = PetscLogEventEnd(mctx->event4,0,0,0,0);CHKERRQ(ierr);
336   PetscFunctionReturn(0);
337 }
338