Date: 2026-03-01T05:25:15+00:00

Actual source code: pdipm.c

      1: #include <../src/tao/constrained/impls/ipm/pdipm.h>

      3: /*
      4:    TaoPDIPMEvaluateFunctionsAndJacobians - Evaluate the objective function f, gradient fx, constraints, and all the Jacobians at current vector

      6:    Collective

      8:    Input Parameter:
      9: +  tao - solver context
     10: -  x - vector at which all objects to be evaluated

     12:    Level: beginner

     14: .seealso: `TAOPDIPM`, `TaoPDIPMUpdateConstraints()`, `TaoPDIPMSetUpBounds()`
     15: */
     16: static PetscErrorCode TaoPDIPMEvaluateFunctionsAndJacobians(Tao tao, Vec x)
     17: {
     18:   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;

     20:   PetscFunctionBegin;
     21:   /* Compute user objective function and gradient */
     22:   PetscCall(TaoComputeObjectiveAndGradient(tao, x, &pdipm->obj, tao->gradient));

     24:   /* Equality constraints and Jacobian */
     25:   if (pdipm->Ng) {
     26:     PetscCall(TaoComputeEqualityConstraints(tao, x, tao->constraints_equality));
     27:     PetscCall(TaoComputeJacobianEquality(tao, x, tao->jacobian_equality, tao->jacobian_equality_pre));
     28:   }

     30:   /* Inequality constraints and Jacobian */
     31:   if (pdipm->Nh) {
     32:     PetscCall(TaoComputeInequalityConstraints(tao, x, tao->constraints_inequality));
     33:     PetscCall(TaoComputeJacobianInequality(tao, x, tao->jacobian_inequality, tao->jacobian_inequality_pre));
     34:   }
     35:   PetscFunctionReturn(PETSC_SUCCESS);
     36: }

     38: /*
     39:   TaoPDIPMUpdateConstraints - Update the vectors ce and ci at x

     41:   Collective

     43:   Input Parameter:
     44: + tao - Tao context
     45: - x - vector at which constraints to be evaluated

     47:    Level: beginner

     49: .seealso: `TAOPDIPM`, `TaoPDIPMEvaluateFunctionsAndJacobians()`
     50: */
     51: static PetscErrorCode TaoPDIPMUpdateConstraints(Tao tao, Vec x)
     52: {
     53:   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
     54:   PetscInt           i, offset, offset1, k, xstart;
     55:   PetscScalar       *carr;
     56:   const PetscInt    *ubptr, *lbptr, *bxptr, *fxptr;
     57:   const PetscScalar *xarr, *xuarr, *xlarr, *garr, *harr;

     59:   PetscFunctionBegin;
     60:   PetscCall(VecGetOwnershipRange(x, &xstart, NULL));
     61:   PetscCall(VecGetArrayRead(x, &xarr));
     62:   PetscCall(VecGetArrayRead(tao->XU, &xuarr));
     63:   PetscCall(VecGetArrayRead(tao->XL, &xlarr));

     65:   /* (1) Update ce vector */
     66:   PetscCall(VecGetArrayWrite(pdipm->ce, &carr));

     68:   if (pdipm->Ng) {
     69:     /* (1.a) Inserting updated g(x) */
     70:     PetscCall(VecGetArrayRead(tao->constraints_equality, &garr));
     71:     PetscCall(PetscMemcpy(carr, garr, pdipm->ng * sizeof(PetscScalar)));
     72:     PetscCall(VecRestoreArrayRead(tao->constraints_equality, &garr));
     73:   }

     75:   /* (1.b) Update xfixed */
     76:   if (pdipm->Nxfixed) {
     77:     offset = pdipm->ng;
     78:     PetscCall(ISGetIndices(pdipm->isxfixed, &fxptr)); /* global indices in x */
     79:     for (k = 0; k < pdipm->nxfixed; k++) {
     80:       i                = fxptr[k] - xstart;
     81:       carr[offset + k] = xarr[i] - xuarr[i];
     82:     }
     83:   }
     84:   PetscCall(VecRestoreArrayWrite(pdipm->ce, &carr));

     86:   /* (2) Update ci vector */
     87:   PetscCall(VecGetArrayWrite(pdipm->ci, &carr));

     89:   if (pdipm->Nh) {
     90:     /* (2.a) Inserting updated h(x) */
     91:     PetscCall(VecGetArrayRead(tao->constraints_inequality, &harr));
     92:     PetscCall(PetscMemcpy(carr, harr, pdipm->nh * sizeof(PetscScalar)));
     93:     PetscCall(VecRestoreArrayRead(tao->constraints_inequality, &harr));
     94:   }

     96:   /* (2.b) Update xub */
     97:   offset = pdipm->nh;
     98:   if (pdipm->Nxub) {
     99:     PetscCall(ISGetIndices(pdipm->isxub, &ubptr));
    100:     for (k = 0; k < pdipm->nxub; k++) {
    101:       i                = ubptr[k] - xstart;
    102:       carr[offset + k] = xuarr[i] - xarr[i];
    103:     }
    104:   }

    106:   if (pdipm->Nxlb) {
    107:     /* (2.c) Update xlb */
    108:     offset += pdipm->nxub;
    109:     PetscCall(ISGetIndices(pdipm->isxlb, &lbptr)); /* global indices in x */
    110:     for (k = 0; k < pdipm->nxlb; k++) {
    111:       i                = lbptr[k] - xstart;
    112:       carr[offset + k] = xarr[i] - xlarr[i];
    113:     }
    114:   }

    116:   if (pdipm->Nxbox) {
    117:     /* (2.d) Update xbox */
    118:     offset += pdipm->nxlb;
    119:     offset1 = offset + pdipm->nxbox;
    120:     PetscCall(ISGetIndices(pdipm->isxbox, &bxptr)); /* global indices in x */
    121:     for (k = 0; k < pdipm->nxbox; k++) {
    122:       i                 = bxptr[k] - xstart; /* local indices in x */
    123:       carr[offset + k]  = xuarr[i] - xarr[i];
    124:       carr[offset1 + k] = xarr[i] - xlarr[i];
    125:     }
    126:   }
    127:   PetscCall(VecRestoreArrayWrite(pdipm->ci, &carr));

    129:   /* Restoring Vectors */
    130:   PetscCall(VecRestoreArrayRead(x, &xarr));
    131:   PetscCall(VecRestoreArrayRead(tao->XU, &xuarr));
    132:   PetscCall(VecRestoreArrayRead(tao->XL, &xlarr));
    133:   PetscFunctionReturn(PETSC_SUCCESS);
    134: }

    136: /*
    137:    TaoPDIPMSetUpBounds - Create upper and lower bound vectors of x

    139:    Collective

    141:    Input Parameter:
    142: .  tao - holds pdipm and XL & XU

    144:    Level: beginner

    146: .seealso: `TAOPDIPM`, `TaoPDIPMUpdateConstraints`
    147: */
    148: static PetscErrorCode TaoPDIPMSetUpBounds(Tao tao)
    149: {
    150:   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
    151:   const PetscScalar *xl, *xu;
    152:   PetscInt           n, *ixlb, *ixub, *ixfixed, *ixfree, *ixbox, i, low, high, idx;
    153:   MPI_Comm           comm;
    154:   PetscInt           sendbuf[5], recvbuf[5];

    156:   PetscFunctionBegin;
    157:   /* Creates upper and lower bounds vectors on x, if not created already */
    158:   PetscCall(TaoComputeVariableBounds(tao));

    160:   PetscCall(VecGetLocalSize(tao->XL, &n));
    161:   PetscCall(PetscMalloc5(n, &ixlb, n, &ixub, n, &ixfree, n, &ixfixed, n, &ixbox));

    163:   PetscCall(VecGetOwnershipRange(tao->XL, &low, &high));
    164:   PetscCall(VecGetArrayRead(tao->XL, &xl));
    165:   PetscCall(VecGetArrayRead(tao->XU, &xu));
    166:   for (i = 0; i < n; i++) {
    167:     idx = low + i;
    168:     if ((PetscRealPart(xl[i]) > PETSC_NINFINITY) && (PetscRealPart(xu[i]) < PETSC_INFINITY)) {
    169:       if (PetscRealPart(xl[i]) == PetscRealPart(xu[i])) {
    170:         ixfixed[pdipm->nxfixed++] = idx;
    171:       } else ixbox[pdipm->nxbox++] = idx;
    172:     } else {
    173:       if ((PetscRealPart(xl[i]) > PETSC_NINFINITY) && (PetscRealPart(xu[i]) >= PETSC_INFINITY)) {
    174:         ixlb[pdipm->nxlb++] = idx;
    175:       } else if ((PetscRealPart(xl[i]) <= PETSC_NINFINITY) && (PetscRealPart(xu[i]) < PETSC_INFINITY)) {
    176:         ixub[pdipm->nxlb++] = idx;
    177:       } else ixfree[pdipm->nxfree++] = idx;
    178:     }
    179:   }
    180:   PetscCall(VecRestoreArrayRead(tao->XL, &xl));
    181:   PetscCall(VecRestoreArrayRead(tao->XU, &xu));

    183:   PetscCall(PetscObjectGetComm((PetscObject)tao, &comm));
    184:   sendbuf[0] = pdipm->nxlb;
    185:   sendbuf[1] = pdipm->nxub;
    186:   sendbuf[2] = pdipm->nxfixed;
    187:   sendbuf[3] = pdipm->nxbox;
    188:   sendbuf[4] = pdipm->nxfree;

    190:   PetscCallMPI(MPIU_Allreduce(sendbuf, recvbuf, 5, MPIU_INT, MPI_SUM, comm));
    191:   pdipm->Nxlb    = recvbuf[0];
    192:   pdipm->Nxub    = recvbuf[1];
    193:   pdipm->Nxfixed = recvbuf[2];
    194:   pdipm->Nxbox   = recvbuf[3];
    195:   pdipm->Nxfree  = recvbuf[4];

    197:   if (pdipm->Nxlb) PetscCall(ISCreateGeneral(comm, pdipm->nxlb, ixlb, PETSC_COPY_VALUES, &pdipm->isxlb));
    198:   if (pdipm->Nxub) PetscCall(ISCreateGeneral(comm, pdipm->nxub, ixub, PETSC_COPY_VALUES, &pdipm->isxub));
    199:   if (pdipm->Nxfixed) PetscCall(ISCreateGeneral(comm, pdipm->nxfixed, ixfixed, PETSC_COPY_VALUES, &pdipm->isxfixed));
    200:   if (pdipm->Nxbox) PetscCall(ISCreateGeneral(comm, pdipm->nxbox, ixbox, PETSC_COPY_VALUES, &pdipm->isxbox));
    201:   if (pdipm->Nxfree) PetscCall(ISCreateGeneral(comm, pdipm->nxfree, ixfree, PETSC_COPY_VALUES, &pdipm->isxfree));
    202:   PetscCall(PetscFree5(ixlb, ixub, ixfixed, ixbox, ixfree));
    203:   PetscFunctionReturn(PETSC_SUCCESS);
    204: }

    206: /*
    207:    TaoPDIPMInitializeSolution - Initialize `TAOPDIPM` solution X = [x; lambdae; lambdai; z].
    208:    X consists of four subvectors in the order [x; lambdae; lambdai; z]. These
    209:      four subvectors need to be initialized and its values copied over to X. Instead
    210:      of copying, we use `VecPlaceArray()`/`VecResetArray()` functions to share the memory locations for
    211:      X and the subvectors

    213:    Collective

    215:    Input Parameter:
    216: .  tao - Tao context

    218:    Level: beginner
    219: */
    220: static PetscErrorCode TaoPDIPMInitializeSolution(Tao tao)
    221: {
    222:   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
    223:   PetscScalar       *Xarr, *z, *lambdai;
    224:   PetscInt           i;
    225:   const PetscScalar *xarr, *h;

    227:   PetscFunctionBegin;
    228:   PetscCall(VecGetArrayWrite(pdipm->X, &Xarr));

    230:   /* Set Initialize X.x = tao->solution */
    231:   PetscCall(VecGetArrayRead(tao->solution, &xarr));
    232:   PetscCall(PetscMemcpy(Xarr, xarr, pdipm->nx * sizeof(PetscScalar)));
    233:   PetscCall(VecRestoreArrayRead(tao->solution, &xarr));

    235:   /* Initialize X.lambdae = 0.0 */
    236:   if (pdipm->lambdae) PetscCall(VecSet(pdipm->lambdae, 0.0));

    238:   /* Initialize X.lambdai = push_init_lambdai, X.z = push_init_slack */
    239:   if (pdipm->Nci) {
    240:     PetscCall(VecSet(pdipm->lambdai, pdipm->push_init_lambdai));
    241:     PetscCall(VecSet(pdipm->z, pdipm->push_init_slack));

    243:     /* Additional modification for X.lambdai and X.z */
    244:     PetscCall(VecGetArrayWrite(pdipm->lambdai, &lambdai));
    245:     PetscCall(VecGetArrayWrite(pdipm->z, &z));
    246:     if (pdipm->Nh) {
    247:       PetscCall(VecGetArrayRead(tao->constraints_inequality, &h));
    248:       for (i = 0; i < pdipm->nh; i++) {
    249:         if (h[i] < -pdipm->push_init_slack) z[i] = -h[i];
    250:         if (pdipm->mu / z[i] > pdipm->push_init_lambdai) lambdai[i] = pdipm->mu / z[i];
    251:       }
    252:       PetscCall(VecRestoreArrayRead(tao->constraints_inequality, &h));
    253:     }
    254:     PetscCall(VecRestoreArrayWrite(pdipm->lambdai, &lambdai));
    255:     PetscCall(VecRestoreArrayWrite(pdipm->z, &z));
    256:   }

    258:   PetscCall(VecRestoreArrayWrite(pdipm->X, &Xarr));
    259:   PetscFunctionReturn(PETSC_SUCCESS);
    260: }

    262: /*
    263:    TaoSNESJacobian_PDIPM - Evaluate the Hessian matrix at X

    265:    Input Parameter:
    266:    snes - SNES context
    267:    X - KKT Vector
    268:    *ctx - pdipm context

    270:    Output Parameter:
    271:    J - Hessian matrix
    272:    Jpre - matrix to build the preconditioner from
    273: */
    274: static PetscErrorCode TaoSNESJacobian_PDIPM(SNES snes, Vec X, Mat J, Mat Jpre, void *ctx)
    275: {
    276:   Tao                tao   = (Tao)ctx;
    277:   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
    278:   PetscInt           i, row, cols[2], Jrstart, rjstart, nc, j;
    279:   const PetscInt    *aj, *ranges, *Jranges, *rranges, *cranges;
    280:   const PetscScalar *Xarr, *aa;
    281:   PetscScalar        vals[2];
    282:   PetscInt           proc, nx_all, *nce_all = pdipm->nce_all;

    284:   PetscFunctionBegin;
    285:   PetscCall(MatGetOwnershipRanges(Jpre, &Jranges));
    286:   PetscCall(MatGetOwnershipRange(Jpre, &Jrstart, NULL));
    287:   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &rranges));
    288:   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &cranges));

    290:   PetscCall(VecGetArrayRead(X, &Xarr));

    292:   /* (1) insert Z and Ci to the 4th block of Jpre -- overwrite existing values */
    293:   if (pdipm->solve_symmetric_kkt) { /* 1 for eq 17 revised pdipm doc 0 for eq 18 (symmetric KKT) */
    294:     vals[0] = 1.0;
    295:     for (i = 0; i < pdipm->nci; i++) {
    296:       row     = Jrstart + pdipm->off_z + i;
    297:       cols[0] = Jrstart + pdipm->off_lambdai + i;
    298:       cols[1] = row;
    299:       vals[1] = Xarr[pdipm->off_lambdai + i] / Xarr[pdipm->off_z + i];
    300:       PetscCall(MatSetValues(Jpre, 1, &row, 2, cols, vals, INSERT_VALUES));
    301:     }
    302:   } else {
    303:     for (i = 0; i < pdipm->nci; i++) {
    304:       row     = Jrstart + pdipm->off_z + i;
    305:       cols[0] = Jrstart + pdipm->off_lambdai + i;
    306:       cols[1] = row;
    307:       vals[0] = Xarr[pdipm->off_z + i];
    308:       vals[1] = Xarr[pdipm->off_lambdai + i];
    309:       PetscCall(MatSetValues(Jpre, 1, &row, 2, cols, vals, INSERT_VALUES));
    310:     }
    311:   }

    313:   /* (2) insert 2nd row block of Jpre: [ grad g, 0, 0, 0] */
    314:   if (pdipm->Ng) {
    315:     PetscCall(MatGetOwnershipRange(tao->jacobian_equality, &rjstart, NULL));
    316:     for (i = 0; i < pdipm->ng; i++) {
    317:       row = Jrstart + pdipm->off_lambdae + i;

    319:       PetscCall(MatGetRow(tao->jacobian_equality, i + rjstart, &nc, &aj, &aa));
    320:       proc = 0;
    321:       for (j = 0; j < nc; j++) {
    322:         while (aj[j] >= cranges[proc + 1]) proc++;
    323:         cols[0] = aj[j] - cranges[proc] + Jranges[proc];
    324:         PetscCall(MatSetValue(Jpre, row, cols[0], aa[j], INSERT_VALUES));
    325:       }
    326:       PetscCall(MatRestoreRow(tao->jacobian_equality, i + rjstart, &nc, &aj, &aa));
    327:       if (pdipm->kkt_pd) {
    328:         /* add shift \delta_c */
    329:         PetscCall(MatSetValue(Jpre, row, row, -pdipm->deltac, INSERT_VALUES));
    330:       }
    331:     }
    332:   }

    334:   /* (3) insert 3rd row block of Jpre: [ -grad h, 0, deltac, I] */
    335:   if (pdipm->Nh) {
    336:     PetscCall(MatGetOwnershipRange(tao->jacobian_inequality, &rjstart, NULL));
    337:     for (i = 0; i < pdipm->nh; i++) {
    338:       row = Jrstart + pdipm->off_lambdai + i;
    339:       PetscCall(MatGetRow(tao->jacobian_inequality, i + rjstart, &nc, &aj, &aa));
    340:       proc = 0;
    341:       for (j = 0; j < nc; j++) {
    342:         while (aj[j] >= cranges[proc + 1]) proc++;
    343:         cols[0] = aj[j] - cranges[proc] + Jranges[proc];
    344:         PetscCall(MatSetValue(Jpre, row, cols[0], -aa[j], INSERT_VALUES));
    345:       }
    346:       PetscCall(MatRestoreRow(tao->jacobian_inequality, i + rjstart, &nc, &aj, &aa));
    347:       if (pdipm->kkt_pd) {
    348:         /* add shift \delta_c */
    349:         PetscCall(MatSetValue(Jpre, row, row, -pdipm->deltac, INSERT_VALUES));
    350:       }
    351:     }
    352:   }

    354:   /* (4) insert 1st row block of Jpre: [Wxx, grad g', -grad h', 0] */
    355:   if (pdipm->Ng) { /* grad g' */
    356:     PetscCall(MatTranspose(tao->jacobian_equality, MAT_REUSE_MATRIX, &pdipm->jac_equality_trans));
    357:   }
    358:   if (pdipm->Nh) { /* grad h' */
    359:     PetscCall(MatTranspose(tao->jacobian_inequality, MAT_REUSE_MATRIX, &pdipm->jac_inequality_trans));
    360:   }

    362:   PetscCall(VecPlaceArray(pdipm->x, Xarr));
    363:   PetscCall(TaoComputeHessian(tao, pdipm->x, tao->hessian, tao->hessian_pre));
    364:   PetscCall(VecResetArray(pdipm->x));

    366:   PetscCall(MatGetOwnershipRange(tao->hessian, &rjstart, NULL));
    367:   for (i = 0; i < pdipm->nx; i++) {
    368:     row = Jrstart + i;

    370:     /* insert Wxx = fxx + ... -- provided by user */
    371:     PetscCall(MatGetRow(tao->hessian, i + rjstart, &nc, &aj, &aa));
    372:     proc = 0;
    373:     for (j = 0; j < nc; j++) {
    374:       while (aj[j] >= cranges[proc + 1]) proc++;
    375:       cols[0] = aj[j] - cranges[proc] + Jranges[proc];
    376:       if (row == cols[0] && pdipm->kkt_pd) {
    377:         /* add shift deltaw to Wxx component */
    378:         PetscCall(MatSetValue(Jpre, row, cols[0], aa[j] + pdipm->deltaw, INSERT_VALUES));
    379:       } else {
    380:         PetscCall(MatSetValue(Jpre, row, cols[0], aa[j], INSERT_VALUES));
    381:       }
    382:     }
    383:     PetscCall(MatRestoreRow(tao->hessian, i + rjstart, &nc, &aj, &aa));

    385:     /* insert grad g' */
    386:     if (pdipm->ng) {
    387:       PetscCall(MatGetRow(pdipm->jac_equality_trans, i + rjstart, &nc, &aj, &aa));
    388:       PetscCall(MatGetOwnershipRanges(tao->jacobian_equality, &ranges));
    389:       proc = 0;
    390:       for (j = 0; j < nc; j++) {
    391:         /* find row ownership of */
    392:         while (aj[j] >= ranges[proc + 1]) proc++;
    393:         nx_all  = rranges[proc + 1] - rranges[proc];
    394:         cols[0] = aj[j] - ranges[proc] + Jranges[proc] + nx_all;
    395:         PetscCall(MatSetValue(Jpre, row, cols[0], aa[j], INSERT_VALUES));
    396:       }
    397:       PetscCall(MatRestoreRow(pdipm->jac_equality_trans, i + rjstart, &nc, &aj, &aa));
    398:     }

    400:     /* insert -grad h' */
    401:     if (pdipm->nh) {
    402:       PetscCall(MatGetRow(pdipm->jac_inequality_trans, i + rjstart, &nc, &aj, &aa));
    403:       PetscCall(MatGetOwnershipRanges(tao->jacobian_inequality, &ranges));
    404:       proc = 0;
    405:       for (j = 0; j < nc; j++) {
    406:         /* find row ownership of */
    407:         while (aj[j] >= ranges[proc + 1]) proc++;
    408:         nx_all  = rranges[proc + 1] - rranges[proc];
    409:         cols[0] = aj[j] - ranges[proc] + Jranges[proc] + nx_all + nce_all[proc];
    410:         PetscCall(MatSetValue(Jpre, row, cols[0], -aa[j], INSERT_VALUES));
    411:       }
    412:       PetscCall(MatRestoreRow(pdipm->jac_inequality_trans, i + rjstart, &nc, &aj, &aa));
    413:     }
    414:   }
    415:   PetscCall(VecRestoreArrayRead(X, &Xarr));

    417:   /* (6) assemble Jpre and J */
    418:   PetscCall(MatAssemblyBegin(Jpre, MAT_FINAL_ASSEMBLY));
    419:   PetscCall(MatAssemblyEnd(Jpre, MAT_FINAL_ASSEMBLY));

    421:   if (J != Jpre) {
    422:     PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
    423:     PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
    424:   }
    425:   PetscFunctionReturn(PETSC_SUCCESS);
    426: }

    428: /*
    429:    TaoSnesFunction_PDIPM - Evaluate KKT function at X

    431:    Input Parameter:
    432:    snes - SNES context
    433:    X - KKT Vector
    434:    *ctx - pdipm

    436:    Output Parameter:
    437:    F - Updated Lagrangian vector
    438: */
    439: static PetscErrorCode TaoSNESFunction_PDIPM(SNES snes, Vec X, Vec F, void *ctx)
    440: {
    441:   Tao                tao   = (Tao)ctx;
    442:   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
    443:   PetscScalar       *Farr;
    444:   Vec                x, L1;
    445:   PetscInt           i;
    446:   const PetscScalar *Xarr, *carr, *zarr, *larr;

    448:   PetscFunctionBegin;
    449:   PetscCall(VecSet(F, 0.0));

    451:   PetscCall(VecGetArrayRead(X, &Xarr));
    452:   PetscCall(VecGetArrayWrite(F, &Farr));

    454:   /* (0) Evaluate f, fx, gradG, gradH at X.x Note: pdipm->x is not changed below */
    455:   x = pdipm->x;
    456:   PetscCall(VecPlaceArray(x, Xarr));
    457:   PetscCall(TaoPDIPMEvaluateFunctionsAndJacobians(tao, x));

    459:   /* Update ce, ci, and Jci at X.x */
    460:   PetscCall(TaoPDIPMUpdateConstraints(tao, x));
    461:   PetscCall(VecResetArray(x));

    463:   /* (1) L1 = fx + (gradG'*DE + Jce_xfixed'*lambdae_xfixed) - (gradH'*DI + Jci_xb'*lambdai_xb) */
    464:   L1 = pdipm->x;
    465:   PetscCall(VecPlaceArray(L1, Farr)); /* L1 = 0.0 */
    466:   if (pdipm->Nci) {
    467:     if (pdipm->Nh) {
    468:       /* L1 += gradH'*DI. Note: tao->DI is not changed below */
    469:       PetscCall(VecPlaceArray(tao->DI, Xarr + pdipm->off_lambdai));
    470:       PetscCall(MatMultTransposeAdd(tao->jacobian_inequality, tao->DI, L1, L1));
    471:       PetscCall(VecResetArray(tao->DI));
    472:     }

    474:     /* L1 += Jci_xb'*lambdai_xb */
    475:     PetscCall(VecPlaceArray(pdipm->lambdai_xb, Xarr + pdipm->off_lambdai + pdipm->nh));
    476:     PetscCall(MatMultTransposeAdd(pdipm->Jci_xb, pdipm->lambdai_xb, L1, L1));
    477:     PetscCall(VecResetArray(pdipm->lambdai_xb));

    479:     /* L1 = - (gradH'*DI + Jci_xb'*lambdai_xb) */
    480:     PetscCall(VecScale(L1, -1.0));
    481:   }

    483:   /* L1 += fx */
    484:   PetscCall(VecAXPY(L1, 1.0, tao->gradient));

    486:   if (pdipm->Nce) {
    487:     if (pdipm->Ng) {
    488:       /* L1 += gradG'*DE. Note: tao->DE is not changed below */
    489:       PetscCall(VecPlaceArray(tao->DE, Xarr + pdipm->off_lambdae));
    490:       PetscCall(MatMultTransposeAdd(tao->jacobian_equality, tao->DE, L1, L1));
    491:       PetscCall(VecResetArray(tao->DE));
    492:     }
    493:     if (pdipm->Nxfixed) {
    494:       /* L1 += Jce_xfixed'*lambdae_xfixed */
    495:       PetscCall(VecPlaceArray(pdipm->lambdae_xfixed, Xarr + pdipm->off_lambdae + pdipm->ng));
    496:       PetscCall(MatMultTransposeAdd(pdipm->Jce_xfixed, pdipm->lambdae_xfixed, L1, L1));
    497:       PetscCall(VecResetArray(pdipm->lambdae_xfixed));
    498:     }
    499:   }
    500:   PetscCall(VecResetArray(L1));

    502:   /* (2) L2 = ce(x) */
    503:   if (pdipm->Nce) {
    504:     PetscCall(VecGetArrayRead(pdipm->ce, &carr));
    505:     for (i = 0; i < pdipm->nce; i++) Farr[pdipm->off_lambdae + i] = carr[i];
    506:     PetscCall(VecRestoreArrayRead(pdipm->ce, &carr));
    507:   }

    509:   if (pdipm->Nci) {
    510:     if (pdipm->solve_symmetric_kkt) {
    511:       /* (3) L3 = z - ci(x);
    512:          (4) L4 = Lambdai * e - mu/z *e  */
    513:       PetscCall(VecGetArrayRead(pdipm->ci, &carr));
    514:       larr = Xarr + pdipm->off_lambdai;
    515:       zarr = Xarr + pdipm->off_z;
    516:       for (i = 0; i < pdipm->nci; i++) {
    517:         Farr[pdipm->off_lambdai + i] = zarr[i] - carr[i];
    518:         Farr[pdipm->off_z + i]       = larr[i] - pdipm->mu / zarr[i];
    519:       }
    520:       PetscCall(VecRestoreArrayRead(pdipm->ci, &carr));
    521:     } else {
    522:       /* (3) L3 = z - ci(x);
    523:          (4) L4 = Z * Lambdai * e - mu * e  */
    524:       PetscCall(VecGetArrayRead(pdipm->ci, &carr));
    525:       larr = Xarr + pdipm->off_lambdai;
    526:       zarr = Xarr + pdipm->off_z;
    527:       for (i = 0; i < pdipm->nci; i++) {
    528:         Farr[pdipm->off_lambdai + i] = zarr[i] - carr[i];
    529:         Farr[pdipm->off_z + i]       = zarr[i] * larr[i] - pdipm->mu;
    530:       }
    531:       PetscCall(VecRestoreArrayRead(pdipm->ci, &carr));
    532:     }
    533:   }

    535:   PetscCall(VecRestoreArrayRead(X, &Xarr));
    536:   PetscCall(VecRestoreArrayWrite(F, &Farr));
    537:   PetscFunctionReturn(PETSC_SUCCESS);
    538: }

    540: /*
    541:   Evaluate F(X); then update tao->gnorm0, tao->step = mu,
    542:   tao->residual = norm2(F_x,F_z) and tao->cnorm = norm2(F_ce,F_ci).
    543: */
    544: static PetscErrorCode TaoSNESFunction_PDIPM_residual(SNES snes, Vec X, Vec F, void *ctx)
    545: {
    546:   Tao                tao   = (Tao)ctx;
    547:   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
    548:   PetscScalar       *Farr, *tmparr;
    549:   Vec                L1;
    550:   PetscInt           i;
    551:   PetscReal          res[2], cnorm[2];
    552:   const PetscScalar *Xarr = NULL;

    554:   PetscFunctionBegin;
    555:   PetscCall(TaoSNESFunction_PDIPM(snes, X, F, (void *)tao));
    556:   PetscCall(VecGetArrayWrite(F, &Farr));
    557:   PetscCall(VecGetArrayRead(X, &Xarr));

    559:   /* compute res[0] = norm2(F_x) */
    560:   L1 = pdipm->x;
    561:   PetscCall(VecPlaceArray(L1, Farr));
    562:   PetscCall(VecNorm(L1, NORM_2, &res[0]));
    563:   PetscCall(VecResetArray(L1));

    565:   /* compute res[1] = norm2(F_z), cnorm[1] = norm2(F_ci) */
    566:   if (pdipm->z) {
    567:     if (pdipm->solve_symmetric_kkt) {
    568:       PetscCall(VecPlaceArray(pdipm->z, Farr + pdipm->off_z));
    569:       if (pdipm->Nci) {
    570:         PetscCall(VecGetArrayWrite(pdipm->z, &tmparr));
    571:         for (i = 0; i < pdipm->nci; i++) tmparr[i] *= Xarr[pdipm->off_z + i];
    572:         PetscCall(VecRestoreArrayWrite(pdipm->z, &tmparr));
    573:       }

    575:       PetscCall(VecNorm(pdipm->z, NORM_2, &res[1]));

    577:       if (pdipm->Nci) {
    578:         PetscCall(VecGetArrayWrite(pdipm->z, &tmparr));
    579:         for (i = 0; i < pdipm->nci; i++) tmparr[i] /= Xarr[pdipm->off_z + i];
    580:         PetscCall(VecRestoreArrayWrite(pdipm->z, &tmparr));
    581:       }
    582:       PetscCall(VecResetArray(pdipm->z));
    583:     } else { /* !solve_symmetric_kkt */
    584:       PetscCall(VecPlaceArray(pdipm->z, Farr + pdipm->off_z));
    585:       PetscCall(VecNorm(pdipm->z, NORM_2, &res[1]));
    586:       PetscCall(VecResetArray(pdipm->z));
    587:     }

    589:     PetscCall(VecPlaceArray(pdipm->ci, Farr + pdipm->off_lambdai));
    590:     PetscCall(VecNorm(pdipm->ci, NORM_2, &cnorm[1]));
    591:     PetscCall(VecResetArray(pdipm->ci));
    592:   } else {
    593:     res[1]   = 0.0;
    594:     cnorm[1] = 0.0;
    595:   }

    597:   /* compute cnorm[0] = norm2(F_ce) */
    598:   if (pdipm->Nce) {
    599:     PetscCall(VecPlaceArray(pdipm->ce, Farr + pdipm->off_lambdae));
    600:     PetscCall(VecNorm(pdipm->ce, NORM_2, &cnorm[0]));
    601:     PetscCall(VecResetArray(pdipm->ce));
    602:   } else cnorm[0] = 0.0;

    604:   PetscCall(VecRestoreArrayWrite(F, &Farr));
    605:   PetscCall(VecRestoreArrayRead(X, &Xarr));

    607:   tao->gnorm0   = tao->residual;
    608:   tao->residual = PetscSqrtReal(res[0] * res[0] + res[1] * res[1]);
    609:   tao->cnorm    = PetscSqrtReal(cnorm[0] * cnorm[0] + cnorm[1] * cnorm[1]);
    610:   tao->step     = pdipm->mu;
    611:   PetscFunctionReturn(PETSC_SUCCESS);
    612: }

    614: /*
    615:   PCPostSetup_PDIPM -- called when the KKT matrix is Cholesky factored for the preconditioner. Checks the inertia of Cholesky factor of the KKT matrix.
    616:   If it does not match the numbers of prime and dual variables, add shifts to the KKT matrix.
    617: */
    618: static PetscErrorCode PCPostSetUp_PDIPM(PC pc)
    619: {
    620:   Tao        tao;
    621:   TAO_PDIPM *pdipm;
    622:   Vec        X;
    623:   SNES       snes;
    624:   KSP        ksp;
    625:   Mat        Factor;
    626:   PetscBool  isCHOL;
    627:   PetscInt   nneg, nzero, npos;

    629:   PetscFunctionBegin;
    630:   PetscCall(PCGetApplicationContext(pc, &tao));
    631:   pdipm = (TAO_PDIPM *)tao->data;
    632:   X     = pdipm->X;
    633:   snes  = pdipm->snes;

    635:   /* Get the inertia of Cholesky factor */
    636:   PetscCall(SNESGetKSP(snes, &ksp));
    637:   PetscCall(KSPGetPC(ksp, &pc));
    638:   PetscCall(PetscObjectTypeCompare((PetscObject)pc, PCCHOLESKY, &isCHOL));
    639:   if (!isCHOL) PetscFunctionReturn(PETSC_SUCCESS);

    641:   PetscCall(PCFactorGetMatrix(pc, &Factor));
    642:   PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));

    644:   if (npos < pdipm->Nx + pdipm->Nci) {
    645:     pdipm->deltaw = PetscMax(pdipm->lastdeltaw / 3, 1.e-4 * PETSC_MACHINE_EPSILON);
    646:     PetscCall(PetscInfo(tao, "Test reduced deltaw=%g; previous MatInertia: nneg %" PetscInt_FMT ", nzero %" PetscInt_FMT ", npos %" PetscInt_FMT "(<%" PetscInt_FMT ")\n", (double)pdipm->deltaw, nneg, nzero, npos, pdipm->Nx + pdipm->Nci));
    647:     PetscCall(TaoSNESJacobian_PDIPM(snes, X, pdipm->K, pdipm->K, tao));
    648:     PetscCall(PCSetPostSetUp(pc, NULL));
    649:     PetscCall(PCSetUp(pc));
    650:     PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));

    652:     if (npos < pdipm->Nx + pdipm->Nci) {
    653:       pdipm->deltaw = pdipm->lastdeltaw;                                           /* in case reduction update does not help, this prevents that step from impacting increasing update */
    654:       while (npos < pdipm->Nx + pdipm->Nci && pdipm->deltaw <= 1. / PETSC_SMALL) { /* increase deltaw */
    655:         PetscCall(PetscInfo(tao, "  deltaw=%g fails, MatInertia: nneg %" PetscInt_FMT ", nzero %" PetscInt_FMT ", npos %" PetscInt_FMT "(<%" PetscInt_FMT ")\n", (double)pdipm->deltaw, nneg, nzero, npos, pdipm->Nx + pdipm->Nci));
    656:         pdipm->deltaw = PetscMin(8 * pdipm->deltaw, PetscPowReal(10, 20));
    657:         PetscCall(TaoSNESJacobian_PDIPM(snes, X, pdipm->K, pdipm->K, tao));
    658:         PetscCall(PCSetUp(pc));
    659:         PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));
    660:       }

    662:       PetscCheck(pdipm->deltaw < 1. / PETSC_SMALL, PetscObjectComm((PetscObject)tao), PETSC_ERR_CONV_FAILED, "Reached maximum delta w will not converge, try different initial x0");

    664:       PetscCall(PetscInfo(tao, "Updated deltaw %g\n", (double)pdipm->deltaw));
    665:       pdipm->lastdeltaw = pdipm->deltaw;
    666:       pdipm->deltaw     = 0.0;
    667:     }
    668:   }

    670:   if (nzero) { /* Jacobian is singular */
    671:     if (pdipm->deltac == 0.0) {
    672:       pdipm->deltac = PETSC_SQRT_MACHINE_EPSILON;
    673:     } else {
    674:       pdipm->deltac = pdipm->deltac * PetscPowReal(pdipm->mu, .25);
    675:     }
    676:     PetscCall(PetscInfo(tao, "Updated deltac=%g, MatInertia: nneg %" PetscInt_FMT ", nzero %" PetscInt_FMT "(!=0), npos %" PetscInt_FMT "\n", (double)pdipm->deltac, nneg, nzero, npos));
    677:     PetscCall(TaoSNESJacobian_PDIPM(snes, X, pdipm->K, pdipm->K, tao));
    678:     PetscCall(PCSetPostSetUp(pc, NULL));
    679:     PetscCall(PCSetUp(pc));
    680:     PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));
    681:   }
    682:   PetscCall(PCSetPostSetUp(pc, PCPostSetUp_PDIPM));
    683:   PetscFunctionReturn(PETSC_SUCCESS);
    684: }

    686: /*
    687:    SNESLineSearch_PDIPM - Custom line search used with PDIPM.

    689:    Collective

    691:    Notes:
    692:    This routine employs a simple backtracking line-search to keep
    693:    the slack variables (z) and inequality constraints Lagrange multipliers
    694:    (lambdai) positive, i.e., z,lambdai >=0. It does this by calculating scalars
    695:    alpha_p and alpha_d to keep z,lambdai non-negative. The decision (x), and the
    696:    slack variables are updated as X = X - alpha_d*dx. The constraint multipliers
    697:    are updated as Lambdai = Lambdai + alpha_p*dLambdai. The barrier parameter mu
    698:    is also updated as mu = mu + z'lambdai/Nci
    699: */
    700: static PetscErrorCode SNESLineSearch_PDIPM(SNESLineSearch linesearch, void *ctx)
    701: {
    702:   Tao                tao   = (Tao)ctx;
    703:   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
    704:   SNES               snes;
    705:   Vec                X, F, Y;
    706:   PetscInt           i, iter;
    707:   PetscReal          alpha_p = 1.0, alpha_d = 1.0, alpha[4];
    708:   PetscScalar       *Xarr, *z, *lambdai, dot, *taosolarr;
    709:   const PetscScalar *dXarr, *dz, *dlambdai;

    711:   PetscFunctionBegin;
    712:   PetscCall(SNESLineSearchGetSNES(linesearch, &snes));
    713:   PetscCall(SNESGetIterationNumber(snes, &iter));

    715:   PetscCall(SNESLineSearchSetReason(linesearch, SNES_LINESEARCH_SUCCEEDED));
    716:   PetscCall(SNESLineSearchGetVecs(linesearch, &X, &F, &Y, NULL, NULL));

    718:   PetscCall(VecGetArrayWrite(X, &Xarr));
    719:   PetscCall(VecGetArrayRead(Y, &dXarr));
    720:   z  = Xarr + pdipm->off_z;
    721:   dz = dXarr + pdipm->off_z;
    722:   for (i = 0; i < pdipm->nci; i++) {
    723:     if (z[i] - dz[i] < 0.0) alpha_p = PetscMin(alpha_p, 0.9999 * z[i] / dz[i]);
    724:   }

    726:   lambdai  = Xarr + pdipm->off_lambdai;
    727:   dlambdai = dXarr + pdipm->off_lambdai;

    729:   for (i = 0; i < pdipm->nci; i++) {
    730:     if (lambdai[i] - dlambdai[i] < 0.0) alpha_d = PetscMin(0.9999 * lambdai[i] / dlambdai[i], alpha_d);
    731:   }

    733:   alpha[0] = alpha_p;
    734:   alpha[1] = alpha_d;
    735:   PetscCall(VecRestoreArrayRead(Y, &dXarr));
    736:   PetscCall(VecRestoreArrayWrite(X, &Xarr));

    738:   /* alpha = min(alpha) over all processes */
    739:   PetscCallMPI(MPIU_Allreduce(alpha, alpha + 2, 2, MPIU_REAL, MPIU_MIN, PetscObjectComm((PetscObject)tao)));

    741:   alpha_p = alpha[2];
    742:   alpha_d = alpha[3];

    744:   /* X = X - alpha * Y */
    745:   PetscCall(VecGetArrayWrite(X, &Xarr));
    746:   PetscCall(VecGetArrayRead(Y, &dXarr));
    747:   for (i = 0; i < pdipm->nx; i++) Xarr[i] -= alpha_p * dXarr[i];
    748:   for (i = 0; i < pdipm->nce; i++) Xarr[i + pdipm->off_lambdae] -= alpha_d * dXarr[i + pdipm->off_lambdae];

    750:   for (i = 0; i < pdipm->nci; i++) {
    751:     Xarr[i + pdipm->off_lambdai] -= alpha_d * dXarr[i + pdipm->off_lambdai];
    752:     Xarr[i + pdipm->off_z] -= alpha_p * dXarr[i + pdipm->off_z];
    753:   }
    754:   PetscCall(VecGetArrayWrite(tao->solution, &taosolarr));
    755:   PetscCall(PetscMemcpy(taosolarr, Xarr, pdipm->nx * sizeof(PetscScalar)));
    756:   PetscCall(VecRestoreArrayWrite(tao->solution, &taosolarr));

    758:   PetscCall(VecRestoreArrayWrite(X, &Xarr));
    759:   PetscCall(VecRestoreArrayRead(Y, &dXarr));

    761:   /* Update mu = mu_update_factor * dot(z,lambdai)/pdipm->nci at updated X */
    762:   if (pdipm->z) PetscCall(VecDot(pdipm->z, pdipm->lambdai, &dot));
    763:   else dot = 0.0;

    765:   /* if (PetscAbsReal(pdipm->gradL) < 0.9*pdipm->mu)  */
    766:   pdipm->mu = pdipm->mu_update_factor * dot / pdipm->Nci;

    768:   /* Update F; get tao->residual and tao->cnorm */
    769:   PetscCall(TaoSNESFunction_PDIPM_residual(snes, X, F, (void *)tao));

    771:   tao->niter++;
    772:   PetscCall(TaoLogConvergenceHistory(tao, pdipm->obj, tao->residual, tao->cnorm, tao->niter));
    773:   PetscCall(TaoMonitor(tao, tao->niter, pdipm->obj, tao->residual, tao->cnorm, pdipm->mu));

    775:   PetscUseTypeMethod(tao, convergencetest, tao->cnvP);
    776:   if (tao->reason) PetscCall(SNESSetConvergedReason(snes, SNES_CONVERGED_FNORM_ABS));
    777:   PetscFunctionReturn(PETSC_SUCCESS);
    778: }

    780: static PetscErrorCode TaoSolve_PDIPM(Tao tao)
    781: {
    782:   TAO_PDIPM     *pdipm = (TAO_PDIPM *)tao->data;
    783:   SNESLineSearch linesearch; /* SNESLineSearch context */
    784:   Vec            dummy;

    786:   PetscFunctionBegin;
    787:   PetscCheck(tao->constraints_equality || tao->constraints_inequality, PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_NULL, "Equality and inequality constraints are not set. Either set them or switch to a different algorithm");

    789:   /* Initialize all variables */
    790:   PetscCall(TaoPDIPMInitializeSolution(tao));

    792:   /* Set linesearch */
    793:   PetscCall(SNESGetLineSearch(pdipm->snes, &linesearch));
    794:   PetscCall(SNESLineSearchSetType(linesearch, SNESLINESEARCHSHELL));
    795:   PetscCall(SNESLineSearchShellSetApply(linesearch, SNESLineSearch_PDIPM, tao));
    796:   PetscCall(SNESLineSearchSetFromOptions(linesearch));

    798:   tao->reason = TAO_CONTINUE_ITERATING;

    800:   /* -tao_monitor for iteration 0 and check convergence */
    801:   PetscCall(VecDuplicate(pdipm->X, &dummy));
    802:   PetscCall(TaoSNESFunction_PDIPM_residual(pdipm->snes, pdipm->X, dummy, (void *)tao));

    804:   PetscCall(TaoLogConvergenceHistory(tao, pdipm->obj, tao->residual, tao->cnorm, tao->niter));
    805:   PetscCall(TaoMonitor(tao, tao->niter, pdipm->obj, tao->residual, tao->cnorm, pdipm->mu));
    806:   PetscCall(VecDestroy(&dummy));
    807:   PetscUseTypeMethod(tao, convergencetest, tao->cnvP);
    808:   if (tao->reason) PetscCall(SNESSetConvergedReason(pdipm->snes, SNES_CONVERGED_FNORM_ABS));

    810:   while (tao->reason == TAO_CONTINUE_ITERATING) {
    811:     SNESConvergedReason reason;
    812:     PetscCall(SNESSolve(pdipm->snes, NULL, pdipm->X));

    814:     /* Check SNES convergence */
    815:     PetscCall(SNESGetConvergedReason(pdipm->snes, &reason));
    816:     if (reason < 0) PetscCall(PetscPrintf(PetscObjectComm((PetscObject)pdipm->snes), "SNES solve did not converged due to reason %s\n", SNESConvergedReasons[reason]));

    818:     /* Check TAO convergence */
    819:     PetscCheck(!PetscIsInfOrNanReal(pdipm->obj), PETSC_COMM_SELF, PETSC_ERR_SUP, "User-provided compute function generated Inf or NaN");
    820:   }
    821:   PetscFunctionReturn(PETSC_SUCCESS);
    822: }

    824: static PetscErrorCode TaoView_PDIPM(Tao tao, PetscViewer viewer)
    825: {
    826:   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;

    828:   PetscFunctionBegin;
    829:   tao->constrained = PETSC_TRUE;
    830:   PetscCall(PetscViewerASCIIPushTab(viewer));
    831:   PetscCall(PetscViewerASCIIPrintf(viewer, "Number of prime=%" PetscInt_FMT ", Number of dual=%" PetscInt_FMT "\n", pdipm->Nx + pdipm->Nci, pdipm->Nce + pdipm->Nci));
    832:   if (pdipm->kkt_pd) PetscCall(PetscViewerASCIIPrintf(viewer, "KKT shifts deltaw=%g, deltac=%g\n", (double)pdipm->deltaw, (double)pdipm->deltac));
    833:   PetscCall(PetscViewerASCIIPopTab(viewer));
    834:   PetscFunctionReturn(PETSC_SUCCESS);
    835: }

    837: static PetscErrorCode TaoSetup_PDIPM(Tao tao)
    838: {
    839:   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
    840:   MPI_Comm           comm;
    841:   PetscMPIInt        size;
    842:   PetscInt           row, col, Jcrstart, Jcrend, k, tmp, nc, proc, *nh_all, *ng_all;
    843:   PetscInt           offset, *xa, *xb, i, j, rstart, rend;
    844:   PetscScalar        one = 1.0, neg_one = -1.0;
    845:   const PetscInt    *cols, *rranges, *cranges, *aj, *ranges;
    846:   const PetscScalar *aa, *Xarr;
    847:   Mat                J;
    848:   Mat                Jce_xfixed_trans, Jci_xb_trans;
    849:   PetscInt          *dnz, *onz, rjstart, nx_all, *nce_all, *Jranges, cols1[2];

    851:   PetscFunctionBegin;
    852:   PetscCall(PetscObjectGetComm((PetscObject)tao, &comm));
    853:   PetscCallMPI(MPI_Comm_size(comm, &size));

    855:   /* (1) Setup Bounds and create Tao vectors */
    856:   PetscCall(TaoPDIPMSetUpBounds(tao));

    858:   if (!tao->gradient) {
    859:     PetscCall(VecDuplicate(tao->solution, &tao->gradient));
    860:     PetscCall(VecDuplicate(tao->solution, &tao->stepdirection));
    861:   }

    863:   /* (2) Get sizes */
    864:   /* Size of vector x - This is set by TaoSetSolution */
    865:   PetscCall(VecGetSize(tao->solution, &pdipm->Nx));
    866:   PetscCall(VecGetLocalSize(tao->solution, &pdipm->nx));

    868:   /* Size of equality constraints and vectors */
    869:   if (tao->constraints_equality) {
    870:     PetscCall(VecGetSize(tao->constraints_equality, &pdipm->Ng));
    871:     PetscCall(VecGetLocalSize(tao->constraints_equality, &pdipm->ng));
    872:   } else {
    873:     pdipm->ng = pdipm->Ng = 0;
    874:   }

    876:   pdipm->nce = pdipm->ng + pdipm->nxfixed;
    877:   pdipm->Nce = pdipm->Ng + pdipm->Nxfixed;

    879:   /* Size of inequality constraints and vectors */
    880:   if (tao->constraints_inequality) {
    881:     PetscCall(VecGetSize(tao->constraints_inequality, &pdipm->Nh));
    882:     PetscCall(VecGetLocalSize(tao->constraints_inequality, &pdipm->nh));
    883:   } else {
    884:     pdipm->nh = pdipm->Nh = 0;
    885:   }

    887:   pdipm->nci = pdipm->nh + pdipm->nxlb + pdipm->nxub + 2 * pdipm->nxbox;
    888:   pdipm->Nci = pdipm->Nh + pdipm->Nxlb + pdipm->Nxub + 2 * pdipm->Nxbox;

    890:   /* Full size of the KKT system to be solved */
    891:   pdipm->n = pdipm->nx + pdipm->nce + 2 * pdipm->nci;
    892:   pdipm->N = pdipm->Nx + pdipm->Nce + 2 * pdipm->Nci;

    894:   /* (3) Offsets for subvectors */
    895:   pdipm->off_lambdae = pdipm->nx;
    896:   pdipm->off_lambdai = pdipm->off_lambdae + pdipm->nce;
    897:   pdipm->off_z       = pdipm->off_lambdai + pdipm->nci;

    899:   /* (4) Create vectors and subvectors */
    900:   /* Ce and Ci vectors */
    901:   PetscCall(VecCreate(comm, &pdipm->ce));
    902:   PetscCall(VecSetSizes(pdipm->ce, pdipm->nce, pdipm->Nce));
    903:   PetscCall(VecSetFromOptions(pdipm->ce));

    905:   PetscCall(VecCreate(comm, &pdipm->ci));
    906:   PetscCall(VecSetSizes(pdipm->ci, pdipm->nci, pdipm->Nci));
    907:   PetscCall(VecSetFromOptions(pdipm->ci));

    909:   /* X=[x; lambdae; lambdai; z] for the big KKT system */
    910:   PetscCall(VecCreate(comm, &pdipm->X));
    911:   PetscCall(VecSetSizes(pdipm->X, pdipm->n, pdipm->N));
    912:   PetscCall(VecSetFromOptions(pdipm->X));

    914:   /* Subvectors; they share local arrays with X */
    915:   PetscCall(VecGetArrayRead(pdipm->X, &Xarr));
    916:   /* x shares local array with X.x */
    917:   if (pdipm->Nx) PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nx, pdipm->Nx, Xarr, &pdipm->x));

    919:   /* lambdae shares local array with X.lambdae */
    920:   if (pdipm->Nce) PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nce, pdipm->Nce, Xarr + pdipm->off_lambdae, &pdipm->lambdae));

    922:   /* tao->DE shares local array with X.lambdae_g */
    923:   if (pdipm->Ng) {
    924:     PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->ng, pdipm->Ng, Xarr + pdipm->off_lambdae, &tao->DE));

    926:     PetscCall(VecCreate(comm, &pdipm->lambdae_xfixed));
    927:     PetscCall(VecSetSizes(pdipm->lambdae_xfixed, pdipm->nxfixed, PETSC_DECIDE));
    928:     PetscCall(VecSetFromOptions(pdipm->lambdae_xfixed));
    929:   }

    931:   if (pdipm->Nci) {
    932:     /* lambdai shares local array with X.lambdai */
    933:     PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nci, pdipm->Nci, Xarr + pdipm->off_lambdai, &pdipm->lambdai));

    935:     /* z for slack variables; it shares local array with X.z */
    936:     PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nci, pdipm->Nci, Xarr + pdipm->off_z, &pdipm->z));
    937:   }

    939:   /* tao->DI which shares local array with X.lambdai_h */
    940:   if (pdipm->Nh) PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nh, pdipm->Nh, Xarr + pdipm->off_lambdai, &tao->DI));
    941:   PetscCall(VecCreate(comm, &pdipm->lambdai_xb));
    942:   PetscCall(VecSetSizes(pdipm->lambdai_xb, pdipm->nci - pdipm->nh, PETSC_DECIDE));
    943:   PetscCall(VecSetFromOptions(pdipm->lambdai_xb));

    945:   PetscCall(VecRestoreArrayRead(pdipm->X, &Xarr));

    947:   /* (5) Create Jacobians Jce_xfixed and Jci */
    948:   /* (5.1) PDIPM Jacobian of equality bounds cebound(x) = J_nxfixed */
    949:   if (pdipm->Nxfixed) {
    950:     /* Create Jce_xfixed */
    951:     PetscCall(MatCreate(comm, &pdipm->Jce_xfixed));
    952:     PetscCall(MatSetSizes(pdipm->Jce_xfixed, pdipm->nxfixed, pdipm->nx, PETSC_DECIDE, pdipm->Nx));
    953:     PetscCall(MatSetFromOptions(pdipm->Jce_xfixed));
    954:     PetscCall(MatSeqAIJSetPreallocation(pdipm->Jce_xfixed, 1, NULL));
    955:     PetscCall(MatMPIAIJSetPreallocation(pdipm->Jce_xfixed, 1, NULL, 1, NULL));

    957:     PetscCall(MatGetOwnershipRange(pdipm->Jce_xfixed, &Jcrstart, &Jcrend));
    958:     PetscCall(ISGetIndices(pdipm->isxfixed, &cols));
    959:     k = 0;
    960:     for (row = Jcrstart; row < Jcrend; row++) {
    961:       PetscCall(MatSetValues(pdipm->Jce_xfixed, 1, &row, 1, cols + k, &one, INSERT_VALUES));
    962:       k++;
    963:     }
    964:     PetscCall(ISRestoreIndices(pdipm->isxfixed, &cols));
    965:     PetscCall(MatAssemblyBegin(pdipm->Jce_xfixed, MAT_FINAL_ASSEMBLY));
    966:     PetscCall(MatAssemblyEnd(pdipm->Jce_xfixed, MAT_FINAL_ASSEMBLY));
    967:   }

    969:   /* (5.2) PDIPM inequality Jacobian Jci = [tao->jacobian_inequality; ...] */
    970:   PetscCall(MatCreate(comm, &pdipm->Jci_xb));
    971:   PetscCall(MatSetSizes(pdipm->Jci_xb, pdipm->nci - pdipm->nh, pdipm->nx, PETSC_DECIDE, pdipm->Nx));
    972:   PetscCall(MatSetFromOptions(pdipm->Jci_xb));
    973:   PetscCall(MatSeqAIJSetPreallocation(pdipm->Jci_xb, 1, NULL));
    974:   PetscCall(MatMPIAIJSetPreallocation(pdipm->Jci_xb, 1, NULL, 1, NULL));

    976:   PetscCall(MatGetOwnershipRange(pdipm->Jci_xb, &Jcrstart, &Jcrend));
    977:   offset = Jcrstart;
    978:   if (pdipm->Nxub) {
    979:     /* Add xub to Jci_xb */
    980:     PetscCall(ISGetIndices(pdipm->isxub, &cols));
    981:     k = 0;
    982:     for (row = offset; row < offset + pdipm->nxub; row++) {
    983:       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &row, 1, cols + k, &neg_one, INSERT_VALUES));
    984:       k++;
    985:     }
    986:     PetscCall(ISRestoreIndices(pdipm->isxub, &cols));
    987:   }

    989:   if (pdipm->Nxlb) {
    990:     /* Add xlb to Jci_xb */
    991:     PetscCall(ISGetIndices(pdipm->isxlb, &cols));
    992:     k = 0;
    993:     offset += pdipm->nxub;
    994:     for (row = offset; row < offset + pdipm->nxlb; row++) {
    995:       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &row, 1, cols + k, &one, INSERT_VALUES));
    996:       k++;
    997:     }
    998:     PetscCall(ISRestoreIndices(pdipm->isxlb, &cols));
    999:   }

    1001:   /* Add xbox to Jci_xb */
    1002:   if (pdipm->Nxbox) {
    1003:     PetscCall(ISGetIndices(pdipm->isxbox, &cols));
    1004:     k = 0;
    1005:     offset += pdipm->nxlb;
    1006:     for (row = offset; row < offset + pdipm->nxbox; row++) {
    1007:       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &row, 1, cols + k, &neg_one, INSERT_VALUES));
    1008:       tmp = row + pdipm->nxbox;
    1009:       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &tmp, 1, cols + k, &one, INSERT_VALUES));
    1010:       k++;
    1011:     }
    1012:     PetscCall(ISRestoreIndices(pdipm->isxbox, &cols));
    1013:   }

    1015:   PetscCall(MatAssemblyBegin(pdipm->Jci_xb, MAT_FINAL_ASSEMBLY));
    1016:   PetscCall(MatAssemblyEnd(pdipm->Jci_xb, MAT_FINAL_ASSEMBLY));
    1017:   /* PetscCall(MatView(pdipm->Jci_xb,PETSC_VIEWER_STDOUT_WORLD)); */

    1019:   /* (6) Set up ISs for PC Fieldsplit */
    1020:   if (pdipm->solve_reduced_kkt) {
    1021:     PetscCall(PetscMalloc2(pdipm->nx + pdipm->nce, &xa, 2 * pdipm->nci, &xb));
    1022:     for (i = 0; i < pdipm->nx + pdipm->nce; i++) xa[i] = i;
    1023:     for (i = 0; i < 2 * pdipm->nci; i++) xb[i] = pdipm->off_lambdai + i;

    1025:     PetscCall(ISCreateGeneral(comm, pdipm->nx + pdipm->nce, xa, PETSC_OWN_POINTER, &pdipm->is1));
    1026:     PetscCall(ISCreateGeneral(comm, 2 * pdipm->nci, xb, PETSC_OWN_POINTER, &pdipm->is2));
    1027:   }

    1029:   /* (7) Gather offsets from all processes */
    1030:   PetscCall(PetscMalloc1(size, &pdipm->nce_all));

    1032:   /* Get rstart of KKT matrix */
    1033:   PetscCallMPI(MPI_Scan(&pdipm->n, &rstart, 1, MPIU_INT, MPI_SUM, comm));
    1034:   rstart -= pdipm->n;

    1036:   PetscCallMPI(MPI_Allgather(&pdipm->nce, 1, MPIU_INT, pdipm->nce_all, 1, MPIU_INT, comm));

    1038:   PetscCall(PetscMalloc3(size, &ng_all, size, &nh_all, size, &Jranges));
    1039:   PetscCallMPI(MPI_Allgather(&rstart, 1, MPIU_INT, Jranges, 1, MPIU_INT, comm));
    1040:   PetscCallMPI(MPI_Allgather(&pdipm->nh, 1, MPIU_INT, nh_all, 1, MPIU_INT, comm));
    1041:   PetscCallMPI(MPI_Allgather(&pdipm->ng, 1, MPIU_INT, ng_all, 1, MPIU_INT, comm));

    1043:   PetscCall(MatGetOwnershipRanges(tao->hessian, &rranges));
    1044:   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &cranges));

    1046:   if (pdipm->Ng) {
    1047:     PetscCall(TaoComputeJacobianEquality(tao, tao->solution, tao->jacobian_equality, tao->jacobian_equality_pre));
    1048:     PetscCall(MatTranspose(tao->jacobian_equality, MAT_INITIAL_MATRIX, &pdipm->jac_equality_trans));
    1049:   }
    1050:   if (pdipm->Nh) {
    1051:     PetscCall(TaoComputeJacobianInequality(tao, tao->solution, tao->jacobian_inequality, tao->jacobian_inequality_pre));
    1052:     PetscCall(MatTranspose(tao->jacobian_inequality, MAT_INITIAL_MATRIX, &pdipm->jac_inequality_trans));
    1053:   }

    1055:   /* Count dnz,onz for preallocation of KKT matrix */
    1056:   nce_all = pdipm->nce_all;

    1058:   if (pdipm->Nxfixed) PetscCall(MatTranspose(pdipm->Jce_xfixed, MAT_INITIAL_MATRIX, &Jce_xfixed_trans));
    1059:   PetscCall(MatTranspose(pdipm->Jci_xb, MAT_INITIAL_MATRIX, &Jci_xb_trans));

    1061:   MatPreallocateBegin(comm, pdipm->n, pdipm->n, dnz, onz);

    1063:   /* 1st row block of KKT matrix: [Wxx; gradCe'; -gradCi'; 0] */
    1064:   PetscCall(TaoPDIPMEvaluateFunctionsAndJacobians(tao, pdipm->x));
    1065:   PetscCall(TaoComputeHessian(tao, tao->solution, tao->hessian, tao->hessian_pre));

    1067:   /* Insert tao->hessian */
    1068:   PetscCall(MatGetOwnershipRange(tao->hessian, &rjstart, NULL));
    1069:   for (i = 0; i < pdipm->nx; i++) {
    1070:     row = rstart + i;

    1072:     PetscCall(MatGetRow(tao->hessian, i + rjstart, &nc, &aj, NULL));
    1073:     proc = 0;
    1074:     for (j = 0; j < nc; j++) {
    1075:       while (aj[j] >= cranges[proc + 1]) proc++;
    1076:       col = aj[j] - cranges[proc] + Jranges[proc];
    1077:       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1078:     }
    1079:     PetscCall(MatRestoreRow(tao->hessian, i + rjstart, &nc, &aj, NULL));

    1081:     if (pdipm->ng) {
    1082:       /* Insert grad g' */
    1083:       PetscCall(MatGetRow(pdipm->jac_equality_trans, i + rjstart, &nc, &aj, NULL));
    1084:       PetscCall(MatGetOwnershipRanges(tao->jacobian_equality, &ranges));
    1085:       proc = 0;
    1086:       for (j = 0; j < nc; j++) {
    1087:         /* find row ownership of */
    1088:         while (aj[j] >= ranges[proc + 1]) proc++;
    1089:         nx_all = rranges[proc + 1] - rranges[proc];
    1090:         col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all;
    1091:         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1092:       }
    1093:       PetscCall(MatRestoreRow(pdipm->jac_equality_trans, i + rjstart, &nc, &aj, NULL));
    1094:     }

    1096:     /* Insert Jce_xfixed^T' */
    1097:     if (pdipm->nxfixed) {
    1098:       PetscCall(MatGetRow(Jce_xfixed_trans, i + rjstart, &nc, &aj, NULL));
    1099:       PetscCall(MatGetOwnershipRanges(pdipm->Jce_xfixed, &ranges));
    1100:       proc = 0;
    1101:       for (j = 0; j < nc; j++) {
    1102:         /* find row ownership of */
    1103:         while (aj[j] >= ranges[proc + 1]) proc++;
    1104:         nx_all = rranges[proc + 1] - rranges[proc];
    1105:         col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all + ng_all[proc];
    1106:         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1107:       }
    1108:       PetscCall(MatRestoreRow(Jce_xfixed_trans, i + rjstart, &nc, &aj, NULL));
    1109:     }

    1111:     if (pdipm->nh) {
    1112:       /* Insert -grad h' */
    1113:       PetscCall(MatGetRow(pdipm->jac_inequality_trans, i + rjstart, &nc, &aj, NULL));
    1114:       PetscCall(MatGetOwnershipRanges(tao->jacobian_inequality, &ranges));
    1115:       proc = 0;
    1116:       for (j = 0; j < nc; j++) {
    1117:         /* find row ownership of */
    1118:         while (aj[j] >= ranges[proc + 1]) proc++;
    1119:         nx_all = rranges[proc + 1] - rranges[proc];
    1120:         col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all + nce_all[proc];
    1121:         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1122:       }
    1123:       PetscCall(MatRestoreRow(pdipm->jac_inequality_trans, i + rjstart, &nc, &aj, NULL));
    1124:     }

    1126:     /* Insert Jci_xb^T' */
    1127:     PetscCall(MatGetRow(Jci_xb_trans, i + rjstart, &nc, &aj, NULL));
    1128:     PetscCall(MatGetOwnershipRanges(pdipm->Jci_xb, &ranges));
    1129:     proc = 0;
    1130:     for (j = 0; j < nc; j++) {
    1131:       /* find row ownership of */
    1132:       while (aj[j] >= ranges[proc + 1]) proc++;
    1133:       nx_all = rranges[proc + 1] - rranges[proc];
    1134:       col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all + nce_all[proc] + nh_all[proc];
    1135:       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1136:     }
    1137:     PetscCall(MatRestoreRow(Jci_xb_trans, i + rjstart, &nc, &aj, NULL));
    1138:   }

    1140:   /* 2nd Row block of KKT matrix: [grad Ce, deltac*I, 0, 0] */
    1141:   if (pdipm->Ng) {
    1142:     PetscCall(MatGetOwnershipRange(tao->jacobian_equality, &rjstart, NULL));
    1143:     for (i = 0; i < pdipm->ng; i++) {
    1144:       row = rstart + pdipm->off_lambdae + i;

    1146:       PetscCall(MatGetRow(tao->jacobian_equality, i + rjstart, &nc, &aj, NULL));
    1147:       proc = 0;
    1148:       for (j = 0; j < nc; j++) {
    1149:         while (aj[j] >= cranges[proc + 1]) proc++;
    1150:         col = aj[j] - cranges[proc] + Jranges[proc];
    1151:         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz)); /* grad g */
    1152:       }
    1153:       PetscCall(MatRestoreRow(tao->jacobian_equality, i + rjstart, &nc, &aj, NULL));
    1154:     }
    1155:   }
    1156:   /* Jce_xfixed */
    1157:   if (pdipm->Nxfixed) {
    1158:     PetscCall(MatGetOwnershipRange(pdipm->Jce_xfixed, &Jcrstart, NULL));
    1159:     for (i = 0; i < (pdipm->nce - pdipm->ng); i++) {
    1160:       row = rstart + pdipm->off_lambdae + pdipm->ng + i;

    1162:       PetscCall(MatGetRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, NULL));
    1163:       PetscCheck(nc == 1, PETSC_COMM_SELF, PETSC_ERR_SUP, "nc != 1");

    1165:       proc = 0;
    1166:       j    = 0;
    1167:       while (cols[j] >= cranges[proc + 1]) proc++;
    1168:       col = cols[j] - cranges[proc] + Jranges[proc];
    1169:       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1170:       PetscCall(MatRestoreRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, NULL));
    1171:     }
    1172:   }

    1174:   /* 3rd Row block of KKT matrix: [ gradCi, 0, deltac*I, -I] */
    1175:   if (pdipm->Nh) {
    1176:     PetscCall(MatGetOwnershipRange(tao->jacobian_inequality, &rjstart, NULL));
    1177:     for (i = 0; i < pdipm->nh; i++) {
    1178:       row = rstart + pdipm->off_lambdai + i;

    1180:       PetscCall(MatGetRow(tao->jacobian_inequality, i + rjstart, &nc, &aj, NULL));
    1181:       proc = 0;
    1182:       for (j = 0; j < nc; j++) {
    1183:         while (aj[j] >= cranges[proc + 1]) proc++;
    1184:         col = aj[j] - cranges[proc] + Jranges[proc];
    1185:         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz)); /* grad h */
    1186:       }
    1187:       PetscCall(MatRestoreRow(tao->jacobian_inequality, i + rjstart, &nc, &aj, NULL));
    1188:     }
    1189:     /* I */
    1190:     for (i = 0; i < pdipm->nh; i++) {
    1191:       row = rstart + pdipm->off_lambdai + i;
    1192:       col = rstart + pdipm->off_z + i;
    1193:       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1194:     }
    1195:   }

    1197:   /* Jci_xb */
    1198:   PetscCall(MatGetOwnershipRange(pdipm->Jci_xb, &Jcrstart, NULL));
    1199:   for (i = 0; i < (pdipm->nci - pdipm->nh); i++) {
    1200:     row = rstart + pdipm->off_lambdai + pdipm->nh + i;

    1202:     PetscCall(MatGetRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, NULL));
    1203:     PetscCheck(nc == 1, PETSC_COMM_SELF, PETSC_ERR_SUP, "nc != 1");
    1204:     proc = 0;
    1205:     for (j = 0; j < nc; j++) {
    1206:       while (cols[j] >= cranges[proc + 1]) proc++;
    1207:       col = cols[j] - cranges[proc] + Jranges[proc];
    1208:       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1209:     }
    1210:     PetscCall(MatRestoreRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, NULL));
    1211:     /* I */
    1212:     col = rstart + pdipm->off_z + pdipm->nh + i;
    1213:     PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1214:   }

    1216:   /* 4-th Row block of KKT matrix: Z and Ci */
    1217:   for (i = 0; i < pdipm->nci; i++) {
    1218:     row      = rstart + pdipm->off_z + i;
    1219:     cols1[0] = rstart + pdipm->off_lambdai + i;
    1220:     cols1[1] = row;
    1221:     PetscCall(MatPreallocateSet(row, 2, cols1, dnz, onz));
    1222:   }

    1224:   /* diagonal entry */
    1225:   for (i = 0; i < pdipm->n; i++) dnz[i]++; /* diagonal entry */

    1227:   /* Create KKT matrix */
    1228:   PetscCall(MatCreate(comm, &J));
    1229:   PetscCall(MatSetSizes(J, pdipm->n, pdipm->n, PETSC_DECIDE, PETSC_DECIDE));
    1230:   PetscCall(MatSetFromOptions(J));
    1231:   PetscCall(MatSeqAIJSetPreallocation(J, 0, dnz));
    1232:   PetscCall(MatMPIAIJSetPreallocation(J, 0, dnz, 0, onz));
    1233:   MatPreallocateEnd(dnz, onz);
    1234:   pdipm->K = J;

    1236:   /* (8) Insert constant entries to  K */
    1237:   /* Set 0.0 to diagonal of K, so that the solver does not complain *about missing diagonal value */
    1238:   PetscCall(MatGetOwnershipRange(J, &rstart, &rend));
    1239:   for (i = rstart; i < rend; i++) PetscCall(MatSetValue(J, i, i, 0.0, INSERT_VALUES));
    1240:   /* In case Wxx has no diagonal entries preset set diagonal to deltaw given */
    1241:   if (pdipm->kkt_pd) {
    1242:     for (i = 0; i < pdipm->nh; i++) {
    1243:       row = rstart + i;
    1244:       PetscCall(MatSetValue(J, row, row, pdipm->deltaw, INSERT_VALUES));
    1245:     }
    1246:   }

    1248:   /* Row block of K: [ grad Ce, 0, 0, 0] */
    1249:   if (pdipm->Nxfixed) {
    1250:     PetscCall(MatGetOwnershipRange(pdipm->Jce_xfixed, &Jcrstart, NULL));
    1251:     for (i = 0; i < (pdipm->nce - pdipm->ng); i++) {
    1252:       row = rstart + pdipm->off_lambdae + pdipm->ng + i;

    1254:       PetscCall(MatGetRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, &aa));
    1255:       proc = 0;
    1256:       for (j = 0; j < nc; j++) {
    1257:         while (cols[j] >= cranges[proc + 1]) proc++;
    1258:         col = cols[j] - cranges[proc] + Jranges[proc];
    1259:         PetscCall(MatSetValue(J, row, col, aa[j], INSERT_VALUES)); /* grad Ce */
    1260:         PetscCall(MatSetValue(J, col, row, aa[j], INSERT_VALUES)); /* grad Ce' */
    1261:       }
    1262:       PetscCall(MatRestoreRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, &aa));
    1263:     }
    1264:   }

    1266:   /* Row block of K: [ -grad Ci, 0, 0, I] */
    1267:   PetscCall(MatGetOwnershipRange(pdipm->Jci_xb, &Jcrstart, NULL));
    1268:   for (i = 0; i < pdipm->nci - pdipm->nh; i++) {
    1269:     row = rstart + pdipm->off_lambdai + pdipm->nh + i;

    1271:     PetscCall(MatGetRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, &aa));
    1272:     proc = 0;
    1273:     for (j = 0; j < nc; j++) {
    1274:       while (cols[j] >= cranges[proc + 1]) proc++;
    1275:       col = cols[j] - cranges[proc] + Jranges[proc];
    1276:       PetscCall(MatSetValue(J, col, row, -aa[j], INSERT_VALUES));
    1277:       PetscCall(MatSetValue(J, row, col, -aa[j], INSERT_VALUES));
    1278:     }
    1279:     PetscCall(MatRestoreRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, &aa));

    1281:     col = rstart + pdipm->off_z + pdipm->nh + i;
    1282:     PetscCall(MatSetValue(J, row, col, 1, INSERT_VALUES));
    1283:   }

    1285:   for (i = 0; i < pdipm->nh; i++) {
    1286:     row = rstart + pdipm->off_lambdai + i;
    1287:     col = rstart + pdipm->off_z + i;
    1288:     PetscCall(MatSetValue(J, row, col, 1, INSERT_VALUES));
    1289:   }

    1291:   /* Row block of K: [ 0, 0, I, ...] */
    1292:   for (i = 0; i < pdipm->nci; i++) {
    1293:     row = rstart + pdipm->off_z + i;
    1294:     col = rstart + pdipm->off_lambdai + i;
    1295:     PetscCall(MatSetValue(J, row, col, 1, INSERT_VALUES));
    1296:   }

    1298:   if (pdipm->Nxfixed) PetscCall(MatDestroy(&Jce_xfixed_trans));
    1299:   PetscCall(MatDestroy(&Jci_xb_trans));
    1300:   PetscCall(PetscFree3(ng_all, nh_all, Jranges));

    1302:   /* (9) Set up nonlinear solver SNES */
    1303:   PetscCall(SNESSetFunction(pdipm->snes, NULL, TaoSNESFunction_PDIPM, (void *)tao));
    1304:   PetscCall(SNESSetJacobian(pdipm->snes, J, J, TaoSNESJacobian_PDIPM, (void *)tao));

    1306:   if (pdipm->solve_reduced_kkt) {
    1307:     PC pc;
    1308:     PetscCall(KSPGetPC(tao->ksp, &pc));
    1309:     PetscCall(PCSetType(pc, PCFIELDSPLIT));
    1310:     PetscCall(PCFieldSplitSetType(pc, PC_COMPOSITE_SCHUR));
    1311:     PetscCall(PCFieldSplitSetIS(pc, "2", pdipm->is2));
    1312:     PetscCall(PCFieldSplitSetIS(pc, "1", pdipm->is1));
    1313:   }
    1314:   PetscCall(SNESSetFromOptions(pdipm->snes));

    1316:   /* (10) Setup PCPostSetUp() for pdipm->solve_symmetric_kkt */
    1317:   if (pdipm->solve_symmetric_kkt) {
    1318:     KSP       ksp;
    1319:     PC        pc;
    1320:     PetscBool isCHOL;

    1322:     PetscCall(SNESGetKSP(pdipm->snes, &ksp));
    1323:     PetscCall(KSPGetPC(ksp, &pc));
    1324:     PetscCall(PCSetPostSetUp(pc, PCPostSetUp_PDIPM));

    1326:     PetscCall(PetscObjectTypeCompare((PetscObject)pc, PCCHOLESKY, &isCHOL));
    1327:     if (isCHOL) {
    1328:       Mat Factor;

    1330:       PetscCheck(PetscDefined(HAVE_MUMPS), PetscObjectComm((PetscObject)tao), PETSC_ERR_SUP, "Requires external package MUMPS");
    1331:       PetscCall(PCFactorGetMatrix(pc, &Factor));
    1332:       PetscCall(MatMumpsSetIcntl(Factor, 24, 1));               /* detection of null pivot rows */
    1333:       if (size > 1) PetscCall(MatMumpsSetIcntl(Factor, 13, 1)); /* parallelism of the root node (enable ScaLAPACK) and its splitting */
    1334:     }
    1335:   }
    1336:   PetscFunctionReturn(PETSC_SUCCESS);
    1337: }

    1339: static PetscErrorCode TaoDestroy_PDIPM(Tao tao)
    1340: {
    1341:   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;

    1343:   PetscFunctionBegin;
    1344:   /* Freeing Vectors assocaiated with KKT (X) */
    1345:   PetscCall(VecDestroy(&pdipm->x));       /* Solution x */
    1346:   PetscCall(VecDestroy(&pdipm->lambdae)); /* Equality constraints lagrangian multiplier*/
    1347:   PetscCall(VecDestroy(&pdipm->lambdai)); /* Inequality constraints lagrangian multiplier*/
    1348:   PetscCall(VecDestroy(&pdipm->z));       /* Slack variables */
    1349:   PetscCall(VecDestroy(&pdipm->X));       /* Big KKT system vector [x; lambdae; lambdai; z] */

    1351:   /* work vectors */
    1352:   PetscCall(VecDestroy(&pdipm->lambdae_xfixed));
    1353:   PetscCall(VecDestroy(&pdipm->lambdai_xb));

    1355:   /* Legrangian equality and inequality Vec */
    1356:   PetscCall(VecDestroy(&pdipm->ce)); /* Vec of equality constraints */
    1357:   PetscCall(VecDestroy(&pdipm->ci)); /* Vec of inequality constraints */

    1359:   /* Matrices */
    1360:   PetscCall(MatDestroy(&pdipm->Jce_xfixed));
    1361:   PetscCall(MatDestroy(&pdipm->Jci_xb)); /* Jacobian of inequality constraints Jci = [tao->jacobian_inequality ; J(nxub); J(nxlb); J(nxbx)] */
    1362:   PetscCall(MatDestroy(&pdipm->K));

    1364:   /* Index Sets */
    1365:   if (pdipm->Nxub) PetscCall(ISDestroy(&pdipm->isxub)); /* Finite upper bound only -inf < x < ub */

    1367:   if (pdipm->Nxlb) PetscCall(ISDestroy(&pdipm->isxlb)); /* Finite lower bound only  lb <= x < inf */

    1369:   if (pdipm->Nxfixed) PetscCall(ISDestroy(&pdipm->isxfixed)); /* Fixed variables         lb =  x = ub */

    1371:   if (pdipm->Nxbox) PetscCall(ISDestroy(&pdipm->isxbox)); /* Boxed variables         lb <= x <= ub */

    1373:   if (pdipm->Nxfree) PetscCall(ISDestroy(&pdipm->isxfree)); /* Free variables        -inf <= x <= inf */

    1375:   if (pdipm->solve_reduced_kkt) {
    1376:     PetscCall(ISDestroy(&pdipm->is1));
    1377:     PetscCall(ISDestroy(&pdipm->is2));
    1378:   }

    1380:   /* SNES */
    1381:   PetscCall(SNESDestroy(&pdipm->snes)); /* Nonlinear solver */
    1382:   PetscCall(PetscFree(pdipm->nce_all));
    1383:   PetscCall(MatDestroy(&pdipm->jac_equality_trans));
    1384:   PetscCall(MatDestroy(&pdipm->jac_inequality_trans));

    1386:   /* Destroy pdipm */
    1387:   PetscCall(PetscFree(tao->data)); /* Holding locations of pdipm */

    1389:   /* Destroy Dual */
    1390:   PetscCall(VecDestroy(&tao->DE)); /* equality dual */
    1391:   PetscCall(VecDestroy(&tao->DI)); /* dinequality dual */
    1392:   PetscFunctionReturn(PETSC_SUCCESS);
    1393: }

    1395: static PetscErrorCode TaoSetFromOptions_PDIPM(Tao tao, PetscOptionItems PetscOptionsObject)
    1396: {
    1397:   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;

    1399:   PetscFunctionBegin;
    1400:   PetscOptionsHeadBegin(PetscOptionsObject, "PDIPM method for constrained optimization");
    1401:   PetscCall(PetscOptionsReal("-tao_pdipm_push_init_slack", "parameter to push initial slack variables away from bounds", NULL, pdipm->push_init_slack, &pdipm->push_init_slack, NULL));
    1402:   PetscCall(PetscOptionsReal("-tao_pdipm_push_init_lambdai", "parameter to push initial (inequality) dual variables away from bounds", NULL, pdipm->push_init_lambdai, &pdipm->push_init_lambdai, NULL));
    1403:   PetscCall(PetscOptionsBool("-tao_pdipm_solve_reduced_kkt", "Solve reduced KKT system using Schur-complement", NULL, pdipm->solve_reduced_kkt, &pdipm->solve_reduced_kkt, NULL));
    1404:   PetscCall(PetscOptionsReal("-tao_pdipm_mu_update_factor", "Update scalar for barrier parameter (mu) update", NULL, pdipm->mu_update_factor, &pdipm->mu_update_factor, NULL));
    1405:   PetscCall(PetscOptionsBool("-tao_pdipm_symmetric_kkt", "Solve non reduced symmetric KKT system", NULL, pdipm->solve_symmetric_kkt, &pdipm->solve_symmetric_kkt, NULL));
    1406:   PetscCall(PetscOptionsBool("-tao_pdipm_kkt_shift_pd", "Add shifts to make KKT matrix positive definite", NULL, pdipm->kkt_pd, &pdipm->kkt_pd, NULL));
    1407:   PetscOptionsHeadEnd();
    1408:   PetscFunctionReturn(PETSC_SUCCESS);
    1409: }

    1411: /*MC
    1412:   TAOPDIPM - Barrier-based primal-dual interior point algorithm for generally constrained optimization.

    1414:   Options Database Keys:
    1415: +   -tao_pdipm_push_init_lambdai - parameter to push initial dual variables away from bounds (> 0)
    1416: .   -tao_pdipm_push_init_slack   - parameter to push initial slack variables away from bounds (> 0)
    1417: .   -tao_pdipm_mu_update_factor  - update scalar for barrier parameter (mu) update (> 0)
    1418: .   -tao_pdipm_symmetric_kkt     - Solve non-reduced symmetric KKT system
    1419: -   -tao_pdipm_kkt_shift_pd      - Add shifts to make KKT matrix positive definite

    1421:   Level: beginner

    1423: .seealso: `TAOPDIPM`, `Tao`, `TaoType`
    1424: M*/

    1426: PETSC_EXTERN PetscErrorCode TaoCreate_PDIPM(Tao tao)
    1427: {
    1428:   TAO_PDIPM *pdipm;
    1429:   PC         pc;

    1431:   PetscFunctionBegin;
    1432:   tao->ops->setup          = TaoSetup_PDIPM;
    1433:   tao->ops->solve          = TaoSolve_PDIPM;
    1434:   tao->ops->setfromoptions = TaoSetFromOptions_PDIPM;
    1435:   tao->ops->view           = TaoView_PDIPM;
    1436:   tao->ops->destroy        = TaoDestroy_PDIPM;

    1438:   PetscCall(PetscNew(&pdipm));
    1439:   tao->data = (void *)pdipm;

    1441:   pdipm->nx = pdipm->Nx = 0;
    1442:   pdipm->nxfixed = pdipm->Nxfixed = 0;
    1443:   pdipm->nxlb = pdipm->Nxlb = 0;
    1444:   pdipm->nxub = pdipm->Nxub = 0;
    1445:   pdipm->nxbox = pdipm->Nxbox = 0;
    1446:   pdipm->nxfree = pdipm->Nxfree = 0;

    1448:   pdipm->ng = pdipm->Ng = pdipm->nce = pdipm->Nce = 0;
    1449:   pdipm->nh = pdipm->Nh = pdipm->nci = pdipm->Nci = 0;
    1450:   pdipm->n = pdipm->N     = 0;
    1451:   pdipm->mu               = 1.0;
    1452:   pdipm->mu_update_factor = 0.1;

    1454:   pdipm->deltaw     = 0.0;
    1455:   pdipm->lastdeltaw = 3 * 1.e-4;
    1456:   pdipm->deltac     = 0.0;
    1457:   pdipm->kkt_pd     = PETSC_FALSE;

    1459:   pdipm->push_init_slack     = 1.0;
    1460:   pdipm->push_init_lambdai   = 1.0;
    1461:   pdipm->solve_reduced_kkt   = PETSC_FALSE;
    1462:   pdipm->solve_symmetric_kkt = PETSC_TRUE;

    1464:   /* Override default settings (unless already changed) */
    1465:   PetscCall(TaoParametersInitialize(tao));
    1466:   PetscObjectParameterSetDefault(tao, max_it, 200);
    1467:   PetscObjectParameterSetDefault(tao, max_funcs, 500);

    1469:   PetscCall(SNESCreate(((PetscObject)tao)->comm, &pdipm->snes));
    1470:   PetscCall(SNESSetOptionsPrefix(pdipm->snes, tao->hdr.prefix));
    1471:   PetscCall(SNESGetKSP(pdipm->snes, &tao->ksp));
    1472:   PetscCall(PetscObjectReference((PetscObject)tao->ksp));
    1473:   PetscCall(KSPGetPC(tao->ksp, &pc));
    1474:   PetscCall(PCSetApplicationContext(pc, (void *)tao));
    1475:   PetscFunctionReturn(PETSC_SUCCESS);
    1476: }