Actual source code: adamat.c
1: #include <petsc/private/matimpl.h>
3: static PetscErrorCode MatCreateADA(Mat, Vec, Vec, Mat *);
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: } _n_TaoMatADACtx;
14: typedef _n_TaoMatADACtx *TaoMatADACtx;
16: static PetscErrorCode MatMult_ADA(Mat mat, Vec a, Vec y)
17: {
18: TaoMatADACtx ctx;
19: PetscReal one = 1.0;
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(PETSC_SUCCESS);
31: }
33: static PetscErrorCode MatMultTranspose_ADA(Mat mat, Vec a, Vec y)
34: {
35: PetscFunctionBegin;
36: PetscCall(MatMult_ADA(mat, a, y));
37: PetscFunctionReturn(PETSC_SUCCESS);
38: }
40: static PetscErrorCode MatDiagonalSet_ADA(Mat M, Vec D, InsertMode mode)
41: {
42: TaoMatADACtx ctx;
43: PetscReal zero = 0.0, one = 1.0;
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(PETSC_SUCCESS);
54: }
56: static PetscErrorCode MatDestroy_ADA(Mat mat)
57: {
58: TaoMatADACtx ctx;
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(PETSC_SUCCESS);
70: }
72: static PetscErrorCode MatShift_ADA(Mat Y, PetscReal a)
73: {
74: TaoMatADACtx ctx;
76: PetscFunctionBegin;
77: PetscCall(MatShellGetContext(Y, &ctx));
78: PetscCall(VecShift(ctx->D2, a));
79: PetscFunctionReturn(PETSC_SUCCESS);
80: }
82: static PetscErrorCode MatDuplicate_ADA(Mat mat, MatDuplicateOption op, Mat *M)
83: {
84: TaoMatADACtx ctx;
85: Mat A2;
86: Vec D1b = NULL, D2b;
88: PetscFunctionBegin;
89: PetscCall(MatShellGetContext(mat, &ctx));
90: PetscCall(MatDuplicate(ctx->A, op, &A2));
91: if (ctx->D1) {
92: PetscCall(VecDuplicate(ctx->D1, &D1b));
93: PetscCall(VecCopy(ctx->D1, D1b));
94: }
95: PetscCall(VecDuplicate(ctx->D2, &D2b));
96: PetscCall(VecCopy(ctx->D2, D2b));
97: PetscCall(MatCreateADA(A2, D1b, D2b, M));
98: if (ctx->D1) PetscCall(PetscObjectDereference((PetscObject)D1b));
99: PetscCall(PetscObjectDereference((PetscObject)D2b));
100: PetscCall(PetscObjectDereference((PetscObject)A2));
101: PetscFunctionReturn(PETSC_SUCCESS);
102: }
104: static PetscErrorCode MatEqual_ADA(Mat A, Mat B, PetscBool *flg)
105: {
106: TaoMatADACtx ctx1, ctx2;
108: PetscFunctionBegin;
109: PetscCall(MatShellGetContext(A, &ctx1));
110: PetscCall(MatShellGetContext(B, &ctx2));
111: PetscCall(VecEqual(ctx1->D2, ctx2->D2, flg));
112: if (*flg == PETSC_TRUE) PetscCall(VecEqual(ctx1->D1, ctx2->D1, flg));
113: if (*flg == PETSC_TRUE) PetscCall(MatEqual(ctx1->A, ctx2->A, flg));
114: PetscFunctionReturn(PETSC_SUCCESS);
115: }
117: static PetscErrorCode MatScale_ADA(Mat mat, PetscReal a)
118: {
119: TaoMatADACtx ctx;
121: PetscFunctionBegin;
122: PetscCall(MatShellGetContext(mat, &ctx));
123: PetscCall(VecScale(ctx->D1, a));
124: if (ctx->D2) PetscCall(VecScale(ctx->D2, a));
125: PetscFunctionReturn(PETSC_SUCCESS);
126: }
128: static PetscErrorCode MatTranspose_ADA(Mat mat, MatReuse reuse, Mat *B)
129: {
130: TaoMatADACtx ctx;
132: PetscFunctionBegin;
133: if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(mat, *B));
134: PetscCall(MatShellGetContext(mat, &ctx));
135: if (reuse == MAT_INITIAL_MATRIX) {
136: PetscCall(MatDuplicate(mat, MAT_COPY_VALUES, B));
137: } else if (reuse == MAT_REUSE_MATRIX) {
138: PetscCall(MatCopy(mat, *B, SAME_NONZERO_PATTERN));
139: } else SETERRQ(PetscObjectComm((PetscObject)mat), PETSC_ERR_SUP, "Does not support inplace transpose");
140: PetscFunctionReturn(PETSC_SUCCESS);
141: }
143: static PetscErrorCode MatADAComputeDiagonal(Mat mat)
144: {
145: PetscInt i, m, n, low, high;
146: PetscScalar *dtemp, *dptr;
147: TaoMatADACtx ctx;
149: PetscFunctionBegin;
150: PetscCall(MatShellGetContext(mat, &ctx));
151: PetscCall(MatGetOwnershipRange(mat, &low, &high));
152: PetscCall(MatGetSize(mat, &m, &n));
154: PetscCall(PetscMalloc1(n, &dtemp));
155: for (i = 0; i < n; i++) {
156: PetscCall(MatGetColumnVector(ctx->A, ctx->W, i));
157: PetscCall(VecPointwiseMult(ctx->W, ctx->W, ctx->W));
158: PetscCall(VecDotBegin(ctx->D1, ctx->W, dtemp + i));
159: }
160: for (i = 0; i < n; i++) PetscCall(VecDotEnd(ctx->D1, ctx->W, dtemp + i));
162: PetscCall(VecGetArray(ctx->ADADiag, &dptr));
163: for (i = low; i < high; i++) dptr[i - low] = dtemp[i];
164: PetscCall(VecRestoreArray(ctx->ADADiag, &dptr));
165: PetscCall(PetscFree(dtemp));
166: PetscFunctionReturn(PETSC_SUCCESS);
167: }
169: static PetscErrorCode MatGetDiagonal_ADA(Mat mat, Vec v)
170: {
171: PetscReal one = 1.0;
172: TaoMatADACtx ctx;
174: PetscFunctionBegin;
175: PetscCall(MatShellGetContext(mat, &ctx));
176: PetscCall(MatADAComputeDiagonal(mat));
177: PetscCall(VecCopy(ctx->ADADiag, v));
178: if (ctx->D2) PetscCall(VecAXPY(v, one, ctx->D2));
179: PetscFunctionReturn(PETSC_SUCCESS);
180: }
182: static PetscErrorCode MatCreateSubMatrix_ADA(Mat mat, IS isrow, IS iscol, MatReuse cll, Mat *newmat)
183: {
184: PetscInt low, high;
185: IS ISrow;
186: Vec D1, D2;
187: Mat Atemp;
188: TaoMatADACtx ctx;
189: PetscBool isequal;
191: PetscFunctionBegin;
192: PetscCall(ISEqual(isrow, iscol, &isequal));
193: PetscCheck(isequal, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only for identical column and row indices");
194: PetscCall(MatShellGetContext(mat, &ctx));
196: PetscCall(MatGetOwnershipRange(ctx->A, &low, &high));
197: PetscCall(ISCreateStride(PetscObjectComm((PetscObject)mat), high - low, low, 1, &ISrow));
198: PetscCall(MatCreateSubMatrix(ctx->A, ISrow, iscol, cll, &Atemp));
199: PetscCall(ISDestroy(&ISrow));
201: if (ctx->D1) {
202: PetscCall(VecDuplicate(ctx->D1, &D1));
203: PetscCall(VecCopy(ctx->D1, D1));
204: } else {
205: D1 = NULL;
206: }
208: if (ctx->D2) {
209: Vec D2sub;
211: PetscCall(VecGetSubVector(ctx->D2, isrow, &D2sub));
212: PetscCall(VecDuplicate(D2sub, &D2));
213: PetscCall(VecCopy(D2sub, D2));
214: PetscCall(VecRestoreSubVector(ctx->D2, isrow, &D2sub));
215: } else {
216: D2 = NULL;
217: }
219: PetscCall(MatCreateADA(Atemp, D1, D2, newmat));
220: PetscCall(MatShellGetContext(*newmat, &ctx));
221: PetscCall(PetscObjectDereference((PetscObject)Atemp));
222: if (ctx->D1) PetscCall(PetscObjectDereference((PetscObject)D1));
223: if (ctx->D2) PetscCall(PetscObjectDereference((PetscObject)D2));
224: PetscFunctionReturn(PETSC_SUCCESS);
225: }
227: static PetscErrorCode MatCreateSubMatrices_ADA(Mat A, PetscInt n, IS *irow, IS *icol, MatReuse scall, Mat **B)
228: {
229: PetscInt i;
231: PetscFunctionBegin;
232: if (scall == MAT_INITIAL_MATRIX) PetscCall(PetscCalloc1(n + 1, B));
233: for (i = 0; i < n; i++) PetscCall(MatCreateSubMatrix_ADA(A, irow[i], icol[i], scall, &(*B)[i]));
234: PetscFunctionReturn(PETSC_SUCCESS);
235: }
237: static PetscErrorCode MatGetColumnVector_ADA(Mat mat, Vec Y, PetscInt col)
238: {
239: PetscInt low, high;
240: PetscScalar zero = 0.0, one = 1.0;
242: PetscFunctionBegin;
243: PetscCall(VecSet(Y, zero));
244: PetscCall(VecGetOwnershipRange(Y, &low, &high));
245: if (col >= low && col < high) PetscCall(VecSetValue(Y, col, one, INSERT_VALUES));
246: PetscCall(VecAssemblyBegin(Y));
247: PetscCall(VecAssemblyEnd(Y));
248: PetscCall(MatMult_ADA(mat, Y, Y));
249: PetscFunctionReturn(PETSC_SUCCESS);
250: }
252: PETSC_INTERN PetscErrorCode MatConvert_ADA(Mat mat, MatType newtype, Mat *NewMat)
253: {
254: PetscMPIInt size;
255: PetscBool sametype, issame, isdense, isseqdense;
256: TaoMatADACtx ctx;
258: PetscFunctionBegin;
259: PetscCall(MatShellGetContext(mat, &ctx));
260: PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size));
262: PetscCall(PetscObjectTypeCompare((PetscObject)mat, newtype, &sametype));
263: PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATSAME, &issame));
264: PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSE, &isdense));
265: PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATSEQDENSE, &isseqdense));
267: if (sametype || issame) PetscCall(MatDuplicate(mat, MAT_COPY_VALUES, NewMat));
268: else if (isdense) {
269: PetscInt i, j, low, high, m, n, M, N;
270: const PetscScalar *dptr;
271: Vec X;
273: PetscCall(VecDuplicate(ctx->D2, &X));
274: PetscCall(MatGetSize(mat, &M, &N));
275: PetscCall(MatGetLocalSize(mat, &m, &n));
276: PetscCall(MatCreateDense(PetscObjectComm((PetscObject)mat), m, m, N, N, NULL, NewMat));
277: PetscCall(MatGetOwnershipRange(*NewMat, &low, &high));
278: for (i = 0; i < M; i++) {
279: PetscCall(MatGetColumnVector_ADA(mat, X, i));
280: PetscCall(VecGetArrayRead(X, &dptr));
281: for (j = 0; j < high - low; j++) PetscCall(MatSetValue(*NewMat, low + j, i, dptr[j], INSERT_VALUES));
282: PetscCall(VecRestoreArrayRead(X, &dptr));
283: }
284: PetscCall(MatAssemblyBegin(*NewMat, MAT_FINAL_ASSEMBLY));
285: PetscCall(MatAssemblyEnd(*NewMat, MAT_FINAL_ASSEMBLY));
286: PetscCall(VecDestroy(&X));
287: } else if (isseqdense && size == 1) {
288: PetscInt i, j, low, high, m, n, M, N;
289: const PetscScalar *dptr;
290: Vec X;
292: PetscCall(VecDuplicate(ctx->D2, &X));
293: PetscCall(MatGetSize(mat, &M, &N));
294: PetscCall(MatGetLocalSize(mat, &m, &n));
295: PetscCall(MatCreateSeqDense(PetscObjectComm((PetscObject)mat), N, N, NULL, NewMat));
296: PetscCall(MatGetOwnershipRange(*NewMat, &low, &high));
297: for (i = 0; i < M; i++) {
298: PetscCall(MatGetColumnVector_ADA(mat, X, i));
299: PetscCall(VecGetArrayRead(X, &dptr));
300: for (j = 0; j < high - low; j++) PetscCall(MatSetValue(*NewMat, low + j, i, dptr[j], INSERT_VALUES));
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 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "No support to convert objects to that type");
307: PetscFunctionReturn(PETSC_SUCCESS);
308: }
310: static PetscErrorCode MatNorm_ADA(Mat mat, NormType type, PetscReal *norm)
311: {
312: TaoMatADACtx ctx;
314: PetscFunctionBegin;
315: PetscCall(MatShellGetContext(mat, &ctx));
316: if (type == NORM_FROBENIUS) {
317: *norm = 1.0;
318: } else if (type == NORM_1 || type == NORM_INFINITY) {
319: *norm = 1.0;
320: } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "No two norm");
321: PetscFunctionReturn(PETSC_SUCCESS);
322: }
324: /*
325: MatCreateADA - Creates a matrix M=A^T D1 A + D2 where D1, D2 are diagonal
327: Collective
329: Input Parameters:
330: + mat - matrix of arbitrary type
331: . d1 - A vector defining a diagonal matrix
332: - d2 - A vector defining a diagonal matrix
334: Output Parameter:
335: . J - New matrix whose operations are defined in terms of mat, D1, and D2.
337: Level: developer
339: Note:
340: The user provides the input data and is responsible for destroying
341: this data after matrix `J` has been destroyed.
343: .seealso: `Mat`, `MatCreate()`
344: */
345: static PetscErrorCode MatCreateADA(Mat mat, Vec d1, Vec d2, Mat *J)
346: {
347: MPI_Comm comm = PetscObjectComm((PetscObject)mat);
348: TaoMatADACtx ctx;
349: PetscInt nloc, n;
351: PetscFunctionBegin;
352: PetscCall(PetscNew(&ctx));
353: ctx->A = mat;
354: ctx->D1 = d1;
355: ctx->D2 = d2;
356: if (d1) {
357: PetscCall(VecDuplicate(d1, &ctx->W));
358: PetscCall(PetscObjectReference((PetscObject)d1));
359: } else {
360: ctx->W = NULL;
361: }
362: if (d2) {
363: PetscCall(VecDuplicate(d2, &ctx->W2));
364: PetscCall(VecDuplicate(d2, &ctx->ADADiag));
365: PetscCall(PetscObjectReference((PetscObject)d2));
366: } else {
367: ctx->W2 = NULL;
368: ctx->ADADiag = NULL;
369: }
371: ctx->GotDiag = 0;
372: PetscCall(PetscObjectReference((PetscObject)mat));
374: PetscCall(VecGetLocalSize(d2, &nloc));
375: PetscCall(VecGetSize(d2, &n));
377: PetscCall(MatCreateShell(comm, nloc, nloc, n, n, ctx, J));
378: PetscCall(MatShellSetManageScalingShifts(*J));
379: PetscCall(MatShellSetOperation(*J, MATOP_MULT, (PetscErrorCodeFn *)MatMult_ADA));
380: PetscCall(MatShellSetOperation(*J, MATOP_DESTROY, (PetscErrorCodeFn *)MatDestroy_ADA));
381: PetscCall(MatShellSetOperation(*J, MATOP_MULT_TRANSPOSE, (PetscErrorCodeFn *)MatMultTranspose_ADA));
382: PetscCall(MatShellSetOperation(*J, MATOP_DIAGONAL_SET, (PetscErrorCodeFn *)MatDiagonalSet_ADA));
383: PetscCall(MatShellSetOperation(*J, MATOP_SHIFT, (PetscErrorCodeFn *)MatShift_ADA));
384: PetscCall(MatShellSetOperation(*J, MATOP_EQUAL, (PetscErrorCodeFn *)MatEqual_ADA));
385: PetscCall(MatShellSetOperation(*J, MATOP_SCALE, (PetscErrorCodeFn *)MatScale_ADA));
386: PetscCall(MatShellSetOperation(*J, MATOP_TRANSPOSE, (PetscErrorCodeFn *)MatTranspose_ADA));
387: PetscCall(MatShellSetOperation(*J, MATOP_GET_DIAGONAL, (PetscErrorCodeFn *)MatGetDiagonal_ADA));
388: PetscCall(MatShellSetOperation(*J, MATOP_CREATE_SUBMATRICES, (PetscErrorCodeFn *)MatCreateSubMatrices_ADA));
389: PetscCall(MatShellSetOperation(*J, MATOP_NORM, (PetscErrorCodeFn *)MatNorm_ADA));
390: PetscCall(MatShellSetOperation(*J, MATOP_DUPLICATE, (PetscErrorCodeFn *)MatDuplicate_ADA));
391: PetscCall(MatShellSetOperation(*J, MATOP_CREATE_SUBMATRIX, (PetscErrorCodeFn *)MatCreateSubMatrix_ADA));
393: PetscCall(MatSetOption(*J, MAT_SYMMETRIC, PETSC_TRUE));
394: PetscFunctionReturn(PETSC_SUCCESS);
395: }