Actual source code: submat.c
petsc-3.14.6 2021-03-30
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 lwork,rwork; /* work vectors inside the scatters */
7: Vec lwork2,rwork2; /* work vectors inside the scatters */
8: VecScatter lrestrict,rprolong;
9: Mat A;
10: } Mat_SubVirtual;
12: static PetscErrorCode MatScale_SubMatrix(Mat N,PetscScalar a)
13: {
14: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
18: MatScale(Na->A,a);
19: return(0);
20: }
22: static PetscErrorCode MatShift_SubMatrix(Mat N,PetscScalar a)
23: {
24: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
28: MatShift(Na->A,a);
29: return(0);
30: }
32: static PetscErrorCode MatDiagonalScale_SubMatrix(Mat N,Vec left,Vec right)
33: {
34: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
38: if (right) {
39: VecZeroEntries(Na->rwork);
40: VecScatterBegin(Na->rprolong,right,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);
41: VecScatterEnd(Na->rprolong,right,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);
42: }
43: if (left) {
44: VecZeroEntries(Na->lwork);
45: VecScatterBegin(Na->lrestrict,left,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);
46: VecScatterEnd(Na->lrestrict,left,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);
47: }
48: MatDiagonalScale(Na->A,left ? Na->lwork : NULL,right ? Na->rwork : NULL);
49: return(0);
50: }
52: static PetscErrorCode MatGetDiagonal_SubMatrix(Mat N,Vec d)
53: {
54: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
58: MatGetDiagonal(Na->A,Na->rwork);
59: VecScatterBegin(Na->rprolong,Na->rwork,d,INSERT_VALUES,SCATTER_REVERSE);
60: VecScatterEnd(Na->rprolong,Na->rwork,d,INSERT_VALUES,SCATTER_REVERSE);
61: return(0);
62: }
64: static PetscErrorCode MatMult_SubMatrix(Mat N,Vec x,Vec y)
65: {
66: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
70: VecZeroEntries(Na->rwork);
71: VecScatterBegin(Na->rprolong,x,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);
72: VecScatterEnd(Na->rprolong,x,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);
73: MatMult(Na->A,Na->rwork,Na->lwork);
74: VecScatterBegin(Na->lrestrict,Na->lwork,y,INSERT_VALUES,SCATTER_FORWARD);
75: VecScatterEnd(Na->lrestrict,Na->lwork,y,INSERT_VALUES,SCATTER_FORWARD);
76: return(0);
77: }
79: static PetscErrorCode MatMultAdd_SubMatrix(Mat N,Vec v1,Vec v2,Vec v3)
80: {
81: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
85: VecZeroEntries(Na->rwork);
86: VecScatterBegin(Na->rprolong,v1,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);
87: VecScatterEnd(Na->rprolong,v1,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);
88: if (v1 == v2) {
89: MatMultAdd(Na->A,Na->rwork,Na->rwork,Na->lwork);
90: } else if (v2 == v3) {
91: VecZeroEntries(Na->lwork);
92: VecScatterBegin(Na->lrestrict,v2,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);
93: VecScatterEnd(Na->lrestrict,v2,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);
94: MatMultAdd(Na->A,Na->rwork,Na->lwork,Na->lwork);
95: } else {
96: if (!Na->lwork2) {
97: VecDuplicate(Na->lwork,&Na->lwork2);
98: } else {
99: VecZeroEntries(Na->lwork2);
100: }
101: VecScatterBegin(Na->lrestrict,v2,Na->lwork2,INSERT_VALUES,SCATTER_REVERSE);
102: VecScatterEnd(Na->lrestrict,v2,Na->lwork2,INSERT_VALUES,SCATTER_REVERSE);
103: MatMultAdd(Na->A,Na->rwork,Na->lwork2,Na->lwork);
104: }
105: VecScatterBegin(Na->lrestrict,Na->lwork,v3,INSERT_VALUES,SCATTER_FORWARD);
106: VecScatterEnd(Na->lrestrict,Na->lwork,v3,INSERT_VALUES,SCATTER_FORWARD);
107: return(0);
108: }
110: static PetscErrorCode MatMultTranspose_SubMatrix(Mat N,Vec x,Vec y)
111: {
112: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
116: VecZeroEntries(Na->lwork);
117: VecScatterBegin(Na->lrestrict,x,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);
118: VecScatterEnd(Na->lrestrict,x,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);
119: MatMultTranspose(Na->A,Na->lwork,Na->rwork);
120: VecScatterBegin(Na->rprolong,Na->rwork,y,INSERT_VALUES,SCATTER_REVERSE);
121: VecScatterEnd(Na->rprolong,Na->rwork,y,INSERT_VALUES,SCATTER_REVERSE);
122: return(0);
123: }
125: static PetscErrorCode MatMultTransposeAdd_SubMatrix(Mat N,Vec v1,Vec v2,Vec v3)
126: {
127: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
131: VecZeroEntries(Na->lwork);
132: VecScatterBegin(Na->lrestrict,v1,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);
133: VecScatterEnd(Na->lrestrict,v1,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);
134: if (v1 == v2) {
135: MatMultTransposeAdd(Na->A,Na->lwork,Na->lwork,Na->rwork);
136: } else if (v2 == v3) {
137: VecZeroEntries(Na->rwork);
138: VecScatterBegin(Na->rprolong,v2,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);
139: VecScatterEnd(Na->rprolong,v2,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);
140: MatMultTransposeAdd(Na->A,Na->lwork,Na->rwork,Na->rwork);
141: } else {
142: if (!Na->rwork2) {
143: VecDuplicate(Na->rwork,&Na->rwork2);
144: } else {
145: VecZeroEntries(Na->rwork2);
146: }
147: VecScatterBegin(Na->rprolong,v2,Na->rwork2,INSERT_VALUES,SCATTER_FORWARD);
148: VecScatterEnd(Na->rprolong,v2,Na->rwork2,INSERT_VALUES,SCATTER_FORWARD);
149: MatMultTransposeAdd(Na->A,Na->lwork,Na->rwork2,Na->rwork);
150: }
151: VecScatterBegin(Na->rprolong,Na->rwork,v3,INSERT_VALUES,SCATTER_REVERSE);
152: VecScatterEnd(Na->rprolong,Na->rwork,v3,INSERT_VALUES,SCATTER_REVERSE);
153: return(0);
154: }
156: static PetscErrorCode MatDestroy_SubMatrix(Mat N)
157: {
158: Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
162: ISDestroy(&Na->isrow);
163: ISDestroy(&Na->iscol);
164: VecDestroy(&Na->lwork);
165: VecDestroy(&Na->rwork);
166: VecDestroy(&Na->lwork2);
167: VecDestroy(&Na->rwork2);
168: VecScatterDestroy(&Na->lrestrict);
169: VecScatterDestroy(&Na->rprolong);
170: MatDestroy(&Na->A);
171: PetscFree(N->data);
172: return(0);
173: }
175: /*@
176: MatCreateSubMatrixVirtual - Creates a virtual matrix that acts as a submatrix
178: Collective on Mat
180: Input Parameters:
181: + A - matrix that we will extract a submatrix of
182: . isrow - rows to be present in the submatrix
183: - iscol - columns to be present in the submatrix
185: Output Parameters:
186: . newmat - new matrix
188: Level: developer
190: Notes:
191: Most will use MatCreateSubMatrix which provides a more efficient representation if it is available.
193: .seealso: MatCreateSubMatrix(), MatSubMatrixVirtualUpdate()
194: @*/
195: PetscErrorCode MatCreateSubMatrixVirtual(Mat A,IS isrow,IS iscol,Mat *newmat)
196: {
197: Vec left,right;
198: PetscInt m,n;
199: Mat N;
200: Mat_SubVirtual *Na;
208: *newmat = NULL;
210: MatCreate(PetscObjectComm((PetscObject)A),&N);
211: ISGetLocalSize(isrow,&m);
212: ISGetLocalSize(iscol,&n);
213: MatSetSizes(N,m,n,PETSC_DETERMINE,PETSC_DETERMINE);
214: PetscObjectChangeTypeName((PetscObject)N,MATSUBMATRIX);
216: PetscNewLog(N,&Na);
217: N->data = (void*)Na;
219: PetscObjectReference((PetscObject)isrow);
220: PetscObjectReference((PetscObject)iscol);
221: Na->isrow = isrow;
222: Na->iscol = iscol;
224: PetscFree(N->defaultvectype);
225: PetscStrallocpy(A->defaultvectype,&N->defaultvectype);
226: /* Do not use MatConvert directly since MatShell has a duplicate operation which does not increase
227: the reference count of the context. This is a problem if A is already of type MATSHELL */
228: MatConvertFrom_Shell(A,MATSHELL,MAT_INITIAL_MATRIX,&Na->A);
230: N->ops->destroy = MatDestroy_SubMatrix;
231: N->ops->mult = MatMult_SubMatrix;
232: N->ops->multadd = MatMultAdd_SubMatrix;
233: N->ops->multtranspose = MatMultTranspose_SubMatrix;
234: N->ops->multtransposeadd = MatMultTransposeAdd_SubMatrix;
235: N->ops->scale = MatScale_SubMatrix;
236: N->ops->diagonalscale = MatDiagonalScale_SubMatrix;
237: N->ops->shift = MatShift_SubMatrix;
238: N->ops->convert = MatConvert_Shell;
239: N->ops->getdiagonal = MatGetDiagonal_SubMatrix;
241: MatSetBlockSizesFromMats(N,A,A);
242: PetscLayoutSetUp(N->rmap);
243: PetscLayoutSetUp(N->cmap);
245: MatCreateVecs(A,&Na->rwork,&Na->lwork);
246: MatCreateVecs(N,&right,&left);
247: VecScatterCreate(Na->lwork,isrow,left,NULL,&Na->lrestrict);
248: VecScatterCreate(right,NULL,Na->rwork,iscol,&Na->rprolong);
249: VecDestroy(&left);
250: VecDestroy(&right);
251: MatSetUp(N);
253: N->assembled = PETSC_TRUE;
254: *newmat = N;
255: return(0);
256: }
259: /*@
260: MatSubMatrixVirtualUpdate - Updates a submatrix
262: Collective on Mat
264: Input Parameters:
265: + N - submatrix to update
266: . A - full matrix in the submatrix
267: . isrow - rows in the update (same as the first time the submatrix was created)
268: - iscol - columns in the update (same as the first time the submatrix was created)
270: Level: developer
272: Notes:
273: Most will use MatCreateSubMatrix which provides a more efficient representation if it is available.
275: .seealso: MatCreateSubMatrixVirtual()
276: @*/
277: PetscErrorCode MatSubMatrixVirtualUpdate(Mat N,Mat A,IS isrow,IS iscol)
278: {
280: PetscBool flg;
281: Mat_SubVirtual *Na;
288: PetscObjectTypeCompare((PetscObject)N,MATSUBMATRIX,&flg);
289: if (!flg) SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_ARG_WRONG,"Matrix has wrong type");
291: Na = (Mat_SubVirtual*)N->data;
292: ISEqual(isrow,Na->isrow,&flg);
293: if (!flg) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"Cannot update submatrix with different row indices");
294: ISEqual(iscol,Na->iscol,&flg);
295: if (!flg) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"Cannot update submatrix with different column indices");
297: PetscFree(N->defaultvectype);
298: PetscStrallocpy(A->defaultvectype,&N->defaultvectype);
299: MatDestroy(&Na->A);
300: /* Do not use MatConvert directly since MatShell has a duplicate operation which does not increase
301: the reference count of the context. This is a problem if A is already of type MATSHELL */
302: MatConvertFrom_Shell(A,MATSHELL,MAT_INITIAL_MATRIX,&Na->A);
303: return(0);
304: }