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(PETSC_SUCCESS); 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(PETSC_SUCCESS); 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(PETSC_SUCCESS); 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(PETSC_SUCCESS); 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(PETSC_SUCCESS); 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(PETSC_SUCCESS); 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(PETSC_SUCCESS); 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(PETSC_SUCCESS); 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(PETSC_SUCCESS); 164 } 165 166 /*@ 167 MatCreateSubMatrixVirtual - Creates a virtual matrix `MATSUBMATRIX` that acts as a submatrix 168 169 Collective 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 Note: 182 Most will use `MatCreateSubMatrix()` which provides a more efficient representation if it is available. 183 184 Developer Note: 185 The `MatType` is `MATSUBMATRIX` but the routines associated have `SubMatrixVirtual` in them, the `MatType` should likely be changed 186 187 .seealso: `MATSUBMATRIX`, `MATLOCALREF`, `MatCreateLocalRef()`, `MatCreateSubMatrix()`, `MatSubMatrixVirtualUpdate()` 188 @*/ 189 PetscErrorCode MatCreateSubMatrixVirtual(Mat A, IS isrow, IS iscol, Mat *newmat) 190 { 191 Vec left, right; 192 PetscInt m, n; 193 Mat N; 194 Mat_SubVirtual *Na; 195 196 PetscFunctionBegin; 197 PetscValidHeaderSpecific(A, MAT_CLASSID, 1); 198 PetscValidHeaderSpecific(isrow, IS_CLASSID, 2); 199 PetscValidHeaderSpecific(iscol, IS_CLASSID, 3); 200 PetscValidPointer(newmat, 4); 201 *newmat = NULL; 202 203 PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &N)); 204 PetscCall(ISGetLocalSize(isrow, &m)); 205 PetscCall(ISGetLocalSize(iscol, &n)); 206 PetscCall(MatSetSizes(N, m, n, PETSC_DETERMINE, PETSC_DETERMINE)); 207 PetscCall(PetscObjectChangeTypeName((PetscObject)N, MATSUBMATRIX)); 208 209 PetscCall(PetscNew(&Na)); 210 N->data = (void *)Na; 211 212 PetscCall(PetscObjectReference((PetscObject)isrow)); 213 PetscCall(PetscObjectReference((PetscObject)iscol)); 214 Na->isrow = isrow; 215 Na->iscol = iscol; 216 217 PetscCall(PetscFree(N->defaultvectype)); 218 PetscCall(PetscStrallocpy(A->defaultvectype, &N->defaultvectype)); 219 /* Do not use MatConvert directly since MatShell has a duplicate operation which does not increase 220 the reference count of the context. This is a problem if A is already of type MATSHELL */ 221 PetscCall(MatConvertFrom_Shell(A, MATSHELL, MAT_INITIAL_MATRIX, &Na->A)); 222 223 N->ops->destroy = MatDestroy_SubMatrix; 224 N->ops->mult = MatMult_SubMatrix; 225 N->ops->multadd = MatMultAdd_SubMatrix; 226 N->ops->multtranspose = MatMultTranspose_SubMatrix; 227 N->ops->multtransposeadd = MatMultTransposeAdd_SubMatrix; 228 N->ops->scale = MatScale_SubMatrix; 229 N->ops->diagonalscale = MatDiagonalScale_SubMatrix; 230 N->ops->shift = MatShift_SubMatrix; 231 N->ops->convert = MatConvert_Shell; 232 N->ops->getdiagonal = MatGetDiagonal_SubMatrix; 233 234 PetscCall(MatSetBlockSizesFromMats(N, A, A)); 235 PetscCall(PetscLayoutSetUp(N->rmap)); 236 PetscCall(PetscLayoutSetUp(N->cmap)); 237 238 PetscCall(MatCreateVecs(A, &Na->rwork, &Na->lwork)); 239 PetscCall(MatCreateVecs(N, &right, &left)); 240 PetscCall(VecScatterCreate(Na->lwork, isrow, left, NULL, &Na->lrestrict)); 241 PetscCall(VecScatterCreate(right, NULL, Na->rwork, iscol, &Na->rprolong)); 242 PetscCall(VecDestroy(&left)); 243 PetscCall(VecDestroy(&right)); 244 PetscCall(MatSetUp(N)); 245 246 N->assembled = PETSC_TRUE; 247 *newmat = N; 248 PetscFunctionReturn(PETSC_SUCCESS); 249 } 250 251 /*MC 252 MATSUBMATRIX - "submatrix" - A matrix type that represents a virtual submatrix of a matrix 253 254 Level: advanced 255 256 Developer Note: 257 This should be named `MATSUBMATRIXVIRTUAL` 258 259 .seealso: `Mat`, `MatCreateSubMatrixVirtual()`, `MatCreateSubMatrixVirtual()`, `MatCreateSubMatrix()` 260 M*/ 261 262 /*@ 263 MatSubMatrixVirtualUpdate - Updates a `MATSUBMATRIX` virtual submatrix 264 265 Collective 266 267 Input Parameters: 268 + N - submatrix to update 269 . A - full matrix in the submatrix 270 . isrow - rows in the update (same as the first time the submatrix was created) 271 - iscol - columns in the update (same as the first time the submatrix was created) 272 273 Level: developer 274 275 Note: 276 Most will use `MatCreateSubMatrix()` which provides a more efficient representation if it is available. 277 278 .seealso: MATSUBMATRIX`, `MatCreateSubMatrixVirtual()` 279 @*/ 280 PetscErrorCode MatSubMatrixVirtualUpdate(Mat N, Mat A, IS isrow, IS iscol) 281 { 282 PetscBool flg; 283 Mat_SubVirtual *Na; 284 285 PetscFunctionBegin; 286 PetscValidHeaderSpecific(N, MAT_CLASSID, 1); 287 PetscValidHeaderSpecific(A, MAT_CLASSID, 2); 288 PetscValidHeaderSpecific(isrow, IS_CLASSID, 3); 289 PetscValidHeaderSpecific(iscol, IS_CLASSID, 4); 290 PetscCall(PetscObjectTypeCompare((PetscObject)N, MATSUBMATRIX, &flg)); 291 PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Matrix has wrong type"); 292 293 Na = (Mat_SubVirtual *)N->data; 294 PetscCall(ISEqual(isrow, Na->isrow, &flg)); 295 PetscCheck(flg, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Cannot update submatrix with different row indices"); 296 PetscCall(ISEqual(iscol, Na->iscol, &flg)); 297 PetscCheck(flg, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Cannot update submatrix with different column indices"); 298 299 PetscCall(PetscFree(N->defaultvectype)); 300 PetscCall(PetscStrallocpy(A->defaultvectype, &N->defaultvectype)); 301 PetscCall(MatDestroy(&Na->A)); 302 /* Do not use MatConvert directly since MatShell has a duplicate operation which does not increase 303 the reference count of the context. This is a problem if A is already of type MATSHELL */ 304 PetscCall(MatConvertFrom_Shell(A, MATSHELL, MAT_INITIAL_MATRIX, &Na->A)); 305 PetscFunctionReturn(PETSC_SUCCESS); 306 } 307