xref: /petsc/src/tao/constrained/impls/ipm/pdipm.c (revision 20f4b53cbb5e9bd9ef12b76a8697d60d197cda17)
1aad13602SShrirang Abhyankar #include <../src/tao/constrained/impls/ipm/pdipm.h>
2aad13602SShrirang Abhyankar 
3aad13602SShrirang Abhyankar /*
4aad13602SShrirang Abhyankar    TaoPDIPMEvaluateFunctionsAndJacobians - Evaluate the objective function f, gradient fx, constraints, and all the Jacobians at current vector
5aad13602SShrirang Abhyankar 
6c3339decSBarry Smith    Collective
7aad13602SShrirang Abhyankar 
8aad13602SShrirang Abhyankar    Input Parameter:
9aad13602SShrirang Abhyankar +  tao - solver context
10aad13602SShrirang Abhyankar -  x - vector at which all objects to be evaluated
11aad13602SShrirang Abhyankar 
12aad13602SShrirang Abhyankar    Level: beginner
13aad13602SShrirang Abhyankar 
14db781477SPatrick Sanan .seealso: `TaoPDIPMUpdateConstraints()`, `TaoPDIPMSetUpBounds()`
15aad13602SShrirang Abhyankar */
16d71ae5a4SJacob Faibussowitsch static PetscErrorCode TaoPDIPMEvaluateFunctionsAndJacobians(Tao tao, Vec x)
17d71ae5a4SJacob Faibussowitsch {
18aad13602SShrirang Abhyankar   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;
19aad13602SShrirang Abhyankar 
20aad13602SShrirang Abhyankar   PetscFunctionBegin;
21aad13602SShrirang Abhyankar   /* Compute user objective function and gradient */
229566063dSJacob Faibussowitsch   PetscCall(TaoComputeObjectiveAndGradient(tao, x, &pdipm->obj, tao->gradient));
23aad13602SShrirang Abhyankar 
24aad13602SShrirang Abhyankar   /* Equality constraints and Jacobian */
25aad13602SShrirang Abhyankar   if (pdipm->Ng) {
269566063dSJacob Faibussowitsch     PetscCall(TaoComputeEqualityConstraints(tao, x, tao->constraints_equality));
279566063dSJacob Faibussowitsch     PetscCall(TaoComputeJacobianEquality(tao, x, tao->jacobian_equality, tao->jacobian_equality_pre));
28aad13602SShrirang Abhyankar   }
29aad13602SShrirang Abhyankar 
30aad13602SShrirang Abhyankar   /* Inequality constraints and Jacobian */
31aad13602SShrirang Abhyankar   if (pdipm->Nh) {
329566063dSJacob Faibussowitsch     PetscCall(TaoComputeInequalityConstraints(tao, x, tao->constraints_inequality));
339566063dSJacob Faibussowitsch     PetscCall(TaoComputeJacobianInequality(tao, x, tao->jacobian_inequality, tao->jacobian_inequality_pre));
34aad13602SShrirang Abhyankar   }
353ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
36aad13602SShrirang Abhyankar }
37aad13602SShrirang Abhyankar 
38aad13602SShrirang Abhyankar /*
39aad13602SShrirang Abhyankar   TaoPDIPMUpdateConstraints - Update the vectors ce and ci at x
40aad13602SShrirang Abhyankar 
41aad13602SShrirang Abhyankar   Collective
42aad13602SShrirang Abhyankar 
43aad13602SShrirang Abhyankar   Input Parameter:
44aad13602SShrirang Abhyankar + tao - Tao context
45a5b23f4aSJose E. Roman - x - vector at which constraints to be evaluated
46aad13602SShrirang Abhyankar 
47aad13602SShrirang Abhyankar    Level: beginner
48aad13602SShrirang Abhyankar 
49db781477SPatrick Sanan .seealso: `TaoPDIPMEvaluateFunctionsAndJacobians()`
50aad13602SShrirang Abhyankar */
51d71ae5a4SJacob Faibussowitsch static PetscErrorCode TaoPDIPMUpdateConstraints(Tao tao, Vec x)
52d71ae5a4SJacob Faibussowitsch {
53aad13602SShrirang Abhyankar   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
54aad13602SShrirang Abhyankar   PetscInt           i, offset, offset1, k, xstart;
55aad13602SShrirang Abhyankar   PetscScalar       *carr;
56aad13602SShrirang Abhyankar   const PetscInt    *ubptr, *lbptr, *bxptr, *fxptr;
57aad13602SShrirang Abhyankar   const PetscScalar *xarr, *xuarr, *xlarr, *garr, *harr;
58aad13602SShrirang Abhyankar 
59aad13602SShrirang Abhyankar   PetscFunctionBegin;
609566063dSJacob Faibussowitsch   PetscCall(VecGetOwnershipRange(x, &xstart, NULL));
619566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(x, &xarr));
629566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(tao->XU, &xuarr));
639566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(tao->XL, &xlarr));
64aad13602SShrirang Abhyankar 
65aad13602SShrirang Abhyankar   /* (1) Update ce vector */
669566063dSJacob Faibussowitsch   PetscCall(VecGetArrayWrite(pdipm->ce, &carr));
67aad13602SShrirang Abhyankar 
688e58fa1dSresundermann   if (pdipm->Ng) {
692da392ccSBarry Smith     /* (1.a) Inserting updated g(x) */
709566063dSJacob Faibussowitsch     PetscCall(VecGetArrayRead(tao->constraints_equality, &garr));
719566063dSJacob Faibussowitsch     PetscCall(PetscMemcpy(carr, garr, pdipm->ng * sizeof(PetscScalar)));
729566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayRead(tao->constraints_equality, &garr));
73aad13602SShrirang Abhyankar   }
74aad13602SShrirang Abhyankar 
75aad13602SShrirang Abhyankar   /* (1.b) Update xfixed */
76aad13602SShrirang Abhyankar   if (pdipm->Nxfixed) {
77aad13602SShrirang Abhyankar     offset = pdipm->ng;
789566063dSJacob Faibussowitsch     PetscCall(ISGetIndices(pdipm->isxfixed, &fxptr)); /* global indices in x */
79aad13602SShrirang Abhyankar     for (k = 0; k < pdipm->nxfixed; k++) {
80aad13602SShrirang Abhyankar       i                = fxptr[k] - xstart;
81aad13602SShrirang Abhyankar       carr[offset + k] = xarr[i] - xuarr[i];
82aad13602SShrirang Abhyankar     }
83aad13602SShrirang Abhyankar   }
849566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayWrite(pdipm->ce, &carr));
85aad13602SShrirang Abhyankar 
86aad13602SShrirang Abhyankar   /* (2) Update ci vector */
879566063dSJacob Faibussowitsch   PetscCall(VecGetArrayWrite(pdipm->ci, &carr));
88aad13602SShrirang Abhyankar 
89aad13602SShrirang Abhyankar   if (pdipm->Nh) {
90aad13602SShrirang Abhyankar     /* (2.a) Inserting updated h(x) */
919566063dSJacob Faibussowitsch     PetscCall(VecGetArrayRead(tao->constraints_inequality, &harr));
929566063dSJacob Faibussowitsch     PetscCall(PetscMemcpy(carr, harr, pdipm->nh * sizeof(PetscScalar)));
939566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayRead(tao->constraints_inequality, &harr));
94aad13602SShrirang Abhyankar   }
95aad13602SShrirang Abhyankar 
96aad13602SShrirang Abhyankar   /* (2.b) Update xub */
97aad13602SShrirang Abhyankar   offset = pdipm->nh;
98aad13602SShrirang Abhyankar   if (pdipm->Nxub) {
999566063dSJacob Faibussowitsch     PetscCall(ISGetIndices(pdipm->isxub, &ubptr));
100aad13602SShrirang Abhyankar     for (k = 0; k < pdipm->nxub; k++) {
101aad13602SShrirang Abhyankar       i                = ubptr[k] - xstart;
102aad13602SShrirang Abhyankar       carr[offset + k] = xuarr[i] - xarr[i];
103aad13602SShrirang Abhyankar     }
104aad13602SShrirang Abhyankar   }
105aad13602SShrirang Abhyankar 
106aad13602SShrirang Abhyankar   if (pdipm->Nxlb) {
107aad13602SShrirang Abhyankar     /* (2.c) Update xlb */
108aad13602SShrirang Abhyankar     offset += pdipm->nxub;
1099566063dSJacob Faibussowitsch     PetscCall(ISGetIndices(pdipm->isxlb, &lbptr)); /* global indices in x */
110aad13602SShrirang Abhyankar     for (k = 0; k < pdipm->nxlb; k++) {
111aad13602SShrirang Abhyankar       i                = lbptr[k] - xstart;
112aad13602SShrirang Abhyankar       carr[offset + k] = xarr[i] - xlarr[i];
113aad13602SShrirang Abhyankar     }
114aad13602SShrirang Abhyankar   }
115aad13602SShrirang Abhyankar 
116aad13602SShrirang Abhyankar   if (pdipm->Nxbox) {
117aad13602SShrirang Abhyankar     /* (2.d) Update xbox */
118aad13602SShrirang Abhyankar     offset += pdipm->nxlb;
119aad13602SShrirang Abhyankar     offset1 = offset + pdipm->nxbox;
1209566063dSJacob Faibussowitsch     PetscCall(ISGetIndices(pdipm->isxbox, &bxptr)); /* global indices in x */
121aad13602SShrirang Abhyankar     for (k = 0; k < pdipm->nxbox; k++) {
122aad13602SShrirang Abhyankar       i                 = bxptr[k] - xstart; /* local indices in x */
123aad13602SShrirang Abhyankar       carr[offset + k]  = xuarr[i] - xarr[i];
124aad13602SShrirang Abhyankar       carr[offset1 + k] = xarr[i] - xlarr[i];
125aad13602SShrirang Abhyankar     }
126aad13602SShrirang Abhyankar   }
1279566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayWrite(pdipm->ci, &carr));
128aad13602SShrirang Abhyankar 
129aad13602SShrirang Abhyankar   /* Restoring Vectors */
1309566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(x, &xarr));
1319566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(tao->XU, &xuarr));
1329566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(tao->XL, &xlarr));
1333ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
134aad13602SShrirang Abhyankar }
135aad13602SShrirang Abhyankar 
136aad13602SShrirang Abhyankar /*
137aad13602SShrirang Abhyankar    TaoPDIPMSetUpBounds - Create upper and lower bound vectors of x
138aad13602SShrirang Abhyankar 
139aad13602SShrirang Abhyankar    Collective
140aad13602SShrirang Abhyankar 
141aad13602SShrirang Abhyankar    Input Parameter:
142aad13602SShrirang Abhyankar .  tao - holds pdipm and XL & XU
143aad13602SShrirang Abhyankar 
144aad13602SShrirang Abhyankar    Level: beginner
145aad13602SShrirang Abhyankar 
146db781477SPatrick Sanan .seealso: `TaoPDIPMUpdateConstraints`
147aad13602SShrirang Abhyankar */
148d71ae5a4SJacob Faibussowitsch static PetscErrorCode TaoPDIPMSetUpBounds(Tao tao)
149d71ae5a4SJacob Faibussowitsch {
150aad13602SShrirang Abhyankar   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
151aad13602SShrirang Abhyankar   const PetscScalar *xl, *xu;
152aad13602SShrirang Abhyankar   PetscInt           n, *ixlb, *ixub, *ixfixed, *ixfree, *ixbox, i, low, high, idx;
153aad13602SShrirang Abhyankar   MPI_Comm           comm;
154aad13602SShrirang Abhyankar   PetscInt           sendbuf[5], recvbuf[5];
155aad13602SShrirang Abhyankar 
156aad13602SShrirang Abhyankar   PetscFunctionBegin;
157aad13602SShrirang Abhyankar   /* Creates upper and lower bounds vectors on x, if not created already */
1589566063dSJacob Faibussowitsch   PetscCall(TaoComputeVariableBounds(tao));
159aad13602SShrirang Abhyankar 
1609566063dSJacob Faibussowitsch   PetscCall(VecGetLocalSize(tao->XL, &n));
1619566063dSJacob Faibussowitsch   PetscCall(PetscMalloc5(n, &ixlb, n, &ixub, n, &ixfree, n, &ixfixed, n, &ixbox));
162aad13602SShrirang Abhyankar 
1639566063dSJacob Faibussowitsch   PetscCall(VecGetOwnershipRange(tao->XL, &low, &high));
1649566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(tao->XL, &xl));
1659566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(tao->XU, &xu));
166aad13602SShrirang Abhyankar   for (i = 0; i < n; i++) {
167aad13602SShrirang Abhyankar     idx = low + i;
168aad13602SShrirang Abhyankar     if ((PetscRealPart(xl[i]) > PETSC_NINFINITY) && (PetscRealPart(xu[i]) < PETSC_INFINITY)) {
169aad13602SShrirang Abhyankar       if (PetscRealPart(xl[i]) == PetscRealPart(xu[i])) {
170aad13602SShrirang Abhyankar         ixfixed[pdipm->nxfixed++] = idx;
171aad13602SShrirang Abhyankar       } else ixbox[pdipm->nxbox++] = idx;
172aad13602SShrirang Abhyankar     } else {
173aad13602SShrirang Abhyankar       if ((PetscRealPart(xl[i]) > PETSC_NINFINITY) && (PetscRealPart(xu[i]) >= PETSC_INFINITY)) {
174aad13602SShrirang Abhyankar         ixlb[pdipm->nxlb++] = idx;
175aad13602SShrirang Abhyankar       } else if ((PetscRealPart(xl[i]) <= PETSC_NINFINITY) && (PetscRealPart(xu[i]) < PETSC_INFINITY)) {
176aad13602SShrirang Abhyankar         ixub[pdipm->nxlb++] = idx;
177aad13602SShrirang Abhyankar       } else ixfree[pdipm->nxfree++] = idx;
178aad13602SShrirang Abhyankar     }
179aad13602SShrirang Abhyankar   }
1809566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(tao->XL, &xl));
1819566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(tao->XU, &xu));
182aad13602SShrirang Abhyankar 
1839566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)tao, &comm));
184aad13602SShrirang Abhyankar   sendbuf[0] = pdipm->nxlb;
185aad13602SShrirang Abhyankar   sendbuf[1] = pdipm->nxub;
186aad13602SShrirang Abhyankar   sendbuf[2] = pdipm->nxfixed;
187aad13602SShrirang Abhyankar   sendbuf[3] = pdipm->nxbox;
188aad13602SShrirang Abhyankar   sendbuf[4] = pdipm->nxfree;
189aad13602SShrirang Abhyankar 
1909566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Allreduce(sendbuf, recvbuf, 5, MPIU_INT, MPI_SUM, comm));
191aad13602SShrirang Abhyankar   pdipm->Nxlb    = recvbuf[0];
192aad13602SShrirang Abhyankar   pdipm->Nxub    = recvbuf[1];
193aad13602SShrirang Abhyankar   pdipm->Nxfixed = recvbuf[2];
194aad13602SShrirang Abhyankar   pdipm->Nxbox   = recvbuf[3];
195aad13602SShrirang Abhyankar   pdipm->Nxfree  = recvbuf[4];
196aad13602SShrirang Abhyankar 
19748a46eb9SPierre Jolivet   if (pdipm->Nxlb) PetscCall(ISCreateGeneral(comm, pdipm->nxlb, ixlb, PETSC_COPY_VALUES, &pdipm->isxlb));
19848a46eb9SPierre Jolivet   if (pdipm->Nxub) PetscCall(ISCreateGeneral(comm, pdipm->nxub, ixub, PETSC_COPY_VALUES, &pdipm->isxub));
19948a46eb9SPierre Jolivet   if (pdipm->Nxfixed) PetscCall(ISCreateGeneral(comm, pdipm->nxfixed, ixfixed, PETSC_COPY_VALUES, &pdipm->isxfixed));
20048a46eb9SPierre Jolivet   if (pdipm->Nxbox) PetscCall(ISCreateGeneral(comm, pdipm->nxbox, ixbox, PETSC_COPY_VALUES, &pdipm->isxbox));
20148a46eb9SPierre Jolivet   if (pdipm->Nxfree) PetscCall(ISCreateGeneral(comm, pdipm->nxfree, ixfree, PETSC_COPY_VALUES, &pdipm->isxfree));
2029566063dSJacob Faibussowitsch   PetscCall(PetscFree5(ixlb, ixub, ixfixed, ixbox, ixfree));
2033ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
204aad13602SShrirang Abhyankar }
205aad13602SShrirang Abhyankar 
206aad13602SShrirang Abhyankar /*
207aad13602SShrirang Abhyankar    TaoPDIPMInitializeSolution - Initialize PDIPM solution X = [x; lambdae; lambdai; z].
208aad13602SShrirang Abhyankar    X consists of four subvectors in the order [x; lambdae; lambdai; z]. These
209aad13602SShrirang Abhyankar      four subvectors need to be initialized and its values copied over to X. Instead
210aad13602SShrirang Abhyankar      of copying, we use VecPlace/ResetArray functions to share the memory locations for
211aad13602SShrirang Abhyankar      X and the subvectors
212aad13602SShrirang Abhyankar 
213aad13602SShrirang Abhyankar    Collective
214aad13602SShrirang Abhyankar 
215aad13602SShrirang Abhyankar    Input Parameter:
216aad13602SShrirang Abhyankar .  tao - Tao context
217aad13602SShrirang Abhyankar 
218aad13602SShrirang Abhyankar    Level: beginner
219aad13602SShrirang Abhyankar */
220d71ae5a4SJacob Faibussowitsch static PetscErrorCode TaoPDIPMInitializeSolution(Tao tao)
221d71ae5a4SJacob Faibussowitsch {
222aad13602SShrirang Abhyankar   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
223aad13602SShrirang Abhyankar   PetscScalar       *Xarr, *z, *lambdai;
224aad13602SShrirang Abhyankar   PetscInt           i;
225aad13602SShrirang Abhyankar   const PetscScalar *xarr, *h;
226aad13602SShrirang Abhyankar 
227aad13602SShrirang Abhyankar   PetscFunctionBegin;
2289566063dSJacob Faibussowitsch   PetscCall(VecGetArrayWrite(pdipm->X, &Xarr));
229aad13602SShrirang Abhyankar 
230aad13602SShrirang Abhyankar   /* Set Initialize X.x = tao->solution */
2319566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(tao->solution, &xarr));
2329566063dSJacob Faibussowitsch   PetscCall(PetscMemcpy(Xarr, xarr, pdipm->nx * sizeof(PetscScalar)));
2339566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(tao->solution, &xarr));
234aad13602SShrirang Abhyankar 
235aad13602SShrirang Abhyankar   /* Initialize X.lambdae = 0.0 */
2361baa6e33SBarry Smith   if (pdipm->lambdae) PetscCall(VecSet(pdipm->lambdae, 0.0));
2377f6ac294SRylee Sundermann 
238aad13602SShrirang Abhyankar   /* Initialize X.lambdai = push_init_lambdai, X.z = push_init_slack */
2397f6ac294SRylee Sundermann   if (pdipm->Nci) {
2409566063dSJacob Faibussowitsch     PetscCall(VecSet(pdipm->lambdai, pdipm->push_init_lambdai));
2419566063dSJacob Faibussowitsch     PetscCall(VecSet(pdipm->z, pdipm->push_init_slack));
242aad13602SShrirang Abhyankar 
243aad13602SShrirang Abhyankar     /* Additional modification for X.lambdai and X.z */
2449566063dSJacob Faibussowitsch     PetscCall(VecGetArrayWrite(pdipm->lambdai, &lambdai));
2459566063dSJacob Faibussowitsch     PetscCall(VecGetArrayWrite(pdipm->z, &z));
246aad13602SShrirang Abhyankar     if (pdipm->Nh) {
2479566063dSJacob Faibussowitsch       PetscCall(VecGetArrayRead(tao->constraints_inequality, &h));
248aad13602SShrirang Abhyankar       for (i = 0; i < pdipm->nh; i++) {
249aad13602SShrirang Abhyankar         if (h[i] < -pdipm->push_init_slack) z[i] = -h[i];
250aad13602SShrirang Abhyankar         if (pdipm->mu / z[i] > pdipm->push_init_lambdai) lambdai[i] = pdipm->mu / z[i];
251aad13602SShrirang Abhyankar       }
2529566063dSJacob Faibussowitsch       PetscCall(VecRestoreArrayRead(tao->constraints_inequality, &h));
253aad13602SShrirang Abhyankar     }
2549566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayWrite(pdipm->lambdai, &lambdai));
2559566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayWrite(pdipm->z, &z));
25652030a5eSPierre Jolivet   }
257aad13602SShrirang Abhyankar 
2589566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayWrite(pdipm->X, &Xarr));
2593ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
260aad13602SShrirang Abhyankar }
261aad13602SShrirang Abhyankar 
262aad13602SShrirang Abhyankar /*
263aad13602SShrirang Abhyankar    TaoSNESJacobian_PDIPM - Evaluate the Hessian matrix at X
264aad13602SShrirang Abhyankar 
265aad13602SShrirang Abhyankar    Input Parameter:
266aad13602SShrirang Abhyankar    snes - SNES context
267aad13602SShrirang Abhyankar    X - KKT Vector
268aad13602SShrirang Abhyankar    *ctx - pdipm context
269aad13602SShrirang Abhyankar 
270aad13602SShrirang Abhyankar    Output Parameter:
271aad13602SShrirang Abhyankar    J - Hessian matrix
272aad13602SShrirang Abhyankar    Jpre - Preconditioner
273aad13602SShrirang Abhyankar */
274d71ae5a4SJacob Faibussowitsch static PetscErrorCode TaoSNESJacobian_PDIPM(SNES snes, Vec X, Mat J, Mat Jpre, void *ctx)
275d71ae5a4SJacob Faibussowitsch {
276aad13602SShrirang Abhyankar   Tao                tao   = (Tao)ctx;
277aad13602SShrirang Abhyankar   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
278aad13602SShrirang Abhyankar   PetscInt           i, row, cols[2], Jrstart, rjstart, nc, j;
279aad13602SShrirang Abhyankar   const PetscInt    *aj, *ranges, *Jranges, *rranges, *cranges;
280aad13602SShrirang Abhyankar   const PetscScalar *Xarr, *aa;
281aad13602SShrirang Abhyankar   PetscScalar        vals[2];
282aad13602SShrirang Abhyankar   PetscInt           proc, nx_all, *nce_all = pdipm->nce_all;
283aad13602SShrirang Abhyankar   MPI_Comm           comm;
284aad13602SShrirang Abhyankar   PetscMPIInt        rank, size;
285aad13602SShrirang Abhyankar 
286aad13602SShrirang Abhyankar   PetscFunctionBegin;
2879566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)snes, &comm));
2889566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(comm, &rank));
2899566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(comm, &size));
290aad13602SShrirang Abhyankar 
2919566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRanges(Jpre, &Jranges));
2929566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(Jpre, &Jrstart, NULL));
2939566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &rranges));
2949566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &cranges));
295aad13602SShrirang Abhyankar 
2969566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(X, &Xarr));
297aad13602SShrirang Abhyankar 
2987f6ac294SRylee Sundermann   /* (1) insert Z and Ci to the 4th block of Jpre -- overwrite existing values */
29912d688e0SRylee Sundermann   if (pdipm->solve_symmetric_kkt) { /* 1 for eq 17 revised pdipm doc 0 for eq 18 (symmetric KKT) */
30012d688e0SRylee Sundermann     vals[0] = 1.0;
30112d688e0SRylee Sundermann     for (i = 0; i < pdipm->nci; i++) {
30212d688e0SRylee Sundermann       row     = Jrstart + pdipm->off_z + i;
30312d688e0SRylee Sundermann       cols[0] = Jrstart + pdipm->off_lambdai + i;
30412d688e0SRylee Sundermann       cols[1] = row;
30512d688e0SRylee Sundermann       vals[1] = Xarr[pdipm->off_lambdai + i] / Xarr[pdipm->off_z + i];
3069566063dSJacob Faibussowitsch       PetscCall(MatSetValues(Jpre, 1, &row, 2, cols, vals, INSERT_VALUES));
30712d688e0SRylee Sundermann     }
30812d688e0SRylee Sundermann   } else {
309aad13602SShrirang Abhyankar     for (i = 0; i < pdipm->nci; i++) {
310aad13602SShrirang Abhyankar       row     = Jrstart + pdipm->off_z + i;
311aad13602SShrirang Abhyankar       cols[0] = Jrstart + pdipm->off_lambdai + i;
312aad13602SShrirang Abhyankar       cols[1] = row;
313aad13602SShrirang Abhyankar       vals[0] = Xarr[pdipm->off_z + i];
314aad13602SShrirang Abhyankar       vals[1] = Xarr[pdipm->off_lambdai + i];
3159566063dSJacob Faibussowitsch       PetscCall(MatSetValues(Jpre, 1, &row, 2, cols, vals, INSERT_VALUES));
316aad13602SShrirang Abhyankar     }
31712d688e0SRylee Sundermann   }
318aad13602SShrirang Abhyankar 
3197f6ac294SRylee Sundermann   /* (2) insert 2nd row block of Jpre: [ grad g, 0, 0, 0] */
320aad13602SShrirang Abhyankar   if (pdipm->Ng) {
3219566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRange(tao->jacobian_equality, &rjstart, NULL));
322aad13602SShrirang Abhyankar     for (i = 0; i < pdipm->ng; i++) {
323aad13602SShrirang Abhyankar       row = Jrstart + pdipm->off_lambdae + i;
324aad13602SShrirang Abhyankar 
3259566063dSJacob Faibussowitsch       PetscCall(MatGetRow(tao->jacobian_equality, i + rjstart, &nc, &aj, &aa));
326aad13602SShrirang Abhyankar       proc = 0;
327aad13602SShrirang Abhyankar       for (j = 0; j < nc; j++) {
328aad13602SShrirang Abhyankar         while (aj[j] >= cranges[proc + 1]) proc++;
329aad13602SShrirang Abhyankar         cols[0] = aj[j] - cranges[proc] + Jranges[proc];
3309566063dSJacob Faibussowitsch         PetscCall(MatSetValue(Jpre, row, cols[0], aa[j], INSERT_VALUES));
331aad13602SShrirang Abhyankar       }
3329566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(tao->jacobian_equality, i + rjstart, &nc, &aj, &aa));
33309ee8bb0SRylee Sundermann       if (pdipm->kkt_pd) {
3347f6ac294SRylee Sundermann         /* add shift \delta_c */
3359566063dSJacob Faibussowitsch         PetscCall(MatSetValue(Jpre, row, row, -pdipm->deltac, INSERT_VALUES));
33609ee8bb0SRylee Sundermann       }
337aad13602SShrirang Abhyankar     }
338aad13602SShrirang Abhyankar   }
339aad13602SShrirang Abhyankar 
340a5b23f4aSJose E. Roman   /* (3) insert 3rd row block of Jpre: [ -grad h, 0, deltac, I] */
341aad13602SShrirang Abhyankar   if (pdipm->Nh) {
3429566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRange(tao->jacobian_inequality, &rjstart, NULL));
343aad13602SShrirang Abhyankar     for (i = 0; i < pdipm->nh; i++) {
344aad13602SShrirang Abhyankar       row = Jrstart + pdipm->off_lambdai + i;
3459566063dSJacob Faibussowitsch       PetscCall(MatGetRow(tao->jacobian_inequality, i + rjstart, &nc, &aj, &aa));
346aad13602SShrirang Abhyankar       proc = 0;
347aad13602SShrirang Abhyankar       for (j = 0; j < nc; j++) {
348aad13602SShrirang Abhyankar         while (aj[j] >= cranges[proc + 1]) proc++;
349aad13602SShrirang Abhyankar         cols[0] = aj[j] - cranges[proc] + Jranges[proc];
3509566063dSJacob Faibussowitsch         PetscCall(MatSetValue(Jpre, row, cols[0], -aa[j], INSERT_VALUES));
351aad13602SShrirang Abhyankar       }
3529566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(tao->jacobian_inequality, i + rjstart, &nc, &aj, &aa));
35309ee8bb0SRylee Sundermann       if (pdipm->kkt_pd) {
3547f6ac294SRylee Sundermann         /* add shift \delta_c */
3559566063dSJacob Faibussowitsch         PetscCall(MatSetValue(Jpre, row, row, -pdipm->deltac, INSERT_VALUES));
35609ee8bb0SRylee Sundermann       }
357aad13602SShrirang Abhyankar     }
358aad13602SShrirang Abhyankar   }
359aad13602SShrirang Abhyankar 
3607f6ac294SRylee Sundermann   /* (4) insert 1st row block of Jpre: [Wxx, grad g', -grad h', 0] */
3617f6ac294SRylee Sundermann   if (pdipm->Ng) { /* grad g' */
3629a8cc538SBarry Smith     PetscCall(MatTranspose(tao->jacobian_equality, MAT_REUSE_MATRIX, &pdipm->jac_equality_trans));
363aad13602SShrirang Abhyankar   }
3647f6ac294SRylee Sundermann   if (pdipm->Nh) { /* grad h' */
3659a8cc538SBarry Smith     PetscCall(MatTranspose(tao->jacobian_inequality, MAT_REUSE_MATRIX, &pdipm->jac_inequality_trans));
366aad13602SShrirang Abhyankar   }
367aad13602SShrirang Abhyankar 
3689566063dSJacob Faibussowitsch   PetscCall(VecPlaceArray(pdipm->x, Xarr));
3699566063dSJacob Faibussowitsch   PetscCall(TaoComputeHessian(tao, pdipm->x, tao->hessian, tao->hessian_pre));
3709566063dSJacob Faibussowitsch   PetscCall(VecResetArray(pdipm->x));
371aad13602SShrirang Abhyankar 
3729566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(tao->hessian, &rjstart, NULL));
373aad13602SShrirang Abhyankar   for (i = 0; i < pdipm->nx; i++) {
374aad13602SShrirang Abhyankar     row = Jrstart + i;
375aad13602SShrirang Abhyankar 
3767f6ac294SRylee Sundermann     /* insert Wxx = fxx + ... -- provided by user */
3779566063dSJacob Faibussowitsch     PetscCall(MatGetRow(tao->hessian, i + rjstart, &nc, &aj, &aa));
378aad13602SShrirang Abhyankar     proc = 0;
379aad13602SShrirang Abhyankar     for (j = 0; j < nc; j++) {
380aad13602SShrirang Abhyankar       while (aj[j] >= cranges[proc + 1]) proc++;
381aad13602SShrirang Abhyankar       cols[0] = aj[j] - cranges[proc] + Jranges[proc];
38209ee8bb0SRylee Sundermann       if (row == cols[0] && pdipm->kkt_pd) {
3837f6ac294SRylee Sundermann         /* add shift deltaw to Wxx component */
3849566063dSJacob Faibussowitsch         PetscCall(MatSetValue(Jpre, row, cols[0], aa[j] + pdipm->deltaw, INSERT_VALUES));
38509ee8bb0SRylee Sundermann       } else {
3869566063dSJacob Faibussowitsch         PetscCall(MatSetValue(Jpre, row, cols[0], aa[j], INSERT_VALUES));
387aad13602SShrirang Abhyankar       }
38809ee8bb0SRylee Sundermann     }
3899566063dSJacob Faibussowitsch     PetscCall(MatRestoreRow(tao->hessian, i + rjstart, &nc, &aj, &aa));
390aad13602SShrirang Abhyankar 
391aad13602SShrirang Abhyankar     /* insert grad g' */
3927f6ac294SRylee Sundermann     if (pdipm->ng) {
3939a8cc538SBarry Smith       PetscCall(MatGetRow(pdipm->jac_equality_trans, i + rjstart, &nc, &aj, &aa));
3949566063dSJacob Faibussowitsch       PetscCall(MatGetOwnershipRanges(tao->jacobian_equality, &ranges));
395aad13602SShrirang Abhyankar       proc = 0;
396aad13602SShrirang Abhyankar       for (j = 0; j < nc; j++) {
397aad13602SShrirang Abhyankar         /* find row ownership of */
398aad13602SShrirang Abhyankar         while (aj[j] >= ranges[proc + 1]) proc++;
399aad13602SShrirang Abhyankar         nx_all  = rranges[proc + 1] - rranges[proc];
400aad13602SShrirang Abhyankar         cols[0] = aj[j] - ranges[proc] + Jranges[proc] + nx_all;
4019566063dSJacob Faibussowitsch         PetscCall(MatSetValue(Jpre, row, cols[0], aa[j], INSERT_VALUES));
402aad13602SShrirang Abhyankar       }
4039a8cc538SBarry Smith       PetscCall(MatRestoreRow(pdipm->jac_equality_trans, i + rjstart, &nc, &aj, &aa));
404aad13602SShrirang Abhyankar     }
405aad13602SShrirang Abhyankar 
406aad13602SShrirang Abhyankar     /* insert -grad h' */
4077f6ac294SRylee Sundermann     if (pdipm->nh) {
4089a8cc538SBarry Smith       PetscCall(MatGetRow(pdipm->jac_inequality_trans, i + rjstart, &nc, &aj, &aa));
4099566063dSJacob Faibussowitsch       PetscCall(MatGetOwnershipRanges(tao->jacobian_inequality, &ranges));
410aad13602SShrirang Abhyankar       proc = 0;
411aad13602SShrirang Abhyankar       for (j = 0; j < nc; j++) {
412aad13602SShrirang Abhyankar         /* find row ownership of */
413aad13602SShrirang Abhyankar         while (aj[j] >= ranges[proc + 1]) proc++;
414aad13602SShrirang Abhyankar         nx_all  = rranges[proc + 1] - rranges[proc];
415aad13602SShrirang Abhyankar         cols[0] = aj[j] - ranges[proc] + Jranges[proc] + nx_all + nce_all[proc];
4169566063dSJacob Faibussowitsch         PetscCall(MatSetValue(Jpre, row, cols[0], -aa[j], INSERT_VALUES));
417aad13602SShrirang Abhyankar       }
4189a8cc538SBarry Smith       PetscCall(MatRestoreRow(pdipm->jac_inequality_trans, i + rjstart, &nc, &aj, &aa));
419aad13602SShrirang Abhyankar     }
420aad13602SShrirang Abhyankar   }
4219566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(X, &Xarr));
422aad13602SShrirang Abhyankar 
423aad13602SShrirang Abhyankar   /* (6) assemble Jpre and J */
4249566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(Jpre, MAT_FINAL_ASSEMBLY));
4259566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(Jpre, MAT_FINAL_ASSEMBLY));
426aad13602SShrirang Abhyankar 
427aad13602SShrirang Abhyankar   if (J != Jpre) {
4289566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
4299566063dSJacob Faibussowitsch     PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
430aad13602SShrirang Abhyankar   }
4313ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
432aad13602SShrirang Abhyankar }
433aad13602SShrirang Abhyankar 
434aad13602SShrirang Abhyankar /*
435aad13602SShrirang Abhyankar    TaoSnesFunction_PDIPM - Evaluate KKT function at X
436aad13602SShrirang Abhyankar 
437aad13602SShrirang Abhyankar    Input Parameter:
438aad13602SShrirang Abhyankar    snes - SNES context
439aad13602SShrirang Abhyankar    X - KKT Vector
440aad13602SShrirang Abhyankar    *ctx - pdipm
441aad13602SShrirang Abhyankar 
442aad13602SShrirang Abhyankar    Output Parameter:
443aad13602SShrirang Abhyankar    F - Updated Lagrangian vector
444aad13602SShrirang Abhyankar */
445d71ae5a4SJacob Faibussowitsch static PetscErrorCode TaoSNESFunction_PDIPM(SNES snes, Vec X, Vec F, void *ctx)
446d71ae5a4SJacob Faibussowitsch {
447aad13602SShrirang Abhyankar   Tao                tao   = (Tao)ctx;
448aad13602SShrirang Abhyankar   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
4497f6ac294SRylee Sundermann   PetscScalar       *Farr;
450aad13602SShrirang Abhyankar   Vec                x, L1;
451aad13602SShrirang Abhyankar   PetscInt           i;
452aad13602SShrirang Abhyankar   const PetscScalar *Xarr, *carr, *zarr, *larr;
453aad13602SShrirang Abhyankar 
454aad13602SShrirang Abhyankar   PetscFunctionBegin;
4559566063dSJacob Faibussowitsch   PetscCall(VecSet(F, 0.0));
456aad13602SShrirang Abhyankar 
4579566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(X, &Xarr));
4589566063dSJacob Faibussowitsch   PetscCall(VecGetArrayWrite(F, &Farr));
459aad13602SShrirang Abhyankar 
4607f6ac294SRylee Sundermann   /* (0) Evaluate f, fx, gradG, gradH at X.x Note: pdipm->x is not changed below */
461aad13602SShrirang Abhyankar   x = pdipm->x;
4629566063dSJacob Faibussowitsch   PetscCall(VecPlaceArray(x, Xarr));
4639566063dSJacob Faibussowitsch   PetscCall(TaoPDIPMEvaluateFunctionsAndJacobians(tao, x));
464aad13602SShrirang Abhyankar 
465aad13602SShrirang Abhyankar   /* Update ce, ci, and Jci at X.x */
4669566063dSJacob Faibussowitsch   PetscCall(TaoPDIPMUpdateConstraints(tao, x));
4679566063dSJacob Faibussowitsch   PetscCall(VecResetArray(x));
468aad13602SShrirang Abhyankar 
469aad13602SShrirang Abhyankar   /* (1) L1 = fx + (gradG'*DE + Jce_xfixed'*lambdae_xfixed) - (gradH'*DI + Jci_xb'*lambdai_xb) */
470aad13602SShrirang Abhyankar   L1 = pdipm->x;
4719566063dSJacob Faibussowitsch   PetscCall(VecPlaceArray(L1, Farr)); /* L1 = 0.0 */
472aad13602SShrirang Abhyankar   if (pdipm->Nci) {
473aad13602SShrirang Abhyankar     if (pdipm->Nh) {
474aad13602SShrirang Abhyankar       /* L1 += gradH'*DI. Note: tao->DI is not changed below */
4759566063dSJacob Faibussowitsch       PetscCall(VecPlaceArray(tao->DI, Xarr + pdipm->off_lambdai));
4769566063dSJacob Faibussowitsch       PetscCall(MatMultTransposeAdd(tao->jacobian_inequality, tao->DI, L1, L1));
4779566063dSJacob Faibussowitsch       PetscCall(VecResetArray(tao->DI));
478aad13602SShrirang Abhyankar     }
479aad13602SShrirang Abhyankar 
480aad13602SShrirang Abhyankar     /* L1 += Jci_xb'*lambdai_xb */
4819566063dSJacob Faibussowitsch     PetscCall(VecPlaceArray(pdipm->lambdai_xb, Xarr + pdipm->off_lambdai + pdipm->nh));
4829566063dSJacob Faibussowitsch     PetscCall(MatMultTransposeAdd(pdipm->Jci_xb, pdipm->lambdai_xb, L1, L1));
4839566063dSJacob Faibussowitsch     PetscCall(VecResetArray(pdipm->lambdai_xb));
484aad13602SShrirang Abhyankar 
4857f6ac294SRylee Sundermann     /* L1 = - (gradH'*DI + Jci_xb'*lambdai_xb) */
4869566063dSJacob Faibussowitsch     PetscCall(VecScale(L1, -1.0));
487aad13602SShrirang Abhyankar   }
488aad13602SShrirang Abhyankar 
489aad13602SShrirang Abhyankar   /* L1 += fx */
4909566063dSJacob Faibussowitsch   PetscCall(VecAXPY(L1, 1.0, tao->gradient));
491aad13602SShrirang Abhyankar 
492aad13602SShrirang Abhyankar   if (pdipm->Nce) {
493aad13602SShrirang Abhyankar     if (pdipm->Ng) {
494aad13602SShrirang Abhyankar       /* L1 += gradG'*DE. Note: tao->DE is not changed below */
4959566063dSJacob Faibussowitsch       PetscCall(VecPlaceArray(tao->DE, Xarr + pdipm->off_lambdae));
4969566063dSJacob Faibussowitsch       PetscCall(MatMultTransposeAdd(tao->jacobian_equality, tao->DE, L1, L1));
4979566063dSJacob Faibussowitsch       PetscCall(VecResetArray(tao->DE));
498aad13602SShrirang Abhyankar     }
499aad13602SShrirang Abhyankar     if (pdipm->Nxfixed) {
500aad13602SShrirang Abhyankar       /* L1 += Jce_xfixed'*lambdae_xfixed */
5019566063dSJacob Faibussowitsch       PetscCall(VecPlaceArray(pdipm->lambdae_xfixed, Xarr + pdipm->off_lambdae + pdipm->ng));
5029566063dSJacob Faibussowitsch       PetscCall(MatMultTransposeAdd(pdipm->Jce_xfixed, pdipm->lambdae_xfixed, L1, L1));
5039566063dSJacob Faibussowitsch       PetscCall(VecResetArray(pdipm->lambdae_xfixed));
504aad13602SShrirang Abhyankar     }
505aad13602SShrirang Abhyankar   }
5069566063dSJacob Faibussowitsch   PetscCall(VecResetArray(L1));
507aad13602SShrirang Abhyankar 
508aad13602SShrirang Abhyankar   /* (2) L2 = ce(x) */
509aad13602SShrirang Abhyankar   if (pdipm->Nce) {
5109566063dSJacob Faibussowitsch     PetscCall(VecGetArrayRead(pdipm->ce, &carr));
511aad13602SShrirang Abhyankar     for (i = 0; i < pdipm->nce; i++) Farr[pdipm->off_lambdae + i] = carr[i];
5129566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayRead(pdipm->ce, &carr));
513aad13602SShrirang Abhyankar   }
514aad13602SShrirang Abhyankar 
515aad13602SShrirang Abhyankar   if (pdipm->Nci) {
51612d688e0SRylee Sundermann     if (pdipm->solve_symmetric_kkt) {
5177f6ac294SRylee Sundermann       /* (3) L3 = z - ci(x);
5187f6ac294SRylee Sundermann          (4) L4 = Lambdai * e - mu/z *e  */
5199566063dSJacob Faibussowitsch       PetscCall(VecGetArrayRead(pdipm->ci, &carr));
52012d688e0SRylee Sundermann       larr = Xarr + pdipm->off_lambdai;
52112d688e0SRylee Sundermann       zarr = Xarr + pdipm->off_z;
52212d688e0SRylee Sundermann       for (i = 0; i < pdipm->nci; i++) {
52312d688e0SRylee Sundermann         Farr[pdipm->off_lambdai + i] = zarr[i] - carr[i];
52412d688e0SRylee Sundermann         Farr[pdipm->off_z + i]       = larr[i] - pdipm->mu / zarr[i];
52512d688e0SRylee Sundermann       }
5269566063dSJacob Faibussowitsch       PetscCall(VecRestoreArrayRead(pdipm->ci, &carr));
52712d688e0SRylee Sundermann     } else {
5287f6ac294SRylee Sundermann       /* (3) L3 = z - ci(x);
5297f6ac294SRylee Sundermann          (4) L4 = Z * Lambdai * e - mu * e  */
5309566063dSJacob Faibussowitsch       PetscCall(VecGetArrayRead(pdipm->ci, &carr));
531aad13602SShrirang Abhyankar       larr = Xarr + pdipm->off_lambdai;
532aad13602SShrirang Abhyankar       zarr = Xarr + pdipm->off_z;
533aad13602SShrirang Abhyankar       for (i = 0; i < pdipm->nci; i++) {
53412d688e0SRylee Sundermann         Farr[pdipm->off_lambdai + i] = zarr[i] - carr[i];
535aad13602SShrirang Abhyankar         Farr[pdipm->off_z + i]       = zarr[i] * larr[i] - pdipm->mu;
536aad13602SShrirang Abhyankar       }
5379566063dSJacob Faibussowitsch       PetscCall(VecRestoreArrayRead(pdipm->ci, &carr));
538aad13602SShrirang Abhyankar     }
53912d688e0SRylee Sundermann   }
540aad13602SShrirang Abhyankar 
5419566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(X, &Xarr));
5429566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayWrite(F, &Farr));
5433ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
5447f6ac294SRylee Sundermann }
545aad13602SShrirang Abhyankar 
5467f6ac294SRylee Sundermann /*
547f560b561SHong Zhang   Evaluate F(X); then update update tao->gnorm0, tao->step = mu,
548f560b561SHong Zhang   tao->residual = norm2(F_x,F_z) and tao->cnorm = norm2(F_ce,F_ci).
5497f6ac294SRylee Sundermann */
550d71ae5a4SJacob Faibussowitsch static PetscErrorCode TaoSNESFunction_PDIPM_residual(SNES snes, Vec X, Vec F, void *ctx)
551d71ae5a4SJacob Faibussowitsch {
5527f6ac294SRylee Sundermann   Tao                tao   = (Tao)ctx;
5537f6ac294SRylee Sundermann   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
5547f6ac294SRylee Sundermann   PetscScalar       *Farr, *tmparr;
5557f6ac294SRylee Sundermann   Vec                L1;
5567f6ac294SRylee Sundermann   PetscInt           i;
5577f6ac294SRylee Sundermann   PetscReal          res[2], cnorm[2];
5587f6ac294SRylee Sundermann   const PetscScalar *Xarr = NULL;
5597f6ac294SRylee Sundermann 
5607f6ac294SRylee Sundermann   PetscFunctionBegin;
5619566063dSJacob Faibussowitsch   PetscCall(TaoSNESFunction_PDIPM(snes, X, F, (void *)tao));
5629566063dSJacob Faibussowitsch   PetscCall(VecGetArrayWrite(F, &Farr));
5639566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(X, &Xarr));
5647f6ac294SRylee Sundermann 
565f560b561SHong Zhang   /* compute res[0] = norm2(F_x) */
5667f6ac294SRylee Sundermann   L1 = pdipm->x;
5679566063dSJacob Faibussowitsch   PetscCall(VecPlaceArray(L1, Farr));
5689566063dSJacob Faibussowitsch   PetscCall(VecNorm(L1, NORM_2, &res[0]));
5699566063dSJacob Faibussowitsch   PetscCall(VecResetArray(L1));
5707f6ac294SRylee Sundermann 
571f560b561SHong Zhang   /* compute res[1] = norm2(F_z), cnorm[1] = norm2(F_ci) */
57252030a5eSPierre Jolivet   if (pdipm->z) {
57312d688e0SRylee Sundermann     if (pdipm->solve_symmetric_kkt) {
5749566063dSJacob Faibussowitsch       PetscCall(VecPlaceArray(pdipm->z, Farr + pdipm->off_z));
57512d688e0SRylee Sundermann       if (pdipm->Nci) {
5769566063dSJacob Faibussowitsch         PetscCall(VecGetArrayWrite(pdipm->z, &tmparr));
57712d688e0SRylee Sundermann         for (i = 0; i < pdipm->nci; i++) tmparr[i] *= Xarr[pdipm->off_z + i];
5789566063dSJacob Faibussowitsch         PetscCall(VecRestoreArrayWrite(pdipm->z, &tmparr));
57912d688e0SRylee Sundermann       }
58012d688e0SRylee Sundermann 
5819566063dSJacob Faibussowitsch       PetscCall(VecNorm(pdipm->z, NORM_2, &res[1]));
5827f6ac294SRylee Sundermann 
58312d688e0SRylee Sundermann       if (pdipm->Nci) {
5849566063dSJacob Faibussowitsch         PetscCall(VecGetArrayWrite(pdipm->z, &tmparr));
585ad540459SPierre Jolivet         for (i = 0; i < pdipm->nci; i++) tmparr[i] /= Xarr[pdipm->off_z + i];
5869566063dSJacob Faibussowitsch         PetscCall(VecRestoreArrayWrite(pdipm->z, &tmparr));
58712d688e0SRylee Sundermann       }
5889566063dSJacob Faibussowitsch       PetscCall(VecResetArray(pdipm->z));
5897f6ac294SRylee Sundermann     } else { /* !solve_symmetric_kkt */
5909566063dSJacob Faibussowitsch       PetscCall(VecPlaceArray(pdipm->z, Farr + pdipm->off_z));
5919566063dSJacob Faibussowitsch       PetscCall(VecNorm(pdipm->z, NORM_2, &res[1]));
5929566063dSJacob Faibussowitsch       PetscCall(VecResetArray(pdipm->z));
59312d688e0SRylee Sundermann     }
594aad13602SShrirang Abhyankar 
5959566063dSJacob Faibussowitsch     PetscCall(VecPlaceArray(pdipm->ci, Farr + pdipm->off_lambdai));
5969566063dSJacob Faibussowitsch     PetscCall(VecNorm(pdipm->ci, NORM_2, &cnorm[1]));
5979566063dSJacob Faibussowitsch     PetscCall(VecResetArray(pdipm->ci));
598f560b561SHong Zhang   } else {
5999371c9d4SSatish Balay     res[1]   = 0.0;
6009371c9d4SSatish Balay     cnorm[1] = 0.0;
601f560b561SHong Zhang   }
6027f6ac294SRylee Sundermann 
603f560b561SHong Zhang   /* compute cnorm[0] = norm2(F_ce) */
6047f6ac294SRylee Sundermann   if (pdipm->Nce) {
6059566063dSJacob Faibussowitsch     PetscCall(VecPlaceArray(pdipm->ce, Farr + pdipm->off_lambdae));
6069566063dSJacob Faibussowitsch     PetscCall(VecNorm(pdipm->ce, NORM_2, &cnorm[0]));
6079566063dSJacob Faibussowitsch     PetscCall(VecResetArray(pdipm->ce));
6087f6ac294SRylee Sundermann   } else cnorm[0] = 0.0;
6097f6ac294SRylee Sundermann 
6109566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayWrite(F, &Farr));
6119566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(X, &Xarr));
612f560b561SHong Zhang 
613f560b561SHong Zhang   tao->gnorm0   = tao->residual;
614f560b561SHong Zhang   tao->residual = PetscSqrtReal(res[0] * res[0] + res[1] * res[1]);
615f560b561SHong Zhang   tao->cnorm    = PetscSqrtReal(cnorm[0] * cnorm[0] + cnorm[1] * cnorm[1]);
616f560b561SHong Zhang   tao->step     = pdipm->mu;
6173ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
618aad13602SShrirang Abhyankar }
619aad13602SShrirang Abhyankar 
620aad13602SShrirang Abhyankar /*
6217f6ac294SRylee Sundermann   KKTAddShifts - Check the inertia of Cholesky factor of KKT matrix.
6227f6ac294SRylee Sundermann   If it does not match the numbers of prime and dual variables, add shifts to the KKT matrix.
623aad13602SShrirang Abhyankar */
624d71ae5a4SJacob Faibussowitsch static PetscErrorCode KKTAddShifts(Tao tao, SNES snes, Vec X)
625d71ae5a4SJacob Faibussowitsch {
626aad13602SShrirang Abhyankar   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;
62709ee8bb0SRylee Sundermann   KSP        ksp;
62809ee8bb0SRylee Sundermann   PC         pc;
62909ee8bb0SRylee Sundermann   Mat        Factor;
63009ee8bb0SRylee Sundermann   PetscBool  isCHOL;
6317f6ac294SRylee Sundermann   PetscInt   nneg, nzero, npos;
632aad13602SShrirang Abhyankar 
633aad13602SShrirang Abhyankar   PetscFunctionBegin;
6347f6ac294SRylee Sundermann   /* Get the inertia of Cholesky factor */
6359566063dSJacob Faibussowitsch   PetscCall(SNESGetKSP(snes, &ksp));
6369566063dSJacob Faibussowitsch   PetscCall(KSPGetPC(ksp, &pc));
6379566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)pc, PCCHOLESKY, &isCHOL));
6383ba16761SJacob Faibussowitsch   if (!isCHOL) PetscFunctionReturn(PETSC_SUCCESS);
63909ee8bb0SRylee Sundermann 
6409566063dSJacob Faibussowitsch   PetscCall(PCFactorGetMatrix(pc, &Factor));
6419566063dSJacob Faibussowitsch   PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));
64209ee8bb0SRylee Sundermann 
64309ee8bb0SRylee Sundermann   if (npos < pdipm->Nx + pdipm->Nci) {
64409ee8bb0SRylee Sundermann     pdipm->deltaw = PetscMax(pdipm->lastdeltaw / 3, 1.e-4 * PETSC_MACHINE_EPSILON);
64563a3b9bcSJacob Faibussowitsch     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));
6469566063dSJacob Faibussowitsch     PetscCall(TaoSNESJacobian_PDIPM(snes, X, pdipm->K, pdipm->K, tao));
6479566063dSJacob Faibussowitsch     PetscCall(PCSetUp(pc));
6489566063dSJacob Faibussowitsch     PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));
64909ee8bb0SRylee Sundermann 
65009ee8bb0SRylee Sundermann     if (npos < pdipm->Nx + pdipm->Nci) {
65109ee8bb0SRylee Sundermann       pdipm->deltaw = pdipm->lastdeltaw;                                           /* in case reduction update does not help, this prevents that step from impacting increasing update */
652f560b561SHong Zhang       while (npos < pdipm->Nx + pdipm->Nci && pdipm->deltaw <= 1. / PETSC_SMALL) { /* increase deltaw */
65363a3b9bcSJacob Faibussowitsch         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));
65409ee8bb0SRylee Sundermann         pdipm->deltaw = PetscMin(8 * pdipm->deltaw, PetscPowReal(10, 20));
6559566063dSJacob Faibussowitsch         PetscCall(TaoSNESJacobian_PDIPM(snes, X, pdipm->K, pdipm->K, tao));
6569566063dSJacob Faibussowitsch         PetscCall(PCSetUp(pc));
6579566063dSJacob Faibussowitsch         PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));
65809ee8bb0SRylee Sundermann       }
65909ee8bb0SRylee Sundermann 
6603c859ba3SBarry Smith       PetscCheck(pdipm->deltaw < 1. / PETSC_SMALL, PetscObjectComm((PetscObject)tao), PETSC_ERR_CONV_FAILED, "Reached maximum delta w will not converge, try different initial x0");
661f560b561SHong Zhang 
6629566063dSJacob Faibussowitsch       PetscCall(PetscInfo(tao, "Updated deltaw %g\n", (double)pdipm->deltaw));
66309ee8bb0SRylee Sundermann       pdipm->lastdeltaw = pdipm->deltaw;
66409ee8bb0SRylee Sundermann       pdipm->deltaw     = 0.0;
66509ee8bb0SRylee Sundermann     }
66609ee8bb0SRylee Sundermann   }
66709ee8bb0SRylee Sundermann 
66809ee8bb0SRylee Sundermann   if (nzero) { /* Jacobian is singular */
66909ee8bb0SRylee Sundermann     if (pdipm->deltac == 0.0) {
6707f6ac294SRylee Sundermann       pdipm->deltac = PETSC_SQRT_MACHINE_EPSILON;
67109ee8bb0SRylee Sundermann     } else {
67209ee8bb0SRylee Sundermann       pdipm->deltac = pdipm->deltac * PetscPowReal(pdipm->mu, .25);
67309ee8bb0SRylee Sundermann     }
67463a3b9bcSJacob Faibussowitsch     PetscCall(PetscInfo(tao, "Updated deltac=%g, MatInertia: nneg %" PetscInt_FMT ", nzero %" PetscInt_FMT "(!=0), npos %" PetscInt_FMT "\n", (double)pdipm->deltac, nneg, nzero, npos));
6759566063dSJacob Faibussowitsch     PetscCall(TaoSNESJacobian_PDIPM(snes, X, pdipm->K, pdipm->K, tao));
6769566063dSJacob Faibussowitsch     PetscCall(PCSetUp(pc));
6779566063dSJacob Faibussowitsch     PetscCall(MatGetInertia(Factor, &nneg, &nzero, &npos));
67809ee8bb0SRylee Sundermann   }
6793ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
6807f6ac294SRylee Sundermann }
6817f6ac294SRylee Sundermann 
6827f6ac294SRylee Sundermann /*
68335cb6cd3SPierre Jolivet   PCPreSolve_PDIPM -- called between MatFactorNumeric() and MatSolve()
6847f6ac294SRylee Sundermann */
685d71ae5a4SJacob Faibussowitsch PetscErrorCode PCPreSolve_PDIPM(PC pc, KSP ksp)
686d71ae5a4SJacob Faibussowitsch {
6877f6ac294SRylee Sundermann   Tao        tao;
6887f6ac294SRylee Sundermann   TAO_PDIPM *pdipm;
6897f6ac294SRylee Sundermann 
6907f6ac294SRylee Sundermann   PetscFunctionBegin;
6919566063dSJacob Faibussowitsch   PetscCall(KSPGetApplicationContext(ksp, &tao));
6927f6ac294SRylee Sundermann   pdipm = (TAO_PDIPM *)tao->data;
6939566063dSJacob Faibussowitsch   PetscCall(KKTAddShifts(tao, pdipm->snes, pdipm->X));
6943ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
6957f6ac294SRylee Sundermann }
6967f6ac294SRylee Sundermann 
6977f6ac294SRylee Sundermann /*
6987f6ac294SRylee Sundermann    SNESLineSearch_PDIPM - Custom line search used with PDIPM.
6997f6ac294SRylee Sundermann 
700c3339decSBarry Smith    Collective
7017f6ac294SRylee Sundermann 
7027f6ac294SRylee Sundermann    Notes:
7037f6ac294SRylee Sundermann    This routine employs a simple backtracking line-search to keep
7047f6ac294SRylee Sundermann    the slack variables (z) and inequality constraints Lagrange multipliers
7057f6ac294SRylee Sundermann    (lambdai) positive, i.e., z,lambdai >=0. It does this by calculating scalars
7067f6ac294SRylee Sundermann    alpha_p and alpha_d to keep z,lambdai non-negative. The decision (x), and the
707f560b561SHong Zhang    slack variables are updated as X = X - alpha_d*dx. The constraint multipliers
7087f6ac294SRylee Sundermann    are updated as Lambdai = Lambdai + alpha_p*dLambdai. The barrier parameter mu
7097f6ac294SRylee Sundermann    is also updated as mu = mu + z'lambdai/Nci
7107f6ac294SRylee Sundermann */
711d71ae5a4SJacob Faibussowitsch static PetscErrorCode SNESLineSearch_PDIPM(SNESLineSearch linesearch, void *ctx)
712d71ae5a4SJacob Faibussowitsch {
7137f6ac294SRylee Sundermann   Tao                tao   = (Tao)ctx;
7147f6ac294SRylee Sundermann   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
7157f6ac294SRylee Sundermann   SNES               snes;
716f560b561SHong Zhang   Vec                X, F, Y;
7177f6ac294SRylee Sundermann   PetscInt           i, iter;
7187f6ac294SRylee Sundermann   PetscReal          alpha_p = 1.0, alpha_d = 1.0, alpha[4];
7197f6ac294SRylee Sundermann   PetscScalar       *Xarr, *z, *lambdai, dot, *taosolarr;
7207f6ac294SRylee Sundermann   const PetscScalar *dXarr, *dz, *dlambdai;
7217f6ac294SRylee Sundermann 
7227f6ac294SRylee Sundermann   PetscFunctionBegin;
7239566063dSJacob Faibussowitsch   PetscCall(SNESLineSearchGetSNES(linesearch, &snes));
7249566063dSJacob Faibussowitsch   PetscCall(SNESGetIterationNumber(snes, &iter));
7257f6ac294SRylee Sundermann 
7269566063dSJacob Faibussowitsch   PetscCall(SNESLineSearchSetReason(linesearch, SNES_LINESEARCH_SUCCEEDED));
7279566063dSJacob Faibussowitsch   PetscCall(SNESLineSearchGetVecs(linesearch, &X, &F, &Y, NULL, NULL));
7287f6ac294SRylee Sundermann 
7299566063dSJacob Faibussowitsch   PetscCall(VecGetArrayWrite(X, &Xarr));
7309566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(Y, &dXarr));
7317f6ac294SRylee Sundermann   z  = Xarr + pdipm->off_z;
7327f6ac294SRylee Sundermann   dz = dXarr + pdipm->off_z;
7337f6ac294SRylee Sundermann   for (i = 0; i < pdipm->nci; i++) {
734f560b561SHong Zhang     if (z[i] - dz[i] < 0.0) alpha_p = PetscMin(alpha_p, 0.9999 * z[i] / dz[i]);
7357f6ac294SRylee Sundermann   }
7367f6ac294SRylee Sundermann 
7377f6ac294SRylee Sundermann   lambdai  = Xarr + pdipm->off_lambdai;
7387f6ac294SRylee Sundermann   dlambdai = dXarr + pdipm->off_lambdai;
7397f6ac294SRylee Sundermann 
7407f6ac294SRylee Sundermann   for (i = 0; i < pdipm->nci; i++) {
741f560b561SHong Zhang     if (lambdai[i] - dlambdai[i] < 0.0) alpha_d = PetscMin(0.9999 * lambdai[i] / dlambdai[i], alpha_d);
7427f6ac294SRylee Sundermann   }
7437f6ac294SRylee Sundermann 
7447f6ac294SRylee Sundermann   alpha[0] = alpha_p;
7457f6ac294SRylee Sundermann   alpha[1] = alpha_d;
7469566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(Y, &dXarr));
7479566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayWrite(X, &Xarr));
7487f6ac294SRylee Sundermann 
7497f6ac294SRylee Sundermann   /* alpha = min(alpha) over all processes */
7509566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Allreduce(alpha, alpha + 2, 2, MPIU_REAL, MPIU_MIN, PetscObjectComm((PetscObject)tao)));
7517f6ac294SRylee Sundermann 
7527f6ac294SRylee Sundermann   alpha_p = alpha[2];
7537f6ac294SRylee Sundermann   alpha_d = alpha[3];
7547f6ac294SRylee Sundermann 
755f560b561SHong Zhang   /* X = X - alpha * Y */
7569566063dSJacob Faibussowitsch   PetscCall(VecGetArrayWrite(X, &Xarr));
7579566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(Y, &dXarr));
7587f6ac294SRylee Sundermann   for (i = 0; i < pdipm->nx; i++) Xarr[i] -= alpha_p * dXarr[i];
759f560b561SHong Zhang   for (i = 0; i < pdipm->nce; i++) Xarr[i + pdipm->off_lambdae] -= alpha_d * dXarr[i + pdipm->off_lambdae];
7607f6ac294SRylee Sundermann 
7617f6ac294SRylee Sundermann   for (i = 0; i < pdipm->nci; i++) {
7627f6ac294SRylee Sundermann     Xarr[i + pdipm->off_lambdai] -= alpha_d * dXarr[i + pdipm->off_lambdai];
7637f6ac294SRylee Sundermann     Xarr[i + pdipm->off_z] -= alpha_p * dXarr[i + pdipm->off_z];
7647f6ac294SRylee Sundermann   }
7659566063dSJacob Faibussowitsch   PetscCall(VecGetArrayWrite(tao->solution, &taosolarr));
7669566063dSJacob Faibussowitsch   PetscCall(PetscMemcpy(taosolarr, Xarr, pdipm->nx * sizeof(PetscScalar)));
7679566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayWrite(tao->solution, &taosolarr));
7687f6ac294SRylee Sundermann 
7699566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayWrite(X, &Xarr));
7709566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(Y, &dXarr));
7717f6ac294SRylee Sundermann 
772f560b561SHong Zhang   /* Update mu = mu_update_factor * dot(z,lambdai)/pdipm->nci at updated X */
7737f6ac294SRylee Sundermann   if (pdipm->z) {
7749566063dSJacob Faibussowitsch     PetscCall(VecDot(pdipm->z, pdipm->lambdai, &dot));
7757f6ac294SRylee Sundermann   } else dot = 0.0;
7767f6ac294SRylee Sundermann 
7777f6ac294SRylee Sundermann   /* if (PetscAbsReal(pdipm->gradL) < 0.9*pdipm->mu)  */
7787f6ac294SRylee Sundermann   pdipm->mu = pdipm->mu_update_factor * dot / pdipm->Nci;
7797f6ac294SRylee Sundermann 
7807f6ac294SRylee Sundermann   /* Update F; get tao->residual and tao->cnorm */
7819566063dSJacob Faibussowitsch   PetscCall(TaoSNESFunction_PDIPM_residual(snes, X, F, (void *)tao));
7827f6ac294SRylee Sundermann 
7837f6ac294SRylee Sundermann   tao->niter++;
7849566063dSJacob Faibussowitsch   PetscCall(TaoLogConvergenceHistory(tao, pdipm->obj, tao->residual, tao->cnorm, tao->niter));
7859566063dSJacob Faibussowitsch   PetscCall(TaoMonitor(tao, tao->niter, pdipm->obj, tao->residual, tao->cnorm, pdipm->mu));
7867f6ac294SRylee Sundermann 
787dbbe0bcdSBarry Smith   PetscUseTypeMethod(tao, convergencetest, tao->cnvP);
7881baa6e33SBarry Smith   if (tao->reason) PetscCall(SNESSetConvergedReason(snes, SNES_CONVERGED_FNORM_ABS));
7893ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
790aad13602SShrirang Abhyankar }
791aad13602SShrirang Abhyankar 
792d71ae5a4SJacob Faibussowitsch PetscErrorCode TaoSolve_PDIPM(Tao tao)
793d71ae5a4SJacob Faibussowitsch {
794aad13602SShrirang Abhyankar   TAO_PDIPM     *pdipm = (TAO_PDIPM *)tao->data;
795aad13602SShrirang Abhyankar   SNESLineSearch linesearch; /* SNESLineSearch context */
796aad13602SShrirang Abhyankar   Vec            dummy;
797aad13602SShrirang Abhyankar 
798aad13602SShrirang Abhyankar   PetscFunctionBegin;
7993c859ba3SBarry Smith   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");
8008e58fa1dSresundermann 
801aad13602SShrirang Abhyankar   /* Initialize all variables */
8029566063dSJacob Faibussowitsch   PetscCall(TaoPDIPMInitializeSolution(tao));
803aad13602SShrirang Abhyankar 
804aad13602SShrirang Abhyankar   /* Set linesearch */
8059566063dSJacob Faibussowitsch   PetscCall(SNESGetLineSearch(pdipm->snes, &linesearch));
8069566063dSJacob Faibussowitsch   PetscCall(SNESLineSearchSetType(linesearch, SNESLINESEARCHSHELL));
8079566063dSJacob Faibussowitsch   PetscCall(SNESLineSearchShellSetUserFunc(linesearch, SNESLineSearch_PDIPM, tao));
8089566063dSJacob Faibussowitsch   PetscCall(SNESLineSearchSetFromOptions(linesearch));
809aad13602SShrirang Abhyankar 
810aad13602SShrirang Abhyankar   tao->reason = TAO_CONTINUE_ITERATING;
811aad13602SShrirang Abhyankar 
812aad13602SShrirang Abhyankar   /* -tao_monitor for iteration 0 and check convergence */
8139566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(pdipm->X, &dummy));
8149566063dSJacob Faibussowitsch   PetscCall(TaoSNESFunction_PDIPM_residual(pdipm->snes, pdipm->X, dummy, (void *)tao));
815aad13602SShrirang Abhyankar 
8169566063dSJacob Faibussowitsch   PetscCall(TaoLogConvergenceHistory(tao, pdipm->obj, tao->residual, tao->cnorm, tao->niter));
8179566063dSJacob Faibussowitsch   PetscCall(TaoMonitor(tao, tao->niter, pdipm->obj, tao->residual, tao->cnorm, pdipm->mu));
8189566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&dummy));
819dbbe0bcdSBarry Smith   PetscUseTypeMethod(tao, convergencetest, tao->cnvP);
82063a3b9bcSJacob Faibussowitsch   if (tao->reason) PetscCall(SNESSetConvergedReason(pdipm->snes, SNES_CONVERGED_FNORM_ABS));
821aad13602SShrirang Abhyankar 
822aad13602SShrirang Abhyankar   while (tao->reason == TAO_CONTINUE_ITERATING) {
823aad13602SShrirang Abhyankar     SNESConvergedReason reason;
8249566063dSJacob Faibussowitsch     PetscCall(SNESSolve(pdipm->snes, NULL, pdipm->X));
825aad13602SShrirang Abhyankar 
826aad13602SShrirang Abhyankar     /* Check SNES convergence */
8279566063dSJacob Faibussowitsch     PetscCall(SNESGetConvergedReason(pdipm->snes, &reason));
82848a46eb9SPierre Jolivet     if (reason < 0) PetscCall(PetscPrintf(PetscObjectComm((PetscObject)pdipm->snes), "SNES solve did not converged due to reason %s\n", SNESConvergedReasons[reason]));
829aad13602SShrirang Abhyankar 
830aad13602SShrirang Abhyankar     /* Check TAO convergence */
8313c859ba3SBarry Smith     PetscCheck(!PetscIsInfOrNanReal(pdipm->obj), PETSC_COMM_SELF, PETSC_ERR_SUP, "User-provided compute function generated Inf or NaN");
832aad13602SShrirang Abhyankar   }
8333ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
834aad13602SShrirang Abhyankar }
835aad13602SShrirang Abhyankar 
836aad13602SShrirang Abhyankar /*
83770c9796eSresundermann   TaoView_PDIPM - View PDIPM
83870c9796eSresundermann 
83970c9796eSresundermann    Input Parameter:
84070c9796eSresundermann     tao - TAO object
84170c9796eSresundermann     viewer - PetscViewer
84270c9796eSresundermann 
84370c9796eSresundermann    Output:
84470c9796eSresundermann */
845d71ae5a4SJacob Faibussowitsch PetscErrorCode TaoView_PDIPM(Tao tao, PetscViewer viewer)
846d71ae5a4SJacob Faibussowitsch {
84770c9796eSresundermann   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;
84870c9796eSresundermann 
84970c9796eSresundermann   PetscFunctionBegin;
85070c9796eSresundermann   tao->constrained = PETSC_TRUE;
8519566063dSJacob Faibussowitsch   PetscCall(PetscViewerASCIIPushTab(viewer));
85263a3b9bcSJacob Faibussowitsch   PetscCall(PetscViewerASCIIPrintf(viewer, "Number of prime=%" PetscInt_FMT ", Number of dual=%" PetscInt_FMT "\n", pdipm->Nx + pdipm->Nci, pdipm->Nce + pdipm->Nci));
85348a46eb9SPierre Jolivet   if (pdipm->kkt_pd) PetscCall(PetscViewerASCIIPrintf(viewer, "KKT shifts deltaw=%g, deltac=%g\n", (double)pdipm->deltaw, (double)pdipm->deltac));
8549566063dSJacob Faibussowitsch   PetscCall(PetscViewerASCIIPopTab(viewer));
8553ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
85670c9796eSresundermann }
85770c9796eSresundermann 
85870c9796eSresundermann /*
859aad13602SShrirang Abhyankar    TaoSetup_PDIPM - Sets up tao and pdipm
860aad13602SShrirang Abhyankar 
861aad13602SShrirang Abhyankar    Input Parameter:
862aad13602SShrirang Abhyankar    tao - TAO object
863aad13602SShrirang Abhyankar 
864aad13602SShrirang Abhyankar    Output:   pdipm - initialized object
865aad13602SShrirang Abhyankar */
866d71ae5a4SJacob Faibussowitsch PetscErrorCode TaoSetup_PDIPM(Tao tao)
867d71ae5a4SJacob Faibussowitsch {
868aad13602SShrirang Abhyankar   TAO_PDIPM         *pdipm = (TAO_PDIPM *)tao->data;
869aad13602SShrirang Abhyankar   MPI_Comm           comm;
870f560b561SHong Zhang   PetscMPIInt        size;
871aad13602SShrirang Abhyankar   PetscInt           row, col, Jcrstart, Jcrend, k, tmp, nc, proc, *nh_all, *ng_all;
872aad13602SShrirang Abhyankar   PetscInt           offset, *xa, *xb, i, j, rstart, rend;
8737f6ac294SRylee Sundermann   PetscScalar        one = 1.0, neg_one = -1.0;
874aad13602SShrirang Abhyankar   const PetscInt    *cols, *rranges, *cranges, *aj, *ranges;
8757f6ac294SRylee Sundermann   const PetscScalar *aa, *Xarr;
8769a8cc538SBarry Smith   Mat                J;
877aad13602SShrirang Abhyankar   Mat                Jce_xfixed_trans, Jci_xb_trans;
878aad13602SShrirang Abhyankar   PetscInt          *dnz, *onz, rjstart, nx_all, *nce_all, *Jranges, cols1[2];
879aad13602SShrirang Abhyankar 
880aad13602SShrirang Abhyankar   PetscFunctionBegin;
8819566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)tao, &comm));
8829566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
883aad13602SShrirang Abhyankar 
884aad13602SShrirang Abhyankar   /* (1) Setup Bounds and create Tao vectors */
8859566063dSJacob Faibussowitsch   PetscCall(TaoPDIPMSetUpBounds(tao));
886aad13602SShrirang Abhyankar 
887aad13602SShrirang Abhyankar   if (!tao->gradient) {
8889566063dSJacob Faibussowitsch     PetscCall(VecDuplicate(tao->solution, &tao->gradient));
8899566063dSJacob Faibussowitsch     PetscCall(VecDuplicate(tao->solution, &tao->stepdirection));
890aad13602SShrirang Abhyankar   }
891aad13602SShrirang Abhyankar 
892aad13602SShrirang Abhyankar   /* (2) Get sizes */
893a82e8c82SStefano Zampini   /* Size of vector x - This is set by TaoSetSolution */
8949566063dSJacob Faibussowitsch   PetscCall(VecGetSize(tao->solution, &pdipm->Nx));
8959566063dSJacob Faibussowitsch   PetscCall(VecGetLocalSize(tao->solution, &pdipm->nx));
896aad13602SShrirang Abhyankar 
897aad13602SShrirang Abhyankar   /* Size of equality constraints and vectors */
898aad13602SShrirang Abhyankar   if (tao->constraints_equality) {
8999566063dSJacob Faibussowitsch     PetscCall(VecGetSize(tao->constraints_equality, &pdipm->Ng));
9009566063dSJacob Faibussowitsch     PetscCall(VecGetLocalSize(tao->constraints_equality, &pdipm->ng));
901aad13602SShrirang Abhyankar   } else {
902aad13602SShrirang Abhyankar     pdipm->ng = pdipm->Ng = 0;
903aad13602SShrirang Abhyankar   }
904aad13602SShrirang Abhyankar 
905aad13602SShrirang Abhyankar   pdipm->nce = pdipm->ng + pdipm->nxfixed;
906aad13602SShrirang Abhyankar   pdipm->Nce = pdipm->Ng + pdipm->Nxfixed;
907aad13602SShrirang Abhyankar 
908aad13602SShrirang Abhyankar   /* Size of inequality constraints and vectors */
909aad13602SShrirang Abhyankar   if (tao->constraints_inequality) {
9109566063dSJacob Faibussowitsch     PetscCall(VecGetSize(tao->constraints_inequality, &pdipm->Nh));
9119566063dSJacob Faibussowitsch     PetscCall(VecGetLocalSize(tao->constraints_inequality, &pdipm->nh));
912aad13602SShrirang Abhyankar   } else {
913aad13602SShrirang Abhyankar     pdipm->nh = pdipm->Nh = 0;
914aad13602SShrirang Abhyankar   }
915aad13602SShrirang Abhyankar 
916aad13602SShrirang Abhyankar   pdipm->nci = pdipm->nh + pdipm->nxlb + pdipm->nxub + 2 * pdipm->nxbox;
917aad13602SShrirang Abhyankar   pdipm->Nci = pdipm->Nh + pdipm->Nxlb + pdipm->Nxub + 2 * pdipm->Nxbox;
918aad13602SShrirang Abhyankar 
919aad13602SShrirang Abhyankar   /* Full size of the KKT system to be solved */
920aad13602SShrirang Abhyankar   pdipm->n = pdipm->nx + pdipm->nce + 2 * pdipm->nci;
921aad13602SShrirang Abhyankar   pdipm->N = pdipm->Nx + pdipm->Nce + 2 * pdipm->Nci;
922aad13602SShrirang Abhyankar 
923aad13602SShrirang Abhyankar   /* (3) Offsets for subvectors */
924aad13602SShrirang Abhyankar   pdipm->off_lambdae = pdipm->nx;
925aad13602SShrirang Abhyankar   pdipm->off_lambdai = pdipm->off_lambdae + pdipm->nce;
926aad13602SShrirang Abhyankar   pdipm->off_z       = pdipm->off_lambdai + pdipm->nci;
927aad13602SShrirang Abhyankar 
928aad13602SShrirang Abhyankar   /* (4) Create vectors and subvectors */
929aad13602SShrirang Abhyankar   /* Ce and Ci vectors */
9309566063dSJacob Faibussowitsch   PetscCall(VecCreate(comm, &pdipm->ce));
9319566063dSJacob Faibussowitsch   PetscCall(VecSetSizes(pdipm->ce, pdipm->nce, pdipm->Nce));
9329566063dSJacob Faibussowitsch   PetscCall(VecSetFromOptions(pdipm->ce));
933aad13602SShrirang Abhyankar 
9349566063dSJacob Faibussowitsch   PetscCall(VecCreate(comm, &pdipm->ci));
9359566063dSJacob Faibussowitsch   PetscCall(VecSetSizes(pdipm->ci, pdipm->nci, pdipm->Nci));
9369566063dSJacob Faibussowitsch   PetscCall(VecSetFromOptions(pdipm->ci));
937aad13602SShrirang Abhyankar 
938aad13602SShrirang Abhyankar   /* X=[x; lambdae; lambdai; z] for the big KKT system */
9399566063dSJacob Faibussowitsch   PetscCall(VecCreate(comm, &pdipm->X));
9409566063dSJacob Faibussowitsch   PetscCall(VecSetSizes(pdipm->X, pdipm->n, pdipm->N));
9419566063dSJacob Faibussowitsch   PetscCall(VecSetFromOptions(pdipm->X));
942aad13602SShrirang Abhyankar 
943aad13602SShrirang Abhyankar   /* Subvectors; they share local arrays with X */
9449566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(pdipm->X, &Xarr));
945aad13602SShrirang Abhyankar   /* x shares local array with X.x */
94648a46eb9SPierre Jolivet   if (pdipm->Nx) PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nx, pdipm->Nx, Xarr, &pdipm->x));
947aad13602SShrirang Abhyankar 
948aad13602SShrirang Abhyankar   /* lambdae shares local array with X.lambdae */
94948a46eb9SPierre Jolivet   if (pdipm->Nce) PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nce, pdipm->Nce, Xarr + pdipm->off_lambdae, &pdipm->lambdae));
950aad13602SShrirang Abhyankar 
951aad13602SShrirang Abhyankar   /* tao->DE shares local array with X.lambdae_g */
952aad13602SShrirang Abhyankar   if (pdipm->Ng) {
9539566063dSJacob Faibussowitsch     PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->ng, pdipm->Ng, Xarr + pdipm->off_lambdae, &tao->DE));
954aad13602SShrirang Abhyankar 
9559566063dSJacob Faibussowitsch     PetscCall(VecCreate(comm, &pdipm->lambdae_xfixed));
9569566063dSJacob Faibussowitsch     PetscCall(VecSetSizes(pdipm->lambdae_xfixed, pdipm->nxfixed, PETSC_DECIDE));
9579566063dSJacob Faibussowitsch     PetscCall(VecSetFromOptions(pdipm->lambdae_xfixed));
958aad13602SShrirang Abhyankar   }
959aad13602SShrirang Abhyankar 
960aad13602SShrirang Abhyankar   if (pdipm->Nci) {
961aad13602SShrirang Abhyankar     /* lambdai shares local array with X.lambdai */
9629566063dSJacob Faibussowitsch     PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nci, pdipm->Nci, Xarr + pdipm->off_lambdai, &pdipm->lambdai));
963aad13602SShrirang Abhyankar 
964aad13602SShrirang Abhyankar     /* z for slack variables; it shares local array with X.z */
9659566063dSJacob Faibussowitsch     PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nci, pdipm->Nci, Xarr + pdipm->off_z, &pdipm->z));
966aad13602SShrirang Abhyankar   }
967aad13602SShrirang Abhyankar 
968aad13602SShrirang Abhyankar   /* tao->DI which shares local array with X.lambdai_h */
96948a46eb9SPierre Jolivet   if (pdipm->Nh) PetscCall(VecCreateMPIWithArray(comm, 1, pdipm->nh, pdipm->Nh, Xarr + pdipm->off_lambdai, &tao->DI));
9709566063dSJacob Faibussowitsch   PetscCall(VecCreate(comm, &pdipm->lambdai_xb));
9719566063dSJacob Faibussowitsch   PetscCall(VecSetSizes(pdipm->lambdai_xb, (pdipm->nci - pdipm->nh), PETSC_DECIDE));
9729566063dSJacob Faibussowitsch   PetscCall(VecSetFromOptions(pdipm->lambdai_xb));
973aad13602SShrirang Abhyankar 
9749566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(pdipm->X, &Xarr));
975aad13602SShrirang Abhyankar 
976aad13602SShrirang Abhyankar   /* (5) Create Jacobians Jce_xfixed and Jci */
977aad13602SShrirang Abhyankar   /* (5.1) PDIPM Jacobian of equality bounds cebound(x) = J_nxfixed */
978aad13602SShrirang Abhyankar   if (pdipm->Nxfixed) {
979aad13602SShrirang Abhyankar     /* Create Jce_xfixed */
9809566063dSJacob Faibussowitsch     PetscCall(MatCreate(comm, &pdipm->Jce_xfixed));
9819566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(pdipm->Jce_xfixed, pdipm->nxfixed, pdipm->nx, PETSC_DECIDE, pdipm->Nx));
9829566063dSJacob Faibussowitsch     PetscCall(MatSetFromOptions(pdipm->Jce_xfixed));
9839566063dSJacob Faibussowitsch     PetscCall(MatSeqAIJSetPreallocation(pdipm->Jce_xfixed, 1, NULL));
9849566063dSJacob Faibussowitsch     PetscCall(MatMPIAIJSetPreallocation(pdipm->Jce_xfixed, 1, NULL, 1, NULL));
985aad13602SShrirang Abhyankar 
9869566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRange(pdipm->Jce_xfixed, &Jcrstart, &Jcrend));
9879566063dSJacob Faibussowitsch     PetscCall(ISGetIndices(pdipm->isxfixed, &cols));
988aad13602SShrirang Abhyankar     k = 0;
989aad13602SShrirang Abhyankar     for (row = Jcrstart; row < Jcrend; row++) {
9909566063dSJacob Faibussowitsch       PetscCall(MatSetValues(pdipm->Jce_xfixed, 1, &row, 1, cols + k, &one, INSERT_VALUES));
991aad13602SShrirang Abhyankar       k++;
992aad13602SShrirang Abhyankar     }
9939566063dSJacob Faibussowitsch     PetscCall(ISRestoreIndices(pdipm->isxfixed, &cols));
9949566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(pdipm->Jce_xfixed, MAT_FINAL_ASSEMBLY));
9959566063dSJacob Faibussowitsch     PetscCall(MatAssemblyEnd(pdipm->Jce_xfixed, MAT_FINAL_ASSEMBLY));
996aad13602SShrirang Abhyankar   }
997aad13602SShrirang Abhyankar 
998aad13602SShrirang Abhyankar   /* (5.2) PDIPM inequality Jacobian Jci = [tao->jacobian_inequality; ...] */
9999566063dSJacob Faibussowitsch   PetscCall(MatCreate(comm, &pdipm->Jci_xb));
10009566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(pdipm->Jci_xb, pdipm->nci - pdipm->nh, pdipm->nx, PETSC_DECIDE, pdipm->Nx));
10019566063dSJacob Faibussowitsch   PetscCall(MatSetFromOptions(pdipm->Jci_xb));
10029566063dSJacob Faibussowitsch   PetscCall(MatSeqAIJSetPreallocation(pdipm->Jci_xb, 1, NULL));
10039566063dSJacob Faibussowitsch   PetscCall(MatMPIAIJSetPreallocation(pdipm->Jci_xb, 1, NULL, 1, NULL));
1004aad13602SShrirang Abhyankar 
10059566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(pdipm->Jci_xb, &Jcrstart, &Jcrend));
1006aad13602SShrirang Abhyankar   offset = Jcrstart;
1007aad13602SShrirang Abhyankar   if (pdipm->Nxub) {
1008aad13602SShrirang Abhyankar     /* Add xub to Jci_xb */
10099566063dSJacob Faibussowitsch     PetscCall(ISGetIndices(pdipm->isxub, &cols));
1010aad13602SShrirang Abhyankar     k = 0;
1011aad13602SShrirang Abhyankar     for (row = offset; row < offset + pdipm->nxub; row++) {
10129566063dSJacob Faibussowitsch       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &row, 1, cols + k, &neg_one, INSERT_VALUES));
1013aad13602SShrirang Abhyankar       k++;
1014aad13602SShrirang Abhyankar     }
10159566063dSJacob Faibussowitsch     PetscCall(ISRestoreIndices(pdipm->isxub, &cols));
1016aad13602SShrirang Abhyankar   }
1017aad13602SShrirang Abhyankar 
1018aad13602SShrirang Abhyankar   if (pdipm->Nxlb) {
1019aad13602SShrirang Abhyankar     /* Add xlb to Jci_xb */
10209566063dSJacob Faibussowitsch     PetscCall(ISGetIndices(pdipm->isxlb, &cols));
1021aad13602SShrirang Abhyankar     k = 0;
1022aad13602SShrirang Abhyankar     offset += pdipm->nxub;
1023aad13602SShrirang Abhyankar     for (row = offset; row < offset + pdipm->nxlb; row++) {
10249566063dSJacob Faibussowitsch       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &row, 1, cols + k, &one, INSERT_VALUES));
1025aad13602SShrirang Abhyankar       k++;
1026aad13602SShrirang Abhyankar     }
10279566063dSJacob Faibussowitsch     PetscCall(ISRestoreIndices(pdipm->isxlb, &cols));
1028aad13602SShrirang Abhyankar   }
1029aad13602SShrirang Abhyankar 
1030aad13602SShrirang Abhyankar   /* Add xbox to Jci_xb */
1031aad13602SShrirang Abhyankar   if (pdipm->Nxbox) {
10329566063dSJacob Faibussowitsch     PetscCall(ISGetIndices(pdipm->isxbox, &cols));
1033aad13602SShrirang Abhyankar     k = 0;
1034aad13602SShrirang Abhyankar     offset += pdipm->nxlb;
1035aad13602SShrirang Abhyankar     for (row = offset; row < offset + pdipm->nxbox; row++) {
10369566063dSJacob Faibussowitsch       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &row, 1, cols + k, &neg_one, INSERT_VALUES));
1037aad13602SShrirang Abhyankar       tmp = row + pdipm->nxbox;
10389566063dSJacob Faibussowitsch       PetscCall(MatSetValues(pdipm->Jci_xb, 1, &tmp, 1, cols + k, &one, INSERT_VALUES));
1039aad13602SShrirang Abhyankar       k++;
1040aad13602SShrirang Abhyankar     }
10419566063dSJacob Faibussowitsch     PetscCall(ISRestoreIndices(pdipm->isxbox, &cols));
1042aad13602SShrirang Abhyankar   }
1043aad13602SShrirang Abhyankar 
10449566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(pdipm->Jci_xb, MAT_FINAL_ASSEMBLY));
10459566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(pdipm->Jci_xb, MAT_FINAL_ASSEMBLY));
10469566063dSJacob Faibussowitsch   /* PetscCall(MatView(pdipm->Jci_xb,PETSC_VIEWER_STDOUT_WORLD)); */
1047aad13602SShrirang Abhyankar 
1048aad13602SShrirang Abhyankar   /* (6) Set up ISs for PC Fieldsplit */
1049aad13602SShrirang Abhyankar   if (pdipm->solve_reduced_kkt) {
10509566063dSJacob Faibussowitsch     PetscCall(PetscMalloc2(pdipm->nx + pdipm->nce, &xa, 2 * pdipm->nci, &xb));
1051aad13602SShrirang Abhyankar     for (i = 0; i < pdipm->nx + pdipm->nce; i++) xa[i] = i;
1052aad13602SShrirang Abhyankar     for (i = 0; i < 2 * pdipm->nci; i++) xb[i] = pdipm->off_lambdai + i;
1053aad13602SShrirang Abhyankar 
10549566063dSJacob Faibussowitsch     PetscCall(ISCreateGeneral(comm, pdipm->nx + pdipm->nce, xa, PETSC_OWN_POINTER, &pdipm->is1));
10559566063dSJacob Faibussowitsch     PetscCall(ISCreateGeneral(comm, 2 * pdipm->nci, xb, PETSC_OWN_POINTER, &pdipm->is2));
1056aad13602SShrirang Abhyankar   }
1057aad13602SShrirang Abhyankar 
1058aad13602SShrirang Abhyankar   /* (7) Gather offsets from all processes */
10599566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(size, &pdipm->nce_all));
1060aad13602SShrirang Abhyankar 
1061aad13602SShrirang Abhyankar   /* Get rstart of KKT matrix */
10629566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Scan(&pdipm->n, &rstart, 1, MPIU_INT, MPI_SUM, comm));
1063aad13602SShrirang Abhyankar   rstart -= pdipm->n;
1064aad13602SShrirang Abhyankar 
10659566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Allgather(&pdipm->nce, 1, MPIU_INT, pdipm->nce_all, 1, MPIU_INT, comm));
1066aad13602SShrirang Abhyankar 
10679566063dSJacob Faibussowitsch   PetscCall(PetscMalloc3(size, &ng_all, size, &nh_all, size, &Jranges));
10689566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Allgather(&rstart, 1, MPIU_INT, Jranges, 1, MPIU_INT, comm));
10699566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Allgather(&pdipm->nh, 1, MPIU_INT, nh_all, 1, MPIU_INT, comm));
10709566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Allgather(&pdipm->ng, 1, MPIU_INT, ng_all, 1, MPIU_INT, comm));
1071aad13602SShrirang Abhyankar 
10729566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRanges(tao->hessian, &rranges));
10739566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRangesColumn(tao->hessian, &cranges));
1074aad13602SShrirang Abhyankar 
1075aad13602SShrirang Abhyankar   if (pdipm->Ng) {
10769566063dSJacob Faibussowitsch     PetscCall(TaoComputeJacobianEquality(tao, tao->solution, tao->jacobian_equality, tao->jacobian_equality_pre));
10779566063dSJacob Faibussowitsch     PetscCall(MatTranspose(tao->jacobian_equality, MAT_INITIAL_MATRIX, &pdipm->jac_equality_trans));
1078aad13602SShrirang Abhyankar   }
1079aad13602SShrirang Abhyankar   if (pdipm->Nh) {
10809566063dSJacob Faibussowitsch     PetscCall(TaoComputeJacobianInequality(tao, tao->solution, tao->jacobian_inequality, tao->jacobian_inequality_pre));
10819566063dSJacob Faibussowitsch     PetscCall(MatTranspose(tao->jacobian_inequality, MAT_INITIAL_MATRIX, &pdipm->jac_inequality_trans));
1082aad13602SShrirang Abhyankar   }
1083aad13602SShrirang Abhyankar 
1084aad13602SShrirang Abhyankar   /* Count dnz,onz for preallocation of KKT matrix */
1085aad13602SShrirang Abhyankar   nce_all = pdipm->nce_all;
1086aad13602SShrirang Abhyankar 
108748a46eb9SPierre Jolivet   if (pdipm->Nxfixed) PetscCall(MatTranspose(pdipm->Jce_xfixed, MAT_INITIAL_MATRIX, &Jce_xfixed_trans));
10889566063dSJacob Faibussowitsch   PetscCall(MatTranspose(pdipm->Jci_xb, MAT_INITIAL_MATRIX, &Jci_xb_trans));
1089aad13602SShrirang Abhyankar 
1090d0609cedSBarry Smith   MatPreallocateBegin(comm, pdipm->n, pdipm->n, dnz, onz);
1091aad13602SShrirang Abhyankar 
1092aad13602SShrirang Abhyankar   /* 1st row block of KKT matrix: [Wxx; gradCe'; -gradCi'; 0] */
10939566063dSJacob Faibussowitsch   PetscCall(TaoPDIPMEvaluateFunctionsAndJacobians(tao, pdipm->x));
10949566063dSJacob Faibussowitsch   PetscCall(TaoComputeHessian(tao, tao->solution, tao->hessian, tao->hessian_pre));
1095aad13602SShrirang Abhyankar 
1096aad13602SShrirang Abhyankar   /* Insert tao->hessian */
10979566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(tao->hessian, &rjstart, NULL));
1098aad13602SShrirang Abhyankar   for (i = 0; i < pdipm->nx; i++) {
1099aad13602SShrirang Abhyankar     row = rstart + i;
1100aad13602SShrirang Abhyankar 
11019566063dSJacob Faibussowitsch     PetscCall(MatGetRow(tao->hessian, i + rjstart, &nc, &aj, NULL));
1102aad13602SShrirang Abhyankar     proc = 0;
1103aad13602SShrirang Abhyankar     for (j = 0; j < nc; j++) {
1104aad13602SShrirang Abhyankar       while (aj[j] >= cranges[proc + 1]) proc++;
1105aad13602SShrirang Abhyankar       col = aj[j] - cranges[proc] + Jranges[proc];
11069566063dSJacob Faibussowitsch       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1107aad13602SShrirang Abhyankar     }
11089566063dSJacob Faibussowitsch     PetscCall(MatRestoreRow(tao->hessian, i + rjstart, &nc, &aj, NULL));
1109aad13602SShrirang Abhyankar 
1110aad13602SShrirang Abhyankar     if (pdipm->ng) {
1111aad13602SShrirang Abhyankar       /* Insert grad g' */
11129a8cc538SBarry Smith       PetscCall(MatGetRow(pdipm->jac_equality_trans, i + rjstart, &nc, &aj, NULL));
11139566063dSJacob Faibussowitsch       PetscCall(MatGetOwnershipRanges(tao->jacobian_equality, &ranges));
1114aad13602SShrirang Abhyankar       proc = 0;
1115aad13602SShrirang Abhyankar       for (j = 0; j < nc; j++) {
1116aad13602SShrirang Abhyankar         /* find row ownership of */
1117aad13602SShrirang Abhyankar         while (aj[j] >= ranges[proc + 1]) proc++;
1118aad13602SShrirang Abhyankar         nx_all = rranges[proc + 1] - rranges[proc];
1119aad13602SShrirang Abhyankar         col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all;
11209566063dSJacob Faibussowitsch         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1121aad13602SShrirang Abhyankar       }
11229a8cc538SBarry Smith       PetscCall(MatRestoreRow(pdipm->jac_equality_trans, i + rjstart, &nc, &aj, NULL));
1123aad13602SShrirang Abhyankar     }
1124aad13602SShrirang Abhyankar 
1125aad13602SShrirang Abhyankar     /* Insert Jce_xfixed^T' */
1126aad13602SShrirang Abhyankar     if (pdipm->nxfixed) {
11279566063dSJacob Faibussowitsch       PetscCall(MatGetRow(Jce_xfixed_trans, i + rjstart, &nc, &aj, NULL));
11289566063dSJacob Faibussowitsch       PetscCall(MatGetOwnershipRanges(pdipm->Jce_xfixed, &ranges));
1129aad13602SShrirang Abhyankar       proc = 0;
1130aad13602SShrirang Abhyankar       for (j = 0; j < nc; j++) {
1131aad13602SShrirang Abhyankar         /* find row ownership of */
1132aad13602SShrirang Abhyankar         while (aj[j] >= ranges[proc + 1]) proc++;
1133aad13602SShrirang Abhyankar         nx_all = rranges[proc + 1] - rranges[proc];
1134aad13602SShrirang Abhyankar         col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all + ng_all[proc];
11359566063dSJacob Faibussowitsch         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1136aad13602SShrirang Abhyankar       }
11379566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(Jce_xfixed_trans, i + rjstart, &nc, &aj, NULL));
1138aad13602SShrirang Abhyankar     }
1139aad13602SShrirang Abhyankar 
1140aad13602SShrirang Abhyankar     if (pdipm->nh) {
1141aad13602SShrirang Abhyankar       /* Insert -grad h' */
11429a8cc538SBarry Smith       PetscCall(MatGetRow(pdipm->jac_inequality_trans, i + rjstart, &nc, &aj, NULL));
11439566063dSJacob Faibussowitsch       PetscCall(MatGetOwnershipRanges(tao->jacobian_inequality, &ranges));
1144aad13602SShrirang Abhyankar       proc = 0;
1145aad13602SShrirang Abhyankar       for (j = 0; j < nc; j++) {
1146aad13602SShrirang Abhyankar         /* find row ownership of */
1147aad13602SShrirang Abhyankar         while (aj[j] >= ranges[proc + 1]) proc++;
1148aad13602SShrirang Abhyankar         nx_all = rranges[proc + 1] - rranges[proc];
1149aad13602SShrirang Abhyankar         col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all + nce_all[proc];
11509566063dSJacob Faibussowitsch         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1151aad13602SShrirang Abhyankar       }
11529a8cc538SBarry Smith       PetscCall(MatRestoreRow(pdipm->jac_inequality_trans, i + rjstart, &nc, &aj, NULL));
1153aad13602SShrirang Abhyankar     }
1154aad13602SShrirang Abhyankar 
1155aad13602SShrirang Abhyankar     /* Insert Jci_xb^T' */
11569566063dSJacob Faibussowitsch     PetscCall(MatGetRow(Jci_xb_trans, i + rjstart, &nc, &aj, NULL));
11579566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRanges(pdipm->Jci_xb, &ranges));
1158aad13602SShrirang Abhyankar     proc = 0;
1159aad13602SShrirang Abhyankar     for (j = 0; j < nc; j++) {
1160aad13602SShrirang Abhyankar       /* find row ownership of */
1161aad13602SShrirang Abhyankar       while (aj[j] >= ranges[proc + 1]) proc++;
1162aad13602SShrirang Abhyankar       nx_all = rranges[proc + 1] - rranges[proc];
1163aad13602SShrirang Abhyankar       col    = aj[j] - ranges[proc] + Jranges[proc] + nx_all + nce_all[proc] + nh_all[proc];
11649566063dSJacob Faibussowitsch       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1165aad13602SShrirang Abhyankar     }
11669566063dSJacob Faibussowitsch     PetscCall(MatRestoreRow(Jci_xb_trans, i + rjstart, &nc, &aj, NULL));
1167aad13602SShrirang Abhyankar   }
1168aad13602SShrirang Abhyankar 
116909ee8bb0SRylee Sundermann   /* 2nd Row block of KKT matrix: [grad Ce, deltac*I, 0, 0] */
1170aad13602SShrirang Abhyankar   if (pdipm->Ng) {
11719566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRange(tao->jacobian_equality, &rjstart, NULL));
1172aad13602SShrirang Abhyankar     for (i = 0; i < pdipm->ng; i++) {
1173aad13602SShrirang Abhyankar       row = rstart + pdipm->off_lambdae + i;
1174aad13602SShrirang Abhyankar 
11759566063dSJacob Faibussowitsch       PetscCall(MatGetRow(tao->jacobian_equality, i + rjstart, &nc, &aj, NULL));
1176aad13602SShrirang Abhyankar       proc = 0;
1177aad13602SShrirang Abhyankar       for (j = 0; j < nc; j++) {
1178aad13602SShrirang Abhyankar         while (aj[j] >= cranges[proc + 1]) proc++;
1179aad13602SShrirang Abhyankar         col = aj[j] - cranges[proc] + Jranges[proc];
11809566063dSJacob Faibussowitsch         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz)); /* grad g */
1181aad13602SShrirang Abhyankar       }
11829566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(tao->jacobian_equality, i + rjstart, &nc, &aj, NULL));
1183aad13602SShrirang Abhyankar     }
1184aad13602SShrirang Abhyankar   }
1185aad13602SShrirang Abhyankar   /* Jce_xfixed */
1186aad13602SShrirang Abhyankar   if (pdipm->Nxfixed) {
11879566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRange(pdipm->Jce_xfixed, &Jcrstart, NULL));
1188aad13602SShrirang Abhyankar     for (i = 0; i < (pdipm->nce - pdipm->ng); i++) {
1189aad13602SShrirang Abhyankar       row = rstart + pdipm->off_lambdae + pdipm->ng + i;
1190aad13602SShrirang Abhyankar 
11919566063dSJacob Faibussowitsch       PetscCall(MatGetRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, NULL));
11923c859ba3SBarry Smith       PetscCheck(nc == 1, PETSC_COMM_SELF, PETSC_ERR_SUP, "nc != 1");
1193aad13602SShrirang Abhyankar 
1194aad13602SShrirang Abhyankar       proc = 0;
1195aad13602SShrirang Abhyankar       j    = 0;
1196aad13602SShrirang Abhyankar       while (cols[j] >= cranges[proc + 1]) proc++;
1197aad13602SShrirang Abhyankar       col = cols[j] - cranges[proc] + Jranges[proc];
11989566063dSJacob Faibussowitsch       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
11999566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, NULL));
1200aad13602SShrirang Abhyankar     }
1201aad13602SShrirang Abhyankar   }
1202aad13602SShrirang Abhyankar 
120309ee8bb0SRylee Sundermann   /* 3rd Row block of KKT matrix: [ gradCi, 0, deltac*I, -I] */
1204aad13602SShrirang Abhyankar   if (pdipm->Nh) {
12059566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRange(tao->jacobian_inequality, &rjstart, NULL));
1206aad13602SShrirang Abhyankar     for (i = 0; i < pdipm->nh; i++) {
1207aad13602SShrirang Abhyankar       row = rstart + pdipm->off_lambdai + i;
1208aad13602SShrirang Abhyankar 
12099566063dSJacob Faibussowitsch       PetscCall(MatGetRow(tao->jacobian_inequality, i + rjstart, &nc, &aj, NULL));
1210aad13602SShrirang Abhyankar       proc = 0;
1211aad13602SShrirang Abhyankar       for (j = 0; j < nc; j++) {
1212aad13602SShrirang Abhyankar         while (aj[j] >= cranges[proc + 1]) proc++;
1213aad13602SShrirang Abhyankar         col = aj[j] - cranges[proc] + Jranges[proc];
12149566063dSJacob Faibussowitsch         PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz)); /* grad h */
1215aad13602SShrirang Abhyankar       }
12169566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(tao->jacobian_inequality, i + rjstart, &nc, &aj, NULL));
1217aad13602SShrirang Abhyankar     }
121812d688e0SRylee Sundermann     /* I */
1219aad13602SShrirang Abhyankar     for (i = 0; i < pdipm->nh; i++) {
1220aad13602SShrirang Abhyankar       row = rstart + pdipm->off_lambdai + i;
1221aad13602SShrirang Abhyankar       col = rstart + pdipm->off_z + i;
12229566063dSJacob Faibussowitsch       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1223aad13602SShrirang Abhyankar     }
1224aad13602SShrirang Abhyankar   }
1225aad13602SShrirang Abhyankar 
1226aad13602SShrirang Abhyankar   /* Jci_xb */
12279566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(pdipm->Jci_xb, &Jcrstart, NULL));
1228aad13602SShrirang Abhyankar   for (i = 0; i < (pdipm->nci - pdipm->nh); i++) {
1229aad13602SShrirang Abhyankar     row = rstart + pdipm->off_lambdai + pdipm->nh + i;
1230aad13602SShrirang Abhyankar 
12319566063dSJacob Faibussowitsch     PetscCall(MatGetRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, NULL));
12323c859ba3SBarry Smith     PetscCheck(nc == 1, PETSC_COMM_SELF, PETSC_ERR_SUP, "nc != 1");
1233aad13602SShrirang Abhyankar     proc = 0;
1234aad13602SShrirang Abhyankar     for (j = 0; j < nc; j++) {
1235aad13602SShrirang Abhyankar       while (cols[j] >= cranges[proc + 1]) proc++;
1236aad13602SShrirang Abhyankar       col = cols[j] - cranges[proc] + Jranges[proc];
12379566063dSJacob Faibussowitsch       PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1238aad13602SShrirang Abhyankar     }
12399566063dSJacob Faibussowitsch     PetscCall(MatRestoreRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, NULL));
124012d688e0SRylee Sundermann     /* I */
1241aad13602SShrirang Abhyankar     col = rstart + pdipm->off_z + pdipm->nh + i;
12429566063dSJacob Faibussowitsch     PetscCall(MatPreallocateSet(row, 1, &col, dnz, onz));
1243aad13602SShrirang Abhyankar   }
1244aad13602SShrirang Abhyankar 
1245aad13602SShrirang Abhyankar   /* 4-th Row block of KKT matrix: Z and Ci */
1246aad13602SShrirang Abhyankar   for (i = 0; i < pdipm->nci; i++) {
1247aad13602SShrirang Abhyankar     row      = rstart + pdipm->off_z + i;
1248aad13602SShrirang Abhyankar     cols1[0] = rstart + pdipm->off_lambdai + i;
1249aad13602SShrirang Abhyankar     cols1[1] = row;
12509566063dSJacob Faibussowitsch     PetscCall(MatPreallocateSet(row, 2, cols1, dnz, onz));
1251aad13602SShrirang Abhyankar   }
1252aad13602SShrirang Abhyankar 
1253aad13602SShrirang Abhyankar   /* diagonal entry */
1254aad13602SShrirang Abhyankar   for (i = 0; i < pdipm->n; i++) dnz[i]++; /* diagonal entry */
1255aad13602SShrirang Abhyankar 
1256aad13602SShrirang Abhyankar   /* Create KKT matrix */
12579566063dSJacob Faibussowitsch   PetscCall(MatCreate(comm, &J));
12589566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(J, pdipm->n, pdipm->n, PETSC_DECIDE, PETSC_DECIDE));
12599566063dSJacob Faibussowitsch   PetscCall(MatSetFromOptions(J));
12609566063dSJacob Faibussowitsch   PetscCall(MatSeqAIJSetPreallocation(J, 0, dnz));
12619566063dSJacob Faibussowitsch   PetscCall(MatMPIAIJSetPreallocation(J, 0, dnz, 0, onz));
1262d0609cedSBarry Smith   MatPreallocateEnd(dnz, onz);
1263aad13602SShrirang Abhyankar   pdipm->K = J;
1264aad13602SShrirang Abhyankar 
1265f560b561SHong Zhang   /* (8) Insert constant entries to  K */
1266aad13602SShrirang Abhyankar   /* Set 0.0 to diagonal of K, so that the solver does not complain *about missing diagonal value */
12679566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(J, &rstart, &rend));
126848a46eb9SPierre Jolivet   for (i = rstart; i < rend; i++) PetscCall(MatSetValue(J, i, i, 0.0, INSERT_VALUES));
126909ee8bb0SRylee Sundermann   /* In case Wxx has no diagonal entries preset set diagonal to deltaw given */
127009ee8bb0SRylee Sundermann   if (pdipm->kkt_pd) {
127109ee8bb0SRylee Sundermann     for (i = 0; i < pdipm->nh; i++) {
127209ee8bb0SRylee Sundermann       row = rstart + i;
12739566063dSJacob Faibussowitsch       PetscCall(MatSetValue(J, row, row, pdipm->deltaw, INSERT_VALUES));
127409ee8bb0SRylee Sundermann     }
127509ee8bb0SRylee Sundermann   }
1276aad13602SShrirang Abhyankar 
1277aad13602SShrirang Abhyankar   /* Row block of K: [ grad Ce, 0, 0, 0] */
1278aad13602SShrirang Abhyankar   if (pdipm->Nxfixed) {
12799566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRange(pdipm->Jce_xfixed, &Jcrstart, NULL));
1280aad13602SShrirang Abhyankar     for (i = 0; i < (pdipm->nce - pdipm->ng); i++) {
1281aad13602SShrirang Abhyankar       row = rstart + pdipm->off_lambdae + pdipm->ng + i;
1282aad13602SShrirang Abhyankar 
12839566063dSJacob Faibussowitsch       PetscCall(MatGetRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, &aa));
1284aad13602SShrirang Abhyankar       proc = 0;
1285aad13602SShrirang Abhyankar       for (j = 0; j < nc; j++) {
1286aad13602SShrirang Abhyankar         while (cols[j] >= cranges[proc + 1]) proc++;
1287aad13602SShrirang Abhyankar         col = cols[j] - cranges[proc] + Jranges[proc];
12889566063dSJacob Faibussowitsch         PetscCall(MatSetValue(J, row, col, aa[j], INSERT_VALUES)); /* grad Ce */
12899566063dSJacob Faibussowitsch         PetscCall(MatSetValue(J, col, row, aa[j], INSERT_VALUES)); /* grad Ce' */
1290aad13602SShrirang Abhyankar       }
12919566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(pdipm->Jce_xfixed, i + Jcrstart, &nc, &cols, &aa));
1292aad13602SShrirang Abhyankar     }
1293aad13602SShrirang Abhyankar   }
1294aad13602SShrirang Abhyankar 
129512d688e0SRylee Sundermann   /* Row block of K: [ -grad Ci, 0, 0, I] */
12969566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(pdipm->Jci_xb, &Jcrstart, NULL));
12972da392ccSBarry Smith   for (i = 0; i < pdipm->nci - pdipm->nh; i++) {
1298aad13602SShrirang Abhyankar     row = rstart + pdipm->off_lambdai + pdipm->nh + i;
1299aad13602SShrirang Abhyankar 
13009566063dSJacob Faibussowitsch     PetscCall(MatGetRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, &aa));
1301aad13602SShrirang Abhyankar     proc = 0;
1302aad13602SShrirang Abhyankar     for (j = 0; j < nc; j++) {
1303aad13602SShrirang Abhyankar       while (cols[j] >= cranges[proc + 1]) proc++;
1304aad13602SShrirang Abhyankar       col = cols[j] - cranges[proc] + Jranges[proc];
13059566063dSJacob Faibussowitsch       PetscCall(MatSetValue(J, col, row, -aa[j], INSERT_VALUES));
13069566063dSJacob Faibussowitsch       PetscCall(MatSetValue(J, row, col, -aa[j], INSERT_VALUES));
1307aad13602SShrirang Abhyankar     }
13089566063dSJacob Faibussowitsch     PetscCall(MatRestoreRow(pdipm->Jci_xb, i + Jcrstart, &nc, &cols, &aa));
1309aad13602SShrirang Abhyankar 
1310aad13602SShrirang Abhyankar     col = rstart + pdipm->off_z + pdipm->nh + i;
13119566063dSJacob Faibussowitsch     PetscCall(MatSetValue(J, row, col, 1, INSERT_VALUES));
1312aad13602SShrirang Abhyankar   }
1313aad13602SShrirang Abhyankar 
1314aad13602SShrirang Abhyankar   for (i = 0; i < pdipm->nh; i++) {
1315aad13602SShrirang Abhyankar     row = rstart + pdipm->off_lambdai + i;
1316aad13602SShrirang Abhyankar     col = rstart + pdipm->off_z + i;
13179566063dSJacob Faibussowitsch     PetscCall(MatSetValue(J, row, col, 1, INSERT_VALUES));
131812d688e0SRylee Sundermann   }
131912d688e0SRylee Sundermann 
132012d688e0SRylee Sundermann   /* Row block of K: [ 0, 0, I, ...] */
132112d688e0SRylee Sundermann   for (i = 0; i < pdipm->nci; i++) {
132212d688e0SRylee Sundermann     row = rstart + pdipm->off_z + i;
132312d688e0SRylee Sundermann     col = rstart + pdipm->off_lambdai + i;
13249566063dSJacob Faibussowitsch     PetscCall(MatSetValue(J, row, col, 1, INSERT_VALUES));
1325aad13602SShrirang Abhyankar   }
1326aad13602SShrirang Abhyankar 
132748a46eb9SPierre Jolivet   if (pdipm->Nxfixed) PetscCall(MatDestroy(&Jce_xfixed_trans));
13289566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&Jci_xb_trans));
13299566063dSJacob Faibussowitsch   PetscCall(PetscFree3(ng_all, nh_all, Jranges));
13307f6ac294SRylee Sundermann 
1331f560b561SHong Zhang   /* (9) Set up nonlinear solver SNES */
13329566063dSJacob Faibussowitsch   PetscCall(SNESSetFunction(pdipm->snes, NULL, TaoSNESFunction_PDIPM, (void *)tao));
13339566063dSJacob Faibussowitsch   PetscCall(SNESSetJacobian(pdipm->snes, J, J, TaoSNESJacobian_PDIPM, (void *)tao));
1334f560b561SHong Zhang 
1335f560b561SHong Zhang   if (pdipm->solve_reduced_kkt) {
1336f560b561SHong Zhang     PC pc;
13379566063dSJacob Faibussowitsch     PetscCall(KSPGetPC(tao->ksp, &pc));
13389566063dSJacob Faibussowitsch     PetscCall(PCSetType(pc, PCFIELDSPLIT));
13399566063dSJacob Faibussowitsch     PetscCall(PCFieldSplitSetType(pc, PC_COMPOSITE_SCHUR));
13409566063dSJacob Faibussowitsch     PetscCall(PCFieldSplitSetIS(pc, "2", pdipm->is2));
13419566063dSJacob Faibussowitsch     PetscCall(PCFieldSplitSetIS(pc, "1", pdipm->is1));
1342f560b561SHong Zhang   }
13439566063dSJacob Faibussowitsch   PetscCall(SNESSetFromOptions(pdipm->snes));
1344f560b561SHong Zhang 
13457f6ac294SRylee Sundermann   /* (10) Setup PCPreSolve() for pdipm->solve_symmetric_kkt */
13467f6ac294SRylee Sundermann   if (pdipm->solve_symmetric_kkt) {
13477f6ac294SRylee Sundermann     KSP       ksp;
13487f6ac294SRylee Sundermann     PC        pc;
1349f560b561SHong Zhang     PetscBool isCHOL;
13509566063dSJacob Faibussowitsch     PetscCall(SNESGetKSP(pdipm->snes, &ksp));
13519566063dSJacob Faibussowitsch     PetscCall(KSPGetPC(ksp, &pc));
13529566063dSJacob Faibussowitsch     PetscCall(PCSetPreSolve(pc, PCPreSolve_PDIPM));
1353f560b561SHong Zhang 
13549566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)pc, PCCHOLESKY, &isCHOL));
1355f560b561SHong Zhang     if (isCHOL) {
1356f560b561SHong Zhang       Mat       Factor;
1357f560b561SHong Zhang       PetscBool isMUMPS;
13589566063dSJacob Faibussowitsch       PetscCall(PCFactorGetMatrix(pc, &Factor));
13599566063dSJacob Faibussowitsch       PetscCall(PetscObjectTypeCompare((PetscObject)Factor, "mumps", &isMUMPS));
1360f560b561SHong Zhang       if (isMUMPS) { /* must set mumps ICNTL(13)=1 and ICNTL(24)=1 to call MatGetInertia() */
1361f560b561SHong Zhang #if defined(PETSC_HAVE_MUMPS)
13629566063dSJacob Faibussowitsch         PetscCall(MatMumpsSetIcntl(Factor, 24, 1)); /* detection of null pivot rows */
13639371c9d4SSatish Balay         if (size > 1) { PetscCall(MatMumpsSetIcntl(Factor, 13, 1)); /* parallelism of the root node (enable ScaLAPACK) and its splitting */ }
1364f560b561SHong Zhang #else
1365f560b561SHong Zhang         SETERRQ(PetscObjectComm((PetscObject)tao), PETSC_ERR_SUP, "Requires external package MUMPS");
1366f560b561SHong Zhang #endif
1367f560b561SHong Zhang       }
1368f560b561SHong Zhang     }
13697f6ac294SRylee Sundermann   }
13703ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1371aad13602SShrirang Abhyankar }
1372aad13602SShrirang Abhyankar 
1373aad13602SShrirang Abhyankar /*
1374aad13602SShrirang Abhyankar    TaoDestroy_PDIPM - Destroys the pdipm object
1375aad13602SShrirang Abhyankar 
1376aad13602SShrirang Abhyankar    Input:
1377aad13602SShrirang Abhyankar    full pdipm
1378aad13602SShrirang Abhyankar 
1379aad13602SShrirang Abhyankar    Output:
1380aad13602SShrirang Abhyankar    Destroyed pdipm
1381aad13602SShrirang Abhyankar */
1382d71ae5a4SJacob Faibussowitsch PetscErrorCode TaoDestroy_PDIPM(Tao tao)
1383d71ae5a4SJacob Faibussowitsch {
1384aad13602SShrirang Abhyankar   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;
1385aad13602SShrirang Abhyankar 
1386aad13602SShrirang Abhyankar   PetscFunctionBegin;
1387aad13602SShrirang Abhyankar   /* Freeing Vectors assocaiated with KKT (X) */
13889566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&pdipm->x));       /* Solution x */
13899566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&pdipm->lambdae)); /* Equality constraints lagrangian multiplier*/
13909566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&pdipm->lambdai)); /* Inequality constraints lagrangian multiplier*/
13919566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&pdipm->z));       /* Slack variables */
13929566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&pdipm->X));       /* Big KKT system vector [x; lambdae; lambdai; z] */
1393aad13602SShrirang Abhyankar 
1394aad13602SShrirang Abhyankar   /* work vectors */
13959566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&pdipm->lambdae_xfixed));
13969566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&pdipm->lambdai_xb));
1397aad13602SShrirang Abhyankar 
1398aad13602SShrirang Abhyankar   /* Legrangian equality and inequality Vec */
13999566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&pdipm->ce)); /* Vec of equality constraints */
14009566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&pdipm->ci)); /* Vec of inequality constraints */
1401aad13602SShrirang Abhyankar 
1402aad13602SShrirang Abhyankar   /* Matrices */
14039566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&pdipm->Jce_xfixed));
14049566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&pdipm->Jci_xb)); /* Jacobian of inequality constraints Jci = [tao->jacobian_inequality ; J(nxub); J(nxlb); J(nxbx)] */
14059566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&pdipm->K));
1406aad13602SShrirang Abhyankar 
1407aad13602SShrirang Abhyankar   /* Index Sets */
14089371c9d4SSatish Balay   if (pdipm->Nxub) { PetscCall(ISDestroy(&pdipm->isxub)); /* Finite upper bound only -inf < x < ub */ }
1409aad13602SShrirang Abhyankar 
14109371c9d4SSatish Balay   if (pdipm->Nxlb) { PetscCall(ISDestroy(&pdipm->isxlb)); /* Finite lower bound only  lb <= x < inf */ }
1411aad13602SShrirang Abhyankar 
14129371c9d4SSatish Balay   if (pdipm->Nxfixed) { PetscCall(ISDestroy(&pdipm->isxfixed)); /* Fixed variables         lb =  x = ub */ }
1413aad13602SShrirang Abhyankar 
14149371c9d4SSatish Balay   if (pdipm->Nxbox) { PetscCall(ISDestroy(&pdipm->isxbox)); /* Boxed variables         lb <= x <= ub */ }
1415aad13602SShrirang Abhyankar 
14169371c9d4SSatish Balay   if (pdipm->Nxfree) { PetscCall(ISDestroy(&pdipm->isxfree)); /* Free variables        -inf <= x <= inf */ }
1417aad13602SShrirang Abhyankar 
1418aad13602SShrirang Abhyankar   if (pdipm->solve_reduced_kkt) {
14199566063dSJacob Faibussowitsch     PetscCall(ISDestroy(&pdipm->is1));
14209566063dSJacob Faibussowitsch     PetscCall(ISDestroy(&pdipm->is2));
1421aad13602SShrirang Abhyankar   }
1422aad13602SShrirang Abhyankar 
1423aad13602SShrirang Abhyankar   /* SNES */
14249566063dSJacob Faibussowitsch   PetscCall(SNESDestroy(&pdipm->snes)); /* Nonlinear solver */
14259566063dSJacob Faibussowitsch   PetscCall(PetscFree(pdipm->nce_all));
14269566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&pdipm->jac_equality_trans));
14279566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&pdipm->jac_inequality_trans));
1428aad13602SShrirang Abhyankar 
1429aad13602SShrirang Abhyankar   /* Destroy pdipm */
14309566063dSJacob Faibussowitsch   PetscCall(PetscFree(tao->data)); /* Holding locations of pdipm */
1431aad13602SShrirang Abhyankar 
1432aad13602SShrirang Abhyankar   /* Destroy Dual */
14339566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&tao->DE)); /* equality dual */
14349566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&tao->DI)); /* dinequality dual */
14353ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1436aad13602SShrirang Abhyankar }
1437aad13602SShrirang Abhyankar 
1438d71ae5a4SJacob Faibussowitsch PetscErrorCode TaoSetFromOptions_PDIPM(Tao tao, PetscOptionItems *PetscOptionsObject)
1439d71ae5a4SJacob Faibussowitsch {
1440aad13602SShrirang Abhyankar   TAO_PDIPM *pdipm = (TAO_PDIPM *)tao->data;
1441aad13602SShrirang Abhyankar 
1442aad13602SShrirang Abhyankar   PetscFunctionBegin;
1443d0609cedSBarry Smith   PetscOptionsHeadBegin(PetscOptionsObject, "PDIPM method for constrained optimization");
14449566063dSJacob Faibussowitsch   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));
14459566063dSJacob Faibussowitsch   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));
14469566063dSJacob Faibussowitsch   PetscCall(PetscOptionsBool("-tao_pdipm_solve_reduced_kkt", "Solve reduced KKT system using Schur-complement", NULL, pdipm->solve_reduced_kkt, &pdipm->solve_reduced_kkt, NULL));
14479566063dSJacob Faibussowitsch   PetscCall(PetscOptionsReal("-tao_pdipm_mu_update_factor", "Update scalar for barrier parameter (mu) update", NULL, pdipm->mu_update_factor, &pdipm->mu_update_factor, NULL));
14489566063dSJacob Faibussowitsch   PetscCall(PetscOptionsBool("-tao_pdipm_symmetric_kkt", "Solve non reduced symmetric KKT system", NULL, pdipm->solve_symmetric_kkt, &pdipm->solve_symmetric_kkt, NULL));
14499566063dSJacob Faibussowitsch   PetscCall(PetscOptionsBool("-tao_pdipm_kkt_shift_pd", "Add shifts to make KKT matrix positive definite", NULL, pdipm->kkt_pd, &pdipm->kkt_pd, NULL));
1450d0609cedSBarry Smith   PetscOptionsHeadEnd();
14513ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1452aad13602SShrirang Abhyankar }
1453aad13602SShrirang Abhyankar 
1454aad13602SShrirang Abhyankar /*MC
1455aad13602SShrirang Abhyankar   TAOPDIPM - Barrier-based primal-dual interior point algorithm for generally constrained optimization.
1456aad13602SShrirang Abhyankar 
1457aad13602SShrirang Abhyankar   Option Database Keys:
1458aad13602SShrirang Abhyankar +   -tao_pdipm_push_init_lambdai - parameter to push initial dual variables away from bounds (> 0)
1459aad13602SShrirang Abhyankar .   -tao_pdipm_push_init_slack - parameter to push initial slack variables away from bounds (> 0)
146012d688e0SRylee Sundermann .   -tao_pdipm_mu_update_factor - update scalar for barrier parameter (mu) update (> 0)
146109ee8bb0SRylee Sundermann .   -tao_pdipm_symmetric_kkt - Solve non-reduced symmetric KKT system
146209ee8bb0SRylee Sundermann -   -tao_pdipm_kkt_shift_pd - Add shifts to make KKT matrix positive definite
1463aad13602SShrirang Abhyankar 
1464aad13602SShrirang Abhyankar   Level: beginner
1465*20f4b53cSBarry Smith 
1466*20f4b53cSBarry Smith .seealso: `Tao`, `TaoType`
1467aad13602SShrirang Abhyankar M*/
1468*20f4b53cSBarry Smith 
1469d71ae5a4SJacob Faibussowitsch PETSC_EXTERN PetscErrorCode TaoCreate_PDIPM(Tao tao)
1470d71ae5a4SJacob Faibussowitsch {
1471aad13602SShrirang Abhyankar   TAO_PDIPM *pdipm;
1472aad13602SShrirang Abhyankar 
1473aad13602SShrirang Abhyankar   PetscFunctionBegin;
1474aad13602SShrirang Abhyankar   tao->ops->setup          = TaoSetup_PDIPM;
1475aad13602SShrirang Abhyankar   tao->ops->solve          = TaoSolve_PDIPM;
1476aad13602SShrirang Abhyankar   tao->ops->setfromoptions = TaoSetFromOptions_PDIPM;
147770c9796eSresundermann   tao->ops->view           = TaoView_PDIPM;
1478aad13602SShrirang Abhyankar   tao->ops->destroy        = TaoDestroy_PDIPM;
1479aad13602SShrirang Abhyankar 
14804dfa11a4SJacob Faibussowitsch   PetscCall(PetscNew(&pdipm));
1481aad13602SShrirang Abhyankar   tao->data = (void *)pdipm;
1482aad13602SShrirang Abhyankar 
1483aad13602SShrirang Abhyankar   pdipm->nx = pdipm->Nx = 0;
1484aad13602SShrirang Abhyankar   pdipm->nxfixed = pdipm->Nxfixed = 0;
1485aad13602SShrirang Abhyankar   pdipm->nxlb = pdipm->Nxlb = 0;
1486aad13602SShrirang Abhyankar   pdipm->nxub = pdipm->Nxub = 0;
1487aad13602SShrirang Abhyankar   pdipm->nxbox = pdipm->Nxbox = 0;
1488aad13602SShrirang Abhyankar   pdipm->nxfree = pdipm->Nxfree = 0;
1489aad13602SShrirang Abhyankar 
1490aad13602SShrirang Abhyankar   pdipm->ng = pdipm->Ng = pdipm->nce = pdipm->Nce = 0;
1491aad13602SShrirang Abhyankar   pdipm->nh = pdipm->Nh = pdipm->nci = pdipm->Nci = 0;
1492aad13602SShrirang Abhyankar   pdipm->n = pdipm->N     = 0;
1493aad13602SShrirang Abhyankar   pdipm->mu               = 1.0;
1494aad13602SShrirang Abhyankar   pdipm->mu_update_factor = 0.1;
1495aad13602SShrirang Abhyankar 
149609ee8bb0SRylee Sundermann   pdipm->deltaw     = 0.0;
149709ee8bb0SRylee Sundermann   pdipm->lastdeltaw = 3 * 1.e-4;
149809ee8bb0SRylee Sundermann   pdipm->deltac     = 0.0;
149909ee8bb0SRylee Sundermann   pdipm->kkt_pd     = PETSC_FALSE;
150009ee8bb0SRylee Sundermann 
1501aad13602SShrirang Abhyankar   pdipm->push_init_slack     = 1.0;
1502aad13602SShrirang Abhyankar   pdipm->push_init_lambdai   = 1.0;
1503aad13602SShrirang Abhyankar   pdipm->solve_reduced_kkt   = PETSC_FALSE;
150412d688e0SRylee Sundermann   pdipm->solve_symmetric_kkt = PETSC_TRUE;
1505aad13602SShrirang Abhyankar 
1506aad13602SShrirang Abhyankar   /* Override default settings (unless already changed) */
1507aad13602SShrirang Abhyankar   if (!tao->max_it_changed) tao->max_it = 200;
1508aad13602SShrirang Abhyankar   if (!tao->max_funcs_changed) tao->max_funcs = 500;
1509aad13602SShrirang Abhyankar 
15109566063dSJacob Faibussowitsch   PetscCall(SNESCreate(((PetscObject)tao)->comm, &pdipm->snes));
15119566063dSJacob Faibussowitsch   PetscCall(SNESSetOptionsPrefix(pdipm->snes, tao->hdr.prefix));
15129566063dSJacob Faibussowitsch   PetscCall(SNESGetKSP(pdipm->snes, &tao->ksp));
15139566063dSJacob Faibussowitsch   PetscCall(PetscObjectReference((PetscObject)tao->ksp));
15149566063dSJacob Faibussowitsch   PetscCall(KSPSetApplicationContext(tao->ksp, (void *)tao));
15153ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1516aad13602SShrirang Abhyankar }
1517