algebra_solver_fp.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_REV_FUNCTOR_FP_SOLVER_HPP
#define STAN_MATH_REV_FUNCTOR_FP_SOLVER_HPP
 
#include <stan/math/rev/meta.hpp>
#include <stan/math/rev/core.hpp>
#include <stan/math/rev/fun/value_of.hpp>
#include <stan/math/rev/functor/algebra_system.hpp>
#include <stan/math/rev/functor/kinsol_data.hpp>
#include <stan/math/prim/err.hpp>
#include <stan/math/prim/fun/mdivide_left.hpp>
#include <stan/math/prim/fun/to_array_1d.hpp>
#include <stan/math/prim/fun/to_vector.hpp>
#include <stan/math/prim/fun/value_of.hpp>
#include <sundials/sundials_context.h>
#include <kinsol/kinsol.h>
#include <sunmatrix/sunmatrix_dense.h>
#include <sunlinsol/sunlinsol_dense.h>
#include <nvector/nvector_serial.h>
#include <algorithm>
#include <iostream>
#include <string>
#include <vector>
 
namespace stan {
namespace math {
 
template <typename F>
struct KinsolFixedPointEnv {
  sundials::Context sundials_context_;
  const F& f_;
  const Eigen::VectorXd y_dummy;
  const Eigen::VectorXd& y_;
  const size_t N_;
  const size_t M_;
  const std::vector<double>& x_r_;
  const std::vector<int>& x_i_;
  std::ostream* msgs_;
  void* mem_;
  N_Vector nv_x_;
  N_Vector nv_u_scal_;
  N_Vector nv_f_scal_;
 
  template <typename T, typename T_u, typename T_f>
  KinsolFixedPointEnv(const F& f, const Eigen::Matrix<T, -1, 1>& x,
                      const Eigen::VectorXd& y, const std::vector<double>& x_r,
                      const std::vector<int>& x_i, std::ostream* msgs,
                      const std::vector<T_u>& u_scale,
                      const std::vector<T_f>& f_scale)
      : sundials_context_(),
        f_(f),
        y_dummy(),
        y_(y),
        N_(x.size()),
        M_(y.size()),
        x_r_(x_r),
        x_i_(x_i),
        msgs_(msgs),
        mem_(KINCreate(sundials_context_)),
        nv_x_(N_VNew_Serial(N_, sundials_context_)),
        nv_u_scal_(N_VNew_Serial(N_, sundials_context_)),
        nv_f_scal_(N_VNew_Serial(N_, sundials_context_)) {
    for (int i = 0; i < N_; ++i) {
      NV_Ith_S(nv_x_, i) = stan::math::value_of(x(i));
      NV_Ith_S(nv_u_scal_, i) = stan::math::value_of(u_scale[i]);
      NV_Ith_S(nv_f_scal_, i) = stan::math::value_of(f_scale[i]);
    }
  }
 
  template <typename T, typename T_u, typename T_f>
  KinsolFixedPointEnv(const F& f, const Eigen::Matrix<T, -1, 1>& x,
                      const Eigen::Matrix<stan::math::var, -1, 1>& y,
                      const std::vector<double>& x_r,
                      const std::vector<int>& x_i, std::ostream* msgs,
                      const std::vector<T_u>& u_scale,
                      const std::vector<T_f>& f_scale)
      : sundials_context_(),
        f_(f),
        y_dummy(stan::math::value_of(y)),
        y_(y_dummy),
        N_(x.size()),
        M_(y.size()),
        x_r_(x_r),
        x_i_(x_i),
        msgs_(msgs),
        mem_(KINCreate(sundials_context_)),
        nv_x_(N_VNew_Serial(N_, sundials_context_)),
        nv_u_scal_(N_VNew_Serial(N_, sundials_context_)),
        nv_f_scal_(N_VNew_Serial(N_, sundials_context_)) {
    for (int i = 0; i < N_; ++i) {
      NV_Ith_S(nv_x_, i) = stan::math::value_of(x(i));
      NV_Ith_S(nv_u_scal_, i) = stan::math::value_of(u_scale[i]);
      NV_Ith_S(nv_f_scal_, i) = stan::math::value_of(f_scale[i]);
    }
  }
 
  ~KinsolFixedPointEnv() {
    N_VDestroy_Serial(nv_x_);
    N_VDestroy_Serial(nv_u_scal_);
    N_VDestroy_Serial(nv_f_scal_);
    KINFree(&mem_);
  }
 
  static int kinsol_f_system(N_Vector x, N_Vector f, void* user_data) {
    auto g = static_cast<const KinsolFixedPointEnv<F>*>(user_data);
    Eigen::VectorXd x_eigen(Eigen::Map<Eigen::VectorXd>(NV_DATA_S(x), g->N_));
    Eigen::Map<Eigen::VectorXd>(N_VGetArrayPointer(f), g->N_)
        = g->f_(x_eigen, g->y_, g->x_r_, g->x_i_, g->msgs_);
    return 0;
  }
};
 
struct FixedPointADJac {
  template <typename F>
  inline Eigen::Matrix<stan::math::var, -1, 1> operator()(
      const Eigen::VectorXd& x, const Eigen::Matrix<stan::math::var, -1, 1>& y,
      KinsolFixedPointEnv<F>& env) {
    using stan::math::precomputed_gradients;
    using stan::math::to_array_1d;
    using stan::math::var;
 
    auto g = [&env](const Eigen::Matrix<var, -1, 1>& par_) {
      Eigen::Matrix<var, -1, 1> x_(par_.head(env.N_));
      Eigen::Matrix<var, -1, 1> y_(par_.tail(env.M_));
      return env.f_(x_, y_, env.x_r_, env.x_i_, env.msgs_);
    };
 
    Eigen::VectorXd theta(x.size() + y.size());
    for (int i = 0; i < env.N_; ++i) {
      theta(i) = x(i);
    }
    for (int i = 0; i < env.M_; ++i) {
      theta(i + env.N_) = env.y_(i);
    }
    Eigen::Matrix<double, -1, 1> fx;
    Eigen::Matrix<double, -1, -1> J_theta;
    stan::math::jacobian(g, theta, fx, J_theta);
    Eigen::MatrixXd A(J_theta.block(0, 0, env.N_, env.N_));
    Eigen::MatrixXd b(J_theta.block(0, env.N_, env.N_, env.M_));
    A = Eigen::MatrixXd::Identity(env.N_, env.N_) - A;
    Eigen::MatrixXd Jxy = A.colPivHouseholderQr().solve(b);
    std::vector<double> gradients(env.M_);
    Eigen::Matrix<var, -1, 1> x_sol(env.N_);
    std::vector<stan::math::var> yv(to_array_1d(y));
    for (int i = 0; i < env.N_; ++i) {
      gradients = to_array_1d(Eigen::VectorXd(Jxy.row(i)));
      x_sol[i] = precomputed_gradients(x(i), yv, gradients);
    }
    return x_sol;
  }
};
 
template <typename fp_env_type, typename fp_jac_type>
struct FixedPointSolver;
 
template <typename F, typename fp_jac_type>
struct FixedPointSolver<KinsolFixedPointEnv<F>, fp_jac_type> {
  void kinsol_solve_fp(Eigen::VectorXd& x, KinsolFixedPointEnv<F>& env,
                       double f_tol, int max_num_steps) {
    int N = env.N_;
    void* mem = env.mem_;
 
    const int default_anderson_depth = 4;
    int anderson_depth = std::min(N, default_anderson_depth);
 
    CHECK_KINSOL_CALL(KINSetNumMaxIters(mem, max_num_steps));
    CHECK_KINSOL_CALL(KINSetMAA(mem, anderson_depth));
    CHECK_KINSOL_CALL(KINInit(mem, &env.kinsol_f_system, env.nv_x_));
    CHECK_KINSOL_CALL(KINSetFuncNormTol(mem, f_tol));
    CHECK_KINSOL_CALL(KINSetUserData(mem, static_cast<void*>(&env)));
    kinsol_check(KINSol(mem, env.nv_x_, KIN_FP, env.nv_u_scal_, env.nv_f_scal_),
                 "KINSol", max_num_steps);
 
    for (int i = 0; i < N; ++i) {
      x(i) = NV_Ith_S(env.nv_x_, i);
    }
  }
 
  template <typename T1>
  Eigen::Matrix<double, -1, 1> solve(const Eigen::Matrix<T1, -1, 1>& x,
                                     const Eigen::Matrix<double, -1, 1>& y,
                                     KinsolFixedPointEnv<F>& env, double f_tol,
                                     int max_num_steps) {
    Eigen::VectorXd xd(stan::math::value_of(x));
    kinsol_solve_fp(xd, env, f_tol, max_num_steps);
    return xd;
  }
 
  template <typename T1>
  Eigen::Matrix<stan::math::var, -1, 1> solve(
      const Eigen::Matrix<T1, -1, 1>& x,
      const Eigen::Matrix<stan::math::var, -1, 1>& y,
      KinsolFixedPointEnv<F>& env, double f_tol, int max_num_steps) {
    using stan::math::value_of;
    using stan::math::var;
 
    // FP solution
    Eigen::VectorXd xd(solve(x, Eigen::VectorXd(), env, f_tol, max_num_steps));
 
    fp_jac_type jac_sol;
    return jac_sol(xd, y, env);
  }
};
 
template <typename F, typename T1, typename T2, typename T_u, typename T_f>
Eigen::Matrix<T2, -1, 1> algebra_solver_fp(
    const F& f, const Eigen::Matrix<T1, -1, 1>& x,
    const Eigen::Matrix<T2, -1, 1>& y, const std::vector<double>& x_r,
    const std::vector<int>& x_i, const std::vector<T_u>& u_scale,
    const std::vector<T_f>& f_scale, std::ostream* msgs = nullptr,
    double f_tol = 1e-8,
    int max_num_steps = 200) {  // NOLINT(runtime/int)
  algebra_solver_check(x, y, x_r, x_i, f_tol, max_num_steps);
  check_nonnegative("algebra_solver", "u_scale", u_scale);
  check_nonnegative("algebra_solver", "f_scale", f_scale);
  check_matching_sizes("algebra_solver", "the algebraic system's output",
                       value_of(f(x, y, x_r, x_i, msgs)),
                       "the vector of unknowns, x,", x);
 
  KinsolFixedPointEnv<F> env(f, x, y, x_r, x_i, msgs, u_scale,
                             f_scale);  // NOLINT
  FixedPointSolver<KinsolFixedPointEnv<F>, FixedPointADJac> fp;
  return fp.solve(x, y, env, f_tol, max_num_steps);
}
 
}  // namespace math
}  // namespace stan
 
#endif