solve_newton.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_REV_FUNCTOR_SOLVE_NEWTON_HPP
#define STAN_MATH_REV_FUNCTOR_SOLVE_NEWTON_HPP
 
#include <stan/math/rev/core.hpp>
#include <stan/math/rev/functor/algebra_system.hpp>
#include <stan/math/rev/functor/kinsol_solve.hpp>
#include <stan/math/prim/err.hpp>
#include <stan/math/prim/fun/eval.hpp>
#include <stan/math/prim/fun/value_of.hpp>
#include <stan/math/prim/functor/algebra_solver_adapter.hpp>
#include <stan/math/prim/functor/apply.hpp>
#include <unsupported/Eigen/NonLinearOptimization>
#include <iostream>
#include <string>
#include <vector>
 
namespace stan {
namespace math {
 
template <typename F, typename T, typename... Args,
          require_eigen_vector_t<T>* = nullptr,
          require_all_st_arithmetic<Args...>* = nullptr>
Eigen::VectorXd solve_newton_tol(const F& f, const T& x,
                                 const double scaling_step_size,
                                 const double function_tolerance,
                                 const int64_t max_num_steps,
                                 std::ostream* const msgs,
                                 const Args&... args) {
  const auto& x_ref = to_ref(value_of(x));
 
  check_nonzero_size("solve_newton", "initial guess", x_ref);
  check_finite("solve_newton", "initial guess", x_ref);
  check_nonnegative("solve_newton", "scaling_step_size", scaling_step_size);
  check_nonnegative("solve_newton", "function_tolerance", function_tolerance);
  check_positive("solve_newton", "max_num_steps", max_num_steps);
 
  return kinsol_solve(f, x_ref, scaling_step_size, function_tolerance,
                      max_num_steps, 1, 10, KIN_LINESEARCH, msgs, args...);
}
 
template <typename F, typename T, typename... T_Args,
          require_eigen_vector_t<T>* = nullptr,
          require_any_st_var<T_Args...>* = nullptr>
Eigen::Matrix<var, Eigen::Dynamic, 1> solve_newton_tol(
    const F& f, const T& x, const double scaling_step_size,
    const double function_tolerance, const int64_t max_num_steps,
    std::ostream* const msgs, const T_Args&... args) {
  const auto& x_ref = to_ref(value_of(x));
  auto arena_args_tuple = make_chainable_ptr(std::make_tuple(eval(args)...));
  auto args_vals_tuple = math::apply(
      [&](const auto&... args) {
        return std::make_tuple(to_ref(value_of(args))...);
      },
      *arena_args_tuple);
 
  check_nonzero_size("solve_newton", "initial guess", x_ref);
  check_finite("solve_newton", "initial guess", x_ref);
  check_nonnegative("solve_newton", "scaling_step_size", scaling_step_size);
  check_nonnegative("solve_newton", "function_tolerance", function_tolerance);
  check_positive("solve_newton", "max_num_steps", max_num_steps);
 
  // Solve the system
  Eigen::VectorXd theta_dbl = math::apply(
      [&](const auto&... vals) {
        return kinsol_solve(f, x_ref, scaling_step_size, function_tolerance,
                            max_num_steps, 1, 10, KIN_LINESEARCH, msgs,
                            vals...);
      },
      args_vals_tuple);
 
  auto f_wrt_x = [&](const auto& x) {
    return math::apply([&](const auto&... args) { return f(x, msgs, args...); },
                       args_vals_tuple);
  };
 
  Eigen::MatrixXd Jf_x;
  Eigen::VectorXd f_x;
 
  jacobian(f_wrt_x, theta_dbl, f_x, Jf_x);
 
  using ret_type = Eigen::Matrix<var, Eigen::Dynamic, -1>;
  arena_t<ret_type> ret = theta_dbl;
  auto Jf_x_T_lu_ptr
      = make_unsafe_chainable_ptr(Jf_x.transpose().partialPivLu());  // Lu
 
  reverse_pass_callback(
      [f, ret, arena_args_tuple, Jf_x_T_lu_ptr, msgs]() mutable {
        Eigen::VectorXd eta = -Jf_x_T_lu_ptr->solve(ret.adj().eval());
 
        // Contract with Jacobian of f with respect to y using a nested reverse
        // autodiff pass.
        {
          nested_rev_autodiff rev;
 
          Eigen::VectorXd ret_val = ret.val();
          auto x_nrad_ = math::apply(
              [&ret_val, &f, msgs](const auto&... args) {
                return eval(f(ret_val, msgs, args...));
              },
              *arena_args_tuple);
          x_nrad_.adj() = eta;
          grad();
        }
      });
 
  return ret_type(ret);
}
 
template <typename F, typename T, typename... T_Args,
          require_eigen_vector_t<T>* = nullptr>
Eigen::Matrix<stan::return_type_t<T_Args...>, Eigen::Dynamic, 1> solve_newton(
    const F& f, const T& x, std::ostream* const msgs, const T_Args&... args) {
  double scaling_step_size = 1e-3;
  double function_tolerance = 1e-6;
  int64_t max_num_steps = 200;
  const auto& args_ref_tuple = std::make_tuple(to_ref(args)...);
  return math::apply(
      [&](const auto&... args_refs) {
        return solve_newton_tol(f, x, scaling_step_size, function_tolerance,
                                max_num_steps, msgs, args_refs...);
      },
      args_ref_tuple);
}
 
template <typename F, typename T1, typename T2,
          require_all_eigen_vector_t<T1, T2>* = nullptr>
Eigen::Matrix<scalar_type_t<T2>, Eigen::Dynamic, 1> algebra_solver_newton(
    const F& f, const T1& x, const T2& y, const std::vector<double>& dat,
    const std::vector<int>& dat_int, std::ostream* const msgs = nullptr,
    const double scaling_step_size = 1e-3,
    const double function_tolerance = 1e-6,
    const long int max_num_steps = 200) {  // NOLINT(runtime/int)
  return solve_newton_tol(algebra_solver_adapter<F>(f), x, scaling_step_size,
                          function_tolerance, max_num_steps, msgs, y, dat,
                          dat_int);
}
 
}  // namespace math
}  // namespace stan
 
#endif