1 #include <petsc/private/matimpl.h> /*I "petscmat.h" I*/ 2 3 typedef struct { 4 Vec diag; 5 PetscBool diag_valid; 6 Vec inv_diag; 7 PetscBool inv_diag_valid; 8 PetscObjectState diag_state, inv_diag_state; 9 PetscInt *col; 10 PetscScalar *val; 11 } Mat_Diagonal; 12 13 static PetscErrorCode MatDiagonalSetUpDiagonal(Mat A) 14 { 15 Mat_Diagonal *ctx = (Mat_Diagonal *)A->data; 16 17 PetscFunctionBegin; 18 if (!ctx->diag_valid) { 19 PetscAssert(ctx->inv_diag_valid, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Neither diagonal nor inverse diagonal is in a valid state"); 20 PetscCall(VecCopy(ctx->inv_diag, ctx->diag)); 21 PetscCall(VecReciprocal(ctx->diag)); 22 ctx->diag_valid = PETSC_TRUE; 23 } 24 PetscFunctionReturn(PETSC_SUCCESS); 25 } 26 27 static PetscErrorCode MatDiagonalSetUpInverseDiagonal(Mat A) 28 { 29 Mat_Diagonal *ctx = (Mat_Diagonal *)A->data; 30 31 PetscFunctionBegin; 32 if (!ctx->inv_diag_valid) { 33 PetscAssert(ctx->diag_valid, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Neither diagonal nor inverse diagonal is in a valid state"); 34 PetscCall(VecCopy(ctx->diag, ctx->inv_diag)); 35 PetscCall(VecReciprocal(ctx->inv_diag)); 36 ctx->inv_diag_valid = PETSC_TRUE; 37 } 38 PetscFunctionReturn(PETSC_SUCCESS); 39 } 40 41 static PetscErrorCode MatAXPY_Diagonal(Mat Y, PetscScalar a, Mat X, MatStructure str) 42 { 43 Mat_Diagonal *yctx = (Mat_Diagonal *)Y->data; 44 Mat_Diagonal *xctx = (Mat_Diagonal *)X->data; 45 46 PetscFunctionBegin; 47 PetscCall(MatDiagonalSetUpDiagonal(Y)); 48 PetscCall(MatDiagonalSetUpDiagonal(X)); 49 PetscCall(VecAXPY(yctx->diag, a, xctx->diag)); 50 yctx->inv_diag_valid = PETSC_FALSE; 51 PetscFunctionReturn(PETSC_SUCCESS); 52 } 53 54 static PetscErrorCode MatGetRow_Diagonal(Mat A, PetscInt row, PetscInt *ncols, PetscInt **cols, PetscScalar **vals) 55 { 56 Mat_Diagonal *mat = (Mat_Diagonal *)A->data; 57 PetscInt rstart = A->rmap->rstart, rend = A->rmap->rend; 58 59 PetscFunctionBegin; 60 PetscCheck(row >= rstart && row < rend, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Only local rows"); 61 if (ncols) *ncols = 1; 62 if (cols) { 63 if (!mat->col) PetscCall(PetscMalloc1(1, &mat->col)); 64 *mat->col = row; 65 *cols = mat->col; 66 } 67 if (vals) { 68 const PetscScalar *v; 69 70 if (!mat->val) PetscCall(PetscMalloc1(1, &mat->val)); 71 PetscCall(VecGetArrayRead(mat->diag, &v)); 72 *mat->val = v[row - rstart]; 73 *vals = mat->val; 74 PetscCall(VecRestoreArrayRead(mat->diag, &v)); 75 } 76 PetscFunctionReturn(PETSC_SUCCESS); 77 } 78 79 static PetscErrorCode MatMult_Diagonal(Mat A, Vec x, Vec y) 80 { 81 Mat_Diagonal *ctx = (Mat_Diagonal *)A->data; 82 83 PetscFunctionBegin; 84 PetscCall(MatDiagonalSetUpDiagonal(A)); 85 PetscCall(VecPointwiseMult(y, ctx->diag, x)); 86 PetscFunctionReturn(PETSC_SUCCESS); 87 } 88 89 static PetscErrorCode MatMultAdd_Diagonal(Mat mat, Vec v1, Vec v2, Vec v3) 90 { 91 Mat_Diagonal *ctx = (Mat_Diagonal *)mat->data; 92 93 PetscFunctionBegin; 94 PetscCall(MatDiagonalSetUpDiagonal(mat)); 95 if (v2 != v3) { 96 PetscCall(VecPointwiseMult(v3, ctx->diag, v1)); 97 PetscCall(VecAXPY(v3, 1.0, v2)); 98 } else { 99 Vec w; 100 PetscCall(VecDuplicate(v3, &w)); 101 PetscCall(VecPointwiseMult(w, ctx->diag, v1)); 102 PetscCall(VecAXPY(v3, 1.0, w)); 103 PetscCall(VecDestroy(&w)); 104 } 105 PetscFunctionReturn(PETSC_SUCCESS); 106 } 107 108 static PetscErrorCode MatNorm_Diagonal(Mat A, NormType type, PetscReal *nrm) 109 { 110 Mat_Diagonal *ctx = (Mat_Diagonal *)A->data; 111 112 PetscFunctionBegin; 113 PetscCall(MatDiagonalSetUpDiagonal(A)); 114 type = (type == NORM_FROBENIUS) ? NORM_2 : NORM_INFINITY; 115 PetscCall(VecNorm(ctx->diag, type, nrm)); 116 PetscFunctionReturn(PETSC_SUCCESS); 117 } 118 119 static PetscErrorCode MatDuplicate_Diagonal(Mat A, MatDuplicateOption op, Mat *B) 120 { 121 Mat_Diagonal *actx = (Mat_Diagonal *)A->data; 122 123 PetscFunctionBegin; 124 PetscCall(MatCreate(PetscObjectComm((PetscObject)A), B)); 125 PetscCall(MatSetSizes(*B, A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N)); 126 PetscCall(MatSetBlockSizesFromMats(*B, A, A)); 127 PetscCall(MatSetType(*B, MATDIAGONAL)); 128 PetscCall(PetscLayoutReference(A->rmap, &(*B)->rmap)); 129 PetscCall(PetscLayoutReference(A->cmap, &(*B)->cmap)); 130 if (op == MAT_COPY_VALUES) { 131 Mat_Diagonal *bctx = (Mat_Diagonal *)(*B)->data; 132 133 PetscCall(MatSetUp(A)); 134 PetscCall(MatSetUp(*B)); 135 PetscCall(VecCopy(actx->diag, bctx->diag)); 136 PetscCall(VecCopy(actx->inv_diag, bctx->inv_diag)); 137 bctx->diag_valid = actx->diag_valid; 138 bctx->inv_diag_valid = actx->inv_diag_valid; 139 } 140 PetscFunctionReturn(PETSC_SUCCESS); 141 } 142 143 /*@ 144 MatDiagonalGetDiagonal - Get the diagonal of a `MATDIAGONAL` 145 146 Input Parameter: 147 . A - the `MATDIAGONAL` 148 149 Output Parameter: 150 . diag - the `Vec` that defines the diagonal 151 152 Level: developer 153 154 Note: 155 The user must call 156 `MatDiagonalRestoreDiagonal()` before using the matrix again. 157 158 For a copy of the diagonal values, rather than a reference, use `MatGetDiagonal()` 159 160 Any changes to the obtained vector immediately change the action of the `Mat`. The matrix can be changed more efficiently by accessing this vector and changing its values, instead of filling a work vector and using `MatDiagonalSet()` 161 162 .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalRestoreDiagonal()`, `MatDiagonalGetInverseDiagonal()`, `MatGetDiagonal()` 163 @*/ 164 PetscErrorCode MatDiagonalGetDiagonal(Mat A, Vec *diag) 165 { 166 PetscFunctionBegin; 167 PetscValidHeaderSpecific(A, MAT_CLASSID, 1); 168 PetscAssertPointer(diag, 2); 169 *diag = NULL; 170 PetscUseMethod(A, "MatDiagonalGetDiagonal_C", (Mat, Vec *), (A, diag)); 171 PetscFunctionReturn(PETSC_SUCCESS); 172 } 173 174 static PetscErrorCode MatDiagonalGetDiagonal_Diagonal(Mat A, Vec *diag) 175 { 176 Mat_Diagonal *ctx = (Mat_Diagonal *)A->data; 177 178 PetscFunctionBegin; 179 PetscCall(MatSetUp(A)); 180 PetscCall(MatDiagonalSetUpDiagonal(A)); 181 *diag = ctx->diag; 182 PetscCall(PetscObjectStateGet((PetscObject)*diag, &ctx->diag_state)); 183 PetscFunctionReturn(PETSC_SUCCESS); 184 } 185 186 /*@ 187 MatDiagonalRestoreDiagonal - Restore the diagonal of a `MATDIAGONAL` 188 189 Input Parameters: 190 + A - the `MATDIAGONAL` 191 - diag - the `Vec` obtained from `MatDiagonalGetDiagonal()` 192 193 Level: developer 194 195 Note: 196 Use `MatDiagonalSet()` to change the values by copy, rather than reference. 197 198 .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalGetDiagonal()` 199 @*/ 200 PetscErrorCode MatDiagonalRestoreDiagonal(Mat A, Vec *diag) 201 { 202 PetscFunctionBegin; 203 PetscValidHeaderSpecific(A, MAT_CLASSID, 1); 204 PetscAssertPointer(diag, 2); 205 PetscUseMethod(A, "MatDiagonalRestoreDiagonal_C", (Mat, Vec *), (A, diag)); 206 PetscFunctionReturn(PETSC_SUCCESS); 207 } 208 209 static PetscErrorCode MatDiagonalRestoreDiagonal_Diagonal(Mat A, Vec *diag) 210 { 211 Mat_Diagonal *ctx = (Mat_Diagonal *)A->data; 212 PetscObjectState diag_state; 213 214 PetscFunctionBegin; 215 PetscCheck(ctx->diag == *diag, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Restored a different diagonal vector"); 216 ctx->diag_valid = PETSC_TRUE; 217 PetscCall(PetscObjectStateGet((PetscObject)*diag, &diag_state)); 218 if (ctx->diag_state != diag_state) { 219 PetscCall(PetscObjectStateIncrease((PetscObject)A)); 220 ctx->inv_diag_valid = PETSC_FALSE; 221 } 222 *diag = NULL; 223 PetscFunctionReturn(PETSC_SUCCESS); 224 } 225 226 /*@ 227 MatDiagonalGetInverseDiagonal - Get the inverse diagonal of a `MATDIAGONAL` 228 229 Input Parameter: 230 . A - the `MATDIAGONAL` 231 232 Output Parameter: 233 . inv_diag - the `Vec` that defines the inverse diagonal 234 235 Level: developer 236 237 Note: 238 The user must call 239 `MatDiagonalRestoreInverseDiagonal()` before using the matrix again. 240 241 If a matrix is created only to call `MatSolve()` (which happens for `MATLMVMDIAGBROYDEN`), 242 using `MatDiagonalGetInverseDiagonal()` and `MatDiagonalRestoreInverseDiagonal()` avoids copies 243 and avoids any call to `VecReciprocal()`. 244 245 .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MATLMVMBROYDEN`, `MatSolve()` 246 @*/ 247 PetscErrorCode MatDiagonalGetInverseDiagonal(Mat A, Vec *inv_diag) 248 { 249 PetscFunctionBegin; 250 PetscValidHeaderSpecific(A, MAT_CLASSID, 1); 251 PetscAssertPointer(inv_diag, 2); 252 *inv_diag = NULL; 253 PetscUseMethod(A, "MatDiagonalGetInverseDiagonal_C", (Mat, Vec *), (A, inv_diag)); 254 PetscFunctionReturn(PETSC_SUCCESS); 255 } 256 257 static PetscErrorCode MatDiagonalGetInverseDiagonal_Diagonal(Mat A, Vec *inv_diag) 258 { 259 Mat_Diagonal *ctx = (Mat_Diagonal *)A->data; 260 261 PetscFunctionBegin; 262 PetscCall(MatSetUp(A)); 263 PetscCall(MatDiagonalSetUpInverseDiagonal(A)); 264 *inv_diag = ctx->inv_diag; 265 PetscCall(PetscObjectStateGet((PetscObject)*inv_diag, &ctx->inv_diag_state)); 266 PetscFunctionReturn(PETSC_SUCCESS); 267 } 268 269 /*@ 270 MatDiagonalRestoreInverseDiagonal - Restore the inverse diagonal of a `MATDIAGONAL` 271 272 Input Parameters: 273 + A - the `MATDIAGONAL` 274 - inv_diag - the `Vec` obtained from `MatDiagonalGetInverseDiagonal()` 275 276 Level: developer 277 278 .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalGetInverseDiagonal()` 279 @*/ 280 PetscErrorCode MatDiagonalRestoreInverseDiagonal(Mat A, Vec *inv_diag) 281 { 282 PetscFunctionBegin; 283 PetscValidHeaderSpecific(A, MAT_CLASSID, 1); 284 PetscAssertPointer(inv_diag, 2); 285 PetscUseMethod(A, "MatDiagonalRestoreInverseDiagonal_C", (Mat, Vec *), (A, inv_diag)); 286 PetscFunctionReturn(PETSC_SUCCESS); 287 } 288 289 static PetscErrorCode MatDiagonalRestoreInverseDiagonal_Diagonal(Mat A, Vec *inv_diag) 290 { 291 Mat_Diagonal *ctx = (Mat_Diagonal *)A->data; 292 PetscObjectState inv_diag_state; 293 294 PetscFunctionBegin; 295 PetscCheck(ctx->inv_diag == *inv_diag, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Restored a different diagonal vector"); 296 ctx->inv_diag_valid = PETSC_TRUE; 297 PetscCall(PetscObjectStateGet((PetscObject)*inv_diag, &inv_diag_state)); 298 if (ctx->inv_diag_state != inv_diag_state) { 299 PetscCall(PetscObjectStateIncrease((PetscObject)A)); 300 ctx->diag_valid = PETSC_FALSE; 301 } 302 *inv_diag = NULL; 303 PetscFunctionReturn(PETSC_SUCCESS); 304 } 305 306 static PetscErrorCode MatPermute_Diagonal(Mat A, IS rowp, IS colp, Mat *B) 307 { 308 Mat_Diagonal *ctx = (Mat_Diagonal *)A->data; 309 Vec v; 310 311 PetscFunctionBegin; 312 PetscCheck(rowp == colp, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_INCOMP, "Row permutation and column permutation must be the same"); 313 PetscCall(VecDuplicate(ctx->diag, &v)); 314 PetscCall(VecCopy(ctx->diag, v)); 315 PetscCall(VecPermute(v, rowp, PETSC_FALSE)); 316 PetscCall(MatCreateDiagonal(v, B)); 317 PetscCall(VecDestroy(&v)); 318 PetscFunctionReturn(PETSC_SUCCESS); 319 } 320 321 static PetscErrorCode MatSetRandom_Diagonal(Mat A, PetscRandom rand) 322 { 323 Vec d; 324 325 PetscFunctionBegin; 326 PetscCall(MatDiagonalGetDiagonal(A, &d)); 327 PetscCall(VecSetRandom(d, rand)); 328 PetscCall(MatDiagonalRestoreDiagonal(A, &d)); 329 PetscFunctionReturn(PETSC_SUCCESS); 330 } 331 332 static PetscErrorCode MatDestroy_Diagonal(Mat mat) 333 { 334 Mat_Diagonal *ctx = (Mat_Diagonal *)mat->data; 335 336 PetscFunctionBegin; 337 PetscCall(VecDestroy(&ctx->diag)); 338 PetscCall(VecDestroy(&ctx->inv_diag)); 339 PetscCall(PetscFree(ctx->col)); 340 PetscCall(PetscFree(ctx->val)); 341 PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalGetDiagonal_C", NULL)); 342 PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalRestoreDiagonal_C", NULL)); 343 PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalGetInverseDiagonal_C", NULL)); 344 PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDiagonalRestoreInverseDiagonal_C", NULL)); 345 PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_diagonal_seqdense_C", NULL)); 346 PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_diagonal_mpidense_C", NULL)); 347 PetscCall(PetscFree(mat->data)); 348 PetscFunctionReturn(PETSC_SUCCESS); 349 } 350 351 static PetscErrorCode MatView_Diagonal(Mat J, PetscViewer viewer) 352 { 353 Mat_Diagonal *ctx = (Mat_Diagonal *)J->data; 354 PetscBool isascii; 355 356 PetscFunctionBegin; 357 PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii)); 358 if (isascii) { 359 PetscViewerFormat format; 360 361 PetscCall(MatDiagonalSetUpDiagonal(J)); 362 PetscCall(PetscViewerGetFormat(viewer, &format)); 363 if (format == PETSC_VIEWER_ASCII_FACTOR_INFO || format == PETSC_VIEWER_ASCII_INFO) { 364 PetscCall(PetscObjectPrintClassNamePrefixType((PetscObject)ctx->diag, viewer)); 365 PetscFunctionReturn(PETSC_SUCCESS); 366 } 367 PetscCall(VecView(ctx->diag, viewer)); 368 } 369 PetscFunctionReturn(PETSC_SUCCESS); 370 } 371 372 static PetscErrorCode MatGetDiagonal_Diagonal(Mat J, Vec x) 373 { 374 Mat_Diagonal *ctx = (Mat_Diagonal *)J->data; 375 376 PetscFunctionBegin; 377 PetscCall(MatDiagonalSetUpDiagonal(J)); 378 PetscCall(VecCopy(ctx->diag, x)); 379 PetscFunctionReturn(PETSC_SUCCESS); 380 } 381 382 static PetscErrorCode MatDiagonalSet_Diagonal(Mat J, Vec D, InsertMode is) 383 { 384 Mat_Diagonal *ctx = (Mat_Diagonal *)J->data; 385 386 PetscFunctionBegin; 387 switch (is) { 388 case ADD_VALUES: 389 case ADD_ALL_VALUES: 390 case ADD_BC_VALUES: 391 PetscCall(MatDiagonalSetUpDiagonal(J)); 392 PetscCall(VecAXPY(ctx->diag, 1.0, D)); 393 ctx->inv_diag_valid = PETSC_FALSE; 394 break; 395 case INSERT_VALUES: 396 case INSERT_BC_VALUES: 397 case INSERT_ALL_VALUES: 398 PetscCall(MatSetUp(J)); 399 PetscCall(VecCopy(D, ctx->diag)); 400 ctx->diag_valid = PETSC_TRUE; 401 ctx->inv_diag_valid = PETSC_FALSE; 402 break; 403 case MAX_VALUES: 404 PetscCall(MatDiagonalSetUpDiagonal(J)); 405 PetscCall(VecPointwiseMax(ctx->diag, D, ctx->diag)); 406 ctx->inv_diag_valid = PETSC_FALSE; 407 break; 408 case MIN_VALUES: 409 PetscCall(MatDiagonalSetUpDiagonal(J)); 410 PetscCall(VecPointwiseMin(ctx->diag, D, ctx->diag)); 411 ctx->inv_diag_valid = PETSC_FALSE; 412 break; 413 case NOT_SET_VALUES: 414 break; 415 } 416 PetscFunctionReturn(PETSC_SUCCESS); 417 } 418 419 static PetscErrorCode MatShift_Diagonal(Mat Y, PetscScalar a) 420 { 421 Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data; 422 423 PetscFunctionBegin; 424 PetscCall(MatDiagonalSetUpDiagonal(Y)); 425 PetscCall(VecShift(ctx->diag, a)); 426 ctx->inv_diag_valid = PETSC_FALSE; 427 PetscFunctionReturn(PETSC_SUCCESS); 428 } 429 430 static PetscErrorCode MatScale_Diagonal(Mat Y, PetscScalar a) 431 { 432 Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data; 433 434 PetscFunctionBegin; 435 PetscCall(MatSetUp(Y)); 436 PetscCall(MatDiagonalSetUpDiagonal(Y)); 437 PetscCall(VecScale(ctx->diag, a)); 438 ctx->inv_diag_valid = PETSC_FALSE; 439 PetscFunctionReturn(PETSC_SUCCESS); 440 } 441 442 static PetscErrorCode MatDiagonalScale_Diagonal(Mat Y, Vec l, Vec r) 443 { 444 Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data; 445 446 PetscFunctionBegin; 447 PetscCall(MatDiagonalSetUpDiagonal(Y)); 448 if (l) { 449 PetscCall(VecPointwiseMult(ctx->diag, ctx->diag, l)); 450 ctx->inv_diag_valid = PETSC_FALSE; 451 } 452 if (r) { 453 PetscCall(VecPointwiseMult(ctx->diag, ctx->diag, r)); 454 ctx->inv_diag_valid = PETSC_FALSE; 455 } 456 PetscFunctionReturn(PETSC_SUCCESS); 457 } 458 459 static PetscErrorCode MatConjugate_Diagonal(Mat Y) 460 { 461 PetscFunctionBegin; 462 Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data; 463 464 if (ctx->diag_valid) { 465 PetscCall(VecConjugate(ctx->diag)); 466 PetscCall(PetscObjectStateGet((PetscObject)ctx->diag, &ctx->diag_state)); 467 } 468 if (ctx->inv_diag_valid) { 469 PetscCall(VecConjugate(ctx->inv_diag)); 470 PetscCall(PetscObjectStateGet((PetscObject)ctx->inv_diag, &ctx->inv_diag_state)); 471 } 472 PetscFunctionReturn(PETSC_SUCCESS); 473 } 474 475 static PetscErrorCode MatTranspose_Diagonal(Mat A, MatReuse reuse, Mat *matout) 476 { 477 PetscFunctionBegin; 478 if (reuse == MAT_INPLACE_MATRIX) { 479 PetscLayout tmplayout = A->rmap; 480 481 A->rmap = A->cmap; 482 A->cmap = tmplayout; 483 } else { 484 Vec diag, newdiag; 485 if (reuse == MAT_INITIAL_MATRIX) { 486 PetscCall(MatDiagonalGetDiagonal(A, &diag)); 487 PetscCall(VecDuplicate(diag, &newdiag)); 488 PetscCall(VecCopy(diag, newdiag)); 489 PetscCall(MatDiagonalRestoreDiagonal(A, &diag)); 490 PetscCall(MatCreateDiagonal(newdiag, matout)); 491 PetscCall(VecDestroy(&newdiag)); 492 } else { 493 PetscCall(MatDiagonalGetDiagonal(A, &diag)); 494 PetscCall(MatDiagonalGetDiagonal(*matout, &newdiag)); 495 PetscCall(VecCopy(diag, newdiag)); 496 PetscCall(MatDiagonalRestoreDiagonal(*matout, &newdiag)); 497 PetscCall(MatDiagonalRestoreDiagonal(A, &diag)); 498 } 499 } 500 PetscFunctionReturn(PETSC_SUCCESS); 501 } 502 503 static PetscErrorCode MatSetUp_Diagonal(Mat A) 504 { 505 Mat_Diagonal *ctx = (Mat_Diagonal *)A->data; 506 507 PetscFunctionBegin; 508 if (!ctx->diag) { 509 PetscCall(PetscLayoutSetUp(A->cmap)); 510 PetscCall(PetscLayoutSetUp(A->rmap)); 511 PetscCall(MatCreateVecs(A, &ctx->diag, NULL)); 512 PetscCall(VecDuplicate(ctx->diag, &ctx->inv_diag)); 513 PetscCall(VecZeroEntries(ctx->diag)); 514 ctx->diag_valid = PETSC_TRUE; 515 } 516 A->assembled = PETSC_TRUE; 517 PetscFunctionReturn(PETSC_SUCCESS); 518 } 519 520 static PetscErrorCode MatZeroEntries_Diagonal(Mat Y) 521 { 522 Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data; 523 524 PetscFunctionBegin; 525 PetscCall(MatSetUp(Y)); 526 PetscCall(VecZeroEntries(ctx->diag)); 527 ctx->diag_valid = PETSC_TRUE; 528 ctx->inv_diag_valid = PETSC_FALSE; 529 PetscFunctionReturn(PETSC_SUCCESS); 530 } 531 532 static PetscErrorCode MatSolve_Diagonal(Mat matin, Vec b, Vec x) 533 { 534 Mat_Diagonal *ctx = (Mat_Diagonal *)matin->data; 535 536 PetscFunctionBegin; 537 PetscCall(MatDiagonalSetUpInverseDiagonal(matin)); 538 PetscCall(VecPointwiseMult(x, b, ctx->inv_diag)); 539 PetscFunctionReturn(PETSC_SUCCESS); 540 } 541 542 static PetscErrorCode MatGetInfo_Diagonal(Mat A, MatInfoType flag, MatInfo *info) 543 { 544 PetscFunctionBegin; 545 info->block_size = 1.0; 546 info->nz_allocated = A->cmap->N; 547 info->nz_used = A->cmap->N; 548 info->nz_unneeded = 0.0; 549 info->assemblies = A->num_ass; 550 info->mallocs = 0.0; 551 info->memory = 0; 552 info->fill_ratio_given = 0; 553 info->fill_ratio_needed = 0; 554 info->factor_mallocs = 0; 555 PetscFunctionReturn(PETSC_SUCCESS); 556 } 557 558 /*@ 559 MatCreateDiagonal - Creates a matrix defined by a given vector along its diagonal. 560 561 Collective 562 563 Input Parameter: 564 . diag - vector for the diagonal 565 566 Output Parameter: 567 . J - the diagonal matrix 568 569 Level: advanced 570 571 Notes: 572 Only supports square matrices with the same number of local rows and columns. 573 574 The input vector `diag` will be referenced internally: any changes to `diag` 575 will affect the matrix `J`. 576 577 .seealso: [](ch_matrices), `Mat`, `MatDestroy()`, `MATCONSTANTDIAGONAL`, `MatScale()`, `MatShift()`, `MatMult()`, `MatGetDiagonal()`, `MatSolve()` 578 `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MatDiagonalRestoreDiagonal()`, `MatDiagonalGetInverseDiagonal()` 579 @*/ 580 PetscErrorCode MatCreateDiagonal(Vec diag, Mat *J) 581 { 582 PetscFunctionBegin; 583 PetscValidHeaderSpecific(diag, VEC_CLASSID, 1); 584 PetscCall(MatCreate(PetscObjectComm((PetscObject)diag), J)); 585 PetscInt m, M; 586 PetscCall(VecGetLocalSize(diag, &m)); 587 PetscCall(VecGetSize(diag, &M)); 588 PetscCall(MatSetSizes(*J, m, m, M, M)); 589 PetscCall(MatSetType(*J, MATDIAGONAL)); 590 591 PetscLayout map; 592 PetscCall(VecGetLayout(diag, &map)); 593 PetscCall(MatSetLayouts(*J, map, map)); 594 Mat_Diagonal *ctx = (Mat_Diagonal *)(*J)->data; 595 PetscCall(PetscObjectReference((PetscObject)diag)); 596 PetscCall(VecDestroy(&ctx->diag)); 597 PetscCall(VecDestroy(&ctx->inv_diag)); 598 ctx->diag = diag; 599 ctx->diag_valid = PETSC_TRUE; 600 ctx->inv_diag_valid = PETSC_FALSE; 601 VecType type; 602 PetscCall(VecDuplicate(ctx->diag, &ctx->inv_diag)); 603 PetscCall(VecGetType(diag, &type)); 604 PetscCall(PetscFree((*J)->defaultvectype)); 605 PetscCall(PetscStrallocpy(type, &(*J)->defaultvectype)); 606 PetscCall(MatSetUp(*J)); 607 ctx->col = NULL; 608 ctx->val = NULL; 609 PetscFunctionReturn(PETSC_SUCCESS); 610 } 611 612 static PetscErrorCode MatProductNumeric_Diagonal_Dense(Mat C) 613 { 614 Mat A, B; 615 Mat_Diagonal *a; 616 PetscScalar *c; 617 const PetscScalar *b, *alpha; 618 PetscInt ldb, ldc; 619 620 PetscFunctionBegin; 621 MatCheckProduct(C, 1); 622 A = C->product->A; 623 B = C->product->B; 624 a = (Mat_Diagonal *)A->data; 625 PetscCall(VecGetArrayRead(a->diag, &alpha)); 626 PetscCall(MatDenseGetLDA(B, &ldb)); 627 PetscCall(MatDenseGetLDA(C, &ldc)); 628 PetscCall(MatDenseGetArrayRead(B, &b)); 629 PetscCall(MatDenseGetArrayWrite(C, &c)); 630 for (PetscInt j = 0; j < B->cmap->N; j++) 631 for (PetscInt i = 0; i < B->rmap->n; i++) c[i + j * ldc] = alpha[i] * b[i + j * ldb]; 632 PetscCall(MatDenseRestoreArrayWrite(C, &c)); 633 PetscCall(MatDenseRestoreArrayRead(B, &b)); 634 PetscCall(VecRestoreArrayRead(a->diag, &alpha)); 635 PetscCall(PetscLogFlops(B->cmap->N * B->rmap->n)); 636 PetscFunctionReturn(PETSC_SUCCESS); 637 } 638 639 static PetscErrorCode MatProductSymbolic_Diagonal_Dense(Mat C) 640 { 641 Mat A, B; 642 PetscInt n, N, m, M; 643 644 PetscFunctionBegin; 645 MatCheckProduct(C, 1); 646 PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty"); 647 A = C->product->A; 648 B = C->product->B; 649 PetscCall(MatDiagonalSetUpDiagonal(A)); 650 PetscCall(MatGetLocalSize(C, &m, &n)); 651 PetscCall(MatGetSize(C, &M, &N)); 652 if (m == PETSC_DECIDE || n == PETSC_DECIDE || M == PETSC_DECIDE || N == PETSC_DECIDE) { 653 PetscCall(MatGetLocalSize(B, NULL, &n)); 654 PetscCall(MatGetSize(B, NULL, &N)); 655 PetscCall(MatGetLocalSize(A, &m, NULL)); 656 PetscCall(MatGetSize(A, &M, NULL)); 657 PetscCall(MatSetSizes(C, m, n, M, N)); 658 } 659 PetscCall(MatSetType(C, ((PetscObject)B)->type_name)); 660 PetscCall(MatSetUp(C)); 661 C->ops->productnumeric = MatProductNumeric_Diagonal_Dense; 662 PetscFunctionReturn(PETSC_SUCCESS); 663 } 664 665 static PetscErrorCode MatProductSetFromOptions_Diagonal_Dense_AB(Mat C) 666 { 667 PetscFunctionBegin; 668 C->ops->productsymbolic = MatProductSymbolic_Diagonal_Dense; 669 PetscFunctionReturn(PETSC_SUCCESS); 670 } 671 672 static PetscErrorCode MatProductSetFromOptions_Diagonal_Dense(Mat C) 673 { 674 Mat_Product *product = C->product; 675 676 PetscFunctionBegin; 677 if (product->type == MATPRODUCT_AB || product->type == MATPRODUCT_AtB) PetscCall(MatProductSetFromOptions_Diagonal_Dense_AB(C)); 678 PetscFunctionReturn(PETSC_SUCCESS); 679 } 680 681 /*MC 682 MATDIAGONAL - MATDIAGONAL = "diagonal" - A diagonal matrix type with the diagonal implemented as a `Vec`. Useful for 683 cases where `VecPointwiseMult()` or `VecPointwiseDivide()` should be thought of as the actions of a linear operator. 684 685 Level: advanced 686 687 .seealso: [](ch_matrices), `Mat`, `MatCreateDiagonal()`, `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MatDiagonalRestoreDiagonal()`, `MatDiagonalGetInverseDiagonal()` 688 M*/ 689 PETSC_INTERN PetscErrorCode MatCreate_Diagonal(Mat A) 690 { 691 Mat_Diagonal *ctx; 692 693 PetscFunctionBegin; 694 PetscCall(PetscNew(&ctx)); 695 A->data = (void *)ctx; 696 697 A->structurally_symmetric = PETSC_BOOL3_TRUE; 698 A->structural_symmetry_eternal = PETSC_TRUE; 699 A->symmetry_eternal = PETSC_TRUE; 700 A->symmetric = PETSC_BOOL3_TRUE; 701 if (!PetscDefined(USE_COMPLEX)) A->hermitian = PETSC_BOOL3_TRUE; 702 703 A->ops->getrow = MatGetRow_Diagonal; 704 A->ops->mult = MatMult_Diagonal; 705 A->ops->multadd = MatMultAdd_Diagonal; 706 A->ops->multtranspose = MatMult_Diagonal; 707 A->ops->multtransposeadd = MatMultAdd_Diagonal; 708 A->ops->norm = MatNorm_Diagonal; 709 A->ops->duplicate = MatDuplicate_Diagonal; 710 A->ops->solve = MatSolve_Diagonal; 711 A->ops->solvetranspose = MatSolve_Diagonal; 712 A->ops->shift = MatShift_Diagonal; 713 A->ops->scale = MatScale_Diagonal; 714 A->ops->diagonalscale = MatDiagonalScale_Diagonal; 715 A->ops->getdiagonal = MatGetDiagonal_Diagonal; 716 A->ops->diagonalset = MatDiagonalSet_Diagonal; 717 A->ops->view = MatView_Diagonal; 718 A->ops->zeroentries = MatZeroEntries_Diagonal; 719 A->ops->destroy = MatDestroy_Diagonal; 720 A->ops->getinfo = MatGetInfo_Diagonal; 721 A->ops->axpy = MatAXPY_Diagonal; 722 A->ops->setup = MatSetUp_Diagonal; 723 A->ops->permute = MatPermute_Diagonal; 724 A->ops->setrandom = MatSetRandom_Diagonal; 725 A->ops->conjugate = MatConjugate_Diagonal; 726 A->ops->transpose = MatTranspose_Diagonal; 727 728 PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalGetDiagonal_C", MatDiagonalGetDiagonal_Diagonal)); 729 PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalRestoreDiagonal_C", MatDiagonalRestoreDiagonal_Diagonal)); 730 PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalGetInverseDiagonal_C", MatDiagonalGetInverseDiagonal_Diagonal)); 731 PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatDiagonalRestoreInverseDiagonal_C", MatDiagonalRestoreInverseDiagonal_Diagonal)); 732 PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_diagonal_seqdense_C", MatProductSetFromOptions_Diagonal_Dense)); 733 PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_diagonal_mpidense_C", MatProductSetFromOptions_Diagonal_Dense)); 734 PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATDIAGONAL)); 735 PetscFunctionReturn(PETSC_SUCCESS); 736 } 737