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