stan-dev · nhuurre · May 5, 2020 · May 5, 2020 · May 12, 2020 · May 12, 2020
diff --git a/stan/math/fwd/meta/operands_and_partials.hpp b/stan/math/fwd/meta/operands_and_partials.hpp
@@ -70,11 +70,11 @@ template <typename Op1, typename Op2, typename Op3, typename Op4, typename Op5,
           typename Dx>
 class operands_and_partials<Op1, Op2, Op3, Op4, Op5, fvar<Dx>> {
  public:
-  internal::ops_partials_edge<Dx, Op1> edge1_;
-  internal::ops_partials_edge<Dx, Op2> edge2_;
-  internal::ops_partials_edge<Dx, Op3> edge3_;
-  internal::ops_partials_edge<Dx, Op4> edge4_;
-  internal::ops_partials_edge<Dx, Op5> edge5_;
+  internal::ops_partials_edge<Dx, std::decay_t<Op1>> edge1_;
+  internal::ops_partials_edge<Dx, std::decay_t<Op2>> edge2_;
+  internal::ops_partials_edge<Dx, std::decay_t<Op3>> edge3_;
+  internal::ops_partials_edge<Dx, std::decay_t<Op4>> edge4_;
+  internal::ops_partials_edge<Dx, std::decay_t<Op5>> edge5_;
   using T_return_type = fvar<Dx>;
   explicit operands_and_partials(const Op1& o1) : edge1_(o1) {}
   operands_and_partials(const Op1& o1, const Op2& o2)

diff --git a/stan/math/prim/err.hpp b/stan/math/prim/err.hpp
@@ -48,6 +48,7 @@
 #include <stan/math/prim/err/domain_error.hpp>
 #include <stan/math/prim/err/domain_error_vec.hpp>
 #include <stan/math/prim/err/elementwise_check.hpp>
+#include <stan/math/prim/err/hmm_check.hpp>
 #include <stan/math/prim/err/invalid_argument.hpp>
 #include <stan/math/prim/err/invalid_argument_vec.hpp>
 #include <stan/math/prim/err/is_cholesky_factor.hpp>

diff --git a/stan/math/prim/err/hmm_check.hpp b/stan/math/prim/err/hmm_check.hpp
@@ -0,0 +1,47 @@
+#ifndef STAN_MATH_PRIM_ERR_HMM_CHECK_HPP
+#define STAN_MATH_PRIM_ERR_HMM_CHECK_HPP
+
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/row.hpp>
+
+namespace stan {
+namespace math {
+
+/**
+ * Check arguments for hidden Markov model functions with a discrete
+ * latent state (lpdf, rng for latent states, and marginal probabilities
+ * for latent sates).
+ *
+ * @tparam T_omega type of the log likelihood matrix
+ * @tparam T_Gamma type of the transition matrix
+ * @tparam T_rho type of the initial guess vector
+ * @param[in] log_omegas log matrix of observational densities.
+ * @param[in] Gamma transition density between hidden states.
+ * @param[in] rho initial state
+ * @param[in] function the name of the function using the arguments.
+ * @throw `std::invalid_argument` if Gamma is not square
+ *         or if the size of rho is not the number of rows of log_omegas.
+ * @throw `std::domain_error` if rho is not a simplex or the rows
+ *         of Gamma are not a simplex.
+ */
+template <typename T_omega, typename T_Gamma, typename T_rho>
+inline void hmm_check(
+    const Eigen::Matrix<T_omega, Eigen::Dynamic, Eigen::Dynamic>& log_omegas,
+    const Eigen::Matrix<T_Gamma, Eigen::Dynamic, Eigen::Dynamic>& Gamma,
+    const Eigen::Matrix<T_rho, Eigen::Dynamic, 1>& rho, const char* function) {
+  int n_states = log_omegas.rows();
+  int n_transitions = log_omegas.cols() - 1;
+
+  check_consistent_size(function, "rho", rho, n_states);
+  check_simplex(function, "rho", rho);
+  check_square(function, "Gamma", Gamma);
+  check_nonzero_size(function, "Gamma", Gamma);
+  check_multiplicable(function, "Gamma", Gamma, "log_omegas", log_omegas);
+  for (int i = 0; i < Gamma.rows(); ++i) {
+    check_simplex(function, "Gamma[i, ]", row(Gamma, i + 1));
+  }
+}
+
+}  // namespace math
+}  // namespace stan
+#endif
diff --git a/stan/math/prim/functor/coupled_ode_system.hpp b/stan/math/prim/functor/coupled_ode_system.hpp
@@ -127,14 +127,16 @@ struct coupled_ode_system_impl<true, F, T_initial, Args...> {
 template <typename F, typename T_initial, typename... Args>
 struct coupled_ode_system
     : public coupled_ode_system_impl<
-          std::is_arithmetic<return_type_t<T_initial, Args...>>::value, F,
-          T_initial, Args...> {
+          std::is_arithmetic<return_type_t<typename F::captured_scalar_t__,
+                                           T_initial, Args...>>::value,
+          F, T_initial, Args...> {
   coupled_ode_system(const F& f,
                      const Eigen::Matrix<T_initial, Eigen::Dynamic, 1>& y0,
                      std::ostream* msgs, const Args&... args)
       : coupled_ode_system_impl<
-            std::is_arithmetic<return_type_t<T_initial, Args...>>::value, F,
-            T_initial, Args...>(f, y0, msgs, args...) {}
+            std::is_arithmetic<return_type_t<typename F::captured_scalar_t__,
+                                             T_initial, Args...>>::value,
+            F, T_initial, Args...>(f, y0, msgs, args...) {}
 };
 
 }  // namespace math

diff --git a/stan/math/prim/functor/integrate_ode_rk45.hpp b/stan/math/prim/functor/integrate_ode_rk45.hpp
@@ -15,7 +15,8 @@ namespace math {
  */
 template <typename F, typename T1, typename T_param, typename T_t0,
           typename T_ts>
-std::vector<std::vector<return_type_t<T1, T_t0, T_ts, T_param>>>
+std::vector<std::vector<
+    return_type_t<typename F::captured_scalar_t__, T1, T_t0, T_ts, T_param>>>
 integrate_ode_rk45(const F& f, const std::vector<T1>& y0, const T_t0& t0,
                    const std::vector<T_ts>& ts,
                    const std::vector<T_param>& theta,
@@ -28,7 +29,9 @@ integrate_ode_rk45(const F& f, const std::vector<T1>& y0, const T_t0& t0,
       = ode_rk45_tol(f_adapted, to_vector(y0), t0, ts, relative_tolerance,
                      absolute_tolerance, max_num_steps, msgs, theta, x, x_int);
 
-  std::vector<std::vector<return_type_t<T1, T_param, T_t0, T_ts>>> y_converted;
+  std::vector<std::vector<
+      return_type_t<typename F::captured_scalar_t__, T1, T_param, T_t0, T_ts>>>
+      y_converted;
   for (size_t i = 0; i < y.size(); ++i)
     y_converted.push_back(to_array_1d(y[i]));
 

diff --git a/stan/math/prim/functor/integrate_ode_std_vector_interface_adapter.hpp b/stan/math/prim/functor/integrate_ode_std_vector_interface_adapter.hpp
@@ -22,8 +22,10 @@ namespace internal {
 template <typename F>
 struct integrate_ode_std_vector_interface_adapter {
   const F f_;
+  const int num_vars__;
 
-  integrate_ode_std_vector_interface_adapter(const F& f) : f_(f) {}
+  integrate_ode_std_vector_interface_adapter(const F& f)
+      : f_(f), num_vars__(f.num_vars__) {}
 
   template <typename T0, typename T1, typename T2>
   auto operator()(const T0& t, const Eigen::Matrix<T1, Eigen::Dynamic, 1>& y,
@@ -32,6 +34,62 @@ struct integrate_ode_std_vector_interface_adapter {
                   const std::vector<int>& x_int) const {
     return to_vector(f_(t, to_array_1d(y), msgs, theta, x, x_int));
   }
+  template <typename v>
+  void save_varis(v** p) const {
+    f_.save_varis(p);
+  }
+  void accumulate_adjoints(double* p) const { f_.accumulate_adjoints(p); }
+  void set_zer_adjoints() const { f_.set_zero_adjoints(); }
+
+  struct DeepCopy_cl__ {
+    const typename F::DeepCopy__ f_;
+    const int num_vars__;
+    struct ValueOf_cl__ {
+      const typename F::ValueOf__ f_;
+      const static int num_vars__ = 0;
+      ValueOf_cl__(const integrate_ode_std_vector_interface_adapter<F>& f)
+          : f_(f.f_) {}
+      ValueOf_cl__(const DeepCopy_cl__& f) : f_(f.f_) {}
+      template <typename T0, typename T1, typename T2>
+      auto operator()(const T0& t,
+                      const Eigen::Matrix<T1, Eigen::Dynamic, 1>& y,
+                      std::ostream* msgs, const std::vector<T2>& theta,
+                      const std::vector<double>& x,
+                      const std::vector<int>& x_int) const {
+        return to_vector(f_(t, to_array_1d(y), msgs, theta, x, x_int));
+      }
+      template <typename v>
+      void save_varis(v** p) const {
+        f_.save_varis(p);
+      }
+      void accumulate_adjoints(double* p) const { f_.accumulate_adjoints(p); }
+      void set_zer_adjoints() const { f_.set_zero_adjoints(); }
+      using captured_scalar_t__ = double;
+      using ValueOf__ = ValueOf_cl__;
+      using DeepCopy__ = ValueOf_cl__;
+    };
+    DeepCopy_cl__(const integrate_ode_std_vector_interface_adapter<F>& f)
+        : f_(f.f_), num_vars__(f.num_vars__) {}
+    template <typename T0, typename T1, typename T2>
+    auto operator()(const T0& t, const Eigen::Matrix<T1, Eigen::Dynamic, 1>& y,
+                    std::ostream* msgs, const std::vector<T2>& theta,
+                    const std::vector<double>& x,
+                    const std::vector<int>& x_int) const {
+      return to_vector(f_(t, to_array_1d(y), msgs, theta, x, x_int));
+    }
+    template <typename v>
+    void save_varis(v** p) const {
+      f_.save_varis(p);
+    }
+    void accumulate_adjoints(double* p) const { f_.accumulate_adjoints(p); }
+    void set_zer_adjoints() const { f_.set_zero_adjoints(); }
+    using captured_scalar_t__ = typename F::captured_scalar_t__;
+    using ValueOf__ = ValueOf_cl__;
+    using DeepCopy__ = DeepCopy_cl__;
+  };
+  using captured_scalar_t__ = typename F::captured_scalar_t__;
+  using DeepCopy__ = DeepCopy_cl__;
+  using ValueOf__ = typename DeepCopy__::ValueOf__;
 };
 
 }  // namespace internal

diff --git a/stan/math/prim/functor/ode_rk45.hpp b/stan/math/prim/functor/ode_rk45.hpp
@@ -51,7 +51,8 @@ namespace math {
  */
 template <typename F, typename T_initial, typename T_t0, typename T_ts,
           typename... Args>
-std::vector<Eigen::Matrix<stan::return_type_t<T_initial, T_t0, T_ts, Args...>,
+std::vector<Eigen::Matrix<stan::return_type_t<typename F::captured_scalar_t__,
+                                              T_initial, T_t0, T_ts, Args...>,
                           Eigen::Dynamic, 1>>
 ode_rk45_tol(const F& f,
              const Eigen::Matrix<T_initial, Eigen::Dynamic, 1>& y0_arg, T_t0 t0,
@@ -89,7 +90,8 @@ ode_rk45_tol(const F& f,
                         absolute_tolerance);
   check_positive("integrate_ode_rk45", "max_num_steps", max_num_steps);
 
-  using return_t = return_type_t<T_initial, T_t0, T_ts, Args...>;
+  using return_t = return_type_t<typename F::captured_scalar_t__, T_initial,
+                                 T_t0, T_ts, Args...>;
   // creates basic or coupled system by template specializations
   coupled_ode_system<F, T_initial_or_t0, Args...> coupled_system(f, y0, msgs,
                                                                  args...);
@@ -172,8 +174,9 @@ ode_rk45_tol(const F& f,
  */
 template <typename F, typename T_initial, typename T_t0, typename T_ts,
           typename... Args>
-std::vector<
-    Eigen::Matrix<stan::return_type_t<T_initial, Args...>, Eigen::Dynamic, 1>>
+std::vector<Eigen::Matrix<
+    stan::return_type_t<typename F::captured_scalar_t__, T_initial, Args...>,
+    Eigen::Dynamic, 1>>
 ode_rk45(const F& f, const Eigen::Matrix<T_initial, Eigen::Dynamic, 1>& y0,
          T_t0 t0, const std::vector<T_ts>& ts, std::ostream* msgs,
          const Args&... args) {

diff --git a/stan/math/prim/functor/ode_store_sensitivities.hpp b/stan/math/prim/functor/ode_store_sensitivities.hpp
@@ -23,10 +23,11 @@ namespace math {
  * @param args Extra arguments passed unmodified through to ODE right hand side
  * @return ODE state
  */
-template <
-    typename F, typename T_y0_t0, typename T_t0, typename T_t, typename... Args,
-    typename
-    = require_all_arithmetic_t<T_y0_t0, T_t0, T_t, scalar_type_t<Args>...>>
+template <typename F, typename T_y0_t0, typename T_t0, typename T_t,
+          typename... Args,
+          typename
+          = require_all_arithmetic_t<typename F::captured_scalar_t__, T_y0_t0,
+                                     T_t0, T_t, scalar_type_t<Args>...>>
 Eigen::VectorXd ode_store_sensitivities(
     const F& f, const Eigen::VectorXd& coupled_state,
     const Eigen::Matrix<T_y0_t0, Eigen::Dynamic, 1>& y0, T_t0 t0, T_t t,

diff --git a/stan/math/prim/meta/as_array_or_scalar.hpp b/stan/math/prim/meta/as_array_or_scalar.hpp
@@ -17,8 +17,8 @@ namespace math {
  * @return Same value.
  */
 template <typename T, typename = require_stan_scalar_t<T>>
-inline const T& as_array_or_scalar(const T& v) {
-  return v;
+inline T as_array_or_scalar(T&& v) {
+  return std::forward<T>(v);
 }
 
 /** \ingroup type_trait

diff --git a/stan/math/prim/meta/broadcast_array.hpp b/stan/math/prim/meta/broadcast_array.hpp
@@ -3,6 +3,7 @@
 
 #include <stan/math/prim/meta/is_eigen.hpp>
 #include <stan/math/prim/meta/promote_scalar_type.hpp>
+#include <stan/math/prim/meta/ref_type.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stdexcept>
 
@@ -28,7 +29,8 @@ class broadcast_array {
    */
   template <typename Y>
   void operator=(const Y& m) {
-    prim_ = m[0];
+    ref_type_t<Y> m_ref = m;
+    prim_ = m_ref[0];
   }
 };
 

diff --git a/stan/math/prim/prob.hpp b/stan/math/prim/prob.hpp
@@ -134,7 +134,9 @@
 #include <stan/math/prim/prob/gumbel_log.hpp>
 #include <stan/math/prim/prob/gumbel_lpdf.hpp>
 #include <stan/math/prim/prob/gumbel_rng.hpp>
-#include <stan/math/prim/prob/hmm_marginal_lpdf.hpp>
+#include <stan/math/prim/prob/hmm_hidden_state_prob.hpp>
+#include <stan/math/prim/prob/hmm_latent_rng.hpp>
+#include <stan/math/prim/prob/hmm_marginal.hpp>
 #include <stan/math/prim/prob/hypergeometric_log.hpp>
 #include <stan/math/prim/prob/hypergeometric_lpmf.hpp>
 #include <stan/math/prim/prob/hypergeometric_rng.hpp>

diff --git a/stan/math/prim/prob/bernoulli_cdf.hpp b/stan/math/prim/prob/bernoulli_cdf.hpp
@@ -27,21 +27,22 @@ namespace math {
 template <typename T_n, typename T_prob>
 return_type_t<T_prob> bernoulli_cdf(const T_n& n, const T_prob& theta) {
   using T_partials_return = partials_return_t<T_n, T_prob>;
+  using T_theta_ref = ref_type_t<T_prob>;
   static const char* function = "bernoulli_cdf";
-  check_finite(function, "Probability parameter", theta);
-  check_bounded(function, "Probability parameter", theta, 0.0, 1.0);
   check_consistent_sizes(function, "Random variable", n,
                          "Probability parameter", theta);
+  T_theta_ref theta_ref = theta;
+  check_bounded(function, "Probability parameter", theta_ref, 0.0, 1.0);
 
   if (size_zero(n, theta)) {
     return 1.0;
   }
 
   T_partials_return P(1.0);
-  operands_and_partials<T_prob> ops_partials(theta);
+  operands_and_partials<T_theta_ref> ops_partials(theta_ref);
 
   scalar_seq_view<T_n> n_vec(n);
-  scalar_seq_view<T_prob> theta_vec(theta);
+  scalar_seq_view<T_theta_ref> theta_vec(theta_ref);
   size_t max_size_seq_view = max_size(n, theta);
 
   // Explicit return for extreme values

diff --git a/stan/math/prim/prob/bernoulli_lccdf.hpp b/stan/math/prim/prob/bernoulli_lccdf.hpp
@@ -30,22 +30,23 @@ namespace math {
 template <typename T_n, typename T_prob>
 return_type_t<T_prob> bernoulli_lccdf(const T_n& n, const T_prob& theta) {
   using T_partials_return = partials_return_t<T_n, T_prob>;
+  using T_theta_ref = ref_type_t<T_prob>;
   using std::log;
   static const char* function = "bernoulli_lccdf";
-  check_finite(function, "Probability parameter", theta);
-  check_bounded(function, "Probability parameter", theta, 0.0, 1.0);
   check_consistent_sizes(function, "Random variable", n,
                          "Probability parameter", theta);
+  T_theta_ref theta_ref = theta;
+  check_bounded(function, "Probability parameter", theta_ref, 0.0, 1.0);
 
   if (size_zero(n, theta)) {
     return 0.0;
   }
 
   T_partials_return P(0.0);
-  operands_and_partials<T_prob> ops_partials(theta);
+  operands_and_partials<T_theta_ref> ops_partials(theta_ref);
 
   scalar_seq_view<T_n> n_vec(n);
-  scalar_seq_view<T_prob> theta_vec(theta);
+  scalar_seq_view<T_theta_ref> theta_vec(theta_ref);
   size_t max_size_seq_view = max_size(n, theta);
 
   // Explicit return for extreme values

diff --git a/stan/math/prim/prob/bernoulli_lcdf.hpp b/stan/math/prim/prob/bernoulli_lcdf.hpp
@@ -30,22 +30,23 @@ namespace math {
 template <typename T_n, typename T_prob>
 return_type_t<T_prob> bernoulli_lcdf(const T_n& n, const T_prob& theta) {
   using T_partials_return = partials_return_t<T_n, T_prob>;
+  using T_theta_ref = ref_type_t<T_prob>;
   using std::log;
   static const char* function = "bernoulli_lcdf";
-  check_finite(function, "Probability parameter", theta);
-  check_bounded(function, "Probability parameter", theta, 0.0, 1.0);
   check_consistent_sizes(function, "Random variable", n,
                          "Probability parameter", theta);
+  T_theta_ref theta_ref = theta;
+  check_bounded(function, "Probability parameter", theta_ref, 0.0, 1.0);
 
   if (size_zero(n, theta)) {
     return 0.0;
   }
 
   T_partials_return P(0.0);
-  operands_and_partials<T_prob> ops_partials(theta);
+  operands_and_partials<T_theta_ref> ops_partials(theta_ref);
 
   scalar_seq_view<T_n> n_vec(n);
-  scalar_seq_view<T_prob> theta_vec(theta);
+  scalar_seq_view<T_theta_ref> theta_vec(theta_ref);
   size_t max_size_seq_view = max_size(n, theta);
 
   // Explicit return for extreme values