1 #include <petsc/private/matimpl.h> /*I "petscmat.h" I*/
2
3 static 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
MatMult_ADA(Mat mat,Vec a,Vec y)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(PETSC_SUCCESS);
31 }
32
MatMultTranspose_ADA(Mat mat,Vec a,Vec y)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 }
39
MatDiagonalSet_ADA(Mat M,Vec D,InsertMode mode)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(PETSC_SUCCESS);
54 }
55
MatDestroy_ADA(Mat mat)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(PETSC_SUCCESS);
70 }
71
MatView_ADA(Mat mat,PetscViewer viewer)72 static PetscErrorCode MatView_ADA(Mat mat, PetscViewer viewer)
73 {
74 PetscFunctionBegin;
75 PetscFunctionReturn(PETSC_SUCCESS);
76 }
77
MatShift_ADA(Mat Y,PetscReal a)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(PETSC_SUCCESS);
86 }
87
MatDuplicate_ADA(Mat mat,MatDuplicateOption op,Mat * M)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(PETSC_SUCCESS);
108 }
109
MatEqual_ADA(Mat A,Mat B,PetscBool * flg)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(PETSC_SUCCESS);
121 }
122
MatScale_ADA(Mat mat,PetscReal a)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(PETSC_SUCCESS);
132 }
133
MatTranspose_ADA(Mat mat,MatReuse reuse,Mat * B)134 static PetscErrorCode MatTranspose_ADA(Mat mat, MatReuse reuse, Mat *B)
135 {
136 TaoMatADACtx ctx;
137
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 }
148
MatADAComputeDiagonal(Mat mat)149 static PetscErrorCode MatADAComputeDiagonal(Mat mat)
150 {
151 PetscInt i, m, n, low, high;
152 PetscScalar *dtemp, *dptr;
153 TaoMatADACtx ctx;
154
155 PetscFunctionBegin;
156 PetscCall(MatShellGetContext(mat, &ctx));
157 PetscCall(MatGetOwnershipRange(mat, &low, &high));
158 PetscCall(MatGetSize(mat, &m, &n));
159
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));
167
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 }
174
MatGetDiagonal_ADA(Mat mat,Vec v)175 static PetscErrorCode MatGetDiagonal_ADA(Mat mat, Vec v)
176 {
177 PetscReal one = 1.0;
178 TaoMatADACtx ctx;
179
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 }
187
MatCreateSubMatrix_ADA(Mat mat,IS isrow,IS iscol,MatReuse cll,Mat * newmat)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;
196
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));
201
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));
206
207 if (ctx->D1) {
208 PetscCall(VecDuplicate(ctx->D1, &D1));
209 PetscCall(VecCopy(ctx->D1, D1));
210 } else {
211 D1 = NULL;
212 }
213
214 if (ctx->D2) {
215 Vec D2sub;
216
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 }
224
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 }
232
MatCreateSubMatrices_ADA(Mat A,PetscInt n,IS * irow,IS * icol,MatReuse scall,Mat ** B)233 static PetscErrorCode MatCreateSubMatrices_ADA(Mat A, PetscInt n, IS *irow, IS *icol, MatReuse scall, Mat **B)
234 {
235 PetscInt i;
236
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 }
242
MatGetColumnVector_ADA(Mat mat,Vec Y,PetscInt col)243 static PetscErrorCode MatGetColumnVector_ADA(Mat mat, Vec Y, PetscInt col)
244 {
245 PetscInt low, high;
246 PetscScalar zero = 0.0, one = 1.0;
247
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 }
257
MatConvert_ADA(Mat mat,MatType newtype,Mat * NewMat)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;
263
264 PetscFunctionBegin;
265 PetscCall(MatShellGetContext(mat, &ctx));
266 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size));
267
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));
272
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;
279
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;
298
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 }
316
MatNorm_ADA(Mat mat,NormType type,PetscReal * norm)317 static PetscErrorCode MatNorm_ADA(Mat mat, NormType type, PetscReal *norm)
318 {
319 TaoMatADACtx ctx;
320
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 }
330
331 /*
332 MatCreateADA - Creates a matrix M=A^T D1 A + D2 where D1, D2 are diagonal
333
334 Collective
335
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
340
341 Output Parameter:
342 . J - New matrix whose operations are defined in terms of mat, D1, and D2.
343
344 Level: developer
345
346 Note:
347 The user provides the input data and is responsible for destroying
348 this data after matrix `J` has been destroyed.
349
350 .seealso: `Mat`, `MatCreate()`
351 */
MatCreateADA(Mat mat,Vec d1,Vec d2,Mat * J)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;
357
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 }
377
378 ctx->GotDiag = 0;
379 PetscCall(PetscObjectReference((PetscObject)mat));
380
381 PetscCall(VecGetLocalSize(d2, &nloc));
382 PetscCall(VecGetSize(d2, &n));
383
384 PetscCall(MatCreateShell(comm, nloc, nloc, n, n, ctx, J));
385 PetscCall(MatShellSetManageScalingShifts(*J));
386 PetscCall(MatShellSetOperation(*J, MATOP_MULT, (PetscErrorCodeFn *)MatMult_ADA));
387 PetscCall(MatShellSetOperation(*J, MATOP_DESTROY, (PetscErrorCodeFn *)MatDestroy_ADA));
388 PetscCall(MatShellSetOperation(*J, MATOP_VIEW, (PetscErrorCodeFn *)MatView_ADA));
389 PetscCall(MatShellSetOperation(*J, MATOP_MULT_TRANSPOSE, (PetscErrorCodeFn *)MatMultTranspose_ADA));
390 PetscCall(MatShellSetOperation(*J, MATOP_DIAGONAL_SET, (PetscErrorCodeFn *)MatDiagonalSet_ADA));
391 PetscCall(MatShellSetOperation(*J, MATOP_SHIFT, (PetscErrorCodeFn *)MatShift_ADA));
392 PetscCall(MatShellSetOperation(*J, MATOP_EQUAL, (PetscErrorCodeFn *)MatEqual_ADA));
393 PetscCall(MatShellSetOperation(*J, MATOP_SCALE, (PetscErrorCodeFn *)MatScale_ADA));
394 PetscCall(MatShellSetOperation(*J, MATOP_TRANSPOSE, (PetscErrorCodeFn *)MatTranspose_ADA));
395 PetscCall(MatShellSetOperation(*J, MATOP_GET_DIAGONAL, (PetscErrorCodeFn *)MatGetDiagonal_ADA));
396 PetscCall(MatShellSetOperation(*J, MATOP_CREATE_SUBMATRICES, (PetscErrorCodeFn *)MatCreateSubMatrices_ADA));
397 PetscCall(MatShellSetOperation(*J, MATOP_NORM, (PetscErrorCodeFn *)MatNorm_ADA));
398 PetscCall(MatShellSetOperation(*J, MATOP_DUPLICATE, (PetscErrorCodeFn *)MatDuplicate_ADA));
399 PetscCall(MatShellSetOperation(*J, MATOP_CREATE_SUBMATRIX, (PetscErrorCodeFn *)MatCreateSubMatrix_ADA));
400
401 PetscCall(MatSetOption(*J, MAT_SYMMETRIC, PETSC_TRUE));
402 PetscFunctionReturn(PETSC_SUCCESS);
403 }
404