xref: /petsc/src/tao/constrained/impls/ipm/pdipm.c (revision 5a884c48ab0c46bab83cd9bb8710f380fa6d8bcf)
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 */
TaoPDIPMEvaluateFunctionsAndJacobians(Tao tao,Vec x)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 */
TaoPDIPMUpdateConstraints(Tao tao,Vec x)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(PetscArraycpy(carr, garr, pdipm->ng));
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(PetscArraycpy(carr, harr, pdipm->nh));
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 */
TaoPDIPMSetUpBounds(Tao tao)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   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];
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 */
TaoPDIPMInitializeSolution(Tao tao)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(PetscArraycpy(Xarr, xarr, pdipm->nx));
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 */
TaoSNESJacobian_PDIPM(SNES snes,Vec X,Mat J,Mat Jpre,PetscCtx ctx)274 static PetscErrorCode TaoSNESJacobian_PDIPM(SNES snes, Vec X, Mat J, Mat Jpre, PetscCtx 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 
284   PetscFunctionBegin;
285   PetscCall(MatGetOwnershipRanges(Jpre, &Jranges));
286   PetscCall(MatGetOwnershipRange(Jpre, &Jrstart, NULL));
287   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &rranges));
288   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &cranges));
289 
290   PetscCall(VecGetArrayRead(X, &Xarr));
291 
292   /* (1) insert Z and Ci to the 4th block of Jpre -- overwrite existing values */
293   if (pdipm->solve_symmetric_kkt) { /* 1 for eq 17 revised pdipm doc 0 for eq 18 (symmetric KKT) */
294     vals[0] = 1.0;
295     for (i = 0; i < pdipm->nci; i++) {
296       row     = Jrstart + pdipm->off_z + i;
297       cols[0] = Jrstart + pdipm->off_lambdai + i;
298       cols[1] = row;
299       vals[1] = Xarr[pdipm->off_lambdai + i] / Xarr[pdipm->off_z + i];
300       PetscCall(MatSetValues(Jpre, 1, &row, 2, cols, vals, INSERT_VALUES));
301     }
302   } else {
303     for (i = 0; i < pdipm->nci; i++) {
304       row     = Jrstart + pdipm->off_z + i;
305       cols[0] = Jrstart + pdipm->off_lambdai + i;
306       cols[1] = row;
307       vals[0] = Xarr[pdipm->off_z + i];
308       vals[1] = Xarr[pdipm->off_lambdai + i];
309       PetscCall(MatSetValues(Jpre, 1, &row, 2, cols, vals, INSERT_VALUES));
310     }
311   }
312 
313   /* (2) insert 2nd row block of Jpre: [ grad g, 0, 0, 0] */
314   if (pdipm->Ng) {
315     PetscCall(MatGetOwnershipRange(tao->jacobian_equality, &rjstart, NULL));
316     for (i = 0; i < pdipm->ng; i++) {
317       row = Jrstart + pdipm->off_lambdae + i;
318 
319       PetscCall(MatGetRow(tao->jacobian_equality, i + rjstart, &nc, &aj, &aa));
320       proc = 0;
321       for (j = 0; j < nc; j++) {
322         while (aj[j] >= cranges[proc + 1]) proc++;
323         cols[0] = aj[j] - cranges[proc] + Jranges[proc];
324         PetscCall(MatSetValue(Jpre, row, cols[0], aa[j], INSERT_VALUES));
325       }
326       PetscCall(MatRestoreRow(tao->jacobian_equality, i + rjstart, &nc, &aj, &aa));
327       if (pdipm->kkt_pd) {
328         /* add shift \delta_c */
329         PetscCall(MatSetValue(Jpre, row, row, -pdipm->deltac, INSERT_VALUES));
330       }
331     }
332   }
333 
334   /* (3) insert 3rd row block of Jpre: [ -grad h, 0, deltac, I] */
335   if (pdipm->Nh) {
336     PetscCall(MatGetOwnershipRange(tao->jacobian_inequality, &rjstart, NULL));
337     for (i = 0; i < pdipm->nh; i++) {
338       row = Jrstart + pdipm->off_lambdai + i;
339       PetscCall(MatGetRow(tao->jacobian_inequality, i + rjstart, &nc, &aj, &aa));
340       proc = 0;
341       for (j = 0; j < nc; j++) {
342         while (aj[j] >= cranges[proc + 1]) proc++;
343         cols[0] = aj[j] - cranges[proc] + Jranges[proc];
344         PetscCall(MatSetValue(Jpre, row, cols[0], -aa[j], INSERT_VALUES));
345       }
346       PetscCall(MatRestoreRow(tao->jacobian_inequality, i + rjstart, &nc, &aj, &aa));
347       if (pdipm->kkt_pd) {
348         /* add shift \delta_c */
349         PetscCall(MatSetValue(Jpre, row, row, -pdipm->deltac, INSERT_VALUES));
350       }
351     }
352   }
353 
354   /* (4) insert 1st row block of Jpre: [Wxx, grad g', -grad h', 0] */
355   if (pdipm->Ng) { /* grad g' */
356     PetscCall(MatTranspose(tao->jacobian_equality, MAT_REUSE_MATRIX, &pdipm->jac_equality_trans));
357   }
358   if (pdipm->Nh) { /* grad h' */
359     PetscCall(MatTranspose(tao->jacobian_inequality, MAT_REUSE_MATRIX, &pdipm->jac_inequality_trans));
360   }
361 
362   PetscCall(VecPlaceArray(pdipm->x, Xarr));
363   PetscCall(TaoComputeHessian(tao, pdipm->x, tao->hessian, tao->hessian_pre));
364   PetscCall(VecResetArray(pdipm->x));
365 
366   PetscCall(MatGetOwnershipRange(tao->hessian, &rjstart, NULL));
367   for (i = 0; i < pdipm->nx; i++) {
368     row = Jrstart + i;
369 
370     /* insert Wxx = fxx + ... -- provided by user */
371     PetscCall(MatGetRow(tao->hessian, i + rjstart, &nc, &aj, &aa));
372     proc = 0;
373     for (j = 0; j < nc; j++) {
374       while (aj[j] >= cranges[proc + 1]) proc++;
375       cols[0] = aj[j] - cranges[proc] + Jranges[proc];
376       if (row == cols[0] && pdipm->kkt_pd) {
377         /* add shift deltaw to Wxx component */
378         PetscCall(MatSetValue(Jpre, row, cols[0], aa[j] + pdipm->deltaw, INSERT_VALUES));
379       } else {
380         PetscCall(MatSetValue(Jpre, row, cols[0], aa[j], INSERT_VALUES));
381       }
382     }
383     PetscCall(MatRestoreRow(tao->hessian, i + rjstart, &nc, &aj, &aa));
384 
385     /* insert grad g' */
386     if (pdipm->ng) {
387       PetscCall(MatGetRow(pdipm->jac_equality_trans, i + rjstart, &nc, &aj, &aa));
388       PetscCall(MatGetOwnershipRanges(tao->jacobian_equality, &ranges));
389       proc = 0;
390       for (j = 0; j < nc; j++) {
391         /* find row ownership of */
392         while (aj[j] >= ranges[proc + 1]) proc++;
393         nx_all  = rranges[proc + 1] - rranges[proc];
394         cols[0] = aj[j] - ranges[proc] + Jranges[proc] + nx_all;
395         PetscCall(MatSetValue(Jpre, row, cols[0], aa[j], INSERT_VALUES));
396       }
397       PetscCall(MatRestoreRow(pdipm->jac_equality_trans, i + rjstart, &nc, &aj, &aa));
398     }
399 
400     /* insert -grad h' */
401     if (pdipm->nh) {
402       PetscCall(MatGetRow(pdipm->jac_inequality_trans, i + rjstart, &nc, &aj, &aa));
403       PetscCall(MatGetOwnershipRanges(tao->jacobian_inequality, &ranges));
404       proc = 0;
405       for (j = 0; j < nc; j++) {
406         /* find row ownership of */
407         while (aj[j] >= ranges[proc + 1]) proc++;
408         nx_all  = rranges[proc + 1] - rranges[proc];
409         cols[0] = aj[j] - ranges[proc] + Jranges[proc] + nx_all + nce_all[proc];
410         PetscCall(MatSetValue(Jpre, row, cols[0], -aa[j], INSERT_VALUES));
411       }
412       PetscCall(MatRestoreRow(pdipm->jac_inequality_trans, i + rjstart, &nc, &aj, &aa));
413     }
414   }
415   PetscCall(VecRestoreArrayRead(X, &Xarr));
416 
417   /* (6) assemble Jpre and J */
418   PetscCall(MatAssemblyBegin(Jpre, MAT_FINAL_ASSEMBLY));
419   PetscCall(MatAssemblyEnd(Jpre, MAT_FINAL_ASSEMBLY));
420 
421   if (J != Jpre) {
422     PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
423     PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
424   }
425   PetscFunctionReturn(PETSC_SUCCESS);
426 }
427 
428 /*
429    TaoSnesFunction_PDIPM - Evaluate KKT function at X
430 
431    Input Parameter:
432    snes - SNES context
433    X - KKT Vector
434    *ctx - pdipm
435 
436    Output Parameter:
437    F - Updated Lagrangian vector
438 */
TaoSNESFunction_PDIPM(SNES snes,Vec X,Vec F,PetscCtx ctx)439 static PetscErrorCode TaoSNESFunction_PDIPM(SNES snes, Vec X, Vec F, PetscCtx ctx)
440 {
441   Tao                tao   = (Tao)ctx;
442   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
443   PetscScalar       *Farr;
444   Vec                x, L1;
445   PetscInt           i;
446   const PetscScalar *Xarr, *carr, *zarr, *larr;
447 
448   PetscFunctionBegin;
449   PetscCall(VecSet(F, 0.0));
450 
451   PetscCall(VecGetArrayRead(X, &Xarr));
452   PetscCall(VecGetArrayWrite(F, &Farr));
453 
454   /* (0) Evaluate f, fx, gradG, gradH at X.x Note: pdipm->x is not changed below */
455   x = pdipm->x;
456   PetscCall(VecPlaceArray(x, Xarr));
457   PetscCall(TaoPDIPMEvaluateFunctionsAndJacobians(tao, x));
458 
459   /* Update ce, ci, and Jci at X.x */
460   PetscCall(TaoPDIPMUpdateConstraints(tao, x));
461   PetscCall(VecResetArray(x));
462 
463   /* (1) L1 = fx + (gradG'*DE + Jce_xfixed'*lambdae_xfixed) - (gradH'*DI + Jci_xb'*lambdai_xb) */
464   L1 = pdipm->x;
465   PetscCall(VecPlaceArray(L1, Farr)); /* L1 = 0.0 */
466   if (pdipm->Nci) {
467     if (pdipm->Nh) {
468       /* L1 += gradH'*DI. Note: tao->DI is not changed below */
469       PetscCall(VecPlaceArray(tao->DI, Xarr + pdipm->off_lambdai));
470       PetscCall(MatMultTransposeAdd(tao->jacobian_inequality, tao->DI, L1, L1));
471       PetscCall(VecResetArray(tao->DI));
472     }
473 
474     /* L1 += Jci_xb'*lambdai_xb */
475     PetscCall(VecPlaceArray(pdipm->lambdai_xb, Xarr + pdipm->off_lambdai + pdipm->nh));
476     PetscCall(MatMultTransposeAdd(pdipm->Jci_xb, pdipm->lambdai_xb, L1, L1));
477     PetscCall(VecResetArray(pdipm->lambdai_xb));
478 
479     /* L1 = - (gradH'*DI + Jci_xb'*lambdai_xb) */
480     PetscCall(VecScale(L1, -1.0));
481   }
482 
483   /* L1 += fx */
484   PetscCall(VecAXPY(L1, 1.0, tao->gradient));
485 
486   if (pdipm->Nce) {
487     if (pdipm->Ng) {
488       /* L1 += gradG'*DE. Note: tao->DE is not changed below */
489       PetscCall(VecPlaceArray(tao->DE, Xarr + pdipm->off_lambdae));
490       PetscCall(MatMultTransposeAdd(tao->jacobian_equality, tao->DE, L1, L1));
491       PetscCall(VecResetArray(tao->DE));
492     }
493     if (pdipm->Nxfixed) {
494       /* L1 += Jce_xfixed'*lambdae_xfixed */
495       PetscCall(VecPlaceArray(pdipm->lambdae_xfixed, Xarr + pdipm->off_lambdae + pdipm->ng));
496       PetscCall(MatMultTransposeAdd(pdipm->Jce_xfixed, pdipm->lambdae_xfixed, L1, L1));
497       PetscCall(VecResetArray(pdipm->lambdae_xfixed));
498     }
499   }
500   PetscCall(VecResetArray(L1));
501 
502   /* (2) L2 = ce(x) */
503   if (pdipm->Nce) {
504     PetscCall(VecGetArrayRead(pdipm->ce, &carr));
505     for (i = 0; i < pdipm->nce; i++) Farr[pdipm->off_lambdae + i] = carr[i];
506     PetscCall(VecRestoreArrayRead(pdipm->ce, &carr));
507   }
508 
509   if (pdipm->Nci) {
510     if (pdipm->solve_symmetric_kkt) {
511       /* (3) L3 = z - ci(x);
512          (4) L4 = Lambdai * e - mu/z *e  */
513       PetscCall(VecGetArrayRead(pdipm->ci, &carr));
514       larr = Xarr + pdipm->off_lambdai;
515       zarr = Xarr + pdipm->off_z;
516       for (i = 0; i < pdipm->nci; i++) {
517         Farr[pdipm->off_lambdai + i] = zarr[i] - carr[i];
518         Farr[pdipm->off_z + i]       = larr[i] - pdipm->mu / zarr[i];
519       }
520       PetscCall(VecRestoreArrayRead(pdipm->ci, &carr));
521     } else {
522       /* (3) L3 = z - ci(x);
523          (4) L4 = Z * Lambdai * e - mu * e  */
524       PetscCall(VecGetArrayRead(pdipm->ci, &carr));
525       larr = Xarr + pdipm->off_lambdai;
526       zarr = Xarr + pdipm->off_z;
527       for (i = 0; i < pdipm->nci; i++) {
528         Farr[pdipm->off_lambdai + i] = zarr[i] - carr[i];
529         Farr[pdipm->off_z + i]       = zarr[i] * larr[i] - pdipm->mu;
530       }
531       PetscCall(VecRestoreArrayRead(pdipm->ci, &carr));
532     }
533   }
534 
535   PetscCall(VecRestoreArrayRead(X, &Xarr));
536   PetscCall(VecRestoreArrayWrite(F, &Farr));
537   PetscFunctionReturn(PETSC_SUCCESS);
538 }
539 
540 /*
541   Evaluate F(X); then update tao->gnorm0, tao->step = mu,
542   tao->residual = norm2(F_x,F_z) and tao->cnorm = norm2(F_ce,F_ci).
543 */
TaoSNESFunction_PDIPM_residual(SNES snes,Vec X,Vec F,PetscCtx ctx)544 static PetscErrorCode TaoSNESFunction_PDIPM_residual(SNES snes, Vec X, Vec F, PetscCtx ctx)
545 {
546   Tao                tao   = (Tao)ctx;
547   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
548   PetscScalar       *Farr, *tmparr;
549   Vec                L1;
550   PetscInt           i;
551   PetscReal          res[2], cnorm[2];
552   const PetscScalar *Xarr = NULL;
553 
554   PetscFunctionBegin;
555   PetscCall(TaoSNESFunction_PDIPM(snes, X, F, (void *)tao));
556   PetscCall(VecGetArrayWrite(F, &Farr));
557   PetscCall(VecGetArrayRead(X, &Xarr));
558 
559   /* compute res[0] = norm2(F_x) */
560   L1 = pdipm->x;
561   PetscCall(VecPlaceArray(L1, Farr));
562   PetscCall(VecNorm(L1, NORM_2, &res[0]));
563   PetscCall(VecResetArray(L1));
564 
565   /* compute res[1] = norm2(F_z), cnorm[1] = norm2(F_ci) */
566   if (pdipm->z) {
567     if (pdipm->solve_symmetric_kkt) {
568       PetscCall(VecPlaceArray(pdipm->z, Farr + pdipm->off_z));
569       if (pdipm->Nci) {
570         PetscCall(VecGetArrayWrite(pdipm->z, &tmparr));
571         for (i = 0; i < pdipm->nci; i++) tmparr[i] *= Xarr[pdipm->off_z + i];
572         PetscCall(VecRestoreArrayWrite(pdipm->z, &tmparr));
573       }
574 
575       PetscCall(VecNorm(pdipm->z, NORM_2, &res[1]));
576 
577       if (pdipm->Nci) {
578         PetscCall(VecGetArrayWrite(pdipm->z, &tmparr));
579         for (i = 0; i < pdipm->nci; i++) tmparr[i] /= Xarr[pdipm->off_z + i];
580         PetscCall(VecRestoreArrayWrite(pdipm->z, &tmparr));
581       }
582       PetscCall(VecResetArray(pdipm->z));
583     } else { /* !solve_symmetric_kkt */
584       PetscCall(VecPlaceArray(pdipm->z, Farr + pdipm->off_z));
585       PetscCall(VecNorm(pdipm->z, NORM_2, &res[1]));
586       PetscCall(VecResetArray(pdipm->z));
587     }
588 
589     PetscCall(VecPlaceArray(pdipm->ci, Farr + pdipm->off_lambdai));
590     PetscCall(VecNorm(pdipm->ci, NORM_2, &cnorm[1]));
591     PetscCall(VecResetArray(pdipm->ci));
592   } else {
593     res[1]   = 0.0;
594     cnorm[1] = 0.0;
595   }
596 
597   /* compute cnorm[0] = norm2(F_ce) */
598   if (pdipm->Nce) {
599     PetscCall(VecPlaceArray(pdipm->ce, Farr + pdipm->off_lambdae));
600     PetscCall(VecNorm(pdipm->ce, NORM_2, &cnorm[0]));
601     PetscCall(VecResetArray(pdipm->ce));
602   } else cnorm[0] = 0.0;
603 
604   PetscCall(VecRestoreArrayWrite(F, &Farr));
605   PetscCall(VecRestoreArrayRead(X, &Xarr));
606 
607   tao->gnorm0   = tao->residual;
608   tao->residual = PetscSqrtReal(res[0] * res[0] + res[1] * res[1]);
609   tao->cnorm    = PetscSqrtReal(cnorm[0] * cnorm[0] + cnorm[1] * cnorm[1]);
610   tao->step     = pdipm->mu;
611   PetscFunctionReturn(PETSC_SUCCESS);
612 }
613 
614 /*
615   PCPostSetup_PDIPM -- called when the KKT matrix is Cholesky factored for the preconditioner. Checks the inertia of Cholesky factor of the KKT matrix.
616   If it does not match the numbers of prime and dual variables, add shifts to the KKT matrix.
617 */
PCPostSetUp_PDIPM(PC pc)618 static PetscErrorCode PCPostSetUp_PDIPM(PC pc)
619 {
620   Tao        tao;
621   TAO_PDIPM *pdipm;
622   Vec        X;
623   SNES       snes;
624   KSP        ksp;
625   Mat        Factor;
626   PetscBool  isCHOL;
627   PetscInt   nneg, nzero, npos;
628 
629   PetscFunctionBegin;
630   PetscCall(PCGetApplicationContext(pc, &tao));
631   pdipm = (TAO_PDIPM *)tao->data;
632   X     = pdipm->X;
633   snes  = pdipm->snes;
634 
635   /* Get the inertia of Cholesky factor */
636   PetscCall(SNESGetKSP(snes, &ksp));
637   PetscCall(KSPGetPC(ksp, &pc));
638   PetscCall(PetscObjectTypeCompare((PetscObject)pc, PCCHOLESKY, &isCHOL));
639   if (!isCHOL) PetscFunctionReturn(PETSC_SUCCESS);
640 
641   PetscCall(PCFactorGetMatrix(pc, &Factor));
642   PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));
643 
644   if (npos < pdipm->Nx + pdipm->Nci) {
645     pdipm->deltaw = PetscMax(pdipm->lastdeltaw / 3, 1.e-4 * PETSC_MACHINE_EPSILON);
646     PetscCall(PetscInfo(tao, "Test reduced deltaw=%g; previous MatInertia: nneg %" PetscInt_FMT ", nzero %" PetscInt_FMT ", npos %" PetscInt_FMT "(<%" PetscInt_FMT ")\n", (double)pdipm->deltaw, nneg, nzero, npos, pdipm->Nx + pdipm->Nci));
647     PetscCall(TaoSNESJacobian_PDIPM(snes, X, pdipm->K, pdipm->K, tao));
648     PetscCall(PCSetPostSetUp(pc, NULL));
649     PetscCall(PCSetUp(pc));
650     PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));
651 
652     if (npos < pdipm->Nx + pdipm->Nci) {
653       pdipm->deltaw = pdipm->lastdeltaw;                                           /* in case reduction update does not help, this prevents that step from impacting increasing update */
654       while (npos < pdipm->Nx + pdipm->Nci && pdipm->deltaw <= 1. / PETSC_SMALL) { /* increase deltaw */
655         PetscCall(PetscInfo(tao, "  deltaw=%g fails, MatInertia: nneg %" PetscInt_FMT ", nzero %" PetscInt_FMT ", npos %" PetscInt_FMT "(<%" PetscInt_FMT ")\n", (double)pdipm->deltaw, nneg, nzero, npos, pdipm->Nx + pdipm->Nci));
656         pdipm->deltaw = PetscMin(8 * pdipm->deltaw, PetscPowReal(10, 20));
657         PetscCall(TaoSNESJacobian_PDIPM(snes, X, pdipm->K, pdipm->K, tao));
658         PetscCall(PCSetUp(pc));
659         PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));
660       }
661 
662       PetscCheck(pdipm->deltaw < 1. / PETSC_SMALL, PetscObjectComm((PetscObject)tao), PETSC_ERR_CONV_FAILED, "Reached maximum delta w will not converge, try different initial x0");
663 
664       PetscCall(PetscInfo(tao, "Updated deltaw %g\n", (double)pdipm->deltaw));
665       pdipm->lastdeltaw = pdipm->deltaw;
666       pdipm->deltaw     = 0.0;
667     }
668   }
669 
670   if (nzero) { /* Jacobian is singular */
671     if (pdipm->deltac == 0.0) {
672       pdipm->deltac = PETSC_SQRT_MACHINE_EPSILON;
673     } else {
674       pdipm->deltac = pdipm->deltac * PetscPowReal(pdipm->mu, .25);
675     }
676     PetscCall(PetscInfo(tao, "Updated deltac=%g, MatInertia: nneg %" PetscInt_FMT ", nzero %" PetscInt_FMT "(!=0), npos %" PetscInt_FMT "\n", (double)pdipm->deltac, nneg, nzero, npos));
677     PetscCall(TaoSNESJacobian_PDIPM(snes, X, pdipm->K, pdipm->K, tao));
678     PetscCall(PCSetPostSetUp(pc, NULL));
679     PetscCall(PCSetUp(pc));
680     PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));
681   }
682   PetscCall(PCSetPostSetUp(pc, PCPostSetUp_PDIPM));
683   PetscFunctionReturn(PETSC_SUCCESS);
684 }
685 
686 /*
687    SNESLineSearch_PDIPM - Custom line search used with PDIPM.
688 
689    Collective
690 
691    Notes:
692    This routine employs a simple backtracking line-search to keep
693    the slack variables (z) and inequality constraints Lagrange multipliers
694    (lambdai) positive, i.e., z,lambdai >=0. It does this by calculating scalars
695    alpha_p and alpha_d to keep z,lambdai non-negative. The decision (x), and the
696    slack variables are updated as X = X - alpha_d*dx. The constraint multipliers
697    are updated as Lambdai = Lambdai + alpha_p*dLambdai. The barrier parameter mu
698    is also updated as mu = mu + z'lambdai/Nci
699 */
SNESLineSearch_PDIPM(SNESLineSearch linesearch,PetscCtx ctx)700 static PetscErrorCode SNESLineSearch_PDIPM(SNESLineSearch linesearch, PetscCtx ctx)
701 {
702   Tao                tao   = (Tao)ctx;
703   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
704   SNES               snes;
705   Vec                X, F, Y;
706   PetscInt           i, iter;
707   PetscReal          alpha_p = 1.0, alpha_d = 1.0, alpha[4];
708   PetscScalar       *Xarr, *z, *lambdai, dot, *taosolarr;
709   const PetscScalar *dXarr, *dz, *dlambdai;
710 
711   PetscFunctionBegin;
712   PetscCall(SNESLineSearchGetSNES(linesearch, &snes));
713   PetscCall(SNESGetIterationNumber(snes, &iter));
714 
715   PetscCall(SNESLineSearchSetReason(linesearch, SNES_LINESEARCH_SUCCEEDED));
716   PetscCall(SNESLineSearchGetVecs(linesearch, &X, &F, &Y, NULL, NULL));
717 
718   PetscCall(VecGetArrayWrite(X, &Xarr));
719   PetscCall(VecGetArrayRead(Y, &dXarr));
720   z  = Xarr + pdipm->off_z;
721   dz = dXarr + pdipm->off_z;
722   for (i = 0; i < pdipm->nci; i++) {
723     if (z[i] - dz[i] < 0.0) alpha_p = PetscMin(alpha_p, 0.9999 * z[i] / dz[i]);
724   }
725 
726   lambdai  = Xarr + pdipm->off_lambdai;
727   dlambdai = dXarr + pdipm->off_lambdai;
728 
729   for (i = 0; i < pdipm->nci; i++) {
730     if (lambdai[i] - dlambdai[i] < 0.0) alpha_d = PetscMin(0.9999 * lambdai[i] / dlambdai[i], alpha_d);
731   }
732 
733   alpha[0] = alpha_p;
734   alpha[1] = alpha_d;
735   PetscCall(VecRestoreArrayRead(Y, &dXarr));
736   PetscCall(VecRestoreArrayWrite(X, &Xarr));
737 
738   /* alpha = min(alpha) over all processes */
739   PetscCallMPI(MPIU_Allreduce(alpha, alpha + 2, 2, MPIU_REAL, MPIU_MIN, PetscObjectComm((PetscObject)tao)));
740 
741   alpha_p = alpha[2];
742   alpha_d = alpha[3];
743 
744   /* X = X - alpha * Y */
745   PetscCall(VecGetArrayWrite(X, &Xarr));
746   PetscCall(VecGetArrayRead(Y, &dXarr));
747   for (i = 0; i < pdipm->nx; i++) Xarr[i] -= alpha_p * dXarr[i];
748   for (i = 0; i < pdipm->nce; i++) Xarr[i + pdipm->off_lambdae] -= alpha_d * dXarr[i + pdipm->off_lambdae];
749 
750   for (i = 0; i < pdipm->nci; i++) {
751     Xarr[i + pdipm->off_lambdai] -= alpha_d * dXarr[i + pdipm->off_lambdai];
752     Xarr[i + pdipm->off_z] -= alpha_p * dXarr[i + pdipm->off_z];
753   }
754   PetscCall(VecGetArrayWrite(tao->solution, &taosolarr));
755   PetscCall(PetscArraycpy(taosolarr, Xarr, pdipm->nx));
756   PetscCall(VecRestoreArrayWrite(tao->solution, &taosolarr));
757 
758   PetscCall(VecRestoreArrayWrite(X, &Xarr));
759   PetscCall(VecRestoreArrayRead(Y, &dXarr));
760 
761   /* Update mu = mu_update_factor * dot(z,lambdai)/pdipm->nci at updated X */
762   if (pdipm->z) PetscCall(VecDot(pdipm->z, pdipm->lambdai, &dot));
763   else dot = 0.0;
764 
765   /* if (PetscAbsReal(pdipm->gradL) < 0.9*pdipm->mu)  */
766   pdipm->mu = pdipm->mu_update_factor * dot / pdipm->Nci;
767 
768   /* Update F; get tao->residual and tao->cnorm */
769   PetscCall(TaoSNESFunction_PDIPM_residual(snes, X, F, (void *)tao));
770 
771   tao->niter++;
772   PetscCall(TaoLogConvergenceHistory(tao, pdipm->obj, tao->residual, tao->cnorm, tao->niter));
773   PetscCall(TaoMonitor(tao, tao->niter, pdipm->obj, tao->residual, tao->cnorm, pdipm->mu));
774 
775   PetscUseTypeMethod(tao, convergencetest, tao->cnvP);
776   if (tao->reason) PetscCall(SNESSetConvergedReason(snes, SNES_CONVERGED_FNORM_ABS));
777   PetscFunctionReturn(PETSC_SUCCESS);
778 }
779 
TaoSolve_PDIPM(Tao tao)780 static PetscErrorCode TaoSolve_PDIPM(Tao tao)
781 {
782   TAO_PDIPM     *pdipm = (TAO_PDIPM *)tao->data;
783   SNESLineSearch linesearch; /* SNESLineSearch context */
784   Vec            dummy;
785 
786   PetscFunctionBegin;
787   PetscCheck(tao->constraints_equality || tao->constraints_inequality, PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_NULL, "Equality and inequality constraints are not set. Either set them or switch to a different algorithm");
788 
789   /* Initialize all variables */
790   PetscCall(TaoPDIPMInitializeSolution(tao));
791 
792   /* Set linesearch */
793   PetscCall(SNESGetLineSearch(pdipm->snes, &linesearch));
794   PetscCall(SNESLineSearchSetType(linesearch, SNESLINESEARCHSHELL));
795   PetscCall(SNESLineSearchShellSetApply(linesearch, SNESLineSearch_PDIPM, tao));
796   PetscCall(SNESLineSearchSetFromOptions(linesearch));
797 
798   tao->reason = TAO_CONTINUE_ITERATING;
799 
800   /* -tao_monitor for iteration 0 and check convergence */
801   PetscCall(VecDuplicate(pdipm->X, &dummy));
802   PetscCall(TaoSNESFunction_PDIPM_residual(pdipm->snes, pdipm->X, dummy, (void *)tao));
803 
804   PetscCall(TaoLogConvergenceHistory(tao, pdipm->obj, tao->residual, tao->cnorm, tao->niter));
805   PetscCall(TaoMonitor(tao, tao->niter, pdipm->obj, tao->residual, tao->cnorm, pdipm->mu));
806   PetscCall(VecDestroy(&dummy));
807   PetscUseTypeMethod(tao, convergencetest, tao->cnvP);
808   if (tao->reason) PetscCall(SNESSetConvergedReason(pdipm->snes, SNES_CONVERGED_FNORM_ABS));
809 
810   while (tao->reason == TAO_CONTINUE_ITERATING) {
811     SNESConvergedReason reason;
812     PetscCall(SNESSolve(pdipm->snes, NULL, pdipm->X));
813 
814     /* Check SNES convergence */
815     PetscCall(SNESGetConvergedReason(pdipm->snes, &reason));
816     if (reason < 0) PetscCall(PetscPrintf(PetscObjectComm((PetscObject)pdipm->snes), "SNES solve did not converged due to reason %s\n", SNESConvergedReasons[reason]));
817 
818     /* Check TAO convergence */
819     PetscCheck(!PetscIsInfOrNanReal(pdipm->obj), PETSC_COMM_SELF, PETSC_ERR_SUP, "User-provided compute function generated infinity or NaN");
820   }
821   PetscFunctionReturn(PETSC_SUCCESS);
822 }
823 
TaoView_PDIPM(Tao tao,PetscViewer viewer)824 static PetscErrorCode TaoView_PDIPM(Tao tao, PetscViewer viewer)
825 {
826   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;
827 
828   PetscFunctionBegin;
829   tao->constrained = PETSC_TRUE;
830   PetscCall(PetscViewerASCIIPushTab(viewer));
831   PetscCall(PetscViewerASCIIPrintf(viewer, "Number of prime=%" PetscInt_FMT ", Number of dual=%" PetscInt_FMT "\n", pdipm->Nx + pdipm->Nci, pdipm->Nce + pdipm->Nci));
832   if (pdipm->kkt_pd) PetscCall(PetscViewerASCIIPrintf(viewer, "KKT shifts deltaw=%g, deltac=%g\n", (double)pdipm->deltaw, (double)pdipm->deltac));
833   PetscCall(PetscViewerASCIIPopTab(viewer));
834   PetscFunctionReturn(PETSC_SUCCESS);
835 }
836 
TaoSetup_PDIPM(Tao tao)837 static PetscErrorCode TaoSetup_PDIPM(Tao tao)
838 {
839   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
840   MPI_Comm           comm;
841   PetscMPIInt        size;
842   PetscInt           row, col, Jcrstart, Jcrend, k, tmp, nc, proc, *nh_all, *ng_all;
843   PetscInt           offset, *xa, *xb, i, j, rstart, rend;
844   PetscScalar        one = 1.0, neg_one = -1.0;
845   const PetscInt    *cols, *rranges, *cranges, *aj, *ranges;
846   const PetscScalar *aa, *Xarr;
847   Mat                J;
848   Mat                Jce_xfixed_trans, Jci_xb_trans;
849   PetscInt          *dnz, *onz, rjstart, nx_all, *nce_all, *Jranges, cols1[2];
850 
851   PetscFunctionBegin;
852   PetscCall(PetscObjectGetComm((PetscObject)tao, &comm));
853   PetscCallMPI(MPI_Comm_size(comm, &size));
854 
855   /* (1) Setup Bounds and create Tao vectors */
856   PetscCall(TaoPDIPMSetUpBounds(tao));
857 
858   if (!tao->gradient) {
859     PetscCall(VecDuplicate(tao->solution, &tao->gradient));
860     PetscCall(VecDuplicate(tao->solution, &tao->stepdirection));
861   }
862 
863   /* (2) Get sizes */
864   /* Size of vector x - This is set by TaoSetSolution */
865   PetscCall(VecGetSize(tao->solution, &pdipm->Nx));
866   PetscCall(VecGetLocalSize(tao->solution, &pdipm->nx));
867 
868   /* Size of equality constraints and vectors */
869   if (tao->constraints_equality) {
870     PetscCall(VecGetSize(tao->constraints_equality, &pdipm->Ng));
871     PetscCall(VecGetLocalSize(tao->constraints_equality, &pdipm->ng));
872   } else {
873     pdipm->ng = pdipm->Ng = 0;
874   }
875 
876   pdipm->nce = pdipm->ng + pdipm->nxfixed;
877   pdipm->Nce = pdipm->Ng + pdipm->Nxfixed;
878 
879   /* Size of inequality constraints and vectors */
880   if (tao->constraints_inequality) {
881     PetscCall(VecGetSize(tao->constraints_inequality, &pdipm->Nh));
882     PetscCall(VecGetLocalSize(tao->constraints_inequality, &pdipm->nh));
883   } else {
884     pdipm->nh = pdipm->Nh = 0;
885   }
886 
887   pdipm->nci = pdipm->nh + pdipm->nxlb + pdipm->nxub + 2 * pdipm->nxbox;
888   pdipm->Nci = pdipm->Nh + pdipm->Nxlb + pdipm->Nxub + 2 * pdipm->Nxbox;
889 
890   /* Full size of the KKT system to be solved */
891   pdipm->n = pdipm->nx + pdipm->nce + 2 * pdipm->nci;
892   pdipm->N = pdipm->Nx + pdipm->Nce + 2 * pdipm->Nci;
893 
894   /* (3) Offsets for subvectors */
895   pdipm->off_lambdae = pdipm->nx;
896   pdipm->off_lambdai = pdipm->off_lambdae + pdipm->nce;
897   pdipm->off_z       = pdipm->off_lambdai + pdipm->nci;
898 
899   /* (4) Create vectors and subvectors */
900   /* Ce and Ci vectors */
901   PetscCall(VecCreate(comm, &pdipm->ce));
902   PetscCall(VecSetSizes(pdipm->ce, pdipm->nce, pdipm->Nce));
903   PetscCall(VecSetFromOptions(pdipm->ce));
904 
905   PetscCall(VecCreate(comm, &pdipm->ci));
906   PetscCall(VecSetSizes(pdipm->ci, pdipm->nci, pdipm->Nci));
907   PetscCall(VecSetFromOptions(pdipm->ci));
908 
909   /* X=[x; lambdae; lambdai; z] for the big KKT system */
910   PetscCall(VecCreate(comm, &pdipm->X));
911   PetscCall(VecSetSizes(pdipm->X, pdipm->n, pdipm->N));
912   PetscCall(VecSetFromOptions(pdipm->X));
913 
914   /* Subvectors; they share local arrays with X */
915   PetscCall(VecGetArrayRead(pdipm->X, &Xarr));
916   /* x shares local array with X.x */
917   if (pdipm->Nx) PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nx, pdipm->Nx, Xarr, &pdipm->x));
918 
919   /* lambdae shares local array with X.lambdae */
920   if (pdipm->Nce) PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nce, pdipm->Nce, Xarr + pdipm->off_lambdae, &pdipm->lambdae));
921 
922   /* tao->DE shares local array with X.lambdae_g */
923   if (pdipm->Ng) {
924     PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->ng, pdipm->Ng, Xarr + pdipm->off_lambdae, &tao->DE));
925 
926     PetscCall(VecCreate(comm, &pdipm->lambdae_xfixed));
927     PetscCall(VecSetSizes(pdipm->lambdae_xfixed, pdipm->nxfixed, PETSC_DECIDE));
928     PetscCall(VecSetFromOptions(pdipm->lambdae_xfixed));
929   }
930 
931   if (pdipm->Nci) {
932     /* lambdai shares local array with X.lambdai */
933     PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nci, pdipm->Nci, Xarr + pdipm->off_lambdai, &pdipm->lambdai));
934 
935     /* z for slack variables; it shares local array with X.z */
936     PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nci, pdipm->Nci, Xarr + pdipm->off_z, &pdipm->z));
937   }
938 
939   /* tao->DI which shares local array with X.lambdai_h */
940   if (pdipm->Nh) PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nh, pdipm->Nh, Xarr + pdipm->off_lambdai, &tao->DI));
941   PetscCall(VecCreate(comm, &pdipm->lambdai_xb));
942   PetscCall(VecSetSizes(pdipm->lambdai_xb, pdipm->nci - pdipm->nh, PETSC_DECIDE));
943   PetscCall(VecSetFromOptions(pdipm->lambdai_xb));
944 
945   PetscCall(VecRestoreArrayRead(pdipm->X, &Xarr));
946 
947   /* (5) Create Jacobians Jce_xfixed and Jci */
948   /* (5.1) PDIPM Jacobian of equality bounds cebound(x) = J_nxfixed */
949   if (pdipm->Nxfixed) {
950     /* Create Jce_xfixed */
951     PetscCall(MatCreate(comm, &pdipm->Jce_xfixed));
952     PetscCall(MatSetSizes(pdipm->Jce_xfixed, pdipm->nxfixed, pdipm->nx, PETSC_DECIDE, pdipm->Nx));
953     PetscCall(MatSetFromOptions(pdipm->Jce_xfixed));
954     PetscCall(MatSeqAIJSetPreallocation(pdipm->Jce_xfixed, 1, NULL));
955     PetscCall(MatMPIAIJSetPreallocation(pdipm->Jce_xfixed, 1, NULL, 1, NULL));
956 
957     PetscCall(MatGetOwnershipRange(pdipm->Jce_xfixed, &Jcrstart, &Jcrend));
958     PetscCall(ISGetIndices(pdipm->isxfixed, &cols));
959     k = 0;
960     for (row = Jcrstart; row < Jcrend; row++) {
961       PetscCall(MatSetValues(pdipm->Jce_xfixed, 1, &row, 1, cols + k, &one, INSERT_VALUES));
962       k++;
963     }
964     PetscCall(ISRestoreIndices(pdipm->isxfixed, &cols));
965     PetscCall(MatAssemblyBegin(pdipm->Jce_xfixed, MAT_FINAL_ASSEMBLY));
966     PetscCall(MatAssemblyEnd(pdipm->Jce_xfixed, MAT_FINAL_ASSEMBLY));
967   }
968 
969   /* (5.2) PDIPM inequality Jacobian Jci = [tao->jacobian_inequality; ...] */
970   PetscCall(MatCreate(comm, &pdipm->Jci_xb));
971   PetscCall(MatSetSizes(pdipm->Jci_xb, pdipm->nci - pdipm->nh, pdipm->nx, PETSC_DECIDE, pdipm->Nx));
972   PetscCall(MatSetFromOptions(pdipm->Jci_xb));
973   PetscCall(MatSeqAIJSetPreallocation(pdipm->Jci_xb, 1, NULL));
974   PetscCall(MatMPIAIJSetPreallocation(pdipm->Jci_xb, 1, NULL, 1, NULL));
975 
976   PetscCall(MatGetOwnershipRange(pdipm->Jci_xb, &Jcrstart, &Jcrend));
977   offset = Jcrstart;
978   if (pdipm->Nxub) {
979     /* Add xub to Jci_xb */
980     PetscCall(ISGetIndices(pdipm->isxub, &cols));
981     k = 0;
982     for (row = offset; row < offset + pdipm->nxub; row++) {
983       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &row, 1, cols + k, &neg_one, INSERT_VALUES));
984       k++;
985     }
986     PetscCall(ISRestoreIndices(pdipm->isxub, &cols));
987   }
988 
989   if (pdipm->Nxlb) {
990     /* Add xlb to Jci_xb */
991     PetscCall(ISGetIndices(pdipm->isxlb, &cols));
992     k = 0;
993     offset += pdipm->nxub;
994     for (row = offset; row < offset + pdipm->nxlb; row++) {
995       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &row, 1, cols + k, &one, INSERT_VALUES));
996       k++;
997     }
998     PetscCall(ISRestoreIndices(pdipm->isxlb, &cols));
999   }
1000 
1001   /* Add xbox to Jci_xb */
1002   if (pdipm->Nxbox) {
1003     PetscCall(ISGetIndices(pdipm->isxbox, &cols));
1004     k = 0;
1005     offset += pdipm->nxlb;
1006     for (row = offset; row < offset + pdipm->nxbox; row++) {
1007       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &row, 1, cols + k, &neg_one, INSERT_VALUES));
1008       tmp = row + pdipm->nxbox;
1009       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &tmp, 1, cols + k, &one, INSERT_VALUES));
1010       k++;
1011     }
1012     PetscCall(ISRestoreIndices(pdipm->isxbox, &cols));
1013   }
1014 
1015   PetscCall(MatAssemblyBegin(pdipm->Jci_xb, MAT_FINAL_ASSEMBLY));
1016   PetscCall(MatAssemblyEnd(pdipm->Jci_xb, MAT_FINAL_ASSEMBLY));
1017   /* PetscCall(MatView(pdipm->Jci_xb,PETSC_VIEWER_STDOUT_WORLD)); */
1018 
1019   /* (6) Set up ISs for PC Fieldsplit */
1020   if (pdipm->solve_reduced_kkt) {
1021     PetscCall(PetscMalloc2(pdipm->nx + pdipm->nce, &xa, 2 * pdipm->nci, &xb));
1022     for (i = 0; i < pdipm->nx + pdipm->nce; i++) xa[i] = i;
1023     for (i = 0; i < 2 * pdipm->nci; i++) xb[i] = pdipm->off_lambdai + i;
1024 
1025     PetscCall(ISCreateGeneral(comm, pdipm->nx + pdipm->nce, xa, PETSC_OWN_POINTER, &pdipm->is1));
1026     PetscCall(ISCreateGeneral(comm, 2 * pdipm->nci, xb, PETSC_OWN_POINTER, &pdipm->is2));
1027   }
1028 
1029   /* (7) Gather offsets from all processes */
1030   PetscCall(PetscMalloc1(size, &pdipm->nce_all));
1031 
1032   /* Get rstart of KKT matrix */
1033   PetscCallMPI(MPI_Scan(&pdipm->n, &rstart, 1, MPIU_INT, MPI_SUM, comm));
1034   rstart -= pdipm->n;
1035 
1036   PetscCallMPI(MPI_Allgather(&pdipm->nce, 1, MPIU_INT, pdipm->nce_all, 1, MPIU_INT, comm));
1037 
1038   PetscCall(PetscMalloc3(size, &ng_all, size, &nh_all, size, &Jranges));
1039   PetscCallMPI(MPI_Allgather(&rstart, 1, MPIU_INT, Jranges, 1, MPIU_INT, comm));
1040   PetscCallMPI(MPI_Allgather(&pdipm->nh, 1, MPIU_INT, nh_all, 1, MPIU_INT, comm));
1041   PetscCallMPI(MPI_Allgather(&pdipm->ng, 1, MPIU_INT, ng_all, 1, MPIU_INT, comm));
1042 
1043   PetscCall(MatGetOwnershipRanges(tao->hessian, &rranges));
1044   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &cranges));
1045 
1046   if (pdipm->Ng) {
1047     PetscCall(TaoComputeJacobianEquality(tao, tao->solution, tao->jacobian_equality, tao->jacobian_equality_pre));
1048     PetscCall(MatTranspose(tao->jacobian_equality, MAT_INITIAL_MATRIX, &pdipm->jac_equality_trans));
1049   }
1050   if (pdipm->Nh) {
1051     PetscCall(TaoComputeJacobianInequality(tao, tao->solution, tao->jacobian_inequality, tao->jacobian_inequality_pre));
1052     PetscCall(MatTranspose(tao->jacobian_inequality, MAT_INITIAL_MATRIX, &pdipm->jac_inequality_trans));
1053   }
1054 
1055   /* Count dnz,onz for preallocation of KKT matrix */
1056   nce_all = pdipm->nce_all;
1057 
1058   if (pdipm->Nxfixed) PetscCall(MatTranspose(pdipm->Jce_xfixed, MAT_INITIAL_MATRIX, &Jce_xfixed_trans));
1059   PetscCall(MatTranspose(pdipm->Jci_xb, MAT_INITIAL_MATRIX, &Jci_xb_trans));
1060 
1061   MatPreallocateBegin(comm, pdipm->n, pdipm->n, dnz, onz);
1062 
1063   /* 1st row block of KKT matrix: [Wxx; gradCe'; -gradCi'; 0] */
1064   PetscCall(TaoPDIPMEvaluateFunctionsAndJacobians(tao, pdipm->x));
1065   PetscCall(TaoComputeHessian(tao, tao->solution, tao->hessian, tao->hessian_pre));
1066 
1067   /* Insert tao->hessian */
1068   PetscCall(MatGetOwnershipRange(tao->hessian, &rjstart, NULL));
1069   for (i = 0; i < pdipm->nx; i++) {
1070     row = rstart + i;
1071 
1072     PetscCall(MatGetRow(tao->hessian, i + rjstart, &nc, &aj, NULL));
1073     proc = 0;
1074     for (j = 0; j < nc; j++) {
1075       while (aj[j] >= cranges[proc + 1]) proc++;
1076       col = aj[j] - cranges[proc] + Jranges[proc];
1077       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1078     }
1079     PetscCall(MatRestoreRow(tao->hessian, i + rjstart, &nc, &aj, NULL));
1080 
1081     if (pdipm->ng) {
1082       /* Insert grad g' */
1083       PetscCall(MatGetRow(pdipm->jac_equality_trans, i + rjstart, &nc, &aj, NULL));
1084       PetscCall(MatGetOwnershipRanges(tao->jacobian_equality, &ranges));
1085       proc = 0;
1086       for (j = 0; j < nc; j++) {
1087         /* find row ownership of */
1088         while (aj[j] >= ranges[proc + 1]) proc++;
1089         nx_all = rranges[proc + 1] - rranges[proc];
1090         col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all;
1091         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1092       }
1093       PetscCall(MatRestoreRow(pdipm->jac_equality_trans, i + rjstart, &nc, &aj, NULL));
1094     }
1095 
1096     /* Insert Jce_xfixed^T' */
1097     if (pdipm->nxfixed) {
1098       PetscCall(MatGetRow(Jce_xfixed_trans, i + rjstart, &nc, &aj, NULL));
1099       PetscCall(MatGetOwnershipRanges(pdipm->Jce_xfixed, &ranges));
1100       proc = 0;
1101       for (j = 0; j < nc; j++) {
1102         /* find row ownership of */
1103         while (aj[j] >= ranges[proc + 1]) proc++;
1104         nx_all = rranges[proc + 1] - rranges[proc];
1105         col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all + ng_all[proc];
1106         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1107       }
1108       PetscCall(MatRestoreRow(Jce_xfixed_trans, i + rjstart, &nc, &aj, NULL));
1109     }
1110 
1111     if (pdipm->nh) {
1112       /* Insert -grad h' */
1113       PetscCall(MatGetRow(pdipm->jac_inequality_trans, i + rjstart, &nc, &aj, NULL));
1114       PetscCall(MatGetOwnershipRanges(tao->jacobian_inequality, &ranges));
1115       proc = 0;
1116       for (j = 0; j < nc; j++) {
1117         /* find row ownership of */
1118         while (aj[j] >= ranges[proc + 1]) proc++;
1119         nx_all = rranges[proc + 1] - rranges[proc];
1120         col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all + nce_all[proc];
1121         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1122       }
1123       PetscCall(MatRestoreRow(pdipm->jac_inequality_trans, i + rjstart, &nc, &aj, NULL));
1124     }
1125 
1126     /* Insert Jci_xb^T' */
1127     PetscCall(MatGetRow(Jci_xb_trans, i + rjstart, &nc, &aj, NULL));
1128     PetscCall(MatGetOwnershipRanges(pdipm->Jci_xb, &ranges));
1129     proc = 0;
1130     for (j = 0; j < nc; j++) {
1131       /* find row ownership of */
1132       while (aj[j] >= ranges[proc + 1]) proc++;
1133       nx_all = rranges[proc + 1] - rranges[proc];
1134       col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all + nce_all[proc] + nh_all[proc];
1135       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1136     }
1137     PetscCall(MatRestoreRow(Jci_xb_trans, i + rjstart, &nc, &aj, NULL));
1138   }
1139 
1140   /* 2nd Row block of KKT matrix: [grad Ce, deltac*I, 0, 0] */
1141   if (pdipm->Ng) {
1142     PetscCall(MatGetOwnershipRange(tao->jacobian_equality, &rjstart, NULL));
1143     for (i = 0; i < pdipm->ng; i++) {
1144       row = rstart + pdipm->off_lambdae + i;
1145 
1146       PetscCall(MatGetRow(tao->jacobian_equality, i + rjstart, &nc, &aj, NULL));
1147       proc = 0;
1148       for (j = 0; j < nc; j++) {
1149         while (aj[j] >= cranges[proc + 1]) proc++;
1150         col = aj[j] - cranges[proc] + Jranges[proc];
1151         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz)); /* grad g */
1152       }
1153       PetscCall(MatRestoreRow(tao->jacobian_equality, i + rjstart, &nc, &aj, NULL));
1154     }
1155   }
1156   /* Jce_xfixed */
1157   if (pdipm->Nxfixed) {
1158     PetscCall(MatGetOwnershipRange(pdipm->Jce_xfixed, &Jcrstart, NULL));
1159     for (i = 0; i < (pdipm->nce - pdipm->ng); i++) {
1160       row = rstart + pdipm->off_lambdae + pdipm->ng + i;
1161 
1162       PetscCall(MatGetRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, NULL));
1163       PetscCheck(nc == 1, PETSC_COMM_SELF, PETSC_ERR_SUP, "nc != 1");
1164 
1165       proc = 0;
1166       j    = 0;
1167       while (cols[j] >= cranges[proc + 1]) proc++;
1168       col = cols[j] - cranges[proc] + Jranges[proc];
1169       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1170       PetscCall(MatRestoreRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, NULL));
1171     }
1172   }
1173 
1174   /* 3rd Row block of KKT matrix: [ gradCi, 0, deltac*I, -I] */
1175   if (pdipm->Nh) {
1176     PetscCall(MatGetOwnershipRange(tao->jacobian_inequality, &rjstart, NULL));
1177     for (i = 0; i < pdipm->nh; i++) {
1178       row = rstart + pdipm->off_lambdai + i;
1179 
1180       PetscCall(MatGetRow(tao->jacobian_inequality, i + rjstart, &nc, &aj, NULL));
1181       proc = 0;
1182       for (j = 0; j < nc; j++) {
1183         while (aj[j] >= cranges[proc + 1]) proc++;
1184         col = aj[j] - cranges[proc] + Jranges[proc];
1185         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz)); /* grad h */
1186       }
1187       PetscCall(MatRestoreRow(tao->jacobian_inequality, i + rjstart, &nc, &aj, NULL));
1188     }
1189     /* I */
1190     for (i = 0; i < pdipm->nh; i++) {
1191       row = rstart + pdipm->off_lambdai + i;
1192       col = rstart + pdipm->off_z + i;
1193       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1194     }
1195   }
1196 
1197   /* Jci_xb */
1198   PetscCall(MatGetOwnershipRange(pdipm->Jci_xb, &Jcrstart, NULL));
1199   for (i = 0; i < (pdipm->nci - pdipm->nh); i++) {
1200     row = rstart + pdipm->off_lambdai + pdipm->nh + i;
1201 
1202     PetscCall(MatGetRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, NULL));
1203     PetscCheck(nc == 1, PETSC_COMM_SELF, PETSC_ERR_SUP, "nc != 1");
1204     proc = 0;
1205     for (j = 0; j < nc; j++) {
1206       while (cols[j] >= cranges[proc + 1]) proc++;
1207       col = cols[j] - cranges[proc] + Jranges[proc];
1208       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1209     }
1210     PetscCall(MatRestoreRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, NULL));
1211     /* I */
1212     col = rstart + pdipm->off_z + pdipm->nh + i;
1213     PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1214   }
1215 
1216   /* 4-th Row block of KKT matrix: Z and Ci */
1217   for (i = 0; i < pdipm->nci; i++) {
1218     row      = rstart + pdipm->off_z + i;
1219     cols1[0] = rstart + pdipm->off_lambdai + i;
1220     cols1[1] = row;
1221     PetscCall(MatPreallocateSet(row, 2, cols1, dnz, onz));
1222   }
1223 
1224   /* diagonal entry */
1225   for (i = 0; i < pdipm->n; i++) dnz[i]++; /* diagonal entry */
1226 
1227   /* Create KKT matrix */
1228   PetscCall(MatCreate(comm, &J));
1229   PetscCall(MatSetSizes(J, pdipm->n, pdipm->n, PETSC_DECIDE, PETSC_DECIDE));
1230   PetscCall(MatSetFromOptions(J));
1231   PetscCall(MatSeqAIJSetPreallocation(J, 0, dnz));
1232   PetscCall(MatMPIAIJSetPreallocation(J, 0, dnz, 0, onz));
1233   MatPreallocateEnd(dnz, onz);
1234   pdipm->K = J;
1235 
1236   /* (8) Insert constant entries to  K */
1237   /* Set 0.0 to diagonal of K, so that the solver does not complain *about missing diagonal value */
1238   PetscCall(MatGetOwnershipRange(J, &rstart, &rend));
1239   for (i = rstart; i < rend; i++) PetscCall(MatSetValue(J, i, i, 0.0, INSERT_VALUES));
1240   /* In case Wxx has no diagonal entries preset set diagonal to deltaw given */
1241   if (pdipm->kkt_pd) {
1242     for (i = 0; i < pdipm->nh; i++) {
1243       row = rstart + i;
1244       PetscCall(MatSetValue(J, row, row, pdipm->deltaw, INSERT_VALUES));
1245     }
1246   }
1247 
1248   /* Row block of K: [ grad Ce, 0, 0, 0] */
1249   if (pdipm->Nxfixed) {
1250     PetscCall(MatGetOwnershipRange(pdipm->Jce_xfixed, &Jcrstart, NULL));
1251     for (i = 0; i < (pdipm->nce - pdipm->ng); i++) {
1252       row = rstart + pdipm->off_lambdae + pdipm->ng + i;
1253 
1254       PetscCall(MatGetRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, &aa));
1255       proc = 0;
1256       for (j = 0; j < nc; j++) {
1257         while (cols[j] >= cranges[proc + 1]) proc++;
1258         col = cols[j] - cranges[proc] + Jranges[proc];
1259         PetscCall(MatSetValue(J, row, col, aa[j], INSERT_VALUES)); /* grad Ce */
1260         PetscCall(MatSetValue(J, col, row, aa[j], INSERT_VALUES)); /* grad Ce' */
1261       }
1262       PetscCall(MatRestoreRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, &aa));
1263     }
1264   }
1265 
1266   /* Row block of K: [ -grad Ci, 0, 0, I] */
1267   PetscCall(MatGetOwnershipRange(pdipm->Jci_xb, &Jcrstart, NULL));
1268   for (i = 0; i < pdipm->nci - pdipm->nh; i++) {
1269     row = rstart + pdipm->off_lambdai + pdipm->nh + i;
1270 
1271     PetscCall(MatGetRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, &aa));
1272     proc = 0;
1273     for (j = 0; j < nc; j++) {
1274       while (cols[j] >= cranges[proc + 1]) proc++;
1275       col = cols[j] - cranges[proc] + Jranges[proc];
1276       PetscCall(MatSetValue(J, col, row, -aa[j], INSERT_VALUES));
1277       PetscCall(MatSetValue(J, row, col, -aa[j], INSERT_VALUES));
1278     }
1279     PetscCall(MatRestoreRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, &aa));
1280 
1281     col = rstart + pdipm->off_z + pdipm->nh + i;
1282     PetscCall(MatSetValue(J, row, col, 1, INSERT_VALUES));
1283   }
1284 
1285   for (i = 0; i < pdipm->nh; i++) {
1286     row = rstart + pdipm->off_lambdai + i;
1287     col = rstart + pdipm->off_z + i;
1288     PetscCall(MatSetValue(J, row, col, 1, INSERT_VALUES));
1289   }
1290 
1291   /* Row block of K: [ 0, 0, I, ...] */
1292   for (i = 0; i < pdipm->nci; i++) {
1293     row = rstart + pdipm->off_z + i;
1294     col = rstart + pdipm->off_lambdai + i;
1295     PetscCall(MatSetValue(J, row, col, 1, INSERT_VALUES));
1296   }
1297 
1298   if (pdipm->Nxfixed) PetscCall(MatDestroy(&Jce_xfixed_trans));
1299   PetscCall(MatDestroy(&Jci_xb_trans));
1300   PetscCall(PetscFree3(ng_all, nh_all, Jranges));
1301 
1302   /* (9) Set up nonlinear solver SNES */
1303   PetscCall(SNESSetFunction(pdipm->snes, NULL, TaoSNESFunction_PDIPM, (void *)tao));
1304   PetscCall(SNESSetJacobian(pdipm->snes, J, J, TaoSNESJacobian_PDIPM, (void *)tao));
1305 
1306   if (pdipm->solve_reduced_kkt) {
1307     PC pc;
1308     PetscCall(KSPGetPC(tao->ksp, &pc));
1309     PetscCall(PCSetType(pc, PCFIELDSPLIT));
1310     PetscCall(PCFieldSplitSetType(pc, PC_COMPOSITE_SCHUR));
1311     PetscCall(PCFieldSplitSetIS(pc, "2", pdipm->is2));
1312     PetscCall(PCFieldSplitSetIS(pc, "1", pdipm->is1));
1313   }
1314   PetscCall(SNESSetFromOptions(pdipm->snes));
1315 
1316   /* (10) Setup PCPostSetUp() for pdipm->solve_symmetric_kkt */
1317   if (pdipm->solve_symmetric_kkt) {
1318     KSP       ksp;
1319     PC        pc;
1320     PetscBool isCHOL;
1321 
1322     PetscCall(SNESGetKSP(pdipm->snes, &ksp));
1323     PetscCall(KSPGetPC(ksp, &pc));
1324     PetscCall(PCSetPostSetUp(pc, PCPostSetUp_PDIPM));
1325 
1326     PetscCall(PetscObjectTypeCompare((PetscObject)pc, PCCHOLESKY, &isCHOL));
1327     if (isCHOL) {
1328       Mat Factor;
1329 
1330       PetscCheck(PetscDefined(HAVE_MUMPS), PetscObjectComm((PetscObject)tao), PETSC_ERR_SUP, "Requires external package MUMPS");
1331       PetscCall(PCFactorGetMatrix(pc, &Factor));
1332       PetscCall(MatMumpsSetIcntl(Factor, 24, 1));               /* detection of null pivot rows */
1333       if (size > 1) PetscCall(MatMumpsSetIcntl(Factor, 13, 1)); /* parallelism of the root node (enable ScaLAPACK) and its splitting */
1334     }
1335   }
1336   PetscFunctionReturn(PETSC_SUCCESS);
1337 }
1338 
TaoDestroy_PDIPM(Tao tao)1339 static PetscErrorCode TaoDestroy_PDIPM(Tao tao)
1340 {
1341   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;
1342 
1343   PetscFunctionBegin;
1344   /* Freeing Vectors assocaiated with KKT (X) */
1345   PetscCall(VecDestroy(&pdipm->x));       /* Solution x */
1346   PetscCall(VecDestroy(&pdipm->lambdae)); /* Equality constraints lagrangian multiplier*/
1347   PetscCall(VecDestroy(&pdipm->lambdai)); /* Inequality constraints lagrangian multiplier*/
1348   PetscCall(VecDestroy(&pdipm->z));       /* Slack variables */
1349   PetscCall(VecDestroy(&pdipm->X));       /* Big KKT system vector [x; lambdae; lambdai; z] */
1350 
1351   /* work vectors */
1352   PetscCall(VecDestroy(&pdipm->lambdae_xfixed));
1353   PetscCall(VecDestroy(&pdipm->lambdai_xb));
1354 
1355   /* Legrangian equality and inequality Vec */
1356   PetscCall(VecDestroy(&pdipm->ce)); /* Vec of equality constraints */
1357   PetscCall(VecDestroy(&pdipm->ci)); /* Vec of inequality constraints */
1358 
1359   /* Matrices */
1360   PetscCall(MatDestroy(&pdipm->Jce_xfixed));
1361   PetscCall(MatDestroy(&pdipm->Jci_xb)); /* Jacobian of inequality constraints Jci = [tao->jacobian_inequality ; J(nxub); J(nxlb); J(nxbx)] */
1362   PetscCall(MatDestroy(&pdipm->K));
1363 
1364   /* Index Sets */
1365   if (pdipm->Nxub) PetscCall(ISDestroy(&pdipm->isxub)); /* Finite upper bound only -inf < x < ub */
1366 
1367   if (pdipm->Nxlb) PetscCall(ISDestroy(&pdipm->isxlb)); /* Finite lower bound only  lb <= x < inf */
1368 
1369   if (pdipm->Nxfixed) PetscCall(ISDestroy(&pdipm->isxfixed)); /* Fixed variables         lb =  x = ub */
1370 
1371   if (pdipm->Nxbox) PetscCall(ISDestroy(&pdipm->isxbox)); /* Boxed variables         lb <= x <= ub */
1372 
1373   if (pdipm->Nxfree) PetscCall(ISDestroy(&pdipm->isxfree)); /* Free variables        -inf <= x <= inf */
1374 
1375   if (pdipm->solve_reduced_kkt) {
1376     PetscCall(ISDestroy(&pdipm->is1));
1377     PetscCall(ISDestroy(&pdipm->is2));
1378   }
1379 
1380   /* SNES */
1381   PetscCall(SNESDestroy(&pdipm->snes)); /* Nonlinear solver */
1382   PetscCall(PetscFree(pdipm->nce_all));
1383   PetscCall(MatDestroy(&pdipm->jac_equality_trans));
1384   PetscCall(MatDestroy(&pdipm->jac_inequality_trans));
1385 
1386   /* Destroy pdipm */
1387   PetscCall(PetscFree(tao->data)); /* Holding locations of pdipm */
1388 
1389   /* Destroy Dual */
1390   PetscCall(VecDestroy(&tao->DE)); /* equality dual */
1391   PetscCall(VecDestroy(&tao->DI)); /* dinequality dual */
1392   PetscFunctionReturn(PETSC_SUCCESS);
1393 }
1394 
TaoSetFromOptions_PDIPM(Tao tao,PetscOptionItems PetscOptionsObject)1395 static PetscErrorCode TaoSetFromOptions_PDIPM(Tao tao, PetscOptionItems PetscOptionsObject)
1396 {
1397   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;
1398 
1399   PetscFunctionBegin;
1400   PetscOptionsHeadBegin(PetscOptionsObject, "PDIPM method for constrained optimization");
1401   PetscCall(PetscOptionsReal("-tao_pdipm_push_init_slack", "parameter to push initial slack variables away from bounds", NULL, pdipm->push_init_slack, &pdipm->push_init_slack, NULL));
1402   PetscCall(PetscOptionsReal("-tao_pdipm_push_init_lambdai", "parameter to push initial (inequality) dual variables away from bounds", NULL, pdipm->push_init_lambdai, &pdipm->push_init_lambdai, NULL));
1403   PetscCall(PetscOptionsBool("-tao_pdipm_solve_reduced_kkt", "Solve reduced KKT system using Schur-complement", NULL, pdipm->solve_reduced_kkt, &pdipm->solve_reduced_kkt, NULL));
1404   PetscCall(PetscOptionsReal("-tao_pdipm_mu_update_factor", "Update scalar for barrier parameter (mu) update", NULL, pdipm->mu_update_factor, &pdipm->mu_update_factor, NULL));
1405   PetscCall(PetscOptionsBool("-tao_pdipm_symmetric_kkt", "Solve non reduced symmetric KKT system", NULL, pdipm->solve_symmetric_kkt, &pdipm->solve_symmetric_kkt, NULL));
1406   PetscCall(PetscOptionsBool("-tao_pdipm_kkt_shift_pd", "Add shifts to make KKT matrix positive definite", NULL, pdipm->kkt_pd, &pdipm->kkt_pd, NULL));
1407   PetscOptionsHeadEnd();
1408   PetscFunctionReturn(PETSC_SUCCESS);
1409 }
1410 
1411 /*MC
1412   TAOPDIPM - Barrier-based primal-dual interior point algorithm for generally constrained optimization.
1413 
1414   Options Database Keys:
1415 +   -tao_pdipm_push_init_lambdai - parameter to push initial dual variables away from bounds (> 0)
1416 .   -tao_pdipm_push_init_slack   - parameter to push initial slack variables away from bounds (> 0)
1417 .   -tao_pdipm_mu_update_factor  - update scalar for barrier parameter (mu) update (> 0)
1418 .   -tao_pdipm_symmetric_kkt     - Solve non-reduced symmetric KKT system
1419 -   -tao_pdipm_kkt_shift_pd      - Add shifts to make KKT matrix positive definite
1420 
1421   Level: beginner
1422 
1423 .seealso: `TAOPDIPM`, `Tao`, `TaoType`
1424 M*/
1425 
TaoCreate_PDIPM(Tao tao)1426 PETSC_EXTERN PetscErrorCode TaoCreate_PDIPM(Tao tao)
1427 {
1428   TAO_PDIPM *pdipm;
1429   PC         pc;
1430 
1431   PetscFunctionBegin;
1432   tao->ops->setup          = TaoSetup_PDIPM;
1433   tao->ops->solve          = TaoSolve_PDIPM;
1434   tao->ops->setfromoptions = TaoSetFromOptions_PDIPM;
1435   tao->ops->view           = TaoView_PDIPM;
1436   tao->ops->destroy        = TaoDestroy_PDIPM;
1437 
1438   PetscCall(PetscNew(&pdipm));
1439   tao->data = (void *)pdipm;
1440 
1441   pdipm->nx = pdipm->Nx = 0;
1442   pdipm->nxfixed = pdipm->Nxfixed = 0;
1443   pdipm->nxlb = pdipm->Nxlb = 0;
1444   pdipm->nxub = pdipm->Nxub = 0;
1445   pdipm->nxbox = pdipm->Nxbox = 0;
1446   pdipm->nxfree = pdipm->Nxfree = 0;
1447 
1448   pdipm->ng = pdipm->Ng = pdipm->nce = pdipm->Nce = 0;
1449   pdipm->nh = pdipm->Nh = pdipm->nci = pdipm->Nci = 0;
1450   pdipm->n = pdipm->N     = 0;
1451   pdipm->mu               = 1.0;
1452   pdipm->mu_update_factor = 0.1;
1453 
1454   pdipm->deltaw     = 0.0;
1455   pdipm->lastdeltaw = 3 * 1.e-4;
1456   pdipm->deltac     = 0.0;
1457   pdipm->kkt_pd     = PETSC_FALSE;
1458 
1459   pdipm->push_init_slack     = 1.0;
1460   pdipm->push_init_lambdai   = 1.0;
1461   pdipm->solve_reduced_kkt   = PETSC_FALSE;
1462   pdipm->solve_symmetric_kkt = PETSC_TRUE;
1463 
1464   /* Override default settings (unless already changed) */
1465   PetscCall(TaoParametersInitialize(tao));
1466   PetscObjectParameterSetDefault(tao, max_it, 200);
1467   PetscObjectParameterSetDefault(tao, max_funcs, 500);
1468 
1469   PetscCall(SNESCreate(((PetscObject)tao)->comm, &pdipm->snes));
1470   PetscCall(SNESSetOptionsPrefix(pdipm->snes, tao->hdr.prefix));
1471   PetscCall(SNESGetKSP(pdipm->snes, &tao->ksp));
1472   PetscCall(PetscObjectReference((PetscObject)tao->ksp));
1473   PetscCall(KSPGetPC(tao->ksp, &pc));
1474   PetscCall(PCSetApplicationContext(pc, (void *)tao));
1475   PetscFunctionReturn(PETSC_SUCCESS);
1476 }
1477