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 PetscErrorCode ierr; 16 17 PetscFunctionBegin; 18 ierr = MatScale(Na->A,a);CHKERRQ(ierr); 19 PetscFunctionReturn(0); 20 } 21 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 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 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 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 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 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 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 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 @*/ 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 @*/ 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