Date: 2025-03-29T05:58:47+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;
    283:   MPI_Comm           comm;
    284:   PetscMPIInt        rank, size;

    286:   PetscFunctionBegin;
    287:   PetscCall(PetscObjectGetComm((PetscObject)snes, &comm));
    288:   PetscCallMPI(MPI_Comm_rank(comm, &rank));
    289:   PetscCallMPI(MPI_Comm_rank(comm, &size));

    291:   PetscCall(MatGetOwnershipRanges(Jpre, &Jranges));
    292:   PetscCall(MatGetOwnershipRange(Jpre, &Jrstart, NULL));
    293:   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &rranges));
    294:   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &cranges));

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

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

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

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

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

    360:   /* (4) insert 1st row block of Jpre: [Wxx, grad g', -grad h', 0] */
    361:   if (pdipm->Ng) { /* grad g' */
    362:     PetscCall(MatTranspose(tao->jacobian_equality, MAT_REUSE_MATRIX, &pdipm->jac_equality_trans));
    363:   }
    364:   if (pdipm->Nh) { /* grad h' */
    365:     PetscCall(MatTranspose(tao->jacobian_inequality, MAT_REUSE_MATRIX, &pdipm->jac_inequality_trans));
    366:   }

    368:   PetscCall(VecPlaceArray(pdipm->x, Xarr));
    369:   PetscCall(TaoComputeHessian(tao, pdipm->x, tao->hessian, tao->hessian_pre));
    370:   PetscCall(VecResetArray(pdipm->x));

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

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

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

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

    423:   /* (6) assemble Jpre and J */
    424:   PetscCall(MatAssemblyBegin(Jpre, MAT_FINAL_ASSEMBLY));
    425:   PetscCall(MatAssemblyEnd(Jpre, MAT_FINAL_ASSEMBLY));

    427:   if (J != Jpre) {
    428:     PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
    429:     PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
    430:   }
    431:   PetscFunctionReturn(PETSC_SUCCESS);
    432: }

    434: /*
    435:    TaoSnesFunction_PDIPM - Evaluate KKT function at X

    437:    Input Parameter:
    438:    snes - SNES context
    439:    X - KKT Vector
    440:    *ctx - pdipm

    442:    Output Parameter:
    443:    F - Updated Lagrangian vector
    444: */
    445: static PetscErrorCode TaoSNESFunction_PDIPM(SNES snes, Vec X, Vec F, void *ctx)
    446: {
    447:   Tao                tao   = (Tao)ctx;
    448:   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
    449:   PetscScalar       *Farr;
    450:   Vec                x, L1;
    451:   PetscInt           i;
    452:   const PetscScalar *Xarr, *carr, *zarr, *larr;

    454:   PetscFunctionBegin;
    455:   PetscCall(VecSet(F, 0.0));

    457:   PetscCall(VecGetArrayRead(X, &Xarr));
    458:   PetscCall(VecGetArrayWrite(F, &Farr));

    460:   /* (0) Evaluate f, fx, gradG, gradH at X.x Note: pdipm->x is not changed below */
    461:   x = pdipm->x;
    462:   PetscCall(VecPlaceArray(x, Xarr));
    463:   PetscCall(TaoPDIPMEvaluateFunctionsAndJacobians(tao, x));

    465:   /* Update ce, ci, and Jci at X.x */
    466:   PetscCall(TaoPDIPMUpdateConstraints(tao, x));
    467:   PetscCall(VecResetArray(x));

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

    480:     /* L1 += Jci_xb'*lambdai_xb */
    481:     PetscCall(VecPlaceArray(pdipm->lambdai_xb, Xarr + pdipm->off_lambdai + pdipm->nh));
    482:     PetscCall(MatMultTransposeAdd(pdipm->Jci_xb, pdipm->lambdai_xb, L1, L1));
    483:     PetscCall(VecResetArray(pdipm->lambdai_xb));

    485:     /* L1 = - (gradH'*DI + Jci_xb'*lambdai_xb) */
    486:     PetscCall(VecScale(L1, -1.0));
    487:   }

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

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

    508:   /* (2) L2 = ce(x) */
    509:   if (pdipm->Nce) {
    510:     PetscCall(VecGetArrayRead(pdipm->ce, &carr));
    511:     for (i = 0; i < pdipm->nce; i++) Farr[pdipm->off_lambdae + i] = carr[i];
    512:     PetscCall(VecRestoreArrayRead(pdipm->ce, &carr));
    513:   }

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

    541:   PetscCall(VecRestoreArrayRead(X, &Xarr));
    542:   PetscCall(VecRestoreArrayWrite(F, &Farr));
    543:   PetscFunctionReturn(PETSC_SUCCESS);
    544: }

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

    560:   PetscFunctionBegin;
    561:   PetscCall(TaoSNESFunction_PDIPM(snes, X, F, (void *)tao));
    562:   PetscCall(VecGetArrayWrite(F, &Farr));
    563:   PetscCall(VecGetArrayRead(X, &Xarr));

    565:   /* compute res[0] = norm2(F_x) */
    566:   L1 = pdipm->x;
    567:   PetscCall(VecPlaceArray(L1, Farr));
    568:   PetscCall(VecNorm(L1, NORM_2, &res[0]));
    569:   PetscCall(VecResetArray(L1));

    571:   /* compute res[1] = norm2(F_z), cnorm[1] = norm2(F_ci) */
    572:   if (pdipm->z) {
    573:     if (pdipm->solve_symmetric_kkt) {
    574:       PetscCall(VecPlaceArray(pdipm->z, Farr + pdipm->off_z));
    575:       if (pdipm->Nci) {
    576:         PetscCall(VecGetArrayWrite(pdipm->z, &tmparr));
    577:         for (i = 0; i < pdipm->nci; i++) tmparr[i] *= Xarr[pdipm->off_z + i];
    578:         PetscCall(VecRestoreArrayWrite(pdipm->z, &tmparr));
    579:       }

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

    583:       if (pdipm->Nci) {
    584:         PetscCall(VecGetArrayWrite(pdipm->z, &tmparr));
    585:         for (i = 0; i < pdipm->nci; i++) tmparr[i] /= Xarr[pdipm->off_z + i];
    586:         PetscCall(VecRestoreArrayWrite(pdipm->z, &tmparr));
    587:       }
    588:       PetscCall(VecResetArray(pdipm->z));
    589:     } else { /* !solve_symmetric_kkt */
    590:       PetscCall(VecPlaceArray(pdipm->z, Farr + pdipm->off_z));
    591:       PetscCall(VecNorm(pdipm->z, NORM_2, &res[1]));
    592:       PetscCall(VecResetArray(pdipm->z));
    593:     }

    595:     PetscCall(VecPlaceArray(pdipm->ci, Farr + pdipm->off_lambdai));
    596:     PetscCall(VecNorm(pdipm->ci, NORM_2, &cnorm[1]));
    597:     PetscCall(VecResetArray(pdipm->ci));
    598:   } else {
    599:     res[1]   = 0.0;
    600:     cnorm[1] = 0.0;
    601:   }

    603:   /* compute cnorm[0] = norm2(F_ce) */
    604:   if (pdipm->Nce) {
    605:     PetscCall(VecPlaceArray(pdipm->ce, Farr + pdipm->off_lambdae));
    606:     PetscCall(VecNorm(pdipm->ce, NORM_2, &cnorm[0]));
    607:     PetscCall(VecResetArray(pdipm->ce));
    608:   } else cnorm[0] = 0.0;

    610:   PetscCall(VecRestoreArrayWrite(F, &Farr));
    611:   PetscCall(VecRestoreArrayRead(X, &Xarr));

    613:   tao->gnorm0   = tao->residual;
    614:   tao->residual = PetscSqrtReal(res[0] * res[0] + res[1] * res[1]);
    615:   tao->cnorm    = PetscSqrtReal(cnorm[0] * cnorm[0] + cnorm[1] * cnorm[1]);
    616:   tao->step     = pdipm->mu;
    617:   PetscFunctionReturn(PETSC_SUCCESS);
    618: }

    620: /*
    621:   KKTAddShifts - Check the inertia of Cholesky factor of KKT matrix.
    622:   If it does not match the numbers of prime and dual variables, add shifts to the KKT matrix.
    623: */
    624: static PetscErrorCode KKTAddShifts(Tao tao, SNES snes, Vec X)
    625: {
    626:   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;
    627:   KSP        ksp;
    628:   PC         pc;
    629:   Mat        Factor;
    630:   PetscBool  isCHOL;
    631:   PetscInt   nneg, nzero, npos;

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

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

    643:   if (npos < pdipm->Nx + pdipm->Nci) {
    644:     pdipm->deltaw = PetscMax(pdipm->lastdeltaw / 3, 1.e-4 * PETSC_MACHINE_EPSILON);
    645:     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));
    646:     PetscCall(TaoSNESJacobian_PDIPM(snes, X, pdipm->K, pdipm->K, tao));
    647:     PetscCall(PCSetUp(pc));
    648:     PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));

    650:     if (npos < pdipm->Nx + pdipm->Nci) {
    651:       pdipm->deltaw = pdipm->lastdeltaw;                                           /* in case reduction update does not help, this prevents that step from impacting increasing update */
    652:       while (npos < pdipm->Nx + pdipm->Nci && pdipm->deltaw <= 1. / PETSC_SMALL) { /* increase deltaw */
    653:         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));
    654:         pdipm->deltaw = PetscMin(8 * pdipm->deltaw, PetscPowReal(10, 20));
    655:         PetscCall(TaoSNESJacobian_PDIPM(snes, X, pdipm->K, pdipm->K, tao));
    656:         PetscCall(PCSetUp(pc));
    657:         PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));
    658:       }

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

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

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

    682: /*
    683:   PCPreSolve_PDIPM -- called between MatFactorNumeric() and MatSolve()
    684: */
    685: static PetscErrorCode PCPreSolve_PDIPM(PC pc, KSP ksp)
    686: {
    687:   Tao        tao;
    688:   TAO_PDIPM *pdipm;

    690:   PetscFunctionBegin;
    691:   PetscCall(KSPGetApplicationContext(ksp, &tao));
    692:   pdipm = (TAO_PDIPM *)tao->data;
    693:   PetscCall(KKTAddShifts(tao, pdipm->snes, pdipm->X));
    694:   PetscFunctionReturn(PETSC_SUCCESS);
    695: }

    697: /*
    698:    SNESLineSearch_PDIPM - Custom line search used with PDIPM.

    700:    Collective

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

    722:   PetscFunctionBegin;
    723:   PetscCall(SNESLineSearchGetSNES(linesearch, &snes));
    724:   PetscCall(SNESGetIterationNumber(snes, &iter));

    726:   PetscCall(SNESLineSearchSetReason(linesearch, SNES_LINESEARCH_SUCCEEDED));
    727:   PetscCall(SNESLineSearchGetVecs(linesearch, &X, &F, &Y, NULL, NULL));

    729:   PetscCall(VecGetArrayWrite(X, &Xarr));
    730:   PetscCall(VecGetArrayRead(Y, &dXarr));
    731:   z  = Xarr + pdipm->off_z;
    732:   dz = dXarr + pdipm->off_z;
    733:   for (i = 0; i < pdipm->nci; i++) {
    734:     if (z[i] - dz[i] < 0.0) alpha_p = PetscMin(alpha_p, 0.9999 * z[i] / dz[i]);
    735:   }

    737:   lambdai  = Xarr + pdipm->off_lambdai;
    738:   dlambdai = dXarr + pdipm->off_lambdai;

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

    744:   alpha[0] = alpha_p;
    745:   alpha[1] = alpha_d;
    746:   PetscCall(VecRestoreArrayRead(Y, &dXarr));
    747:   PetscCall(VecRestoreArrayWrite(X, &Xarr));

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

    752:   alpha_p = alpha[2];
    753:   alpha_d = alpha[3];

    755:   /* X = X - alpha * Y */
    756:   PetscCall(VecGetArrayWrite(X, &Xarr));
    757:   PetscCall(VecGetArrayRead(Y, &dXarr));
    758:   for (i = 0; i < pdipm->nx; i++) Xarr[i] -= alpha_p * dXarr[i];
    759:   for (i = 0; i < pdipm->nce; i++) Xarr[i + pdipm->off_lambdae] -= alpha_d * dXarr[i + pdipm->off_lambdae];

    761:   for (i = 0; i < pdipm->nci; i++) {
    762:     Xarr[i + pdipm->off_lambdai] -= alpha_d * dXarr[i + pdipm->off_lambdai];
    763:     Xarr[i + pdipm->off_z] -= alpha_p * dXarr[i + pdipm->off_z];
    764:   }
    765:   PetscCall(VecGetArrayWrite(tao->solution, &taosolarr));
    766:   PetscCall(PetscMemcpy(taosolarr, Xarr, pdipm->nx * sizeof(PetscScalar)));
    767:   PetscCall(VecRestoreArrayWrite(tao->solution, &taosolarr));

    769:   PetscCall(VecRestoreArrayWrite(X, &Xarr));
    770:   PetscCall(VecRestoreArrayRead(Y, &dXarr));

    772:   /* Update mu = mu_update_factor * dot(z,lambdai)/pdipm->nci at updated X */
    773:   if (pdipm->z) {
    774:     PetscCall(VecDot(pdipm->z, pdipm->lambdai, &dot));
    775:   } else dot = 0.0;

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

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

    783:   tao->niter++;
    784:   PetscCall(TaoLogConvergenceHistory(tao, pdipm->obj, tao->residual, tao->cnorm, tao->niter));
    785:   PetscCall(TaoMonitor(tao, tao->niter, pdipm->obj, tao->residual, tao->cnorm, pdipm->mu));

    787:   PetscUseTypeMethod(tao, convergencetest, tao->cnvP);
    788:   if (tao->reason) PetscCall(SNESSetConvergedReason(snes, SNES_CONVERGED_FNORM_ABS));
    789:   PetscFunctionReturn(PETSC_SUCCESS);
    790: }

    792: static PetscErrorCode TaoSolve_PDIPM(Tao tao)
    793: {
    794:   TAO_PDIPM     *pdipm = (TAO_PDIPM *)tao->data;
    795:   SNESLineSearch linesearch; /* SNESLineSearch context */
    796:   Vec            dummy;

    798:   PetscFunctionBegin;
    799:   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");

    801:   /* Initialize all variables */
    802:   PetscCall(TaoPDIPMInitializeSolution(tao));

    804:   /* Set linesearch */
    805:   PetscCall(SNESGetLineSearch(pdipm->snes, &linesearch));
    806:   PetscCall(SNESLineSearchSetType(linesearch, SNESLINESEARCHSHELL));
    807:   PetscCall(SNESLineSearchShellSetApply(linesearch, SNESLineSearch_PDIPM, tao));
    808:   PetscCall(SNESLineSearchSetFromOptions(linesearch));

    810:   tao->reason = TAO_CONTINUE_ITERATING;

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

    816:   PetscCall(TaoLogConvergenceHistory(tao, pdipm->obj, tao->residual, tao->cnorm, tao->niter));
    817:   PetscCall(TaoMonitor(tao, tao->niter, pdipm->obj, tao->residual, tao->cnorm, pdipm->mu));
    818:   PetscCall(VecDestroy(&dummy));
    819:   PetscUseTypeMethod(tao, convergencetest, tao->cnvP);
    820:   if (tao->reason) PetscCall(SNESSetConvergedReason(pdipm->snes, SNES_CONVERGED_FNORM_ABS));

    822:   while (tao->reason == TAO_CONTINUE_ITERATING) {
    823:     SNESConvergedReason reason;
    824:     PetscCall(SNESSolve(pdipm->snes, NULL, pdipm->X));

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

    830:     /* Check TAO convergence */
    831:     PetscCheck(!PetscIsInfOrNanReal(pdipm->obj), PETSC_COMM_SELF, PETSC_ERR_SUP, "User-provided compute function generated Inf or NaN");
    832:   }
    833:   PetscFunctionReturn(PETSC_SUCCESS);
    834: }

    836: static PetscErrorCode TaoView_PDIPM(Tao tao, PetscViewer viewer)
    837: {
    838:   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;

    840:   PetscFunctionBegin;
    841:   tao->constrained = PETSC_TRUE;
    842:   PetscCall(PetscViewerASCIIPushTab(viewer));
    843:   PetscCall(PetscViewerASCIIPrintf(viewer, "Number of prime=%" PetscInt_FMT ", Number of dual=%" PetscInt_FMT "\n", pdipm->Nx + pdipm->Nci, pdipm->Nce + pdipm->Nci));
    844:   if (pdipm->kkt_pd) PetscCall(PetscViewerASCIIPrintf(viewer, "KKT shifts deltaw=%g, deltac=%g\n", (double)pdipm->deltaw, (double)pdipm->deltac));
    845:   PetscCall(PetscViewerASCIIPopTab(viewer));
    846:   PetscFunctionReturn(PETSC_SUCCESS);
    847: }

    849: static PetscErrorCode TaoSetup_PDIPM(Tao tao)
    850: {
    851:   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
    852:   MPI_Comm           comm;
    853:   PetscMPIInt        size;
    854:   PetscInt           row, col, Jcrstart, Jcrend, k, tmp, nc, proc, *nh_all, *ng_all;
    855:   PetscInt           offset, *xa, *xb, i, j, rstart, rend;
    856:   PetscScalar        one = 1.0, neg_one = -1.0;
    857:   const PetscInt    *cols, *rranges, *cranges, *aj, *ranges;
    858:   const PetscScalar *aa, *Xarr;
    859:   Mat                J;
    860:   Mat                Jce_xfixed_trans, Jci_xb_trans;
    861:   PetscInt          *dnz, *onz, rjstart, nx_all, *nce_all, *Jranges, cols1[2];

    863:   PetscFunctionBegin;
    864:   PetscCall(PetscObjectGetComm((PetscObject)tao, &comm));
    865:   PetscCallMPI(MPI_Comm_size(comm, &size));

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

    870:   if (!tao->gradient) {
    871:     PetscCall(VecDuplicate(tao->solution, &tao->gradient));
    872:     PetscCall(VecDuplicate(tao->solution, &tao->stepdirection));
    873:   }

    875:   /* (2) Get sizes */
    876:   /* Size of vector x - This is set by TaoSetSolution */
    877:   PetscCall(VecGetSize(tao->solution, &pdipm->Nx));
    878:   PetscCall(VecGetLocalSize(tao->solution, &pdipm->nx));

    880:   /* Size of equality constraints and vectors */
    881:   if (tao->constraints_equality) {
    882:     PetscCall(VecGetSize(tao->constraints_equality, &pdipm->Ng));
    883:     PetscCall(VecGetLocalSize(tao->constraints_equality, &pdipm->ng));
    884:   } else {
    885:     pdipm->ng = pdipm->Ng = 0;
    886:   }

    888:   pdipm->nce = pdipm->ng + pdipm->nxfixed;
    889:   pdipm->Nce = pdipm->Ng + pdipm->Nxfixed;

    891:   /* Size of inequality constraints and vectors */
    892:   if (tao->constraints_inequality) {
    893:     PetscCall(VecGetSize(tao->constraints_inequality, &pdipm->Nh));
    894:     PetscCall(VecGetLocalSize(tao->constraints_inequality, &pdipm->nh));
    895:   } else {
    896:     pdipm->nh = pdipm->Nh = 0;
    897:   }

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

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

    906:   /* (3) Offsets for subvectors */
    907:   pdipm->off_lambdae = pdipm->nx;
    908:   pdipm->off_lambdai = pdipm->off_lambdae + pdipm->nce;
    909:   pdipm->off_z       = pdipm->off_lambdai + pdipm->nci;

    911:   /* (4) Create vectors and subvectors */
    912:   /* Ce and Ci vectors */
    913:   PetscCall(VecCreate(comm, &pdipm->ce));
    914:   PetscCall(VecSetSizes(pdipm->ce, pdipm->nce, pdipm->Nce));
    915:   PetscCall(VecSetFromOptions(pdipm->ce));

    917:   PetscCall(VecCreate(comm, &pdipm->ci));
    918:   PetscCall(VecSetSizes(pdipm->ci, pdipm->nci, pdipm->Nci));
    919:   PetscCall(VecSetFromOptions(pdipm->ci));

    921:   /* X=[x; lambdae; lambdai; z] for the big KKT system */
    922:   PetscCall(VecCreate(comm, &pdipm->X));
    923:   PetscCall(VecSetSizes(pdipm->X, pdipm->n, pdipm->N));
    924:   PetscCall(VecSetFromOptions(pdipm->X));

    926:   /* Subvectors; they share local arrays with X */
    927:   PetscCall(VecGetArrayRead(pdipm->X, &Xarr));
    928:   /* x shares local array with X.x */
    929:   if (pdipm->Nx) PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nx, pdipm->Nx, Xarr, &pdipm->x));

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

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

    938:     PetscCall(VecCreate(comm, &pdipm->lambdae_xfixed));
    939:     PetscCall(VecSetSizes(pdipm->lambdae_xfixed, pdipm->nxfixed, PETSC_DECIDE));
    940:     PetscCall(VecSetFromOptions(pdipm->lambdae_xfixed));
    941:   }

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

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

    951:   /* tao->DI which shares local array with X.lambdai_h */
    952:   if (pdipm->Nh) PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nh, pdipm->Nh, Xarr + pdipm->off_lambdai, &tao->DI));
    953:   PetscCall(VecCreate(comm, &pdipm->lambdai_xb));
    954:   PetscCall(VecSetSizes(pdipm->lambdai_xb, pdipm->nci - pdipm->nh, PETSC_DECIDE));
    955:   PetscCall(VecSetFromOptions(pdipm->lambdai_xb));

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

    959:   /* (5) Create Jacobians Jce_xfixed and Jci */
    960:   /* (5.1) PDIPM Jacobian of equality bounds cebound(x) = J_nxfixed */
    961:   if (pdipm->Nxfixed) {
    962:     /* Create Jce_xfixed */
    963:     PetscCall(MatCreate(comm, &pdipm->Jce_xfixed));
    964:     PetscCall(MatSetSizes(pdipm->Jce_xfixed, pdipm->nxfixed, pdipm->nx, PETSC_DECIDE, pdipm->Nx));
    965:     PetscCall(MatSetFromOptions(pdipm->Jce_xfixed));
    966:     PetscCall(MatSeqAIJSetPreallocation(pdipm->Jce_xfixed, 1, NULL));
    967:     PetscCall(MatMPIAIJSetPreallocation(pdipm->Jce_xfixed, 1, NULL, 1, NULL));

    969:     PetscCall(MatGetOwnershipRange(pdipm->Jce_xfixed, &Jcrstart, &Jcrend));
    970:     PetscCall(ISGetIndices(pdipm->isxfixed, &cols));
    971:     k = 0;
    972:     for (row = Jcrstart; row < Jcrend; row++) {
    973:       PetscCall(MatSetValues(pdipm->Jce_xfixed, 1, &row, 1, cols + k, &one, INSERT_VALUES));
    974:       k++;
    975:     }
    976:     PetscCall(ISRestoreIndices(pdipm->isxfixed, &cols));
    977:     PetscCall(MatAssemblyBegin(pdipm->Jce_xfixed, MAT_FINAL_ASSEMBLY));
    978:     PetscCall(MatAssemblyEnd(pdipm->Jce_xfixed, MAT_FINAL_ASSEMBLY));
    979:   }

    981:   /* (5.2) PDIPM inequality Jacobian Jci = [tao->jacobian_inequality; ...] */
    982:   PetscCall(MatCreate(comm, &pdipm->Jci_xb));
    983:   PetscCall(MatSetSizes(pdipm->Jci_xb, pdipm->nci - pdipm->nh, pdipm->nx, PETSC_DECIDE, pdipm->Nx));
    984:   PetscCall(MatSetFromOptions(pdipm->Jci_xb));
    985:   PetscCall(MatSeqAIJSetPreallocation(pdipm->Jci_xb, 1, NULL));
    986:   PetscCall(MatMPIAIJSetPreallocation(pdipm->Jci_xb, 1, NULL, 1, NULL));

    988:   PetscCall(MatGetOwnershipRange(pdipm->Jci_xb, &Jcrstart, &Jcrend));
    989:   offset = Jcrstart;
    990:   if (pdipm->Nxub) {
    991:     /* Add xub to Jci_xb */
    992:     PetscCall(ISGetIndices(pdipm->isxub, &cols));
    993:     k = 0;
    994:     for (row = offset; row < offset + pdipm->nxub; row++) {
    995:       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &row, 1, cols + k, &neg_one, INSERT_VALUES));
    996:       k++;
    997:     }
    998:     PetscCall(ISRestoreIndices(pdipm->isxub, &cols));
    999:   }

    1001:   if (pdipm->Nxlb) {
    1002:     /* Add xlb to Jci_xb */
    1003:     PetscCall(ISGetIndices(pdipm->isxlb, &cols));
    1004:     k = 0;
    1005:     offset += pdipm->nxub;
    1006:     for (row = offset; row < offset + pdipm->nxlb; row++) {
    1007:       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &row, 1, cols + k, &one, INSERT_VALUES));
    1008:       k++;
    1009:     }
    1010:     PetscCall(ISRestoreIndices(pdipm->isxlb, &cols));
    1011:   }

    1013:   /* Add xbox to Jci_xb */
    1014:   if (pdipm->Nxbox) {
    1015:     PetscCall(ISGetIndices(pdipm->isxbox, &cols));
    1016:     k = 0;
    1017:     offset += pdipm->nxlb;
    1018:     for (row = offset; row < offset + pdipm->nxbox; row++) {
    1019:       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &row, 1, cols + k, &neg_one, INSERT_VALUES));
    1020:       tmp = row + pdipm->nxbox;
    1021:       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &tmp, 1, cols + k, &one, INSERT_VALUES));
    1022:       k++;
    1023:     }
    1024:     PetscCall(ISRestoreIndices(pdipm->isxbox, &cols));
    1025:   }

    1027:   PetscCall(MatAssemblyBegin(pdipm->Jci_xb, MAT_FINAL_ASSEMBLY));
    1028:   PetscCall(MatAssemblyEnd(pdipm->Jci_xb, MAT_FINAL_ASSEMBLY));
    1029:   /* PetscCall(MatView(pdipm->Jci_xb,PETSC_VIEWER_STDOUT_WORLD)); */

    1031:   /* (6) Set up ISs for PC Fieldsplit */
    1032:   if (pdipm->solve_reduced_kkt) {
    1033:     PetscCall(PetscMalloc2(pdipm->nx + pdipm->nce, &xa, 2 * pdipm->nci, &xb));
    1034:     for (i = 0; i < pdipm->nx + pdipm->nce; i++) xa[i] = i;
    1035:     for (i = 0; i < 2 * pdipm->nci; i++) xb[i] = pdipm->off_lambdai + i;

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

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

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

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

    1050:   PetscCall(PetscMalloc3(size, &ng_all, size, &nh_all, size, &Jranges));
    1051:   PetscCallMPI(MPI_Allgather(&rstart, 1, MPIU_INT, Jranges, 1, MPIU_INT, comm));
    1052:   PetscCallMPI(MPI_Allgather(&pdipm->nh, 1, MPIU_INT, nh_all, 1, MPIU_INT, comm));
    1053:   PetscCallMPI(MPI_Allgather(&pdipm->ng, 1, MPIU_INT, ng_all, 1, MPIU_INT, comm));

    1055:   PetscCall(MatGetOwnershipRanges(tao->hessian, &rranges));
    1056:   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &cranges));

    1058:   if (pdipm->Ng) {
    1059:     PetscCall(TaoComputeJacobianEquality(tao, tao->solution, tao->jacobian_equality, tao->jacobian_equality_pre));
    1060:     PetscCall(MatTranspose(tao->jacobian_equality, MAT_INITIAL_MATRIX, &pdipm->jac_equality_trans));
    1061:   }
    1062:   if (pdipm->Nh) {
    1063:     PetscCall(TaoComputeJacobianInequality(tao, tao->solution, tao->jacobian_inequality, tao->jacobian_inequality_pre));
    1064:     PetscCall(MatTranspose(tao->jacobian_inequality, MAT_INITIAL_MATRIX, &pdipm->jac_inequality_trans));
    1065:   }

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

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

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

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

    1079:   /* Insert tao->hessian */
    1080:   PetscCall(MatGetOwnershipRange(tao->hessian, &rjstart, NULL));
    1081:   for (i = 0; i < pdipm->nx; i++) {
    1082:     row = rstart + i;

    1084:     PetscCall(MatGetRow(tao->hessian, i + rjstart, &nc, &aj, NULL));
    1085:     proc = 0;
    1086:     for (j = 0; j < nc; j++) {
    1087:       while (aj[j] >= cranges[proc + 1]) proc++;
    1088:       col = aj[j] - cranges[proc] + Jranges[proc];
    1089:       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1090:     }
    1091:     PetscCall(MatRestoreRow(tao->hessian, i + rjstart, &nc, &aj, NULL));

    1093:     if (pdipm->ng) {
    1094:       /* Insert grad g' */
    1095:       PetscCall(MatGetRow(pdipm->jac_equality_trans, i + rjstart, &nc, &aj, NULL));
    1096:       PetscCall(MatGetOwnershipRanges(tao->jacobian_equality, &ranges));
    1097:       proc = 0;
    1098:       for (j = 0; j < nc; j++) {
    1099:         /* find row ownership of */
    1100:         while (aj[j] >= ranges[proc + 1]) proc++;
    1101:         nx_all = rranges[proc + 1] - rranges[proc];
    1102:         col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all;
    1103:         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1104:       }
    1105:       PetscCall(MatRestoreRow(pdipm->jac_equality_trans, i + rjstart, &nc, &aj, NULL));
    1106:     }

    1108:     /* Insert Jce_xfixed^T' */
    1109:     if (pdipm->nxfixed) {
    1110:       PetscCall(MatGetRow(Jce_xfixed_trans, i + rjstart, &nc, &aj, NULL));
    1111:       PetscCall(MatGetOwnershipRanges(pdipm->Jce_xfixed, &ranges));
    1112:       proc = 0;
    1113:       for (j = 0; j < nc; j++) {
    1114:         /* find row ownership of */
    1115:         while (aj[j] >= ranges[proc + 1]) proc++;
    1116:         nx_all = rranges[proc + 1] - rranges[proc];
    1117:         col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all + ng_all[proc];
    1118:         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1119:       }
    1120:       PetscCall(MatRestoreRow(Jce_xfixed_trans, i + rjstart, &nc, &aj, NULL));
    1121:     }

    1123:     if (pdipm->nh) {
    1124:       /* Insert -grad h' */
    1125:       PetscCall(MatGetRow(pdipm->jac_inequality_trans, i + rjstart, &nc, &aj, NULL));
    1126:       PetscCall(MatGetOwnershipRanges(tao->jacobian_inequality, &ranges));
    1127:       proc = 0;
    1128:       for (j = 0; j < nc; j++) {
    1129:         /* find row ownership of */
    1130:         while (aj[j] >= ranges[proc + 1]) proc++;
    1131:         nx_all = rranges[proc + 1] - rranges[proc];
    1132:         col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all + nce_all[proc];
    1133:         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1134:       }
    1135:       PetscCall(MatRestoreRow(pdipm->jac_inequality_trans, i + rjstart, &nc, &aj, NULL));
    1136:     }

    1138:     /* Insert Jci_xb^T' */
    1139:     PetscCall(MatGetRow(Jci_xb_trans, i + rjstart, &nc, &aj, NULL));
    1140:     PetscCall(MatGetOwnershipRanges(pdipm->Jci_xb, &ranges));
    1141:     proc = 0;
    1142:     for (j = 0; j < nc; j++) {
    1143:       /* find row ownership of */
    1144:       while (aj[j] >= ranges[proc + 1]) proc++;
    1145:       nx_all = rranges[proc + 1] - rranges[proc];
    1146:       col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all + nce_all[proc] + nh_all[proc];
    1147:       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1148:     }
    1149:     PetscCall(MatRestoreRow(Jci_xb_trans, i + rjstart, &nc, &aj, NULL));
    1150:   }

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

    1158:       PetscCall(MatGetRow(tao->jacobian_equality, i + rjstart, &nc, &aj, NULL));
    1159:       proc = 0;
    1160:       for (j = 0; j < nc; j++) {
    1161:         while (aj[j] >= cranges[proc + 1]) proc++;
    1162:         col = aj[j] - cranges[proc] + Jranges[proc];
    1163:         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz)); /* grad g */
    1164:       }
    1165:       PetscCall(MatRestoreRow(tao->jacobian_equality, i + rjstart, &nc, &aj, NULL));
    1166:     }
    1167:   }
    1168:   /* Jce_xfixed */
    1169:   if (pdipm->Nxfixed) {
    1170:     PetscCall(MatGetOwnershipRange(pdipm->Jce_xfixed, &Jcrstart, NULL));
    1171:     for (i = 0; i < (pdipm->nce - pdipm->ng); i++) {
    1172:       row = rstart + pdipm->off_lambdae + pdipm->ng + i;

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

    1177:       proc = 0;
    1178:       j    = 0;
    1179:       while (cols[j] >= cranges[proc + 1]) proc++;
    1180:       col = cols[j] - cranges[proc] + Jranges[proc];
    1181:       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1182:       PetscCall(MatRestoreRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, NULL));
    1183:     }
    1184:   }

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

    1192:       PetscCall(MatGetRow(tao->jacobian_inequality, i + rjstart, &nc, &aj, NULL));
    1193:       proc = 0;
    1194:       for (j = 0; j < nc; j++) {
    1195:         while (aj[j] >= cranges[proc + 1]) proc++;
    1196:         col = aj[j] - cranges[proc] + Jranges[proc];
    1197:         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz)); /* grad h */
    1198:       }
    1199:       PetscCall(MatRestoreRow(tao->jacobian_inequality, i + rjstart, &nc, &aj, NULL));
    1200:     }
    1201:     /* I */
    1202:     for (i = 0; i < pdipm->nh; i++) {
    1203:       row = rstart + pdipm->off_lambdai + i;
    1204:       col = rstart + pdipm->off_z + i;
    1205:       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1206:     }
    1207:   }

    1209:   /* Jci_xb */
    1210:   PetscCall(MatGetOwnershipRange(pdipm->Jci_xb, &Jcrstart, NULL));
    1211:   for (i = 0; i < (pdipm->nci - pdipm->nh); i++) {
    1212:     row = rstart + pdipm->off_lambdai + pdipm->nh + i;

    1214:     PetscCall(MatGetRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, NULL));
    1215:     PetscCheck(nc == 1, PETSC_COMM_SELF, PETSC_ERR_SUP, "nc != 1");
    1216:     proc = 0;
    1217:     for (j = 0; j < nc; j++) {
    1218:       while (cols[j] >= cranges[proc + 1]) proc++;
    1219:       col = cols[j] - cranges[proc] + Jranges[proc];
    1220:       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1221:     }
    1222:     PetscCall(MatRestoreRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, NULL));
    1223:     /* I */
    1224:     col = rstart + pdipm->off_z + pdipm->nh + i;
    1225:     PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
    1226:   }

    1228:   /* 4-th Row block of KKT matrix: Z and Ci */
    1229:   for (i = 0; i < pdipm->nci; i++) {
    1230:     row      = rstart + pdipm->off_z + i;
    1231:     cols1[0] = rstart + pdipm->off_lambdai + i;
    1232:     cols1[1] = row;
    1233:     PetscCall(MatPreallocateSet(row, 2, cols1, dnz, onz));
    1234:   }

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

    1239:   /* Create KKT matrix */
    1240:   PetscCall(MatCreate(comm, &J));
    1241:   PetscCall(MatSetSizes(J, pdipm->n, pdipm->n, PETSC_DECIDE, PETSC_DECIDE));
    1242:   PetscCall(MatSetFromOptions(J));
    1243:   PetscCall(MatSeqAIJSetPreallocation(J, 0, dnz));
    1244:   PetscCall(MatMPIAIJSetPreallocation(J, 0, dnz, 0, onz));
    1245:   MatPreallocateEnd(dnz, onz);
    1246:   pdipm->K = J;

    1248:   /* (8) Insert constant entries to  K */
    1249:   /* Set 0.0 to diagonal of K, so that the solver does not complain *about missing diagonal value */
    1250:   PetscCall(MatGetOwnershipRange(J, &rstart, &rend));
    1251:   for (i = rstart; i < rend; i++) PetscCall(MatSetValue(J, i, i, 0.0, INSERT_VALUES));
    1252:   /* In case Wxx has no diagonal entries preset set diagonal to deltaw given */
    1253:   if (pdipm->kkt_pd) {
    1254:     for (i = 0; i < pdipm->nh; i++) {
    1255:       row = rstart + i;
    1256:       PetscCall(MatSetValue(J, row, row, pdipm->deltaw, INSERT_VALUES));
    1257:     }
    1258:   }

    1260:   /* Row block of K: [ grad Ce, 0, 0, 0] */
    1261:   if (pdipm->Nxfixed) {
    1262:     PetscCall(MatGetOwnershipRange(pdipm->Jce_xfixed, &Jcrstart, NULL));
    1263:     for (i = 0; i < (pdipm->nce - pdipm->ng); i++) {
    1264:       row = rstart + pdipm->off_lambdae + pdipm->ng + i;

    1266:       PetscCall(MatGetRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, &aa));
    1267:       proc = 0;
    1268:       for (j = 0; j < nc; j++) {
    1269:         while (cols[j] >= cranges[proc + 1]) proc++;
    1270:         col = cols[j] - cranges[proc] + Jranges[proc];
    1271:         PetscCall(MatSetValue(J, row, col, aa[j], INSERT_VALUES)); /* grad Ce */
    1272:         PetscCall(MatSetValue(J, col, row, aa[j], INSERT_VALUES)); /* grad Ce' */
    1273:       }
    1274:       PetscCall(MatRestoreRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, &aa));
    1275:     }
    1276:   }

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

    1283:     PetscCall(MatGetRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, &aa));
    1284:     proc = 0;
    1285:     for (j = 0; j < nc; j++) {
    1286:       while (cols[j] >= cranges[proc + 1]) proc++;
    1287:       col = cols[j] - cranges[proc] + Jranges[proc];
    1288:       PetscCall(MatSetValue(J, col, row, -aa[j], INSERT_VALUES));
    1289:       PetscCall(MatSetValue(J, row, col, -aa[j], INSERT_VALUES));
    1290:     }
    1291:     PetscCall(MatRestoreRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, &aa));

    1293:     col = rstart + pdipm->off_z + pdipm->nh + i;
    1294:     PetscCall(MatSetValue(J, row, col, 1, INSERT_VALUES));
    1295:   }

    1297:   for (i = 0; i < pdipm->nh; i++) {
    1298:     row = rstart + pdipm->off_lambdai + i;
    1299:     col = rstart + pdipm->off_z + i;
    1300:     PetscCall(MatSetValue(J, row, col, 1, INSERT_VALUES));
    1301:   }

    1303:   /* Row block of K: [ 0, 0, I, ...] */
    1304:   for (i = 0; i < pdipm->nci; i++) {
    1305:     row = rstart + pdipm->off_z + i;
    1306:     col = rstart + pdipm->off_lambdai + i;
    1307:     PetscCall(MatSetValue(J, row, col, 1, INSERT_VALUES));
    1308:   }

    1310:   if (pdipm->Nxfixed) PetscCall(MatDestroy(&Jce_xfixed_trans));
    1311:   PetscCall(MatDestroy(&Jci_xb_trans));
    1312:   PetscCall(PetscFree3(ng_all, nh_all, Jranges));

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

    1318:   if (pdipm->solve_reduced_kkt) {
    1319:     PC pc;
    1320:     PetscCall(KSPGetPC(tao->ksp, &pc));
    1321:     PetscCall(PCSetType(pc, PCFIELDSPLIT));
    1322:     PetscCall(PCFieldSplitSetType(pc, PC_COMPOSITE_SCHUR));
    1323:     PetscCall(PCFieldSplitSetIS(pc, "2", pdipm->is2));
    1324:     PetscCall(PCFieldSplitSetIS(pc, "1", pdipm->is1));
    1325:   }
    1326:   PetscCall(SNESSetFromOptions(pdipm->snes));

    1328:   /* (10) Setup PCPreSolve() for pdipm->solve_symmetric_kkt */
    1329:   if (pdipm->solve_symmetric_kkt) {
    1330:     KSP       ksp;
    1331:     PC        pc;
    1332:     PetscBool isCHOL;
    1333:     PetscCall(SNESGetKSP(pdipm->snes, &ksp));
    1334:     PetscCall(KSPGetPC(ksp, &pc));
    1335:     PetscCall(PCSetPreSolve(pc, PCPreSolve_PDIPM));

    1337:     PetscCall(PetscObjectTypeCompare((PetscObject)pc, PCCHOLESKY, &isCHOL));
    1338:     if (isCHOL) {
    1339:       Mat       Factor;
    1340:       PetscBool isMUMPS;
    1341:       PetscCall(PCFactorGetMatrix(pc, &Factor));
    1342:       PetscCall(PetscObjectTypeCompare((PetscObject)Factor, "mumps", &isMUMPS));
    1343:       if (isMUMPS) { /* must set mumps ICNTL(13)=1 and ICNTL(24)=1 to call MatGetInertia() */
    1344: #if defined(PETSC_HAVE_MUMPS)
    1345:         PetscCall(MatMumpsSetIcntl(Factor, 24, 1)); /* detection of null pivot rows */
    1346:         if (size > 1) { PetscCall(MatMumpsSetIcntl(Factor, 13, 1)); /* parallelism of the root node (enable ScaLAPACK) and its splitting */ }
    1347: #else
    1348:         SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_SUP, "Requires external package MUMPS");
    1349: #endif
    1350:       }
    1351:     }
    1352:   }
    1353:   PetscFunctionReturn(PETSC_SUCCESS);
    1354: }

    1356: static PetscErrorCode TaoDestroy_PDIPM(Tao tao)
    1357: {
    1358:   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;

    1360:   PetscFunctionBegin;
    1361:   /* Freeing Vectors assocaiated with KKT (X) */
    1362:   PetscCall(VecDestroy(&pdipm->x));       /* Solution x */
    1363:   PetscCall(VecDestroy(&pdipm->lambdae)); /* Equality constraints lagrangian multiplier*/
    1364:   PetscCall(VecDestroy(&pdipm->lambdai)); /* Inequality constraints lagrangian multiplier*/
    1365:   PetscCall(VecDestroy(&pdipm->z));       /* Slack variables */
    1366:   PetscCall(VecDestroy(&pdipm->X));       /* Big KKT system vector [x; lambdae; lambdai; z] */

    1368:   /* work vectors */
    1369:   PetscCall(VecDestroy(&pdipm->lambdae_xfixed));
    1370:   PetscCall(VecDestroy(&pdipm->lambdai_xb));

    1372:   /* Legrangian equality and inequality Vec */
    1373:   PetscCall(VecDestroy(&pdipm->ce)); /* Vec of equality constraints */
    1374:   PetscCall(VecDestroy(&pdipm->ci)); /* Vec of inequality constraints */

    1376:   /* Matrices */
    1377:   PetscCall(MatDestroy(&pdipm->Jce_xfixed));
    1378:   PetscCall(MatDestroy(&pdipm->Jci_xb)); /* Jacobian of inequality constraints Jci = [tao->jacobian_inequality ; J(nxub); J(nxlb); J(nxbx)] */
    1379:   PetscCall(MatDestroy(&pdipm->K));

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

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

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

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

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

    1392:   if (pdipm->solve_reduced_kkt) {
    1393:     PetscCall(ISDestroy(&pdipm->is1));
    1394:     PetscCall(ISDestroy(&pdipm->is2));
    1395:   }

    1397:   /* SNES */
    1398:   PetscCall(SNESDestroy(&pdipm->snes)); /* Nonlinear solver */
    1399:   PetscCall(PetscFree(pdipm->nce_all));
    1400:   PetscCall(MatDestroy(&pdipm->jac_equality_trans));
    1401:   PetscCall(MatDestroy(&pdipm->jac_inequality_trans));

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

    1406:   /* Destroy Dual */
    1407:   PetscCall(VecDestroy(&tao->DE)); /* equality dual */
    1408:   PetscCall(VecDestroy(&tao->DI)); /* dinequality dual */
    1409:   PetscFunctionReturn(PETSC_SUCCESS);
    1410: }

    1412: static PetscErrorCode TaoSetFromOptions_PDIPM(Tao tao, PetscOptionItems PetscOptionsObject)
    1413: {
    1414:   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;

    1416:   PetscFunctionBegin;
    1417:   PetscOptionsHeadBegin(PetscOptionsObject, "PDIPM method for constrained optimization");
    1418:   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));
    1419:   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));
    1420:   PetscCall(PetscOptionsBool("-tao_pdipm_solve_reduced_kkt", "Solve reduced KKT system using Schur-complement", NULL, pdipm->solve_reduced_kkt, &pdipm->solve_reduced_kkt, NULL));
    1421:   PetscCall(PetscOptionsReal("-tao_pdipm_mu_update_factor", "Update scalar for barrier parameter (mu) update", NULL, pdipm->mu_update_factor, &pdipm->mu_update_factor, NULL));
    1422:   PetscCall(PetscOptionsBool("-tao_pdipm_symmetric_kkt", "Solve non reduced symmetric KKT system", NULL, pdipm->solve_symmetric_kkt, &pdipm->solve_symmetric_kkt, NULL));
    1423:   PetscCall(PetscOptionsBool("-tao_pdipm_kkt_shift_pd", "Add shifts to make KKT matrix positive definite", NULL, pdipm->kkt_pd, &pdipm->kkt_pd, NULL));
    1424:   PetscOptionsHeadEnd();
    1425:   PetscFunctionReturn(PETSC_SUCCESS);
    1426: }

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

    1431:   Options Database Keys:
    1432: +   -tao_pdipm_push_init_lambdai - parameter to push initial dual variables away from bounds (> 0)
    1433: .   -tao_pdipm_push_init_slack - parameter to push initial slack variables away from bounds (> 0)
    1434: .   -tao_pdipm_mu_update_factor - update scalar for barrier parameter (mu) update (> 0)
    1435: .   -tao_pdipm_symmetric_kkt - Solve non-reduced symmetric KKT system
    1436: -   -tao_pdipm_kkt_shift_pd - Add shifts to make KKT matrix positive definite

    1438:   Level: beginner

    1440: .seealso: `TAOPDIPM`, `Tao`, `TaoType`
    1441: M*/

    1443: PETSC_EXTERN PetscErrorCode TaoCreate_PDIPM(Tao tao)
    1444: {
    1445:   TAO_PDIPM *pdipm;

    1447:   PetscFunctionBegin;
    1448:   tao->ops->setup          = TaoSetup_PDIPM;
    1449:   tao->ops->solve          = TaoSolve_PDIPM;
    1450:   tao->ops->setfromoptions = TaoSetFromOptions_PDIPM;
    1451:   tao->ops->view           = TaoView_PDIPM;
    1452:   tao->ops->destroy        = TaoDestroy_PDIPM;

    1454:   PetscCall(PetscNew(&pdipm));
    1455:   tao->data = (void *)pdipm;

    1457:   pdipm->nx = pdipm->Nx = 0;
    1458:   pdipm->nxfixed = pdipm->Nxfixed = 0;
    1459:   pdipm->nxlb = pdipm->Nxlb = 0;
    1460:   pdipm->nxub = pdipm->Nxub = 0;
    1461:   pdipm->nxbox = pdipm->Nxbox = 0;
    1462:   pdipm->nxfree = pdipm->Nxfree = 0;

    1464:   pdipm->ng = pdipm->Ng = pdipm->nce = pdipm->Nce = 0;
    1465:   pdipm->nh = pdipm->Nh = pdipm->nci = pdipm->Nci = 0;
    1466:   pdipm->n = pdipm->N     = 0;
    1467:   pdipm->mu               = 1.0;
    1468:   pdipm->mu_update_factor = 0.1;

    1470:   pdipm->deltaw     = 0.0;
    1471:   pdipm->lastdeltaw = 3 * 1.e-4;
    1472:   pdipm->deltac     = 0.0;
    1473:   pdipm->kkt_pd     = PETSC_FALSE;

    1475:   pdipm->push_init_slack     = 1.0;
    1476:   pdipm->push_init_lambdai   = 1.0;
    1477:   pdipm->solve_reduced_kkt   = PETSC_FALSE;
    1478:   pdipm->solve_symmetric_kkt = PETSC_TRUE;

    1480:   /* Override default settings (unless already changed) */
    1481:   PetscCall(TaoParametersInitialize(tao));
    1482:   PetscObjectParameterSetDefault(tao, max_it, 200);
    1483:   PetscObjectParameterSetDefault(tao, max_funcs, 500);

    1485:   PetscCall(SNESCreate(((PetscObject)tao)->comm, &pdipm->snes));
    1486:   PetscCall(SNESSetOptionsPrefix(pdipm->snes, tao->hdr.prefix));
    1487:   PetscCall(SNESGetKSP(pdipm->snes, &tao->ksp));
    1488:   PetscCall(PetscObjectReference((PetscObject)tao->ksp));
    1489:   PetscCall(KSPSetApplicationContext(tao->ksp, (void *)tao));
    1490:   PetscFunctionReturn(PETSC_SUCCESS);
    1491: }