Actual source code: matnest.c

      1: #include <../src/mat/impls/nest/matnestimpl.h>
      2: #include <../src/mat/impls/aij/seq/aij.h>
      3: #include <../src/mat/impls/shell/shell.h>
      4: #include <petscsf.h>

      6: static PetscErrorCode MatSetUp_NestIS_Private(Mat, PetscInt, const IS[], PetscInt, const IS[]);
      7: static PetscErrorCode MatCreateVecs_Nest(Mat, Vec *, Vec *);
      8: static PetscErrorCode MatReset_Nest(Mat);

     10: PETSC_INTERN PetscErrorCode MatConvert_Nest_IS(Mat, MatType, MatReuse, Mat *);

     12: /* private functions */
     13: static PetscErrorCode MatNestGetSizes_Private(Mat A, PetscInt *m, PetscInt *n, PetscInt *M, PetscInt *N)
     14: {
     15:   Mat_Nest *bA = (Mat_Nest *)A->data;
     16:   PetscInt  i, j;

     18:   PetscFunctionBegin;
     19:   *m = *n = *M = *N = 0;
     20:   for (i = 0; i < bA->nr; i++) { /* rows */
     21:     PetscInt sm, sM;
     22:     PetscCall(ISGetLocalSize(bA->isglobal.row[i], &sm));
     23:     PetscCall(ISGetSize(bA->isglobal.row[i], &sM));
     24:     *m += sm;
     25:     *M += sM;
     26:   }
     27:   for (j = 0; j < bA->nc; j++) { /* cols */
     28:     PetscInt sn, sN;
     29:     PetscCall(ISGetLocalSize(bA->isglobal.col[j], &sn));
     30:     PetscCall(ISGetSize(bA->isglobal.col[j], &sN));
     31:     *n += sn;
     32:     *N += sN;
     33:   }
     34:   PetscFunctionReturn(PETSC_SUCCESS);
     35: }

     37: /* operations */
     38: static PetscErrorCode MatMult_Nest(Mat A, Vec x, Vec y)
     39: {
     40:   Mat_Nest *bA = (Mat_Nest *)A->data;
     41:   Vec      *bx = bA->right, *by = bA->left;
     42:   PetscInt  i, j, nr = bA->nr, nc = bA->nc;

     44:   PetscFunctionBegin;
     45:   for (i = 0; i < nr; i++) PetscCall(VecGetSubVector(y, bA->isglobal.row[i], &by[i]));
     46:   for (i = 0; i < nc; i++) PetscCall(VecGetSubVector(x, bA->isglobal.col[i], &bx[i]));
     47:   for (i = 0; i < nr; i++) {
     48:     PetscCall(VecZeroEntries(by[i]));
     49:     for (j = 0; j < nc; j++) {
     50:       if (!bA->m[i][j]) continue;
     51:       /* y[i] <- y[i] + A[i][j] * x[j] */
     52:       PetscCall(MatMultAdd(bA->m[i][j], bx[j], by[i], by[i]));
     53:     }
     54:   }
     55:   for (i = 0; i < nr; i++) PetscCall(VecRestoreSubVector(y, bA->isglobal.row[i], &by[i]));
     56:   for (i = 0; i < nc; i++) PetscCall(VecRestoreSubVector(x, bA->isglobal.col[i], &bx[i]));
     57:   PetscFunctionReturn(PETSC_SUCCESS);
     58: }

     60: static PetscErrorCode MatMultAdd_Nest(Mat A, Vec x, Vec y, Vec z)
     61: {
     62:   Mat_Nest *bA = (Mat_Nest *)A->data;
     63:   Vec      *bx = bA->right, *bz = bA->left;
     64:   PetscInt  i, j, nr = bA->nr, nc = bA->nc;

     66:   PetscFunctionBegin;
     67:   for (i = 0; i < nr; i++) PetscCall(VecGetSubVector(z, bA->isglobal.row[i], &bz[i]));
     68:   for (i = 0; i < nc; i++) PetscCall(VecGetSubVector(x, bA->isglobal.col[i], &bx[i]));
     69:   for (i = 0; i < nr; i++) {
     70:     if (y != z) {
     71:       Vec by;
     72:       PetscCall(VecGetSubVector(y, bA->isglobal.row[i], &by));
     73:       PetscCall(VecCopy(by, bz[i]));
     74:       PetscCall(VecRestoreSubVector(y, bA->isglobal.row[i], &by));
     75:     }
     76:     for (j = 0; j < nc; j++) {
     77:       if (!bA->m[i][j]) continue;
     78:       /* y[i] <- y[i] + A[i][j] * x[j] */
     79:       PetscCall(MatMultAdd(bA->m[i][j], bx[j], bz[i], bz[i]));
     80:     }
     81:   }
     82:   for (i = 0; i < nr; i++) PetscCall(VecRestoreSubVector(z, bA->isglobal.row[i], &bz[i]));
     83:   for (i = 0; i < nc; i++) PetscCall(VecRestoreSubVector(x, bA->isglobal.col[i], &bx[i]));
     84:   PetscFunctionReturn(PETSC_SUCCESS);
     85: }

     87: typedef struct {
     88:   Mat         *workC;      /* array of Mat with specific containers depending on the underlying MatMatMult implementation */
     89:   PetscScalar *tarray;     /* buffer for storing all temporary products A[i][j] B[j] */
     90:   PetscInt    *dm, *dn, k; /* displacements and number of submatrices */
     91: } Nest_Dense;

     93: static PetscErrorCode MatProductNumeric_Nest_Dense(Mat C)
     94: {
     95:   Mat_Nest          *bA;
     96:   Nest_Dense        *contents;
     97:   Mat                viewB, viewC, productB, workC;
     98:   const PetscScalar *barray;
     99:   PetscScalar       *carray;
    100:   PetscInt           i, j, M, N, nr, nc, ldb, ldc;
    101:   Mat                A, B;

    103:   PetscFunctionBegin;
    104:   MatCheckProduct(C, 1);
    105:   A = C->product->A;
    106:   B = C->product->B;
    107:   PetscCall(MatGetSize(B, NULL, &N));
    108:   if (!N) {
    109:     PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
    110:     PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
    111:     PetscFunctionReturn(PETSC_SUCCESS);
    112:   }
    113:   contents = (Nest_Dense *)C->product->data;
    114:   PetscCheck(contents, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
    115:   bA = (Mat_Nest *)A->data;
    116:   nr = bA->nr;
    117:   nc = bA->nc;
    118:   PetscCall(MatDenseGetLDA(B, &ldb));
    119:   PetscCall(MatDenseGetLDA(C, &ldc));
    120:   PetscCall(MatZeroEntries(C));
    121:   PetscCall(MatDenseGetArrayRead(B, &barray));
    122:   PetscCall(MatDenseGetArray(C, &carray));
    123:   for (i = 0; i < nr; i++) {
    124:     PetscCall(ISGetSize(bA->isglobal.row[i], &M));
    125:     PetscCall(MatCreateDense(PetscObjectComm((PetscObject)A), contents->dm[i + 1] - contents->dm[i], PETSC_DECIDE, M, N, PetscSafePointerPlusOffset(carray, contents->dm[i]), &viewC));
    126:     PetscCall(MatDenseSetLDA(viewC, ldc));
    127:     for (j = 0; j < nc; j++) {
    128:       if (!bA->m[i][j]) continue;
    129:       PetscCall(ISGetSize(bA->isglobal.col[j], &M));
    130:       PetscCall(MatCreateDense(PetscObjectComm((PetscObject)A), contents->dn[j + 1] - contents->dn[j], PETSC_DECIDE, M, N, PetscSafePointerPlusOffset((PetscScalar *)barray, contents->dn[j]), &viewB));
    131:       PetscCall(MatDenseSetLDA(viewB, ldb));

    133:       /* MatMatMultNumeric(bA->m[i][j],viewB,contents->workC[i*nc + j]); */
    134:       workC             = contents->workC[i * nc + j];
    135:       productB          = workC->product->B;
    136:       workC->product->B = viewB; /* use newly created dense matrix viewB */
    137:       PetscCall(MatProductNumeric(workC));
    138:       PetscCall(MatDestroy(&viewB));
    139:       workC->product->B = productB; /* resume original B */

    141:       /* C[i] <- workC + C[i] */
    142:       PetscCall(MatAXPY(viewC, 1.0, contents->workC[i * nc + j], SAME_NONZERO_PATTERN));
    143:     }
    144:     PetscCall(MatDestroy(&viewC));
    145:   }
    146:   PetscCall(MatDenseRestoreArray(C, &carray));
    147:   PetscCall(MatDenseRestoreArrayRead(B, &barray));

    149:   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
    150:   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
    151:   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
    152:   PetscFunctionReturn(PETSC_SUCCESS);
    153: }

    155: static PetscErrorCode MatNest_DenseDestroy(void *ctx)
    156: {
    157:   Nest_Dense *contents = (Nest_Dense *)ctx;
    158:   PetscInt    i;

    160:   PetscFunctionBegin;
    161:   PetscCall(PetscFree(contents->tarray));
    162:   for (i = 0; i < contents->k; i++) PetscCall(MatDestroy(contents->workC + i));
    163:   PetscCall(PetscFree3(contents->dm, contents->dn, contents->workC));
    164:   PetscCall(PetscFree(contents));
    165:   PetscFunctionReturn(PETSC_SUCCESS);
    166: }

    168: static PetscErrorCode MatProductSymbolic_Nest_Dense(Mat C)
    169: {
    170:   Mat_Nest          *bA;
    171:   Mat                viewB, workC;
    172:   const PetscScalar *barray;
    173:   PetscInt           i, j, M, N, m, n, nr, nc, maxm = 0, ldb;
    174:   Nest_Dense        *contents = NULL;
    175:   PetscBool          cisdense;
    176:   Mat                A, B;
    177:   PetscReal          fill;

    179:   PetscFunctionBegin;
    180:   MatCheckProduct(C, 1);
    181:   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
    182:   A    = C->product->A;
    183:   B    = C->product->B;
    184:   fill = C->product->fill;
    185:   bA   = (Mat_Nest *)A->data;
    186:   nr   = bA->nr;
    187:   nc   = bA->nc;
    188:   PetscCall(MatGetLocalSize(C, &m, &n));
    189:   PetscCall(MatGetSize(C, &M, &N));
    190:   if (m == PETSC_DECIDE || n == PETSC_DECIDE || M == PETSC_DECIDE || N == PETSC_DECIDE) {
    191:     PetscCall(MatGetLocalSize(B, NULL, &n));
    192:     PetscCall(MatGetSize(B, NULL, &N));
    193:     PetscCall(MatGetLocalSize(A, &m, NULL));
    194:     PetscCall(MatGetSize(A, &M, NULL));
    195:     PetscCall(MatSetSizes(C, m, n, M, N));
    196:   }
    197:   PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATSEQDENSE, MATMPIDENSE, MATSEQDENSECUDA, MATMPIDENSECUDA, ""));
    198:   if (!cisdense) PetscCall(MatSetType(C, ((PetscObject)B)->type_name));
    199:   PetscCall(MatSetUp(C));
    200:   if (!N) {
    201:     C->ops->productnumeric = MatProductNumeric_Nest_Dense;
    202:     PetscFunctionReturn(PETSC_SUCCESS);
    203:   }

    205:   PetscCall(PetscNew(&contents));
    206:   C->product->data    = contents;
    207:   C->product->destroy = MatNest_DenseDestroy;
    208:   PetscCall(PetscCalloc3(nr + 1, &contents->dm, nc + 1, &contents->dn, nr * nc, &contents->workC));
    209:   contents->k = nr * nc;
    210:   for (i = 0; i < nr; i++) {
    211:     PetscCall(ISGetLocalSize(bA->isglobal.row[i], contents->dm + i + 1));
    212:     maxm = PetscMax(maxm, contents->dm[i + 1]);
    213:     contents->dm[i + 1] += contents->dm[i];
    214:   }
    215:   for (i = 0; i < nc; i++) {
    216:     PetscCall(ISGetLocalSize(bA->isglobal.col[i], contents->dn + i + 1));
    217:     contents->dn[i + 1] += contents->dn[i];
    218:   }
    219:   PetscCall(PetscMalloc1(maxm * N, &contents->tarray));
    220:   PetscCall(MatDenseGetLDA(B, &ldb));
    221:   PetscCall(MatGetSize(B, NULL, &N));
    222:   PetscCall(MatDenseGetArrayRead(B, &barray));
    223:   /* loops are permuted compared to MatMatMultNumeric so that viewB is created only once per column of A */
    224:   for (j = 0; j < nc; j++) {
    225:     PetscCall(ISGetSize(bA->isglobal.col[j], &M));
    226:     PetscCall(MatCreateDense(PetscObjectComm((PetscObject)A), contents->dn[j + 1] - contents->dn[j], PETSC_DECIDE, M, N, PetscSafePointerPlusOffset((PetscScalar *)barray, contents->dn[j]), &viewB));
    227:     PetscCall(MatDenseSetLDA(viewB, ldb));
    228:     for (i = 0; i < nr; i++) {
    229:       if (!bA->m[i][j]) continue;
    230:       /* MatMatMultSymbolic may attach a specific container (depending on MatType of bA->m[i][j]) to workC[i][j] */

    232:       PetscCall(MatProductCreate(bA->m[i][j], viewB, NULL, &contents->workC[i * nc + j]));
    233:       workC = contents->workC[i * nc + j];
    234:       PetscCall(MatProductSetType(workC, MATPRODUCT_AB));
    235:       PetscCall(MatProductSetAlgorithm(workC, "default"));
    236:       PetscCall(MatProductSetFill(workC, fill));
    237:       PetscCall(MatProductSetFromOptions(workC));
    238:       PetscCall(MatProductSymbolic(workC));

    240:       /* since tarray will be shared by all Mat */
    241:       PetscCall(MatSeqDenseSetPreallocation(workC, contents->tarray));
    242:       PetscCall(MatMPIDenseSetPreallocation(workC, contents->tarray));
    243:     }
    244:     PetscCall(MatDestroy(&viewB));
    245:   }
    246:   PetscCall(MatDenseRestoreArrayRead(B, &barray));

    248:   C->ops->productnumeric = MatProductNumeric_Nest_Dense;
    249:   PetscFunctionReturn(PETSC_SUCCESS);
    250: }

    252: static PetscErrorCode MatProductSetFromOptions_Nest_Dense(Mat C)
    253: {
    254:   Mat_Product *product = C->product;

    256:   PetscFunctionBegin;
    257:   if (product->type == MATPRODUCT_AB) C->ops->productsymbolic = MatProductSymbolic_Nest_Dense;
    258:   PetscFunctionReturn(PETSC_SUCCESS);
    259: }

    261: static PetscErrorCode MatMultTransposeKernel_Nest(Mat A, Vec x, Vec y, PetscBool herm)
    262: {
    263:   Mat_Nest *bA = (Mat_Nest *)A->data;
    264:   Vec      *bx = bA->left, *by = bA->right;
    265:   PetscInt  i, j, nr = bA->nr, nc = bA->nc;

    267:   PetscFunctionBegin;
    268:   for (i = 0; i < nr; i++) PetscCall(VecGetSubVector(x, bA->isglobal.row[i], &bx[i]));
    269:   for (i = 0; i < nc; i++) PetscCall(VecGetSubVector(y, bA->isglobal.col[i], &by[i]));
    270:   for (j = 0; j < nc; j++) {
    271:     PetscCall(VecZeroEntries(by[j]));
    272:     for (i = 0; i < nr; i++) {
    273:       if (!bA->m[i][j]) continue;
    274:       if (herm) PetscCall(MatMultHermitianTransposeAdd(bA->m[i][j], bx[i], by[j], by[j])); /* y[j] <- y[j] + (A[i][j])^H * x[i] */
    275:       else PetscCall(MatMultTransposeAdd(bA->m[i][j], bx[i], by[j], by[j]));               /* y[j] <- y[j] + (A[i][j])^T * x[i] */
    276:     }
    277:   }
    278:   for (i = 0; i < nr; i++) PetscCall(VecRestoreSubVector(x, bA->isglobal.row[i], &bx[i]));
    279:   for (i = 0; i < nc; i++) PetscCall(VecRestoreSubVector(y, bA->isglobal.col[i], &by[i]));
    280:   PetscFunctionReturn(PETSC_SUCCESS);
    281: }

    283: static PetscErrorCode MatMultTranspose_Nest(Mat A, Vec x, Vec y)
    284: {
    285:   PetscFunctionBegin;
    286:   PetscCall(MatMultTransposeKernel_Nest(A, x, y, PETSC_FALSE));
    287:   PetscFunctionReturn(PETSC_SUCCESS);
    288: }

    290: static PetscErrorCode MatMultHermitianTranspose_Nest(Mat A, Vec x, Vec y)
    291: {
    292:   PetscFunctionBegin;
    293:   PetscCall(MatMultTransposeKernel_Nest(A, x, y, PETSC_TRUE));
    294:   PetscFunctionReturn(PETSC_SUCCESS);
    295: }

    297: static PetscErrorCode MatMultTransposeAddKernel_Nest(Mat A, Vec x, Vec y, Vec z, PetscBool herm)
    298: {
    299:   Mat_Nest *bA = (Mat_Nest *)A->data;
    300:   Vec      *bx = bA->left, *bz = bA->right;
    301:   PetscInt  i, j, nr = bA->nr, nc = bA->nc;

    303:   PetscFunctionBegin;
    304:   for (i = 0; i < nr; i++) PetscCall(VecGetSubVector(x, bA->isglobal.row[i], &bx[i]));
    305:   for (i = 0; i < nc; i++) PetscCall(VecGetSubVector(z, bA->isglobal.col[i], &bz[i]));
    306:   for (j = 0; j < nc; j++) {
    307:     if (y != z) {
    308:       Vec by;
    309:       PetscCall(VecGetSubVector(y, bA->isglobal.col[j], &by));
    310:       PetscCall(VecCopy(by, bz[j]));
    311:       PetscCall(VecRestoreSubVector(y, bA->isglobal.col[j], &by));
    312:     }
    313:     for (i = 0; i < nr; i++) {
    314:       if (!bA->m[i][j]) continue;
    315:       if (herm) PetscCall(MatMultHermitianTransposeAdd(bA->m[i][j], bx[i], bz[j], bz[j])); /* z[j] <- y[j] + (A[i][j])^H * x[i] */
    316:       else PetscCall(MatMultTransposeAdd(bA->m[i][j], bx[i], bz[j], bz[j]));               /* z[j] <- y[j] + (A[i][j])^T * x[i] */
    317:     }
    318:   }
    319:   for (i = 0; i < nr; i++) PetscCall(VecRestoreSubVector(x, bA->isglobal.row[i], &bx[i]));
    320:   for (i = 0; i < nc; i++) PetscCall(VecRestoreSubVector(z, bA->isglobal.col[i], &bz[i]));
    321:   PetscFunctionReturn(PETSC_SUCCESS);
    322: }

    324: static PetscErrorCode MatMultTransposeAdd_Nest(Mat A, Vec x, Vec y, Vec z)
    325: {
    326:   PetscFunctionBegin;
    327:   PetscCall(MatMultTransposeAddKernel_Nest(A, x, y, z, PETSC_FALSE));
    328:   PetscFunctionReturn(PETSC_SUCCESS);
    329: }

    331: static PetscErrorCode MatMultHermitianTransposeAdd_Nest(Mat A, Vec x, Vec y, Vec z)
    332: {
    333:   PetscFunctionBegin;
    334:   PetscCall(MatMultTransposeAddKernel_Nest(A, x, y, z, PETSC_TRUE));
    335:   PetscFunctionReturn(PETSC_SUCCESS);
    336: }

    338: static PetscErrorCode MatTranspose_Nest(Mat A, MatReuse reuse, Mat *B)
    339: {
    340:   Mat_Nest *bA = (Mat_Nest *)A->data, *bC;
    341:   Mat       C;
    342:   PetscInt  i, j, nr = bA->nr, nc = bA->nc;

    344:   PetscFunctionBegin;
    345:   if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(A, *B));
    346:   PetscCheck(reuse != MAT_INPLACE_MATRIX || nr == nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_SIZ, "Square nested matrix only for in-place");

    348:   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_INPLACE_MATRIX) {
    349:     Mat *subs;
    350:     IS  *is_row, *is_col;

    352:     PetscCall(PetscCalloc1(nr * nc, &subs));
    353:     PetscCall(PetscMalloc2(nr, &is_row, nc, &is_col));
    354:     PetscCall(MatNestGetISs(A, is_row, is_col));
    355:     if (reuse == MAT_INPLACE_MATRIX) {
    356:       for (i = 0; i < nr; i++) {
    357:         for (j = 0; j < nc; j++) subs[i + nr * j] = bA->m[i][j];
    358:       }
    359:     }

    361:     PetscCall(MatCreateNest(PetscObjectComm((PetscObject)A), nc, is_col, nr, is_row, subs, &C));
    362:     PetscCall(PetscFree(subs));
    363:     PetscCall(PetscFree2(is_row, is_col));
    364:   } else {
    365:     C = *B;
    366:   }

    368:   bC = (Mat_Nest *)C->data;
    369:   for (i = 0; i < nr; i++) {
    370:     for (j = 0; j < nc; j++) {
    371:       if (bA->m[i][j]) {
    372:         PetscCall(MatTranspose(bA->m[i][j], reuse, &bC->m[j][i]));
    373:       } else {
    374:         bC->m[j][i] = NULL;
    375:       }
    376:     }
    377:   }

    379:   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_REUSE_MATRIX) {
    380:     *B = C;
    381:   } else {
    382:     PetscCall(MatHeaderMerge(A, &C));
    383:   }
    384:   PetscFunctionReturn(PETSC_SUCCESS);
    385: }

    387: static PetscErrorCode MatNestDestroyISList(PetscInt n, IS **list)
    388: {
    389:   IS      *lst = *list;
    390:   PetscInt i;

    392:   PetscFunctionBegin;
    393:   if (!lst) PetscFunctionReturn(PETSC_SUCCESS);
    394:   for (i = 0; i < n; i++)
    395:     if (lst[i]) PetscCall(ISDestroy(&lst[i]));
    396:   PetscCall(PetscFree(lst));
    397:   *list = NULL;
    398:   PetscFunctionReturn(PETSC_SUCCESS);
    399: }

    401: static PetscErrorCode MatReset_Nest(Mat A)
    402: {
    403:   Mat_Nest *vs = (Mat_Nest *)A->data;
    404:   PetscInt  i, j;

    406:   PetscFunctionBegin;
    407:   /* release the matrices and the place holders */
    408:   PetscCall(MatNestDestroyISList(vs->nr, &vs->isglobal.row));
    409:   PetscCall(MatNestDestroyISList(vs->nc, &vs->isglobal.col));
    410:   PetscCall(MatNestDestroyISList(vs->nr, &vs->islocal.row));
    411:   PetscCall(MatNestDestroyISList(vs->nc, &vs->islocal.col));

    413:   PetscCall(PetscFree(vs->row_len));
    414:   PetscCall(PetscFree(vs->col_len));
    415:   PetscCall(PetscFree(vs->nnzstate));

    417:   PetscCall(PetscFree2(vs->left, vs->right));

    419:   /* release the matrices and the place holders */
    420:   if (vs->m) {
    421:     for (i = 0; i < vs->nr; i++) {
    422:       for (j = 0; j < vs->nc; j++) PetscCall(MatDestroy(&vs->m[i][j]));
    423:     }
    424:     PetscCall(PetscFree(vs->m[0]));
    425:     PetscCall(PetscFree(vs->m));
    426:   }

    428:   /* restore defaults */
    429:   vs->nr            = 0;
    430:   vs->nc            = 0;
    431:   vs->splitassembly = PETSC_FALSE;
    432:   PetscFunctionReturn(PETSC_SUCCESS);
    433: }

    435: static PetscErrorCode MatDestroy_Nest(Mat A)
    436: {
    437:   PetscFunctionBegin;
    438:   PetscCall(MatReset_Nest(A));
    439:   PetscCall(PetscFree(A->data));
    440:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSubMat_C", NULL));
    441:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetSubMat_C", NULL));
    442:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSubMats_C", NULL));
    443:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSize_C", NULL));
    444:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetISs_C", NULL));
    445:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetLocalISs_C", NULL));
    446:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetVecType_C", NULL));
    447:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetSubMats_C", NULL));
    448:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_mpiaij_C", NULL));
    449:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_seqaij_C", NULL));
    450:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_aij_C", NULL));
    451:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_is_C", NULL));
    452:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_mpidense_C", NULL));
    453:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_seqdense_C", NULL));
    454:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_nest_seqdense_C", NULL));
    455:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_nest_mpidense_C", NULL));
    456:   PetscFunctionReturn(PETSC_SUCCESS);
    457: }

    459: static PetscErrorCode MatMissingDiagonal_Nest(Mat mat, PetscBool *missing, PetscInt *dd)
    460: {
    461:   Mat_Nest *vs = (Mat_Nest *)mat->data;
    462:   PetscInt  i;

    464:   PetscFunctionBegin;
    465:   if (dd) *dd = 0;
    466:   if (!vs->nr) {
    467:     *missing = PETSC_TRUE;
    468:     PetscFunctionReturn(PETSC_SUCCESS);
    469:   }
    470:   *missing = PETSC_FALSE;
    471:   for (i = 0; i < vs->nr && !*missing; i++) {
    472:     *missing = PETSC_TRUE;
    473:     if (vs->m[i][i]) {
    474:       PetscCall(MatMissingDiagonal(vs->m[i][i], missing, NULL));
    475:       PetscCheck(!*missing || !dd, PetscObjectComm((PetscObject)mat), PETSC_ERR_SUP, "First missing entry not yet implemented");
    476:     }
    477:   }
    478:   PetscFunctionReturn(PETSC_SUCCESS);
    479: }

    481: static PetscErrorCode MatAssemblyBegin_Nest(Mat A, MatAssemblyType type)
    482: {
    483:   Mat_Nest *vs = (Mat_Nest *)A->data;
    484:   PetscInt  i, j;
    485:   PetscBool nnzstate = PETSC_FALSE;

    487:   PetscFunctionBegin;
    488:   for (i = 0; i < vs->nr; i++) {
    489:     for (j = 0; j < vs->nc; j++) {
    490:       PetscObjectState subnnzstate = 0;
    491:       if (vs->m[i][j]) {
    492:         PetscCall(MatAssemblyBegin(vs->m[i][j], type));
    493:         if (!vs->splitassembly) {
    494:           /* Note: split assembly will fail if the same block appears more than once (even indirectly through a nested
    495:            * sub-block). This could be fixed by adding a flag to Mat so that there was a way to check if a Mat was
    496:            * already performing an assembly, but the result would by more complicated and appears to offer less
    497:            * potential for diagnostics and correctness checking. Split assembly should be fixed once there is an
    498:            * interface for libraries to make asynchronous progress in "user-defined non-blocking collectives".
    499:            */
    500:           PetscCall(MatAssemblyEnd(vs->m[i][j], type));
    501:           PetscCall(MatGetNonzeroState(vs->m[i][j], &subnnzstate));
    502:         }
    503:       }
    504:       nnzstate                     = (PetscBool)(nnzstate || vs->nnzstate[i * vs->nc + j] != subnnzstate);
    505:       vs->nnzstate[i * vs->nc + j] = subnnzstate;
    506:     }
    507:   }
    508:   if (nnzstate) A->nonzerostate++;
    509:   PetscFunctionReturn(PETSC_SUCCESS);
    510: }

    512: static PetscErrorCode MatAssemblyEnd_Nest(Mat A, MatAssemblyType type)
    513: {
    514:   Mat_Nest *vs = (Mat_Nest *)A->data;
    515:   PetscInt  i, j;

    517:   PetscFunctionBegin;
    518:   for (i = 0; i < vs->nr; i++) {
    519:     for (j = 0; j < vs->nc; j++) {
    520:       if (vs->m[i][j]) {
    521:         if (vs->splitassembly) PetscCall(MatAssemblyEnd(vs->m[i][j], type));
    522:       }
    523:     }
    524:   }
    525:   PetscFunctionReturn(PETSC_SUCCESS);
    526: }

    528: static PetscErrorCode MatNestFindNonzeroSubMatRow(Mat A, PetscInt row, Mat *B)
    529: {
    530:   Mat_Nest *vs = (Mat_Nest *)A->data;
    531:   PetscInt  j;
    532:   Mat       sub;

    534:   PetscFunctionBegin;
    535:   sub = (row < vs->nc) ? vs->m[row][row] : (Mat)NULL; /* Prefer to find on the diagonal */
    536:   for (j = 0; !sub && j < vs->nc; j++) sub = vs->m[row][j];
    537:   if (sub) PetscCall(MatSetUp(sub)); /* Ensure that the sizes are available */
    538:   *B = sub;
    539:   PetscFunctionReturn(PETSC_SUCCESS);
    540: }

    542: static PetscErrorCode MatNestFindNonzeroSubMatCol(Mat A, PetscInt col, Mat *B)
    543: {
    544:   Mat_Nest *vs = (Mat_Nest *)A->data;
    545:   PetscInt  i;
    546:   Mat       sub;

    548:   PetscFunctionBegin;
    549:   sub = (col < vs->nr) ? vs->m[col][col] : (Mat)NULL; /* Prefer to find on the diagonal */
    550:   for (i = 0; !sub && i < vs->nr; i++) sub = vs->m[i][col];
    551:   if (sub) PetscCall(MatSetUp(sub)); /* Ensure that the sizes are available */
    552:   *B = sub;
    553:   PetscFunctionReturn(PETSC_SUCCESS);
    554: }

    556: static PetscErrorCode MatNestFindISRange(Mat A, PetscInt n, const IS list[], IS is, PetscInt *begin, PetscInt *end)
    557: {
    558:   PetscInt  i, j, size, m;
    559:   PetscBool flg;
    560:   IS        out, concatenate[2];

    562:   PetscFunctionBegin;
    563:   PetscAssertPointer(list, 3);
    565:   if (begin) {
    566:     PetscAssertPointer(begin, 5);
    567:     *begin = -1;
    568:   }
    569:   if (end) {
    570:     PetscAssertPointer(end, 6);
    571:     *end = -1;
    572:   }
    573:   for (i = 0; i < n; i++) {
    574:     if (!list[i]) continue;
    575:     PetscCall(ISEqualUnsorted(list[i], is, &flg));
    576:     if (flg) {
    577:       if (begin) *begin = i;
    578:       if (end) *end = i + 1;
    579:       PetscFunctionReturn(PETSC_SUCCESS);
    580:     }
    581:   }
    582:   PetscCall(ISGetSize(is, &size));
    583:   for (i = 0; i < n - 1; i++) {
    584:     if (!list[i]) continue;
    585:     m = 0;
    586:     PetscCall(ISConcatenate(PetscObjectComm((PetscObject)A), 2, list + i, &out));
    587:     PetscCall(ISGetSize(out, &m));
    588:     for (j = i + 2; j < n && m < size; j++) {
    589:       if (list[j]) {
    590:         concatenate[0] = out;
    591:         concatenate[1] = list[j];
    592:         PetscCall(ISConcatenate(PetscObjectComm((PetscObject)A), 2, concatenate, &out));
    593:         PetscCall(ISDestroy(concatenate));
    594:         PetscCall(ISGetSize(out, &m));
    595:       }
    596:     }
    597:     if (m == size) {
    598:       PetscCall(ISEqualUnsorted(out, is, &flg));
    599:       if (flg) {
    600:         if (begin) *begin = i;
    601:         if (end) *end = j;
    602:         PetscCall(ISDestroy(&out));
    603:         PetscFunctionReturn(PETSC_SUCCESS);
    604:       }
    605:     }
    606:     PetscCall(ISDestroy(&out));
    607:   }
    608:   PetscFunctionReturn(PETSC_SUCCESS);
    609: }

    611: static PetscErrorCode MatNestFillEmptyMat_Private(Mat A, PetscInt i, PetscInt j, Mat *B)
    612: {
    613:   Mat_Nest *vs = (Mat_Nest *)A->data;
    614:   PetscInt  lr, lc;

    616:   PetscFunctionBegin;
    617:   PetscCall(MatCreate(PetscObjectComm((PetscObject)A), B));
    618:   PetscCall(ISGetLocalSize(vs->isglobal.row[i], &lr));
    619:   PetscCall(ISGetLocalSize(vs->isglobal.col[j], &lc));
    620:   PetscCall(MatSetSizes(*B, lr, lc, PETSC_DECIDE, PETSC_DECIDE));
    621:   PetscCall(MatSetType(*B, MATAIJ));
    622:   PetscCall(MatSeqAIJSetPreallocation(*B, 0, NULL));
    623:   PetscCall(MatMPIAIJSetPreallocation(*B, 0, NULL, 0, NULL));
    624:   PetscCall(MatSetUp(*B));
    625:   PetscCall(MatSetOption(*B, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
    626:   PetscCall(MatAssemblyBegin(*B, MAT_FINAL_ASSEMBLY));
    627:   PetscCall(MatAssemblyEnd(*B, MAT_FINAL_ASSEMBLY));
    628:   PetscFunctionReturn(PETSC_SUCCESS);
    629: }

    631: static PetscErrorCode MatNestGetBlock_Private(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *B)
    632: {
    633:   Mat_Nest  *vs = (Mat_Nest *)A->data;
    634:   Mat       *a;
    635:   PetscInt   i, j, k, l, nr = rend - rbegin, nc = cend - cbegin;
    636:   char       keyname[256];
    637:   PetscBool *b;
    638:   PetscBool  flg;

    640:   PetscFunctionBegin;
    641:   *B = NULL;
    642:   PetscCall(PetscSNPrintf(keyname, sizeof(keyname), "NestBlock_%" PetscInt_FMT "-%" PetscInt_FMT "x%" PetscInt_FMT "-%" PetscInt_FMT, rbegin, rend, cbegin, cend));
    643:   PetscCall(PetscObjectQuery((PetscObject)A, keyname, (PetscObject *)B));
    644:   if (*B) PetscFunctionReturn(PETSC_SUCCESS);

    646:   PetscCall(PetscMalloc2(nr * nc, &a, nr * nc, &b));
    647:   for (i = 0; i < nr; i++) {
    648:     for (j = 0; j < nc; j++) {
    649:       a[i * nc + j] = vs->m[rbegin + i][cbegin + j];
    650:       b[i * nc + j] = PETSC_FALSE;
    651:     }
    652:   }
    653:   if (nc != vs->nc && nr != vs->nr) {
    654:     for (i = 0; i < nr; i++) {
    655:       for (j = 0; j < nc; j++) {
    656:         flg = PETSC_FALSE;
    657:         for (k = 0; (k < nr && !flg); k++) {
    658:           if (a[j + k * nc]) flg = PETSC_TRUE;
    659:         }
    660:         if (flg) {
    661:           flg = PETSC_FALSE;
    662:           for (l = 0; (l < nc && !flg); l++) {
    663:             if (a[i * nc + l]) flg = PETSC_TRUE;
    664:           }
    665:         }
    666:         if (!flg) {
    667:           b[i * nc + j] = PETSC_TRUE;
    668:           PetscCall(MatNestFillEmptyMat_Private(A, rbegin + i, cbegin + j, a + i * nc + j));
    669:         }
    670:       }
    671:     }
    672:   }
    673:   PetscCall(MatCreateNest(PetscObjectComm((PetscObject)A), nr, nr != vs->nr ? NULL : vs->isglobal.row, nc, nc != vs->nc ? NULL : vs->isglobal.col, a, B));
    674:   for (i = 0; i < nr; i++) {
    675:     for (j = 0; j < nc; j++) {
    676:       if (b[i * nc + j]) PetscCall(MatDestroy(a + i * nc + j));
    677:     }
    678:   }
    679:   PetscCall(PetscFree2(a, b));
    680:   (*B)->assembled = A->assembled;
    681:   PetscCall(PetscObjectCompose((PetscObject)A, keyname, (PetscObject)*B));
    682:   PetscCall(PetscObjectDereference((PetscObject)*B)); /* Leave the only remaining reference in the composition */
    683:   PetscFunctionReturn(PETSC_SUCCESS);
    684: }

    686: static PetscErrorCode MatNestFindSubMat(Mat A, struct MatNestISPair *is, IS isrow, IS iscol, Mat *B)
    687: {
    688:   Mat_Nest *vs = (Mat_Nest *)A->data;
    689:   PetscInt  rbegin, rend, cbegin, cend;

    691:   PetscFunctionBegin;
    692:   PetscCall(MatNestFindISRange(A, vs->nr, is->row, isrow, &rbegin, &rend));
    693:   PetscCall(MatNestFindISRange(A, vs->nc, is->col, iscol, &cbegin, &cend));
    694:   if (rend == rbegin + 1 && cend == cbegin + 1) {
    695:     if (!vs->m[rbegin][cbegin]) PetscCall(MatNestFillEmptyMat_Private(A, rbegin, cbegin, vs->m[rbegin] + cbegin));
    696:     *B = vs->m[rbegin][cbegin];
    697:   } else if (rbegin != -1 && cbegin != -1) {
    698:     PetscCall(MatNestGetBlock_Private(A, rbegin, rend, cbegin, cend, B));
    699:   } else SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_INCOMP, "Could not find index set");
    700:   PetscFunctionReturn(PETSC_SUCCESS);
    701: }

    703: /*
    704:    TODO: This does not actually returns a submatrix we can modify
    705: */
    706: static PetscErrorCode MatCreateSubMatrix_Nest(Mat A, IS isrow, IS iscol, MatReuse reuse, Mat *B)
    707: {
    708:   Mat_Nest *vs = (Mat_Nest *)A->data;
    709:   Mat       sub;

    711:   PetscFunctionBegin;
    712:   PetscCall(MatNestFindSubMat(A, &vs->isglobal, isrow, iscol, &sub));
    713:   switch (reuse) {
    714:   case MAT_INITIAL_MATRIX:
    715:     if (sub) PetscCall(PetscObjectReference((PetscObject)sub));
    716:     *B = sub;
    717:     break;
    718:   case MAT_REUSE_MATRIX:
    719:     PetscCheck(sub == *B, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Submatrix was not used before in this call");
    720:     break;
    721:   case MAT_IGNORE_MATRIX: /* Nothing to do */
    722:     break;
    723:   case MAT_INPLACE_MATRIX: /* Nothing to do */
    724:     SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "MAT_INPLACE_MATRIX is not supported yet");
    725:   }
    726:   PetscFunctionReturn(PETSC_SUCCESS);
    727: }

    729: static PetscErrorCode MatGetLocalSubMatrix_Nest(Mat A, IS isrow, IS iscol, Mat *B)
    730: {
    731:   Mat_Nest *vs = (Mat_Nest *)A->data;
    732:   Mat       sub;

    734:   PetscFunctionBegin;
    735:   PetscCall(MatNestFindSubMat(A, &vs->islocal, isrow, iscol, &sub));
    736:   /* We allow the submatrix to be NULL, perhaps it would be better for the user to return an empty matrix instead */
    737:   if (sub) PetscCall(PetscObjectReference((PetscObject)sub));
    738:   *B = sub;
    739:   PetscFunctionReturn(PETSC_SUCCESS);
    740: }

    742: static PetscErrorCode MatRestoreLocalSubMatrix_Nest(Mat A, IS isrow, IS iscol, Mat *B)
    743: {
    744:   Mat_Nest *vs = (Mat_Nest *)A->data;
    745:   Mat       sub;

    747:   PetscFunctionBegin;
    748:   PetscCall(MatNestFindSubMat(A, &vs->islocal, isrow, iscol, &sub));
    749:   PetscCheck(*B == sub, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Local submatrix has not been gotten");
    750:   if (sub) {
    751:     PetscCheck(((PetscObject)sub)->refct > 1, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Local submatrix has had reference count decremented too many times");
    752:     PetscCall(MatDestroy(B));
    753:   }
    754:   PetscFunctionReturn(PETSC_SUCCESS);
    755: }

    757: static PetscErrorCode MatGetDiagonal_Nest(Mat A, Vec v)
    758: {
    759:   Mat_Nest *bA = (Mat_Nest *)A->data;
    760:   PetscInt  i;

    762:   PetscFunctionBegin;
    763:   for (i = 0; i < bA->nr; i++) {
    764:     Vec bv;
    765:     PetscCall(VecGetSubVector(v, bA->isglobal.row[i], &bv));
    766:     if (bA->m[i][i]) {
    767:       PetscCall(MatGetDiagonal(bA->m[i][i], bv));
    768:     } else {
    769:       PetscCall(VecSet(bv, 0.0));
    770:     }
    771:     PetscCall(VecRestoreSubVector(v, bA->isglobal.row[i], &bv));
    772:   }
    773:   PetscFunctionReturn(PETSC_SUCCESS);
    774: }

    776: static PetscErrorCode MatDiagonalScale_Nest(Mat A, Vec l, Vec r)
    777: {
    778:   Mat_Nest *bA = (Mat_Nest *)A->data;
    779:   Vec       bl, *br;
    780:   PetscInt  i, j;

    782:   PetscFunctionBegin;
    783:   PetscCall(PetscCalloc1(bA->nc, &br));
    784:   if (r) {
    785:     for (j = 0; j < bA->nc; j++) PetscCall(VecGetSubVector(r, bA->isglobal.col[j], &br[j]));
    786:   }
    787:   bl = NULL;
    788:   for (i = 0; i < bA->nr; i++) {
    789:     if (l) PetscCall(VecGetSubVector(l, bA->isglobal.row[i], &bl));
    790:     for (j = 0; j < bA->nc; j++) {
    791:       if (bA->m[i][j]) PetscCall(MatDiagonalScale(bA->m[i][j], bl, br[j]));
    792:     }
    793:     if (l) PetscCall(VecRestoreSubVector(l, bA->isglobal.row[i], &bl));
    794:   }
    795:   if (r) {
    796:     for (j = 0; j < bA->nc; j++) PetscCall(VecRestoreSubVector(r, bA->isglobal.col[j], &br[j]));
    797:   }
    798:   PetscCall(PetscFree(br));
    799:   PetscFunctionReturn(PETSC_SUCCESS);
    800: }

    802: static PetscErrorCode MatScale_Nest(Mat A, PetscScalar a)
    803: {
    804:   Mat_Nest *bA = (Mat_Nest *)A->data;
    805:   PetscInt  i, j;

    807:   PetscFunctionBegin;
    808:   for (i = 0; i < bA->nr; i++) {
    809:     for (j = 0; j < bA->nc; j++) {
    810:       if (bA->m[i][j]) PetscCall(MatScale(bA->m[i][j], a));
    811:     }
    812:   }
    813:   PetscFunctionReturn(PETSC_SUCCESS);
    814: }

    816: static PetscErrorCode MatShift_Nest(Mat A, PetscScalar a)
    817: {
    818:   Mat_Nest *bA = (Mat_Nest *)A->data;
    819:   PetscInt  i;
    820:   PetscBool nnzstate = PETSC_FALSE;

    822:   PetscFunctionBegin;
    823:   for (i = 0; i < bA->nr; i++) {
    824:     PetscObjectState subnnzstate = 0;
    825:     PetscCheck(bA->m[i][i], PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "No support for shifting an empty diagonal block, insert a matrix in block (%" PetscInt_FMT ",%" PetscInt_FMT ")", i, i);
    826:     PetscCall(MatShift(bA->m[i][i], a));
    827:     PetscCall(MatGetNonzeroState(bA->m[i][i], &subnnzstate));
    828:     nnzstate                     = (PetscBool)(nnzstate || bA->nnzstate[i * bA->nc + i] != subnnzstate);
    829:     bA->nnzstate[i * bA->nc + i] = subnnzstate;
    830:   }
    831:   if (nnzstate) A->nonzerostate++;
    832:   PetscFunctionReturn(PETSC_SUCCESS);
    833: }

    835: static PetscErrorCode MatDiagonalSet_Nest(Mat A, Vec D, InsertMode is)
    836: {
    837:   Mat_Nest *bA = (Mat_Nest *)A->data;
    838:   PetscInt  i;
    839:   PetscBool nnzstate = PETSC_FALSE;

    841:   PetscFunctionBegin;
    842:   for (i = 0; i < bA->nr; i++) {
    843:     PetscObjectState subnnzstate = 0;
    844:     Vec              bv;
    845:     PetscCall(VecGetSubVector(D, bA->isglobal.row[i], &bv));
    846:     if (bA->m[i][i]) {
    847:       PetscCall(MatDiagonalSet(bA->m[i][i], bv, is));
    848:       PetscCall(MatGetNonzeroState(bA->m[i][i], &subnnzstate));
    849:     }
    850:     PetscCall(VecRestoreSubVector(D, bA->isglobal.row[i], &bv));
    851:     nnzstate                     = (PetscBool)(nnzstate || bA->nnzstate[i * bA->nc + i] != subnnzstate);
    852:     bA->nnzstate[i * bA->nc + i] = subnnzstate;
    853:   }
    854:   if (nnzstate) A->nonzerostate++;
    855:   PetscFunctionReturn(PETSC_SUCCESS);
    856: }

    858: static PetscErrorCode MatSetRandom_Nest(Mat A, PetscRandom rctx)
    859: {
    860:   Mat_Nest *bA = (Mat_Nest *)A->data;
    861:   PetscInt  i, j;

    863:   PetscFunctionBegin;
    864:   for (i = 0; i < bA->nr; i++) {
    865:     for (j = 0; j < bA->nc; j++) {
    866:       if (bA->m[i][j]) PetscCall(MatSetRandom(bA->m[i][j], rctx));
    867:     }
    868:   }
    869:   PetscFunctionReturn(PETSC_SUCCESS);
    870: }

    872: static PetscErrorCode MatCreateVecs_Nest(Mat A, Vec *right, Vec *left)
    873: {
    874:   Mat_Nest *bA = (Mat_Nest *)A->data;
    875:   Vec      *L, *R;
    876:   MPI_Comm  comm;
    877:   PetscInt  i, j;

    879:   PetscFunctionBegin;
    880:   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
    881:   if (right) {
    882:     /* allocate R */
    883:     PetscCall(PetscMalloc1(bA->nc, &R));
    884:     /* Create the right vectors */
    885:     for (j = 0; j < bA->nc; j++) {
    886:       for (i = 0; i < bA->nr; i++) {
    887:         if (bA->m[i][j]) {
    888:           PetscCall(MatCreateVecs(bA->m[i][j], &R[j], NULL));
    889:           break;
    890:         }
    891:       }
    892:       PetscCheck(i != bA->nr, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Mat(Nest) contains a null column.");
    893:     }
    894:     PetscCall(VecCreateNest(comm, bA->nc, bA->isglobal.col, R, right));
    895:     /* hand back control to the nest vector */
    896:     for (j = 0; j < bA->nc; j++) PetscCall(VecDestroy(&R[j]));
    897:     PetscCall(PetscFree(R));
    898:   }

    900:   if (left) {
    901:     /* allocate L */
    902:     PetscCall(PetscMalloc1(bA->nr, &L));
    903:     /* Create the left vectors */
    904:     for (i = 0; i < bA->nr; i++) {
    905:       for (j = 0; j < bA->nc; j++) {
    906:         if (bA->m[i][j]) {
    907:           PetscCall(MatCreateVecs(bA->m[i][j], NULL, &L[i]));
    908:           break;
    909:         }
    910:       }
    911:       PetscCheck(j != bA->nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Mat(Nest) contains a null row.");
    912:     }

    914:     PetscCall(VecCreateNest(comm, bA->nr, bA->isglobal.row, L, left));
    915:     for (i = 0; i < bA->nr; i++) PetscCall(VecDestroy(&L[i]));

    917:     PetscCall(PetscFree(L));
    918:   }
    919:   PetscFunctionReturn(PETSC_SUCCESS);
    920: }

    922: static PetscErrorCode MatView_Nest(Mat A, PetscViewer viewer)
    923: {
    924:   Mat_Nest *bA = (Mat_Nest *)A->data;
    925:   PetscBool isascii, viewSub = PETSC_FALSE;
    926:   PetscInt  i, j;

    928:   PetscFunctionBegin;
    929:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
    930:   if (isascii) {
    931:     PetscCall(PetscOptionsGetBool(((PetscObject)A)->options, ((PetscObject)A)->prefix, "-mat_view_nest_sub", &viewSub, NULL));
    932:     PetscCall(PetscViewerASCIIPrintf(viewer, "Matrix object:\n"));
    933:     PetscCall(PetscViewerASCIIPushTab(viewer));
    934:     PetscCall(PetscViewerASCIIPrintf(viewer, "type=nest, rows=%" PetscInt_FMT ", cols=%" PetscInt_FMT "\n", bA->nr, bA->nc));

    936:     PetscCall(PetscViewerASCIIPrintf(viewer, "MatNest structure:\n"));
    937:     for (i = 0; i < bA->nr; i++) {
    938:       for (j = 0; j < bA->nc; j++) {
    939:         MatType   type;
    940:         char      name[256] = "", prefix[256] = "";
    941:         PetscInt  NR, NC;
    942:         PetscBool isNest = PETSC_FALSE;

    944:         if (!bA->m[i][j]) {
    945:           PetscCall(PetscViewerASCIIPrintf(viewer, "(%" PetscInt_FMT ",%" PetscInt_FMT ") : NULL\n", i, j));
    946:           continue;
    947:         }
    948:         PetscCall(MatGetSize(bA->m[i][j], &NR, &NC));
    949:         PetscCall(MatGetType(bA->m[i][j], &type));
    950:         if (((PetscObject)bA->m[i][j])->name) PetscCall(PetscSNPrintf(name, sizeof(name), "name=\"%s\", ", ((PetscObject)bA->m[i][j])->name));
    951:         if (((PetscObject)bA->m[i][j])->prefix) PetscCall(PetscSNPrintf(prefix, sizeof(prefix), "prefix=\"%s\", ", ((PetscObject)bA->m[i][j])->prefix));
    952:         PetscCall(PetscObjectTypeCompare((PetscObject)bA->m[i][j], MATNEST, &isNest));

    954:         PetscCall(PetscViewerASCIIPrintf(viewer, "(%" PetscInt_FMT ",%" PetscInt_FMT ") : %s%stype=%s, rows=%" PetscInt_FMT ", cols=%" PetscInt_FMT "\n", i, j, name, prefix, type, NR, NC));

    956:         if (isNest || viewSub) {
    957:           PetscCall(PetscViewerASCIIPushTab(viewer)); /* push1 */
    958:           PetscCall(MatView(bA->m[i][j], viewer));
    959:           PetscCall(PetscViewerASCIIPopTab(viewer)); /* pop1 */
    960:         }
    961:       }
    962:     }
    963:     PetscCall(PetscViewerASCIIPopTab(viewer)); /* pop0 */
    964:   }
    965:   PetscFunctionReturn(PETSC_SUCCESS);
    966: }

    968: static PetscErrorCode MatZeroEntries_Nest(Mat A)
    969: {
    970:   Mat_Nest *bA = (Mat_Nest *)A->data;
    971:   PetscInt  i, j;

    973:   PetscFunctionBegin;
    974:   for (i = 0; i < bA->nr; i++) {
    975:     for (j = 0; j < bA->nc; j++) {
    976:       if (!bA->m[i][j]) continue;
    977:       PetscCall(MatZeroEntries(bA->m[i][j]));
    978:     }
    979:   }
    980:   PetscFunctionReturn(PETSC_SUCCESS);
    981: }

    983: static PetscErrorCode MatCopy_Nest(Mat A, Mat B, MatStructure str)
    984: {
    985:   Mat_Nest *bA = (Mat_Nest *)A->data, *bB = (Mat_Nest *)B->data;
    986:   PetscInt  i, j, nr = bA->nr, nc = bA->nc;
    987:   PetscBool nnzstate = PETSC_FALSE;

    989:   PetscFunctionBegin;
    990:   PetscCheck(nr == bB->nr && nc == bB->nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_INCOMP, "Cannot copy a Mat_Nest of block size (%" PetscInt_FMT ",%" PetscInt_FMT ") to a Mat_Nest of block size (%" PetscInt_FMT ",%" PetscInt_FMT ")", bB->nr, bB->nc, nr, nc);
    991:   for (i = 0; i < nr; i++) {
    992:     for (j = 0; j < nc; j++) {
    993:       PetscObjectState subnnzstate = 0;
    994:       if (bA->m[i][j] && bB->m[i][j]) {
    995:         PetscCall(MatCopy(bA->m[i][j], bB->m[i][j], str));
    996:       } else PetscCheck(!bA->m[i][j] && !bB->m[i][j], PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_INCOMP, "Matrix block does not exist at %" PetscInt_FMT ",%" PetscInt_FMT, i, j);
    997:       PetscCall(MatGetNonzeroState(bB->m[i][j], &subnnzstate));
    998:       nnzstate                 = (PetscBool)(nnzstate || bB->nnzstate[i * nc + j] != subnnzstate);
    999:       bB->nnzstate[i * nc + j] = subnnzstate;
    1000:     }
    1001:   }
    1002:   if (nnzstate) B->nonzerostate++;
    1003:   PetscFunctionReturn(PETSC_SUCCESS);
    1004: }

    1006: static PetscErrorCode MatAXPY_Nest(Mat Y, PetscScalar a, Mat X, MatStructure str)
    1007: {
    1008:   Mat_Nest *bY = (Mat_Nest *)Y->data, *bX = (Mat_Nest *)X->data;
    1009:   PetscInt  i, j, nr = bY->nr, nc = bY->nc;
    1010:   PetscBool nnzstate = PETSC_FALSE;

    1012:   PetscFunctionBegin;
    1013:   PetscCheck(nr == bX->nr && nc == bX->nc, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_INCOMP, "Cannot AXPY a MatNest of block size (%" PetscInt_FMT ",%" PetscInt_FMT ") with a MatNest of block size (%" PetscInt_FMT ",%" PetscInt_FMT ")", bX->nr, bX->nc, nr, nc);
    1014:   for (i = 0; i < nr; i++) {
    1015:     for (j = 0; j < nc; j++) {
    1016:       PetscObjectState subnnzstate = 0;
    1017:       if (bY->m[i][j] && bX->m[i][j]) {
    1018:         PetscCall(MatAXPY(bY->m[i][j], a, bX->m[i][j], str));
    1019:       } else if (bX->m[i][j]) {
    1020:         Mat M;

    1022:         PetscCheck(str == DIFFERENT_NONZERO_PATTERN || str == UNKNOWN_NONZERO_PATTERN, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_INCOMP, "Matrix block does not exist at %" PetscInt_FMT ",%" PetscInt_FMT ". Use DIFFERENT_NONZERO_PATTERN or UNKNOWN_NONZERO_PATTERN", i, j);
    1023:         PetscCall(MatDuplicate(bX->m[i][j], MAT_COPY_VALUES, &M));
    1024:         PetscCall(MatNestSetSubMat(Y, i, j, M));
    1025:         PetscCall(MatDestroy(&M));
    1026:       }
    1027:       if (bY->m[i][j]) PetscCall(MatGetNonzeroState(bY->m[i][j], &subnnzstate));
    1028:       nnzstate                 = (PetscBool)(nnzstate || bY->nnzstate[i * nc + j] != subnnzstate);
    1029:       bY->nnzstate[i * nc + j] = subnnzstate;
    1030:     }
    1031:   }
    1032:   if (nnzstate) Y->nonzerostate++;
    1033:   PetscFunctionReturn(PETSC_SUCCESS);
    1034: }

    1036: static PetscErrorCode MatDuplicate_Nest(Mat A, MatDuplicateOption op, Mat *B)
    1037: {
    1038:   Mat_Nest *bA = (Mat_Nest *)A->data;
    1039:   Mat      *b;
    1040:   PetscInt  i, j, nr = bA->nr, nc = bA->nc;

    1042:   PetscFunctionBegin;
    1043:   PetscCall(PetscMalloc1(nr * nc, &b));
    1044:   for (i = 0; i < nr; i++) {
    1045:     for (j = 0; j < nc; j++) {
    1046:       if (bA->m[i][j]) {
    1047:         PetscCall(MatDuplicate(bA->m[i][j], op, &b[i * nc + j]));
    1048:       } else {
    1049:         b[i * nc + j] = NULL;
    1050:       }
    1051:     }
    1052:   }
    1053:   PetscCall(MatCreateNest(PetscObjectComm((PetscObject)A), nr, bA->isglobal.row, nc, bA->isglobal.col, b, B));
    1054:   /* Give the new MatNest exclusive ownership */
    1055:   for (i = 0; i < nr * nc; i++) PetscCall(MatDestroy(&b[i]));
    1056:   PetscCall(PetscFree(b));

    1058:   PetscCall(MatAssemblyBegin(*B, MAT_FINAL_ASSEMBLY));
    1059:   PetscCall(MatAssemblyEnd(*B, MAT_FINAL_ASSEMBLY));
    1060:   PetscFunctionReturn(PETSC_SUCCESS);
    1061: }

    1063: /* nest api */
    1064: static PetscErrorCode MatNestGetSubMat_Nest(Mat A, PetscInt idxm, PetscInt jdxm, Mat *mat)
    1065: {
    1066:   Mat_Nest *bA = (Mat_Nest *)A->data;

    1068:   PetscFunctionBegin;
    1069:   PetscCheck(idxm < bA->nr, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Row too large: row %" PetscInt_FMT " max %" PetscInt_FMT, idxm, bA->nr - 1);
    1070:   PetscCheck(jdxm < bA->nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Col too large: row %" PetscInt_FMT " max %" PetscInt_FMT, jdxm, bA->nc - 1);
    1071:   *mat = bA->m[idxm][jdxm];
    1072:   PetscFunctionReturn(PETSC_SUCCESS);
    1073: }

    1075: /*@
    1076:   MatNestGetSubMat - Returns a single, sub-matrix from a `MATNEST`

    1078:   Not Collective

    1080:   Input Parameters:
    1081: + A    - `MATNEST` matrix
    1082: . idxm - index of the matrix within the nest matrix
    1083: - jdxm - index of the matrix within the nest matrix

    1085:   Output Parameter:
    1086: . sub - matrix at index `idxm`, `jdxm` within the nest matrix

    1088:   Level: developer

    1090: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSize()`, `MatNestGetSubMats()`, `MatCreateNest()`, `MatNestSetSubMat()`,
    1091:           `MatNestGetLocalISs()`, `MatNestGetISs()`
    1092: @*/
    1093: PetscErrorCode MatNestGetSubMat(Mat A, PetscInt idxm, PetscInt jdxm, Mat *sub)
    1094: {
    1095:   PetscFunctionBegin;
    1099:   PetscAssertPointer(sub, 4);
    1100:   PetscUseMethod(A, "MatNestGetSubMat_C", (Mat, PetscInt, PetscInt, Mat *), (A, idxm, jdxm, sub));
    1101:   PetscFunctionReturn(PETSC_SUCCESS);
    1102: }

    1104: static PetscErrorCode MatNestSetSubMat_Nest(Mat A, PetscInt idxm, PetscInt jdxm, Mat mat)
    1105: {
    1106:   Mat_Nest *bA = (Mat_Nest *)A->data;
    1107:   PetscInt  m, n, M, N, mi, ni, Mi, Ni;

    1109:   PetscFunctionBegin;
    1110:   PetscCheck(idxm < bA->nr, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Row too large: row %" PetscInt_FMT " max %" PetscInt_FMT, idxm, bA->nr - 1);
    1111:   PetscCheck(jdxm < bA->nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Col too large: row %" PetscInt_FMT " max %" PetscInt_FMT, jdxm, bA->nc - 1);
    1112:   if (mat) {
    1113:     PetscCall(MatGetLocalSize(mat, &m, &n));
    1114:     PetscCall(MatGetSize(mat, &M, &N));
    1115:     PetscCall(ISGetLocalSize(bA->isglobal.row[idxm], &mi));
    1116:     PetscCall(ISGetSize(bA->isglobal.row[idxm], &Mi));
    1117:     PetscCall(ISGetLocalSize(bA->isglobal.col[jdxm], &ni));
    1118:     PetscCall(ISGetSize(bA->isglobal.col[jdxm], &Ni));
    1119:     PetscCheck(M == Mi && N == Ni, PetscObjectComm((PetscObject)mat), PETSC_ERR_ARG_INCOMP, "Submatrix dimension (%" PetscInt_FMT ",%" PetscInt_FMT ") incompatible with nest block (%" PetscInt_FMT ",%" PetscInt_FMT ")", M, N, Mi, Ni);
    1120:     PetscCheck(m == mi && n == ni, PetscObjectComm((PetscObject)mat), PETSC_ERR_ARG_INCOMP, "Submatrix local dimension (%" PetscInt_FMT ",%" PetscInt_FMT ") incompatible with nest block (%" PetscInt_FMT ",%" PetscInt_FMT ")", m, n, mi, ni);
    1121:   }

    1123:   /* do not increase object state */
    1124:   if (mat == bA->m[idxm][jdxm]) PetscFunctionReturn(PETSC_SUCCESS);

    1126:   PetscCall(PetscObjectReference((PetscObject)mat));
    1127:   PetscCall(MatDestroy(&bA->m[idxm][jdxm]));
    1128:   bA->m[idxm][jdxm] = mat;
    1129:   PetscCall(PetscObjectStateIncrease((PetscObject)A));
    1130:   if (mat) PetscCall(MatGetNonzeroState(mat, &bA->nnzstate[idxm * bA->nc + jdxm]));
    1131:   else bA->nnzstate[idxm * bA->nc + jdxm] = 0;
    1132:   A->nonzerostate++;
    1133:   PetscFunctionReturn(PETSC_SUCCESS);
    1134: }

    1136: /*@
    1137:   MatNestSetSubMat - Set a single submatrix in the `MATNEST`

    1139:   Logically Collective

    1141:   Input Parameters:
    1142: + A    - `MATNEST` matrix
    1143: . idxm - index of the matrix within the nest matrix
    1144: . jdxm - index of the matrix within the nest matrix
    1145: - sub  - matrix at index `idxm`, `jdxm` within the nest matrix

    1147:   Level: developer

    1149:   Notes:
    1150:   The new submatrix must have the same size and communicator as that block of the nest.

    1152:   This increments the reference count of the submatrix.

    1154: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestSetSubMats()`, `MatNestGetSubMats()`, `MatNestGetLocalISs()`, `MatCreateNest()`,
    1155:           `MatNestGetSubMat()`, `MatNestGetISs()`, `MatNestGetSize()`
    1156: @*/
    1157: PetscErrorCode MatNestSetSubMat(Mat A, PetscInt idxm, PetscInt jdxm, Mat sub)
    1158: {
    1159:   PetscFunctionBegin;
    1164:   PetscTryMethod(A, "MatNestSetSubMat_C", (Mat, PetscInt, PetscInt, Mat), (A, idxm, jdxm, sub));
    1165:   PetscFunctionReturn(PETSC_SUCCESS);
    1166: }

    1168: static PetscErrorCode MatNestGetSubMats_Nest(Mat A, PetscInt *M, PetscInt *N, Mat ***mat)
    1169: {
    1170:   Mat_Nest *bA = (Mat_Nest *)A->data;

    1172:   PetscFunctionBegin;
    1173:   if (M) *M = bA->nr;
    1174:   if (N) *N = bA->nc;
    1175:   if (mat) *mat = bA->m;
    1176:   PetscFunctionReturn(PETSC_SUCCESS);
    1177: }

    1179: /*@C
    1180:   MatNestGetSubMats - Returns the entire two dimensional array of matrices defining a `MATNEST` matrix.

    1182:   Not Collective

    1184:   Input Parameter:
    1185: . A - nest matrix

    1187:   Output Parameters:
    1188: + M   - number of submatrix rows in the nest matrix
    1189: . N   - number of submatrix columns in the nest matrix
    1190: - mat - array of matrices

    1192:   Level: developer

    1194:   Note:
    1195:   The user should not free the array `mat`.

    1197:   Fortran Notes:
    1198:   This routine has a calling sequence `call MatNestGetSubMats(A, M, N, mat, ierr)`
    1199:   where the space allocated for the optional argument `mat` is assumed large enough (if provided).
    1200:   Matrices in `mat` are returned in row-major order, see `MatCreateNest()` for an example.

    1202: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSize()`, `MatNestGetSubMat()`, `MatNestGetLocalISs()`, `MatCreateNest()`,
    1203:           `MatNestSetSubMats()`, `MatNestGetISs()`, `MatNestSetSubMat()`
    1204: @*/
    1205: PetscErrorCode MatNestGetSubMats(Mat A, PetscInt *M, PetscInt *N, Mat ***mat)
    1206: {
    1207:   PetscFunctionBegin;
    1209:   PetscUseMethod(A, "MatNestGetSubMats_C", (Mat, PetscInt *, PetscInt *, Mat ***), (A, M, N, mat));
    1210:   PetscFunctionReturn(PETSC_SUCCESS);
    1211: }

    1213: static PetscErrorCode MatNestGetSize_Nest(Mat A, PetscInt *M, PetscInt *N)
    1214: {
    1215:   Mat_Nest *bA = (Mat_Nest *)A->data;

    1217:   PetscFunctionBegin;
    1218:   if (M) *M = bA->nr;
    1219:   if (N) *N = bA->nc;
    1220:   PetscFunctionReturn(PETSC_SUCCESS);
    1221: }

    1223: /*@
    1224:   MatNestGetSize - Returns the size of the `MATNEST` matrix.

    1226:   Not Collective

    1228:   Input Parameter:
    1229: . A - `MATNEST` matrix

    1231:   Output Parameters:
    1232: + M - number of rows in the nested mat
    1233: - N - number of cols in the nested mat

    1235:   Level: developer

    1237:   Note:
    1238:   `size` refers to the number of submatrices in the row and column directions of the nested matrix

    1240: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSubMat()`, `MatNestGetSubMats()`, `MatCreateNest()`, `MatNestGetLocalISs()`,
    1241:           `MatNestGetISs()`
    1242: @*/
    1243: PetscErrorCode MatNestGetSize(Mat A, PetscInt *M, PetscInt *N)
    1244: {
    1245:   PetscFunctionBegin;
    1247:   PetscUseMethod(A, "MatNestGetSize_C", (Mat, PetscInt *, PetscInt *), (A, M, N));
    1248:   PetscFunctionReturn(PETSC_SUCCESS);
    1249: }

    1251: static PetscErrorCode MatNestGetISs_Nest(Mat A, IS rows[], IS cols[])
    1252: {
    1253:   Mat_Nest *vs = (Mat_Nest *)A->data;
    1254:   PetscInt  i;

    1256:   PetscFunctionBegin;
    1257:   if (rows)
    1258:     for (i = 0; i < vs->nr; i++) rows[i] = vs->isglobal.row[i];
    1259:   if (cols)
    1260:     for (i = 0; i < vs->nc; i++) cols[i] = vs->isglobal.col[i];
    1261:   PetscFunctionReturn(PETSC_SUCCESS);
    1262: }

    1264: /*@
    1265:   MatNestGetISs - Returns the index sets partitioning the row and column spaces of a `MATNEST`

    1267:   Not Collective

    1269:   Input Parameter:
    1270: . A - `MATNEST` matrix

    1272:   Output Parameters:
    1273: + rows - array of row index sets (pass `NULL` to ignore)
    1274: - cols - array of column index sets (pass `NULL` to ignore)

    1276:   Level: advanced

    1278:   Note:
    1279:   The user must have allocated arrays of the correct size. The reference count is not increased on the returned `IS`s.

    1281: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSubMat()`, `MatNestGetSubMats()`, `MatNestGetSize()`, `MatNestGetLocalISs()`,
    1282:           `MatCreateNest()`, `MatNestSetSubMats()`
    1283: @*/
    1284: PetscErrorCode MatNestGetISs(Mat A, IS rows[], IS cols[])
    1285: {
    1286:   PetscFunctionBegin;
    1288:   PetscUseMethod(A, "MatNestGetISs_C", (Mat, IS[], IS[]), (A, rows, cols));
    1289:   PetscFunctionReturn(PETSC_SUCCESS);
    1290: }

    1292: static PetscErrorCode MatNestGetLocalISs_Nest(Mat A, IS rows[], IS cols[])
    1293: {
    1294:   Mat_Nest *vs = (Mat_Nest *)A->data;
    1295:   PetscInt  i;

    1297:   PetscFunctionBegin;
    1298:   if (rows)
    1299:     for (i = 0; i < vs->nr; i++) rows[i] = vs->islocal.row[i];
    1300:   if (cols)
    1301:     for (i = 0; i < vs->nc; i++) cols[i] = vs->islocal.col[i];
    1302:   PetscFunctionReturn(PETSC_SUCCESS);
    1303: }

    1305: /*@
    1306:   MatNestGetLocalISs - Returns the index sets partitioning the row and column spaces of a `MATNEST`

    1308:   Not Collective

    1310:   Input Parameter:
    1311: . A - `MATNEST` matrix

    1313:   Output Parameters:
    1314: + rows - array of row index sets (pass `NULL` to ignore)
    1315: - cols - array of column index sets (pass `NULL` to ignore)

    1317:   Level: advanced

    1319:   Note:
    1320:   The user must have allocated arrays of the correct size. The reference count is not increased on the returned `IS`s.

    1322: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSubMat()`, `MatNestGetSubMats()`, `MatNestGetSize()`, `MatNestGetISs()`, `MatCreateNest()`,
    1323:           `MatNestSetSubMats()`, `MatNestSetSubMat()`
    1324: @*/
    1325: PetscErrorCode MatNestGetLocalISs(Mat A, IS rows[], IS cols[])
    1326: {
    1327:   PetscFunctionBegin;
    1329:   PetscUseMethod(A, "MatNestGetLocalISs_C", (Mat, IS[], IS[]), (A, rows, cols));
    1330:   PetscFunctionReturn(PETSC_SUCCESS);
    1331: }

    1333: static PetscErrorCode MatNestSetVecType_Nest(Mat A, VecType vtype)
    1334: {
    1335:   PetscBool flg;

    1337:   PetscFunctionBegin;
    1338:   PetscCall(PetscStrcmp(vtype, VECNEST, &flg));
    1339:   /* In reality, this only distinguishes VECNEST and "other" */
    1340:   if (flg) A->ops->getvecs = MatCreateVecs_Nest;
    1341:   else A->ops->getvecs = NULL;
    1342:   PetscFunctionReturn(PETSC_SUCCESS);
    1343: }

    1345: /*@
    1346:   MatNestSetVecType - Sets the type of `Vec` returned by `MatCreateVecs()`

    1348:   Not Collective

    1350:   Input Parameters:
    1351: + A     - `MATNEST` matrix
    1352: - vtype - `VecType` to use for creating vectors

    1354:   Level: developer

    1356: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreateVecs()`, `MatCreateNest()`, `VecType`
    1357: @*/
    1358: PetscErrorCode MatNestSetVecType(Mat A, VecType vtype)
    1359: {
    1360:   PetscFunctionBegin;
    1362:   PetscTryMethod(A, "MatNestSetVecType_C", (Mat, VecType), (A, vtype));
    1363:   PetscFunctionReturn(PETSC_SUCCESS);
    1364: }

    1366: static PetscErrorCode MatNestSetSubMats_Nest(Mat A, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[], const Mat a[])
    1367: {
    1368:   Mat_Nest *s = (Mat_Nest *)A->data;
    1369:   PetscInt  i, j, m, n, M, N;
    1370:   PetscBool cong, isstd, sametype = PETSC_FALSE;
    1371:   VecType   vtype, type;

    1373:   PetscFunctionBegin;
    1374:   PetscCall(MatReset_Nest(A));

    1376:   s->nr = nr;
    1377:   s->nc = nc;

    1379:   /* Create space for submatrices */
    1380:   PetscCall(PetscMalloc1(nr, &s->m));
    1381:   PetscCall(PetscMalloc1(nr * nc, &s->m[0]));
    1382:   for (i = 0; i < nr; i++) {
    1383:     s->m[i] = s->m[0] + i * nc;
    1384:     for (j = 0; j < nc; j++) {
    1385:       s->m[i][j] = a ? a[i * nc + j] : NULL;
    1386:       PetscCall(PetscObjectReference((PetscObject)s->m[i][j]));
    1387:     }
    1388:   }
    1389:   PetscCall(MatGetVecType(A, &vtype));
    1390:   PetscCall(PetscStrcmp(vtype, VECSTANDARD, &isstd));
    1391:   if (isstd) {
    1392:     /* check if all blocks have the same vectype */
    1393:     vtype = NULL;
    1394:     for (i = 0; i < nr; i++) {
    1395:       for (j = 0; j < nc; j++) {
    1396:         if (s->m[i][j]) {
    1397:           if (!vtype) { /* first visited block */
    1398:             PetscCall(MatGetVecType(s->m[i][j], &vtype));
    1399:             sametype = PETSC_TRUE;
    1400:           } else if (sametype) {
    1401:             PetscCall(MatGetVecType(s->m[i][j], &type));
    1402:             PetscCall(PetscStrcmp(vtype, type, &sametype));
    1403:           }
    1404:         }
    1405:       }
    1406:     }
    1407:     if (sametype) { /* propagate vectype */
    1408:       PetscCall(MatSetVecType(A, vtype));
    1409:     }
    1410:   }

    1412:   PetscCall(MatSetUp_NestIS_Private(A, nr, is_row, nc, is_col));

    1414:   PetscCall(PetscMalloc1(nr, &s->row_len));
    1415:   PetscCall(PetscMalloc1(nc, &s->col_len));
    1416:   for (i = 0; i < nr; i++) s->row_len[i] = -1;
    1417:   for (j = 0; j < nc; j++) s->col_len[j] = -1;

    1419:   PetscCall(PetscCalloc1(nr * nc, &s->nnzstate));
    1420:   for (i = 0; i < nr; i++) {
    1421:     for (j = 0; j < nc; j++) {
    1422:       if (s->m[i][j]) PetscCall(MatGetNonzeroState(s->m[i][j], &s->nnzstate[i * nc + j]));
    1423:     }
    1424:   }

    1426:   PetscCall(MatNestGetSizes_Private(A, &m, &n, &M, &N));

    1428:   PetscCall(PetscLayoutSetSize(A->rmap, M));
    1429:   PetscCall(PetscLayoutSetLocalSize(A->rmap, m));
    1430:   PetscCall(PetscLayoutSetSize(A->cmap, N));
    1431:   PetscCall(PetscLayoutSetLocalSize(A->cmap, n));

    1433:   PetscCall(PetscLayoutSetUp(A->rmap));
    1434:   PetscCall(PetscLayoutSetUp(A->cmap));

    1436:   /* disable operations that are not supported for non-square matrices,
    1437:      or matrices for which is_row != is_col  */
    1438:   PetscCall(MatHasCongruentLayouts(A, &cong));
    1439:   if (cong && nr != nc) cong = PETSC_FALSE;
    1440:   if (cong) {
    1441:     for (i = 0; cong && i < nr; i++) PetscCall(ISEqualUnsorted(s->isglobal.row[i], s->isglobal.col[i], &cong));
    1442:   }
    1443:   if (!cong) {
    1444:     A->ops->missingdiagonal = NULL;
    1445:     A->ops->getdiagonal     = NULL;
    1446:     A->ops->shift           = NULL;
    1447:     A->ops->diagonalset     = NULL;
    1448:   }

    1450:   PetscCall(PetscCalloc2(nr, &s->left, nc, &s->right));
    1451:   PetscCall(PetscObjectStateIncrease((PetscObject)A));
    1452:   A->nonzerostate++;
    1453:   PetscFunctionReturn(PETSC_SUCCESS);
    1454: }

    1456: /*@
    1457:   MatNestSetSubMats - Sets the nested submatrices in a `MATNEST`

    1459:   Collective

    1461:   Input Parameters:
    1462: + A      - `MATNEST` matrix
    1463: . nr     - number of nested row blocks
    1464: . is_row - index sets for each nested row block, or `NULL` to make contiguous
    1465: . nc     - number of nested column blocks
    1466: . is_col - index sets for each nested column block, or `NULL` to make contiguous
    1467: - a      - array of nr*nc submatrices, or `NULL`

    1469:   Level: advanced

    1471:   Notes:
    1472:   This always resets any block matrix information previously set.
    1473:   Pass `NULL` in the corresponding entry of `a` for an empty block.

    1475:   In both C and Fortran, `a` must be a row-major order array containing the matrices. See
    1476:   `MatCreateNest()` for an example.

    1478: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreateNest()`, `MatNestSetSubMat()`, `MatNestGetSubMat()`, `MatNestGetSubMats()`
    1479: @*/
    1480: PetscErrorCode MatNestSetSubMats(Mat A, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[], const Mat a[])
    1481: {
    1482:   PetscFunctionBegin;
    1485:   PetscCheck(nr >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Number of rows cannot be negative");
    1486:   if (nr && is_row) {
    1487:     PetscAssertPointer(is_row, 3);
    1489:   }
    1491:   PetscCheck(nc >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Number of columns cannot be negative");
    1492:   if (nc && is_col) {
    1493:     PetscAssertPointer(is_col, 5);
    1495:   }
    1496:   PetscTryMethod(A, "MatNestSetSubMats_C", (Mat, PetscInt, const IS[], PetscInt, const IS[], const Mat[]), (A, nr, is_row, nc, is_col, a));
    1497:   PetscFunctionReturn(PETSC_SUCCESS);
    1498: }

    1500: static PetscErrorCode MatNestCreateAggregateL2G_Private(Mat A, PetscInt n, const IS islocal[], const IS isglobal[], PetscBool colflg, ISLocalToGlobalMapping *ltog)
    1501: {
    1502:   PetscBool flg;
    1503:   PetscInt  i, j, m, mi, *ix;

    1505:   PetscFunctionBegin;
    1506:   *ltog = NULL;
    1507:   for (i = 0, m = 0, flg = PETSC_FALSE; i < n; i++) {
    1508:     if (islocal[i]) {
    1509:       PetscCall(ISGetLocalSize(islocal[i], &mi));
    1510:       flg = PETSC_TRUE; /* We found a non-trivial entry */
    1511:     } else {
    1512:       PetscCall(ISGetLocalSize(isglobal[i], &mi));
    1513:     }
    1514:     m += mi;
    1515:   }
    1516:   if (!flg) PetscFunctionReturn(PETSC_SUCCESS);

    1518:   PetscCall(PetscMalloc1(m, &ix));
    1519:   for (i = 0, m = 0; i < n; i++) {
    1520:     ISLocalToGlobalMapping smap = NULL;
    1521:     Mat                    sub  = NULL;
    1522:     PetscSF                sf;
    1523:     PetscLayout            map;
    1524:     const PetscInt        *ix2;

    1526:     if (!colflg) {
    1527:       PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
    1528:     } else {
    1529:       PetscCall(MatNestFindNonzeroSubMatCol(A, i, &sub));
    1530:     }
    1531:     if (sub) {
    1532:       if (!colflg) {
    1533:         PetscCall(MatGetLocalToGlobalMapping(sub, &smap, NULL));
    1534:       } else {
    1535:         PetscCall(MatGetLocalToGlobalMapping(sub, NULL, &smap));
    1536:       }
    1537:     }
    1538:     /*
    1539:        Now we need to extract the monolithic global indices that correspond to the given split global indices.
    1540:        In many/most cases, we only want MatGetLocalSubMatrix() to work, in which case we only need to know the size of the local spaces.
    1541:     */
    1542:     PetscCall(ISGetIndices(isglobal[i], &ix2));
    1543:     if (islocal[i]) {
    1544:       PetscInt *ilocal, *iremote;
    1545:       PetscInt  mil, nleaves;

    1547:       PetscCall(ISGetLocalSize(islocal[i], &mi));
    1548:       PetscCheck(smap, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing local to global map");
    1549:       for (j = 0; j < mi; j++) ix[m + j] = j;
    1550:       PetscCall(ISLocalToGlobalMappingApply(smap, mi, ix + m, ix + m));

    1552:       /* PetscSFSetGraphLayout does not like negative indices */
    1553:       PetscCall(PetscMalloc2(mi, &ilocal, mi, &iremote));
    1554:       for (j = 0, nleaves = 0; j < mi; j++) {
    1555:         if (ix[m + j] < 0) continue;
    1556:         ilocal[nleaves]  = j;
    1557:         iremote[nleaves] = ix[m + j];
    1558:         nleaves++;
    1559:       }
    1560:       PetscCall(ISGetLocalSize(isglobal[i], &mil));
    1561:       PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)A), &sf));
    1562:       PetscCall(PetscLayoutCreate(PetscObjectComm((PetscObject)A), &map));
    1563:       PetscCall(PetscLayoutSetLocalSize(map, mil));
    1564:       PetscCall(PetscLayoutSetUp(map));
    1565:       PetscCall(PetscSFSetGraphLayout(sf, map, nleaves, ilocal, PETSC_USE_POINTER, iremote));
    1566:       PetscCall(PetscLayoutDestroy(&map));
    1567:       PetscCall(PetscSFBcastBegin(sf, MPIU_INT, ix2, ix + m, MPI_REPLACE));
    1568:       PetscCall(PetscSFBcastEnd(sf, MPIU_INT, ix2, ix + m, MPI_REPLACE));
    1569:       PetscCall(PetscSFDestroy(&sf));
    1570:       PetscCall(PetscFree2(ilocal, iremote));
    1571:     } else {
    1572:       PetscCall(ISGetLocalSize(isglobal[i], &mi));
    1573:       for (j = 0; j < mi; j++) ix[m + j] = ix2[i];
    1574:     }
    1575:     PetscCall(ISRestoreIndices(isglobal[i], &ix2));
    1576:     m += mi;
    1577:   }
    1578:   PetscCall(ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)A), 1, m, ix, PETSC_OWN_POINTER, ltog));
    1579:   PetscFunctionReturn(PETSC_SUCCESS);
    1580: }

    1582: /* If an IS was provided, there is nothing Nest needs to do, otherwise Nest will build a strided IS */
    1583: /*
    1584:   nprocessors = NP
    1585:   Nest x^T = ((g_0,g_1,...g_nprocs-1), (h_0,h_1,...h_NP-1))
    1586:        proc 0: => (g_0,h_0,)
    1587:        proc 1: => (g_1,h_1,)
    1588:        ...
    1589:        proc nprocs-1: => (g_NP-1,h_NP-1,)

    1591:             proc 0:                      proc 1:                    proc nprocs-1:
    1592:     is[0] = (0,1,2,...,nlocal(g_0)-1)  (0,1,...,nlocal(g_1)-1)  (0,1,...,nlocal(g_NP-1))

    1594:             proc 0:
    1595:     is[1] = (nlocal(g_0),nlocal(g_0)+1,...,nlocal(g_0)+nlocal(h_0)-1)
    1596:             proc 1:
    1597:     is[1] = (nlocal(g_1),nlocal(g_1)+1,...,nlocal(g_1)+nlocal(h_1)-1)

    1599:             proc NP-1:
    1600:     is[1] = (nlocal(g_NP-1),nlocal(g_NP-1)+1,...,nlocal(g_NP-1)+nlocal(h_NP-1)-1)
    1601: */
    1602: static PetscErrorCode MatSetUp_NestIS_Private(Mat A, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[])
    1603: {
    1604:   Mat_Nest *vs = (Mat_Nest *)A->data;
    1605:   PetscInt  i, j, offset, n, nsum, bs;
    1606:   Mat       sub = NULL;

    1608:   PetscFunctionBegin;
    1609:   PetscCall(PetscMalloc1(nr, &vs->isglobal.row));
    1610:   PetscCall(PetscMalloc1(nc, &vs->isglobal.col));
    1611:   if (is_row) { /* valid IS is passed in */
    1612:     /* refs on is[] are incremented */
    1613:     for (i = 0; i < vs->nr; i++) {
    1614:       PetscCall(PetscObjectReference((PetscObject)is_row[i]));

    1616:       vs->isglobal.row[i] = is_row[i];
    1617:     }
    1618:   } else { /* Create the ISs by inspecting sizes of a submatrix in each row */
    1619:     nsum = 0;
    1620:     for (i = 0; i < vs->nr; i++) { /* Add up the local sizes to compute the aggregate offset */
    1621:       PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
    1622:       PetscCheck(sub, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "No nonzero submatrix in row %" PetscInt_FMT, i);
    1623:       PetscCall(MatGetLocalSize(sub, &n, NULL));
    1624:       PetscCheck(n >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Sizes have not yet been set for submatrix");
    1625:       nsum += n;
    1626:     }
    1627:     PetscCallMPI(MPI_Scan(&nsum, &offset, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)A)));
    1628:     offset -= nsum;
    1629:     for (i = 0; i < vs->nr; i++) {
    1630:       PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
    1631:       PetscCall(MatGetLocalSize(sub, &n, NULL));
    1632:       PetscCall(MatGetBlockSizes(sub, &bs, NULL));
    1633:       PetscCall(ISCreateStride(PetscObjectComm((PetscObject)sub), n, offset, 1, &vs->isglobal.row[i]));
    1634:       PetscCall(ISSetBlockSize(vs->isglobal.row[i], bs));
    1635:       offset += n;
    1636:     }
    1637:   }

    1639:   if (is_col) { /* valid IS is passed in */
    1640:     /* refs on is[] are incremented */
    1641:     for (j = 0; j < vs->nc; j++) {
    1642:       PetscCall(PetscObjectReference((PetscObject)is_col[j]));

    1644:       vs->isglobal.col[j] = is_col[j];
    1645:     }
    1646:   } else { /* Create the ISs by inspecting sizes of a submatrix in each column */
    1647:     offset = A->cmap->rstart;
    1648:     nsum   = 0;
    1649:     for (j = 0; j < vs->nc; j++) {
    1650:       PetscCall(MatNestFindNonzeroSubMatCol(A, j, &sub));
    1651:       PetscCheck(sub, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "No nonzero submatrix in column %" PetscInt_FMT, i);
    1652:       PetscCall(MatGetLocalSize(sub, NULL, &n));
    1653:       PetscCheck(n >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Sizes have not yet been set for submatrix");
    1654:       nsum += n;
    1655:     }
    1656:     PetscCallMPI(MPI_Scan(&nsum, &offset, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)A)));
    1657:     offset -= nsum;
    1658:     for (j = 0; j < vs->nc; j++) {
    1659:       PetscCall(MatNestFindNonzeroSubMatCol(A, j, &sub));
    1660:       PetscCall(MatGetLocalSize(sub, NULL, &n));
    1661:       PetscCall(MatGetBlockSizes(sub, NULL, &bs));
    1662:       PetscCall(ISCreateStride(PetscObjectComm((PetscObject)sub), n, offset, 1, &vs->isglobal.col[j]));
    1663:       PetscCall(ISSetBlockSize(vs->isglobal.col[j], bs));
    1664:       offset += n;
    1665:     }
    1666:   }

    1668:   /* Set up the local ISs */
    1669:   PetscCall(PetscMalloc1(vs->nr, &vs->islocal.row));
    1670:   PetscCall(PetscMalloc1(vs->nc, &vs->islocal.col));
    1671:   for (i = 0, offset = 0; i < vs->nr; i++) {
    1672:     IS                     isloc;
    1673:     ISLocalToGlobalMapping rmap = NULL;
    1674:     PetscInt               nlocal, bs;
    1675:     PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
    1676:     if (sub) PetscCall(MatGetLocalToGlobalMapping(sub, &rmap, NULL));
    1677:     if (rmap) {
    1678:       PetscCall(MatGetBlockSizes(sub, &bs, NULL));
    1679:       PetscCall(ISLocalToGlobalMappingGetSize(rmap, &nlocal));
    1680:       PetscCall(ISCreateStride(PETSC_COMM_SELF, nlocal, offset, 1, &isloc));
    1681:       PetscCall(ISSetBlockSize(isloc, bs));
    1682:     } else {
    1683:       nlocal = 0;
    1684:       isloc  = NULL;
    1685:     }
    1686:     vs->islocal.row[i] = isloc;
    1687:     offset += nlocal;
    1688:   }
    1689:   for (i = 0, offset = 0; i < vs->nc; i++) {
    1690:     IS                     isloc;
    1691:     ISLocalToGlobalMapping cmap = NULL;
    1692:     PetscInt               nlocal, bs;
    1693:     PetscCall(MatNestFindNonzeroSubMatCol(A, i, &sub));
    1694:     if (sub) PetscCall(MatGetLocalToGlobalMapping(sub, NULL, &cmap));
    1695:     if (cmap) {
    1696:       PetscCall(MatGetBlockSizes(sub, NULL, &bs));
    1697:       PetscCall(ISLocalToGlobalMappingGetSize(cmap, &nlocal));
    1698:       PetscCall(ISCreateStride(PETSC_COMM_SELF, nlocal, offset, 1, &isloc));
    1699:       PetscCall(ISSetBlockSize(isloc, bs));
    1700:     } else {
    1701:       nlocal = 0;
    1702:       isloc  = NULL;
    1703:     }
    1704:     vs->islocal.col[i] = isloc;
    1705:     offset += nlocal;
    1706:   }

    1708:   /* Set up the aggregate ISLocalToGlobalMapping */
    1709:   {
    1710:     ISLocalToGlobalMapping rmap, cmap;
    1711:     PetscCall(MatNestCreateAggregateL2G_Private(A, vs->nr, vs->islocal.row, vs->isglobal.row, PETSC_FALSE, &rmap));
    1712:     PetscCall(MatNestCreateAggregateL2G_Private(A, vs->nc, vs->islocal.col, vs->isglobal.col, PETSC_TRUE, &cmap));
    1713:     if (rmap && cmap) PetscCall(MatSetLocalToGlobalMapping(A, rmap, cmap));
    1714:     PetscCall(ISLocalToGlobalMappingDestroy(&rmap));
    1715:     PetscCall(ISLocalToGlobalMappingDestroy(&cmap));
    1716:   }

    1718:   if (PetscDefined(USE_DEBUG)) {
    1719:     for (i = 0; i < vs->nr; i++) {
    1720:       for (j = 0; j < vs->nc; j++) {
    1721:         PetscInt m, n, M, N, mi, ni, Mi, Ni;
    1722:         Mat      B = vs->m[i][j];
    1723:         if (!B) continue;
    1724:         PetscCall(MatGetSize(B, &M, &N));
    1725:         PetscCall(MatGetLocalSize(B, &m, &n));
    1726:         PetscCall(ISGetSize(vs->isglobal.row[i], &Mi));
    1727:         PetscCall(ISGetSize(vs->isglobal.col[j], &Ni));
    1728:         PetscCall(ISGetLocalSize(vs->isglobal.row[i], &mi));
    1729:         PetscCall(ISGetLocalSize(vs->isglobal.col[j], &ni));
    1730:         PetscCheck(M == Mi && N == Ni, PetscObjectComm((PetscObject)sub), PETSC_ERR_ARG_INCOMP, "Global sizes (%" PetscInt_FMT ",%" PetscInt_FMT ") of nested submatrix (%" PetscInt_FMT ",%" PetscInt_FMT ") do not agree with space defined by index sets (%" PetscInt_FMT ",%" PetscInt_FMT ")", M, N, i, j, Mi, Ni);
    1731:         PetscCheck(m == mi && n == ni, PetscObjectComm((PetscObject)sub), PETSC_ERR_ARG_INCOMP, "Local sizes (%" PetscInt_FMT ",%" PetscInt_FMT ") of nested submatrix (%" PetscInt_FMT ",%" PetscInt_FMT ") do not agree with space defined by index sets (%" PetscInt_FMT ",%" PetscInt_FMT ")", m, n, i, j, mi, ni);
    1732:       }
    1733:     }
    1734:   }

    1736:   /* Set A->assembled if all non-null blocks are currently assembled */
    1737:   for (i = 0; i < vs->nr; i++) {
    1738:     for (j = 0; j < vs->nc; j++) {
    1739:       if (vs->m[i][j] && !vs->m[i][j]->assembled) PetscFunctionReturn(PETSC_SUCCESS);
    1740:     }
    1741:   }
    1742:   A->assembled = PETSC_TRUE;
    1743:   PetscFunctionReturn(PETSC_SUCCESS);
    1744: }

    1746: /*@C
    1747:   MatCreateNest - Creates a new `MATNEST` matrix containing several nested submatrices, each stored separately

    1749:   Collective

    1751:   Input Parameters:
    1752: + comm   - Communicator for the new `MATNEST`
    1753: . nr     - number of nested row blocks
    1754: . is_row - index sets for each nested row block, or `NULL` to make contiguous
    1755: . nc     - number of nested column blocks
    1756: . is_col - index sets for each nested column block, or `NULL` to make contiguous
    1757: - a      - array of nr*nc submatrices, empty submatrices can be passed using `NULL`

    1759:   Output Parameter:
    1760: . B - new matrix

    1762:   Note:
    1763:   In both C and Fortran, `a` must be a row-major order array holding references to the matrices.
    1764:   For instance, to represent the matrix
    1765:   $\begin{bmatrix} A_{11} & A_{12} \\ A_{21} & A_{22}\end{bmatrix}$
    1766:   one should use `Mat a[4]={A11,A12,A21,A22}`.

    1768:   Level: advanced

    1770: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreate()`, `VecCreateNest()`, `DMCreateMatrix()`, `MatNestSetSubMat()`,
    1771:           `MatNestGetSubMat()`, `MatNestGetLocalISs()`, `MatNestGetSize()`,
    1772:           `MatNestGetISs()`, `MatNestSetSubMats()`, `MatNestGetSubMats()`
    1773: @*/
    1774: PetscErrorCode MatCreateNest(MPI_Comm comm, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[], const Mat a[], Mat *B)
    1775: {
    1776:   PetscFunctionBegin;
    1777:   PetscCall(MatCreate(comm, B));
    1778:   PetscCall(MatSetType(*B, MATNEST));
    1779:   (*B)->preallocated = PETSC_TRUE;
    1780:   PetscCall(MatNestSetSubMats(*B, nr, is_row, nc, is_col, a));
    1781:   PetscFunctionReturn(PETSC_SUCCESS);
    1782: }

    1784: static PetscErrorCode MatConvert_Nest_SeqAIJ_fast(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
    1785: {
    1786:   Mat_Nest     *nest = (Mat_Nest *)A->data;
    1787:   Mat          *trans;
    1788:   PetscScalar **avv;
    1789:   PetscScalar  *vv;
    1790:   PetscInt    **aii, **ajj;
    1791:   PetscInt     *ii, *jj, *ci;
    1792:   PetscInt      nr, nc, nnz, i, j;
    1793:   PetscBool     done;

    1795:   PetscFunctionBegin;
    1796:   PetscCall(MatGetSize(A, &nr, &nc));
    1797:   if (reuse == MAT_REUSE_MATRIX) {
    1798:     PetscInt rnr;

    1800:     PetscCall(MatGetRowIJ(*newmat, 0, PETSC_FALSE, PETSC_FALSE, &rnr, (const PetscInt **)&ii, (const PetscInt **)&jj, &done));
    1801:     PetscCheck(done, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "MatGetRowIJ");
    1802:     PetscCheck(rnr == nr, PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "Cannot reuse matrix, wrong number of rows");
    1803:     PetscCall(MatSeqAIJGetArray(*newmat, &vv));
    1804:   }
    1805:   /* extract CSR for nested SeqAIJ matrices */
    1806:   nnz = 0;
    1807:   PetscCall(PetscCalloc4(nest->nr * nest->nc, &aii, nest->nr * nest->nc, &ajj, nest->nr * nest->nc, &avv, nest->nr * nest->nc, &trans));
    1808:   for (i = 0; i < nest->nr; ++i) {
    1809:     for (j = 0; j < nest->nc; ++j) {
    1810:       Mat B = nest->m[i][j];
    1811:       if (B) {
    1812:         PetscScalar *naa;
    1813:         PetscInt    *nii, *njj, nnr;
    1814:         PetscBool    istrans;

    1816:         PetscCall(PetscObjectTypeCompare((PetscObject)B, MATTRANSPOSEVIRTUAL, &istrans));
    1817:         if (istrans) {
    1818:           Mat Bt;

    1820:           PetscCall(MatTransposeGetMat(B, &Bt));
    1821:           PetscCall(MatTranspose(Bt, MAT_INITIAL_MATRIX, &trans[i * nest->nc + j]));
    1822:           B = trans[i * nest->nc + j];
    1823:         } else {
    1824:           PetscCall(PetscObjectTypeCompare((PetscObject)B, MATHERMITIANTRANSPOSEVIRTUAL, &istrans));
    1825:           if (istrans) {
    1826:             Mat Bt;

    1828:             PetscCall(MatHermitianTransposeGetMat(B, &Bt));
    1829:             PetscCall(MatHermitianTranspose(Bt, MAT_INITIAL_MATRIX, &trans[i * nest->nc + j]));
    1830:             B = trans[i * nest->nc + j];
    1831:           }
    1832:         }
    1833:         PetscCall(MatGetRowIJ(B, 0, PETSC_FALSE, PETSC_FALSE, &nnr, (const PetscInt **)&nii, (const PetscInt **)&njj, &done));
    1834:         PetscCheck(done, PetscObjectComm((PetscObject)B), PETSC_ERR_PLIB, "MatGetRowIJ");
    1835:         PetscCall(MatSeqAIJGetArray(B, &naa));
    1836:         nnz += nii[nnr];

    1838:         aii[i * nest->nc + j] = nii;
    1839:         ajj[i * nest->nc + j] = njj;
    1840:         avv[i * nest->nc + j] = naa;
    1841:       }
    1842:     }
    1843:   }
    1844:   if (reuse != MAT_REUSE_MATRIX) {
    1845:     PetscCall(PetscMalloc1(nr + 1, &ii));
    1846:     PetscCall(PetscMalloc1(nnz, &jj));
    1847:     PetscCall(PetscMalloc1(nnz, &vv));
    1848:   } else {
    1849:     PetscCheck(nnz == ii[nr], PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "Cannot reuse matrix, wrong number of nonzeros");
    1850:   }

    1852:   /* new row pointer */
    1853:   PetscCall(PetscArrayzero(ii, nr + 1));
    1854:   for (i = 0; i < nest->nr; ++i) {
    1855:     PetscInt ncr, rst;

    1857:     PetscCall(ISStrideGetInfo(nest->isglobal.row[i], &rst, NULL));
    1858:     PetscCall(ISGetLocalSize(nest->isglobal.row[i], &ncr));
    1859:     for (j = 0; j < nest->nc; ++j) {
    1860:       if (aii[i * nest->nc + j]) {
    1861:         PetscInt *nii = aii[i * nest->nc + j];
    1862:         PetscInt  ir;

    1864:         for (ir = rst; ir < ncr + rst; ++ir) {
    1865:           ii[ir + 1] += nii[1] - nii[0];
    1866:           nii++;
    1867:         }
    1868:       }
    1869:     }
    1870:   }
    1871:   for (i = 0; i < nr; i++) ii[i + 1] += ii[i];

    1873:   /* construct CSR for the new matrix */
    1874:   PetscCall(PetscCalloc1(nr, &ci));
    1875:   for (i = 0; i < nest->nr; ++i) {
    1876:     PetscInt ncr, rst;

    1878:     PetscCall(ISStrideGetInfo(nest->isglobal.row[i], &rst, NULL));
    1879:     PetscCall(ISGetLocalSize(nest->isglobal.row[i], &ncr));
    1880:     for (j = 0; j < nest->nc; ++j) {
    1881:       if (aii[i * nest->nc + j]) {
    1882:         PetscScalar *nvv = avv[i * nest->nc + j], vscale = 1.0, vshift = 0.0;
    1883:         PetscInt    *nii = aii[i * nest->nc + j];
    1884:         PetscInt    *njj = ajj[i * nest->nc + j];
    1885:         PetscInt     ir, cst;

    1887:         if (trans[i * nest->nc + j]) {
    1888:           vscale = ((Mat_Shell *)nest->m[i][j]->data)->vscale;
    1889:           vshift = ((Mat_Shell *)nest->m[i][j]->data)->vshift;
    1890:         }
    1891:         PetscCall(ISStrideGetInfo(nest->isglobal.col[j], &cst, NULL));
    1892:         for (ir = rst; ir < ncr + rst; ++ir) {
    1893:           PetscInt ij, rsize = nii[1] - nii[0], ist = ii[ir] + ci[ir];

    1895:           for (ij = 0; ij < rsize; ij++) {
    1896:             jj[ist + ij] = *njj + cst;
    1897:             vv[ist + ij] = vscale * *nvv;
    1898:             if (PetscUnlikely(vshift != 0.0 && *njj == ir - rst)) vv[ist + ij] += vshift;
    1899:             njj++;
    1900:             nvv++;
    1901:           }
    1902:           ci[ir] += rsize;
    1903:           nii++;
    1904:         }
    1905:       }
    1906:     }
    1907:   }
    1908:   PetscCall(PetscFree(ci));

    1910:   /* restore info */
    1911:   for (i = 0; i < nest->nr; ++i) {
    1912:     for (j = 0; j < nest->nc; ++j) {
    1913:       Mat B = nest->m[i][j];
    1914:       if (B) {
    1915:         PetscInt nnr = 0, k = i * nest->nc + j;

    1917:         B = (trans[k] ? trans[k] : B);
    1918:         PetscCall(MatRestoreRowIJ(B, 0, PETSC_FALSE, PETSC_FALSE, &nnr, (const PetscInt **)&aii[k], (const PetscInt **)&ajj[k], &done));
    1919:         PetscCheck(done, PetscObjectComm((PetscObject)B), PETSC_ERR_PLIB, "MatRestoreRowIJ");
    1920:         PetscCall(MatSeqAIJRestoreArray(B, &avv[k]));
    1921:         PetscCall(MatDestroy(&trans[k]));
    1922:       }
    1923:     }
    1924:   }
    1925:   PetscCall(PetscFree4(aii, ajj, avv, trans));

    1927:   /* finalize newmat */
    1928:   if (reuse == MAT_INITIAL_MATRIX) {
    1929:     PetscCall(MatCreateSeqAIJWithArrays(PetscObjectComm((PetscObject)A), nr, nc, ii, jj, vv, newmat));
    1930:   } else if (reuse == MAT_INPLACE_MATRIX) {
    1931:     Mat B;

    1933:     PetscCall(MatCreateSeqAIJWithArrays(PetscObjectComm((PetscObject)A), nr, nc, ii, jj, vv, &B));
    1934:     PetscCall(MatHeaderReplace(A, &B));
    1935:   }
    1936:   PetscCall(MatAssemblyBegin(*newmat, MAT_FINAL_ASSEMBLY));
    1937:   PetscCall(MatAssemblyEnd(*newmat, MAT_FINAL_ASSEMBLY));
    1938:   {
    1939:     Mat_SeqAIJ *a = (Mat_SeqAIJ *)((*newmat)->data);
    1940:     a->free_a     = PETSC_TRUE;
    1941:     a->free_ij    = PETSC_TRUE;
    1942:   }
    1943:   PetscFunctionReturn(PETSC_SUCCESS);
    1944: }

    1946: PETSC_INTERN PetscErrorCode MatAXPY_Dense_Nest(Mat Y, PetscScalar a, Mat X)
    1947: {
    1948:   Mat_Nest *nest = (Mat_Nest *)X->data;
    1949:   PetscInt  i, j, k, rstart;
    1950:   PetscBool flg;

    1952:   PetscFunctionBegin;
    1953:   /* Fill by row */
    1954:   for (j = 0; j < nest->nc; ++j) {
    1955:     /* Using global column indices and ISAllGather() is not scalable. */
    1956:     IS              bNis;
    1957:     PetscInt        bN;
    1958:     const PetscInt *bNindices;
    1959:     PetscCall(ISAllGather(nest->isglobal.col[j], &bNis));
    1960:     PetscCall(ISGetSize(bNis, &bN));
    1961:     PetscCall(ISGetIndices(bNis, &bNindices));
    1962:     for (i = 0; i < nest->nr; ++i) {
    1963:       Mat             B = nest->m[i][j], D = NULL;
    1964:       PetscInt        bm, br;
    1965:       const PetscInt *bmindices;
    1966:       if (!B) continue;
    1967:       PetscCall(PetscObjectTypeCompareAny((PetscObject)B, &flg, MATTRANSPOSEVIRTUAL, MATHERMITIANTRANSPOSEVIRTUAL, ""));
    1968:       if (flg) {
    1969:         PetscTryMethod(B, "MatTransposeGetMat_C", (Mat, Mat *), (B, &D));
    1970:         PetscTryMethod(B, "MatHermitianTransposeGetMat_C", (Mat, Mat *), (B, &D));
    1971:         PetscCall(MatConvert(B, ((PetscObject)D)->type_name, MAT_INITIAL_MATRIX, &D));
    1972:         B = D;
    1973:       }
    1974:       PetscCall(PetscObjectTypeCompareAny((PetscObject)B, &flg, MATSEQSBAIJ, MATMPISBAIJ, ""));
    1975:       if (flg) {
    1976:         if (D) PetscCall(MatConvert(D, MATBAIJ, MAT_INPLACE_MATRIX, &D));
    1977:         else PetscCall(MatConvert(B, MATBAIJ, MAT_INITIAL_MATRIX, &D));
    1978:         B = D;
    1979:       }
    1980:       PetscCall(ISGetLocalSize(nest->isglobal.row[i], &bm));
    1981:       PetscCall(ISGetIndices(nest->isglobal.row[i], &bmindices));
    1982:       PetscCall(MatGetOwnershipRange(B, &rstart, NULL));
    1983:       for (br = 0; br < bm; ++br) {
    1984:         PetscInt           row = bmindices[br], brncols, *cols;
    1985:         const PetscInt    *brcols;
    1986:         const PetscScalar *brcoldata;
    1987:         PetscScalar       *vals = NULL;
    1988:         PetscCall(MatGetRow(B, br + rstart, &brncols, &brcols, &brcoldata));
    1989:         PetscCall(PetscMalloc1(brncols, &cols));
    1990:         for (k = 0; k < brncols; k++) cols[k] = bNindices[brcols[k]];
    1991:         /*
    1992:           Nest blocks are required to be nonoverlapping -- otherwise nest and monolithic index layouts wouldn't match.
    1993:           Thus, we could use INSERT_VALUES, but I prefer ADD_VALUES.
    1994:          */
    1995:         if (a != 1.0) {
    1996:           PetscCall(PetscMalloc1(brncols, &vals));
    1997:           for (k = 0; k < brncols; k++) vals[k] = a * brcoldata[k];
    1998:           PetscCall(MatSetValues(Y, 1, &row, brncols, cols, vals, ADD_VALUES));
    1999:           PetscCall(PetscFree(vals));
    2000:         } else {
    2001:           PetscCall(MatSetValues(Y, 1, &row, brncols, cols, brcoldata, ADD_VALUES));
    2002:         }
    2003:         PetscCall(MatRestoreRow(B, br + rstart, &brncols, &brcols, &brcoldata));
    2004:         PetscCall(PetscFree(cols));
    2005:       }
    2006:       if (D) PetscCall(MatDestroy(&D));
    2007:       PetscCall(ISRestoreIndices(nest->isglobal.row[i], &bmindices));
    2008:     }
    2009:     PetscCall(ISRestoreIndices(bNis, &bNindices));
    2010:     PetscCall(ISDestroy(&bNis));
    2011:   }
    2012:   PetscCall(MatAssemblyBegin(Y, MAT_FINAL_ASSEMBLY));
    2013:   PetscCall(MatAssemblyEnd(Y, MAT_FINAL_ASSEMBLY));
    2014:   PetscFunctionReturn(PETSC_SUCCESS);
    2015: }

    2017: static PetscErrorCode MatConvert_Nest_AIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
    2018: {
    2019:   Mat_Nest   *nest = (Mat_Nest *)A->data;
    2020:   PetscInt    m, n, M, N, i, j, k, *dnnz, *onnz = NULL, rstart, cstart, cend;
    2021:   PetscMPIInt size;
    2022:   Mat         C;

    2024:   PetscFunctionBegin;
    2025:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
    2026:   if (size == 1) { /* look for a special case with SeqAIJ matrices and strided-1, contiguous, blocks */
    2027:     PetscInt  nf;
    2028:     PetscBool fast;

    2030:     PetscCall(PetscStrcmp(newtype, MATAIJ, &fast));
    2031:     if (!fast) PetscCall(PetscStrcmp(newtype, MATSEQAIJ, &fast));
    2032:     for (i = 0; i < nest->nr && fast; ++i) {
    2033:       for (j = 0; j < nest->nc && fast; ++j) {
    2034:         Mat B = nest->m[i][j];
    2035:         if (B) {
    2036:           PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQAIJ, &fast));
    2037:           if (!fast) {
    2038:             PetscBool istrans;

    2040:             PetscCall(PetscObjectTypeCompare((PetscObject)B, MATTRANSPOSEVIRTUAL, &istrans));
    2041:             if (istrans) {
    2042:               Mat Bt;

    2044:               PetscCall(MatTransposeGetMat(B, &Bt));
    2045:               PetscCall(PetscObjectTypeCompare((PetscObject)Bt, MATSEQAIJ, &fast));
    2046:             } else {
    2047:               PetscCall(PetscObjectTypeCompare((PetscObject)B, MATHERMITIANTRANSPOSEVIRTUAL, &istrans));
    2048:               if (istrans) {
    2049:                 Mat Bt;

    2051:                 PetscCall(MatHermitianTransposeGetMat(B, &Bt));
    2052:                 PetscCall(PetscObjectTypeCompare((PetscObject)Bt, MATSEQAIJ, &fast));
    2053:               }
    2054:             }
    2055:             if (fast) fast = (PetscBool)(!((Mat_Shell *)B->data)->zrows && !((Mat_Shell *)B->data)->zcols && !((Mat_Shell *)B->data)->axpy && !((Mat_Shell *)B->data)->left && !((Mat_Shell *)B->data)->right && !((Mat_Shell *)B->data)->dshift);
    2056:           }
    2057:         }
    2058:       }
    2059:     }
    2060:     for (i = 0, nf = 0; i < nest->nr && fast; ++i) {
    2061:       PetscCall(PetscObjectTypeCompare((PetscObject)nest->isglobal.row[i], ISSTRIDE, &fast));
    2062:       if (fast) {
    2063:         PetscInt f, s;

    2065:         PetscCall(ISStrideGetInfo(nest->isglobal.row[i], &f, &s));
    2066:         if (f != nf || s != 1) {
    2067:           fast = PETSC_FALSE;
    2068:         } else {
    2069:           PetscCall(ISGetSize(nest->isglobal.row[i], &f));
    2070:           nf += f;
    2071:         }
    2072:       }
    2073:     }
    2074:     for (i = 0, nf = 0; i < nest->nc && fast; ++i) {
    2075:       PetscCall(PetscObjectTypeCompare((PetscObject)nest->isglobal.col[i], ISSTRIDE, &fast));
    2076:       if (fast) {
    2077:         PetscInt f, s;

    2079:         PetscCall(ISStrideGetInfo(nest->isglobal.col[i], &f, &s));
    2080:         if (f != nf || s != 1) {
    2081:           fast = PETSC_FALSE;
    2082:         } else {
    2083:           PetscCall(ISGetSize(nest->isglobal.col[i], &f));
    2084:           nf += f;
    2085:         }
    2086:       }
    2087:     }
    2088:     if (fast) {
    2089:       PetscCall(MatConvert_Nest_SeqAIJ_fast(A, newtype, reuse, newmat));
    2090:       PetscFunctionReturn(PETSC_SUCCESS);
    2091:     }
    2092:   }
    2093:   PetscCall(MatGetSize(A, &M, &N));
    2094:   PetscCall(MatGetLocalSize(A, &m, &n));
    2095:   PetscCall(MatGetOwnershipRangeColumn(A, &cstart, &cend));
    2096:   if (reuse == MAT_REUSE_MATRIX) C = *newmat;
    2097:   else {
    2098:     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &C));
    2099:     PetscCall(MatSetType(C, newtype));
    2100:     PetscCall(MatSetSizes(C, m, n, M, N));
    2101:   }
    2102:   PetscCall(PetscMalloc1(2 * m, &dnnz));
    2103:   if (m) {
    2104:     onnz = dnnz + m;
    2105:     for (k = 0; k < m; k++) {
    2106:       dnnz[k] = 0;
    2107:       onnz[k] = 0;
    2108:     }
    2109:   }
    2110:   for (j = 0; j < nest->nc; ++j) {
    2111:     IS              bNis;
    2112:     PetscInt        bN;
    2113:     const PetscInt *bNindices;
    2114:     PetscBool       flg;
    2115:     /* Using global column indices and ISAllGather() is not scalable. */
    2116:     PetscCall(ISAllGather(nest->isglobal.col[j], &bNis));
    2117:     PetscCall(ISGetSize(bNis, &bN));
    2118:     PetscCall(ISGetIndices(bNis, &bNindices));
    2119:     for (i = 0; i < nest->nr; ++i) {
    2120:       PetscSF         bmsf;
    2121:       PetscSFNode    *iremote;
    2122:       Mat             B = nest->m[i][j], D = NULL;
    2123:       PetscInt        bm, *sub_dnnz, *sub_onnz, br;
    2124:       const PetscInt *bmindices;
    2125:       if (!B) continue;
    2126:       PetscCall(ISGetLocalSize(nest->isglobal.row[i], &bm));
    2127:       PetscCall(ISGetIndices(nest->isglobal.row[i], &bmindices));
    2128:       PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)A), &bmsf));
    2129:       PetscCall(PetscMalloc1(bm, &iremote));
    2130:       PetscCall(PetscMalloc1(bm, &sub_dnnz));
    2131:       PetscCall(PetscMalloc1(bm, &sub_onnz));
    2132:       for (k = 0; k < bm; ++k) {
    2133:         sub_dnnz[k] = 0;
    2134:         sub_onnz[k] = 0;
    2135:       }
    2136:       PetscCall(PetscObjectTypeCompareAny((PetscObject)B, &flg, MATTRANSPOSEVIRTUAL, MATHERMITIANTRANSPOSEVIRTUAL, ""));
    2137:       if (flg) {
    2138:         PetscTryMethod(B, "MatTransposeGetMat_C", (Mat, Mat *), (B, &D));
    2139:         PetscTryMethod(B, "MatHermitianTransposeGetMat_C", (Mat, Mat *), (B, &D));
    2140:         PetscCall(MatConvert(B, ((PetscObject)D)->type_name, MAT_INITIAL_MATRIX, &D));
    2141:         B = D;
    2142:       }
    2143:       PetscCall(PetscObjectTypeCompareAny((PetscObject)B, &flg, MATSEQSBAIJ, MATMPISBAIJ, ""));
    2144:       if (flg) {
    2145:         if (D) PetscCall(MatConvert(D, MATBAIJ, MAT_INPLACE_MATRIX, &D));
    2146:         else PetscCall(MatConvert(B, MATBAIJ, MAT_INITIAL_MATRIX, &D));
    2147:         B = D;
    2148:       }
    2149:       /*
    2150:        Locate the owners for all of the locally-owned global row indices for this row block.
    2151:        These determine the roots of PetscSF used to communicate preallocation data to row owners.
    2152:        The roots correspond to the dnnz and onnz entries; thus, there are two roots per row.
    2153:        */
    2154:       PetscCall(MatGetOwnershipRange(B, &rstart, NULL));
    2155:       for (br = 0; br < bm; ++br) {
    2156:         PetscInt        row = bmindices[br], brncols, col;
    2157:         const PetscInt *brcols;
    2158:         PetscInt        rowrel   = 0; /* row's relative index on its owner rank */
    2159:         PetscMPIInt     rowowner = 0;
    2160:         PetscCall(PetscLayoutFindOwnerIndex(A->rmap, row, &rowowner, &rowrel));
    2161:         /* how many roots  */
    2162:         iremote[br].rank  = rowowner;
    2163:         iremote[br].index = rowrel; /* edge from bmdnnz to dnnz */
    2164:         /* get nonzero pattern */
    2165:         PetscCall(MatGetRow(B, br + rstart, &brncols, &brcols, NULL));
    2166:         for (k = 0; k < brncols; k++) {
    2167:           col = bNindices[brcols[k]];
    2168:           if (col >= A->cmap->range[rowowner] && col < A->cmap->range[rowowner + 1]) {
    2169:             sub_dnnz[br]++;
    2170:           } else {
    2171:             sub_onnz[br]++;
    2172:           }
    2173:         }
    2174:         PetscCall(MatRestoreRow(B, br + rstart, &brncols, &brcols, NULL));
    2175:       }
    2176:       if (D) PetscCall(MatDestroy(&D));
    2177:       PetscCall(ISRestoreIndices(nest->isglobal.row[i], &bmindices));
    2178:       /* bsf will have to take care of disposing of bedges. */
    2179:       PetscCall(PetscSFSetGraph(bmsf, m, bm, NULL, PETSC_OWN_POINTER, iremote, PETSC_OWN_POINTER));
    2180:       PetscCall(PetscSFReduceBegin(bmsf, MPIU_INT, sub_dnnz, dnnz, MPI_SUM));
    2181:       PetscCall(PetscSFReduceEnd(bmsf, MPIU_INT, sub_dnnz, dnnz, MPI_SUM));
    2182:       PetscCall(PetscSFReduceBegin(bmsf, MPIU_INT, sub_onnz, onnz, MPI_SUM));
    2183:       PetscCall(PetscSFReduceEnd(bmsf, MPIU_INT, sub_onnz, onnz, MPI_SUM));
    2184:       PetscCall(PetscFree(sub_dnnz));
    2185:       PetscCall(PetscFree(sub_onnz));
    2186:       PetscCall(PetscSFDestroy(&bmsf));
    2187:     }
    2188:     PetscCall(ISRestoreIndices(bNis, &bNindices));
    2189:     PetscCall(ISDestroy(&bNis));
    2190:   }
    2191:   /* Resize preallocation if overestimated */
    2192:   for (i = 0; i < m; i++) {
    2193:     dnnz[i] = PetscMin(dnnz[i], A->cmap->n);
    2194:     onnz[i] = PetscMin(onnz[i], A->cmap->N - A->cmap->n);
    2195:   }
    2196:   PetscCall(MatSeqAIJSetPreallocation(C, 0, dnnz));
    2197:   PetscCall(MatMPIAIJSetPreallocation(C, 0, dnnz, 0, onnz));
    2198:   PetscCall(PetscFree(dnnz));
    2199:   PetscCall(MatAXPY_Dense_Nest(C, 1.0, A));
    2200:   if (reuse == MAT_INPLACE_MATRIX) {
    2201:     PetscCall(MatHeaderReplace(A, &C));
    2202:   } else *newmat = C;
    2203:   PetscFunctionReturn(PETSC_SUCCESS);
    2204: }

    2206: static PetscErrorCode MatConvert_Nest_Dense(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
    2207: {
    2208:   Mat      B;
    2209:   PetscInt m, n, M, N;

    2211:   PetscFunctionBegin;
    2212:   PetscCall(MatGetSize(A, &M, &N));
    2213:   PetscCall(MatGetLocalSize(A, &m, &n));
    2214:   if (reuse == MAT_REUSE_MATRIX) {
    2215:     B = *newmat;
    2216:     PetscCall(MatZeroEntries(B));
    2217:   } else {
    2218:     PetscCall(MatCreateDense(PetscObjectComm((PetscObject)A), m, PETSC_DECIDE, M, N, NULL, &B));
    2219:   }
    2220:   PetscCall(MatAXPY_Dense_Nest(B, 1.0, A));
    2221:   if (reuse == MAT_INPLACE_MATRIX) {
    2222:     PetscCall(MatHeaderReplace(A, &B));
    2223:   } else if (reuse == MAT_INITIAL_MATRIX) *newmat = B;
    2224:   PetscFunctionReturn(PETSC_SUCCESS);
    2225: }

    2227: static PetscErrorCode MatHasOperation_Nest(Mat mat, MatOperation op, PetscBool *has)
    2228: {
    2229:   Mat_Nest    *bA = (Mat_Nest *)mat->data;
    2230:   MatOperation opAdd;
    2231:   PetscInt     i, j, nr = bA->nr, nc = bA->nc;
    2232:   PetscBool    flg;

    2234:   PetscFunctionBegin;
    2235:   *has = PETSC_FALSE;
    2236:   if (op == MATOP_MULT || op == MATOP_MULT_ADD || op == MATOP_MULT_TRANSPOSE || op == MATOP_MULT_TRANSPOSE_ADD) {
    2237:     opAdd = (op == MATOP_MULT || op == MATOP_MULT_ADD ? MATOP_MULT_ADD : MATOP_MULT_TRANSPOSE_ADD);
    2238:     for (j = 0; j < nc; j++) {
    2239:       for (i = 0; i < nr; i++) {
    2240:         if (!bA->m[i][j]) continue;
    2241:         PetscCall(MatHasOperation(bA->m[i][j], opAdd, &flg));
    2242:         if (!flg) PetscFunctionReturn(PETSC_SUCCESS);
    2243:       }
    2244:     }
    2245:   }
    2246:   if (((void **)mat->ops)[op]) *has = PETSC_TRUE;
    2247:   PetscFunctionReturn(PETSC_SUCCESS);
    2248: }

    2250: /*MC
    2251:   MATNEST -  "nest" - Matrix type consisting of nested submatrices, each stored separately.

    2253:   Level: intermediate

    2255:   Notes:
    2256:   This matrix type permits scalable use of `PCFIELDSPLIT` and avoids the large memory costs of extracting submatrices.
    2257:   It allows the use of symmetric and block formats for parts of multi-physics simulations.
    2258:   It is usually used with `DMCOMPOSITE` and `DMCreateMatrix()`

    2260:   Each of the submatrices lives on the same MPI communicator as the original nest matrix (though they can have zero
    2261:   rows/columns on some processes.) Thus this is not meant for cases where the submatrices live on far fewer processes
    2262:   than the nest matrix.

    2264: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreate()`, `MatType`, `MatCreateNest()`, `MatNestSetSubMat()`, `MatNestGetSubMat()`,
    2265:           `VecCreateNest()`, `DMCreateMatrix()`, `DMCOMPOSITE`, `MatNestSetVecType()`, `MatNestGetLocalISs()`,
    2266:           `MatNestGetISs()`, `MatNestSetSubMats()`, `MatNestGetSubMats()`
    2267: M*/
    2268: PETSC_EXTERN PetscErrorCode MatCreate_Nest(Mat A)
    2269: {
    2270:   Mat_Nest *s;

    2272:   PetscFunctionBegin;
    2273:   PetscCall(PetscNew(&s));
    2274:   A->data = (void *)s;

    2276:   s->nr            = -1;
    2277:   s->nc            = -1;
    2278:   s->m             = NULL;
    2279:   s->splitassembly = PETSC_FALSE;

    2281:   PetscCall(PetscMemzero(A->ops, sizeof(*A->ops)));

    2283:   A->ops->mult                      = MatMult_Nest;
    2284:   A->ops->multadd                   = MatMultAdd_Nest;
    2285:   A->ops->multtranspose             = MatMultTranspose_Nest;
    2286:   A->ops->multtransposeadd          = MatMultTransposeAdd_Nest;
    2287:   A->ops->transpose                 = MatTranspose_Nest;
    2288:   A->ops->multhermitiantranspose    = MatMultHermitianTranspose_Nest;
    2289:   A->ops->multhermitiantransposeadd = MatMultHermitianTransposeAdd_Nest;
    2290:   A->ops->assemblybegin             = MatAssemblyBegin_Nest;
    2291:   A->ops->assemblyend               = MatAssemblyEnd_Nest;
    2292:   A->ops->zeroentries               = MatZeroEntries_Nest;
    2293:   A->ops->copy                      = MatCopy_Nest;
    2294:   A->ops->axpy                      = MatAXPY_Nest;
    2295:   A->ops->duplicate                 = MatDuplicate_Nest;
    2296:   A->ops->createsubmatrix           = MatCreateSubMatrix_Nest;
    2297:   A->ops->destroy                   = MatDestroy_Nest;
    2298:   A->ops->view                      = MatView_Nest;
    2299:   A->ops->getvecs                   = NULL; /* Use VECNEST by calling MatNestSetVecType(A,VECNEST) */
    2300:   A->ops->getlocalsubmatrix         = MatGetLocalSubMatrix_Nest;
    2301:   A->ops->restorelocalsubmatrix     = MatRestoreLocalSubMatrix_Nest;
    2302:   A->ops->getdiagonal               = MatGetDiagonal_Nest;
    2303:   A->ops->diagonalscale             = MatDiagonalScale_Nest;
    2304:   A->ops->scale                     = MatScale_Nest;
    2305:   A->ops->shift                     = MatShift_Nest;
    2306:   A->ops->diagonalset               = MatDiagonalSet_Nest;
    2307:   A->ops->setrandom                 = MatSetRandom_Nest;
    2308:   A->ops->hasoperation              = MatHasOperation_Nest;
    2309:   A->ops->missingdiagonal           = MatMissingDiagonal_Nest;

    2311:   A->spptr     = NULL;
    2312:   A->assembled = PETSC_FALSE;

    2314:   /* expose Nest api's */
    2315:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSubMat_C", MatNestGetSubMat_Nest));
    2316:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetSubMat_C", MatNestSetSubMat_Nest));
    2317:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSubMats_C", MatNestGetSubMats_Nest));
    2318:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSize_C", MatNestGetSize_Nest));
    2319:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetISs_C", MatNestGetISs_Nest));
    2320:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetLocalISs_C", MatNestGetLocalISs_Nest));
    2321:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetVecType_C", MatNestSetVecType_Nest));
    2322:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetSubMats_C", MatNestSetSubMats_Nest));
    2323:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_mpiaij_C", MatConvert_Nest_AIJ));
    2324:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_seqaij_C", MatConvert_Nest_AIJ));
    2325:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_aij_C", MatConvert_Nest_AIJ));
    2326:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_is_C", MatConvert_Nest_IS));
    2327:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_mpidense_C", MatConvert_Nest_Dense));
    2328:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_seqdense_C", MatConvert_Nest_Dense));
    2329:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_nest_seqdense_C", MatProductSetFromOptions_Nest_Dense));
    2330:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_nest_mpidense_C", MatProductSetFromOptions_Nest_Dense));

    2332:   PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATNEST));
    2333:   PetscFunctionReturn(PETSC_SUCCESS);
    2334: }