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