Date: 2026-07-30T21:14:15+00:00

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(PetscCtxRt 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++) PetscCall(ISDestroy(&lst[i]));
    395:   PetscCall(PetscFree(lst));
    396:   *list = NULL;
    397:   PetscFunctionReturn(PETSC_SUCCESS);
    398: }

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

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

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

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

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

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

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

    458: static PetscErrorCode MatAssemblyBegin_Nest(Mat A, MatAssemblyType type)
    459: {
    460:   Mat_Nest *vs = (Mat_Nest *)A->data;
    461:   PetscInt  i, j;
    462:   PetscBool nnzstate = PETSC_FALSE;

    464:   PetscFunctionBegin;
    465:   for (i = 0; i < vs->nr; i++) {
    466:     for (j = 0; j < vs->nc; j++) {
    467:       PetscObjectState subnnzstate = 0;
    468:       if (vs->m[i][j]) {
    469:         PetscCall(MatAssemblyBegin(vs->m[i][j], type));
    470:         if (!vs->splitassembly) {
    471:           /* Note: split assembly will fail if the same block appears more than once (even indirectly through a nested
    472:            * sub-block). This could be fixed by adding a flag to Mat so that there was a way to check if a Mat was
    473:            * already performing an assembly, but the result would by more complicated and appears to offer less
    474:            * potential for diagnostics and correctness checking. Split assembly should be fixed once there is an
    475:            * interface for libraries to make asynchronous progress in "user-defined non-blocking collectives".
    476:            */
    477:           PetscCall(MatAssemblyEnd(vs->m[i][j], type));
    478:           PetscCall(MatGetNonzeroState(vs->m[i][j], &subnnzstate));
    479:         }
    480:       }
    481:       nnzstate                     = (PetscBool)(nnzstate || vs->nnzstate[i * vs->nc + j] != subnnzstate);
    482:       vs->nnzstate[i * vs->nc + j] = subnnzstate;
    483:     }
    484:   }
    485:   if (nnzstate) A->nonzerostate++;
    486:   PetscFunctionReturn(PETSC_SUCCESS);
    487: }

    489: static PetscErrorCode MatAssemblyEnd_Nest(Mat A, MatAssemblyType type)
    490: {
    491:   Mat_Nest *vs = (Mat_Nest *)A->data;
    492:   PetscInt  i, j;

    494:   PetscFunctionBegin;
    495:   for (i = 0; i < vs->nr; i++) {
    496:     for (j = 0; j < vs->nc; j++) {
    497:       if (vs->m[i][j]) {
    498:         if (vs->splitassembly) PetscCall(MatAssemblyEnd(vs->m[i][j], type));
    499:       }
    500:     }
    501:   }
    502:   PetscFunctionReturn(PETSC_SUCCESS);
    503: }

    505: static PetscErrorCode MatNestFindNonzeroSubMatRow(Mat A, PetscInt row, Mat *B)
    506: {
    507:   Mat_Nest *vs = (Mat_Nest *)A->data;
    508:   PetscInt  j;
    509:   Mat       sub;

    511:   PetscFunctionBegin;
    512:   sub = (row < vs->nc) ? vs->m[row][row] : (Mat)NULL; /* Prefer to find on the diagonal */
    513:   for (j = 0; !sub && j < vs->nc; j++) sub = vs->m[row][j];
    514:   if (sub) PetscCall(MatSetUp(sub)); /* Ensure that the sizes are available */
    515:   *B = sub;
    516:   PetscFunctionReturn(PETSC_SUCCESS);
    517: }

    519: static PetscErrorCode MatNestFindNonzeroSubMatCol(Mat A, PetscInt col, Mat *B)
    520: {
    521:   Mat_Nest *vs = (Mat_Nest *)A->data;
    522:   PetscInt  i;
    523:   Mat       sub;

    525:   PetscFunctionBegin;
    526:   sub = (col < vs->nr) ? vs->m[col][col] : (Mat)NULL; /* Prefer to find on the diagonal */
    527:   for (i = 0; !sub && i < vs->nr; i++) sub = vs->m[i][col];
    528:   if (sub) PetscCall(MatSetUp(sub)); /* Ensure that the sizes are available */
    529:   *B = sub;
    530:   PetscFunctionReturn(PETSC_SUCCESS);
    531: }

    533: static PetscErrorCode MatNestFindISRange(Mat A, PetscInt n, const IS list[], IS is, PetscInt *begin, PetscInt *end)
    534: {
    535:   PetscInt  i, j, size, m;
    536:   PetscBool flg;
    537:   IS        out, concatenate[2];

    539:   PetscFunctionBegin;
    540:   PetscAssertPointer(list, 3);
    542:   if (begin) {
    543:     PetscAssertPointer(begin, 5);
    544:     *begin = -1;
    545:   }
    546:   if (end) {
    547:     PetscAssertPointer(end, 6);
    548:     *end = -1;
    549:   }
    550:   for (i = 0; i < n; i++) {
    551:     if (!list[i]) continue;
    552:     PetscCall(ISEqualUnsorted(list[i], is, &flg));
    553:     if (flg) {
    554:       if (begin) *begin = i;
    555:       if (end) *end = i + 1;
    556:       PetscFunctionReturn(PETSC_SUCCESS);
    557:     }
    558:   }
    559:   PetscCall(ISGetSize(is, &size));
    560:   for (i = 0; i < n - 1; i++) {
    561:     if (!list[i]) continue;
    562:     m = 0;
    563:     PetscCall(ISConcatenate(PetscObjectComm((PetscObject)A), 2, list + i, &out));
    564:     PetscCall(ISGetSize(out, &m));
    565:     for (j = i + 2; j < n && m < size; j++) {
    566:       if (list[j]) {
    567:         concatenate[0] = out;
    568:         concatenate[1] = list[j];
    569:         PetscCall(ISConcatenate(PetscObjectComm((PetscObject)A), 2, concatenate, &out));
    570:         PetscCall(ISDestroy(concatenate));
    571:         PetscCall(ISGetSize(out, &m));
    572:       }
    573:     }
    574:     if (m == size) {
    575:       PetscCall(ISEqualUnsorted(out, is, &flg));
    576:       if (flg) {
    577:         if (begin) *begin = i;
    578:         if (end) *end = j;
    579:         PetscCall(ISDestroy(&out));
    580:         PetscFunctionReturn(PETSC_SUCCESS);
    581:       }
    582:     }
    583:     PetscCall(ISDestroy(&out));
    584:   }
    585:   PetscFunctionReturn(PETSC_SUCCESS);
    586: }

    588: static PetscErrorCode MatNestFillEmptyMat_Private(Mat A, PetscInt i, PetscInt j, Mat *B)
    589: {
    590:   Mat_Nest *vs = (Mat_Nest *)A->data;
    591:   PetscInt  lr, lc;

    593:   PetscFunctionBegin;
    594:   PetscCall(MatCreate(PetscObjectComm((PetscObject)A), B));
    595:   PetscCall(ISGetLocalSize(vs->isglobal.row[i], &lr));
    596:   PetscCall(ISGetLocalSize(vs->isglobal.col[j], &lc));
    597:   PetscCall(MatSetSizes(*B, lr, lc, PETSC_DECIDE, PETSC_DECIDE));
    598:   PetscCall(MatSetType(*B, MATAIJ));
    599:   PetscCall(MatSeqAIJSetPreallocation(*B, 0, NULL));
    600:   PetscCall(MatMPIAIJSetPreallocation(*B, 0, NULL, 0, NULL));
    601:   PetscCall(MatSetUp(*B));
    602:   PetscCall(MatSetOption(*B, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
    603:   PetscCall(MatAssemblyBegin(*B, MAT_FINAL_ASSEMBLY));
    604:   PetscCall(MatAssemblyEnd(*B, MAT_FINAL_ASSEMBLY));
    605:   PetscFunctionReturn(PETSC_SUCCESS);
    606: }

    608: static PetscErrorCode MatNestGetBlock_Private(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *B)
    609: {
    610:   Mat_Nest  *vs = (Mat_Nest *)A->data;
    611:   Mat       *a;
    612:   PetscInt   i, j, k, l, nr = rend - rbegin, nc = cend - cbegin;
    613:   char       keyname[256];
    614:   PetscBool *b;
    615:   PetscBool  flg;

    617:   PetscFunctionBegin;
    618:   *B = NULL;
    619:   PetscCall(PetscSNPrintf(keyname, sizeof(keyname), "NestBlock_%" PetscInt_FMT "-%" PetscInt_FMT "x%" PetscInt_FMT "-%" PetscInt_FMT, rbegin, rend, cbegin, cend));
    620:   PetscCall(PetscObjectQuery((PetscObject)A, keyname, (PetscObject *)B));
    621:   if (*B) PetscFunctionReturn(PETSC_SUCCESS);

    623:   PetscCall(PetscMalloc2(nr * nc, &a, nr * nc, &b));
    624:   for (i = 0; i < nr; i++) {
    625:     for (j = 0; j < nc; j++) {
    626:       a[i * nc + j] = vs->m[rbegin + i][cbegin + j];
    627:       b[i * nc + j] = PETSC_FALSE;
    628:     }
    629:   }
    630:   if (nc != vs->nc && nr != vs->nr) {
    631:     for (i = 0; i < nr; i++) {
    632:       for (j = 0; j < nc; j++) {
    633:         flg = PETSC_FALSE;
    634:         for (k = 0; (k < nr && !flg); k++) {
    635:           if (a[j + k * nc]) flg = PETSC_TRUE;
    636:         }
    637:         if (flg) {
    638:           flg = PETSC_FALSE;
    639:           for (l = 0; (l < nc && !flg); l++) {
    640:             if (a[i * nc + l]) flg = PETSC_TRUE;
    641:           }
    642:         }
    643:         if (!flg) {
    644:           b[i * nc + j] = PETSC_TRUE;
    645:           PetscCall(MatNestFillEmptyMat_Private(A, rbegin + i, cbegin + j, a + i * nc + j));
    646:         }
    647:       }
    648:     }
    649:   }
    650:   PetscCall(MatCreateNest(PetscObjectComm((PetscObject)A), nr, nr != vs->nr ? NULL : vs->isglobal.row, nc, nc != vs->nc ? NULL : vs->isglobal.col, a, B));
    651:   for (i = 0; i < nr; i++) {
    652:     for (j = 0; j < nc; j++) {
    653:       if (b[i * nc + j]) PetscCall(MatDestroy(a + i * nc + j));
    654:     }
    655:   }
    656:   PetscCall(PetscFree2(a, b));
    657:   (*B)->assembled = A->assembled;
    658:   PetscCall(PetscObjectCompose((PetscObject)A, keyname, (PetscObject)*B));
    659:   PetscCall(PetscObjectDereference((PetscObject)*B)); /* Leave the only remaining reference in the composition */
    660:   PetscFunctionReturn(PETSC_SUCCESS);
    661: }

    663: static PetscErrorCode MatNestFindSubMat(Mat A, struct MatNestISPair *is, IS isrow, IS iscol, Mat *B)
    664: {
    665:   Mat_Nest *vs = (Mat_Nest *)A->data;
    666:   PetscInt  rbegin, rend, cbegin, cend;

    668:   PetscFunctionBegin;
    669:   PetscCall(MatNestFindISRange(A, vs->nr, is->row, isrow, &rbegin, &rend));
    670:   PetscCall(MatNestFindISRange(A, vs->nc, is->col, iscol, &cbegin, &cend));
    671:   if (rend == rbegin + 1 && cend == cbegin + 1) {
    672:     if (!vs->m[rbegin][cbegin]) PetscCall(MatNestFillEmptyMat_Private(A, rbegin, cbegin, vs->m[rbegin] + cbegin));
    673:     *B = vs->m[rbegin][cbegin];
    674:   } else if (rbegin != -1 && cbegin != -1) {
    675:     PetscCall(MatNestGetBlock_Private(A, rbegin, rend, cbegin, cend, B));
    676:   } else SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_INCOMP, "Could not find index set");
    677:   PetscFunctionReturn(PETSC_SUCCESS);
    678: }

    680: /*
    681:    TODO: This does not actually returns a submatrix we can modify
    682: */
    683: static PetscErrorCode MatCreateSubMatrix_Nest(Mat A, IS isrow, IS iscol, MatReuse reuse, Mat *B)
    684: {
    685:   Mat_Nest *vs = (Mat_Nest *)A->data;
    686:   Mat       sub;

    688:   PetscFunctionBegin;
    689:   PetscCall(MatNestFindSubMat(A, &vs->isglobal, isrow, iscol, &sub));
    690:   switch (reuse) {
    691:   case MAT_INITIAL_MATRIX:
    692:     PetscCall(PetscObjectReference((PetscObject)sub));
    693:     if (sub) PetscCall(PetscObjectStateIncrease((PetscObject)sub));
    694:     *B = sub;
    695:     break;
    696:   case MAT_REUSE_MATRIX:
    697:     PetscCheck(sub == *B, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Submatrix was not used before in this call");
    698:     if (sub) PetscCall(PetscObjectStateIncrease((PetscObject)sub));
    699:     break;
    700:   default:
    701:     break;
    702:   }
    703:   PetscFunctionReturn(PETSC_SUCCESS);
    704: }

    706: static PetscErrorCode MatGetLocalSubMatrix_Nest(Mat A, IS isrow, IS iscol, Mat *B)
    707: {
    708:   Mat_Nest *vs = (Mat_Nest *)A->data;
    709:   Mat       sub;

    711:   PetscFunctionBegin;
    712:   PetscCall(MatNestFindSubMat(A, &vs->islocal, isrow, iscol, &sub));
    713:   /* We allow the submatrix to be NULL, perhaps it would be better for the user to return an empty matrix instead */
    714:   PetscCall(PetscObjectReference((PetscObject)sub));
    715:   *B = sub;
    716:   PetscFunctionReturn(PETSC_SUCCESS);
    717: }

    719: static PetscErrorCode MatRestoreLocalSubMatrix_Nest(Mat A, IS isrow, IS iscol, Mat *B)
    720: {
    721:   Mat_Nest *vs = (Mat_Nest *)A->data;
    722:   Mat       sub;

    724:   PetscFunctionBegin;
    725:   PetscCall(MatNestFindSubMat(A, &vs->islocal, isrow, iscol, &sub));
    726:   PetscCheck(*B == sub, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Local submatrix has not been gotten");
    727:   if (sub) {
    728:     PetscCheck(((PetscObject)sub)->refct > 1, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Local submatrix has had reference count decremented too many times");
    729:     PetscCall(MatDestroy(B));
    730:   }
    731:   PetscFunctionReturn(PETSC_SUCCESS);
    732: }

    734: static PetscErrorCode MatGetDiagonal_Nest(Mat A, Vec v)
    735: {
    736:   Mat_Nest *bA = (Mat_Nest *)A->data;
    737:   PetscInt  i;

    739:   PetscFunctionBegin;
    740:   for (i = 0; i < bA->nr; i++) {
    741:     Vec bv;
    742:     PetscCall(VecGetSubVector(v, bA->isglobal.row[i], &bv));
    743:     if (bA->m[i][i]) PetscCall(MatGetDiagonal(bA->m[i][i], bv));
    744:     else PetscCall(VecSet(bv, 0.0));
    745:     PetscCall(VecRestoreSubVector(v, bA->isglobal.row[i], &bv));
    746:   }
    747:   PetscFunctionReturn(PETSC_SUCCESS);
    748: }

    750: static PetscErrorCode MatDiagonalScale_Nest(Mat A, Vec l, Vec r)
    751: {
    752:   Mat_Nest *bA = (Mat_Nest *)A->data;
    753:   Vec       bl, *br;
    754:   PetscInt  i, j;

    756:   PetscFunctionBegin;
    757:   PetscCall(PetscCalloc1(bA->nc, &br));
    758:   if (r) {
    759:     for (j = 0; j < bA->nc; j++) PetscCall(VecGetSubVector(r, bA->isglobal.col[j], &br[j]));
    760:   }
    761:   bl = NULL;
    762:   for (i = 0; i < bA->nr; i++) {
    763:     if (l) PetscCall(VecGetSubVector(l, bA->isglobal.row[i], &bl));
    764:     for (j = 0; j < bA->nc; j++) {
    765:       if (bA->m[i][j]) PetscCall(MatDiagonalScale(bA->m[i][j], bl, br[j]));
    766:     }
    767:     if (l) PetscCall(VecRestoreSubVector(l, bA->isglobal.row[i], &bl));
    768:   }
    769:   if (r) {
    770:     for (j = 0; j < bA->nc; j++) PetscCall(VecRestoreSubVector(r, bA->isglobal.col[j], &br[j]));
    771:   }
    772:   PetscCall(PetscFree(br));
    773:   PetscFunctionReturn(PETSC_SUCCESS);
    774: }

    776: static PetscErrorCode MatScale_Nest(Mat A, PetscScalar a)
    777: {
    778:   Mat_Nest *bA = (Mat_Nest *)A->data;
    779:   PetscInt  i, j;

    781:   PetscFunctionBegin;
    782:   for (i = 0; i < bA->nr; i++) {
    783:     for (j = 0; j < bA->nc; j++) {
    784:       if (bA->m[i][j]) PetscCall(MatScale(bA->m[i][j], a));
    785:     }
    786:   }
    787:   PetscFunctionReturn(PETSC_SUCCESS);
    788: }

    790: static PetscErrorCode MatShift_Nest(Mat A, PetscScalar a)
    791: {
    792:   Mat_Nest *bA = (Mat_Nest *)A->data;
    793:   PetscInt  i;
    794:   PetscBool nnzstate = PETSC_FALSE;

    796:   PetscFunctionBegin;
    797:   for (i = 0; i < bA->nr; i++) {
    798:     PetscObjectState subnnzstate = 0;
    799:     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);
    800:     PetscCall(MatShift(bA->m[i][i], a));
    801:     PetscCall(MatGetNonzeroState(bA->m[i][i], &subnnzstate));
    802:     nnzstate                     = (PetscBool)(nnzstate || bA->nnzstate[i * bA->nc + i] != subnnzstate);
    803:     bA->nnzstate[i * bA->nc + i] = subnnzstate;
    804:   }
    805:   if (nnzstate) A->nonzerostate++;
    806:   PetscFunctionReturn(PETSC_SUCCESS);
    807: }

    809: static PetscErrorCode MatDiagonalSet_Nest(Mat A, Vec D, InsertMode is)
    810: {
    811:   Mat_Nest *bA = (Mat_Nest *)A->data;
    812:   PetscInt  i;
    813:   PetscBool nnzstate = PETSC_FALSE;

    815:   PetscFunctionBegin;
    816:   for (i = 0; i < bA->nr; i++) {
    817:     PetscObjectState subnnzstate = 0;
    818:     Vec              bv;
    819:     PetscCall(VecGetSubVector(D, bA->isglobal.row[i], &bv));
    820:     if (bA->m[i][i]) {
    821:       PetscCall(MatDiagonalSet(bA->m[i][i], bv, is));
    822:       PetscCall(MatGetNonzeroState(bA->m[i][i], &subnnzstate));
    823:     }
    824:     PetscCall(VecRestoreSubVector(D, bA->isglobal.row[i], &bv));
    825:     nnzstate                     = (PetscBool)(nnzstate || bA->nnzstate[i * bA->nc + i] != subnnzstate);
    826:     bA->nnzstate[i * bA->nc + i] = subnnzstate;
    827:   }
    828:   if (nnzstate) A->nonzerostate++;
    829:   PetscFunctionReturn(PETSC_SUCCESS);
    830: }

    832: static PetscErrorCode MatSetRandom_Nest(Mat A, PetscRandom rctx)
    833: {
    834:   Mat_Nest *bA = (Mat_Nest *)A->data;
    835:   PetscInt  i, j;

    837:   PetscFunctionBegin;
    838:   for (i = 0; i < bA->nr; i++) {
    839:     for (j = 0; j < bA->nc; j++) {
    840:       if (bA->m[i][j]) PetscCall(MatSetRandom(bA->m[i][j], rctx));
    841:     }
    842:   }
    843:   PetscFunctionReturn(PETSC_SUCCESS);
    844: }

    846: static PetscErrorCode MatCreateVecs_Nest(Mat A, Vec *right, Vec *left)
    847: {
    848:   Mat_Nest *bA = (Mat_Nest *)A->data;
    849:   Vec      *L, *R;
    850:   MPI_Comm  comm;
    851:   PetscInt  i, j;

    853:   PetscFunctionBegin;
    854:   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
    855:   if (right) {
    856:     /* allocate R */
    857:     PetscCall(PetscMalloc1(bA->nc, &R));
    858:     /* Create the right vectors */
    859:     for (j = 0; j < bA->nc; j++) {
    860:       for (i = 0; i < bA->nr; i++) {
    861:         if (bA->m[i][j]) {
    862:           PetscCall(MatCreateVecs(bA->m[i][j], &R[j], NULL));
    863:           break;
    864:         }
    865:       }
    866:       PetscCheck(i != bA->nr, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Mat(Nest) contains a null column.");
    867:     }
    868:     PetscCall(VecCreateNest(comm, bA->nc, bA->isglobal.col, R, right));
    869:     /* hand back control to the nest vector */
    870:     for (j = 0; j < bA->nc; j++) PetscCall(VecDestroy(&R[j]));
    871:     PetscCall(PetscFree(R));
    872:   }

    874:   if (left) {
    875:     /* allocate L */
    876:     PetscCall(PetscMalloc1(bA->nr, &L));
    877:     /* Create the left vectors */
    878:     for (i = 0; i < bA->nr; i++) {
    879:       for (j = 0; j < bA->nc; j++) {
    880:         if (bA->m[i][j]) {
    881:           PetscCall(MatCreateVecs(bA->m[i][j], NULL, &L[i]));
    882:           break;
    883:         }
    884:       }
    885:       PetscCheck(j != bA->nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Mat(Nest) contains a null row.");
    886:     }

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

    891:     PetscCall(PetscFree(L));
    892:   }
    893:   PetscFunctionReturn(PETSC_SUCCESS);
    894: }

    896: static PetscErrorCode MatView_Nest(Mat A, PetscViewer viewer)
    897: {
    898:   Mat_Nest *bA = (Mat_Nest *)A->data;
    899:   PetscBool isascii, viewSub = PETSC_FALSE;
    900:   PetscInt  i, j;

    902:   PetscFunctionBegin;
    903:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
    904:   if (isascii) {
    905:     PetscViewerFormat format;

    907:     PetscCall(PetscViewerGetFormat(viewer, &format));
    908:     if (format == PETSC_VIEWER_ASCII_MATLAB) {
    909:       Mat T;

    911:       PetscCall(MatConvert(A, MATAIJ, MAT_INITIAL_MATRIX, &T));
    912:       PetscCall(MatView(T, viewer));
    913:       PetscCall(MatDestroy(&T));
    914:       PetscFunctionReturn(PETSC_SUCCESS);
    915:     }
    916:     PetscCall(PetscOptionsGetBool(((PetscObject)A)->options, ((PetscObject)A)->prefix, "-mat_view_nest_sub", &viewSub, NULL));
    917:     PetscCall(PetscViewerASCIIPushTab(viewer));
    918:     PetscCall(PetscViewerASCIIPrintf(viewer, "MatNest, rows=%" PetscInt_FMT ", cols=%" PetscInt_FMT ", structure:\n", bA->nr, bA->nc));
    919:     for (i = 0; i < bA->nr; i++) {
    920:       for (j = 0; j < bA->nc; j++) {
    921:         MatType   type;
    922:         char      name[256] = "", prefix[256] = "";
    923:         PetscInt  NR, NC;
    924:         PetscBool isNest = PETSC_FALSE;

    926:         if (!bA->m[i][j]) {
    927:           PetscCall(PetscViewerASCIIPrintf(viewer, "(%" PetscInt_FMT ",%" PetscInt_FMT ") : NULL\n", i, j));
    928:           continue;
    929:         }
    930:         PetscCall(MatGetSize(bA->m[i][j], &NR, &NC));
    931:         PetscCall(MatGetType(bA->m[i][j], &type));
    932:         if (((PetscObject)bA->m[i][j])->name) PetscCall(PetscSNPrintf(name, sizeof(name), "name=\"%s\", ", ((PetscObject)bA->m[i][j])->name));
    933:         if (((PetscObject)bA->m[i][j])->prefix) PetscCall(PetscSNPrintf(prefix, sizeof(prefix), "prefix=\"%s\", ", ((PetscObject)bA->m[i][j])->prefix));
    934:         PetscCall(PetscObjectTypeCompare((PetscObject)bA->m[i][j], MATNEST, &isNest));

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

    938:         if (isNest || viewSub) {
    939:           PetscCall(PetscViewerASCIIPushTab(viewer)); /* push1 */
    940:           PetscCall(MatView(bA->m[i][j], viewer));
    941:           PetscCall(PetscViewerASCIIPopTab(viewer)); /* pop1 */
    942:         }
    943:       }
    944:     }
    945:     PetscCall(PetscViewerASCIIPopTab(viewer)); /* pop0 */
    946:   }
    947:   PetscFunctionReturn(PETSC_SUCCESS);
    948: }

    950: static PetscErrorCode MatZeroEntries_Nest(Mat A)
    951: {
    952:   Mat_Nest *bA = (Mat_Nest *)A->data;
    953:   PetscInt  i, j;

    955:   PetscFunctionBegin;
    956:   for (i = 0; i < bA->nr; i++) {
    957:     for (j = 0; j < bA->nc; j++) {
    958:       if (!bA->m[i][j]) continue;
    959:       PetscCall(MatZeroEntries(bA->m[i][j]));
    960:     }
    961:   }
    962:   PetscFunctionReturn(PETSC_SUCCESS);
    963: }

    965: static PetscErrorCode MatCopy_Nest(Mat A, Mat B, MatStructure str)
    966: {
    967:   Mat_Nest *bA = (Mat_Nest *)A->data, *bB = (Mat_Nest *)B->data;
    968:   PetscInt  i, j, nr = bA->nr, nc = bA->nc;
    969:   PetscBool nnzstate = PETSC_FALSE;

    971:   PetscFunctionBegin;
    972:   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);
    973:   for (i = 0; i < nr; i++) {
    974:     for (j = 0; j < nc; j++) {
    975:       PetscObjectState subnnzstate = 0;
    976:       if (bA->m[i][j] && bB->m[i][j]) {
    977:         PetscCall(MatCopy(bA->m[i][j], bB->m[i][j], str));
    978:         PetscCall(MatGetNonzeroState(bB->m[i][j], &subnnzstate));
    979:         nnzstate                 = (PetscBool)(nnzstate || bB->nnzstate[i * nc + j] != subnnzstate);
    980:         bB->nnzstate[i * nc + j] = subnnzstate;
    981:       } else if (bA->m[i][j]) { // bB->m[i][j] is NULL
    982:         Mat M;

    984:         PetscCheck(str == DIFFERENT_NONZERO_PATTERN || str == UNKNOWN_NONZERO_PATTERN, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_INCOMP, "Matrix block does not exist at %" PetscInt_FMT ",%" PetscInt_FMT ". Use DIFFERENT_NONZERO_PATTERN or UNKNOWN_NONZERO_PATTERN", i, j);
    985:         PetscCall(MatDuplicate(bA->m[i][j], MAT_COPY_VALUES, &M));
    986:         PetscCall(MatNestSetSubMat(B, i, j, M));
    987:         PetscCall(MatDestroy(&M));
    988:       } else if (bB->m[i][j]) { // bA->m[i][j] is NULL
    989:         PetscCheck(str == DIFFERENT_NONZERO_PATTERN || str == SUBSET_NONZERO_PATTERN || str == UNKNOWN_NONZERO_PATTERN, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_INCOMP, "Matrix block does not exist at %" PetscInt_FMT ",%" PetscInt_FMT ". Use DIFFERENT_NONZERO_PATTERN, SUBSET_NONZERO_PATTERN or UNKNOWN_NONZERO_PATTERN", i, j);
    990:         PetscCall(MatNestSetSubMat(B, i, j, NULL));
    991:       }
    992:     }
    993:   }
    994:   if (nnzstate) B->nonzerostate++;
    995:   PetscFunctionReturn(PETSC_SUCCESS);
    996: }

    998: static PetscErrorCode MatAXPY_Nest(Mat Y, PetscScalar a, Mat X, MatStructure str)
    999: {
    1000:   Mat_Nest *bY = (Mat_Nest *)Y->data, *bX = (Mat_Nest *)X->data;
    1001:   PetscInt  i, j, nr = bY->nr, nc = bY->nc;
    1002:   PetscBool nnzstate = PETSC_FALSE;

    1004:   PetscFunctionBegin;
    1005:   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);
    1006:   for (i = 0; i < nr; i++) {
    1007:     for (j = 0; j < nc; j++) {
    1008:       PetscObjectState subnnzstate = 0;
    1009:       if (bY->m[i][j] && bX->m[i][j]) {
    1010:         PetscCall(MatAXPY(bY->m[i][j], a, bX->m[i][j], str));
    1011:       } else if (bX->m[i][j]) {
    1012:         Mat M;

    1014:         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);
    1015:         PetscCall(MatDuplicate(bX->m[i][j], MAT_COPY_VALUES, &M));
    1016:         PetscCall(MatScale(M, a));
    1017:         PetscCall(MatNestSetSubMat(Y, i, j, M));
    1018:         PetscCall(MatDestroy(&M));
    1019:       }
    1020:       if (bY->m[i][j]) PetscCall(MatGetNonzeroState(bY->m[i][j], &subnnzstate));
    1021:       nnzstate                 = (PetscBool)(nnzstate || bY->nnzstate[i * nc + j] != subnnzstate);
    1022:       bY->nnzstate[i * nc + j] = subnnzstate;
    1023:     }
    1024:   }
    1025:   if (nnzstate) Y->nonzerostate++;
    1026:   PetscFunctionReturn(PETSC_SUCCESS);
    1027: }

    1029: static PetscErrorCode MatDuplicate_Nest(Mat A, MatDuplicateOption op, Mat *B)
    1030: {
    1031:   Mat_Nest *bA = (Mat_Nest *)A->data;
    1032:   Mat      *b;
    1033:   PetscInt  i, j, nr = bA->nr, nc = bA->nc;

    1035:   PetscFunctionBegin;
    1036:   PetscCall(PetscMalloc1(nr * nc, &b));
    1037:   for (i = 0; i < nr; i++) {
    1038:     for (j = 0; j < nc; j++) {
    1039:       if (bA->m[i][j]) PetscCall(MatDuplicate(bA->m[i][j], op, &b[i * nc + j]));
    1040:       else b[i * nc + j] = NULL;
    1041:     }
    1042:   }
    1043:   PetscCall(MatCreateNest(PetscObjectComm((PetscObject)A), nr, bA->isglobal.row, nc, bA->isglobal.col, b, B));
    1044:   /* Give the new MatNest exclusive ownership */
    1045:   for (i = 0; i < nr * nc; i++) PetscCall(MatDestroy(&b[i]));
    1046:   PetscCall(PetscFree(b));

    1048:   PetscCall(MatAssemblyBegin(*B, MAT_FINAL_ASSEMBLY));
    1049:   PetscCall(MatAssemblyEnd(*B, MAT_FINAL_ASSEMBLY));
    1050:   PetscFunctionReturn(PETSC_SUCCESS);
    1051: }

    1053: /* nest api */
    1054: static PetscErrorCode MatNestGetSubMat_Nest(Mat A, PetscInt idxm, PetscInt jdxm, Mat *mat)
    1055: {
    1056:   Mat_Nest *bA = (Mat_Nest *)A->data;

    1058:   PetscFunctionBegin;
    1059:   PetscCheck(idxm < bA->nr, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Row too large: row %" PetscInt_FMT " max %" PetscInt_FMT, idxm, bA->nr - 1);
    1060:   PetscCheck(jdxm < bA->nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Col too large: row %" PetscInt_FMT " max %" PetscInt_FMT, jdxm, bA->nc - 1);
    1061:   *mat = bA->m[idxm][jdxm];
    1062:   PetscFunctionReturn(PETSC_SUCCESS);
    1063: }

    1065: /*@
    1066:   MatNestGetSubMat - Returns a single, sub-matrix from a `MATNEST`

    1068:   Not Collective

    1070:   Input Parameters:
    1071: + A    - `MATNEST` matrix
    1072: . idxm - index of the matrix within the nest matrix
    1073: - jdxm - index of the matrix within the nest matrix

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

    1078:   Level: developer

    1080: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSize()`, `MatNestGetSubMats()`, `MatCreateNest()`, `MatNestSetSubMat()`,
    1081:           `MatNestGetLocalISs()`, `MatNestGetISs()`
    1082: @*/
    1083: PetscErrorCode MatNestGetSubMat(Mat A, PetscInt idxm, PetscInt jdxm, Mat *sub)
    1084: {
    1085:   PetscFunctionBegin;
    1089:   PetscAssertPointer(sub, 4);
    1090:   PetscUseMethod(A, "MatNestGetSubMat_C", (Mat, PetscInt, PetscInt, Mat *), (A, idxm, jdxm, sub));
    1091:   PetscFunctionReturn(PETSC_SUCCESS);
    1092: }

    1094: static PetscErrorCode MatNestSetSubMat_Nest(Mat A, PetscInt idxm, PetscInt jdxm, Mat mat)
    1095: {
    1096:   Mat_Nest *bA = (Mat_Nest *)A->data;
    1097:   PetscInt  m, n, M, N, mi, ni, Mi, Ni;

    1099:   PetscFunctionBegin;
    1100:   PetscCheck(idxm < bA->nr, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Row too large: row %" PetscInt_FMT " max %" PetscInt_FMT, idxm, bA->nr - 1);
    1101:   PetscCheck(jdxm < bA->nc, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Col too large: row %" PetscInt_FMT " max %" PetscInt_FMT, jdxm, bA->nc - 1);
    1102:   if (mat) {
    1103:     PetscCall(MatGetLocalSize(mat, &m, &n));
    1104:     PetscCall(MatGetSize(mat, &M, &N));
    1105:     PetscCall(ISGetLocalSize(bA->isglobal.row[idxm], &mi));
    1106:     PetscCall(ISGetSize(bA->isglobal.row[idxm], &Mi));
    1107:     PetscCall(ISGetLocalSize(bA->isglobal.col[jdxm], &ni));
    1108:     PetscCall(ISGetSize(bA->isglobal.col[jdxm], &Ni));
    1109:     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);
    1110:     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);
    1111:   }

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

    1116:   PetscCall(PetscObjectReference((PetscObject)mat));
    1117:   PetscCall(MatDestroy(&bA->m[idxm][jdxm]));
    1118:   bA->m[idxm][jdxm] = mat;
    1119:   PetscCall(PetscObjectStateIncrease((PetscObject)A));
    1120:   if (mat) PetscCall(MatGetNonzeroState(mat, &bA->nnzstate[idxm * bA->nc + jdxm]));
    1121:   else bA->nnzstate[idxm * bA->nc + jdxm] = 0;
    1122:   A->nonzerostate++;
    1123:   PetscFunctionReturn(PETSC_SUCCESS);
    1124: }

    1126: /*@
    1127:   MatNestSetSubMat - Set a single submatrix in the `MATNEST`

    1129:   Logically Collective

    1131:   Input Parameters:
    1132: + A    - `MATNEST` matrix
    1133: . idxm - index of the matrix within the nest matrix
    1134: . jdxm - index of the matrix within the nest matrix
    1135: - sub  - matrix at index `idxm`, `jdxm` within the nest matrix

    1137:   Level: developer

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

    1142:   This increments the reference count of the submatrix.

    1144: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestSetSubMats()`, `MatNestGetSubMats()`, `MatNestGetLocalISs()`, `MatCreateNest()`,
    1145:           `MatNestGetSubMat()`, `MatNestGetISs()`, `MatNestGetSize()`
    1146: @*/
    1147: PetscErrorCode MatNestSetSubMat(Mat A, PetscInt idxm, PetscInt jdxm, Mat sub)
    1148: {
    1149:   PetscFunctionBegin;
    1154:   PetscTryMethod(A, "MatNestSetSubMat_C", (Mat, PetscInt, PetscInt, Mat), (A, idxm, jdxm, sub));
    1155:   PetscFunctionReturn(PETSC_SUCCESS);
    1156: }

    1158: static PetscErrorCode MatNestGetSubMats_Nest(Mat A, PetscInt *M, PetscInt *N, Mat ***mat)
    1159: {
    1160:   Mat_Nest *bA = (Mat_Nest *)A->data;

    1162:   PetscFunctionBegin;
    1163:   if (M) *M = bA->nr;
    1164:   if (N) *N = bA->nc;
    1165:   if (mat) *mat = bA->m;
    1166:   PetscFunctionReturn(PETSC_SUCCESS);
    1167: }

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

    1172:   Not Collective

    1174:   Input Parameter:
    1175: . A - nest matrix

    1177:   Output Parameters:
    1178: + M   - number of submatrix rows in the nest matrix
    1179: . N   - number of submatrix columns in the nest matrix
    1180: - mat - array of matrices

    1182:   Level: developer

    1184:   Note:
    1185:   The user should not free the array `mat`.

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

    1192: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSize()`, `MatNestGetSubMat()`, `MatNestGetLocalISs()`, `MatCreateNest()`,
    1193:           `MatNestSetSubMats()`, `MatNestGetISs()`, `MatNestSetSubMat()`
    1194: @*/
    1195: PetscErrorCode MatNestGetSubMats(Mat A, PetscInt *M, PetscInt *N, Mat ***mat)
    1196: {
    1197:   PetscFunctionBegin;
    1199:   PetscUseMethod(A, "MatNestGetSubMats_C", (Mat, PetscInt *, PetscInt *, Mat ***), (A, M, N, mat));
    1200:   PetscFunctionReturn(PETSC_SUCCESS);
    1201: }

    1203: static PetscErrorCode MatNestGetSize_Nest(Mat A, PetscInt *M, PetscInt *N)
    1204: {
    1205:   Mat_Nest *bA = (Mat_Nest *)A->data;

    1207:   PetscFunctionBegin;
    1208:   if (M) *M = bA->nr;
    1209:   if (N) *N = bA->nc;
    1210:   PetscFunctionReturn(PETSC_SUCCESS);
    1211: }

    1213: /*@
    1214:   MatNestGetSize - Returns the size of the `MATNEST` matrix.

    1216:   Not Collective

    1218:   Input Parameter:
    1219: . A - `MATNEST` matrix

    1221:   Output Parameters:
    1222: + M - number of rows in the nested mat
    1223: - N - number of cols in the nested mat

    1225:   Level: developer

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

    1230: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSubMat()`, `MatNestGetSubMats()`, `MatCreateNest()`, `MatNestGetLocalISs()`,
    1231:           `MatNestGetISs()`
    1232: @*/
    1233: PetscErrorCode MatNestGetSize(Mat A, PetscInt *M, PetscInt *N)
    1234: {
    1235:   PetscFunctionBegin;
    1237:   PetscUseMethod(A, "MatNestGetSize_C", (Mat, PetscInt *, PetscInt *), (A, M, N));
    1238:   PetscFunctionReturn(PETSC_SUCCESS);
    1239: }

    1241: static PetscErrorCode MatNestGetISs_Nest(Mat A, IS rows[], IS cols[])
    1242: {
    1243:   Mat_Nest *vs = (Mat_Nest *)A->data;
    1244:   PetscInt  i;

    1246:   PetscFunctionBegin;
    1247:   if (rows)
    1248:     for (i = 0; i < vs->nr; i++) rows[i] = vs->isglobal.row[i];
    1249:   if (cols)
    1250:     for (i = 0; i < vs->nc; i++) cols[i] = vs->isglobal.col[i];
    1251:   PetscFunctionReturn(PETSC_SUCCESS);
    1252: }

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

    1257:   Not Collective

    1259:   Input Parameter:
    1260: . A - `MATNEST` matrix

    1262:   Output Parameters:
    1263: + rows - array of row index sets (pass `NULL` to ignore)
    1264: - cols - array of column index sets (pass `NULL` to ignore)

    1266:   Level: advanced

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

    1271: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSubMat()`, `MatNestGetSubMats()`, `MatNestGetSize()`, `MatNestGetLocalISs()`,
    1272:           `MatCreateNest()`, `MatNestSetSubMats()`
    1273: @*/
    1274: PetscErrorCode MatNestGetISs(Mat A, IS rows[], IS cols[])
    1275: {
    1276:   PetscFunctionBegin;
    1278:   PetscUseMethod(A, "MatNestGetISs_C", (Mat, IS[], IS[]), (A, rows, cols));
    1279:   PetscFunctionReturn(PETSC_SUCCESS);
    1280: }

    1282: static PetscErrorCode MatNestGetLocalISs_Nest(Mat A, IS rows[], IS cols[])
    1283: {
    1284:   Mat_Nest *vs = (Mat_Nest *)A->data;
    1285:   PetscInt  i;

    1287:   PetscFunctionBegin;
    1288:   if (rows)
    1289:     for (i = 0; i < vs->nr; i++) rows[i] = vs->islocal.row[i];
    1290:   if (cols)
    1291:     for (i = 0; i < vs->nc; i++) cols[i] = vs->islocal.col[i];
    1292:   PetscFunctionReturn(PETSC_SUCCESS);
    1293: }

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

    1298:   Not Collective

    1300:   Input Parameter:
    1301: . A - `MATNEST` matrix

    1303:   Output Parameters:
    1304: + rows - array of row index sets (pass `NULL` to ignore)
    1305: - cols - array of column index sets (pass `NULL` to ignore)

    1307:   Level: advanced

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

    1312: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatNestGetSubMat()`, `MatNestGetSubMats()`, `MatNestGetSize()`, `MatNestGetISs()`, `MatCreateNest()`,
    1313:           `MatNestSetSubMats()`, `MatNestSetSubMat()`
    1314: @*/
    1315: PetscErrorCode MatNestGetLocalISs(Mat A, IS rows[], IS cols[])
    1316: {
    1317:   PetscFunctionBegin;
    1319:   PetscUseMethod(A, "MatNestGetLocalISs_C", (Mat, IS[], IS[]), (A, rows, cols));
    1320:   PetscFunctionReturn(PETSC_SUCCESS);
    1321: }

    1323: static PetscErrorCode MatNestSetVecType_Nest(Mat A, VecType vtype)
    1324: {
    1325:   PetscBool flg;

    1327:   PetscFunctionBegin;
    1328:   PetscCall(PetscStrcmp(vtype, VECNEST, &flg));
    1329:   /* In reality, this only distinguishes VECNEST and "other" */
    1330:   if (flg) A->ops->getvecs = MatCreateVecs_Nest;
    1331:   else A->ops->getvecs = NULL;
    1332:   PetscFunctionReturn(PETSC_SUCCESS);
    1333: }

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

    1338:   Not Collective

    1340:   Input Parameters:
    1341: + A     - `MATNEST` matrix
    1342: - vtype - `VecType` to use for creating vectors

    1344:   Level: developer

    1346: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreateVecs()`, `MatCreateNest()`, `VecType`
    1347: @*/
    1348: PetscErrorCode MatNestSetVecType(Mat A, VecType vtype)
    1349: {
    1350:   PetscFunctionBegin;
    1352:   PetscTryMethod(A, "MatNestSetVecType_C", (Mat, VecType), (A, vtype));
    1353:   PetscFunctionReturn(PETSC_SUCCESS);
    1354: }

    1356: static PetscErrorCode MatNestSetSubMats_Nest(Mat A, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[], const Mat a[])
    1357: {
    1358:   Mat_Nest *s = (Mat_Nest *)A->data;
    1359:   PetscInt  i, j, m, n, M, N;
    1360:   PetscBool cong, isstd, sametype = PETSC_FALSE;
    1361:   VecType   vtype, type;

    1363:   PetscFunctionBegin;
    1364:   PetscCall(MatReset_Nest(A));

    1366:   s->nr = nr;
    1367:   s->nc = nc;

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

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

    1404:   PetscCall(PetscMalloc1(nr, &s->row_len));
    1405:   PetscCall(PetscMalloc1(nc, &s->col_len));
    1406:   for (i = 0; i < nr; i++) s->row_len[i] = -1;
    1407:   for (j = 0; j < nc; j++) s->col_len[j] = -1;

    1409:   PetscCall(PetscCalloc1(nr * nc, &s->nnzstate));
    1410:   for (i = 0; i < nr; i++) {
    1411:     for (j = 0; j < nc; j++) {
    1412:       if (s->m[i][j]) PetscCall(MatGetNonzeroState(s->m[i][j], &s->nnzstate[i * nc + j]));
    1413:     }
    1414:   }

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

    1418:   PetscCall(PetscLayoutSetSize(A->rmap, M));
    1419:   PetscCall(PetscLayoutSetLocalSize(A->rmap, m));
    1420:   PetscCall(PetscLayoutSetSize(A->cmap, N));
    1421:   PetscCall(PetscLayoutSetLocalSize(A->cmap, n));

    1423:   PetscCall(PetscLayoutSetUp(A->rmap));
    1424:   PetscCall(PetscLayoutSetUp(A->cmap));

    1426:   /* disable operations that are not supported for non-square matrices,
    1427:      or matrices for which is_row != is_col  */
    1428:   PetscCall(MatHasCongruentLayouts(A, &cong));
    1429:   if (cong && nr != nc) cong = PETSC_FALSE;
    1430:   if (cong) {
    1431:     for (i = 0; cong && i < nr; i++) PetscCall(ISEqualUnsorted(s->isglobal.row[i], s->isglobal.col[i], &cong));
    1432:   }
    1433:   if (!cong) {
    1434:     A->ops->getdiagonal = NULL;
    1435:     A->ops->shift       = NULL;
    1436:     A->ops->diagonalset = NULL;
    1437:   }

    1439:   PetscCall(PetscCalloc2(nr, &s->left, nc, &s->right));
    1440:   PetscCall(PetscObjectStateIncrease((PetscObject)A));
    1441:   A->nonzerostate++;
    1442:   PetscFunctionReturn(PETSC_SUCCESS);
    1443: }

    1445: /*@
    1446:   MatNestSetSubMats - Sets the nested submatrices in a `MATNEST`

    1448:   Collective

    1450:   Input Parameters:
    1451: + A      - `MATNEST` matrix
    1452: . nr     - number of nested row blocks
    1453: . is_row - index sets for each nested row block, or `NULL` to make contiguous
    1454: . nc     - number of nested column blocks
    1455: . is_col - index sets for each nested column block, or `NULL` to make contiguous
    1456: - a      - array of $ nr \times nc$ submatrices, or `NULL`

    1458:   Level: advanced

    1460:   Notes:
    1461:   This always resets any block matrix information previously set.

    1463:   Pass `NULL` in the corresponding entry of `a` for an empty block.

    1465:   In both C and Fortran, `a` must be a one-dimensional array representing a two-dimensional row-major order array containing the matrices. See
    1466:   `MatCreateNest()` for an example.

    1468:   Fortran Note:
    1469:   Pass `PETSC_NULL_MAT` in the corresponding entry of `a` for an empty block

    1471: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreateNest()`, `MatNestSetSubMat()`, `MatNestGetSubMat()`, `MatNestGetSubMats()`
    1472: @*/
    1473: PetscErrorCode MatNestSetSubMats(Mat A, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[], const Mat a[]) PeNSS
    1474: {
    1475:   PetscFunctionBegin;
    1478:   PetscCheck(nr >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Number of rows cannot be negative");
    1479:   if (nr && is_row) {
    1480:     PetscAssertPointer(is_row, 3);
    1482:   }
    1484:   PetscCheck(nc >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_OUTOFRANGE, "Number of columns cannot be negative");
    1485:   if (nc && is_col) {
    1486:     PetscAssertPointer(is_col, 5);
    1488:   }
    1489:   PetscTryMethod(A, "MatNestSetSubMats_C", (Mat, PetscInt, const IS[], PetscInt, const IS[], const Mat[]), (A, nr, is_row, nc, is_col, a));
    1490:   PetscFunctionReturn(PETSC_SUCCESS);
    1491: }

    1493: static PetscErrorCode MatNestCreateAggregateL2G_Private(Mat A, PetscInt n, const IS islocal[], const IS isglobal[], PetscBool colflg, ISLocalToGlobalMapping *ltog)
    1494: {
    1495:   PetscBool flg;
    1496:   PetscInt  i, j, m, mi, *ix;

    1498:   PetscFunctionBegin;
    1499:   *ltog = NULL;
    1500:   for (i = 0, m = 0, flg = PETSC_FALSE; i < n; i++) {
    1501:     if (islocal[i]) {
    1502:       PetscCall(ISGetLocalSize(islocal[i], &mi));
    1503:       flg = PETSC_TRUE; /* We found a non-trivial entry */
    1504:     } else {
    1505:       PetscCall(ISGetLocalSize(isglobal[i], &mi));
    1506:     }
    1507:     m += mi;
    1508:   }
    1509:   if (!flg) PetscFunctionReturn(PETSC_SUCCESS);

    1511:   PetscCall(PetscMalloc1(m, &ix));
    1512:   for (i = 0, m = 0; i < n; i++) {
    1513:     ISLocalToGlobalMapping smap = NULL;
    1514:     Mat                    sub  = NULL;
    1515:     PetscSF                sf;
    1516:     PetscLayout            map;
    1517:     const PetscInt        *ix2;

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

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

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

    1572: /* If an IS was provided, there is nothing Nest needs to do, otherwise Nest will build a strided IS */
    1573: /*
    1574:   nprocessors = NP
    1575:   Nest x^T = ((g_0,g_1,...g_nprocs-1), (h_0,h_1,...h_NP-1))
    1576:        proc 0: => (g_0,h_0,)
    1577:        proc 1: => (g_1,h_1,)
    1578:        ...
    1579:        proc nprocs-1: => (g_NP-1,h_NP-1,)

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

    1584:             proc 0:
    1585:     is[1] = (nlocal(g_0),nlocal(g_0)+1,...,nlocal(g_0)+nlocal(h_0)-1)
    1586:             proc 1:
    1587:     is[1] = (nlocal(g_1),nlocal(g_1)+1,...,nlocal(g_1)+nlocal(h_1)-1)

    1589:             proc NP-1:
    1590:     is[1] = (nlocal(g_NP-1),nlocal(g_NP-1)+1,...,nlocal(g_NP-1)+nlocal(h_NP-1)-1)
    1591: */
    1592: static PetscErrorCode MatSetUp_NestIS_Private(Mat A, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[])
    1593: {
    1594:   Mat_Nest *vs = (Mat_Nest *)A->data;
    1595:   PetscInt  i, j, offset, n, nsum, bs;
    1596:   Mat       sub = NULL;

    1598:   PetscFunctionBegin;
    1599:   PetscCall(PetscMalloc1(nr, &vs->isglobal.row));
    1600:   PetscCall(PetscMalloc1(nc, &vs->isglobal.col));
    1601:   if (is_row) { /* valid IS is passed in */
    1602:     /* refs on is[] are incremented */
    1603:     for (i = 0; i < vs->nr; i++) {
    1604:       PetscCall(PetscObjectReference((PetscObject)is_row[i]));
    1605:       vs->isglobal.row[i] = is_row[i];
    1606:     }
    1607:   } else { /* Create the ISs by inspecting sizes of a submatrix in each row */
    1608:     nsum = 0;
    1609:     for (i = 0; i < vs->nr; i++) { /* Add up the local sizes to compute the aggregate offset */
    1610:       PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
    1611:       PetscCheck(sub, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "No nonzero submatrix in row %" PetscInt_FMT, i);
    1612:       PetscCall(MatGetLocalSize(sub, &n, NULL));
    1613:       PetscCheck(n >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Sizes have not yet been set for submatrix");
    1614:       nsum += n;
    1615:     }
    1616:     PetscCallMPI(MPI_Scan(&nsum, &offset, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)A)));
    1617:     offset -= nsum;
    1618:     for (i = 0; i < vs->nr; i++) {
    1619:       PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
    1620:       PetscCall(MatGetLocalSize(sub, &n, NULL));
    1621:       PetscCall(MatGetBlockSizes(sub, &bs, NULL));
    1622:       PetscCall(ISCreateStride(PetscObjectComm((PetscObject)sub), n, offset, 1, &vs->isglobal.row[i]));
    1623:       PetscCall(ISSetBlockSize(vs->isglobal.row[i], bs));
    1624:       offset += n;
    1625:     }
    1626:   }

    1628:   if (is_col) { /* valid IS is passed in */
    1629:     /* refs on is[] are incremented */
    1630:     for (j = 0; j < vs->nc; j++) {
    1631:       PetscCall(PetscObjectReference((PetscObject)is_col[j]));
    1632:       vs->isglobal.col[j] = is_col[j];
    1633:     }
    1634:   } else { /* Create the ISs by inspecting sizes of a submatrix in each column */
    1635:     offset = A->cmap->rstart;
    1636:     nsum   = 0;
    1637:     for (j = 0; j < vs->nc; j++) {
    1638:       PetscCall(MatNestFindNonzeroSubMatCol(A, j, &sub));
    1639:       PetscCheck(sub, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "No nonzero submatrix in column %" PetscInt_FMT, i);
    1640:       PetscCall(MatGetLocalSize(sub, NULL, &n));
    1641:       PetscCheck(n >= 0, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Sizes have not yet been set for submatrix");
    1642:       nsum += n;
    1643:     }
    1644:     PetscCallMPI(MPI_Scan(&nsum, &offset, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)A)));
    1645:     offset -= nsum;
    1646:     for (j = 0; j < vs->nc; j++) {
    1647:       PetscCall(MatNestFindNonzeroSubMatCol(A, j, &sub));
    1648:       PetscCall(MatGetLocalSize(sub, NULL, &n));
    1649:       PetscCall(MatGetBlockSizes(sub, NULL, &bs));
    1650:       PetscCall(ISCreateStride(PetscObjectComm((PetscObject)sub), n, offset, 1, &vs->isglobal.col[j]));
    1651:       PetscCall(ISSetBlockSize(vs->isglobal.col[j], bs));
    1652:       offset += n;
    1653:     }
    1654:   }

    1656:   /* Set up the local ISs */
    1657:   PetscCall(PetscMalloc1(vs->nr, &vs->islocal.row));
    1658:   PetscCall(PetscMalloc1(vs->nc, &vs->islocal.col));
    1659:   for (i = 0, offset = 0; i < vs->nr; i++) {
    1660:     IS                     isloc;
    1661:     ISLocalToGlobalMapping rmap = NULL;
    1662:     PetscInt               nlocal, bs;
    1663:     PetscCall(MatNestFindNonzeroSubMatRow(A, i, &sub));
    1664:     if (sub) PetscCall(MatGetLocalToGlobalMapping(sub, &rmap, NULL));
    1665:     if (rmap) {
    1666:       PetscCall(MatGetBlockSizes(sub, &bs, NULL));
    1667:       PetscCall(ISLocalToGlobalMappingGetSize(rmap, &nlocal));
    1668:       PetscCall(ISCreateStride(PETSC_COMM_SELF, nlocal, offset, 1, &isloc));
    1669:       PetscCall(ISSetBlockSize(isloc, bs));
    1670:     } else {
    1671:       nlocal = 0;
    1672:       isloc  = NULL;
    1673:     }
    1674:     vs->islocal.row[i] = isloc;
    1675:     offset += nlocal;
    1676:   }
    1677:   for (i = 0, offset = 0; i < vs->nc; i++) {
    1678:     IS                     isloc;
    1679:     ISLocalToGlobalMapping cmap = NULL;
    1680:     PetscInt               nlocal, bs;
    1681:     PetscCall(MatNestFindNonzeroSubMatCol(A, i, &sub));
    1682:     if (sub) PetscCall(MatGetLocalToGlobalMapping(sub, NULL, &cmap));
    1683:     if (cmap) {
    1684:       PetscCall(MatGetBlockSizes(sub, NULL, &bs));
    1685:       PetscCall(ISLocalToGlobalMappingGetSize(cmap, &nlocal));
    1686:       PetscCall(ISCreateStride(PETSC_COMM_SELF, nlocal, offset, 1, &isloc));
    1687:       PetscCall(ISSetBlockSize(isloc, bs));
    1688:     } else {
    1689:       nlocal = 0;
    1690:       isloc  = NULL;
    1691:     }
    1692:     vs->islocal.col[i] = isloc;
    1693:     offset += nlocal;
    1694:   }

    1696:   /* Set up the aggregate ISLocalToGlobalMapping */
    1697:   {
    1698:     ISLocalToGlobalMapping rmap, cmap;
    1699:     PetscCall(MatNestCreateAggregateL2G_Private(A, vs->nr, vs->islocal.row, vs->isglobal.row, PETSC_FALSE, &rmap));
    1700:     PetscCall(MatNestCreateAggregateL2G_Private(A, vs->nc, vs->islocal.col, vs->isglobal.col, PETSC_TRUE, &cmap));
    1701:     if (rmap && cmap) PetscCall(MatSetLocalToGlobalMapping(A, rmap, cmap));
    1702:     PetscCall(ISLocalToGlobalMappingDestroy(&rmap));
    1703:     PetscCall(ISLocalToGlobalMappingDestroy(&cmap));
    1704:   }

    1706:   if (PetscDefined(USE_DEBUG)) {
    1707:     for (i = 0; i < vs->nr; i++) {
    1708:       for (j = 0; j < vs->nc; j++) {
    1709:         PetscInt m, n, M, N, mi, ni, Mi, Ni;
    1710:         Mat      B = vs->m[i][j];
    1711:         if (!B) continue;
    1712:         PetscCall(MatGetSize(B, &M, &N));
    1713:         PetscCall(MatGetLocalSize(B, &m, &n));
    1714:         PetscCall(ISGetSize(vs->isglobal.row[i], &Mi));
    1715:         PetscCall(ISGetSize(vs->isglobal.col[j], &Ni));
    1716:         PetscCall(ISGetLocalSize(vs->isglobal.row[i], &mi));
    1717:         PetscCall(ISGetLocalSize(vs->isglobal.col[j], &ni));
    1718:         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);
    1719:         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);
    1720:       }
    1721:     }
    1722:   }

    1724:   /* Set A->assembled if all non-null blocks are currently assembled */
    1725:   for (i = 0; i < vs->nr; i++) {
    1726:     for (j = 0; j < vs->nc; j++) {
    1727:       if (vs->m[i][j] && !vs->m[i][j]->assembled) PetscFunctionReturn(PETSC_SUCCESS);
    1728:     }
    1729:   }
    1730:   A->assembled = PETSC_TRUE;
    1731:   PetscFunctionReturn(PETSC_SUCCESS);
    1732: }

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

    1737:   Collective

    1739:   Input Parameters:
    1740: + comm   - Communicator for the new `MATNEST`
    1741: . nr     - number of nested row blocks
    1742: . is_row - index sets for each nested row block, or `NULL` to make contiguous
    1743: . nc     - number of nested column blocks
    1744: . is_col - index sets for each nested column block, or `NULL` to make contiguous
    1745: - a      - array of $nr \times nc$ submatrices, empty submatrices can be passed using `NULL`

    1747:   Output Parameter:
    1748: . B - new matrix

    1750:   Level: advanced

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

    1758:   Fortran Note:
    1759:   Pass `PETSC_NULL_MAT` in the corresponding entry of `a` for an empty block

    1761: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreate()`, `VecCreateNest()`, `DMCreateMatrix()`, `MatNestSetSubMat()`,
    1762:           `MatNestGetSubMat()`, `MatNestGetLocalISs()`, `MatNestGetSize()`,
    1763:           `MatNestGetISs()`, `MatNestSetSubMats()`, `MatNestGetSubMats()`
    1764: @*/
    1765: PetscErrorCode MatCreateNest(MPI_Comm comm, PetscInt nr, const IS is_row[], PetscInt nc, const IS is_col[], const Mat a[], Mat *B) PeNSS
    1766: {
    1767:   PetscFunctionBegin;
    1768:   PetscCall(MatCreate(comm, B));
    1769:   PetscCall(MatSetType(*B, MATNEST));
    1770:   (*B)->preallocated = PETSC_TRUE;
    1771:   PetscCall(MatNestSetSubMats(*B, nr, is_row, nc, is_col, a));
    1772:   PetscFunctionReturn(PETSC_SUCCESS);
    1773: }

    1775: static PetscErrorCode MatConvert_Nest_SeqAIJ_fast(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
    1776: {
    1777:   Mat_Nest     *nest = (Mat_Nest *)A->data;
    1778:   Mat          *trans;
    1779:   PetscScalar **avv;
    1780:   PetscScalar  *vv;
    1781:   PetscInt    **aii, **ajj;
    1782:   PetscInt     *ii, *jj, *ci;
    1783:   PetscInt      nr, nc, nnz, i, j;
    1784:   PetscBool     done;

    1786:   PetscFunctionBegin;
    1787:   PetscCall(MatGetSize(A, &nr, &nc));
    1788:   if (reuse == MAT_REUSE_MATRIX) {
    1789:     PetscInt rnr;

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

    1807:         PetscCall(PetscObjectTypeCompare((PetscObject)B, MATTRANSPOSEVIRTUAL, &istrans));
    1808:         if (istrans) {
    1809:           Mat Bt;

    1811:           PetscCall(MatTransposeGetMat(B, &Bt));
    1812:           PetscCall(MatTranspose(Bt, MAT_INITIAL_MATRIX, &trans[i * nest->nc + j]));
    1813:           B = trans[i * nest->nc + j];
    1814:         } else {
    1815:           PetscCall(PetscObjectTypeCompare((PetscObject)B, MATHERMITIANTRANSPOSEVIRTUAL, &istrans));
    1816:           if (istrans) {
    1817:             Mat Bt;

    1819:             PetscCall(MatHermitianTransposeGetMat(B, &Bt));
    1820:             PetscCall(MatHermitianTranspose(Bt, MAT_INITIAL_MATRIX, &trans[i * nest->nc + j]));
    1821:             B = trans[i * nest->nc + j];
    1822:           }
    1823:         }
    1824:         PetscCall(MatGetRowIJ(B, 0, PETSC_FALSE, PETSC_FALSE, &nnr, (const PetscInt **)&nii, (const PetscInt **)&njj, &done));
    1825:         PetscCheck(done, PetscObjectComm((PetscObject)B), PETSC_ERR_PLIB, "MatGetRowIJ");
    1826:         PetscCall(MatSeqAIJGetArray(B, &naa));
    1827:         nnz += nii[nnr];

    1829:         aii[i * nest->nc + j] = nii;
    1830:         ajj[i * nest->nc + j] = njj;
    1831:         avv[i * nest->nc + j] = naa;
    1832:       }
    1833:     }
    1834:   }
    1835:   if (reuse != MAT_REUSE_MATRIX) {
    1836:     PetscCall(PetscMalloc1(nr + 1, &ii));
    1837:     PetscCall(PetscMalloc1(nnz, &jj));
    1838:     PetscCall(PetscMalloc1(nnz, &vv));
    1839:   } else {
    1840:     PetscCheck(nnz == ii[nr], PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "Cannot reuse matrix, wrong number of nonzeros");
    1841:   }

    1843:   /* new row pointer */
    1844:   PetscCall(PetscArrayzero(ii, nr + 1));
    1845:   for (i = 0; i < nest->nr; ++i) {
    1846:     PetscInt ncr, rst;

    1848:     PetscCall(ISStrideGetInfo(nest->isglobal.row[i], &rst, NULL));
    1849:     PetscCall(ISGetLocalSize(nest->isglobal.row[i], &ncr));
    1850:     for (j = 0; j < nest->nc; ++j) {
    1851:       if (aii[i * nest->nc + j]) {
    1852:         PetscInt *nii = aii[i * nest->nc + j];
    1853:         PetscInt  ir;

    1855:         for (ir = rst; ir < ncr + rst; ++ir) {
    1856:           ii[ir + 1] += nii[1] - nii[0];
    1857:           nii++;
    1858:         }
    1859:       }
    1860:     }
    1861:   }
    1862:   for (i = 0; i < nr; i++) ii[i + 1] += ii[i];

    1864:   /* construct CSR for the new matrix */
    1865:   PetscCall(PetscCalloc1(nr, &ci));
    1866:   for (i = 0; i < nest->nr; ++i) {
    1867:     PetscInt ncr, rst;

    1869:     PetscCall(ISStrideGetInfo(nest->isglobal.row[i], &rst, NULL));
    1870:     PetscCall(ISGetLocalSize(nest->isglobal.row[i], &ncr));
    1871:     for (j = 0; j < nest->nc; ++j) {
    1872:       if (aii[i * nest->nc + j]) {
    1873:         PetscScalar *nvv = avv[i * nest->nc + j], vscale = 1.0, vshift = 0.0;
    1874:         PetscInt    *nii = aii[i * nest->nc + j];
    1875:         PetscInt    *njj = ajj[i * nest->nc + j];
    1876:         PetscInt     ir, cst;

    1878:         if (trans[i * nest->nc + j]) {
    1879:           vscale = ((Mat_Shell *)nest->m[i][j]->data)->vscale;
    1880:           vshift = ((Mat_Shell *)nest->m[i][j]->data)->vshift;
    1881:         }
    1882:         PetscCall(ISStrideGetInfo(nest->isglobal.col[j], &cst, NULL));
    1883:         for (ir = rst; ir < ncr + rst; ++ir) {
    1884:           PetscInt ij, rsize = nii[1] - nii[0], ist = ii[ir] + ci[ir];

    1886:           for (ij = 0; ij < rsize; ij++) {
    1887:             jj[ist + ij] = *njj + cst;
    1888:             vv[ist + ij] = vscale * *nvv;
    1889:             if (PetscUnlikely(vshift != 0.0 && *njj == ir - rst)) vv[ist + ij] += vshift;
    1890:             njj++;
    1891:             nvv++;
    1892:           }
    1893:           ci[ir] += rsize;
    1894:           nii++;
    1895:         }
    1896:       }
    1897:     }
    1898:   }
    1899:   PetscCall(PetscFree(ci));

    1901:   /* restore info */
    1902:   for (i = 0; i < nest->nr; ++i) {
    1903:     for (j = 0; j < nest->nc; ++j) {
    1904:       Mat B = nest->m[i][j];
    1905:       if (B) {
    1906:         PetscInt nnr = 0, k = i * nest->nc + j;

    1908:         B = (trans[k] ? trans[k] : B);
    1909:         PetscCall(MatRestoreRowIJ(B, 0, PETSC_FALSE, PETSC_FALSE, &nnr, (const PetscInt **)&aii[k], (const PetscInt **)&ajj[k], &done));
    1910:         PetscCheck(done, PetscObjectComm((PetscObject)B), PETSC_ERR_PLIB, "MatRestoreRowIJ");
    1911:         PetscCall(MatSeqAIJRestoreArray(B, &avv[k]));
    1912:         PetscCall(MatDestroy(&trans[k]));
    1913:       }
    1914:     }
    1915:   }
    1916:   PetscCall(PetscFree4(aii, ajj, avv, trans));

    1918:   /* finalize newmat */
    1919:   if (reuse == MAT_INITIAL_MATRIX) {
    1920:     PetscCall(MatCreateSeqAIJWithArrays(PetscObjectComm((PetscObject)A), nr, nc, ii, jj, vv, newmat));
    1921:   } else if (reuse == MAT_INPLACE_MATRIX) {
    1922:     Mat B;

    1924:     PetscCall(MatCreateSeqAIJWithArrays(PetscObjectComm((PetscObject)A), nr, nc, ii, jj, vv, &B));
    1925:     PetscCall(MatHeaderReplace(A, &B));
    1926:   }
    1927:   PetscCall(MatAssemblyBegin(*newmat, MAT_FINAL_ASSEMBLY));
    1928:   PetscCall(MatAssemblyEnd(*newmat, MAT_FINAL_ASSEMBLY));
    1929:   {
    1930:     Mat_SeqAIJ *a = (Mat_SeqAIJ *)((*newmat)->data);
    1931:     a->free_a     = PETSC_TRUE;
    1932:     a->free_ij    = PETSC_TRUE;
    1933:   }
    1934:   PetscFunctionReturn(PETSC_SUCCESS);
    1935: }

    1937: PETSC_INTERN PetscErrorCode MatAXPY_Dense_Nest(Mat Y, PetscScalar a, Mat X)
    1938: {
    1939:   Mat_Nest *nest = (Mat_Nest *)X->data;
    1940:   PetscInt  i, j, k, rstart;
    1941:   PetscBool flg;

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

    2008: static PetscErrorCode MatConvert_Nest_AIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
    2009: {
    2010:   Mat_Nest   *nest = (Mat_Nest *)A->data;
    2011:   PetscInt    m, n, M, N, i, j, k, *dnnz, *onnz = NULL, rstart, cstart, cend;
    2012:   PetscMPIInt size;
    2013:   Mat         C;

    2015:   PetscFunctionBegin;
    2016:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
    2017:   if (size == 1) { /* look for a special case with SeqAIJ matrices and strided-1, contiguous, blocks */
    2018:     PetscInt  nf;
    2019:     PetscBool fast;

    2021:     PetscCall(PetscStrcmp(newtype, MATAIJ, &fast));
    2022:     if (!fast) PetscCall(PetscStrcmp(newtype, MATSEQAIJ, &fast));
    2023:     for (i = 0; i < nest->nr && fast; ++i) {
    2024:       for (j = 0; j < nest->nc && fast; ++j) {
    2025:         Mat B = nest->m[i][j];
    2026:         if (B) {
    2027:           PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQAIJ, &fast));
    2028:           if (!fast) {
    2029:             PetscBool istrans;

    2031:             PetscCall(PetscObjectTypeCompare((PetscObject)B, MATTRANSPOSEVIRTUAL, &istrans));
    2032:             if (istrans) {
    2033:               Mat Bt;

    2035:               PetscCall(MatTransposeGetMat(B, &Bt));
    2036:               PetscCall(PetscObjectTypeCompare((PetscObject)Bt, MATSEQAIJ, &fast));
    2037:             } else {
    2038:               PetscCall(PetscObjectTypeCompare((PetscObject)B, MATHERMITIANTRANSPOSEVIRTUAL, &istrans));
    2039:               if (istrans) {
    2040:                 Mat Bt;

    2042:                 PetscCall(MatHermitianTransposeGetMat(B, &Bt));
    2043:                 PetscCall(PetscObjectTypeCompare((PetscObject)Bt, MATSEQAIJ, &fast));
    2044:               }
    2045:             }
    2046:             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);
    2047:           }
    2048:         }
    2049:       }
    2050:     }
    2051:     for (i = 0, nf = 0; i < nest->nr && fast; ++i) {
    2052:       PetscCall(PetscObjectTypeCompare((PetscObject)nest->isglobal.row[i], ISSTRIDE, &fast));
    2053:       if (fast) {
    2054:         PetscInt f, s;

    2056:         PetscCall(ISStrideGetInfo(nest->isglobal.row[i], &f, &s));
    2057:         if (f != nf || s != 1) {
    2058:           fast = PETSC_FALSE;
    2059:         } else {
    2060:           PetscCall(ISGetSize(nest->isglobal.row[i], &f));
    2061:           nf += f;
    2062:         }
    2063:       }
    2064:     }
    2065:     for (i = 0, nf = 0; i < nest->nc && fast; ++i) {
    2066:       PetscCall(PetscObjectTypeCompare((PetscObject)nest->isglobal.col[i], ISSTRIDE, &fast));
    2067:       if (fast) {
    2068:         PetscInt f, s;

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

    2196: static PetscErrorCode MatConvert_Nest_Dense(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
    2197: {
    2198:   Mat      B;
    2199:   PetscInt m, n, M, N;

    2201:   PetscFunctionBegin;
    2202:   PetscCall(MatGetSize(A, &M, &N));
    2203:   PetscCall(MatGetLocalSize(A, &m, &n));
    2204:   if (reuse == MAT_REUSE_MATRIX) {
    2205:     B = *newmat;
    2206:     PetscCall(MatZeroEntries(B));
    2207:   } else {
    2208:     PetscCall(MatCreateDense(PetscObjectComm((PetscObject)A), m, PETSC_DECIDE, M, N, NULL, &B));
    2209:   }
    2210:   PetscCall(MatAXPY_Dense_Nest(B, 1.0, A));
    2211:   if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &B));
    2212:   else if (reuse == MAT_INITIAL_MATRIX) *newmat = B;
    2213:   PetscFunctionReturn(PETSC_SUCCESS);
    2214: }

    2216: static PetscErrorCode MatHasOperation_Nest(Mat mat, MatOperation op, PetscBool *has)
    2217: {
    2218:   Mat_Nest    *bA = (Mat_Nest *)mat->data;
    2219:   MatOperation opAdd;
    2220:   PetscInt     i, j, nr = bA->nr, nc = bA->nc;
    2221:   PetscBool    flg;

    2223:   PetscFunctionBegin;
    2224:   *has = PETSC_FALSE;
    2225:   if (op == MATOP_MULT || op == MATOP_MULT_ADD || op == MATOP_MULT_TRANSPOSE || op == MATOP_MULT_TRANSPOSE_ADD) {
    2226:     opAdd = (op == MATOP_MULT || op == MATOP_MULT_ADD ? MATOP_MULT_ADD : MATOP_MULT_TRANSPOSE_ADD);
    2227:     for (j = 0; j < nc; j++) {
    2228:       for (i = 0; i < nr; i++) {
    2229:         if (!bA->m[i][j]) continue;
    2230:         PetscCall(MatHasOperation(bA->m[i][j], opAdd, &flg));
    2231:         if (!flg) PetscFunctionReturn(PETSC_SUCCESS);
    2232:       }
    2233:     }
    2234:   }
    2235:   if (((void **)mat->ops)[op]) *has = PETSC_TRUE;
    2236:   PetscFunctionReturn(PETSC_SUCCESS);
    2237: }

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

    2242:   Level: intermediate

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

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

    2253: .seealso: [](ch_matrices), `Mat`, `MATNEST`, `MatCreate()`, `MatType`, `MatCreateNest()`, `MatNestSetSubMat()`, `MatNestGetSubMat()`,
    2254:           `VecCreateNest()`, `DMCreateMatrix()`, `DMCOMPOSITE`, `MatNestSetVecType()`, `MatNestGetLocalISs()`,
    2255:           `MatNestGetISs()`, `MatNestSetSubMats()`, `MatNestGetSubMats()`
    2256: M*/
    2257: PETSC_EXTERN PetscErrorCode MatCreate_Nest(Mat A)
    2258: {
    2259:   Mat_Nest *s;

    2261:   PetscFunctionBegin;
    2262:   PetscCall(PetscNew(&s));
    2263:   A->data = (void *)s;

    2265:   s->nr            = -1;
    2266:   s->nc            = -1;
    2267:   s->m             = NULL;
    2268:   s->splitassembly = PETSC_FALSE;

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

    2272:   A->ops->mult                      = MatMult_Nest;
    2273:   A->ops->multadd                   = MatMultAdd_Nest;
    2274:   A->ops->multtranspose             = MatMultTranspose_Nest;
    2275:   A->ops->multtransposeadd          = MatMultTransposeAdd_Nest;
    2276:   A->ops->transpose                 = MatTranspose_Nest;
    2277:   A->ops->multhermitiantranspose    = MatMultHermitianTranspose_Nest;
    2278:   A->ops->multhermitiantransposeadd = MatMultHermitianTransposeAdd_Nest;
    2279:   A->ops->assemblybegin             = MatAssemblyBegin_Nest;
    2280:   A->ops->assemblyend               = MatAssemblyEnd_Nest;
    2281:   A->ops->zeroentries               = MatZeroEntries_Nest;
    2282:   A->ops->copy                      = MatCopy_Nest;
    2283:   A->ops->axpy                      = MatAXPY_Nest;
    2284:   A->ops->duplicate                 = MatDuplicate_Nest;
    2285:   A->ops->createsubmatrix           = MatCreateSubMatrix_Nest;
    2286:   A->ops->destroy                   = MatDestroy_Nest;
    2287:   A->ops->view                      = MatView_Nest;
    2288:   A->ops->getvecs                   = NULL; /* Use VECNEST by calling MatNestSetVecType(A,VECNEST) */
    2289:   A->ops->getlocalsubmatrix         = MatGetLocalSubMatrix_Nest;
    2290:   A->ops->restorelocalsubmatrix     = MatRestoreLocalSubMatrix_Nest;
    2291:   A->ops->getdiagonal               = MatGetDiagonal_Nest;
    2292:   A->ops->diagonalscale             = MatDiagonalScale_Nest;
    2293:   A->ops->scale                     = MatScale_Nest;
    2294:   A->ops->shift                     = MatShift_Nest;
    2295:   A->ops->diagonalset               = MatDiagonalSet_Nest;
    2296:   A->ops->setrandom                 = MatSetRandom_Nest;
    2297:   A->ops->hasoperation              = MatHasOperation_Nest;

    2299:   A->spptr     = NULL;
    2300:   A->assembled = PETSC_FALSE;

    2302:   /* expose Nest api's */
    2303:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSubMat_C", MatNestGetSubMat_Nest));
    2304:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetSubMat_C", MatNestSetSubMat_Nest));
    2305:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSubMats_C", MatNestGetSubMats_Nest));
    2306:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetSize_C", MatNestGetSize_Nest));
    2307:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetISs_C", MatNestGetISs_Nest));
    2308:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestGetLocalISs_C", MatNestGetLocalISs_Nest));
    2309:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetVecType_C", MatNestSetVecType_Nest));
    2310:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatNestSetSubMats_C", MatNestSetSubMats_Nest));
    2311:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_mpiaij_C", MatConvert_Nest_AIJ));
    2312:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_seqaij_C", MatConvert_Nest_AIJ));
    2313:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_aij_C", MatConvert_Nest_AIJ));
    2314:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_is_C", MatConvert_Nest_IS));
    2315:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_mpidense_C", MatConvert_Nest_Dense));
    2316:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_nest_seqdense_C", MatConvert_Nest_Dense));
    2317:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_nest_seqdense_C", MatProductSetFromOptions_Nest_Dense));
    2318:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_nest_mpidense_C", MatProductSetFromOptions_Nest_Dense));

    2320:   PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATNEST));
    2321:   PetscFunctionReturn(PETSC_SUCCESS);
    2322: }