Actual source code: submat.c
petsc-3.10.5 2019-03-28
2: #include <petsc/private/matimpl.h>
4: typedef struct {
5: IS isrow,iscol; /* rows and columns in submatrix, only used to check consistency */
6: Vec left,right; /* optional scaling */
7: Vec olwork,orwork; /* work vectors outside the scatters, only touched by PreScale and only created if needed*/
8: Vec lwork,rwork; /* work vectors inside the scatters */
9: Vec dshift;
10: VecScatter lrestrict,rprolong;
11: Mat A;
12: PetscScalar vscale, axpy_vscale;
13: PetscScalar vshift, axpy_vshift;
14: } Mat_SubVirtual;
16: static PetscErrorCode PreScaleLeft(Mat N,Vec x,Vec *xx)
17: {
18: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
22: if (!Na->left) {
23: *xx = x;
24: } else {
25: if (!Na->olwork) {
26: VecDuplicate(Na->left,&Na->olwork);
27: }
28: VecPointwiseMult(Na->olwork,x,Na->left);
29: *xx = Na->olwork;
30: }
31: return(0);
32: }
34: static PetscErrorCode PreScaleRight(Mat N,Vec x,Vec *xx)
35: {
36: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
40: if (!Na->right) {
41: *xx = x;
42: } else {
43: if (!Na->orwork) {
44: VecDuplicate(Na->right,&Na->orwork);
45: }
46: VecPointwiseMult(Na->orwork,x,Na->right);
47: *xx = Na->orwork;
48: }
49: return(0);
50: }
52: static PetscErrorCode PostScaleLeft(Mat N,Vec x)
53: {
54: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
58: if (Na->left) {
59: VecPointwiseMult(x,x,Na->left);
60: }
61: return(0);
62: }
64: static PetscErrorCode PostScaleRight(Mat N,Vec x)
65: {
66: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
70: if (Na->right) {
71: VecPointwiseMult(x,x,Na->right);
72: }
73: return(0);
74: }
76: /*
77: Y = vscale*Y + diag(dshift)*X + vshift*X
79: On input Y already contains A*x
80: */
81: static PetscErrorCode MatSubmatShiftAndScale(Mat A,Vec X,Vec Y)
82: {
83: Mat_SubVirtual *Na = (Mat_SubVirtual*)A->data;
87: if (Na->dshift) { /* get arrays because there is no VecPointwiseMultAdd() */
88: PetscInt i,m;
89: const PetscScalar *x,*d;
90: PetscScalar *y;
91: VecGetLocalSize(X,&m);
92: VecGetArrayRead(Na->dshift,&d);
93: VecGetArrayRead(X,&x);
94: VecGetArray(Y,&y);
95: for (i=0; i<m; i++) y[i] = Na->vscale*y[i] + d[i]*x[i];
96: VecRestoreArrayRead(Na->dshift,&d);
97: VecRestoreArrayRead(X,&x);
98: VecRestoreArray(Y,&y);
99: } else {
100: VecScale(Y,Na->vscale);
101: }
102: if (Na->vshift != 0.0) {VecAXPY(Y,Na->vshift,X);} /* if test is for non-square matrices */
103: return(0);
104: }
106: static PetscErrorCode MatScale_SubMatrix(Mat N,PetscScalar a)
107: {
108: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
112: Na->vscale *= a;
113: Na->vshift *= a;
114: if (Na->dshift) {
115: VecScale(Na->dshift,a);
116: }
117: Na->axpy_vscale *= a;
118: return(0);
119: }
121: static PetscErrorCode MatShift_SubMatrix(Mat N,PetscScalar a)
122: {
123: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
127: if (Na->left || Na->right) {
128: if (!Na->dshift) {
129: VecDuplicate(Na->left ? Na->left : Na->right, &Na->dshift);
130: VecSet(Na->dshift,a);
131: } else {
132: if (Na->left) {VecPointwiseMult(Na->dshift,Na->dshift,Na->left);}
133: if (Na->right) {VecPointwiseMult(Na->dshift,Na->dshift,Na->right);}
134: VecShift(Na->dshift,a);
135: }
136: if (Na->left) {VecPointwiseDivide(Na->dshift,Na->dshift,Na->left);}
137: if (Na->right) {VecPointwiseDivide(Na->dshift,Na->dshift,Na->right);}
138: } else Na->vshift += a;
139: return(0);
140: }
142: static PetscErrorCode MatDiagonalScale_SubMatrix(Mat N,Vec left,Vec right)
143: {
144: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
148: if (left) {
149: if (!Na->left) {
150: VecDuplicate(left,&Na->left);
151: VecCopy(left,Na->left);
152: } else {
153: VecPointwiseMult(Na->left,left,Na->left);
154: }
155: }
156: if (right) {
157: if (!Na->right) {
158: VecDuplicate(right,&Na->right);
159: VecCopy(right,Na->right);
160: } else {
161: VecPointwiseMult(Na->right,right,Na->right);
162: }
163: }
164: return(0);
165: }
167: static PetscErrorCode MatMult_SubMatrix(Mat N,Vec x,Vec y)
168: {
169: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
170: Vec xx = 0;
174: PreScaleRight(N,x,&xx);
175: VecZeroEntries(Na->rwork);
176: VecScatterBegin(Na->rprolong,xx,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);
177: VecScatterEnd (Na->rprolong,xx,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);
178: MatMult(Na->A,Na->rwork,Na->lwork);
179: VecScatterBegin(Na->lrestrict,Na->lwork,y,INSERT_VALUES,SCATTER_FORWARD);
180: VecScatterEnd (Na->lrestrict,Na->lwork,y,INSERT_VALUES,SCATTER_FORWARD);
181: MatSubmatShiftAndScale(N,xx,y);
182: PostScaleLeft(N,y);
183: return(0);
184: }
186: static PetscErrorCode MatMultAdd_SubMatrix(Mat N,Vec v1,Vec v2,Vec v3)
187: {
188: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
189: Vec xx = 0;
193: PreScaleRight(N,v1,&xx);
194: VecZeroEntries(Na->rwork);
195: VecScatterBegin(Na->rprolong,xx,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);
196: VecScatterEnd (Na->rprolong,xx,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);
197: MatMult(Na->A,Na->rwork,Na->lwork);
198: if (v2 == v3) {
199: if (!Na->olwork) {VecDuplicate(v3,&Na->olwork);}
200: VecScatterBegin(Na->lrestrict,Na->lwork,Na->olwork,INSERT_VALUES,SCATTER_FORWARD);
201: VecScatterEnd (Na->lrestrict,Na->lwork,Na->olwork,INSERT_VALUES,SCATTER_FORWARD);
202: MatSubmatShiftAndScale(N,xx,Na->olwork);
203: PostScaleLeft(N,Na->olwork);
204: VecAXPY(v3,1.0,Na->olwork);
205: } else {
206: VecScatterBegin(Na->lrestrict,Na->lwork,v3,INSERT_VALUES,SCATTER_FORWARD);
207: VecScatterEnd (Na->lrestrict,Na->lwork,v3,INSERT_VALUES,SCATTER_FORWARD);
208: MatSubmatShiftAndScale(N,xx,v3);
209: PostScaleLeft(N,v3);
210: VecAXPY(v3,1.0,v2);
211: }
212: return(0);
213: }
215: static PetscErrorCode MatMultTranspose_SubMatrix(Mat N,Vec x,Vec y)
216: {
217: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
218: Vec xx = 0;
222: PreScaleLeft(N,x,&xx);
223: VecZeroEntries(Na->lwork);
224: VecScatterBegin(Na->lrestrict,xx,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);
225: VecScatterEnd (Na->lrestrict,xx,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);
226: MatMultTranspose(Na->A,Na->lwork,Na->rwork);
227: VecScatterBegin(Na->rprolong,Na->rwork,y,INSERT_VALUES,SCATTER_REVERSE);
228: VecScatterEnd (Na->rprolong,Na->rwork,y,INSERT_VALUES,SCATTER_REVERSE);
229: MatSubmatShiftAndScale(N,xx,y);
230: PostScaleRight(N,y);
231: return(0);
232: }
234: static PetscErrorCode MatMultTransposeAdd_SubMatrix(Mat N,Vec v1,Vec v2,Vec v3)
235: {
236: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
237: Vec xx = 0;
241: PreScaleLeft(N,v1,&xx);
242: VecZeroEntries(Na->lwork);
243: VecScatterBegin(Na->lrestrict,xx,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);
244: VecScatterEnd (Na->lrestrict,xx,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);
245: MatMultTranspose(Na->A,Na->lwork,Na->rwork);
246: if (v2 == v3) {
247: if (!Na->orwork) {VecDuplicate(v3,&Na->orwork);}
248: VecScatterBegin(Na->rprolong,Na->rwork,Na->orwork,INSERT_VALUES,SCATTER_REVERSE);
249: VecScatterEnd (Na->rprolong,Na->rwork,Na->orwork,INSERT_VALUES,SCATTER_REVERSE);
250: MatSubmatShiftAndScale(N,xx,Na->orwork);
251: PostScaleRight(N,Na->orwork);
252: VecAXPY(v3,1.0,Na->orwork);
253: } else {
254: VecScatterBegin(Na->rprolong,Na->rwork,v3,INSERT_VALUES,SCATTER_REVERSE);
255: VecScatterEnd (Na->rprolong,Na->rwork,v3,INSERT_VALUES,SCATTER_REVERSE);
256: MatSubmatShiftAndScale(N,xx,v3);
257: PostScaleRight(N,v3);
258: VecAXPY(v3,1.0,v2);
259: }
260: return(0);
261: }
263: static PetscErrorCode MatDestroy_SubMatrix(Mat N)
264: {
265: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
269: ISDestroy(&Na->isrow);
270: ISDestroy(&Na->iscol);
271: VecDestroy(&Na->left);
272: VecDestroy(&Na->right);
273: VecDestroy(&Na->olwork);
274: VecDestroy(&Na->orwork);
275: VecDestroy(&Na->lwork);
276: VecDestroy(&Na->rwork);
277: VecDestroy(&Na->dshift);
278: VecScatterDestroy(&Na->lrestrict);
279: VecScatterDestroy(&Na->rprolong);
280: MatDestroy(&Na->A);
281: PetscFree(N->data);
282: return(0);
283: }
285: /*@
286: MatCreateSubMatrixVirtual - Creates a virtual matrix that acts as a submatrix
288: Collective on Mat
290: Input Parameters:
291: + A - matrix that we will extract a submatrix of
292: . isrow - rows to be present in the submatrix
293: - iscol - columns to be present in the submatrix
295: Output Parameters:
296: . newmat - new matrix
298: Level: developer
300: Notes:
301: Most will use MatCreateSubMatrix which provides a more efficient representation if it is available.
303: .seealso: MatCreateSubMatrix(), MatSubMatrixVirtualUpdate()
304: @*/
305: PetscErrorCode MatCreateSubMatrixVirtual(Mat A,IS isrow,IS iscol,Mat *newmat)
306: {
307: Vec left,right;
308: PetscInt m,n;
309: Mat N;
310: Mat_SubVirtual *Na;
318: *newmat = 0;
320: MatCreate(PetscObjectComm((PetscObject)A),&N);
321: ISGetLocalSize(isrow,&m);
322: ISGetLocalSize(iscol,&n);
323: MatSetSizes(N,m,n,PETSC_DETERMINE,PETSC_DETERMINE);
324: PetscObjectChangeTypeName((PetscObject)N,MATSUBMATRIX);
326: PetscNewLog(N,&Na);
327: N->data = (void*)Na;
328: PetscObjectReference((PetscObject)A);
329: PetscObjectReference((PetscObject)isrow);
330: PetscObjectReference((PetscObject)iscol);
331: Na->A = A;
332: Na->isrow = isrow;
333: Na->iscol = iscol;
334: Na->vscale = 1.0;
335: Na->vshift = 0.0;
337: N->ops->destroy = MatDestroy_SubMatrix;
338: N->ops->mult = MatMult_SubMatrix;
339: N->ops->multadd = MatMultAdd_SubMatrix;
340: N->ops->multtranspose = MatMultTranspose_SubMatrix;
341: N->ops->multtransposeadd = MatMultTransposeAdd_SubMatrix;
342: N->ops->scale = MatScale_SubMatrix;
343: N->ops->diagonalscale = MatDiagonalScale_SubMatrix;
344: N->ops->shift = MatShift_SubMatrix;
346: MatSetBlockSizesFromMats(N,A,A);
347: PetscLayoutSetUp(N->rmap);
348: PetscLayoutSetUp(N->cmap);
350: MatCreateVecs(A,&Na->rwork,&Na->lwork);
351: VecCreate(PetscObjectComm((PetscObject)isrow),&left);
352: VecCreate(PetscObjectComm((PetscObject)iscol),&right);
353: VecSetSizes(left,m,PETSC_DETERMINE);
354: VecSetSizes(right,n,PETSC_DETERMINE);
355: VecSetUp(left);
356: VecSetUp(right);
357: VecScatterCreate(Na->lwork,isrow,left,NULL,&Na->lrestrict);
358: VecScatterCreate(right,NULL,Na->rwork,iscol,&Na->rprolong);
359: VecDestroy(&left);
360: VecDestroy(&right);
362: N->assembled = PETSC_TRUE;
364: MatSetUp(N);
366: *newmat = N;
367: return(0);
368: }
371: /*@
372: MatSubMatrixVirtualUpdate - Updates a submatrix
374: Collective on Mat
376: Input Parameters:
377: + N - submatrix to update
378: . A - full matrix in the submatrix
379: . isrow - rows in the update (same as the first time the submatrix was created)
380: - iscol - columns in the update (same as the first time the submatrix was created)
382: Level: developer
384: Notes:
385: Most will use MatCreateSubMatrix which provides a more efficient representation if it is available.
387: .seealso: MatCreateSubMatrixVirtual()
388: @*/
389: PetscErrorCode MatSubMatrixVirtualUpdate(Mat N,Mat A,IS isrow,IS iscol)
390: {
392: PetscBool flg;
393: Mat_SubVirtual *Na;
400: PetscObjectTypeCompare((PetscObject)N,MATSUBMATRIX,&flg);
401: if (!flg) SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_ARG_WRONG,"Matrix has wrong type");
403: Na = (Mat_SubVirtual*)N->data;
404: ISEqual(isrow,Na->isrow,&flg);
405: if (!flg) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"Cannot update submatrix with different row indices");
406: ISEqual(iscol,Na->iscol,&flg);
407: if (!flg) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"Cannot update submatrix with different column indices");
409: PetscObjectReference((PetscObject)A);
410: MatDestroy(&Na->A);
411: Na->A = A;
413: Na->vshift = 0.0;
414: Na->vscale = 1.0;
415: VecDestroy(&Na->left);
416: VecDestroy(&Na->right);
417: return(0);
418: }