Actual source code: adamat.c

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

  3: PETSC_INTERN 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:   MatShellGetContext(mat, &ctx);
 22:   MatMult(ctx->A, a, ctx->W);
 23:   if (ctx->D1) VecPointwiseMult(ctx->W, ctx->D1, ctx->W);
 24:   MatMultTranspose(ctx->A, ctx->W, y);
 25:   if (ctx->D2) {
 26:     VecPointwiseMult(ctx->W2, ctx->D2, a);
 27:     VecAXPY(y, one, ctx->W2);
 28:   }
 29:   return 0;
 30: }

 32: static PetscErrorCode MatMultTranspose_ADA(Mat mat, Vec a, Vec y)
 33: {
 34:   MatMult_ADA(mat, a, y);
 35:   return 0;
 36: }

 38: static PetscErrorCode MatDiagonalSet_ADA(Mat M, Vec D, InsertMode mode)
 39: {
 40:   TaoMatADACtx ctx;
 41:   PetscReal    zero = 0.0, one = 1.0;

 44:   MatShellGetContext(M, &ctx);
 45:   if (!ctx->D2) {
 46:     VecDuplicate(D, &ctx->D2);
 47:     VecSet(ctx->D2, zero);
 48:   }
 49:   VecAXPY(ctx->D2, one, D);
 50:   return 0;
 51: }

 53: static PetscErrorCode MatDestroy_ADA(Mat mat)
 54: {
 55:   TaoMatADACtx ctx;

 57:   MatShellGetContext(mat, &ctx);
 58:   VecDestroy(&ctx->W);
 59:   VecDestroy(&ctx->W2);
 60:   VecDestroy(&ctx->ADADiag);
 61:   MatDestroy(&ctx->A);
 62:   VecDestroy(&ctx->D1);
 63:   VecDestroy(&ctx->D2);
 64:   PetscFree(ctx);
 65:   return 0;
 66: }

 68: static PetscErrorCode MatView_ADA(Mat mat, PetscViewer viewer)
 69: {
 70:   return 0;
 71: }

 73: static PetscErrorCode MatShift_ADA(Mat Y, PetscReal a)
 74: {
 75:   TaoMatADACtx ctx;

 77:   MatShellGetContext(Y, &ctx);
 78:   VecShift(ctx->D2, a);
 79:   return 0;
 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:   MatShellGetContext(mat, &ctx);
 89:   MatDuplicate(ctx->A, op, &A2);
 90:   if (ctx->D1) {
 91:     VecDuplicate(ctx->D1, &D1b);
 92:     VecCopy(ctx->D1, D1b);
 93:   }
 94:   VecDuplicate(ctx->D2, &D2b);
 95:   VecCopy(ctx->D2, D2b);
 96:   MatCreateADA(A2, D1b, D2b, M);
 97:   if (ctx->D1) PetscObjectDereference((PetscObject)D1b);
 98:   PetscObjectDereference((PetscObject)D2b);
 99:   PetscObjectDereference((PetscObject)A2);
100:   return 0;
101: }

103: static PetscErrorCode MatEqual_ADA(Mat A, Mat B, PetscBool *flg)
104: {
105:   TaoMatADACtx ctx1, ctx2;

107:   MatShellGetContext(A, &ctx1);
108:   MatShellGetContext(B, &ctx2);
109:   VecEqual(ctx1->D2, ctx2->D2, flg);
110:   if (*flg == PETSC_TRUE) VecEqual(ctx1->D1, ctx2->D1, flg);
111:   if (*flg == PETSC_TRUE) MatEqual(ctx1->A, ctx2->A, flg);
112:   return 0;
113: }

115: static PetscErrorCode MatScale_ADA(Mat mat, PetscReal a)
116: {
117:   TaoMatADACtx ctx;

119:   MatShellGetContext(mat, &ctx);
120:   VecScale(ctx->D1, a);
121:   if (ctx->D2) VecScale(ctx->D2, a);
122:   return 0;
123: }

125: static PetscErrorCode MatTranspose_ADA(Mat mat, MatReuse reuse, Mat *B)
126: {
127:   TaoMatADACtx ctx;

129:   if (reuse == MAT_REUSE_MATRIX) MatTransposeCheckNonzeroState_Private(mat, *B);
130:   MatShellGetContext(mat, &ctx);
131:   if (reuse == MAT_INITIAL_MATRIX) {
132:     MatDuplicate(mat, MAT_COPY_VALUES, B);
133:   } else if (reuse == MAT_REUSE_MATRIX) {
134:     MatCopy(mat, *B, SAME_NONZERO_PATTERN);
135:   } else SETERRQ(PetscObjectComm((PetscObject)mat), PETSC_ERR_SUP, "Does not support inplace transpose");
136:   return 0;
137: }

139: static PetscErrorCode MatADAComputeDiagonal(Mat mat)
140: {
141:   PetscInt     i, m, n, low, high;
142:   PetscScalar *dtemp, *dptr;
143:   TaoMatADACtx ctx;

145:   MatShellGetContext(mat, &ctx);
146:   MatGetOwnershipRange(mat, &low, &high);
147:   MatGetSize(mat, &m, &n);

149:   PetscMalloc1(n, &dtemp);
150:   for (i = 0; i < n; i++) {
151:     MatGetColumnVector(ctx->A, ctx->W, i);
152:     VecPointwiseMult(ctx->W, ctx->W, ctx->W);
153:     VecDotBegin(ctx->D1, ctx->W, dtemp + i);
154:   }
155:   for (i = 0; i < n; i++) VecDotEnd(ctx->D1, ctx->W, dtemp + i);

157:   VecGetArray(ctx->ADADiag, &dptr);
158:   for (i = low; i < high; i++) dptr[i - low] = dtemp[i];
159:   VecRestoreArray(ctx->ADADiag, &dptr);
160:   PetscFree(dtemp);
161:   return 0;
162: }

164: static PetscErrorCode MatGetDiagonal_ADA(Mat mat, Vec v)
165: {
166:   PetscReal    one = 1.0;
167:   TaoMatADACtx ctx;

169:   MatShellGetContext(mat, &ctx);
170:   MatADAComputeDiagonal(mat);
171:   VecCopy(ctx->ADADiag, v);
172:   if (ctx->D2) VecAXPY(v, one, ctx->D2);
173:   return 0;
174: }

176: static PetscErrorCode MatCreateSubMatrix_ADA(Mat mat, IS isrow, IS iscol, MatReuse cll, Mat *newmat)
177: {
178:   PetscInt     low, high;
179:   IS           ISrow;
180:   Vec          D1, D2;
181:   Mat          Atemp;
182:   TaoMatADACtx ctx;
183:   PetscBool    isequal;

185:   ISEqual(isrow, iscol, &isequal);
187:   MatShellGetContext(mat, &ctx);

189:   MatGetOwnershipRange(ctx->A, &low, &high);
190:   ISCreateStride(PetscObjectComm((PetscObject)mat), high - low, low, 1, &ISrow);
191:   MatCreateSubMatrix(ctx->A, ISrow, iscol, cll, &Atemp);
192:   ISDestroy(&ISrow);

194:   if (ctx->D1) {
195:     VecDuplicate(ctx->D1, &D1);
196:     VecCopy(ctx->D1, D1);
197:   } else {
198:     D1 = NULL;
199:   }

201:   if (ctx->D2) {
202:     Vec D2sub;

204:     VecGetSubVector(ctx->D2, isrow, &D2sub);
205:     VecDuplicate(D2sub, &D2);
206:     VecCopy(D2sub, D2);
207:     VecRestoreSubVector(ctx->D2, isrow, &D2sub);
208:   } else {
209:     D2 = NULL;
210:   }

212:   MatCreateADA(Atemp, D1, D2, newmat);
213:   MatShellGetContext(*newmat, &ctx);
214:   PetscObjectDereference((PetscObject)Atemp);
215:   if (ctx->D1) PetscObjectDereference((PetscObject)D1);
216:   if (ctx->D2) PetscObjectDereference((PetscObject)D2);
217:   return 0;
218: }

220: static PetscErrorCode MatCreateSubMatrices_ADA(Mat A, PetscInt n, IS *irow, IS *icol, MatReuse scall, Mat **B)
221: {
222:   PetscInt i;

224:   if (scall == MAT_INITIAL_MATRIX) PetscCalloc1(n + 1, B);
225:   for (i = 0; i < n; i++) MatCreateSubMatrix_ADA(A, irow[i], icol[i], scall, &(*B)[i]);
226:   return 0;
227: }

229: static PetscErrorCode MatGetColumnVector_ADA(Mat mat, Vec Y, PetscInt col)
230: {
231:   PetscInt    low, high;
232:   PetscScalar zero = 0.0, one = 1.0;

234:   VecSet(Y, zero);
235:   VecGetOwnershipRange(Y, &low, &high);
236:   if (col >= low && col < high) VecSetValue(Y, col, one, INSERT_VALUES);
237:   VecAssemblyBegin(Y);
238:   VecAssemblyEnd(Y);
239:   MatMult_ADA(mat, Y, Y);
240:   return 0;
241: }

243: PETSC_INTERN PetscErrorCode MatConvert_ADA(Mat mat, MatType newtype, Mat *NewMat)
244: {
245:   PetscMPIInt  size;
246:   PetscBool    sametype, issame, isdense, isseqdense;
247:   TaoMatADACtx ctx;

249:   MatShellGetContext(mat, &ctx);
250:   MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size);

252:   PetscObjectTypeCompare((PetscObject)mat, newtype, &sametype);
253:   PetscObjectTypeCompare((PetscObject)mat, MATSAME, &issame);
254:   PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSE, &isdense);
255:   PetscObjectTypeCompare((PetscObject)mat, MATSEQDENSE, &isseqdense);

257:   if (sametype || issame) {
258:     MatDuplicate(mat, MAT_COPY_VALUES, NewMat);
259:   } else if (isdense) {
260:     PetscInt           i, j, low, high, m, n, M, N;
261:     const PetscScalar *dptr;
262:     Vec                X;

264:     VecDuplicate(ctx->D2, &X);
265:     MatGetSize(mat, &M, &N);
266:     MatGetLocalSize(mat, &m, &n);
267:     MatCreateDense(PetscObjectComm((PetscObject)mat), m, m, N, N, NULL, NewMat);
268:     MatGetOwnershipRange(*NewMat, &low, &high);
269:     for (i = 0; i < M; i++) {
270:       MatGetColumnVector_ADA(mat, X, i);
271:       VecGetArrayRead(X, &dptr);
272:       for (j = 0; j < high - low; j++) MatSetValue(*NewMat, low + j, i, dptr[j], INSERT_VALUES);
273:       VecRestoreArrayRead(X, &dptr);
274:     }
275:     MatAssemblyBegin(*NewMat, MAT_FINAL_ASSEMBLY);
276:     MatAssemblyEnd(*NewMat, MAT_FINAL_ASSEMBLY);
277:     VecDestroy(&X);
278:   } else if (isseqdense && size == 1) {
279:     PetscInt           i, j, low, high, m, n, M, N;
280:     const PetscScalar *dptr;
281:     Vec                X;

283:     VecDuplicate(ctx->D2, &X);
284:     MatGetSize(mat, &M, &N);
285:     MatGetLocalSize(mat, &m, &n);
286:     MatCreateSeqDense(PetscObjectComm((PetscObject)mat), N, N, NULL, NewMat);
287:     MatGetOwnershipRange(*NewMat, &low, &high);
288:     for (i = 0; i < M; i++) {
289:       MatGetColumnVector_ADA(mat, X, i);
290:       VecGetArrayRead(X, &dptr);
291:       for (j = 0; j < high - low; j++) MatSetValue(*NewMat, low + j, i, dptr[j], INSERT_VALUES);
292:       VecRestoreArrayRead(X, &dptr);
293:     }
294:     MatAssemblyBegin(*NewMat, MAT_FINAL_ASSEMBLY);
295:     MatAssemblyEnd(*NewMat, MAT_FINAL_ASSEMBLY);
296:     VecDestroy(&X);
297:   } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "No support to convert objects to that type");
298:   return 0;
299: }

301: static PetscErrorCode MatNorm_ADA(Mat mat, NormType type, PetscReal *norm)
302: {
303:   TaoMatADACtx ctx;

305:   MatShellGetContext(mat, &ctx);
306:   if (type == NORM_FROBENIUS) {
307:     *norm = 1.0;
308:   } else if (type == NORM_1 || type == NORM_INFINITY) {
309:     *norm = 1.0;
310:   } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "No two norm");
311:   return 0;
312: }

314: /*@C
315:    MatCreateADA - Creates a matrix M=A^T D1 A + D2 where D1, D2 are diagonal

317:    Collective

319:    Input Parameters:
320: +  mat - matrix of arbitrary type
321: .  d1 - A vector defining a diagonal matrix
322: -  d2 - A vector defining a diagonal matrix

324:    Output Parameters:
325: .  J - New matrix whose operations are defined in terms of mat, D1, and D2.

327:    Level: developer

329:    Note:
330:    The user provides the input data and is responsible for destroying
331:    this data after matrix J has been destroyed.

333: .seealso: `Mat`, `MatCreate()`
334: @*/
335: PetscErrorCode MatCreateADA(Mat mat, Vec d1, Vec d2, Mat *J)
336: {
337:   MPI_Comm     comm = PetscObjectComm((PetscObject)mat);
338:   TaoMatADACtx ctx;
339:   PetscInt     nloc, n;

341:   PetscNew(&ctx);
342:   ctx->A  = mat;
343:   ctx->D1 = d1;
344:   ctx->D2 = d2;
345:   if (d1) {
346:     VecDuplicate(d1, &ctx->W);
347:     PetscObjectReference((PetscObject)d1);
348:   } else {
349:     ctx->W = NULL;
350:   }
351:   if (d2) {
352:     VecDuplicate(d2, &ctx->W2);
353:     VecDuplicate(d2, &ctx->ADADiag);
354:     PetscObjectReference((PetscObject)d2);
355:   } else {
356:     ctx->W2      = NULL;
357:     ctx->ADADiag = NULL;
358:   }

360:   ctx->GotDiag = 0;
361:   PetscObjectReference((PetscObject)mat);

363:   VecGetLocalSize(d2, &nloc);
364:   VecGetSize(d2, &n);

366:   MatCreateShell(comm, nloc, nloc, n, n, ctx, J);
367:   MatShellSetManageScalingShifts(*J);
368:   MatShellSetOperation(*J, MATOP_MULT, (void (*)(void))MatMult_ADA);
369:   MatShellSetOperation(*J, MATOP_DESTROY, (void (*)(void))MatDestroy_ADA);
370:   MatShellSetOperation(*J, MATOP_VIEW, (void (*)(void))MatView_ADA);
371:   MatShellSetOperation(*J, MATOP_MULT_TRANSPOSE, (void (*)(void))MatMultTranspose_ADA);
372:   MatShellSetOperation(*J, MATOP_DIAGONAL_SET, (void (*)(void))MatDiagonalSet_ADA);
373:   MatShellSetOperation(*J, MATOP_SHIFT, (void (*)(void))MatShift_ADA);
374:   MatShellSetOperation(*J, MATOP_EQUAL, (void (*)(void))MatEqual_ADA);
375:   MatShellSetOperation(*J, MATOP_SCALE, (void (*)(void))MatScale_ADA);
376:   MatShellSetOperation(*J, MATOP_TRANSPOSE, (void (*)(void))MatTranspose_ADA);
377:   MatShellSetOperation(*J, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiagonal_ADA);
378:   MatShellSetOperation(*J, MATOP_CREATE_SUBMATRICES, (void (*)(void))MatCreateSubMatrices_ADA);
379:   MatShellSetOperation(*J, MATOP_NORM, (void (*)(void))MatNorm_ADA);
380:   MatShellSetOperation(*J, MATOP_DUPLICATE, (void (*)(void))MatDuplicate_ADA);
381:   MatShellSetOperation(*J, MATOP_CREATE_SUBMATRIX, (void (*)(void))MatCreateSubMatrix_ADA);

383:   MatSetOption(*J, MAT_SYMMETRIC, PETSC_TRUE);
384:   return 0;
385: }