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