Actual source code: diagonal.c

  1: #include <petsc/private/matimpl.h>

  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;

 13: static PetscErrorCode MatDiagonalSetUpDiagonal(Mat A)
 14: {
 15:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

 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: }

 27: static PetscErrorCode MatDiagonalSetUpInverseDiagonal(Mat A)
 28: {
 29:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

 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: }

 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;

 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: }

 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;

 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;

 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: }

 79: static PetscErrorCode MatMult_Diagonal(Mat A, Vec x, Vec y)
 80: {
 81:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

 83:   PetscFunctionBegin;
 84:   PetscCall(MatDiagonalSetUpDiagonal(A));
 85:   PetscCall(VecPointwiseMult(y, ctx->diag, x));
 86:   PetscFunctionReturn(PETSC_SUCCESS);
 87: }

 89: static PetscErrorCode MatMultAdd_Diagonal(Mat mat, Vec v1, Vec v2, Vec v3)
 90: {
 91:   Mat_Diagonal *ctx = (Mat_Diagonal *)mat->data;

 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: }

108: static PetscErrorCode MatNorm_Diagonal(Mat A, NormType type, PetscReal *nrm)
109: {
110:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

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: }

119: static PetscErrorCode MatDuplicate_Diagonal(Mat A, MatDuplicateOption op, Mat *B)
120: {
121:   Mat_Diagonal *actx = (Mat_Diagonal *)A->data;

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;

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: }

143: /*@
144:   MatDiagonalGetDiagonal - Get the diagonal of a `MATDIAGONAL`

146:   Input Parameter:
147: . A - the `MATDIAGONAL`

149:   Output Parameter:
150: . diag - the `Vec` that defines the diagonal

152:   Level: developer

154:   Note:
155:   The user must call
156:   `MatDiagonalRestoreDiagonal()` before using the matrix again.

158:   For a copy of the diagonal values, rather than a reference, use `MatGetDiagonal()`

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()`

162: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalRestoreDiagonal()`, `MatDiagonalGetInverseDiagonal()`, `MatGetDiagonal()`
163: @*/
164: PetscErrorCode MatDiagonalGetDiagonal(Mat A, Vec *diag)
165: {
166:   PetscFunctionBegin;
168:   PetscAssertPointer(diag, 2);
169:   *diag = NULL;
170:   PetscUseMethod(A, "MatDiagonalGetDiagonal_C", (Mat, Vec *), (A, diag));
171:   PetscFunctionReturn(PETSC_SUCCESS);
172: }

174: static PetscErrorCode MatDiagonalGetDiagonal_Diagonal(Mat A, Vec *diag)
175: {
176:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

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: }

186: /*@
187:   MatDiagonalRestoreDiagonal - Restore the diagonal of a `MATDIAGONAL`

189:   Input Parameters:
190: + A    - the `MATDIAGONAL`
191: - diag - the `Vec` obtained from `MatDiagonalGetDiagonal()`

193:   Level: developer

195:   Note:
196:   Use `MatDiagonalSet()` to change the values by copy, rather than reference.

198: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalGetDiagonal()`
199: @*/
200: PetscErrorCode MatDiagonalRestoreDiagonal(Mat A, Vec *diag)
201: {
202:   PetscFunctionBegin;
204:   PetscAssertPointer(diag, 2);
205:   PetscUseMethod(A, "MatDiagonalRestoreDiagonal_C", (Mat, Vec *), (A, diag));
206:   PetscFunctionReturn(PETSC_SUCCESS);
207: }

209: static PetscErrorCode MatDiagonalRestoreDiagonal_Diagonal(Mat A, Vec *diag)
210: {
211:   Mat_Diagonal    *ctx = (Mat_Diagonal *)A->data;
212:   PetscObjectState diag_state;

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: }

226: /*@
227:   MatDiagonalGetInverseDiagonal - Get the inverse diagonal of a `MATDIAGONAL`

229:   Input Parameter:
230: . A - the `MATDIAGONAL`

232:   Output Parameter:
233: . inv_diag - the `Vec` that defines the inverse diagonal

235:   Level: developer

237:   Note:
238:   The user must call
239:   `MatDiagonalRestoreInverseDiagonal()` before using the matrix again.

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()`.

245: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalRestoreInverseDiagonal()`, `MatDiagonalGetDiagonal()`, `MATLMVMBROYDEN`, `MatSolve()`
246: @*/
247: PetscErrorCode MatDiagonalGetInverseDiagonal(Mat A, Vec *inv_diag)
248: {
249:   PetscFunctionBegin;
251:   PetscAssertPointer(inv_diag, 2);
252:   *inv_diag = NULL;
253:   PetscUseMethod(A, "MatDiagonalGetInverseDiagonal_C", (Mat, Vec *), (A, inv_diag));
254:   PetscFunctionReturn(PETSC_SUCCESS);
255: }

257: static PetscErrorCode MatDiagonalGetInverseDiagonal_Diagonal(Mat A, Vec *inv_diag)
258: {
259:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

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: }

269: /*@
270:   MatDiagonalRestoreInverseDiagonal - Restore the inverse diagonal of a `MATDIAGONAL`

272:   Input Parameters:
273: + A        - the `MATDIAGONAL`
274: - inv_diag - the `Vec` obtained from `MatDiagonalGetInverseDiagonal()`

276:   Level: developer

278: .seealso: [](ch_matrices), `MATDIAGONAL`, `MatCreateDiagonal()`, `MatDiagonalGetInverseDiagonal()`
279: @*/
280: PetscErrorCode MatDiagonalRestoreInverseDiagonal(Mat A, Vec *inv_diag)
281: {
282:   PetscFunctionBegin;
284:   PetscAssertPointer(inv_diag, 2);
285:   PetscUseMethod(A, "MatDiagonalRestoreInverseDiagonal_C", (Mat, Vec *), (A, inv_diag));
286:   PetscFunctionReturn(PETSC_SUCCESS);
287: }

289: static PetscErrorCode MatDiagonalRestoreInverseDiagonal_Diagonal(Mat A, Vec *inv_diag)
290: {
291:   Mat_Diagonal    *ctx = (Mat_Diagonal *)A->data;
292:   PetscObjectState inv_diag_state;

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: }

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;

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: }

321: static PetscErrorCode MatSetRandom_Diagonal(Mat A, PetscRandom rand)
322: {
323:   Vec d;

325:   PetscFunctionBegin;
326:   PetscCall(MatDiagonalGetDiagonal(A, &d));
327:   PetscCall(VecSetRandom(d, rand));
328:   PetscCall(MatDiagonalRestoreDiagonal(A, &d));
329:   PetscFunctionReturn(PETSC_SUCCESS);
330: }

332: static PetscErrorCode MatDestroy_Diagonal(Mat mat)
333: {
334:   Mat_Diagonal *ctx = (Mat_Diagonal *)mat->data;

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: }

351: static PetscErrorCode MatView_Diagonal(Mat J, PetscViewer viewer)
352: {
353:   Mat_Diagonal *ctx = (Mat_Diagonal *)J->data;
354:   PetscBool     isascii;

356:   PetscFunctionBegin;
357:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
358:   if (isascii) {
359:     PetscViewerFormat format;

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: }

372: static PetscErrorCode MatGetDiagonal_Diagonal(Mat J, Vec x)
373: {
374:   Mat_Diagonal *ctx = (Mat_Diagonal *)J->data;

376:   PetscFunctionBegin;
377:   PetscCall(MatDiagonalSetUpDiagonal(J));
378:   PetscCall(VecCopy(ctx->diag, x));
379:   PetscFunctionReturn(PETSC_SUCCESS);
380: }

382: static PetscErrorCode MatDiagonalSet_Diagonal(Mat J, Vec D, InsertMode is)
383: {
384:   Mat_Diagonal *ctx = (Mat_Diagonal *)J->data;

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: }

419: static PetscErrorCode MatShift_Diagonal(Mat Y, PetscScalar a)
420: {
421:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

423:   PetscFunctionBegin;
424:   PetscCall(MatDiagonalSetUpDiagonal(Y));
425:   PetscCall(VecShift(ctx->diag, a));
426:   ctx->inv_diag_valid = PETSC_FALSE;
427:   PetscFunctionReturn(PETSC_SUCCESS);
428: }

430: static PetscErrorCode MatScale_Diagonal(Mat Y, PetscScalar a)
431: {
432:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

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: }

442: static PetscErrorCode MatDiagonalScale_Diagonal(Mat Y, Vec l, Vec r)
443: {
444:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

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: }

459: static PetscErrorCode MatConjugate_Diagonal(Mat Y)
460: {
461:   PetscFunctionBegin;
462:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

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: }

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;

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: }

503: static PetscErrorCode MatSetUp_Diagonal(Mat A)
504: {
505:   Mat_Diagonal *ctx = (Mat_Diagonal *)A->data;

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: }

520: static PetscErrorCode MatZeroEntries_Diagonal(Mat Y)
521: {
522:   Mat_Diagonal *ctx = (Mat_Diagonal *)Y->data;

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: }

532: static PetscErrorCode MatSolve_Diagonal(Mat matin, Vec b, Vec x)
533: {
534:   Mat_Diagonal *ctx = (Mat_Diagonal *)matin->data;

536:   PetscFunctionBegin;
537:   PetscCall(MatDiagonalSetUpInverseDiagonal(matin));
538:   PetscCall(VecPointwiseMult(x, b, ctx->inv_diag));
539:   PetscFunctionReturn(PETSC_SUCCESS);
540: }

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: }

558: /*@
559:   MatCreateDiagonal - Creates a matrix defined by a given vector along its diagonal.

561:   Collective

563:   Input Parameter:
564: . diag - vector for the diagonal

566:   Output Parameter:
567: . J - the diagonal matrix

569:   Level: advanced

571:   Notes:
572:   Only supports square matrices with the same number of local rows and columns.

574:   The input vector `diag` will be referenced internally: any changes to `diag`
575:   will affect the matrix `J`.

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;
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));

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: }

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;

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: }

639: static PetscErrorCode MatProductSymbolic_Diagonal_Dense(Mat C)
640: {
641:   Mat      A, B;
642:   PetscInt n, N, m, M;

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: }

665: static PetscErrorCode MatProductSetFromOptions_Diagonal_Dense_AB(Mat C)
666: {
667:   PetscFunctionBegin;
668:   C->ops->productsymbolic = MatProductSymbolic_Diagonal_Dense;
669:   PetscFunctionReturn(PETSC_SUCCESS);
670: }

672: static PetscErrorCode MatProductSetFromOptions_Diagonal_Dense(Mat C)
673: {
674:   Mat_Product *product = C->product;

676:   PetscFunctionBegin;
677:   if (product->type == MATPRODUCT_AB || product->type == MATPRODUCT_AtB) PetscCall(MatProductSetFromOptions_Diagonal_Dense_AB(C));
678:   PetscFunctionReturn(PETSC_SUCCESS);
679: }

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.

685:   Level: advanced

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;

693:   PetscFunctionBegin;
694:   PetscCall(PetscNew(&ctx));
695:   A->data = (void *)ctx;

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;

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;

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: }