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