1
2 #include <petsc/private/matimpl.h> /*I "petscmat.h" I*/
3
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;
11
MatScale_SubMatrix(Mat N,PetscScalar a)12 static PetscErrorCode MatScale_SubMatrix(Mat N,PetscScalar a)
13 {
14 Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
15 PetscErrorCode ierr;
16
17 PetscFunctionBegin;
18 ierr = MatScale(Na->A,a);CHKERRQ(ierr);
19 PetscFunctionReturn(0);
20 }
21
MatShift_SubMatrix(Mat N,PetscScalar a)22 static PetscErrorCode MatShift_SubMatrix(Mat N,PetscScalar a)
23 {
24 Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
25 PetscErrorCode ierr;
26
27 PetscFunctionBegin;
28 ierr = MatShift(Na->A,a);CHKERRQ(ierr);
29 PetscFunctionReturn(0);
30 }
31
MatDiagonalScale_SubMatrix(Mat N,Vec left,Vec right)32 static PetscErrorCode MatDiagonalScale_SubMatrix(Mat N,Vec left,Vec right)
33 {
34 Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
35 PetscErrorCode ierr;
36
37 PetscFunctionBegin;
38 if (right) {
39 ierr = VecZeroEntries(Na->rwork);CHKERRQ(ierr);
40 ierr = VecScatterBegin(Na->rprolong,right,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
41 ierr = VecScatterEnd(Na->rprolong,right,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
42 }
43 if (left) {
44 ierr = VecZeroEntries(Na->lwork);CHKERRQ(ierr);
45 ierr = VecScatterBegin(Na->lrestrict,left,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
46 ierr = VecScatterEnd(Na->lrestrict,left,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
47 }
48 ierr = MatDiagonalScale(Na->A,left ? Na->lwork : NULL,right ? Na->rwork : NULL);CHKERRQ(ierr);
49 PetscFunctionReturn(0);
50 }
51
MatGetDiagonal_SubMatrix(Mat N,Vec d)52 static PetscErrorCode MatGetDiagonal_SubMatrix(Mat N,Vec d)
53 {
54 Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
55 PetscErrorCode ierr;
56
57 PetscFunctionBegin;
58 ierr = MatGetDiagonal(Na->A,Na->rwork);CHKERRQ(ierr);
59 ierr = VecScatterBegin(Na->rprolong,Na->rwork,d,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
60 ierr = VecScatterEnd(Na->rprolong,Na->rwork,d,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
61 PetscFunctionReturn(0);
62 }
63
MatMult_SubMatrix(Mat N,Vec x,Vec y)64 static PetscErrorCode MatMult_SubMatrix(Mat N,Vec x,Vec y)
65 {
66 Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
67 PetscErrorCode ierr;
68
69 PetscFunctionBegin;
70 ierr = VecZeroEntries(Na->rwork);CHKERRQ(ierr);
71 ierr = VecScatterBegin(Na->rprolong,x,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
72 ierr = VecScatterEnd(Na->rprolong,x,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
73 ierr = MatMult(Na->A,Na->rwork,Na->lwork);CHKERRQ(ierr);
74 ierr = VecScatterBegin(Na->lrestrict,Na->lwork,y,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
75 ierr = VecScatterEnd(Na->lrestrict,Na->lwork,y,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
76 PetscFunctionReturn(0);
77 }
78
MatMultAdd_SubMatrix(Mat N,Vec v1,Vec v2,Vec v3)79 static PetscErrorCode MatMultAdd_SubMatrix(Mat N,Vec v1,Vec v2,Vec v3)
80 {
81 Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
82 PetscErrorCode ierr;
83
84 PetscFunctionBegin;
85 ierr = VecZeroEntries(Na->rwork);CHKERRQ(ierr);
86 ierr = VecScatterBegin(Na->rprolong,v1,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
87 ierr = VecScatterEnd(Na->rprolong,v1,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
88 if (v1 == v2) {
89 ierr = MatMultAdd(Na->A,Na->rwork,Na->rwork,Na->lwork);CHKERRQ(ierr);
90 } else if (v2 == v3) {
91 ierr = VecZeroEntries(Na->lwork);CHKERRQ(ierr);
92 ierr = VecScatterBegin(Na->lrestrict,v2,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
93 ierr = VecScatterEnd(Na->lrestrict,v2,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
94 ierr = MatMultAdd(Na->A,Na->rwork,Na->lwork,Na->lwork);CHKERRQ(ierr);
95 } else {
96 if (!Na->lwork2) {
97 ierr = VecDuplicate(Na->lwork,&Na->lwork2);CHKERRQ(ierr);
98 } else {
99 ierr = VecZeroEntries(Na->lwork2);CHKERRQ(ierr);
100 }
101 ierr = VecScatterBegin(Na->lrestrict,v2,Na->lwork2,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
102 ierr = VecScatterEnd(Na->lrestrict,v2,Na->lwork2,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
103 ierr = MatMultAdd(Na->A,Na->rwork,Na->lwork2,Na->lwork);CHKERRQ(ierr);
104 }
105 ierr = VecScatterBegin(Na->lrestrict,Na->lwork,v3,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
106 ierr = VecScatterEnd(Na->lrestrict,Na->lwork,v3,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
107 PetscFunctionReturn(0);
108 }
109
MatMultTranspose_SubMatrix(Mat N,Vec x,Vec y)110 static PetscErrorCode MatMultTranspose_SubMatrix(Mat N,Vec x,Vec y)
111 {
112 Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
113 PetscErrorCode ierr;
114
115 PetscFunctionBegin;
116 ierr = VecZeroEntries(Na->lwork);CHKERRQ(ierr);
117 ierr = VecScatterBegin(Na->lrestrict,x,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
118 ierr = VecScatterEnd(Na->lrestrict,x,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
119 ierr = MatMultTranspose(Na->A,Na->lwork,Na->rwork);CHKERRQ(ierr);
120 ierr = VecScatterBegin(Na->rprolong,Na->rwork,y,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
121 ierr = VecScatterEnd(Na->rprolong,Na->rwork,y,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
122 PetscFunctionReturn(0);
123 }
124
MatMultTransposeAdd_SubMatrix(Mat N,Vec v1,Vec v2,Vec v3)125 static PetscErrorCode MatMultTransposeAdd_SubMatrix(Mat N,Vec v1,Vec v2,Vec v3)
126 {
127 Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
128 PetscErrorCode ierr;
129
130 PetscFunctionBegin;
131 ierr = VecZeroEntries(Na->lwork);CHKERRQ(ierr);
132 ierr = VecScatterBegin(Na->lrestrict,v1,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
133 ierr = VecScatterEnd(Na->lrestrict,v1,Na->lwork,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
134 if (v1 == v2) {
135 ierr = MatMultTransposeAdd(Na->A,Na->lwork,Na->lwork,Na->rwork);CHKERRQ(ierr);
136 } else if (v2 == v3) {
137 ierr = VecZeroEntries(Na->rwork);CHKERRQ(ierr);
138 ierr = VecScatterBegin(Na->rprolong,v2,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
139 ierr = VecScatterEnd(Na->rprolong,v2,Na->rwork,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
140 ierr = MatMultTransposeAdd(Na->A,Na->lwork,Na->rwork,Na->rwork);CHKERRQ(ierr);
141 } else {
142 if (!Na->rwork2) {
143 ierr = VecDuplicate(Na->rwork,&Na->rwork2);CHKERRQ(ierr);
144 } else {
145 ierr = VecZeroEntries(Na->rwork2);CHKERRQ(ierr);
146 }
147 ierr = VecScatterBegin(Na->rprolong,v2,Na->rwork2,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
148 ierr = VecScatterEnd(Na->rprolong,v2,Na->rwork2,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
149 ierr = MatMultTransposeAdd(Na->A,Na->lwork,Na->rwork2,Na->rwork);CHKERRQ(ierr);
150 }
151 ierr = VecScatterBegin(Na->rprolong,Na->rwork,v3,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
152 ierr = VecScatterEnd(Na->rprolong,Na->rwork,v3,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
153 PetscFunctionReturn(0);
154 }
155
MatDestroy_SubMatrix(Mat N)156 static PetscErrorCode MatDestroy_SubMatrix(Mat N)
157 {
158 Mat_SubVirtual *Na = (Mat_SubVirtual*)N->data;
159 PetscErrorCode ierr;
160
161 PetscFunctionBegin;
162 ierr = ISDestroy(&Na->isrow);CHKERRQ(ierr);
163 ierr = ISDestroy(&Na->iscol);CHKERRQ(ierr);
164 ierr = VecDestroy(&Na->lwork);CHKERRQ(ierr);
165 ierr = VecDestroy(&Na->rwork);CHKERRQ(ierr);
166 ierr = VecDestroy(&Na->lwork2);CHKERRQ(ierr);
167 ierr = VecDestroy(&Na->rwork2);CHKERRQ(ierr);
168 ierr = VecScatterDestroy(&Na->lrestrict);CHKERRQ(ierr);
169 ierr = VecScatterDestroy(&Na->rprolong);CHKERRQ(ierr);
170 ierr = MatDestroy(&Na->A);CHKERRQ(ierr);
171 ierr = PetscFree(N->data);CHKERRQ(ierr);
172 PetscFunctionReturn(0);
173 }
174
175 /*@
176 MatCreateSubMatrixVirtual - Creates a virtual matrix that acts as a submatrix
177
178 Collective on Mat
179
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
184
185 Output Parameters:
186 . newmat - new matrix
187
188 Level: developer
189
190 Notes:
191 Most will use MatCreateSubMatrix which provides a more efficient representation if it is available.
192
193 .seealso: MatCreateSubMatrix(), MatSubMatrixVirtualUpdate()
194 @*/
MatCreateSubMatrixVirtual(Mat A,IS isrow,IS iscol,Mat * newmat)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;
201 PetscErrorCode ierr;
202
203 PetscFunctionBegin;
204 PetscValidHeaderSpecific(A,MAT_CLASSID,1);
205 PetscValidHeaderSpecific(isrow,IS_CLASSID,2);
206 PetscValidHeaderSpecific(iscol,IS_CLASSID,3);
207 PetscValidPointer(newmat,4);
208 *newmat = NULL;
209
210 ierr = MatCreate(PetscObjectComm((PetscObject)A),&N);CHKERRQ(ierr);
211 ierr = ISGetLocalSize(isrow,&m);CHKERRQ(ierr);
212 ierr = ISGetLocalSize(iscol,&n);CHKERRQ(ierr);
213 ierr = MatSetSizes(N,m,n,PETSC_DETERMINE,PETSC_DETERMINE);CHKERRQ(ierr);
214 ierr = PetscObjectChangeTypeName((PetscObject)N,MATSUBMATRIX);CHKERRQ(ierr);
215
216 ierr = PetscNewLog(N,&Na);CHKERRQ(ierr);
217 N->data = (void*)Na;
218
219 ierr = PetscObjectReference((PetscObject)isrow);CHKERRQ(ierr);
220 ierr = PetscObjectReference((PetscObject)iscol);CHKERRQ(ierr);
221 Na->isrow = isrow;
222 Na->iscol = iscol;
223
224 ierr = PetscFree(N->defaultvectype);CHKERRQ(ierr);
225 ierr = PetscStrallocpy(A->defaultvectype,&N->defaultvectype);CHKERRQ(ierr);
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 ierr = MatConvertFrom_Shell(A,MATSHELL,MAT_INITIAL_MATRIX,&Na->A);CHKERRQ(ierr);
229
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;
240
241 ierr = MatSetBlockSizesFromMats(N,A,A);CHKERRQ(ierr);
242 ierr = PetscLayoutSetUp(N->rmap);CHKERRQ(ierr);
243 ierr = PetscLayoutSetUp(N->cmap);CHKERRQ(ierr);
244
245 ierr = MatCreateVecs(A,&Na->rwork,&Na->lwork);CHKERRQ(ierr);
246 ierr = MatCreateVecs(N,&right,&left);CHKERRQ(ierr);
247 ierr = VecScatterCreate(Na->lwork,isrow,left,NULL,&Na->lrestrict);CHKERRQ(ierr);
248 ierr = VecScatterCreate(right,NULL,Na->rwork,iscol,&Na->rprolong);CHKERRQ(ierr);
249 ierr = VecDestroy(&left);CHKERRQ(ierr);
250 ierr = VecDestroy(&right);CHKERRQ(ierr);
251 ierr = MatSetUp(N);CHKERRQ(ierr);
252
253 N->assembled = PETSC_TRUE;
254 *newmat = N;
255 PetscFunctionReturn(0);
256 }
257
258
259 /*@
260 MatSubMatrixVirtualUpdate - Updates a submatrix
261
262 Collective on Mat
263
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)
269
270 Level: developer
271
272 Notes:
273 Most will use MatCreateSubMatrix which provides a more efficient representation if it is available.
274
275 .seealso: MatCreateSubMatrixVirtual()
276 @*/
MatSubMatrixVirtualUpdate(Mat N,Mat A,IS isrow,IS iscol)277 PetscErrorCode MatSubMatrixVirtualUpdate(Mat N,Mat A,IS isrow,IS iscol)
278 {
279 PetscErrorCode ierr;
280 PetscBool flg;
281 Mat_SubVirtual *Na;
282
283 PetscFunctionBegin;
284 PetscValidHeaderSpecific(N,MAT_CLASSID,1);
285 PetscValidHeaderSpecific(A,MAT_CLASSID,2);
286 PetscValidHeaderSpecific(isrow,IS_CLASSID,3);
287 PetscValidHeaderSpecific(iscol,IS_CLASSID,4);
288 ierr = PetscObjectTypeCompare((PetscObject)N,MATSUBMATRIX,&flg);CHKERRQ(ierr);
289 if (!flg) SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_ARG_WRONG,"Matrix has wrong type");
290
291 Na = (Mat_SubVirtual*)N->data;
292 ierr = ISEqual(isrow,Na->isrow,&flg);CHKERRQ(ierr);
293 if (!flg) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"Cannot update submatrix with different row indices");
294 ierr = ISEqual(iscol,Na->iscol,&flg);CHKERRQ(ierr);
295 if (!flg) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"Cannot update submatrix with different column indices");
296
297 ierr = PetscFree(N->defaultvectype);CHKERRQ(ierr);
298 ierr = PetscStrallocpy(A->defaultvectype,&N->defaultvectype);CHKERRQ(ierr);
299 ierr = MatDestroy(&Na->A);CHKERRQ(ierr);
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 ierr = MatConvertFrom_Shell(A,MATSHELL,MAT_INITIAL_MATRIX,&Na->A);CHKERRQ(ierr);
303 PetscFunctionReturn(0);
304 }
305