1 2 #include <petsc/private/matimpl.h> /*I "petscmat.h" I*/ 3 4 const char *const MatCompositeMergeTypes[] = {"left","right","MatCompositeMergeType","MAT_COMPOSITE_",NULL}; 5 6 typedef struct _Mat_CompositeLink *Mat_CompositeLink; 7 struct _Mat_CompositeLink { 8 Mat mat; 9 Vec work; 10 Mat_CompositeLink next,prev; 11 }; 12 13 typedef struct { 14 MatCompositeType type; 15 Mat_CompositeLink head,tail; 16 Vec work; 17 PetscScalar scale; /* scale factor supplied with MatScale() */ 18 Vec left,right; /* left and right diagonal scaling provided with MatDiagonalScale() */ 19 Vec leftwork,rightwork,leftwork2,rightwork2; /* Two pairs of working vectors */ 20 PetscInt nmat; 21 PetscBool merge; 22 MatCompositeMergeType mergetype; 23 MatStructure structure; 24 25 PetscScalar *scalings; 26 PetscBool merge_mvctx; /* Whether need to merge mvctx of component matrices */ 27 Vec *lvecs; /* [nmat] Basically, they are Mvctx->lvec of each component matrix */ 28 PetscScalar *larray; /* [len] Data arrays of lvecs[] are stored consecutively in larray */ 29 PetscInt len; /* Length of larray[] */ 30 Vec gvec; /* Union of lvecs[] without duplicated entries */ 31 PetscInt *location; /* A map that maps entries in garray[] to larray[] */ 32 VecScatter Mvctx; 33 } Mat_Composite; 34 35 PetscErrorCode MatDestroy_Composite(Mat mat) 36 { 37 Mat_Composite *shell = (Mat_Composite*)mat->data; 38 Mat_CompositeLink next = shell->head,oldnext; 39 PetscInt i; 40 41 PetscFunctionBegin; 42 while (next) { 43 PetscCall(MatDestroy(&next->mat)); 44 if (next->work && (!next->next || next->work != next->next->work)) { 45 PetscCall(VecDestroy(&next->work)); 46 } 47 oldnext = next; 48 next = next->next; 49 PetscCall(PetscFree(oldnext)); 50 } 51 PetscCall(VecDestroy(&shell->work)); 52 PetscCall(VecDestroy(&shell->left)); 53 PetscCall(VecDestroy(&shell->right)); 54 PetscCall(VecDestroy(&shell->leftwork)); 55 PetscCall(VecDestroy(&shell->rightwork)); 56 PetscCall(VecDestroy(&shell->leftwork2)); 57 PetscCall(VecDestroy(&shell->rightwork2)); 58 59 if (shell->Mvctx) { 60 for (i=0; i<shell->nmat; i++) PetscCall(VecDestroy(&shell->lvecs[i])); 61 PetscCall(PetscFree3(shell->location,shell->larray,shell->lvecs)); 62 PetscCall(PetscFree(shell->larray)); 63 PetscCall(VecDestroy(&shell->gvec)); 64 PetscCall(VecScatterDestroy(&shell->Mvctx)); 65 } 66 67 PetscCall(PetscFree(shell->scalings)); 68 PetscCall(PetscFree(mat->data)); 69 PetscFunctionReturn(0); 70 } 71 72 PetscErrorCode MatMult_Composite_Multiplicative(Mat A,Vec x,Vec y) 73 { 74 Mat_Composite *shell = (Mat_Composite*)A->data; 75 Mat_CompositeLink next = shell->head; 76 Vec in,out; 77 PetscScalar scale; 78 PetscInt i; 79 80 PetscFunctionBegin; 81 PetscCheck(next,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()"); 82 in = x; 83 if (shell->right) { 84 if (!shell->rightwork) { 85 PetscCall(VecDuplicate(shell->right,&shell->rightwork)); 86 } 87 PetscCall(VecPointwiseMult(shell->rightwork,shell->right,in)); 88 in = shell->rightwork; 89 } 90 while (next->next) { 91 if (!next->work) { /* should reuse previous work if the same size */ 92 PetscCall(MatCreateVecs(next->mat,NULL,&next->work)); 93 } 94 out = next->work; 95 PetscCall(MatMult(next->mat,in,out)); 96 in = out; 97 next = next->next; 98 } 99 PetscCall(MatMult(next->mat,in,y)); 100 if (shell->left) { 101 PetscCall(VecPointwiseMult(y,shell->left,y)); 102 } 103 scale = shell->scale; 104 if (shell->scalings) {for (i=0; i<shell->nmat; i++) scale *= shell->scalings[i];} 105 PetscCall(VecScale(y,scale)); 106 PetscFunctionReturn(0); 107 } 108 109 PetscErrorCode MatMultTranspose_Composite_Multiplicative(Mat A,Vec x,Vec y) 110 { 111 Mat_Composite *shell = (Mat_Composite*)A->data; 112 Mat_CompositeLink tail = shell->tail; 113 Vec in,out; 114 PetscScalar scale; 115 PetscInt i; 116 117 PetscFunctionBegin; 118 PetscCheck(tail,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()"); 119 in = x; 120 if (shell->left) { 121 if (!shell->leftwork) { 122 PetscCall(VecDuplicate(shell->left,&shell->leftwork)); 123 } 124 PetscCall(VecPointwiseMult(shell->leftwork,shell->left,in)); 125 in = shell->leftwork; 126 } 127 while (tail->prev) { 128 if (!tail->prev->work) { /* should reuse previous work if the same size */ 129 PetscCall(MatCreateVecs(tail->mat,NULL,&tail->prev->work)); 130 } 131 out = tail->prev->work; 132 PetscCall(MatMultTranspose(tail->mat,in,out)); 133 in = out; 134 tail = tail->prev; 135 } 136 PetscCall(MatMultTranspose(tail->mat,in,y)); 137 if (shell->right) { 138 PetscCall(VecPointwiseMult(y,shell->right,y)); 139 } 140 141 scale = shell->scale; 142 if (shell->scalings) {for (i=0; i<shell->nmat; i++) scale *= shell->scalings[i];} 143 PetscCall(VecScale(y,scale)); 144 PetscFunctionReturn(0); 145 } 146 147 PetscErrorCode MatMult_Composite(Mat mat,Vec x,Vec y) 148 { 149 Mat_Composite *shell = (Mat_Composite*)mat->data; 150 Mat_CompositeLink cur = shell->head; 151 Vec in,y2,xin; 152 Mat A,B; 153 PetscInt i,j,k,n,nuniq,lo,hi,mid,*gindices,*buf,*tmp,tot; 154 const PetscScalar *vals; 155 const PetscInt *garray; 156 IS ix,iy; 157 PetscBool match; 158 159 PetscFunctionBegin; 160 PetscCheck(cur,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()"); 161 in = x; 162 if (shell->right) { 163 if (!shell->rightwork) { 164 PetscCall(VecDuplicate(shell->right,&shell->rightwork)); 165 } 166 PetscCall(VecPointwiseMult(shell->rightwork,shell->right,in)); 167 in = shell->rightwork; 168 } 169 170 /* Try to merge Mvctx when instructed but not yet done. We did not do it in MatAssemblyEnd() since at that time 171 we did not know whether mat is ADDITIVE or MULTIPLICATIVE. Only now we are assured mat is ADDITIVE and 172 it is legal to merge Mvctx, because all component matrices have the same size. 173 */ 174 if (shell->merge_mvctx && !shell->Mvctx) { 175 /* Currently only implemented for MATMPIAIJ */ 176 for (cur=shell->head; cur; cur=cur->next) { 177 PetscCall(PetscObjectTypeCompare((PetscObject)cur->mat,MATMPIAIJ,&match)); 178 if (!match) { 179 shell->merge_mvctx = PETSC_FALSE; 180 goto skip_merge_mvctx; 181 } 182 } 183 184 /* Go through matrices first time to count total number of nonzero off-diag columns (may have dups) */ 185 tot = 0; 186 for (cur=shell->head; cur; cur=cur->next) { 187 PetscCall(MatMPIAIJGetSeqAIJ(cur->mat,NULL,&B,NULL)); 188 PetscCall(MatGetLocalSize(B,NULL,&n)); 189 tot += n; 190 } 191 PetscCall(PetscMalloc3(tot,&shell->location,tot,&shell->larray,shell->nmat,&shell->lvecs)); 192 shell->len = tot; 193 194 /* Go through matrices second time to sort off-diag columns and remove dups */ 195 PetscCall(PetscMalloc1(tot,&gindices)); /* No Malloc2() since we will give one to petsc and free the other */ 196 PetscCall(PetscMalloc1(tot,&buf)); 197 nuniq = 0; /* Number of unique nonzero columns */ 198 for (cur=shell->head; cur; cur=cur->next) { 199 PetscCall(MatMPIAIJGetSeqAIJ(cur->mat,NULL,&B,&garray)); 200 PetscCall(MatGetLocalSize(B,NULL,&n)); 201 /* Merge pre-sorted garray[0,n) and gindices[0,nuniq) to buf[] */ 202 i = j = k = 0; 203 while (i < n && j < nuniq) { 204 if (garray[i] < gindices[j]) buf[k++] = garray[i++]; 205 else if (garray[i] > gindices[j]) buf[k++] = gindices[j++]; 206 else {buf[k++] = garray[i++]; j++;} 207 } 208 /* Copy leftover in garray[] or gindices[] */ 209 if (i < n) { 210 PetscCall(PetscArraycpy(buf+k,garray+i,n-i)); 211 nuniq = k + n-i; 212 } else if (j < nuniq) { 213 PetscCall(PetscArraycpy(buf+k,gindices+j,nuniq-j)); 214 nuniq = k + nuniq-j; 215 } else nuniq = k; 216 /* Swap gindices and buf to merge garray of the next matrix */ 217 tmp = gindices; 218 gindices = buf; 219 buf = tmp; 220 } 221 PetscCall(PetscFree(buf)); 222 223 /* Go through matrices third time to build a map from gindices[] to garray[] */ 224 tot = 0; 225 for (cur=shell->head,j=0; cur; cur=cur->next,j++) { /* j-th matrix */ 226 PetscCall(MatMPIAIJGetSeqAIJ(cur->mat,NULL,&B,&garray)); 227 PetscCall(MatGetLocalSize(B,NULL,&n)); 228 PetscCall(VecCreateSeqWithArray(PETSC_COMM_SELF,1,n,NULL,&shell->lvecs[j])); 229 /* This is an optimized PetscFindInt(garray[i],nuniq,gindices,&shell->location[tot+i]), using the fact that garray[] is also sorted */ 230 lo = 0; 231 for (i=0; i<n; i++) { 232 hi = nuniq; 233 while (hi - lo > 1) { 234 mid = lo + (hi - lo)/2; 235 if (garray[i] < gindices[mid]) hi = mid; 236 else lo = mid; 237 } 238 shell->location[tot+i] = lo; /* gindices[lo] = garray[i] */ 239 lo++; /* Since garray[i+1] > garray[i], we can safely advance lo */ 240 } 241 tot += n; 242 } 243 244 /* Build merged Mvctx */ 245 PetscCall(ISCreateGeneral(PETSC_COMM_SELF,nuniq,gindices,PETSC_OWN_POINTER,&ix)); 246 PetscCall(ISCreateStride(PETSC_COMM_SELF,nuniq,0,1,&iy)); 247 PetscCall(VecCreateMPIWithArray(PetscObjectComm((PetscObject)mat),1,mat->cmap->n,mat->cmap->N,NULL,&xin)); 248 PetscCall(VecCreateSeq(PETSC_COMM_SELF,nuniq,&shell->gvec)); 249 PetscCall(VecScatterCreate(xin,ix,shell->gvec,iy,&shell->Mvctx)); 250 PetscCall(VecDestroy(&xin)); 251 PetscCall(ISDestroy(&ix)); 252 PetscCall(ISDestroy(&iy)); 253 } 254 255 skip_merge_mvctx: 256 PetscCall(VecSet(y,0)); 257 if (!shell->leftwork2) PetscCall(VecDuplicate(y,&shell->leftwork2)); 258 y2 = shell->leftwork2; 259 260 if (shell->Mvctx) { /* Have a merged Mvctx */ 261 /* Suppose we want to compute y = sMx, where s is the scaling factor and A, B are matrix M's diagonal/off-diagonal part. We could do 262 in y = s(Ax1 + Bx2) or y = sAx1 + sBx2. The former incurs less FLOPS than the latter, but the latter provides an oppertunity to 263 overlap communication/computation since we can do sAx1 while communicating x2. Here, we use the former approach. 264 */ 265 PetscCall(VecScatterBegin(shell->Mvctx,in,shell->gvec,INSERT_VALUES,SCATTER_FORWARD)); 266 PetscCall(VecScatterEnd(shell->Mvctx,in,shell->gvec,INSERT_VALUES,SCATTER_FORWARD)); 267 268 PetscCall(VecGetArrayRead(shell->gvec,&vals)); 269 for (i=0; i<shell->len; i++) shell->larray[i] = vals[shell->location[i]]; 270 PetscCall(VecRestoreArrayRead(shell->gvec,&vals)); 271 272 for (cur=shell->head,tot=i=0; cur; cur=cur->next,i++) { /* i-th matrix */ 273 PetscCall(MatMPIAIJGetSeqAIJ(cur->mat,&A,&B,NULL)); 274 PetscCall((*A->ops->mult)(A,in,y2)); 275 PetscCall(MatGetLocalSize(B,NULL,&n)); 276 PetscCall(VecPlaceArray(shell->lvecs[i],&shell->larray[tot])); 277 PetscCall((*B->ops->multadd)(B,shell->lvecs[i],y2,y2)); 278 PetscCall(VecResetArray(shell->lvecs[i])); 279 PetscCall(VecAXPY(y,(shell->scalings ? shell->scalings[i] : 1.0),y2)); 280 tot += n; 281 } 282 } else { 283 if (shell->scalings) { 284 for (cur=shell->head,i=0; cur; cur=cur->next,i++) { 285 PetscCall(MatMult(cur->mat,in,y2)); 286 PetscCall(VecAXPY(y,shell->scalings[i],y2)); 287 } 288 } else { 289 for (cur=shell->head; cur; cur=cur->next) PetscCall(MatMultAdd(cur->mat,in,y,y)); 290 } 291 } 292 293 if (shell->left) PetscCall(VecPointwiseMult(y,shell->left,y)); 294 PetscCall(VecScale(y,shell->scale)); 295 PetscFunctionReturn(0); 296 } 297 298 PetscErrorCode MatMultTranspose_Composite(Mat A,Vec x,Vec y) 299 { 300 Mat_Composite *shell = (Mat_Composite*)A->data; 301 Mat_CompositeLink next = shell->head; 302 Vec in,y2 = NULL; 303 PetscInt i; 304 305 PetscFunctionBegin; 306 PetscCheck(next,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()"); 307 in = x; 308 if (shell->left) { 309 if (!shell->leftwork) { 310 PetscCall(VecDuplicate(shell->left,&shell->leftwork)); 311 } 312 PetscCall(VecPointwiseMult(shell->leftwork,shell->left,in)); 313 in = shell->leftwork; 314 } 315 316 PetscCall(MatMultTranspose(next->mat,in,y)); 317 if (shell->scalings) { 318 PetscCall(VecScale(y,shell->scalings[0])); 319 if (!shell->rightwork2) PetscCall(VecDuplicate(y,&shell->rightwork2)); 320 y2 = shell->rightwork2; 321 } 322 i = 1; 323 while ((next = next->next)) { 324 if (!shell->scalings) PetscCall(MatMultTransposeAdd(next->mat,in,y,y)); 325 else { 326 PetscCall(MatMultTranspose(next->mat,in,y2)); 327 PetscCall(VecAXPY(y,shell->scalings[i++],y2)); 328 } 329 } 330 if (shell->right) { 331 PetscCall(VecPointwiseMult(y,shell->right,y)); 332 } 333 PetscCall(VecScale(y,shell->scale)); 334 PetscFunctionReturn(0); 335 } 336 337 PetscErrorCode MatMultAdd_Composite(Mat A,Vec x,Vec y,Vec z) 338 { 339 Mat_Composite *shell = (Mat_Composite*)A->data; 340 341 PetscFunctionBegin; 342 if (y != z) { 343 PetscCall(MatMult(A,x,z)); 344 PetscCall(VecAXPY(z,1.0,y)); 345 } else { 346 if (!shell->leftwork) { 347 PetscCall(VecDuplicate(z,&shell->leftwork)); 348 } 349 PetscCall(MatMult(A,x,shell->leftwork)); 350 PetscCall(VecCopy(y,z)); 351 PetscCall(VecAXPY(z,1.0,shell->leftwork)); 352 } 353 PetscFunctionReturn(0); 354 } 355 356 PetscErrorCode MatMultTransposeAdd_Composite(Mat A,Vec x,Vec y, Vec z) 357 { 358 Mat_Composite *shell = (Mat_Composite*)A->data; 359 360 PetscFunctionBegin; 361 if (y != z) { 362 PetscCall(MatMultTranspose(A,x,z)); 363 PetscCall(VecAXPY(z,1.0,y)); 364 } else { 365 if (!shell->rightwork) { 366 PetscCall(VecDuplicate(z,&shell->rightwork)); 367 } 368 PetscCall(MatMultTranspose(A,x,shell->rightwork)); 369 PetscCall(VecCopy(y,z)); 370 PetscCall(VecAXPY(z,1.0,shell->rightwork)); 371 } 372 PetscFunctionReturn(0); 373 } 374 375 PetscErrorCode MatGetDiagonal_Composite(Mat A,Vec v) 376 { 377 Mat_Composite *shell = (Mat_Composite*)A->data; 378 Mat_CompositeLink next = shell->head; 379 PetscInt i; 380 381 PetscFunctionBegin; 382 PetscCheck(next,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()"); 383 PetscCheck(!shell->right && !shell->left,PETSC_COMM_SELF,PETSC_ERR_SUP,"Cannot get diagonal if left or right scaling"); 384 385 PetscCall(MatGetDiagonal(next->mat,v)); 386 if (shell->scalings) PetscCall(VecScale(v,shell->scalings[0])); 387 388 if (next->next && !shell->work) { 389 PetscCall(VecDuplicate(v,&shell->work)); 390 } 391 i = 1; 392 while ((next = next->next)) { 393 PetscCall(MatGetDiagonal(next->mat,shell->work)); 394 PetscCall(VecAXPY(v,(shell->scalings ? shell->scalings[i++] : 1.0),shell->work)); 395 } 396 PetscCall(VecScale(v,shell->scale)); 397 PetscFunctionReturn(0); 398 } 399 400 PetscErrorCode MatAssemblyEnd_Composite(Mat Y,MatAssemblyType t) 401 { 402 Mat_Composite *shell = (Mat_Composite*)Y->data; 403 404 PetscFunctionBegin; 405 if (shell->merge) { 406 PetscCall(MatCompositeMerge(Y)); 407 } 408 PetscFunctionReturn(0); 409 } 410 411 PetscErrorCode MatScale_Composite(Mat inA,PetscScalar alpha) 412 { 413 Mat_Composite *a = (Mat_Composite*)inA->data; 414 415 PetscFunctionBegin; 416 a->scale *= alpha; 417 PetscFunctionReturn(0); 418 } 419 420 PetscErrorCode MatDiagonalScale_Composite(Mat inA,Vec left,Vec right) 421 { 422 Mat_Composite *a = (Mat_Composite*)inA->data; 423 424 PetscFunctionBegin; 425 if (left) { 426 if (!a->left) { 427 PetscCall(VecDuplicate(left,&a->left)); 428 PetscCall(VecCopy(left,a->left)); 429 } else { 430 PetscCall(VecPointwiseMult(a->left,left,a->left)); 431 } 432 } 433 if (right) { 434 if (!a->right) { 435 PetscCall(VecDuplicate(right,&a->right)); 436 PetscCall(VecCopy(right,a->right)); 437 } else { 438 PetscCall(VecPointwiseMult(a->right,right,a->right)); 439 } 440 } 441 PetscFunctionReturn(0); 442 } 443 444 PetscErrorCode MatSetFromOptions_Composite(PetscOptionItems *PetscOptionsObject,Mat A) 445 { 446 Mat_Composite *a = (Mat_Composite*)A->data; 447 448 PetscFunctionBegin; 449 PetscOptionsHeadBegin(PetscOptionsObject,"MATCOMPOSITE options"); 450 PetscCall(PetscOptionsBool("-mat_composite_merge","Merge at MatAssemblyEnd","MatCompositeMerge",a->merge,&a->merge,NULL)); 451 PetscCall(PetscOptionsEnum("-mat_composite_merge_type","Set composite merge direction","MatCompositeSetMergeType",MatCompositeMergeTypes,(PetscEnum)a->mergetype,(PetscEnum*)&a->mergetype,NULL)); 452 PetscCall(PetscOptionsBool("-mat_composite_merge_mvctx","Merge MatMult() vecscat contexts","MatCreateComposite",a->merge_mvctx,&a->merge_mvctx,NULL)); 453 PetscOptionsHeadEnd(); 454 PetscFunctionReturn(0); 455 } 456 457 /*@ 458 MatCreateComposite - Creates a matrix as the sum or product of one or more matrices 459 460 Collective 461 462 Input Parameters: 463 + comm - MPI communicator 464 . nmat - number of matrices to put in 465 - mats - the matrices 466 467 Output Parameter: 468 . mat - the matrix 469 470 Options Database Keys: 471 + -mat_composite_merge - merge in MatAssemblyEnd() 472 . -mat_composite_merge_mvctx - merge Mvctx of component matrices to optimize communication in MatMult() for ADDITIVE matrices 473 - -mat_composite_merge_type - set merge direction 474 475 Level: advanced 476 477 Notes: 478 Alternative construction 479 $ MatCreate(comm,&mat); 480 $ MatSetSizes(mat,m,n,M,N); 481 $ MatSetType(mat,MATCOMPOSITE); 482 $ MatCompositeAddMat(mat,mats[0]); 483 $ .... 484 $ MatCompositeAddMat(mat,mats[nmat-1]); 485 $ MatAssemblyBegin(mat,MAT_FINAL_ASSEMBLY); 486 $ MatAssemblyEnd(mat,MAT_FINAL_ASSEMBLY); 487 488 For the multiplicative form the product is mat[nmat-1]*mat[nmat-2]*....*mat[0] 489 490 .seealso: `MatDestroy()`, `MatMult()`, `MatCompositeAddMat()`, `MatCompositeGetMat()`, `MatCompositeMerge()`, `MatCompositeSetType()`, `MATCOMPOSITE` 491 492 @*/ 493 PetscErrorCode MatCreateComposite(MPI_Comm comm,PetscInt nmat,const Mat *mats,Mat *mat) 494 { 495 PetscInt m,n,M,N,i; 496 497 PetscFunctionBegin; 498 PetscCheck(nmat >= 1,PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Must pass in at least one matrix"); 499 PetscValidPointer(mat,4); 500 501 PetscCall(MatGetLocalSize(mats[0],PETSC_IGNORE,&n)); 502 PetscCall(MatGetLocalSize(mats[nmat-1],&m,PETSC_IGNORE)); 503 PetscCall(MatGetSize(mats[0],PETSC_IGNORE,&N)); 504 PetscCall(MatGetSize(mats[nmat-1],&M,PETSC_IGNORE)); 505 PetscCall(MatCreate(comm,mat)); 506 PetscCall(MatSetSizes(*mat,m,n,M,N)); 507 PetscCall(MatSetType(*mat,MATCOMPOSITE)); 508 for (i=0; i<nmat; i++) { 509 PetscCall(MatCompositeAddMat(*mat,mats[i])); 510 } 511 PetscCall(MatAssemblyBegin(*mat,MAT_FINAL_ASSEMBLY)); 512 PetscCall(MatAssemblyEnd(*mat,MAT_FINAL_ASSEMBLY)); 513 PetscFunctionReturn(0); 514 } 515 516 static PetscErrorCode MatCompositeAddMat_Composite(Mat mat,Mat smat) 517 { 518 Mat_Composite *shell = (Mat_Composite*)mat->data; 519 Mat_CompositeLink ilink,next = shell->head; 520 521 PetscFunctionBegin; 522 PetscCall(PetscNewLog(mat,&ilink)); 523 ilink->next = NULL; 524 PetscCall(PetscObjectReference((PetscObject)smat)); 525 ilink->mat = smat; 526 527 if (!next) shell->head = ilink; 528 else { 529 while (next->next) { 530 next = next->next; 531 } 532 next->next = ilink; 533 ilink->prev = next; 534 } 535 shell->tail = ilink; 536 shell->nmat += 1; 537 538 /* Retain the old scalings (if any) and expand it with a 1.0 for the newly added matrix */ 539 if (shell->scalings) { 540 PetscCall(PetscRealloc(sizeof(PetscScalar)*shell->nmat,&shell->scalings)); 541 shell->scalings[shell->nmat-1] = 1.0; 542 } 543 PetscFunctionReturn(0); 544 } 545 546 /*@ 547 MatCompositeAddMat - Add another matrix to a composite matrix. 548 549 Collective on Mat 550 551 Input Parameters: 552 + mat - the composite matrix 553 - smat - the partial matrix 554 555 Level: advanced 556 557 .seealso: `MatCreateComposite()`, `MatCompositeGetMat()`, `MATCOMPOSITE` 558 @*/ 559 PetscErrorCode MatCompositeAddMat(Mat mat,Mat smat) 560 { 561 PetscFunctionBegin; 562 PetscValidHeaderSpecific(mat,MAT_CLASSID,1); 563 PetscValidHeaderSpecific(smat,MAT_CLASSID,2); 564 PetscUseMethod(mat,"MatCompositeAddMat_C",(Mat,Mat),(mat,smat)); 565 PetscFunctionReturn(0); 566 } 567 568 static PetscErrorCode MatCompositeSetType_Composite(Mat mat,MatCompositeType type) 569 { 570 Mat_Composite *b = (Mat_Composite*)mat->data; 571 572 PetscFunctionBegin; 573 b->type = type; 574 if (type == MAT_COMPOSITE_MULTIPLICATIVE) { 575 mat->ops->getdiagonal = NULL; 576 mat->ops->mult = MatMult_Composite_Multiplicative; 577 mat->ops->multtranspose = MatMultTranspose_Composite_Multiplicative; 578 b->merge_mvctx = PETSC_FALSE; 579 } else { 580 mat->ops->getdiagonal = MatGetDiagonal_Composite; 581 mat->ops->mult = MatMult_Composite; 582 mat->ops->multtranspose = MatMultTranspose_Composite; 583 } 584 PetscFunctionReturn(0); 585 } 586 587 /*@ 588 MatCompositeSetType - Indicates if the matrix is defined as the sum of a set of matrices or the product. 589 590 Logically Collective on Mat 591 592 Input Parameters: 593 . mat - the composite matrix 594 595 Level: advanced 596 597 .seealso: `MatDestroy()`, `MatMult()`, `MatCompositeAddMat()`, `MatCreateComposite()`, `MatCompositeGetType()`, `MATCOMPOSITE` 598 599 @*/ 600 PetscErrorCode MatCompositeSetType(Mat mat,MatCompositeType type) 601 { 602 PetscFunctionBegin; 603 PetscValidHeaderSpecific(mat,MAT_CLASSID,1); 604 PetscValidLogicalCollectiveEnum(mat,type,2); 605 PetscUseMethod(mat,"MatCompositeSetType_C",(Mat,MatCompositeType),(mat,type)); 606 PetscFunctionReturn(0); 607 } 608 609 static PetscErrorCode MatCompositeGetType_Composite(Mat mat,MatCompositeType *type) 610 { 611 Mat_Composite *b = (Mat_Composite*)mat->data; 612 613 PetscFunctionBegin; 614 *type = b->type; 615 PetscFunctionReturn(0); 616 } 617 618 /*@ 619 MatCompositeGetType - Returns type of composite. 620 621 Not Collective 622 623 Input Parameter: 624 . mat - the composite matrix 625 626 Output Parameter: 627 . type - type of composite 628 629 Level: advanced 630 631 .seealso: `MatCreateComposite()`, `MatCompositeSetType()`, `MATCOMPOSITE` 632 633 @*/ 634 PetscErrorCode MatCompositeGetType(Mat mat,MatCompositeType *type) 635 { 636 PetscFunctionBegin; 637 PetscValidHeaderSpecific(mat,MAT_CLASSID,1); 638 PetscValidPointer(type,2); 639 PetscUseMethod(mat,"MatCompositeGetType_C",(Mat,MatCompositeType*),(mat,type)); 640 PetscFunctionReturn(0); 641 } 642 643 static PetscErrorCode MatCompositeSetMatStructure_Composite(Mat mat,MatStructure str) 644 { 645 Mat_Composite *b = (Mat_Composite*)mat->data; 646 647 PetscFunctionBegin; 648 b->structure = str; 649 PetscFunctionReturn(0); 650 } 651 652 /*@ 653 MatCompositeSetMatStructure - Indicates structure of matrices in the composite matrix. 654 655 Not Collective 656 657 Input Parameters: 658 + mat - the composite matrix 659 - str - either SAME_NONZERO_PATTERN, DIFFERENT_NONZERO_PATTERN (default) or SUBSET_NONZERO_PATTERN 660 661 Level: advanced 662 663 Notes: 664 Information about the matrices structure is used in MatCompositeMerge() for additive composite matrix. 665 666 .seealso: `MatAXPY()`, `MatCreateComposite()`, `MatCompositeMerge()` `MatCompositeGetMatStructure()`, `MATCOMPOSITE` 667 668 @*/ 669 PetscErrorCode MatCompositeSetMatStructure(Mat mat,MatStructure str) 670 { 671 PetscFunctionBegin; 672 PetscValidHeaderSpecific(mat,MAT_CLASSID,1); 673 PetscUseMethod(mat,"MatCompositeSetMatStructure_C",(Mat,MatStructure),(mat,str)); 674 PetscFunctionReturn(0); 675 } 676 677 static PetscErrorCode MatCompositeGetMatStructure_Composite(Mat mat,MatStructure *str) 678 { 679 Mat_Composite *b = (Mat_Composite*)mat->data; 680 681 PetscFunctionBegin; 682 *str = b->structure; 683 PetscFunctionReturn(0); 684 } 685 686 /*@ 687 MatCompositeGetMatStructure - Returns the structure of matrices in the composite matrix. 688 689 Not Collective 690 691 Input Parameter: 692 . mat - the composite matrix 693 694 Output Parameter: 695 . str - structure of the matrices 696 697 Level: advanced 698 699 .seealso: `MatCreateComposite()`, `MatCompositeSetMatStructure()`, `MATCOMPOSITE` 700 701 @*/ 702 PetscErrorCode MatCompositeGetMatStructure(Mat mat,MatStructure *str) 703 { 704 PetscFunctionBegin; 705 PetscValidHeaderSpecific(mat,MAT_CLASSID,1); 706 PetscValidPointer(str,2); 707 PetscUseMethod(mat,"MatCompositeGetMatStructure_C",(Mat,MatStructure*),(mat,str)); 708 PetscFunctionReturn(0); 709 } 710 711 static PetscErrorCode MatCompositeSetMergeType_Composite(Mat mat,MatCompositeMergeType type) 712 { 713 Mat_Composite *shell = (Mat_Composite*)mat->data; 714 715 PetscFunctionBegin; 716 shell->mergetype = type; 717 PetscFunctionReturn(0); 718 } 719 720 /*@ 721 MatCompositeSetMergeType - Sets order of MatCompositeMerge(). 722 723 Logically Collective on Mat 724 725 Input Parameters: 726 + mat - the composite matrix 727 - type - MAT_COMPOSITE_MERGE RIGHT (default) to start merge from right with the first added matrix (mat[0]), 728 MAT_COMPOSITE_MERGE_LEFT to start merge from left with the last added matrix (mat[nmat-1]) 729 730 Level: advanced 731 732 Notes: 733 The resulting matrix is the same regardles of the MergeType. Only the order of operation is changed. 734 If set to MAT_COMPOSITE_MERGE_RIGHT the order of the merge is mat[nmat-1]*(mat[nmat-2]*(...*(mat[1]*mat[0]))) 735 otherwise the order is (((mat[nmat-1]*mat[nmat-2])*mat[nmat-3])*...)*mat[0]. 736 737 .seealso: `MatCreateComposite()`, `MatCompositeMerge()`, `MATCOMPOSITE` 738 739 @*/ 740 PetscErrorCode MatCompositeSetMergeType(Mat mat,MatCompositeMergeType type) 741 { 742 PetscFunctionBegin; 743 PetscValidHeaderSpecific(mat,MAT_CLASSID,1); 744 PetscValidLogicalCollectiveEnum(mat,type,2); 745 PetscUseMethod(mat,"MatCompositeSetMergeType_C",(Mat,MatCompositeMergeType),(mat,type)); 746 PetscFunctionReturn(0); 747 } 748 749 static PetscErrorCode MatCompositeMerge_Composite(Mat mat) 750 { 751 Mat_Composite *shell = (Mat_Composite*)mat->data; 752 Mat_CompositeLink next = shell->head, prev = shell->tail; 753 Mat tmat,newmat; 754 Vec left,right; 755 PetscScalar scale; 756 PetscInt i; 757 758 PetscFunctionBegin; 759 PetscCheck(next,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Must provide at least one matrix with MatCompositeAddMat()"); 760 scale = shell->scale; 761 if (shell->type == MAT_COMPOSITE_ADDITIVE) { 762 if (shell->mergetype == MAT_COMPOSITE_MERGE_RIGHT) { 763 i = 0; 764 PetscCall(MatDuplicate(next->mat,MAT_COPY_VALUES,&tmat)); 765 if (shell->scalings) PetscCall(MatScale(tmat,shell->scalings[i++])); 766 while ((next = next->next)) { 767 PetscCall(MatAXPY(tmat,(shell->scalings ? shell->scalings[i++] : 1.0),next->mat,shell->structure)); 768 } 769 } else { 770 i = shell->nmat-1; 771 PetscCall(MatDuplicate(prev->mat,MAT_COPY_VALUES,&tmat)); 772 if (shell->scalings) PetscCall(MatScale(tmat,shell->scalings[i--])); 773 while ((prev = prev->prev)) { 774 PetscCall(MatAXPY(tmat,(shell->scalings ? shell->scalings[i--] : 1.0),prev->mat,shell->structure)); 775 } 776 } 777 } else { 778 if (shell->mergetype == MAT_COMPOSITE_MERGE_RIGHT) { 779 PetscCall(MatDuplicate(next->mat,MAT_COPY_VALUES,&tmat)); 780 while ((next = next->next)) { 781 PetscCall(MatMatMult(next->mat,tmat,MAT_INITIAL_MATRIX,PETSC_DECIDE,&newmat)); 782 PetscCall(MatDestroy(&tmat)); 783 tmat = newmat; 784 } 785 } else { 786 PetscCall(MatDuplicate(prev->mat,MAT_COPY_VALUES,&tmat)); 787 while ((prev = prev->prev)) { 788 PetscCall(MatMatMult(tmat,prev->mat,MAT_INITIAL_MATRIX,PETSC_DECIDE,&newmat)); 789 PetscCall(MatDestroy(&tmat)); 790 tmat = newmat; 791 } 792 } 793 if (shell->scalings) {for (i=0; i<shell->nmat; i++) scale *= shell->scalings[i];} 794 } 795 796 if ((left = shell->left)) PetscCall(PetscObjectReference((PetscObject)left)); 797 if ((right = shell->right)) PetscCall(PetscObjectReference((PetscObject)right)); 798 799 PetscCall(MatHeaderReplace(mat,&tmat)); 800 801 PetscCall(MatDiagonalScale(mat,left,right)); 802 PetscCall(MatScale(mat,scale)); 803 PetscCall(VecDestroy(&left)); 804 PetscCall(VecDestroy(&right)); 805 PetscFunctionReturn(0); 806 } 807 808 /*@ 809 MatCompositeMerge - Given a composite matrix, replaces it with a "regular" matrix 810 by summing or computing the product of all the matrices inside the composite matrix. 811 812 Collective 813 814 Input Parameter: 815 . mat - the composite matrix 816 817 Options Database Keys: 818 + -mat_composite_merge - merge in MatAssemblyEnd() 819 - -mat_composite_merge_type - set merge direction 820 821 Level: advanced 822 823 Notes: 824 The MatType of the resulting matrix will be the same as the MatType of the FIRST 825 matrix in the composite matrix. 826 827 .seealso: `MatDestroy()`, `MatMult()`, `MatCompositeAddMat()`, `MatCreateComposite()`, `MatCompositeSetMatStructure()`, `MatCompositeSetMergeType()`, `MATCOMPOSITE` 828 829 @*/ 830 PetscErrorCode MatCompositeMerge(Mat mat) 831 { 832 PetscFunctionBegin; 833 PetscValidHeaderSpecific(mat,MAT_CLASSID,1); 834 PetscUseMethod(mat,"MatCompositeMerge_C",(Mat),(mat)); 835 PetscFunctionReturn(0); 836 } 837 838 static PetscErrorCode MatCompositeGetNumberMat_Composite(Mat mat,PetscInt *nmat) 839 { 840 Mat_Composite *shell = (Mat_Composite*)mat->data; 841 842 PetscFunctionBegin; 843 *nmat = shell->nmat; 844 PetscFunctionReturn(0); 845 } 846 847 /*@ 848 MatCompositeGetNumberMat - Returns the number of matrices in the composite matrix. 849 850 Not Collective 851 852 Input Parameter: 853 . mat - the composite matrix 854 855 Output Parameter: 856 . nmat - number of matrices in the composite matrix 857 858 Level: advanced 859 860 .seealso: `MatCreateComposite()`, `MatCompositeGetMat()`, `MATCOMPOSITE` 861 862 @*/ 863 PetscErrorCode MatCompositeGetNumberMat(Mat mat,PetscInt *nmat) 864 { 865 PetscFunctionBegin; 866 PetscValidHeaderSpecific(mat,MAT_CLASSID,1); 867 PetscValidIntPointer(nmat,2); 868 PetscUseMethod(mat,"MatCompositeGetNumberMat_C",(Mat,PetscInt*),(mat,nmat)); 869 PetscFunctionReturn(0); 870 } 871 872 static PetscErrorCode MatCompositeGetMat_Composite(Mat mat,PetscInt i,Mat *Ai) 873 { 874 Mat_Composite *shell = (Mat_Composite*)mat->data; 875 Mat_CompositeLink ilink; 876 PetscInt k; 877 878 PetscFunctionBegin; 879 PetscCheck(i < shell->nmat,PetscObjectComm((PetscObject)mat),PETSC_ERR_ARG_OUTOFRANGE,"index out of range: %" PetscInt_FMT " >= %" PetscInt_FMT,i,shell->nmat); 880 ilink = shell->head; 881 for (k=0; k<i; k++) { 882 ilink = ilink->next; 883 } 884 *Ai = ilink->mat; 885 PetscFunctionReturn(0); 886 } 887 888 /*@ 889 MatCompositeGetMat - Returns the ith matrix from the composite matrix. 890 891 Logically Collective on Mat 892 893 Input Parameters: 894 + mat - the composite matrix 895 - i - the number of requested matrix 896 897 Output Parameter: 898 . Ai - ith matrix in composite 899 900 Level: advanced 901 902 .seealso: `MatCreateComposite()`, `MatCompositeGetNumberMat()`, `MatCompositeAddMat()`, `MATCOMPOSITE` 903 904 @*/ 905 PetscErrorCode MatCompositeGetMat(Mat mat,PetscInt i,Mat *Ai) 906 { 907 PetscFunctionBegin; 908 PetscValidHeaderSpecific(mat,MAT_CLASSID,1); 909 PetscValidLogicalCollectiveInt(mat,i,2); 910 PetscValidPointer(Ai,3); 911 PetscUseMethod(mat,"MatCompositeGetMat_C",(Mat,PetscInt,Mat*),(mat,i,Ai)); 912 PetscFunctionReturn(0); 913 } 914 915 PetscErrorCode MatCompositeSetScalings_Composite(Mat mat,const PetscScalar *scalings) 916 { 917 Mat_Composite *shell = (Mat_Composite*)mat->data; 918 PetscInt nmat; 919 920 PetscFunctionBegin; 921 PetscCall(MatCompositeGetNumberMat(mat,&nmat)); 922 if (!shell->scalings) PetscCall(PetscMalloc1(nmat,&shell->scalings)); 923 PetscCall(PetscArraycpy(shell->scalings,scalings,nmat)); 924 PetscFunctionReturn(0); 925 } 926 927 /*@ 928 MatCompositeSetScalings - Sets separate scaling factors for component matrices. 929 930 Logically Collective on Mat 931 932 Input Parameters: 933 + mat - the composite matrix 934 - scalings - array of scaling factors with scalings[i] being factor of i-th matrix, for i in [0, nmat) 935 936 Level: advanced 937 938 .seealso: `MatScale()`, `MatDiagonalScale()`, `MATCOMPOSITE` 939 940 @*/ 941 PetscErrorCode MatCompositeSetScalings(Mat mat,const PetscScalar *scalings) 942 { 943 PetscFunctionBegin; 944 PetscValidHeaderSpecific(mat,MAT_CLASSID,1); 945 PetscValidScalarPointer(scalings,2); 946 PetscValidLogicalCollectiveScalar(mat,*scalings,2); 947 PetscUseMethod(mat,"MatCompositeSetScalings_C",(Mat,const PetscScalar*),(mat,scalings)); 948 PetscFunctionReturn(0); 949 } 950 951 static struct _MatOps MatOps_Values = {NULL, 952 NULL, 953 NULL, 954 MatMult_Composite, 955 MatMultAdd_Composite, 956 /* 5*/ MatMultTranspose_Composite, 957 MatMultTransposeAdd_Composite, 958 NULL, 959 NULL, 960 NULL, 961 /* 10*/ NULL, 962 NULL, 963 NULL, 964 NULL, 965 NULL, 966 /* 15*/ NULL, 967 NULL, 968 MatGetDiagonal_Composite, 969 MatDiagonalScale_Composite, 970 NULL, 971 /* 20*/ NULL, 972 MatAssemblyEnd_Composite, 973 NULL, 974 NULL, 975 /* 24*/ NULL, 976 NULL, 977 NULL, 978 NULL, 979 NULL, 980 /* 29*/ NULL, 981 NULL, 982 NULL, 983 NULL, 984 NULL, 985 /* 34*/ NULL, 986 NULL, 987 NULL, 988 NULL, 989 NULL, 990 /* 39*/ NULL, 991 NULL, 992 NULL, 993 NULL, 994 NULL, 995 /* 44*/ NULL, 996 MatScale_Composite, 997 MatShift_Basic, 998 NULL, 999 NULL, 1000 /* 49*/ NULL, 1001 NULL, 1002 NULL, 1003 NULL, 1004 NULL, 1005 /* 54*/ NULL, 1006 NULL, 1007 NULL, 1008 NULL, 1009 NULL, 1010 /* 59*/ NULL, 1011 MatDestroy_Composite, 1012 NULL, 1013 NULL, 1014 NULL, 1015 /* 64*/ NULL, 1016 NULL, 1017 NULL, 1018 NULL, 1019 NULL, 1020 /* 69*/ NULL, 1021 NULL, 1022 NULL, 1023 NULL, 1024 NULL, 1025 /* 74*/ NULL, 1026 NULL, 1027 MatSetFromOptions_Composite, 1028 NULL, 1029 NULL, 1030 /* 79*/ NULL, 1031 NULL, 1032 NULL, 1033 NULL, 1034 NULL, 1035 /* 84*/ NULL, 1036 NULL, 1037 NULL, 1038 NULL, 1039 NULL, 1040 /* 89*/ NULL, 1041 NULL, 1042 NULL, 1043 NULL, 1044 NULL, 1045 /* 94*/ NULL, 1046 NULL, 1047 NULL, 1048 NULL, 1049 NULL, 1050 /*99*/ NULL, 1051 NULL, 1052 NULL, 1053 NULL, 1054 NULL, 1055 /*104*/ NULL, 1056 NULL, 1057 NULL, 1058 NULL, 1059 NULL, 1060 /*109*/ NULL, 1061 NULL, 1062 NULL, 1063 NULL, 1064 NULL, 1065 /*114*/ NULL, 1066 NULL, 1067 NULL, 1068 NULL, 1069 NULL, 1070 /*119*/ NULL, 1071 NULL, 1072 NULL, 1073 NULL, 1074 NULL, 1075 /*124*/ NULL, 1076 NULL, 1077 NULL, 1078 NULL, 1079 NULL, 1080 /*129*/ NULL, 1081 NULL, 1082 NULL, 1083 NULL, 1084 NULL, 1085 /*134*/ NULL, 1086 NULL, 1087 NULL, 1088 NULL, 1089 NULL, 1090 /*139*/ NULL, 1091 NULL, 1092 NULL, 1093 NULL, 1094 NULL, 1095 /*144*/NULL, 1096 NULL, 1097 NULL, 1098 NULL 1099 }; 1100 1101 /*MC 1102 MATCOMPOSITE - A matrix defined by the sum (or product) of one or more matrices. 1103 The matrices need to have a correct size and parallel layout for the sum or product to be valid. 1104 1105 Notes: 1106 to use the product of the matrices call MatCompositeSetType(mat,MAT_COMPOSITE_MULTIPLICATIVE); 1107 1108 Level: advanced 1109 1110 .seealso: `MatCreateComposite()`, `MatCompositeSetScalings()`, `MatCompositeAddMat()`, `MatSetType()`, `MatCompositeSetType()`, `MatCompositeGetType()`, `MatCompositeSetMatStructure()`, `MatCompositeGetMatStructure()`, `MatCompositeMerge()`, `MatCompositeSetMergeType()`, `MatCompositeGetNumberMat()`, `MatCompositeGetMat()` 1111 M*/ 1112 1113 PETSC_EXTERN PetscErrorCode MatCreate_Composite(Mat A) 1114 { 1115 Mat_Composite *b; 1116 1117 PetscFunctionBegin; 1118 PetscCall(PetscNewLog(A,&b)); 1119 A->data = (void*)b; 1120 PetscCall(PetscMemcpy(A->ops,&MatOps_Values,sizeof(struct _MatOps))); 1121 1122 PetscCall(PetscLayoutSetUp(A->rmap)); 1123 PetscCall(PetscLayoutSetUp(A->cmap)); 1124 1125 A->assembled = PETSC_TRUE; 1126 A->preallocated = PETSC_TRUE; 1127 b->type = MAT_COMPOSITE_ADDITIVE; 1128 b->scale = 1.0; 1129 b->nmat = 0; 1130 b->merge = PETSC_FALSE; 1131 b->mergetype = MAT_COMPOSITE_MERGE_RIGHT; 1132 b->structure = DIFFERENT_NONZERO_PATTERN; 1133 b->merge_mvctx = PETSC_TRUE; 1134 1135 PetscCall(PetscObjectComposeFunction((PetscObject)A,"MatCompositeAddMat_C",MatCompositeAddMat_Composite)); 1136 PetscCall(PetscObjectComposeFunction((PetscObject)A,"MatCompositeSetType_C",MatCompositeSetType_Composite)); 1137 PetscCall(PetscObjectComposeFunction((PetscObject)A,"MatCompositeGetType_C",MatCompositeGetType_Composite)); 1138 PetscCall(PetscObjectComposeFunction((PetscObject)A,"MatCompositeSetMergeType_C",MatCompositeSetMergeType_Composite)); 1139 PetscCall(PetscObjectComposeFunction((PetscObject)A,"MatCompositeSetMatStructure_C",MatCompositeSetMatStructure_Composite)); 1140 PetscCall(PetscObjectComposeFunction((PetscObject)A,"MatCompositeGetMatStructure_C",MatCompositeGetMatStructure_Composite)); 1141 PetscCall(PetscObjectComposeFunction((PetscObject)A,"MatCompositeMerge_C",MatCompositeMerge_Composite)); 1142 PetscCall(PetscObjectComposeFunction((PetscObject)A,"MatCompositeGetNumberMat_C",MatCompositeGetNumberMat_Composite)); 1143 PetscCall(PetscObjectComposeFunction((PetscObject)A,"MatCompositeGetMat_C",MatCompositeGetMat_Composite)); 1144 PetscCall(PetscObjectComposeFunction((PetscObject)A,"MatCompositeSetScalings_C",MatCompositeSetScalings_Composite)); 1145 1146 PetscCall(PetscObjectChangeTypeName((PetscObject)A,MATCOMPOSITE)); 1147 PetscFunctionReturn(0); 1148 } 1149