tao_null_space.cpp

Go to the documentation of this file.

/** \file tao_null_space.cpp
 * \brief Test null-space TAO optimization of a scaled epigraph constraint.
 */
 
#include <MoFEM.hpp>
#include <limits>
 
using namespace MoFEM;
 
static char help[] =
    "Test the null-space TAO solver on an epigraph problem.\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 MPI_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);
}
 
struct Problem {
  PetscReal scale;
  PetscInt objectiveCalls = 0;
  PetscInt gradientCalls = 0;
  PetscInt constraintCalls = 0;
  PetscInt jacobianCalls = 0;
 
  MoFEMErrorCode evaluate(Vec x, PetscReal *value, Vec gradient) {
    MoFEMFunctionBegin;
    std::array<PetscScalar, 3> values;
    CHKERR readValues(x, values);
    const auto p1 = PetscRealPart(values[0]);
    const auto p2 = PetscRealPart(values[1]);
    if (value) {
      ++objectiveCalls;
      *value = 0.5 * ((p1 - 3) * (p1 - 3) + (p2 - 4) * (p2 - 4)) +
               PetscRealPart(values[2]);
    }
    if (gradient) {
      ++gradientCalls;
      CHKERR writeValues(gradient,
                         std::array<PetscScalar, 3>{p1 - 3, p2 - 4, 1});
    }
    MoFEMFunctionReturn(0);
  }
 
  static PetscErrorCode constraints(Tao, Vec x, Vec c, void *ctx) {
    MoFEMFunctionBegin;
    auto *problem_ptr = static_cast<Problem *>(ctx);
    ++problem_ptr->constraintCalls;
    std::array<PetscScalar, 3> values;
    CHKERR readValues(x, values);
    const auto radius =
        std::hypot(PetscRealPart(values[0]), PetscRealPart(values[1]));
    CHKERR writeValues(
        c, std::array<PetscScalar, 2>{problem_ptr->scale * (values[2] - radius),
                                      6 - values[2]});
    MoFEMFunctionReturn(0);
  }
 
  static PetscErrorCode jacobian(Tao, Vec x, Mat j, Mat j_pre, void *ctx) {
    MoFEMFunctionBegin;
    auto *problem_ptr = static_cast<Problem *>(ctx);
    ++problem_ptr->jacobianCalls;
    std::array<PetscScalar, 3> values;
    CHKERR readValues(x, values);
    const auto radius =
        std::hypot(PetscRealPart(values[0]), PetscRealPart(values[1]));
    const auto scale = problem_ptr->scale;
    const PetscScalar rows[2][3] = {{radius ? -scale * values[0] / radius : 0,
                                     radius ? -scale * values[1] / radius : 0,
                                     scale},
                                    {0, 0, -1}};
    const PetscInt columns[] = {0, 1, 2};
    for (auto matrix : {j, j_pre}) {
      PetscInt first, last;
      CHKERR MatGetOwnershipRange(matrix, &first, &last);
      for (auto row = first; row != last; ++row)
        CHKERR MatSetValues(matrix, 1, &row, 3, columns, rows[row],
                            INSERT_VALUES);
      CHKERR MatAssemblyBegin(matrix, MAT_FINAL_ASSEMBLY);
      CHKERR MatAssemblyEnd(matrix, MAT_FINAL_ASSEMBLY);
    }
    MoFEMFunctionReturn(0);
  }
};
 
MoFEMErrorCode checkUnavailableSubproblem(Tao inner, Vec x, Vec gradient) {
  MoFEMFunctionBegin;
  PetscReal value;
  CHKERR PetscPushErrorHandler(PetscReturnErrorHandler, nullptr);
  const auto objective_error = TaoComputeObjective(inner, x, &value);
  const auto gradient_error = TaoComputeGradient(inner, x, gradient);
  CHKERR PetscPopErrorHandler();
  // PETSc's compute routines leave the input read lock on error.
  if (gradient_error)
    CHKERR VecLockReadPop(x);
  if (objective_error)
    CHKERR VecLockReadPop(x);
  if (objective_error != PETSC_ERR_ARG_WRONGSTATE ||
      gradient_error != PETSC_ERR_ARG_WRONGSTATE)
    SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
            "Inner callbacks require a live parent and a frozen outer iterate");
  MoFEMFunctionReturn(0);
}
 
MoFEMErrorCode checkLineSearchRequests(Tao inner, Vec start, Problem &problem) {
  MoFEMFunctionBegin;
  auto trial = vectorDuplicate(start);
  auto direction = vectorDuplicate(start);
  auto gradient = vectorDuplicate(start);
  auto expected_gradient = vectorDuplicate(start);
  TaoLineSearch line_search;
  CHKERR TaoGetLineSearch(inner, &line_search);
  for (const auto type : {TAOLINESEARCHARMIJO, "nullspace"}) {
    CHKERR TaoLineSearchSetType(line_search, type);
    CHKERR TaoLineSearchReset(line_search);
    CHKERR TaoLineSearchUseTaoRoutines(line_search, inner);
    CHKERR TaoLineSearchSetInitialStepLength(line_search, 1);
    CHKERR VecCopy(start, trial);
    PetscReal value, step;
    CHKERR TaoComputeObjectiveAndGradient(inner, trial, &value, gradient);
    // Force backtracking so rejected values outnumber derivative evaluations.
    CHKERR VecCopy(gradient, direction);
    CHKERR VecScale(direction, -100);
    const auto objective_calls = problem.objectiveCalls;
    const auto gradient_calls = problem.gradientCalls;
    const auto jacobian_calls = problem.jacobianCalls;
    TaoLineSearchConvergedReason reason;
    CHKERR TaoLineSearchApply(line_search, trial, &value, gradient, direction,
                              &step, &reason);
    if (reason <= 0 || !(step < 1) ||
        problem.objectiveCalls <= objective_calls + 1 ||
        problem.gradientCalls != gradient_calls + 1 ||
        problem.jacobianCalls != jacobian_calls + 1)
      SETERRQ(PETSC_COMM_WORLD, MOFEM_DATA_INCONSISTENCY,
              "%s backtracking must evaluate derivatives only at the accepted "
              "trial: objective/gradient/Jacobian calls %d/%d/%d, step %g",
              type, static_cast<int>(problem.objectiveCalls - objective_calls),
              static_cast<int>(problem.gradientCalls - gradient_calls),
              static_cast<int>(problem.jacobianCalls - jacobian_calls),
              static_cast<double>(step));
    PetscReal expected_value, gradient_error;
    CHKERR TaoComputeObjectiveAndGradient(inner, trial, &expected_value,
                                          expected_gradient);
    CHKERR VecAXPY(expected_gradient, -1, gradient);
    CHKERR VecNorm(expected_gradient, NORM_INFINITY, &gradient_error);
    if (std::abs(value - expected_value) > 1e-12 || gradient_error > 1e-12)
      SETERRQ(PETSC_COMM_WORLD, MOFEM_DATA_INCONSISTENCY,
              "%s backtracking returned inconsistent value and gradient",
              type);
  }
  MoFEMFunctionReturn(0);
}
 
MoFEMErrorCode checkFrozenMerit(Tao inner, Vec solution, Problem &problem,
                                const bool combined) {
  MoFEMFunctionBegin;
  auto trial = vectorDuplicate(solution);
  auto direction = vectorDuplicate(solution);
  auto gradient = vectorDuplicate(solution);
  CHKERR VecCopy(solution, trial);
  CHKERR writeValues(direction, std::array<PetscScalar, 3>{0.2, -0.3, 0.1});
  CHKERR VecAXPY(trial, 0.01, direction);
  const auto objective_calls = problem.objectiveCalls;
  const auto gradient_calls = problem.gradientCalls;
  const auto constraint_calls = problem.constraintCalls;
  const auto jacobian_calls = problem.jacobianCalls;
  PetscReal value;
  CHKERR TaoComputeObjective(inner, trial, &value);
  if (problem.objectiveCalls != objective_calls + 1 ||
      problem.constraintCalls != constraint_calls + 1 ||
      problem.jacobianCalls != jacobian_calls ||
      (!combined && problem.gradientCalls != gradient_calls))
    SETERRQ(PETSC_COMM_WORLD, MOFEM_DATA_INCONSISTENCY,
            "Inner objective requested unnecessary derivatives");
  CHKERR TaoComputeGradient(inner, trial, gradient);
  if (problem.gradientCalls != gradient_calls + (combined ? 2 : 1) ||
      problem.jacobianCalls != jacobian_calls + 1)
    SETERRQ(PETSC_COMM_WORLD, MOFEM_DATA_INCONSISTENCY,
            "Inner gradient did not evaluate the required derivatives");
  auto combined_gradient = vectorDuplicate(gradient);
  PetscReal combined_value, gradient_error;
  CHKERR TaoComputeObjectiveAndGradient(inner, trial, &combined_value,
                                        combined_gradient);
  CHKERR VecAXPY(combined_gradient, -1, gradient);
  CHKERR VecNorm(combined_gradient, NORM_INFINITY, &gradient_error);
  if (std::abs(value - combined_value) > 1e-12 || gradient_error > 1e-12)
    SETERRQ(PETSC_COMM_WORLD, MOFEM_DATA_INCONSISTENCY,
            "Separate and combined frozen merit evaluations differ");
  PetscScalar derivative;
  CHKERR VecDot(gradient, direction, &derivative);
  constexpr PetscReal epsilon = 1e-6;
  PetscReal plus, minus;
  CHKERR VecAXPY(trial, epsilon, direction);
  CHKERR TaoComputeObjective(inner, trial, &plus);
  CHKERR VecAXPY(trial, -2 * epsilon, direction);
  CHKERR TaoComputeObjective(inner, trial, &minus);
  if (std::abs((plus - minus) / (2 * epsilon) - PetscRealPart(derivative)) >
      2e-6 * (1 + PetscAbsScalar(derivative)))
    SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
            "Inner gradient is not the derivative of its frozen scalar merit");
  if (!combined)
    CHKERR checkLineSearchRequests(inner, trial, problem);
  MoFEMFunctionReturn(0);
}
 
MoFEMErrorCode solveProblem(const std::array<PetscScalar, 3> &start,
                            const PetscReal scale, const bool combined) {
  MoFEMFunctionBegin;
  Problem problem{scale};
  auto x = createVectorMPI(PETSC_COMM_WORLD, PETSC_DECIDE, 3);
  auto g = vectorDuplicate(x);
  auto c = createVectorMPI(PETSC_COMM_WORLD, PETSC_DECIDE, 2);
  auto dm = createDM(PETSC_COMM_WORLD, DMSHELL);
  CHKERR writeValues(x, start);
  PetscInt local_rows, local_columns;
  CHKERR VecGetLocalSize(c, &local_rows);
  CHKERR VecGetLocalSize(x, &local_columns);
  Mat matrix;
  CHKERR MatCreateAIJ(PETSC_COMM_WORLD, local_rows, local_columns, 2, 3,
                      local_columns, nullptr, 3 - local_columns, nullptr,
                      &matrix);
  SmartPetscObj<Mat> jacobian(matrix);
  CHKERR MatCreateAIJ(PETSC_COMM_WORLD, local_rows, local_columns, 2, 3,
                      local_columns, nullptr, 3 - local_columns, nullptr,
                      &matrix);
  SmartPetscObj<Mat> preconditioner(matrix);
  auto tao = createTao(PETSC_COMM_WORLD);
  CHKERR TaoSetType(tao, combined ? TAOSHELL : TAOALMM);
  CHKERR TaoSetOptionsPrefix(tao, "test_");
  CHKERR TaoSetSolution(tao, x);
  if (combined) {
    CHKERR ::TaoSetObjectiveAndGradient(
        tao, g,
        [](Tao, Vec x, PetscReal *f, Vec g, void *ctx) {
          return static_cast<Problem *>(ctx)->evaluate(x, f, g);
        },
        &problem);
  } else {
    CHKERR ::TaoSetObjective(
        tao,
        [](Tao, Vec x, PetscReal *f, void *ctx) {
          return static_cast<Problem *>(ctx)->evaluate(x, f, nullptr);
        },
        &problem);
    CHKERR ::TaoSetGradient(
        tao, g,
        [](Tao, Vec x, Vec g, void *ctx) {
          return static_cast<Problem *>(ctx)->evaluate(x, nullptr, g);
        },
        &problem);
  }
  CHKERR TaoSetInequalityConstraintsRoutine(tao, c, Problem::constraints,
                                            &problem);
  CHKERR TaoSetJacobianInequalityRoutine(tao, jacobian, preconditioner,
                                         Problem::jacobian, &problem);
  if (combined)
    CHKERR TaoSetType(tao, TAONULLSPACE);
  else
    CHKERR TaoSetFromOptions(tao);
  CHKERR TaoNullSpaceSetConstraints(tao, dm, c, jacobian, preconditioner);
  CHKERR TaoSetMaximumIterations(tao, 300);
  CHKERR TaoSetTolerances(tao, 1e-7, 0, 0);
  CHKERR TaoSetConstraintTolerances(tao, 1e-7, 0);
  CHKERR TaoSetFromOptions(tao);
  PetscBool matches;
  CHKERR PetscObjectTypeCompare(getPetscObject(tao.get()), TAONULLSPACE,
                                &matches);
  if (!matches)
    SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY, "Incorrect TAO type");
  Tao inner;
  CHKERR TaoNullSpaceGetSubsolver(tao, &inner);
  const char *prefix;
  CHKERR TaoGetOptionsPrefix(inner, &prefix);
  CHKERR PetscStrcmp(prefix, "test_tao_nullspace_", &matches);
  if (!matches)
    SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
            "Inner TAO did not inherit its options prefix");
  CHKERR TaoSetUp(tao);
  Vec sub_solution, sub_gradient, multipliers;
  CHKERR TaoGetSolution(inner, &sub_solution);
  CHKERR TaoGetGradient(inner, &sub_gradient, nullptr, nullptr);
  CHKERR TaoNullSpaceGetMultipliers(tao, &multipliers);
  CHKERR VecEqual(x, sub_solution, &matches);
  if (!matches || sub_solution == x.get() || sub_gradient == g.get())
    SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
            "Inner TAO needs an independent copy of the parent state");
  CHKERR checkUnavailableSubproblem(inner, x, g);
  CHKERR TaoSolve(tao);
  TaoConvergedReason reason;
  CHKERR TaoGetConvergedReason(tao, &reason);
  if (reason <= 0)
    SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
            "Epigraph solve did not converge (scale %g, reason %d)",
            static_cast<double>(scale), static_cast<int>(reason));
 
  PetscBool check_recycling = PETSC_FALSE;
  CHKERR PetscOptionsHasName(nullptr, prefix, "-tao_lmvm_recycle",
                              &check_recycling);
  if (check_recycling)
    CHKERR TaoView(inner, PETSC_VIEWER_STDOUT_WORLD);
 
  std::array<PetscScalar, 3> solution;
  std::array<PetscScalar, 2> dual;
  CHKERR readValues(x, solution);
  CHKERR readValues(multipliers, dual);
  constexpr PetscReal tolerance = 5e-6;
  const PetscReal radius =
      std::hypot(PetscRealPart(solution[0]), PetscRealPart(solution[1]));
  if (PetscAbsScalar(solution[0] - 2.4) > tolerance ||
      PetscAbsScalar(solution[1] - 3.2) > tolerance ||
      PetscAbsScalar(solution[2] - 4) > tolerance ||
      PetscAbsScalar(scale * dual[0] - 1) > tolerance ||
      PetscAbsScalar(dual[1]) > tolerance ||
      PetscRealPart(solution[2]) - radius < -tolerance ||
      PetscRealPart(6 - solution[2]) < -tolerance)
    SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
            "Incorrect epigraph optimum (scale %g): x=(%g,%g,%g), "
            "scaled multipliers=(%g,%g)",
            static_cast<double>(scale), static_cast<double>(solution[0]),
            static_cast<double>(solution[1]), static_cast<double>(solution[2]),
            static_cast<double>(scale * dual[0]), static_cast<double>(dual[1]));
  PetscReal objective, stationarity;
  CHKERR TaoComputeObjectiveAndGradient(tao, x, &objective, g);
  CHKERR TaoComputeJacobianInequality(tao, x, jacobian, preconditioner);
  auto residual = vectorDuplicate(x);
  CHKERR MatMultTranspose(jacobian, multipliers, residual);
  CHKERR VecAYPX(residual, -1, g);
  CHKERR VecNorm(residual, NORM_INFINITY, &stationarity);
  if (std::abs(objective - 4.5) > tolerance || stationarity > tolerance ||
      !problem.objectiveCalls || !problem.gradientCalls ||
      !problem.constraintCalls || !problem.jacobianCalls)
    SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
            "Epigraph objective or constrained stationarity is incorrect");
  CHKERR checkFrozenMerit(inner, x, problem, combined);
  CHKERR PetscPrintf(
      PETSC_COMM_WORLD,
      "Epigraph scale %g start (%g,%g,%g): x=(%g,%g,%g), f=%g, KKT=%g\n",
      static_cast<double>(scale), static_cast<double>(start[0]),
      static_cast<double>(start[1]), static_cast<double>(start[2]),
      static_cast<double>(solution[0]), static_cast<double>(solution[1]),
      static_cast<double>(solution[2]), static_cast<double>(objective),
      static_cast<double>(stationarity));
 
  SmartPetscObj<Tao> retained_inner(inner, true);
  SmartPetscObj<Vec> retained_multipliers(multipliers, true);
  if (combined)
    tao.reset();
  else
    CHKERR TaoSetType(tao, TAOLMVM);
  if (retained_inner.use_count() != 1 ||
      retained_multipliers.use_count() != 1 || dm.use_count() != 1)
    SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
            "Parent destruction did not release null-space solver resources");
  CHKERR checkUnavailableSubproblem(inner, x, g);
  MoFEMFunctionReturn(0);
}
 
struct ActivationProblem {
  PetscReal scale;
 
  static PetscErrorCode objective(Tao, Vec x, PetscReal *f, Vec g, void *) {
    MoFEMFunctionBegin;
    std::array<PetscScalar, 2> values;
    CHKERR readValues(x, values);
    const auto k = PetscRealPart(values[0] * values[0] + values[1]);
    *f = k + 0.5 * k * k;
    CHKERR writeValues(
        g, std::array<PetscScalar, 2>{2 * values[0] * (1 + k), 1 + k});
    MoFEMFunctionReturn(0);
  }
 
  static PetscErrorCode constraints(Tao, Vec x, Vec c, void *ctx) {
    MoFEMFunctionBegin;
    const auto scale = static_cast<ActivationProblem *>(ctx)->scale;
    std::array<PetscScalar, 2> values;
    CHKERR readValues(x, values);
    CHKERR VecSet(c, scale * (values[0] * values[0] + values[1]));
    MoFEMFunctionReturn(0);
  }
 
  static PetscErrorCode jacobian(Tao, Vec x, Mat j, Mat, void *ctx) {
    MoFEMFunctionBegin;
    const auto scale = static_cast<ActivationProblem *>(ctx)->scale;
    std::array<PetscScalar, 2> values;
    CHKERR readValues(x, values);
    const PetscInt columns[] = {0, 1};
    const PetscScalar row_values[] = {2 * scale * values[0], scale};
    PetscInt first, last;
    CHKERR MatGetOwnershipRange(j, &first, &last);
    for (auto row = first; row != last; ++row)
      CHKERR MatSetValues(j, 1, &row, 2, columns, row_values, INSERT_VALUES);
    CHKERR MatAssemblyBegin(j, MAT_FINAL_ASSEMBLY);
    CHKERR MatAssemblyEnd(j, MAT_FINAL_ASSEMBLY);
    MoFEMFunctionReturn(0);
  }
};
 
MoFEMErrorCode checkCandidateActivation() {
  MoFEMFunctionBegin;
  PetscReal active_distance = 1e-6;
  CHKERR PetscOptionsGetReal(nullptr, "activation_",
                             "-tao_nullspace_active_tol", &active_distance,
                             nullptr);
  for (const auto scale : {0.1, 1.0, 10.0}) {
    ActivationProblem problem{scale};
    auto x = createVectorMPI(PETSC_COMM_WORLD, PETSC_DECIDE, 2);
    auto g = vectorDuplicate(x);
    auto sub_gradient = vectorDuplicate(x);
    auto c = createVectorMPI(PETSC_COMM_WORLD, PETSC_DECIDE, 1);
    PetscInt local_rows, local_columns;
    CHKERR VecGetLocalSize(c, &local_rows);
    CHKERR VecGetLocalSize(x, &local_columns);
    Mat matrix;
    CHKERR MatCreateAIJ(PETSC_COMM_WORLD, local_rows, local_columns, 1, 2,
                        local_columns, nullptr, 2 - local_columns, nullptr,
                        &matrix);
    SmartPetscObj<Mat> jacobian(matrix);
    auto tao = createTao(PETSC_COMM_WORLD);
    CHKERR TaoSetType(tao, TAONULLSPACE);
    CHKERR TaoSetOptionsPrefix(tao, "activation_");
    CHKERR TaoSetSolution(tao, x);
    CHKERR ::TaoSetObjectiveAndGradient(tao, g, ActivationProblem::objective,
                                        &problem);
    CHKERR TaoSetInequalityConstraintsRoutine(
        tao, c, ActivationProblem::constraints, &problem);
    CHKERR TaoSetJacobianInequalityRoutine(
        tao, jacobian, jacobian, ActivationProblem::jacobian, &problem);
    CHKERR TaoNullSpaceSetConstraints(tao, nullptr, c, jacobian);
    CHKERR TaoSetFromOptions(tao);
    // Freeze the current epoch without taking an optimization step.
    CHKERR TaoSetConvergenceTest(
        tao,
        [](Tao tao, void *) {
          return TaoSetConvergedReason(tao, TAO_CONVERGED_USER);
        },
        nullptr);
    // Reusing the solver checks that nonlinear row norms refresh each epoch.
    // On two ranks, the only row also contains an off-rank column.
    for (const auto first_value : {0.5, 1.0}) {
      const auto norm_squared = 4 * first_value * first_value + 1;
      for (const auto distance_fraction : {0.95, 1.05}) {
        const auto k =
            distance_fraction * active_distance * std::sqrt(norm_squared);
        CHKERR writeValues(x, std::array<PetscScalar, 2>{
                                  first_value, k - first_value * first_value});
        CHKERR TaoSolve(tao);
        Vec multipliers;
        Tao inner;
        CHKERR TaoNullSpaceGetMultipliers(tao, &multipliers);
        CHKERR TaoNullSpaceGetSubsolver(tao, &inner);
        std::array<PetscScalar, 1> lambda;
        std::array<PetscScalar, 2> merit_gradient;
        CHKERR readValues(multipliers, lambda);
        PetscReal merit;
        CHKERR TaoComputeObjectiveAndGradient(inner, x, &merit, sub_gradient);
        CHKERR readValues(sub_gradient, merit_gradient);
        const bool active = distance_fraction < 1;
        const auto expected_multiplier = active ? 1 + k : 0;
        const auto expected_merit =
            active ? -0.5 * k * k + 0.5 * k * k / norm_squared
                   : k + 0.5 * k * k;
        const auto gradient_factor = active ? k / norm_squared : 1 + k;
        constexpr PetscReal tolerance = 1e-9;
        if (PetscAbsScalar(scale * lambda[0] - expected_multiplier) >
                tolerance ||
            std::abs(merit - expected_merit) > tolerance ||
            PetscAbsScalar(merit_gradient[0] -
                           2 * first_value * gradient_factor) > tolerance ||
            PetscAbsScalar(merit_gradient[1] - gradient_factor) > tolerance)
          SETERRQ(PETSC_COMM_WORLD, MOFEM_DATA_INCONSISTENCY,
                  "Candidate activation or restoration changed with row "
                  "scaling: scale %g, x0 %g, normalized distance %g h",
                  static_cast<double>(scale), static_cast<double>(first_value),
                  static_cast<double>(distance_fraction));
      }
    }
  }
  MoFEMFunctionReturn(0);
}
 
struct LineSearchProblem {
  bool objectiveDomain;
  PetscReal domainUpper;
  PetscInt rejectedObjectives = 0;
  PetscReal objectivePoint = 0;
  PetscReal constraintPoint = 0;
  PetscReal jacobianPoint = 0;
 
  static PetscErrorCode objective(Tao, Vec x, PetscReal *f, Vec g, void *ctx) {
    MoFEMFunctionBegin;
    auto &problem = *static_cast<LineSearchProblem *>(ctx);
    std::array<PetscScalar, 1> values;
    CHKERR readValues(x, values);
    problem.objectivePoint = PetscRealPart(values[0]);
    if (problem.objectiveDomain &&
        problem.objectivePoint > problem.domainUpper) {
      ++problem.rejectedObjectives;
      *f = PETSC_INFINITY;
      CHKERR VecSet(g, std::numeric_limits<PetscReal>::quiet_NaN());
      MoFEMFunctionReturnHot(0);
    }
    *f = 0.5 * PetscSqr(problem.objectivePoint - 2);
    CHKERR VecSet(g, problem.objectivePoint - 2);
    MoFEMFunctionReturn(0);
  }
 
  static PetscErrorCode constraints(Tao, Vec x, Vec c, void *ctx) {
    MoFEMFunctionBegin;
    auto &problem = *static_cast<LineSearchProblem *>(ctx);
    std::array<PetscScalar, 1> values;
    CHKERR readValues(x, values);
    problem.constraintPoint = PetscRealPart(values[0]);
    if (problem.objectiveDomain &&
        problem.constraintPoint > problem.domainUpper)
      SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
              "Constraints evaluated after an invalid objective");
    CHKERR VecSet(c, 1 - problem.constraintPoint);
    MoFEMFunctionReturn(0);
  }
 
  static PetscErrorCode jacobian(Tao, Vec x, Mat j, Mat, void *ctx) {
    MoFEMFunctionBegin;
    auto &problem = *static_cast<LineSearchProblem *>(ctx);
    std::array<PetscScalar, 1> values;
    CHKERR readValues(x, values);
    problem.jacobianPoint = PetscRealPart(values[0]);
    if (problem.objectiveDomain && problem.jacobianPoint > problem.domainUpper)
      SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
              "Jacobian evaluated after an invalid objective");
    PetscInt first, last;
    CHKERR MatGetOwnershipRange(j, &first, &last);
    for (auto row = first; row != last; ++row)
      CHKERR MatSetValue(j, row, 0, -1, INSERT_VALUES);
    CHKERR MatAssemblyBegin(j, MAT_FINAL_ASSEMBLY);
    CHKERR MatAssemblyEnd(j, MAT_FINAL_ASSEMBLY);
    MoFEMFunctionReturn(0);
  }
};
 
MoFEMErrorCode checkObjectiveDomainLineSearch(const bool exhaust_budget) {
  MoFEMFunctionBegin;
  // With gtol=0.1, no x <= 1 satisfies strong Wolfe for this descent ray.
  // An objective-domain rejection needs an Armijo step with its true derivative.
  LineSearchProblem problem{true, exhaust_budget ? 0.0 : 1.0};
  auto x = createVectorMPI(PETSC_COMM_WORLD, PETSC_DECIDE, 1);
  auto g = vectorDuplicate(x);
  auto c = vectorDuplicate(x);
  auto dm = createDM(PETSC_COMM_WORLD, DMSHELL);
  CHKERR VecZeroEntries(x);
  PetscInt local_size;
  CHKERR VecGetLocalSize(x, &local_size);
  Mat matrix;
  CHKERR MatCreateAIJ(PETSC_COMM_WORLD, local_size, local_size, 1, 1,
                      local_size, nullptr, 1 - local_size, nullptr, &matrix);
  SmartPetscObj<Mat> jacobian(matrix);
  auto tao = createTao(PETSC_COMM_WORLD);
  CHKERR TaoSetType(tao, TAONULLSPACE);
  CHKERR TaoSetOptionsPrefix(tao, "domain_");
  CHKERR TaoSetSolution(tao, x);
  CHKERR ::TaoSetObjectiveAndGradient(tao, g, LineSearchProblem::objective,
                                      &problem);
  CHKERR TaoSetInequalityConstraintsRoutine(
      tao, c, LineSearchProblem::constraints, &problem);
  CHKERR TaoSetJacobianInequalityRoutine(tao, jacobian, jacobian,
                                         LineSearchProblem::jacobian, &problem);
  CHKERR TaoNullSpaceSetConstraints(tao, dm, c, jacobian);
  CHKERR TaoSetMaximumIterations(tao, 1);
  CHKERR TaoSetTolerances(tao, 1e-10, 0, 0);
  CHKERR TaoSetConstraintTolerances(tao, 1e-10, 0);
  CHKERR TaoSetFromOptions(tao);
  CHKERR TaoSolve(tao);
  TaoConvergedReason reason;
  CHKERR TaoGetConvergedReason(tao, &reason);
  std::array<PetscScalar, 1> solution, constraints;
  CHKERR readValues(x, solution);
  CHKERR readValues(c, constraints);
  const auto accepted = PetscRealPart(solution[0]);
  constexpr PetscReal tolerance = 1e-10;
  if (!problem.rejectedObjectives ||
      std::abs(problem.objectivePoint - accepted) > tolerance ||
      std::abs(problem.constraintPoint - accepted) > tolerance ||
      std::abs(problem.jacobianPoint - accepted) > tolerance ||
      std::abs(PetscRealPart(constraints[0]) - (1 - accepted)) > tolerance)
    SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
            "Objective-domain line search did not preserve valid application "
            "state");
  if (exhaust_budget) {
    if (reason != TAO_DIVERGED_LS_FAILURE || accepted != 0)
      SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
              "Exhausted line search must fail and restore the initial state");
  } else {
    if (!(accepted > 0) || accepted > 1 + tolerance ||
        (reason <= 0 && reason != TAO_DIVERGED_MAXITS))
      SETERRQ(
          PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
          "Objective-domain line search did not accept a valid descent step");
    Tao inner;
    TaoLineSearch line_search;
    CHKERR TaoNullSpaceGetSubsolver(tao, &inner);
    CHKERR TaoGetLineSearch(inner, &line_search);
    Vec ls_x, ls_g;
    CHKERR TaoGetSolution(inner, &ls_x);
    CHKERR TaoGetGradient(inner, &ls_g, nullptr, nullptr);
    PetscReal objective, step;
    TaoLineSearchConvergedReason ls_reason;
    CHKERR TaoLineSearchGetSolution(line_search, ls_x, &objective, ls_g, &step,
                                    &ls_reason);
    std::array<PetscScalar, 1> ls_solution, ls_gradient, direction;
    CHKERR readValues(ls_x, ls_solution);
    CHKERR readValues(ls_g, ls_gradient);
    Vec step_direction;
    CHKERR TaoLineSearchGetStepDirection(line_search, &step_direction);
    CHKERR readValues(step_direction, direction);
    if (ls_reason != TAOLINESEARCH_SUCCESS_USER || !(step > 0) ||
        PetscAbsScalar(ls_solution[0] - accepted) > tolerance ||
        PetscAbsScalar(step * direction[0] - accepted) > tolerance ||
        PetscAbsScalar(ls_gradient[0] - (accepted - 2)) > tolerance ||
        std::abs(objective - 0.5 * PetscSqr(accepted - 2)) > tolerance ||
        objective > 2 - 1e-4 * 2 * accepted)
      SETERRQ(
          PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
          "Domain fallback returned inconsistent objective, gradient or step");
  }
  CHKERR PetscPrintf(
      PETSC_COMM_WORLD,
      "Objective-domain line search budget exhaustion %d: "
      "x=%g, rejected objectives=%d, reason=%d\n",
      static_cast<int>(exhaust_budget), static_cast<double>(accepted),
      static_cast<int>(problem.rejectedObjectives), static_cast<int>(reason));
  MoFEMFunctionReturn(0);
}
 
MoFEMErrorCode checkAlgorithmOneStep(const bool exhaust_budget) {
  MoFEMFunctionBegin;
  LineSearchProblem problem{false, 0};
  auto x = createVectorMPI(PETSC_COMM_WORLD, PETSC_DECIDE, 1);
  auto g = vectorDuplicate(x);
  auto c = vectorDuplicate(x);
  CHKERR VecZeroEntries(x);
  PetscInt local_size;
  CHKERR VecGetLocalSize(x, &local_size);
  Mat matrix;
  CHKERR MatCreateAIJ(PETSC_COMM_WORLD, local_size, local_size, 1, 1,
                      local_size, nullptr, 1 - local_size, nullptr, &matrix);
  SmartPetscObj<Mat> jacobian(matrix);
  auto tao = createTao(PETSC_COMM_WORLD);
  CHKERR TaoSetType(tao, TAONULLSPACE);
  CHKERR TaoSetOptionsPrefix(tao, "flow_");
  CHKERR TaoSetSolution(tao, x);
  CHKERR ::TaoSetObjectiveAndGradient(tao, g, LineSearchProblem::objective,
                                      &problem);
  CHKERR TaoSetInequalityConstraintsRoutine(
      tao, c, LineSearchProblem::constraints, &problem);
  CHKERR TaoSetJacobianInequalityRoutine(tao, jacobian, jacobian,
                                         LineSearchProblem::jacobian, &problem);
  CHKERR TaoNullSpaceSetConstraints(tao, nullptr, c, jacobian);
  CHKERR TaoSetTolerances(tao, 1e-10, 0, 0);
  CHKERR TaoSetConstraintTolerances(tao, 1e-10, 0);
  CHKERR TaoSetMaximumIterations(tao, exhaust_budget ? 1 : 2);
  Tao inner;
  CHKERR TaoNullSpaceGetSubsolver(tao, &inner);
  // Returning to shell must recreate its line search and flow callback.
  CHKERR TaoSetType(inner, TAOLMVM);
  CHKERR TaoSetType(inner, TAOSHELL);
  CHKERR TaoSetFromOptions(tao);
  TaoLineSearch line_search;
  CHKERR TaoGetLineSearch(inner, &line_search);
  if (!line_search)
    SETERRQ(
        PETSC_COMM_WORLD, MOFEM_DATA_INCONSISTENCY,
        "Returning the inner TAO to shell did not recreate its line search");
  for (const auto type : {TAOLINESEARCHARMIJO, TAOLINESEARCHMT}) {
    CHKERR TaoLineSearchSetType(line_search, type);
    CHKERR TaoLineSearchSetType(line_search, "nullspace");
    // The previous implementation's reset hook must not survive the switch.
    CHKERR TaoLineSearchReset(line_search);
  }
  CHKERR TaoLineSearchSetFromOptions(line_search);
  if (!exhaust_budget) {
    CHKERR TaoMonitorSet(
        tao,
        [](Tao tao, void *) {
          MoFEMFunctionBegin;
          PetscInt iteration;
          CHKERR TaoGetIterationNumber(tao, &iteration);
          if (iteration == 1) {
            Vec solution;
            CHKERR TaoGetSolution(tao, &solution);
            std::array<PetscScalar, 1> values;
            CHKERR readValues(solution, values);
            // The first -grad(Phi) step crosses the previously inactive row.
            if (PetscAbsScalar(values[0] - 2) > 1e-12)
              SETERRQ(PETSC_COMM_WORLD, MOFEM_DATA_INCONSISTENCY,
                      "The first null-space flow step was blocked by an "
                      "inactive inequality");
          }
          MoFEMFunctionReturn(0);
        },
        nullptr, nullptr);
  }
  CHKERR TaoSolve(tao);
  TaoConvergedReason reason;
  CHKERR TaoGetConvergedReason(tao, &reason);
  std::array<PetscScalar, 1> solution, dual;
  Vec multipliers;
  CHKERR TaoNullSpaceGetMultipliers(tao, &multipliers);
  CHKERR readValues(x, solution);
  CHKERR readValues(multipliers, dual);
  const auto expected_solution = exhaust_budget ? 6.0 : 1.5;
  const auto expected_multiplier = exhaust_budget ? 0.0 : 0.5;
  if (reason != TAO_DIVERGED_MAXITS ||
      PetscAbsScalar(solution[0] - expected_solution) > 1e-10 ||
      PetscAbsScalar(dual[0] - expected_multiplier) > 1e-10)
    SETERRQ(PETSC_COMM_WORLD, MOFEM_DATA_INCONSISTENCY,
            "Incorrect null-space flow step or refreshed active set: "
            "budget %d, x=%g, lambda=%g, reason=%d",
            static_cast<int>(exhaust_budget),
            static_cast<double>(PetscRealPart(solution[0])),
            static_cast<double>(PetscRealPart(dual[0])),
            static_cast<int>(reason));
 
  if (exhaust_budget) {
    Vec inner_x, inner_g;
    CHKERR TaoGetSolution(inner, &inner_x);
    CHKERR TaoGetGradient(inner, &inner_g, nullptr, nullptr);
    PetscReal objective, step;
    TaoLineSearchConvergedReason ls_reason;
    CHKERR TaoLineSearchGetSolution(line_search, inner_x, &objective, inner_g,
                                    &step, &ls_reason);
    std::array<PetscScalar, 1> gradient;
    CHKERR readValues(inner_g, gradient);
    // With one trial and stepinit=3, x=6 raises the finite merit from 2 to 8.
    if (ls_reason != TAOLINESEARCH_SUCCESS_USER ||
        std::abs(objective - 8) > 1e-12 || std::abs(step - 3) > 1e-12 ||
        PetscAbsScalar(gradient[0] - 4) > 1e-12)
      SETERRQ(PETSC_COMM_WORLD, MOFEM_DATA_INCONSISTENCY,
              "Algorithm 1 did not accept the final finite trial at its "
              "evaluation limit with its actual merit and gradient");
  }
  CHKERR PetscPrintf(PETSC_COMM_WORLD,
                     "Algorithm 1 budget fallback %d: x=%g, lambda=%g\n",
                     static_cast<int>(exhaust_budget),
                     static_cast<double>(PetscRealPart(solution[0])),
                     static_cast<double>(PetscRealPart(dual[0])));
  MoFEMFunctionReturn(0);
}
 
} // namespace
 
int main(int argc, char *argv[]) {
  MoFEM::Core::Initialize(&argc, &argv, nullptr, help);
  try {
    CHKERR TaoRegisterNullSpace();
    PetscBool activation = PETSC_FALSE, domain = PETSC_FALSE,
              domain_budget = PETSC_FALSE, flow = PETSC_FALSE,
              flow_budget = PETSC_FALSE;
    CHKERR PetscOptionsGetBool(nullptr, nullptr, "-test_nullspace_activation",
                               &activation, nullptr);
    CHKERR PetscOptionsGetBool(nullptr, nullptr, "-test_nullspace_domain",
                               &domain, nullptr);
    CHKERR PetscOptionsGetBool(nullptr, nullptr,
                               "-test_nullspace_domain_budget", &domain_budget,
                               nullptr);
    CHKERR PetscOptionsGetBool(nullptr, nullptr, "-test_nullspace_flow", &flow,
                               nullptr);
    CHKERR PetscOptionsGetBool(nullptr, nullptr, "-test_nullspace_flow_budget",
                               &flow_budget, nullptr);
    if (activation) {
      CHKERR checkCandidateActivation();
    } else if (domain_budget) {
      CHKERR checkObjectiveDomainLineSearch(true);
    } else if (domain) {
      CHKERR checkObjectiveDomainLineSearch(false);
    } else if (flow || flow_budget) {
      CHKERR checkAlgorithmOneStep(flow_budget);
    } else {
      const std::array<std::array<PetscScalar, 3>, 3> starts = {
          {{1, 1, 3}, {3, 4, 2}, {1, 1, 6}}};
      for (const auto scale : {1.0, 0.1, 10.0})
        for (size_t i = 0; i != starts.size(); ++i)
          CHKERR solveProblem(starts[i], scale, i != 1);
    }
  }
  CATCH_ERRORS;
  CHKERR MoFEM::Core::Finalize();
  return 0;
}