1 #include <petsc/private/matimpl.h> /*I "petscmat.h" I*/ 2 3 PETSC_INTERN PetscErrorCode MatCreateADA(Mat,Vec, Vec, Mat*); 4 5 typedef struct { 6 Mat A; 7 Vec D1; 8 Vec D2; 9 Vec W; 10 Vec W2; 11 Vec ADADiag; 12 PetscInt GotDiag; 13 } _p_TaoMatADACtx; 14 typedef _p_TaoMatADACtx* TaoMatADACtx; 15 16 static PetscErrorCode MatMult_ADA(Mat mat,Vec a,Vec y) 17 { 18 TaoMatADACtx ctx; 19 PetscReal one = 1.0; 20 21 PetscFunctionBegin; 22 PetscCall(MatShellGetContext(mat,&ctx)); 23 PetscCall(MatMult(ctx->A,a,ctx->W)); 24 if (ctx->D1) PetscCall(VecPointwiseMult(ctx->W,ctx->D1,ctx->W)); 25 PetscCall(MatMultTranspose(ctx->A,ctx->W,y)); 26 if (ctx->D2) { 27 PetscCall(VecPointwiseMult(ctx->W2, ctx->D2, a)); 28 PetscCall(VecAXPY(y, one, ctx->W2)); 29 } 30 PetscFunctionReturn(0); 31 } 32 33 static PetscErrorCode MatMultTranspose_ADA(Mat mat,Vec a,Vec y) 34 { 35 PetscFunctionBegin; 36 PetscCall(MatMult_ADA(mat,a,y)); 37 PetscFunctionReturn(0); 38 } 39 40 static PetscErrorCode MatDiagonalSet_ADA(Mat M,Vec D, InsertMode mode) 41 { 42 TaoMatADACtx ctx; 43 PetscReal zero=0.0,one = 1.0; 44 45 PetscFunctionBegin; 46 PetscCheck(mode != INSERT_VALUES,PetscObjectComm((PetscObject)M),PETSC_ERR_SUP,"Cannot insert diagonal entries of this matrix type, can only add"); 47 PetscCall(MatShellGetContext(M,&ctx)); 48 if (!ctx->D2) { 49 PetscCall(VecDuplicate(D,&ctx->D2)); 50 PetscCall(VecSet(ctx->D2, zero)); 51 } 52 PetscCall(VecAXPY(ctx->D2, one, D)); 53 PetscFunctionReturn(0); 54 } 55 56 static PetscErrorCode MatDestroy_ADA(Mat mat) 57 { 58 TaoMatADACtx ctx; 59 60 PetscFunctionBegin; 61 PetscCall(MatShellGetContext(mat,&ctx)); 62 PetscCall(VecDestroy(&ctx->W)); 63 PetscCall(VecDestroy(&ctx->W2)); 64 PetscCall(VecDestroy(&ctx->ADADiag)); 65 PetscCall(MatDestroy(&ctx->A)); 66 PetscCall(VecDestroy(&ctx->D1)); 67 PetscCall(VecDestroy(&ctx->D2)); 68 PetscCall(PetscFree(ctx)); 69 PetscFunctionReturn(0); 70 } 71 72 static PetscErrorCode MatView_ADA(Mat mat,PetscViewer viewer) 73 { 74 PetscFunctionBegin; 75 PetscFunctionReturn(0); 76 } 77 78 static PetscErrorCode MatShift_ADA(Mat Y, PetscReal a) 79 { 80 TaoMatADACtx ctx; 81 82 PetscFunctionBegin; 83 PetscCall(MatShellGetContext(Y,&ctx)); 84 PetscCall(VecShift(ctx->D2,a)); 85 PetscFunctionReturn(0); 86 } 87 88 static PetscErrorCode MatDuplicate_ADA(Mat mat,MatDuplicateOption op,Mat *M) 89 { 90 TaoMatADACtx ctx; 91 Mat A2; 92 Vec D1b=NULL,D2b; 93 94 PetscFunctionBegin; 95 PetscCall(MatShellGetContext(mat,&ctx)); 96 PetscCall(MatDuplicate(ctx->A,op,&A2)); 97 if (ctx->D1) { 98 PetscCall(VecDuplicate(ctx->D1,&D1b)); 99 PetscCall(VecCopy(ctx->D1,D1b)); 100 } 101 PetscCall(VecDuplicate(ctx->D2,&D2b)); 102 PetscCall(VecCopy(ctx->D2,D2b)); 103 PetscCall(MatCreateADA(A2,D1b,D2b,M)); 104 if (ctx->D1) PetscCall(PetscObjectDereference((PetscObject)D1b)); 105 PetscCall(PetscObjectDereference((PetscObject)D2b)); 106 PetscCall(PetscObjectDereference((PetscObject)A2)); 107 PetscFunctionReturn(0); 108 } 109 110 static PetscErrorCode MatEqual_ADA(Mat A,Mat B,PetscBool *flg) 111 { 112 TaoMatADACtx ctx1,ctx2; 113 114 PetscFunctionBegin; 115 PetscCall(MatShellGetContext(A,&ctx1)); 116 PetscCall(MatShellGetContext(B,&ctx2)); 117 PetscCall(VecEqual(ctx1->D2,ctx2->D2,flg)); 118 if (*flg==PETSC_TRUE) PetscCall(VecEqual(ctx1->D1,ctx2->D1,flg)); 119 if (*flg==PETSC_TRUE) PetscCall(MatEqual(ctx1->A,ctx2->A,flg)); 120 PetscFunctionReturn(0); 121 } 122 123 static PetscErrorCode MatScale_ADA(Mat mat, PetscReal a) 124 { 125 TaoMatADACtx ctx; 126 127 PetscFunctionBegin; 128 PetscCall(MatShellGetContext(mat,&ctx)); 129 PetscCall(VecScale(ctx->D1,a)); 130 if (ctx->D2) PetscCall(VecScale(ctx->D2,a)); 131 PetscFunctionReturn(0); 132 } 133 134 static PetscErrorCode MatTranspose_ADA(Mat mat,MatReuse reuse,Mat *B) 135 { 136 TaoMatADACtx ctx; 137 138 PetscFunctionBegin; 139 if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(mat,*B)); 140 PetscCall(MatShellGetContext(mat,&ctx)); 141 if (reuse == MAT_INITIAL_MATRIX) { 142 PetscCall(MatDuplicate(mat,MAT_COPY_VALUES,B)); 143 } else if (reuse == MAT_REUSE_MATRIX) { 144 PetscCall(MatCopy(mat,*B,SAME_NONZERO_PATTERN)); 145 } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"Does not support inplace transpose"); 146 PetscFunctionReturn(0); 147 } 148 149 static PetscErrorCode MatADAComputeDiagonal(Mat mat) 150 { 151 PetscInt i,m,n,low,high; 152 PetscScalar *dtemp,*dptr; 153 TaoMatADACtx ctx; 154 155 PetscFunctionBegin; 156 PetscCall(MatShellGetContext(mat,&ctx)); 157 PetscCall(MatGetOwnershipRange(mat, &low, &high)); 158 PetscCall(MatGetSize(mat,&m,&n)); 159 160 PetscCall(PetscMalloc1(n,&dtemp)); 161 for (i=0; i<n; i++) { 162 PetscCall(MatGetColumnVector(ctx->A, ctx->W, i)); 163 PetscCall(VecPointwiseMult(ctx->W,ctx->W,ctx->W)); 164 PetscCall(VecDotBegin(ctx->D1, ctx->W,dtemp+i)); 165 } 166 for (i=0; i<n; i++) { 167 PetscCall(VecDotEnd(ctx->D1, ctx->W,dtemp+i)); 168 } 169 170 PetscCall(VecGetArray(ctx->ADADiag,&dptr)); 171 for (i=low; i<high; i++) { 172 dptr[i-low]= dtemp[i]; 173 } 174 PetscCall(VecRestoreArray(ctx->ADADiag,&dptr)); 175 PetscCall(PetscFree(dtemp)); 176 PetscFunctionReturn(0); 177 } 178 179 static PetscErrorCode MatGetDiagonal_ADA(Mat mat,Vec v) 180 { 181 PetscReal one=1.0; 182 TaoMatADACtx ctx; 183 184 PetscFunctionBegin; 185 PetscCall(MatShellGetContext(mat,&ctx)); 186 PetscCall(MatADAComputeDiagonal(mat)); 187 PetscCall(VecCopy(ctx->ADADiag,v)); 188 if (ctx->D2) PetscCall(VecAXPY(v, one, ctx->D2)); 189 PetscFunctionReturn(0); 190 } 191 192 static PetscErrorCode MatCreateSubMatrix_ADA(Mat mat,IS isrow,IS iscol,MatReuse cll, Mat *newmat) 193 { 194 PetscInt low,high; 195 IS ISrow; 196 Vec D1,D2; 197 Mat Atemp; 198 TaoMatADACtx ctx; 199 PetscBool isequal; 200 201 PetscFunctionBegin; 202 PetscCall(ISEqual(isrow,iscol,&isequal)); 203 PetscCheck(isequal,PETSC_COMM_SELF,PETSC_ERR_SUP,"Only for identical column and row indices"); 204 PetscCall(MatShellGetContext(mat,&ctx)); 205 206 PetscCall(MatGetOwnershipRange(ctx->A,&low,&high)); 207 PetscCall(ISCreateStride(PetscObjectComm((PetscObject)mat),high-low,low,1,&ISrow)); 208 PetscCall(MatCreateSubMatrix(ctx->A,ISrow,iscol,cll,&Atemp)); 209 PetscCall(ISDestroy(&ISrow)); 210 211 if (ctx->D1) { 212 PetscCall(VecDuplicate(ctx->D1,&D1)); 213 PetscCall(VecCopy(ctx->D1,D1)); 214 } else { 215 D1 = NULL; 216 } 217 218 if (ctx->D2) { 219 Vec D2sub; 220 221 PetscCall(VecGetSubVector(ctx->D2,isrow,&D2sub)); 222 PetscCall(VecDuplicate(D2sub,&D2)); 223 PetscCall(VecCopy(D2sub,D2)); 224 PetscCall(VecRestoreSubVector(ctx->D2,isrow,&D2sub)); 225 } else { 226 D2 = NULL; 227 } 228 229 PetscCall(MatCreateADA(Atemp,D1,D2,newmat)); 230 PetscCall(MatShellGetContext(*newmat,&ctx)); 231 PetscCall(PetscObjectDereference((PetscObject)Atemp)); 232 if (ctx->D1) PetscCall(PetscObjectDereference((PetscObject)D1)); 233 if (ctx->D2) PetscCall(PetscObjectDereference((PetscObject)D2)); 234 PetscFunctionReturn(0); 235 } 236 237 static PetscErrorCode MatCreateSubMatrices_ADA(Mat A,PetscInt n, IS *irow,IS *icol,MatReuse scall,Mat **B) 238 { 239 PetscInt i; 240 241 PetscFunctionBegin; 242 if (scall == MAT_INITIAL_MATRIX) { 243 PetscCall(PetscCalloc1(n+1,B)); 244 } 245 for (i=0; i<n; i++) { 246 PetscCall(MatCreateSubMatrix_ADA(A,irow[i],icol[i],scall,&(*B)[i])); 247 } 248 PetscFunctionReturn(0); 249 } 250 251 static PetscErrorCode MatGetColumnVector_ADA(Mat mat,Vec Y, PetscInt col) 252 { 253 PetscInt low,high; 254 PetscScalar zero=0.0,one=1.0; 255 256 PetscFunctionBegin; 257 PetscCall(VecSet(Y, zero)); 258 PetscCall(VecGetOwnershipRange(Y,&low,&high)); 259 if (col>=low && col<high) { 260 PetscCall(VecSetValue(Y,col,one,INSERT_VALUES)); 261 } 262 PetscCall(VecAssemblyBegin(Y)); 263 PetscCall(VecAssemblyEnd(Y)); 264 PetscCall(MatMult_ADA(mat,Y,Y)); 265 PetscFunctionReturn(0); 266 } 267 268 PETSC_INTERN PetscErrorCode MatConvert_ADA(Mat mat,MatType newtype,Mat *NewMat) 269 { 270 PetscMPIInt size; 271 PetscBool sametype, issame, isdense, isseqdense; 272 TaoMatADACtx ctx; 273 274 PetscFunctionBegin; 275 PetscCall(MatShellGetContext(mat,&ctx)); 276 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size)); 277 278 PetscCall(PetscObjectTypeCompare((PetscObject)mat,newtype,&sametype)); 279 PetscCall(PetscObjectTypeCompare((PetscObject)mat,MATSAME,&issame)); 280 PetscCall(PetscObjectTypeCompare((PetscObject)mat,MATMPIDENSE,&isdense)); 281 PetscCall(PetscObjectTypeCompare((PetscObject)mat,MATSEQDENSE,&isseqdense)); 282 283 if (sametype || issame) { 284 PetscCall(MatDuplicate(mat,MAT_COPY_VALUES,NewMat)); 285 } else if (isdense) { 286 PetscInt i,j,low,high,m,n,M,N; 287 const PetscScalar *dptr; 288 Vec X; 289 290 PetscCall(VecDuplicate(ctx->D2,&X)); 291 PetscCall(MatGetSize(mat,&M,&N)); 292 PetscCall(MatGetLocalSize(mat,&m,&n)); 293 PetscCall(MatCreateDense(PetscObjectComm((PetscObject)mat),m,m,N,N,NULL,NewMat)); 294 PetscCall(MatGetOwnershipRange(*NewMat,&low,&high)); 295 for (i=0;i<M;i++) { 296 PetscCall(MatGetColumnVector_ADA(mat,X,i)); 297 PetscCall(VecGetArrayRead(X,&dptr)); 298 for (j=0; j<high-low; j++) { 299 PetscCall(MatSetValue(*NewMat,low+j,i,dptr[j],INSERT_VALUES)); 300 } 301 PetscCall(VecRestoreArrayRead(X,&dptr)); 302 } 303 PetscCall(MatAssemblyBegin(*NewMat,MAT_FINAL_ASSEMBLY)); 304 PetscCall(MatAssemblyEnd(*NewMat,MAT_FINAL_ASSEMBLY)); 305 PetscCall(VecDestroy(&X)); 306 } else if (isseqdense && size==1) { 307 PetscInt i,j,low,high,m,n,M,N; 308 const PetscScalar *dptr; 309 Vec X; 310 311 PetscCall(VecDuplicate(ctx->D2,&X)); 312 PetscCall(MatGetSize(mat,&M,&N)); 313 PetscCall(MatGetLocalSize(mat,&m,&n)); 314 PetscCall(MatCreateSeqDense(PetscObjectComm((PetscObject)mat),N,N,NULL,NewMat)); 315 PetscCall(MatGetOwnershipRange(*NewMat,&low,&high)); 316 for (i=0;i<M;i++) { 317 PetscCall(MatGetColumnVector_ADA(mat,X,i)); 318 PetscCall(VecGetArrayRead(X,&dptr)); 319 for (j=0; j<high-low; j++) { 320 PetscCall(MatSetValue(*NewMat,low+j,i,dptr[j],INSERT_VALUES)); 321 } 322 PetscCall(VecRestoreArrayRead(X,&dptr)); 323 } 324 PetscCall(MatAssemblyBegin(*NewMat,MAT_FINAL_ASSEMBLY)); 325 PetscCall(MatAssemblyEnd(*NewMat,MAT_FINAL_ASSEMBLY)); 326 PetscCall(VecDestroy(&X)); 327 } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"No support to convert objects to that type"); 328 PetscFunctionReturn(0); 329 } 330 331 static PetscErrorCode MatNorm_ADA(Mat mat,NormType type,PetscReal *norm) 332 { 333 TaoMatADACtx ctx; 334 335 PetscFunctionBegin; 336 PetscCall(MatShellGetContext(mat,&ctx)); 337 if (type == NORM_FROBENIUS) { 338 *norm = 1.0; 339 } else if (type == NORM_1 || type == NORM_INFINITY) { 340 *norm = 1.0; 341 } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"No two norm"); 342 PetscFunctionReturn(0); 343 } 344 345 /*@C 346 MatCreateADA - Creates a matrix M=A^T D1 A + D2 where D1, D2 are diagonal 347 348 Collective on matrix 349 350 Input Parameters: 351 + mat - matrix of arbitrary type 352 . d1 - A vector defining a diagonal matrix 353 - d2 - A vector defining a diagonal matrix 354 355 Output Parameters: 356 . J - New matrix whose operations are defined in terms of mat, D1, and D2. 357 358 Notes: 359 The user provides the input data and is responsible for destroying 360 this data after matrix J has been destroyed. 361 362 Level: developer 363 364 .seealso: `MatCreate()` 365 @*/ 366 PetscErrorCode MatCreateADA(Mat mat,Vec d1, Vec d2, Mat *J) 367 { 368 MPI_Comm comm = PetscObjectComm((PetscObject)mat); 369 TaoMatADACtx ctx; 370 PetscInt nloc,n; 371 372 PetscFunctionBegin; 373 PetscCall(PetscNew(&ctx)); 374 ctx->A=mat; 375 ctx->D1=d1; 376 ctx->D2=d2; 377 if (d1) { 378 PetscCall(VecDuplicate(d1,&ctx->W)); 379 PetscCall(PetscObjectReference((PetscObject)d1)); 380 } else { 381 ctx->W = NULL; 382 } 383 if (d2) { 384 PetscCall(VecDuplicate(d2,&ctx->W2)); 385 PetscCall(VecDuplicate(d2,&ctx->ADADiag)); 386 PetscCall(PetscObjectReference((PetscObject)d2)); 387 } else { 388 ctx->W2 = NULL; 389 ctx->ADADiag = NULL; 390 } 391 392 ctx->GotDiag = 0; 393 PetscCall(PetscObjectReference((PetscObject)mat)); 394 395 PetscCall(VecGetLocalSize(d2,&nloc)); 396 PetscCall(VecGetSize(d2,&n)); 397 398 PetscCall(MatCreateShell(comm,nloc,nloc,n,n,ctx,J)); 399 PetscCall(MatShellSetManageScalingShifts(*J)); 400 PetscCall(MatShellSetOperation(*J,MATOP_MULT,(void(*)(void))MatMult_ADA)); 401 PetscCall(MatShellSetOperation(*J,MATOP_DESTROY,(void(*)(void))MatDestroy_ADA)); 402 PetscCall(MatShellSetOperation(*J,MATOP_VIEW,(void(*)(void))MatView_ADA)); 403 PetscCall(MatShellSetOperation(*J,MATOP_MULT_TRANSPOSE,(void(*)(void))MatMultTranspose_ADA)); 404 PetscCall(MatShellSetOperation(*J,MATOP_DIAGONAL_SET,(void(*)(void))MatDiagonalSet_ADA)); 405 PetscCall(MatShellSetOperation(*J,MATOP_SHIFT,(void(*)(void))MatShift_ADA)); 406 PetscCall(MatShellSetOperation(*J,MATOP_EQUAL,(void(*)(void))MatEqual_ADA)); 407 PetscCall(MatShellSetOperation(*J,MATOP_SCALE,(void(*)(void))MatScale_ADA)); 408 PetscCall(MatShellSetOperation(*J,MATOP_TRANSPOSE,(void(*)(void))MatTranspose_ADA)); 409 PetscCall(MatShellSetOperation(*J,MATOP_GET_DIAGONAL,(void(*)(void))MatGetDiagonal_ADA)); 410 PetscCall(MatShellSetOperation(*J,MATOP_CREATE_SUBMATRICES,(void(*)(void))MatCreateSubMatrices_ADA)); 411 PetscCall(MatShellSetOperation(*J,MATOP_NORM,(void(*)(void))MatNorm_ADA)); 412 PetscCall(MatShellSetOperation(*J,MATOP_DUPLICATE,(void(*)(void))MatDuplicate_ADA)); 413 PetscCall(MatShellSetOperation(*J,MATOP_CREATE_SUBMATRIX,(void(*)(void))MatCreateSubMatrix_ADA)); 414 415 PetscCall(PetscLogObjectParent((PetscObject)(*J),(PetscObject)ctx->W)); 416 PetscCall(PetscLogObjectParent((PetscObject)mat,(PetscObject)(*J))); 417 418 PetscCall(MatSetOption(*J,MAT_SYMMETRIC,PETSC_TRUE)); 419 PetscFunctionReturn(0); 420 } 421