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