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 */ 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 */ 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 */ 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 */ 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