TaoDiagonalNNLS.cpp

Go to the documentation of this file.

/** \file TaoDiagonalNNLS.cpp
 * \brief Direct TAO shell for NNLS with orthogonal constraint gradients.
 */
 
#include <petsc/private/taoimpl.h>
#include <TaoDiagonalNNLS.hpp>
 
#include <memory>
 
namespace MoFEM {
namespace {
 
constexpr char diagonalNNLSContextKey[] = "_MoFEM_TaoDiagonalNNLS";
 
struct TaoDiagonalNNLS {
  Tao parentTao = nullptr;
  SmartPetscObj<Mat> jacobian;
  SmartPetscObj<Vec> target;
  SmartPetscObj<Vec> residual;
  SmartPetscObj<Vec> gradient;
  SmartPetscObj<Vec> diagonal;
  SmartPetscObj<Vec> work;
  PetscErrorCode (*destroyShell)(Tao) = nullptr;
 
  MoFEMErrorCode evaluate(Vec x, PetscReal *f, Vec g) {
    MoFEMFunctionBegin;
    CHKERR MatMultTranspose(jacobian, x, residual);
    CHKERR VecAXPY(residual, -1, target);
    if (f) {
      PetscScalar squared;
      CHKERR VecDot(residual, residual, &squared);
      *f = 0.5 * PetscRealPart(squared);
    }
    if (g)
      CHKERR MatMult(jacobian, residual, g);
    MoFEMFunctionReturn(0);
  }
 
  /** Monitor the actual NNLS KKT residual, including active lower bounds. */
  MoFEMErrorCode monitor(Vec x, PetscInt iteration, PetscReal step) {
    MoFEMFunctionBegin;
    PetscReal objective;
    CHKERR TaoComputeObjectiveAndGradient(parentTao, x, &objective, gradient);
    CHKERR VecCopy(gradient, work);
    const PetscScalar *values;
    PetscScalar *projected;
    PetscInt size;
    CHKERR VecGetLocalSize(x, &size);
    CHKERR VecGetArrayRead(x, &values);
    CHKERR VecGetArray(work, &projected);
    for (PetscInt i = 0; i != size; ++i)
      if (PetscRealPart(values[i]) == 0)
        projected[i] = std::min(PetscRealPart(projected[i]), PetscReal(0));
    CHKERR VecRestoreArray(work, &projected);
    CHKERR VecRestoreArrayRead(x, &values);
    PetscReal norm;
    CHKERR VecNorm(work, NORM_2, &norm);
    if (!std::isfinite(objective) || !std::isfinite(norm))
      SETERRQ(PetscObjectComm(getPetscObject(parentTao)), PETSC_ERR_FP,
              "Diagonal NNLS objective or projected gradient is not finite");
    CHKERR TaoSetIterationNumber(parentTao, iteration);
    CHKERR TaoLogConvergenceHistory(parentTao, objective, norm, 0, 0);
    CHKERR TaoMonitor(parentTao, iteration, objective, norm, 0, step);
    CHKERR (*parentTao->ops->convergencetest)(parentTao, parentTao->cnvP);
    MoFEMFunctionReturn(0);
  }
 
  MoFEMErrorCode solve(Tao tao) {
    MoFEMFunctionBegin;
    const auto comm = PetscObjectComm(getPetscObject(tao));
    if (!jacobian || !target)
      SETERRQ(comm, PETSC_ERR_ARG_WRONGSTATE,
              "Call TaoDiagonalNNLSSetProblem before TaoSolve");
    if (tao->eq_constrained || tao->ineq_constrained ||
        tao->ops->computeconstraints)
      SETERRQ(comm, PETSC_ERR_SUP,
              "Diagonal NNLS only supports nonnegativity constraints");
    Vec x;
    CHKERR TaoGetSolution(tao, &x);
    PetscInt rows, columns, local_rows, local_columns;
    PetscInt size, local_size, target_size, target_local_size;
    CHKERR MatGetSize(jacobian, &rows, &columns);
    CHKERR MatGetLocalSize(jacobian, &local_rows, &local_columns);
    CHKERR VecGetSize(x, &size);
    CHKERR VecGetLocalSize(x, &local_size);
    CHKERR VecGetSize(target, &target_size);
    CHKERR VecGetLocalSize(target, &target_local_size);
    PetscInt compatible = rows == size && columns == target_size &&
                          local_rows == local_size &&
                          local_columns == target_local_size;
    CHKERR MPIU_Allreduce(MPI_IN_PLACE, &compatible, 1, MPIU_INT, MPI_MIN,
                          comm);
    if (!compatible)
      SETERRQ(comm, PETSC_ERR_ARG_SIZ,
              "Diagonal NNLS solution/target must use the Jacobian row/column "
              "layouts");
    if (tao->bounded || tao->ops->computebounds) {
      CHKERR TaoComputeVariableBounds(tao);
      compatible = 1;
      for (const auto bound : {tao->XL, tao->XU}) {
        if (!bound)
          continue;
        PetscInt bound_size, bound_local_size;
        CHKERR VecGetSize(bound, &bound_size);
        CHKERR VecGetLocalSize(bound, &bound_local_size);
        if (bound_size != size || bound_local_size != local_size)
          compatible = 0;
      }
      CHKERR MPIU_Allreduce(MPI_IN_PLACE, &compatible, 1, MPIU_INT, MPI_MIN,
                            comm);
      if (!compatible)
        SETERRQ(comm, PETSC_ERR_ARG_SIZ,
                "Diagonal NNLS bounds must use the solution vector layout");
      const PetscScalar *lower = nullptr, *upper = nullptr;
      if (tao->XL)
        CHKERR VecGetArrayRead(tao->XL, &lower);
      if (tao->XU)
        CHKERR VecGetArrayRead(tao->XU, &upper);
      for (PetscInt i = 0; i != local_size; ++i)
        if ((lower && PetscRealPart(lower[i]) != 0) ||
            (upper && !(PetscRealPart(upper[i]) >= PETSC_INFINITY)))
          compatible = 0;
      if (tao->XU)
        CHKERR VecRestoreArrayRead(tao->XU, &upper);
      if (tao->XL)
        CHKERR VecRestoreArrayRead(tao->XL, &lower);
      CHKERR MPIU_Allreduce(MPI_IN_PLACE, &compatible, 1, MPIU_INT, MPI_MIN,
                            comm);
      if (!compatible)
        SETERRQ(comm, PETSC_ERR_SUP,
                "Diagonal NNLS requires lower bounds zero and upper bounds "
                "positive infinity");
    }
    if (rows) {
      Mat matrix;
      CHKERR MatMatTransposeMult(jacobian, jacobian, MAT_INITIAL_MATRIX,
                                 PETSC_DEFAULT, &matrix);
      SmartPetscObj<Mat> gram(matrix);
      CHKERR MatGetDiagonal(gram, diagonal);
      const PetscScalar *entries;
      CHKERR VecGetArrayRead(diagonal, &entries);
      compatible = 1;
      for (PetscInt i = 0; i != local_size; ++i)
        if (!std::isfinite(PetscRealPart(entries[i])) ||
            !(PetscRealPart(entries[i]) > 0))
          compatible = 0;
      CHKERR VecRestoreArrayRead(diagonal, &entries);
      CHKERR MPIU_Allreduce(MPI_IN_PLACE, &compatible, 1, MPIU_INT, MPI_MIN,
                            comm);
      if (!compatible)
        SETERRQ(comm, PETSC_ERR_ARG_OUTOFRANGE,
                "Diagonal NNLS requires nonzero finite Jacobian rows");
      // Normalize before checking coupling: small-volume rows must receive
      // the same relative check as large-volume rows.
      CHKERR VecCopy(diagonal, work);
      CHKERR VecSqrtAbs(work);
      CHKERR VecReciprocal(work);
      CHKERR MatDiagonalScale(gram, work, work);
      CHKERR MatShift(gram, -1);
      PetscReal coupling;
      CHKERR MatNorm(gram, NORM_INFINITY, &coupling);
      if (!std::isfinite(coupling) || coupling > 100 * PETSC_MACHINE_EPSILON)
        SETERRQ(comm, PETSC_ERR_ARG_INCOMP,
                "Diagonal NNLS requires mutually orthogonal Jacobian rows; "
                "normalized Gram coupling is %g. Use blmvm for coupled rows",
                static_cast<double>(coupling));
    }
    CHKERR TaoSetConvergedReason(tao, TAO_CONTINUE_ITERATING);
    CHKERR VecZeroEntries(work);
    CHKERR VecPointwiseMax(x, x, work);
    CHKERR monitor(x, 0, 1);
    if (tao->reason != TAO_CONTINUE_ITERATING)
      MoFEMFunctionReturnHot(0);
    auto previous = vectorDuplicate(x);
    CHKERR VecCopy(x, previous);
    CHKERR MatMult(jacobian, target, x);
    CHKERR VecPointwiseDivide(x, x, diagonal);
    CHKERR VecZeroEntries(work);
    CHKERR VecPointwiseMax(x, x, work);
    CHKERR VecWAXPY(work, -1, previous, x);
    PetscReal step;
    CHKERR VecNorm(work, NORM_2, &step);
    CHKERR monitor(x, 1, step);
    if (tao->reason == TAO_CONTINUE_ITERATING ||
        tao->reason == TAO_CONVERGED_STEPTOL) {
      // Roundoff or tolerated Gram coupling can prevent a requested residual
      // tolerance. Report that failure instead of declaring a direct success.
      CHKERR TaoSetConvergedReason(tao, TAO_DIVERGED_USER);
      CHKERR PetscInfo(tao, "Diagonal NNLS direct update did not satisfy the "
                            "requested optimality tolerance\n");
    }
    MoFEMFunctionReturn(0);
  }
};
 
MoFEMErrorCode getDiagonalNNLSContext(Tao tao, TaoDiagonalNNLS *&ctx) {
  MoFEMFunctionBegin;
  PetscBool matches;
  CHKERR PetscObjectTypeCompare(getPetscObject(tao), TAODIAGONALNNLS, &matches);
  if (!matches)
    SETERRQ(PetscObjectComm(getPetscObject(tao)), PETSC_ERR_ARG_WRONG,
            "TAO type must be diagonal");
  PetscContainer container = nullptr;
  CHKERR PetscObjectQuery(getPetscObject(tao), diagonalNNLSContextKey,
                          reinterpret_cast<PetscObject *>(&container));
  if (!container)
    SETERRQ(PetscObjectComm(getPetscObject(tao)), PETSC_ERR_ARG_WRONGSTATE,
            "Diagonal NNLS context is not attached");
  void *ptr;
  CHKERR PetscContainerGetPointer(container, &ptr);
  ctx = static_cast<TaoDiagonalNNLS *>(ptr);
  MoFEMFunctionReturn(0);
}
 
MoFEMErrorCode createTaoDiagonalNNLS(Tao tao) {
  MoFEMFunctionBegin;
  CHKERR PetscObjectChangeTypeName(getPetscObject(tao), nullptr);
  CHKERR TaoSetType(tao, TAOSHELL);
  const auto comm = PetscObjectComm(getPetscObject(tao));
  auto ctx = std::make_unique<TaoDiagonalNNLS>();
  ctx->parentTao = tao;
  ctx->destroyShell = tao->ops->destroy;
  PetscContainer container;
  CHKERR PetscContainerCreate(comm, &container);
  CHKERR PetscContainerSetPointer(container, ctx.get());
  CHKERR PetscContainerSetUserDestroy(container, [](void *ptr) {
    MoFEMFunctionBegin;
    auto *ctx = static_cast<TaoDiagonalNNLS *>(ptr);
    const auto tao = ctx->parentTao;
    if (tao->user_objP == ctx) {
      tao->ops->computeobjective = nullptr;
      tao->user_objP = nullptr;
    }
    if (tao->user_gradP == ctx) {
      tao->ops->computegradient = nullptr;
      tao->user_gradP = nullptr;
    }
    if (tao->user_objgradP == ctx) {
      tao->ops->computeobjectiveandgradient = nullptr;
      tao->user_objgradP = nullptr;
    }
    delete ctx;
    MoFEMFunctionReturn(0);
  });
  ctx.release();
  const auto error = PetscObjectCompose(
      getPetscObject(tao), diagonalNNLSContextKey, getPetscObject(container));
  CHKERR PetscContainerDestroy(&container);
  CHKERR error;
  tao->ops->destroy = [](Tao tao) {
    MoFEMFunctionBegin;
    TaoDiagonalNNLS *ctx;
    CHKERR getDiagonalNNLSContext(tao, ctx);
    const auto destroy_shell = ctx->destroyShell;
    CHKERR PetscObjectCompose(getPetscObject(tao), diagonalNNLSContextKey,
                              nullptr);
    CHKERR destroy_shell(tao);
    MoFEMFunctionReturn(0);
  };
  tao->ops->view = [](Tao tao, PetscViewer viewer) {
    MoFEMFunctionBegin;
    PetscBool ascii;
    CHKERR PetscObjectTypeCompare(getPetscObject(viewer), PETSCVIEWERASCII,
                                  &ascii);
    if (ascii)
      CHKERR PetscViewerASCIIPrintf(
          viewer, "  Direct nonnegative least squares with orthogonal rows\n"
                  "  One diagonal update; no line search\n");
    MoFEMFunctionReturn(0);
  };
  CHKERR TaoShellSetSolve(tao, [](Tao tao) {
    MoFEMFunctionBegin;
    TaoDiagonalNNLS *ctx;
    CHKERR getDiagonalNNLSContext(tao, ctx);
    CHKERR ctx->solve(tao);
    MoFEMFunctionReturn(0);
  });
  MoFEMFunctionReturn(0);
}
 
} // namespace
 
MoFEMErrorCode TaoRegisterDiagonalNNLS() {
  MoFEMFunctionBegin;
  CHKERR TaoRegister(TAODIAGONALNNLS, createTaoDiagonalNNLS);
  MoFEMFunctionReturn(0);
}
 
MoFEMErrorCode TaoDiagonalNNLSSetProblem(Tao tao, Mat jacobian, Vec target) {
  MoFEMFunctionBegin;
  TaoDiagonalNNLS *ctx;
  CHKERR getDiagonalNNLSContext(tao, ctx);
  ctx->jacobian = SmartPetscObj<Mat>(jacobian, true);
  ctx->target = SmartPetscObj<Vec>(target, true);
  Vec row, column;
  CHKERR MatCreateVecs(jacobian, &column, &row);
  ctx->residual = SmartPetscObj<Vec>(column);
  ctx->gradient = SmartPetscObj<Vec>(row);
  ctx->diagonal = vectorDuplicate(ctx->gradient);
  ctx->work = vectorDuplicate(ctx->gradient);
  CHKERR ::TaoSetObjective(
      tao,
      [](Tao, Vec x, PetscReal *f, void *ptr) {
        return static_cast<TaoDiagonalNNLS *>(ptr)->evaluate(x, f, nullptr);
      },
      ctx);
  CHKERR ::TaoSetGradient(
      tao, ctx->work,
      [](Tao, Vec x, Vec g, void *ptr) {
        return static_cast<TaoDiagonalNNLS *>(ptr)->evaluate(x, nullptr, g);
      },
      ctx);
  CHKERR ::TaoSetObjectiveAndGradient(
      tao, ctx->work,
      [](Tao, Vec x, PetscReal *f, Vec g, void *ptr) {
        return static_cast<TaoDiagonalNNLS *>(ptr)->evaluate(x, f, g);
      },
      ctx);
  MoFEMFunctionReturn(0);
}
 
} // namespace MoFEM