tao_diagonal_nnls.cpp

Go to the documentation of this file.

/** \file tao_diagonal_nnls.cpp
 * \brief Test diagonal NNLS and its use for null-space multipliers.
 */
 
#include <MoFEM.hpp>
 
using namespace MoFEM;
 
static char help[] = "Test diagonal nonnegative least squares.\n\n";
 
namespace {
 
template <size_t N>
MoFEMErrorCode readValues(Vec vec, std::array<PetscScalar, N> &values) {
  MoFEMFunctionBegin;
  values.fill(0);
  PetscInt first, last;
  const PetscScalar *data;
  CHKERR VecGetOwnershipRange(vec, &first, &last);
  CHKERR VecGetArrayRead(vec, &data);
  for (auto i = first; i != last; ++i)
    values[i] = data[i - first];
  CHKERR VecRestoreArrayRead(vec, &data);
  CHKERR MPIU_Allreduce(MPI_IN_PLACE, values.data(), N, MPIU_SCALAR, MPI_SUM,
                        PetscObjectComm(getPetscObject(vec)));
  MoFEMFunctionReturn(0);
}
 
template <size_t N>
MoFEMErrorCode writeValues(Vec vec, const std::array<PetscScalar, N> &values) {
  MoFEMFunctionBegin;
  PetscInt first, last;
  PetscScalar *data;
  CHKERR VecGetOwnershipRange(vec, &first, &last);
  CHKERR VecGetArray(vec, &data);
  for (auto i = first; i != last; ++i)
    data[i - first] = values[i];
  CHKERR VecRestoreArray(vec, &data);
  MoFEMFunctionReturn(0);
}
 
SmartPetscObj<Mat> createJacobian(Vec rows, Vec columns) {
  PetscInt local_rows, local_columns, global_rows, global_columns;
  CHK_THROW_MESSAGE(VecGetLocalSize(rows, &local_rows), "Get local row size");
  CHK_THROW_MESSAGE(VecGetLocalSize(columns, &local_columns),
                    "Get local column size");
  CHK_THROW_MESSAGE(VecGetSize(rows, &global_rows), "Get row size");
  CHK_THROW_MESSAGE(VecGetSize(columns, &global_columns), "Get column size");
  Mat matrix;
  CHK_THROW_MESSAGE(MatCreateAIJ(PETSC_COMM_WORLD, local_rows, local_columns,
                                 global_rows, global_columns, local_columns,
                                 nullptr, global_columns - local_columns,
                                 nullptr, &matrix),
                    "Create Jacobian");
  return SmartPetscObj<Mat>(matrix);
}
 
template <size_t N>
MoFEMErrorCode
writeRows(Mat matrix, const std::array<std::array<PetscScalar, 3>, N> &rows) {
  MoFEMFunctionBegin;
  PetscInt first, last;
  const PetscInt columns[] = {0, 1, 2};
  CHKERR MatGetOwnershipRange(matrix, &first, &last);
  for (auto row = first; row != last; ++row)
    CHKERR MatSetValues(matrix, 1, &row, 3, columns, rows[row].data(),
                        INSERT_VALUES);
  CHKERR MatAssemblyBegin(matrix, MAT_FINAL_ASSEMBLY);
  CHKERR MatAssemblyEnd(matrix, MAT_FINAL_ASSEMBLY);
  MoFEMFunctionReturn(0);
}
 
struct LeastSquares {
  Mat jacobian;
  Vec target;
  SmartPetscObj<Vec> residual;
 
  static PetscErrorCode evaluate(Tao, Vec x, PetscReal *f, Vec g, void *ctx) {
    MoFEMFunctionBegin;
    auto &problem = *static_cast<LeastSquares *>(ctx);
    CHKERR MatMultTranspose(problem.jacobian, x, problem.residual);
    CHKERR VecAXPY(problem.residual, -1, problem.target);
    PetscScalar squared;
    CHKERR VecDot(problem.residual, problem.residual, &squared);
    *f = 0.5 * PetscRealPart(squared);
    CHKERR MatMult(problem.jacobian, problem.residual, g);
    MoFEMFunctionReturn(0);
  }
};
 
MoFEMErrorCode checkDirectSolve() {
  MoFEMFunctionBegin;
  auto x = createVectorMPI(PETSC_COMM_WORLD, PETSC_DECIDE, 2);
  auto nonnegative = vectorDuplicate(x);
  CHKERR VecZeroEntries(nonnegative);
  auto target = createVectorMPI(PETSC_COMM_WORLD, PETSC_DECIDE, 3);
  auto jacobian = createJacobian(x, target);
  auto tao = createTao(PETSC_COMM_WORLD);
  CHKERR TaoSetType(tao, TAOSHELL);
  CHKERR TaoSetType(tao, TAODIAGONALNNLS);
  CHKERR TaoSetOptionsPrefix(tao, "direct_");
  CHKERR TaoSetSolution(tao, x);
  CHKERR TaoSetFromOptions(tao);
 
  struct Case {
    std::array<PetscScalar, 3> target;
    std::array<PetscScalar, 2> scales, expected;
    PetscReal objective;
  };
  const std::array<Case, 4> cases = {{{{3, 4, 1}, {1, 1}, {1, 0}, 0.5},
                                      {{7, 1, 2}, {1, 1}, {1, 0}, 14.5},
                                      {{0, 0, 0}, {1, 1}, {0, 0}, 0},
                                      {{-1, 7, 2}, {0.1, 10}, {10, 0.1}, 2}}};
  for (size_t i = 0; i != cases.size(); ++i) {
    const auto &test = cases[i];
    // Orthogonal rows deliberately overlap in their control-vector support.
    CHKERR writeRows(jacobian,
                     std::array<std::array<PetscScalar, 3>, 2>{
                         {{3 * test.scales[0], 4 * test.scales[0], 0},
                          {-4 * test.scales[1], 3 * test.scales[1], 0}}});
    SmartPetscObj<Vec> previous_target;
    if (i == 2) {
      previous_target = target;
      target = vectorDuplicate(target);
    }
    CHKERR writeValues(target, test.target);
    if (i == 0 || i == 2)
      CHKERR TaoDiagonalNNLSSetProblem(tao, jacobian, target);
    if (previous_target && previous_target.use_count() != 1)
      SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
              "Rebinding did not release the previous NNLS target");
    if (i == 1)
      CHKERR TaoSetVariableBounds(tao, nonnegative, nullptr);
    CHKERR VecSet(x, 7);
    CHKERR TaoSolve(tao);
    PetscReal objective, projected_gradient;
    TaoConvergedReason reason;
    CHKERR TaoGetSolutionStatus(tao, nullptr, &objective, &projected_gradient,
                                nullptr, nullptr, &reason);
    Vec stored_gradient;
    PetscReal stored_norm;
    CHKERR TaoGetGradient(tao, &stored_gradient, nullptr, nullptr);
    CHKERR VecNorm(stored_gradient, NORM_2, &stored_norm);
    std::array<PetscScalar, 2> values;
    CHKERR readValues(x, values);
    if (reason <= 0 || !std::isfinite(stored_norm) ||
        std::abs(stored_norm - projected_gradient) > 1e-10 ||
        std::abs(objective - test.objective) > 1e-10 ||
        projected_gradient > 1e-10 ||
        PetscAbsScalar(values[0] - test.expected[0]) > 1e-10 ||
        PetscAbsScalar(values[1] - test.expected[1]) > 1e-10)
      SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
              "Incorrect diagonal NNLS solution or convergence status");
 
    auto reference = createTao(PETSC_COMM_WORLD);
    auto reference_x = vectorDuplicate(x);
    auto gradient = vectorDuplicate(x);
    auto lower = vectorDuplicate(x);
    auto upper = vectorDuplicate(x);
    CHKERR VecZeroEntries(reference_x);
    CHKERR VecZeroEntries(lower);
    CHKERR VecSet(upper, PETSC_INFINITY);
    LeastSquares problem{jacobian, target, vectorDuplicate(target)};
    CHKERR TaoSetType(reference, TAOBLMVM);
    CHKERR TaoSetSolution(reference, reference_x);
    CHKERR TaoSetVariableBounds(reference, lower, upper);
    CHKERR ::TaoSetObjectiveAndGradient(reference, gradient,
                                        LeastSquares::evaluate, &problem);
    CHKERR TaoSetTolerances(reference, 1e-10, 0, 0);
    CHKERR TaoSetMaximumIterations(reference, 1000);
    CHKERR TaoSolve(reference);
    CHKERR TaoGetConvergedReason(reference, &reason);
    CHKERR VecAXPY(reference_x, -1, x);
    PetscReal difference;
    CHKERR VecNorm(reference_x, NORM_INFINITY, &difference);
    if (reason <= 0 || difference > 1e-7)
      SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
              "Diagonal NNLS disagrees with BLMVM");
  }
  auto short_bound = createVectorMPI(PETSC_COMM_WORLD, PETSC_DECIDE, 1);
  CHKERR VecZeroEntries(short_bound);
  for (const auto wrong_size : {true, false}) {
    if (!wrong_size)
      CHKERR VecSet(nonnegative, 1);
    CHKERR PetscPushErrorHandler(PetscReturnErrorHandler, nullptr);
    auto error = TaoSetVariableBounds(
        tao, wrong_size ? short_bound.get() : nonnegative.get(), nullptr);
    if (!error)
      error = TaoSolve(tao);
    CHKERR PetscPopErrorHandler();
    if (error != (wrong_size ? PETSC_ERR_ARG_SIZ : PETSC_ERR_SUP))
      SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
              "Diagonal NNLS did not reject unsupported bounds");
  }
  CHKERR VecZeroEntries(nonnegative);
  CHKERR TaoSetVariableBounds(tao, nonnegative, nullptr);
  if (jacobian.use_count() < 2 || target.use_count() < 2)
    SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
            "NNLS did not retain its problem data");
  CHKERR TaoSetType(tao, TAOLMVM);
  if (jacobian.use_count() != 1 || target.use_count() != 1)
    SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
            "Changing TAO type did not release NNLS problem data");
  CHKERR PetscPrintf(PETSC_COMM_WORLD,
                     "Diagonal NNLS: four analytic cases agree with BLMVM\n");
  MoFEMFunctionReturn(0);
}
 
MoFEMErrorCode checkInvalidRows(const bool zero_row) {
  MoFEMFunctionBegin;
  auto x = createVectorMPI(PETSC_COMM_WORLD, PETSC_DECIDE, 2);
  auto target = createVectorMPI(PETSC_COMM_WORLD, PETSC_DECIDE, 3);
  auto jacobian = createJacobian(x, target);
  CHKERR VecZeroEntries(x);
  CHKERR VecSet(target, 1);
  CHKERR writeRows(jacobian,
                   std::array<std::array<PetscScalar, 3>, 2>{
                       {zero_row ? std::array<PetscScalar, 3>{0, 0, 0}
                                 : std::array<PetscScalar, 3>{1, 1, 0},
                        {0, 1, 1}}});
  auto tao = createTao(PETSC_COMM_WORLD);
  CHKERR TaoSetType(tao, TAODIAGONALNNLS);
  CHKERR TaoSetSolution(tao, x);
  CHKERR PetscPushErrorHandler(PetscReturnErrorHandler, nullptr);
  auto error = TaoDiagonalNNLSSetProblem(tao, jacobian, target);
  if (!error)
    error = TaoSolve(tao);
  CHKERR PetscPopErrorHandler();
  if (error != (zero_row ? PETSC_ERR_ARG_OUTOFRANGE : PETSC_ERR_ARG_INCOMP))
    SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
            "Diagonal NNLS did not reject unsupported Jacobian rows");
  MoFEMFunctionReturn(0);
}
 
MoFEMErrorCode checkNullSpaceIntegration() {
  MoFEMFunctionBegin;
  auto x = createVectorMPI(PETSC_COMM_WORLD, PETSC_DECIDE, 3);
  auto gradient = vectorDuplicate(x);
  auto c = createVectorMPI(PETSC_COMM_WORLD, PETSC_DECIDE, 1);
  auto jacobian = createJacobian(c, x);
  auto dm = createDM(PETSC_COMM_WORLD, DMSHELL);
  auto tao = createTao(PETSC_COMM_WORLD);
  CHKERR writeValues(x, std::array<PetscScalar, 3>{1, 1, 3});
  CHKERR TaoSetType(tao, TAONULLSPACE);
  CHKERR TaoSetOptionsPrefix(tao, "cell_");
  CHKERR TaoSetSolution(tao, x);
  CHKERR ::TaoSetObjectiveAndGradient(
      tao, gradient,
      [](Tao, Vec x, PetscReal *f, Vec g, void *) {
        MoFEMFunctionBegin;
        std::array<PetscScalar, 3> p;
        CHKERR readValues(x, p);
        *f = 0.5 * (PetscSqr(p[0] - 3) + PetscSqr(p[1] - 4)) + p[2];
        CHKERR writeValues(g,
                           std::array<PetscScalar, 3>{p[0] - 3, p[1] - 4, 1});
        MoFEMFunctionReturn(0);
      },
      nullptr);
  CHKERR TaoSetInequalityConstraintsRoutine(
      tao, c,
      [](Tao, Vec x, Vec c, void *) {
        MoFEMFunctionBegin;
        std::array<PetscScalar, 3> p;
        CHKERR readValues(x, p);
        CHKERR VecSet(c, p[2] - std::hypot(p[0], p[1]));
        MoFEMFunctionReturn(0);
      },
      nullptr);
  CHKERR TaoSetJacobianInequalityRoutine(
      tao, jacobian, jacobian,
      [](Tao, Vec x, Mat j, Mat, void *) {
        MoFEMFunctionBegin;
        std::array<PetscScalar, 3> p;
        CHKERR readValues(x, p);
        const auto radius = std::hypot(p[0], p[1]);
        CHKERR writeRows(j, std::array<std::array<PetscScalar, 3>, 1>{
                                {{-p[0] / radius, -p[1] / radius, 1}}});
        MoFEMFunctionReturn(0);
      },
      nullptr);
  CHKERR TaoNullSpaceSetConstraints(tao, dm, c, jacobian);
  CHKERR TaoSetMaximumIterations(tao, 300);
  CHKERR TaoSetTolerances(tao, 1e-7, 0, 0);
  CHKERR TaoSetConstraintTolerances(tao, 1e-7, 0);
  CHKERR TaoSetFromOptions(tao);
  CHKERR TaoSolve(tao);
  PetscReal objective, residual;
  TaoConvergedReason reason;
  CHKERR TaoGetSolutionStatus(tao, nullptr, &objective, &residual, nullptr,
                              nullptr, &reason);
  std::array<PetscScalar, 3> values;
  CHKERR readValues(x, values);
  Vec multipliers;
  CHKERR TaoNullSpaceGetMultipliers(tao, &multipliers);
  std::array<PetscScalar, 1> dual, constraint;
  CHKERR readValues(multipliers, dual);
  CHKERR readValues(c, constraint);
  if (reason <= 0 || residual > 1e-7 || std::abs(objective - 4.5) > 5e-6 ||
      PetscAbsScalar(values[0] - 2.4) > 5e-6 ||
      PetscAbsScalar(values[1] - 3.2) > 5e-6 ||
      PetscAbsScalar(values[2] - 4) > 5e-6 ||
      PetscAbsScalar(dual[0] - 1) > 5e-6 ||
      PetscRealPart(constraint[0]) < -1e-7)
    SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
            "Incorrect null-space optimum with diagonal multiplier solver");
  CHKERR PetscPrintf(
      PETSC_COMM_WORLD,
      "Null-space diagonal multipliers: x=(%g,%g,%g), KKT=%g\n",
      static_cast<double>(values[0]), static_cast<double>(values[1]),
      static_cast<double>(values[2]), static_cast<double>(residual));
  MoFEMFunctionReturn(0);
}
 
} // namespace
 
int main(int argc, char *argv[]) {
  MoFEM::Core::Initialize(&argc, &argv, nullptr, help);
  try {
    CHKERR TaoRegisterDiagonalNNLS();
    CHKERR TaoRegisterNullSpace();
    CHKERR checkDirectSolve();
    CHKERR checkInvalidRows(false);
    CHKERR checkInvalidRows(true);
    CHKERR checkNullSpaceIntegration();
  }
  CATCH_ERRORS;
  CHKERR MoFEM::Core::Finalize();
  return 0;
}