xref: /petsc/src/tao/constrained/impls/ipm/pdipm.c (revision 98d129c30f3ee9fdddc40fdbc5a989b7be64f888)
1 #include <../src/tao/constrained/impls/ipm/pdipm.h>
2 
3 /*
4    TaoPDIPMEvaluateFunctionsAndJacobians - Evaluate the objective function f, gradient fx, constraints, and all the Jacobians at current vector
5 
6    Collective
7 
8    Input Parameter:
9 +  tao - solver context
10 -  x - vector at which all objects to be evaluated
11 
12    Level: beginner
13 
14 .seealso: `TAOPDIPM`, `TaoPDIPMUpdateConstraints()`, `TaoPDIPMSetUpBounds()`
15 */
16 static PetscErrorCode TaoPDIPMEvaluateFunctionsAndJacobians(Tao tao, Vec x)
17 {
18   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;
19 
20   PetscFunctionBegin;
21   /* Compute user objective function and gradient */
22   PetscCall(TaoComputeObjectiveAndGradient(tao, x, &pdipm->obj, tao->gradient));
23 
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   }
29 
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 }
37 
38 /*
39   TaoPDIPMUpdateConstraints - Update the vectors ce and ci at x
40 
41   Collective
42 
43   Input Parameter:
44 + tao - Tao context
45 - x - vector at which constraints to be evaluated
46 
47    Level: beginner
48 
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;
58 
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));
64 
65   /* (1) Update ce vector */
66   PetscCall(VecGetArrayWrite(pdipm->ce, &carr));
67 
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   }
74 
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));
85 
86   /* (2) Update ci vector */
87   PetscCall(VecGetArrayWrite(pdipm->ci, &carr));
88 
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   }
95 
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   }
105 
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   }
115 
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));
128 
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 }
135 
136 /*
137    TaoPDIPMSetUpBounds - Create upper and lower bound vectors of x
138 
139    Collective
140 
141    Input Parameter:
142 .  tao - holds pdipm and XL & XU
143 
144    Level: beginner
145 
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];
155 
156   PetscFunctionBegin;
157   /* Creates upper and lower bounds vectors on x, if not created already */
158   PetscCall(TaoComputeVariableBounds(tao));
159 
160   PetscCall(VecGetLocalSize(tao->XL, &n));
161   PetscCall(PetscMalloc5(n, &ixlb, n, &ixub, n, &ixfree, n, &ixfixed, n, &ixbox));
162 
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));
182 
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;
189 
190   PetscCall(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];
196 
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 }
205 
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
212 
213    Collective
214 
215    Input Parameter:
216 .  tao - Tao context
217 
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;
226 
227   PetscFunctionBegin;
228   PetscCall(VecGetArrayWrite(pdipm->X, &Xarr));
229 
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));
234 
235   /* Initialize X.lambdae = 0.0 */
236   if (pdipm->lambdae) PetscCall(VecSet(pdipm->lambdae, 0.0));
237 
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));
242 
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   }
257 
258   PetscCall(VecRestoreArrayWrite(pdipm->X, &Xarr));
259   PetscFunctionReturn(PETSC_SUCCESS);
260 }
261 
262 /*
263    TaoSNESJacobian_PDIPM - Evaluate the Hessian matrix at X
264 
265    Input Parameter:
266    snes - SNES context
267    X - KKT Vector
268    *ctx - pdipm context
269 
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;
285 
286   PetscFunctionBegin;
287   PetscCall(PetscObjectGetComm((PetscObject)snes, &comm));
288   PetscCallMPI(MPI_Comm_rank(comm, &rank));
289   PetscCallMPI(MPI_Comm_rank(comm, &size));
290 
291   PetscCall(MatGetOwnershipRanges(Jpre, &Jranges));
292   PetscCall(MatGetOwnershipRange(Jpre, &Jrstart, NULL));
293   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &rranges));
294   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &cranges));
295 
296   PetscCall(VecGetArrayRead(X, &Xarr));
297 
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   }
318 
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;
324 
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   }
339 
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   }
359 
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   }
367 
368   PetscCall(VecPlaceArray(pdipm->x, Xarr));
369   PetscCall(TaoComputeHessian(tao, pdipm->x, tao->hessian, tao->hessian_pre));
370   PetscCall(VecResetArray(pdipm->x));
371 
372   PetscCall(MatGetOwnershipRange(tao->hessian, &rjstart, NULL));
373   for (i = 0; i < pdipm->nx; i++) {
374     row = Jrstart + i;
375 
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));
390 
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     }
405 
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));
422 
423   /* (6) assemble Jpre and J */
424   PetscCall(MatAssemblyBegin(Jpre, MAT_FINAL_ASSEMBLY));
425   PetscCall(MatAssemblyEnd(Jpre, MAT_FINAL_ASSEMBLY));
426 
427   if (J != Jpre) {
428     PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
429     PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
430   }
431   PetscFunctionReturn(PETSC_SUCCESS);
432 }
433 
434 /*
435    TaoSnesFunction_PDIPM - Evaluate KKT function at X
436 
437    Input Parameter:
438    snes - SNES context
439    X - KKT Vector
440    *ctx - pdipm
441 
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;
453 
454   PetscFunctionBegin;
455   PetscCall(VecSet(F, 0.0));
456 
457   PetscCall(VecGetArrayRead(X, &Xarr));
458   PetscCall(VecGetArrayWrite(F, &Farr));
459 
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));
464 
465   /* Update ce, ci, and Jci at X.x */
466   PetscCall(TaoPDIPMUpdateConstraints(tao, x));
467   PetscCall(VecResetArray(x));
468 
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     }
479 
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));
484 
485     /* L1 = - (gradH'*DI + Jci_xb'*lambdai_xb) */
486     PetscCall(VecScale(L1, -1.0));
487   }
488 
489   /* L1 += fx */
490   PetscCall(VecAXPY(L1, 1.0, tao->gradient));
491 
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));
507 
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   }
514 
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   }
540 
541   PetscCall(VecRestoreArrayRead(X, &Xarr));
542   PetscCall(VecRestoreArrayWrite(F, &Farr));
543   PetscFunctionReturn(PETSC_SUCCESS);
544 }
545 
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;
559 
560   PetscFunctionBegin;
561   PetscCall(TaoSNESFunction_PDIPM(snes, X, F, (void *)tao));
562   PetscCall(VecGetArrayWrite(F, &Farr));
563   PetscCall(VecGetArrayRead(X, &Xarr));
564 
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));
570 
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       }
580 
581       PetscCall(VecNorm(pdipm->z, NORM_2, &res[1]));
582 
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     }
594 
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   }
602 
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;
609 
610   PetscCall(VecRestoreArrayWrite(F, &Farr));
611   PetscCall(VecRestoreArrayRead(X, &Xarr));
612 
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 }
619 
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;
632 
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);
639 
640   PetscCall(PCFactorGetMatrix(pc, &Factor));
641   PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));
642 
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));
649 
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       }
659 
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");
661 
662       PetscCall(PetscInfo(tao, "Updated deltaw %g\n", (double)pdipm->deltaw));
663       pdipm->lastdeltaw = pdipm->deltaw;
664       pdipm->deltaw     = 0.0;
665     }
666   }
667 
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 }
681 
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;
689 
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 }
696 
697 /*
698    SNESLineSearch_PDIPM - Custom line search used with PDIPM.
699 
700    Collective
701 
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;
721 
722   PetscFunctionBegin;
723   PetscCall(SNESLineSearchGetSNES(linesearch, &snes));
724   PetscCall(SNESGetIterationNumber(snes, &iter));
725 
726   PetscCall(SNESLineSearchSetReason(linesearch, SNES_LINESEARCH_SUCCEEDED));
727   PetscCall(SNESLineSearchGetVecs(linesearch, &X, &F, &Y, NULL, NULL));
728 
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   }
736 
737   lambdai  = Xarr + pdipm->off_lambdai;
738   dlambdai = dXarr + pdipm->off_lambdai;
739 
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   }
743 
744   alpha[0] = alpha_p;
745   alpha[1] = alpha_d;
746   PetscCall(VecRestoreArrayRead(Y, &dXarr));
747   PetscCall(VecRestoreArrayWrite(X, &Xarr));
748 
749   /* alpha = min(alpha) over all processes */
750   PetscCall(MPIU_Allreduce(alpha, alpha + 2, 2, MPIU_REAL, MPIU_MIN, PetscObjectComm((PetscObject)tao)));
751 
752   alpha_p = alpha[2];
753   alpha_d = alpha[3];
754 
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];
760 
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));
768 
769   PetscCall(VecRestoreArrayWrite(X, &Xarr));
770   PetscCall(VecRestoreArrayRead(Y, &dXarr));
771 
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;
776 
777   /* if (PetscAbsReal(pdipm->gradL) < 0.9*pdipm->mu)  */
778   pdipm->mu = pdipm->mu_update_factor * dot / pdipm->Nci;
779 
780   /* Update F; get tao->residual and tao->cnorm */
781   PetscCall(TaoSNESFunction_PDIPM_residual(snes, X, F, (void *)tao));
782 
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));
786 
787   PetscUseTypeMethod(tao, convergencetest, tao->cnvP);
788   if (tao->reason) PetscCall(SNESSetConvergedReason(snes, SNES_CONVERGED_FNORM_ABS));
789   PetscFunctionReturn(PETSC_SUCCESS);
790 }
791 
792 static PetscErrorCode TaoSolve_PDIPM(Tao tao)
793 {
794   TAO_PDIPM     *pdipm = (TAO_PDIPM *)tao->data;
795   SNESLineSearch linesearch; /* SNESLineSearch context */
796   Vec            dummy;
797 
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");
800 
801   /* Initialize all variables */
802   PetscCall(TaoPDIPMInitializeSolution(tao));
803 
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));
809 
810   tao->reason = TAO_CONTINUE_ITERATING;
811 
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));
815 
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));
821 
822   while (tao->reason == TAO_CONTINUE_ITERATING) {
823     SNESConvergedReason reason;
824     PetscCall(SNESSolve(pdipm->snes, NULL, pdipm->X));
825 
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]));
829 
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 }
835 
836 static PetscErrorCode TaoView_PDIPM(Tao tao, PetscViewer viewer)
837 {
838   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;
839 
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 }
848 
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];
862 
863   PetscFunctionBegin;
864   PetscCall(PetscObjectGetComm((PetscObject)tao, &comm));
865   PetscCallMPI(MPI_Comm_size(comm, &size));
866 
867   /* (1) Setup Bounds and create Tao vectors */
868   PetscCall(TaoPDIPMSetUpBounds(tao));
869 
870   if (!tao->gradient) {
871     PetscCall(VecDuplicate(tao->solution, &tao->gradient));
872     PetscCall(VecDuplicate(tao->solution, &tao->stepdirection));
873   }
874 
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));
879 
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   }
887 
888   pdipm->nce = pdipm->ng + pdipm->nxfixed;
889   pdipm->Nce = pdipm->Ng + pdipm->Nxfixed;
890 
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   }
898 
899   pdipm->nci = pdipm->nh + pdipm->nxlb + pdipm->nxub + 2 * pdipm->nxbox;
900   pdipm->Nci = pdipm->Nh + pdipm->Nxlb + pdipm->Nxub + 2 * pdipm->Nxbox;
901 
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;
905 
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;
910 
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));
916 
917   PetscCall(VecCreate(comm, &pdipm->ci));
918   PetscCall(VecSetSizes(pdipm->ci, pdipm->nci, pdipm->Nci));
919   PetscCall(VecSetFromOptions(pdipm->ci));
920 
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));
925 
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));
930 
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));
933 
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));
937 
938     PetscCall(VecCreate(comm, &pdipm->lambdae_xfixed));
939     PetscCall(VecSetSizes(pdipm->lambdae_xfixed, pdipm->nxfixed, PETSC_DECIDE));
940     PetscCall(VecSetFromOptions(pdipm->lambdae_xfixed));
941   }
942 
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));
946 
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   }
950 
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));
956 
957   PetscCall(VecRestoreArrayRead(pdipm->X, &Xarr));
958 
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));
968 
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   }
980 
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));
987 
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   }
1000 
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   }
1012 
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   }
1026 
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)); */
1030 
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;
1036 
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   }
1040 
1041   /* (7) Gather offsets from all processes */
1042   PetscCall(PetscMalloc1(size, &pdipm->nce_all));
1043 
1044   /* Get rstart of KKT matrix */
1045   PetscCallMPI(MPI_Scan(&pdipm->n, &rstart, 1, MPIU_INT, MPI_SUM, comm));
1046   rstart -= pdipm->n;
1047 
1048   PetscCallMPI(MPI_Allgather(&pdipm->nce, 1, MPIU_INT, pdipm->nce_all, 1, MPIU_INT, comm));
1049 
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));
1054 
1055   PetscCall(MatGetOwnershipRanges(tao->hessian, &rranges));
1056   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &cranges));
1057 
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   }
1066 
1067   /* Count dnz,onz for preallocation of KKT matrix */
1068   nce_all = pdipm->nce_all;
1069 
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));
1072 
1073   MatPreallocateBegin(comm, pdipm->n, pdipm->n, dnz, onz);
1074 
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));
1078 
1079   /* Insert tao->hessian */
1080   PetscCall(MatGetOwnershipRange(tao->hessian, &rjstart, NULL));
1081   for (i = 0; i < pdipm->nx; i++) {
1082     row = rstart + i;
1083 
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));
1092 
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     }
1107 
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     }
1122 
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     }
1137 
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   }
1151 
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;
1157 
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;
1173 
1174       PetscCall(MatGetRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, NULL));
1175       PetscCheck(nc == 1, PETSC_COMM_SELF, PETSC_ERR_SUP, "nc != 1");
1176 
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   }
1185 
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;
1191 
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   }
1208 
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;
1213 
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   }
1227 
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   }
1235 
1236   /* diagonal entry */
1237   for (i = 0; i < pdipm->n; i++) dnz[i]++; /* diagonal entry */
1238 
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;
1247 
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   }
1259 
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;
1265 
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   }
1277 
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;
1282 
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));
1292 
1293     col = rstart + pdipm->off_z + pdipm->nh + i;
1294     PetscCall(MatSetValue(J, row, col, 1, INSERT_VALUES));
1295   }
1296 
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   }
1302 
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   }
1309 
1310   if (pdipm->Nxfixed) PetscCall(MatDestroy(&Jce_xfixed_trans));
1311   PetscCall(MatDestroy(&Jci_xb_trans));
1312   PetscCall(PetscFree3(ng_all, nh_all, Jranges));
1313 
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));
1317 
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));
1327 
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));
1336 
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 }
1355 
1356 static PetscErrorCode TaoDestroy_PDIPM(Tao tao)
1357 {
1358   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;
1359 
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] */
1367 
1368   /* work vectors */
1369   PetscCall(VecDestroy(&pdipm->lambdae_xfixed));
1370   PetscCall(VecDestroy(&pdipm->lambdai_xb));
1371 
1372   /* Legrangian equality and inequality Vec */
1373   PetscCall(VecDestroy(&pdipm->ce)); /* Vec of equality constraints */
1374   PetscCall(VecDestroy(&pdipm->ci)); /* Vec of inequality constraints */
1375 
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));
1380 
1381   /* Index Sets */
1382   if (pdipm->Nxub) { PetscCall(ISDestroy(&pdipm->isxub)); /* Finite upper bound only -inf < x < ub */ }
1383 
1384   if (pdipm->Nxlb) { PetscCall(ISDestroy(&pdipm->isxlb)); /* Finite lower bound only  lb <= x < inf */ }
1385 
1386   if (pdipm->Nxfixed) { PetscCall(ISDestroy(&pdipm->isxfixed)); /* Fixed variables         lb =  x = ub */ }
1387 
1388   if (pdipm->Nxbox) { PetscCall(ISDestroy(&pdipm->isxbox)); /* Boxed variables         lb <= x <= ub */ }
1389 
1390   if (pdipm->Nxfree) { PetscCall(ISDestroy(&pdipm->isxfree)); /* Free variables        -inf <= x <= inf */ }
1391 
1392   if (pdipm->solve_reduced_kkt) {
1393     PetscCall(ISDestroy(&pdipm->is1));
1394     PetscCall(ISDestroy(&pdipm->is2));
1395   }
1396 
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));
1402 
1403   /* Destroy pdipm */
1404   PetscCall(PetscFree(tao->data)); /* Holding locations of pdipm */
1405 
1406   /* Destroy Dual */
1407   PetscCall(VecDestroy(&tao->DE)); /* equality dual */
1408   PetscCall(VecDestroy(&tao->DI)); /* dinequality dual */
1409   PetscFunctionReturn(PETSC_SUCCESS);
1410 }
1411 
1412 static PetscErrorCode TaoSetFromOptions_PDIPM(Tao tao, PetscOptionItems *PetscOptionsObject)
1413 {
1414   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;
1415 
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 }
1427 
1428 /*MC
1429   TAOPDIPM - Barrier-based primal-dual interior point algorithm for generally constrained optimization.
1430 
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
1437 
1438   Level: beginner
1439 
1440 .seealso: `TAOPDIPM`, `Tao`, `TaoType`
1441 M*/
1442 
1443 PETSC_EXTERN PetscErrorCode TaoCreate_PDIPM(Tao tao)
1444 {
1445   TAO_PDIPM *pdipm;
1446 
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;
1453 
1454   PetscCall(PetscNew(&pdipm));
1455   tao->data = (void *)pdipm;
1456 
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;
1463 
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;
1469 
1470   pdipm->deltaw     = 0.0;
1471   pdipm->lastdeltaw = 3 * 1.e-4;
1472   pdipm->deltac     = 0.0;
1473   pdipm->kkt_pd     = PETSC_FALSE;
1474 
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;
1479 
1480   /* Override default settings (unless already changed) */
1481   if (!tao->max_it_changed) tao->max_it = 200;
1482   if (!tao->max_funcs_changed) tao->max_funcs = 500;
1483 
1484   PetscCall(SNESCreate(((PetscObject)tao)->comm, &pdipm->snes));
1485   PetscCall(SNESSetOptionsPrefix(pdipm->snes, tao->hdr.prefix));
1486   PetscCall(SNESGetKSP(pdipm->snes, &tao->ksp));
1487   PetscCall(PetscObjectReference((PetscObject)tao->ksp));
1488   PetscCall(KSPSetApplicationContext(tao->ksp, (void *)tao));
1489   PetscFunctionReturn(PETSC_SUCCESS);
1490 }
1491