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) 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 PetscCall(MatShellGetContext(mat,&ctx)); 140 if (reuse == MAT_INITIAL_MATRIX) { 141 PetscCall(MatDuplicate(mat,MAT_COPY_VALUES,B)); 142 } else if (reuse == MAT_REUSE_MATRIX) { 143 PetscCall(MatCopy(mat,*B,SAME_NONZERO_PATTERN)); 144 } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"Does not support inplace transpose"); 145 PetscFunctionReturn(0); 146 } 147 148 static PetscErrorCode MatADAComputeDiagonal(Mat mat) 149 { 150 PetscInt i,m,n,low,high; 151 PetscScalar *dtemp,*dptr; 152 TaoMatADACtx ctx; 153 154 PetscFunctionBegin; 155 PetscCall(MatShellGetContext(mat,&ctx)); 156 PetscCall(MatGetOwnershipRange(mat, &low, &high)); 157 PetscCall(MatGetSize(mat,&m,&n)); 158 159 PetscCall(PetscMalloc1(n,&dtemp)); 160 for (i=0; i<n; i++) { 161 PetscCall(MatGetColumnVector(ctx->A, ctx->W, i)); 162 PetscCall(VecPointwiseMult(ctx->W,ctx->W,ctx->W)); 163 PetscCall(VecDotBegin(ctx->D1, ctx->W,dtemp+i)); 164 } 165 for (i=0; i<n; i++) { 166 PetscCall(VecDotEnd(ctx->D1, ctx->W,dtemp+i)); 167 } 168 169 PetscCall(VecGetArray(ctx->ADADiag,&dptr)); 170 for (i=low; i<high; i++) { 171 dptr[i-low]= dtemp[i]; 172 } 173 PetscCall(VecRestoreArray(ctx->ADADiag,&dptr)); 174 PetscCall(PetscFree(dtemp)); 175 PetscFunctionReturn(0); 176 } 177 178 static PetscErrorCode MatGetDiagonal_ADA(Mat mat,Vec v) 179 { 180 PetscReal one=1.0; 181 TaoMatADACtx ctx; 182 183 PetscFunctionBegin; 184 PetscCall(MatShellGetContext(mat,&ctx)); 185 PetscCall(MatADAComputeDiagonal(mat)); 186 PetscCall(VecCopy(ctx->ADADiag,v)); 187 if (ctx->D2) PetscCall(VecAXPY(v, one, ctx->D2)); 188 PetscFunctionReturn(0); 189 } 190 191 static PetscErrorCode MatCreateSubMatrix_ADA(Mat mat,IS isrow,IS iscol,MatReuse cll, Mat *newmat) 192 { 193 PetscInt low,high; 194 IS ISrow; 195 Vec D1,D2; 196 Mat Atemp; 197 TaoMatADACtx ctx; 198 PetscBool isequal; 199 200 PetscFunctionBegin; 201 PetscCall(ISEqual(isrow,iscol,&isequal)); 202 PetscCheck(isequal,PETSC_COMM_SELF,PETSC_ERR_SUP,"Only for identical column and row indices"); 203 PetscCall(MatShellGetContext(mat,&ctx)); 204 205 PetscCall(MatGetOwnershipRange(ctx->A,&low,&high)); 206 PetscCall(ISCreateStride(PetscObjectComm((PetscObject)mat),high-low,low,1,&ISrow)); 207 PetscCall(MatCreateSubMatrix(ctx->A,ISrow,iscol,cll,&Atemp)); 208 PetscCall(ISDestroy(&ISrow)); 209 210 if (ctx->D1) { 211 PetscCall(VecDuplicate(ctx->D1,&D1)); 212 PetscCall(VecCopy(ctx->D1,D1)); 213 } else { 214 D1 = NULL; 215 } 216 217 if (ctx->D2) { 218 Vec D2sub; 219 220 PetscCall(VecGetSubVector(ctx->D2,isrow,&D2sub)); 221 PetscCall(VecDuplicate(D2sub,&D2)); 222 PetscCall(VecCopy(D2sub,D2)); 223 PetscCall(VecRestoreSubVector(ctx->D2,isrow,&D2sub)); 224 } else { 225 D2 = NULL; 226 } 227 228 PetscCall(MatCreateADA(Atemp,D1,D2,newmat)); 229 PetscCall(MatShellGetContext(*newmat,&ctx)); 230 PetscCall(PetscObjectDereference((PetscObject)Atemp)); 231 if (ctx->D1) PetscCall(PetscObjectDereference((PetscObject)D1)); 232 if (ctx->D2) PetscCall(PetscObjectDereference((PetscObject)D2)); 233 PetscFunctionReturn(0); 234 } 235 236 static PetscErrorCode MatCreateSubMatrices_ADA(Mat A,PetscInt n, IS *irow,IS *icol,MatReuse scall,Mat **B) 237 { 238 PetscInt i; 239 240 PetscFunctionBegin; 241 if (scall == MAT_INITIAL_MATRIX) { 242 PetscCall(PetscCalloc1(n+1,B)); 243 } 244 for (i=0; i<n; i++) { 245 PetscCall(MatCreateSubMatrix_ADA(A,irow[i],icol[i],scall,&(*B)[i])); 246 } 247 PetscFunctionReturn(0); 248 } 249 250 static PetscErrorCode MatGetColumnVector_ADA(Mat mat,Vec Y, PetscInt col) 251 { 252 PetscInt low,high; 253 PetscScalar zero=0.0,one=1.0; 254 255 PetscFunctionBegin; 256 PetscCall(VecSet(Y, zero)); 257 PetscCall(VecGetOwnershipRange(Y,&low,&high)); 258 if (col>=low && col<high) { 259 PetscCall(VecSetValue(Y,col,one,INSERT_VALUES)); 260 } 261 PetscCall(VecAssemblyBegin(Y)); 262 PetscCall(VecAssemblyEnd(Y)); 263 PetscCall(MatMult_ADA(mat,Y,Y)); 264 PetscFunctionReturn(0); 265 } 266 267 PETSC_INTERN PetscErrorCode MatConvert_ADA(Mat mat,MatType newtype,Mat *NewMat) 268 { 269 PetscMPIInt size; 270 PetscBool sametype, issame, isdense, isseqdense; 271 TaoMatADACtx ctx; 272 273 PetscFunctionBegin; 274 PetscCall(MatShellGetContext(mat,&ctx)); 275 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size)); 276 277 PetscCall(PetscObjectTypeCompare((PetscObject)mat,newtype,&sametype)); 278 PetscCall(PetscObjectTypeCompare((PetscObject)mat,MATSAME,&issame)); 279 PetscCall(PetscObjectTypeCompare((PetscObject)mat,MATMPIDENSE,&isdense)); 280 PetscCall(PetscObjectTypeCompare((PetscObject)mat,MATSEQDENSE,&isseqdense)); 281 282 if (sametype || issame) { 283 PetscCall(MatDuplicate(mat,MAT_COPY_VALUES,NewMat)); 284 } else if (isdense) { 285 PetscInt i,j,low,high,m,n,M,N; 286 const PetscScalar *dptr; 287 Vec X; 288 289 PetscCall(VecDuplicate(ctx->D2,&X)); 290 PetscCall(MatGetSize(mat,&M,&N)); 291 PetscCall(MatGetLocalSize(mat,&m,&n)); 292 PetscCall(MatCreateDense(PetscObjectComm((PetscObject)mat),m,m,N,N,NULL,NewMat)); 293 PetscCall(MatGetOwnershipRange(*NewMat,&low,&high)); 294 for (i=0;i<M;i++) { 295 PetscCall(MatGetColumnVector_ADA(mat,X,i)); 296 PetscCall(VecGetArrayRead(X,&dptr)); 297 for (j=0; j<high-low; j++) { 298 PetscCall(MatSetValue(*NewMat,low+j,i,dptr[j],INSERT_VALUES)); 299 } 300 PetscCall(VecRestoreArrayRead(X,&dptr)); 301 } 302 PetscCall(MatAssemblyBegin(*NewMat,MAT_FINAL_ASSEMBLY)); 303 PetscCall(MatAssemblyEnd(*NewMat,MAT_FINAL_ASSEMBLY)); 304 PetscCall(VecDestroy(&X)); 305 } else if (isseqdense && size==1) { 306 PetscInt i,j,low,high,m,n,M,N; 307 const PetscScalar *dptr; 308 Vec X; 309 310 PetscCall(VecDuplicate(ctx->D2,&X)); 311 PetscCall(MatGetSize(mat,&M,&N)); 312 PetscCall(MatGetLocalSize(mat,&m,&n)); 313 PetscCall(MatCreateSeqDense(PetscObjectComm((PetscObject)mat),N,N,NULL,NewMat)); 314 PetscCall(MatGetOwnershipRange(*NewMat,&low,&high)); 315 for (i=0;i<M;i++) { 316 PetscCall(MatGetColumnVector_ADA(mat,X,i)); 317 PetscCall(VecGetArrayRead(X,&dptr)); 318 for (j=0; j<high-low; j++) { 319 PetscCall(MatSetValue(*NewMat,low+j,i,dptr[j],INSERT_VALUES)); 320 } 321 PetscCall(VecRestoreArrayRead(X,&dptr)); 322 } 323 PetscCall(MatAssemblyBegin(*NewMat,MAT_FINAL_ASSEMBLY)); 324 PetscCall(MatAssemblyEnd(*NewMat,MAT_FINAL_ASSEMBLY)); 325 PetscCall(VecDestroy(&X)); 326 } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"No support to convert objects to that type"); 327 PetscFunctionReturn(0); 328 } 329 330 static PetscErrorCode MatNorm_ADA(Mat mat,NormType type,PetscReal *norm) 331 { 332 TaoMatADACtx ctx; 333 334 PetscFunctionBegin; 335 PetscCall(MatShellGetContext(mat,&ctx)); 336 if (type == NORM_FROBENIUS) { 337 *norm = 1.0; 338 } else if (type == NORM_1 || type == NORM_INFINITY) { 339 *norm = 1.0; 340 } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"No two norm"); 341 PetscFunctionReturn(0); 342 } 343 344 /*@C 345 MatCreateADA - Creates a matrix M=A^T D1 A + D2 where D1, D2 are diagonal 346 347 Collective on matrix 348 349 Input Parameters: 350 + mat - matrix of arbitrary type 351 . d1 - A vector defining a diagonal matrix 352 - d2 - A vector defining a diagonal matrix 353 354 Output Parameters: 355 . J - New matrix whose operations are defined in terms of mat, D1, and D2. 356 357 Notes: 358 The user provides the input data and is responsible for destroying 359 this data after matrix J has been destroyed. 360 361 Level: developer 362 363 .seealso: `MatCreate()` 364 @*/ 365 PetscErrorCode MatCreateADA(Mat mat,Vec d1, Vec d2, Mat *J) 366 { 367 MPI_Comm comm = PetscObjectComm((PetscObject)mat); 368 TaoMatADACtx ctx; 369 PetscInt nloc,n; 370 371 PetscFunctionBegin; 372 PetscCall(PetscNew(&ctx)); 373 ctx->A=mat; 374 ctx->D1=d1; 375 ctx->D2=d2; 376 if (d1) { 377 PetscCall(VecDuplicate(d1,&ctx->W)); 378 PetscCall(PetscObjectReference((PetscObject)d1)); 379 } else { 380 ctx->W = NULL; 381 } 382 if (d2) { 383 PetscCall(VecDuplicate(d2,&ctx->W2)); 384 PetscCall(VecDuplicate(d2,&ctx->ADADiag)); 385 PetscCall(PetscObjectReference((PetscObject)d2)); 386 } else { 387 ctx->W2 = NULL; 388 ctx->ADADiag = NULL; 389 } 390 391 ctx->GotDiag = 0; 392 PetscCall(PetscObjectReference((PetscObject)mat)); 393 394 PetscCall(VecGetLocalSize(d2,&nloc)); 395 PetscCall(VecGetSize(d2,&n)); 396 397 PetscCall(MatCreateShell(comm,nloc,nloc,n,n,ctx,J)); 398 PetscCall(MatShellSetManageScalingShifts(*J)); 399 PetscCall(MatShellSetOperation(*J,MATOP_MULT,(void(*)(void))MatMult_ADA)); 400 PetscCall(MatShellSetOperation(*J,MATOP_DESTROY,(void(*)(void))MatDestroy_ADA)); 401 PetscCall(MatShellSetOperation(*J,MATOP_VIEW,(void(*)(void))MatView_ADA)); 402 PetscCall(MatShellSetOperation(*J,MATOP_MULT_TRANSPOSE,(void(*)(void))MatMultTranspose_ADA)); 403 PetscCall(MatShellSetOperation(*J,MATOP_DIAGONAL_SET,(void(*)(void))MatDiagonalSet_ADA)); 404 PetscCall(MatShellSetOperation(*J,MATOP_SHIFT,(void(*)(void))MatShift_ADA)); 405 PetscCall(MatShellSetOperation(*J,MATOP_EQUAL,(void(*)(void))MatEqual_ADA)); 406 PetscCall(MatShellSetOperation(*J,MATOP_SCALE,(void(*)(void))MatScale_ADA)); 407 PetscCall(MatShellSetOperation(*J,MATOP_TRANSPOSE,(void(*)(void))MatTranspose_ADA)); 408 PetscCall(MatShellSetOperation(*J,MATOP_GET_DIAGONAL,(void(*)(void))MatGetDiagonal_ADA)); 409 PetscCall(MatShellSetOperation(*J,MATOP_CREATE_SUBMATRICES,(void(*)(void))MatCreateSubMatrices_ADA)); 410 PetscCall(MatShellSetOperation(*J,MATOP_NORM,(void(*)(void))MatNorm_ADA)); 411 PetscCall(MatShellSetOperation(*J,MATOP_DUPLICATE,(void(*)(void))MatDuplicate_ADA)); 412 PetscCall(MatShellSetOperation(*J,MATOP_CREATE_SUBMATRIX,(void(*)(void))MatCreateSubMatrix_ADA)); 413 414 PetscCall(PetscLogObjectParent((PetscObject)(*J),(PetscObject)ctx->W)); 415 PetscCall(PetscLogObjectParent((PetscObject)mat,(PetscObject)(*J))); 416 417 PetscCall(MatSetOption(*J,MAT_SYMMETRIC,PETSC_TRUE)); 418 PetscFunctionReturn(0); 419 } 420