xref: /petsc/src/tao/constrained/impls/almm/almmutils.c (revision 20f4b53cbb5e9bd9ef12b76a8697d60d197cda17)

#include <../src/tao/constrained/impls/almm/almm.h> /*I "petsctao.h" I*/
#include <petsctao.h>
#include <petsc/private/petscimpl.h>
#include <petsc/private/vecimpl.h>

/*@
   TaoALMMGetType - Retrieve the augmented Lagrangian formulation type for the subproblem.

   Input Parameters:
.  tao - the `Tao` context for the `TAOALMM` solver

   Output Parameters:
.  type - augmented Lagragrangian type

   Level: advanced

.seealso: `Tao`, `TAOALMM`, `TaoALMMSetType()`, `TaoALMMType`
@*/
PetscErrorCode TaoALMMGetType(Tao tao, TaoALMMType *type)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(tao, TAO_CLASSID, 1);
  PetscValidPointer(type, 2);
  PetscUseMethod(tao, "TaoALMMGetType_C", (Tao, TaoALMMType *), (tao, type));
  PetscFunctionReturn(PETSC_SUCCESS);
}

PetscErrorCode TaoALMMGetType_Private(Tao tao, TaoALMMType *type)
{
  TAO_ALMM *auglag = (TAO_ALMM *)tao->data;

  PetscFunctionBegin;
  *type = auglag->type;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   TaoALMMSetType - Determine the augmented Lagrangian formulation type for the subproblem.

   Input Parameters:
+  tao - the `Tao` context for the `TAOALMM` solver
-  type - augmented Lagragrangian type

   Level: advanced

.seealso: `Tao`, `TAOALMM`, `TaoALMMGetType()`, `TaoALMMType`
@*/
PetscErrorCode TaoALMMSetType(Tao tao, TaoALMMType type)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(tao, TAO_CLASSID, 1);
  PetscTryMethod(tao, "TaoALMMSetType_C", (Tao, TaoALMMType), (tao, type));
  PetscFunctionReturn(PETSC_SUCCESS);
}

PetscErrorCode TaoALMMSetType_Private(Tao tao, TaoALMMType type)
{
  TAO_ALMM *auglag = (TAO_ALMM *)tao->data;

  PetscFunctionBegin;
  PetscCheck(!tao->setupcalled, PetscObjectComm((PetscObject)tao), PETSC_ERR_ORDER, "TaoALMMSetType() must be called before TaoSetUp()");
  auglag->type = type;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   TaoALMMGetSubsolver - Retrieve the subsolver being used by `TAOALMM`.

   Input Parameters:
.  tao - the `Tao` context for the `TAOALMM` solver

   Output Parameter:
.  subsolver - the `Tao` context for the subsolver

   Level: advanced

.seealso: `Tao`, `TAOALMM`, `TaoALMMSetSubsolver()`
@*/
PetscErrorCode TaoALMMGetSubsolver(Tao tao, Tao *subsolver)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(tao, TAO_CLASSID, 1);
  PetscValidPointer(subsolver, 2);
  PetscUseMethod(tao, "TaoALMMGetSubsolver_C", (Tao, Tao *), (tao, subsolver));
  PetscFunctionReturn(PETSC_SUCCESS);
}

PetscErrorCode TaoALMMGetSubsolver_Private(Tao tao, Tao *subsolver)
{
  TAO_ALMM *auglag = (TAO_ALMM *)tao->data;

  PetscFunctionBegin;
  *subsolver = auglag->subsolver;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   TaoALMMSetSubsolver - Changes the subsolver inside `TAOALMM` with the user provided one.

   Input Parameters:
+  tao - the `Tao` context for the `TAOALMM` solver
-  subsolver - the Tao context for the subsolver

   Level: advanced

   Note:
   This is not recommended, instead call `TaoALMMGetSubsolver()` and set the type as desired.

.seealso: `Tao`, `TAOALMM`, `TaoALMMGetSubsolver()`
@*/
PetscErrorCode TaoALMMSetSubsolver(Tao tao, Tao subsolver)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(tao, TAO_CLASSID, 1);
  PetscValidHeaderSpecific(subsolver, TAO_CLASSID, 2);
  PetscTryMethod(tao, "TaoALMMSetSubsolver_C", (Tao, Tao), (tao, subsolver));
  PetscFunctionReturn(PETSC_SUCCESS);
}

PetscErrorCode TaoALMMSetSubsolver_Private(Tao tao, Tao subsolver)
{
  TAO_ALMM *auglag = (TAO_ALMM *)tao->data;
  PetscBool compatible;

  PetscFunctionBegin;
  if (subsolver == auglag->subsolver) PetscFunctionReturn(PETSC_SUCCESS);
  if (tao->bounded) {
    PetscCall(PetscObjectTypeCompareAny((PetscObject)subsolver, &compatible, TAOSHELL, TAOBNCG, TAOBQNLS, TAOBQNKLS, TAOBQNKTR, TAOBQNKTL, ""));
    PetscCheck(compatible, PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_INCOMP, "Subsolver must be a bound-constrained first-order method");
  } else {
    PetscCall(PetscObjectTypeCompareAny((PetscObject)subsolver, &compatible, TAOSHELL, TAOCG, TAOLMVM, TAOBNCG, TAOBQNLS, TAOBQNKLS, TAOBQNKTR, TAOBQNKTL, ""));
    PetscCheck(compatible, PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_INCOMP, "Subsolver must be a first-order method");
  }
  PetscCheck(compatible, PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_INCOMP, "Subsolver must be a first-order method");
  PetscCall(PetscObjectReference((PetscObject)subsolver));
  PetscCall(TaoDestroy(&auglag->subsolver));
  auglag->subsolver = subsolver;
  if (tao->setupcalled) {
    PetscCall(TaoSetSolution(auglag->subsolver, auglag->P));
    PetscCall(TaoSetObjective(auglag->subsolver, TaoALMMSubsolverObjective_Private, (void *)auglag));
    PetscCall(TaoSetObjectiveAndGradient(auglag->subsolver, NULL, TaoALMMSubsolverObjectiveAndGradient_Private, (void *)auglag));
    PetscCall(TaoSetVariableBounds(auglag->subsolver, auglag->PL, auglag->PU));
  }
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   TaoALMMGetMultipliers - Retrieve a pointer to the Lagrange multipliers.

   Input Parameters:
.  tao - the `Tao` context for the `TAOALMM` solver

   Output Parameters:
.  Y - vector of Lagrange multipliers

   Level: advanced

   Notes:
   For problems with both equality and inequality constraints,
   the multipliers are combined together as Y = (Ye, Yi). Users
   can recover copies of the subcomponents using index sets
   provided by `TaoALMMGetDualIS()` and use `VecGetSubVector()`.

.seealso: `TAOALMM`, `Tao`, `TaoALMMSetMultipliers()`, `TaoALMMGetDualIS()`
@*/
PetscErrorCode TaoALMMGetMultipliers(Tao tao, Vec *Y)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(tao, TAO_CLASSID, 1);
  PetscValidPointer(Y, 2);
  PetscUseMethod(tao, "TaoALMMGetMultipliers_C", (Tao, Vec *), (tao, Y));
  PetscFunctionReturn(PETSC_SUCCESS);
}

PetscErrorCode TaoALMMGetMultipliers_Private(Tao tao, Vec *Y)
{
  TAO_ALMM *auglag = (TAO_ALMM *)tao->data;

  PetscFunctionBegin;
  PetscCheck(tao->setupcalled, PetscObjectComm((PetscObject)tao), PETSC_ERR_ORDER, "TaoSetUp() must be called first for scatters to be constructed");
  *Y = auglag->Y;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   TaoALMMSetMultipliers - Set user-defined Lagrange multipliers.

   Input Parameters:
+  tao - the `Tao` context for the `TAOALMM` solver
-  Y - vector of Lagrange multipliers

   Level: advanced

   Notes:
   The vector type and parallel layout must match the equality and inequality constraints.

   The vector must have a local size equal to the sum of the local sizes for the constraint vectors, and a
   global size equal to the sum of the global sizes of the constraint vectors.

   This routine is only useful if the user wants to change the
   parallel distribution of the combined dual vector in problems that
   feature both equality and inequality constraints. For other tasks,
   it is strongly recommended that the user retrieve the dual vector
   created by the solver using `TaoALMMGetMultipliers()`.

.seealso: `TAOALMM`, `Tao`, `TaoALMMGetMultipliers()`
@*/
PetscErrorCode TaoALMMSetMultipliers(Tao tao, Vec Y)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(tao, TAO_CLASSID, 1);
  PetscValidHeaderSpecific(Y, VEC_CLASSID, 2);
  PetscTryMethod(tao, "TaoALMMSetMultipliers_C", (Tao, Vec), (tao, Y));
  PetscFunctionReturn(PETSC_SUCCESS);
}

PetscErrorCode TaoALMMSetMultipliers_Private(Tao tao, Vec Y)
{
  TAO_ALMM *auglag = (TAO_ALMM *)tao->data;
  VecType   Ytype;
  PetscInt  Nuser, Neq, Nineq, N;
  PetscBool same = PETSC_FALSE;

  PetscFunctionBegin;
  /* no-op if user provides vector from TaoALMMGetMultipliers() */
  if (Y == auglag->Y) PetscFunctionReturn(PETSC_SUCCESS);
  /* make sure vector type is same as equality and inequality constraints */
  if (tao->eq_constrained) {
    PetscCall(VecGetType(tao->constraints_equality, &Ytype));
  } else {
    PetscCall(VecGetType(tao->constraints_inequality, &Ytype));
  }
  PetscCall(PetscObjectTypeCompare((PetscObject)Y, Ytype, &same));
  PetscCheck(same, PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_INCOMP, "Given vector for multipliers is not the same type as constraint vectors");
  /* make sure global size matches sum of equality and inequality */
  if (tao->eq_constrained) {
    PetscCall(VecGetSize(tao->constraints_equality, &Neq));
  } else {
    Neq = 0;
  }
  if (tao->ineq_constrained) {
    PetscCall(VecGetSize(tao->constraints_inequality, &Nineq));
  } else {
    Nineq = 0;
  }
  N = Neq + Nineq;
  PetscCall(VecGetSize(Y, &Nuser));
  PetscCheck(Nuser == N, PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_INCOMP, "Given vector has wrong global size");
  /* if there is only one type of constraint, then we need the local size to match too */
  if (Neq == 0) {
    PetscCall(VecGetLocalSize(tao->constraints_inequality, &Nineq));
    PetscCall(VecGetLocalSize(Y, &Nuser));
    PetscCheck(Nuser == Nineq, PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_INCOMP, "Given vector has wrong local size");
  }
  if (Nineq == 0) {
    PetscCall(VecGetLocalSize(tao->constraints_equality, &Neq));
    PetscCall(VecGetLocalSize(Y, &Nuser));
    PetscCheck(Nuser == Neq, PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_INCOMP, "Given vector has wrong local size");
  }
  /* if we got here, the given vector is compatible so we can replace the current one */
  PetscCall(PetscObjectReference((PetscObject)Y));
  PetscCall(VecDestroy(&auglag->Y));
  auglag->Y = Y;
  /* if there are both types of constraints and the solver has already been set up,
     then we need to regenerate VecScatter objects for the new combined dual vector */
  if (tao->setupcalled && tao->eq_constrained && tao->ineq_constrained) {
    PetscCall(VecDestroy(&auglag->C));
    PetscCall(VecDuplicate(auglag->Y, &auglag->C));
    PetscCall(VecScatterDestroy(&auglag->Yscatter[0]));
    PetscCall(VecScatterCreate(auglag->Y, auglag->Yis[0], auglag->Ye, NULL, &auglag->Yscatter[0]));
    PetscCall(VecScatterDestroy(&auglag->Yscatter[1]));
    PetscCall(VecScatterCreate(auglag->Y, auglag->Yis[1], auglag->Yi, NULL, &auglag->Yscatter[1]));
  }
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   TaoALMMGetPrimalIS - Retrieve the index set that identifies optimization
                        and slack variable components of the subsolver's solution vector.

   Input Parameter:
.  tao - the `Tao` context for the `TAOALMM` solver

   Output Parameters:
+  opt_is - index set associated with the optimization variables (`NULL` if not needed)
-  slack_is - index set associated with the slack variables (`NULL` if not needed)

   Level: advanced

.seealso: `TAOALMM`, `Tao`, `IS`, `TaoALMMGetPrimalVector()`
@*/
PetscErrorCode TaoALMMGetPrimalIS(Tao tao, IS *opt_is, IS *slack_is)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(tao, TAO_CLASSID, 1);
  PetscUseMethod(tao, "TaoALMMGetPrimalIS_C", (Tao, IS *, IS *), (tao, opt_is, slack_is));
  PetscFunctionReturn(PETSC_SUCCESS);
}

PetscErrorCode TaoALMMGetPrimalIS_Private(Tao tao, IS *opt_is, IS *slack_is)
{
  TAO_ALMM *auglag = (TAO_ALMM *)tao->data;

  PetscFunctionBegin;
  PetscCheck(tao->ineq_constrained, PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_WRONGSTATE, "Primal space has index sets only for inequality constrained problems");
  PetscCheck(tao->setupcalled, PetscObjectComm((PetscObject)tao), PETSC_ERR_ORDER, "TaoSetUp() must be called first for index sets to be constructed");
  if (opt_is) *opt_is = auglag->Pis[0];
  if (slack_is) *slack_is = auglag->Pis[1];
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   TaoALMMGetDualIS - Retrieve the index set that identifies equality
                      and inequality constraint components of the dual vector returned
                      by `TaoALMMGetMultipliers()`.

   Input Parameter:
.  tao - the Tao context for the `TAOALMM` solver

   Output Parameters:
+  eq_is - index set associated with the equality constraints (`NULL` if not needed)
-  ineq_is - index set associated with the inequality constraints (`NULL` if not needed)

   Level: advanced

.seealso: `TAOALMM`, `Tao`, `TaoALMMGetMultipliers()`
@*/
PetscErrorCode TaoALMMGetDualIS(Tao tao, IS *eq_is, IS *ineq_is)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(tao, TAO_CLASSID, 1);
  PetscUseMethod(tao, "TaoALMMGetDualIS_C", (Tao, IS *, IS *), (tao, eq_is, ineq_is));
  PetscFunctionReturn(PETSC_SUCCESS);
}

PetscErrorCode TaoALMMGetDualIS_Private(Tao tao, IS *eq_is, IS *ineq_is)
{
  TAO_ALMM *auglag = (TAO_ALMM *)tao->data;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(tao, TAO_CLASSID, 1);
  PetscCheck(tao->ineq_constrained && tao->ineq_constrained, PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_WRONGSTATE, "Dual space has index sets only when problem has both equality and inequality constraints");
  PetscCheck(tao->setupcalled, PetscObjectComm((PetscObject)tao), PETSC_ERR_ORDER, "TaoSetUp() must be called first for index sets to be constructed");
  if (eq_is) *eq_is = auglag->Yis[0];
  if (ineq_is) *ineq_is = auglag->Yis[1];
  PetscFunctionReturn(PETSC_SUCCESS);
}