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