Merge pull request #1848 from bstatcomp/generalize_fun_s_t

Generalize functions starting with s and t

Author: t4c1

stan/math/fwd/fun/softmax.hpp

@@ -8,36 +8,39 @@
 namespace stan {
 namespace math {
 
-template <typename T>
-inline Eigen::Matrix<fvar<T>, Eigen::Dynamic, 1> softmax(
-    const Eigen::Matrix<fvar<T>, Eigen::Dynamic, 1>& alpha) {
+template <typename ColVec,
+          require_eigen_col_vector_vt<is_fvar, ColVec>* = nullptr>
+inline auto softmax(const ColVec& alpha) {
   using Eigen::Dynamic;
   using Eigen::Matrix;
+  using T = typename value_type_t<ColVec>::Scalar;
 
   Matrix<T, Dynamic, 1> alpha_t(alpha.size());
   for (int k = 0; k < alpha.size(); ++k) {
-    alpha_t(k) = alpha(k).val_;
+    alpha_t.coeffRef(k) = alpha.coeff(k).val_;
   }
 
   Matrix<T, Dynamic, 1> softmax_alpha_t = softmax(alpha_t);
 
   Matrix<fvar<T>, Dynamic, 1> softmax_alpha(alpha.size());
   for (int k = 0; k < alpha.size(); ++k) {
-    softmax_alpha(k).val_ = softmax_alpha_t(k);
-    softmax_alpha(k).d_ = 0;
+    softmax_alpha.coeffRef(k).val_ = softmax_alpha_t.coeff(k);
+    softmax_alpha.coeffRef(k).d_ = 0;
   }
 
   for (int m = 0; m < alpha.size(); ++m) {
     T negative_alpha_m_d_times_softmax_alpha_t_m
-        = -alpha(m).d_ * softmax_alpha_t(m);
+        = -alpha.coeff(m).d_ * softmax_alpha_t.coeff(m);
     for (int k = 0; k < alpha.size(); ++k) {
       if (m == k) {
-        softmax_alpha(k).d_
-            += softmax_alpha_t(k)
-               * (alpha(m).d_ + negative_alpha_m_d_times_softmax_alpha_t_m);
+        softmax_alpha.coeffRef(k).d_
+            += softmax_alpha_t.coeff(k)
+               * (alpha.coeff(m).d_
+                  + negative_alpha_m_d_times_softmax_alpha_t_m);
       } else {
-        softmax_alpha(k).d_
-            += negative_alpha_m_d_times_softmax_alpha_t_m * softmax_alpha_t(k);
+        softmax_alpha.coeffRef(k).d_
+            += softmax_alpha_t.coeff(k)
+               * negative_alpha_m_d_times_softmax_alpha_t_m;
       }
     }
   }

stan/math/fwd/fun/trace_quad_form.hpp

@@ -11,29 +11,16 @@
 namespace stan {
 namespace math {
 
-template <int RA, int CA, int RB, int CB, typename T>
-inline fvar<T> trace_quad_form(const Eigen::Matrix<fvar<T>, RA, CA> &A,
-                               const Eigen::Matrix<fvar<T>, RB, CB> &B) {
+template <typename EigMat1, typename EigMat2,
+          require_all_eigen_t<EigMat1, EigMat2>* = nullptr,
+          require_any_vt_fvar<EigMat1, EigMat2>* = nullptr>
+inline return_type_t<EigMat1, EigMat2> trace_quad_form(const EigMat1& A,
+                                                       const EigMat2& B) {
   check_square("trace_quad_form", "A", A);
   check_multiplicable("trace_quad_form", "A", A, "B", B);
-  return trace(multiply(transpose(B), multiply(A, B)));
+  return B.cwiseProduct(multiply(A, B)).sum();
 }
 
-template <int RA, int CA, int RB, int CB, typename T>
-inline fvar<T> trace_quad_form(const Eigen::Matrix<fvar<T>, RA, CA> &A,
-                               const Eigen::Matrix<double, RB, CB> &B) {
-  check_square("trace_quad_form", "A", A);
-  check_multiplicable("trace_quad_form", "A", A, "B", B);
-  return trace(multiply(transpose(B), multiply(A, B)));
-}
-
-template <int RA, int CA, int RB, int CB, typename T>
-inline fvar<T> trace_quad_form(const Eigen::Matrix<double, RA, CA> &A,
-                               const Eigen::Matrix<fvar<T>, RB, CB> &B) {
-  check_square("trace_quad_form", "A", A);
-  check_multiplicable("trace_quad_form", "A", A, "B", B);
-  return trace(multiply(transpose(B), multiply(A, B)));
-}
 }  // namespace math
 }  // namespace stan
 

stan/math/prim/fun/scale_matrix_exp_multiply.hpp

@@ -16,16 +16,19 @@ namespace math {
  *
  * Specialized for double values for efficiency.
  *
- * @tparam Cb number of columns in matrix B, can be Eigen::Dynamic
+ * @tparam EigMat1 type of the first matrix
+ * @tparam EigMat2 type of the second matrix
+ *
  * @param[in] A Matrix
  * @param[in] B Matrix
  * @param[in] t double
- * @return exponential of At multiplies B
+ * @return exponential of At multiplied by B
  */
-template <int Cb>
-inline Eigen::Matrix<double, -1, Cb> scale_matrix_exp_multiply(
-    const double& t, const Eigen::MatrixXd& A,
-    const Eigen::Matrix<double, -1, Cb>& B) {
+template <typename EigMat1, typename EigMat2,
+          require_all_eigen_t<EigMat1, EigMat2>* = nullptr,
+          require_all_vt_same<double, EigMat1, EigMat2>* = nullptr>
+inline Eigen::Matrix<double, Eigen::Dynamic, EigMat2::ColsAtCompileTime>
+scale_matrix_exp_multiply(const double& t, const EigMat1& A, const EigMat2& B) {
   check_square("scale_matrix_exp_multiply", "input matrix", A);
   check_multiplicable("scale_matrix_exp_multiply", "A", A, "B", B);
   if (A.size() == 0) {
@@ -39,20 +42,22 @@ inline Eigen::Matrix<double, -1, Cb> scale_matrix_exp_multiply(
  * Return product of exp(At) and B, where A is a NxN matrix,
  * B is a NxCb matrix and t is a scalar.
  *
- * Generic implementation when arguments are not double.
+ * Generic implementation when arguments are not all double.
  *
- * @tparam Ta scalar type matrix A
- * @tparam Tb scalar type matrix B
- * @tparam Cb number of columns in matrix B, can be Eigen::Dynamic
+ * @tparam Tt type of \c t
+ * @tparam EigMat1 type of the first matrix
+ * @tparam EigMat2 type of the second matrix
  * @param[in] A Matrix
  * @param[in] B Matrix
  * @param[in] t double
- * @return exponential of At multiplies B
+ * @return exponential of At multiplied by B
  */
-template <typename Tt, typename Ta, typename Tb, int Cb>
-inline Eigen::Matrix<return_type_t<Tt, Ta, Tb>, -1, Cb>
-scale_matrix_exp_multiply(const Tt& t, const Eigen::Matrix<Ta, -1, -1>& A,
-                          const Eigen::Matrix<Tb, -1, Cb>& B) {
+template <typename Tt, typename EigMat1, typename EigMat2,
+          require_all_eigen_t<EigMat1, EigMat2>* = nullptr,
+          require_any_not_vt_same<double, Tt, EigMat1, EigMat2>* = nullptr>
+inline Eigen::Matrix<return_type_t<Tt, EigMat1, EigMat2>, Eigen::Dynamic,
+                     EigMat2::ColsAtCompileTime>
+scale_matrix_exp_multiply(const Tt& t, const EigMat1& A, const EigMat2& B) {
   check_square("scale_matrix_exp_multiply", "input matrix", A);
   check_multiplicable("scale_matrix_exp_multiply", "A", A, "B", B);
   if (A.size() == 0) {

stan/math/prim/fun/sd.hpp

@@ -13,34 +13,17 @@ namespace math {
 
 /**
  * Returns the unbiased sample standard deviation of the
- * coefficients in the specified column vector.
+ * coefficients in the specified std vector, column vector, row vector, or
+ * matrix.
  *
- * @tparam T type of elements in the vector
- * @param v Specified vector.
- * @return Sample variance of vector.
- */
-template <typename T>
-inline return_type_t<T> sd(const std::vector<T>& v) {
-  check_nonzero_size("sd", "v", v);
-  if (v.size() == 1) {
-    return 0.0;
-  }
-  return sqrt(variance(v));
-}
-
-/**
- * Returns the unbiased sample standard deviation of the
- * coefficients in the specified vector, row vector, or matrix.
- *
- * @tparam T type of elements in the vector, row vector, or matrix
- * @tparam R number of rows, can be Eigen::Dynamic
- * @tparam C number of columns, can be Eigen::Dynamic
+ * @tparam T type of the container
  *
- * @param m Specified vector, row vector or matrix.
+ * @param m Specified container.
  * @return Sample variance.
  */
-template <typename T, int R, int C>
-inline return_type_t<T> sd(const Eigen::Matrix<T, R, C>& m) {
+template <typename T, require_container_t<T>* = nullptr,
+          require_not_vt_var<T>* = nullptr>
+inline return_type_t<T> sd(const T& m) {
   using std::sqrt;
   check_nonzero_size("sd", "m", m);
   if (m.size() == 1) {

stan/math/prim/fun/segment.hpp

@@ -9,38 +9,23 @@ namespace stan {
 namespace math {
 
 /**
- * Return the specified number of elements as a vector starting
- * from the specified element - 1 of the specified vector.
+ * Return the specified number of elements as a row/column vector starting
+ * from the specified element - 1 of the specified row/column vector.
  *
- * @tparam T type of elements in the vector
+ * @tparam T type of the vector
  */
-template <typename T>
-inline Eigen::Matrix<T, Eigen::Dynamic, 1> segment(
-    const Eigen::Matrix<T, Eigen::Dynamic, 1>& v, size_t i, size_t n) {
+template <typename EigVec, require_eigen_vector_t<EigVec>* = nullptr>
+inline plain_type_t<EigVec> segment(const EigVec& v, size_t i, size_t n) {
   check_greater("segment", "n", i, 0.0);
-  check_less_or_equal("segment", "n", i, static_cast<size_t>(v.rows()));
+  check_less_or_equal("segment", "n", i, static_cast<size_t>(v.size()));
   if (n != 0) {
     check_greater("segment", "n", i + n - 1, 0.0);
     check_less_or_equal("segment", "n", i + n - 1,
-                        static_cast<size_t>(v.rows()));
+                        static_cast<size_t>(v.size()));
   }
   return v.segment(i - 1, n);
 }
 
-template <typename T>
-inline Eigen::Matrix<T, 1, Eigen::Dynamic> segment(
-    const Eigen::Matrix<T, 1, Eigen::Dynamic>& v, size_t i, size_t n) {
-  check_greater("segment", "n", i, 0.0);
-  check_less_or_equal("segment", "n", i, static_cast<size_t>(v.cols()));
-  if (n != 0) {
-    check_greater("segment", "n", i + n - 1, 0.0);
-    check_less_or_equal("segment", "n", i + n - 1,
-                        static_cast<size_t>(v.cols()));
-  }
-
-  return v.segment(i - 1, n);
-}
-
 template <typename T>
 std::vector<T> segment(const std::vector<T>& sv, size_t i, size_t n) {
   check_greater("segment", "i", i, 0.0);

stan/math/prim/fun/simplex_constrain.hpp

@@ -20,28 +20,27 @@ namespace math {
  *
  * The transform is based on a centered stick-breaking process.
  *
- * @tparam T type of elements in the vector
+ * @tparam ColVec type of the vector
  * @param y Free vector input of dimensionality K - 1.
  * @return Simplex of dimensionality K.
  */
-template <typename T>
-Eigen::Matrix<T, Eigen::Dynamic, 1> simplex_constrain(
-    const Eigen::Matrix<T, Eigen::Dynamic, 1>& y) {
+template <typename ColVec, require_eigen_col_vector_t<ColVec>* = nullptr>
+auto simplex_constrain(const ColVec& y) {
   // cut & paste simplex_constrain(Eigen::Matrix, T) w/o Jacobian
   using Eigen::Dynamic;
   using Eigen::Matrix;
   using std::log;
-  using size_type = index_type_t<Matrix<T, Dynamic, 1>>;
+  using T = value_type_t<ColVec>;
 
   int Km1 = y.size();
   Matrix<T, Dynamic, 1> x(Km1 + 1);
   T stick_len(1.0);
-  for (size_type k = 0; k < Km1; ++k) {
-    T z_k(inv_logit(y(k) - log(Km1 - k)));
-    x(k) = stick_len * z_k;
-    stick_len -= x(k);
+  for (Eigen::Index k = 0; k < Km1; ++k) {
+    T z_k = inv_logit(y.coeff(k) - log(Km1 - k));
+    x.coeffRef(k) = stick_len * z_k;
+    stick_len -= x.coeff(k);
   }
-  x(Km1) = stick_len;
+  x.coeffRef(Km1) = stick_len;
   return x;
 }
 
@@ -53,34 +52,32 @@ Eigen::Matrix<T, Eigen::Dynamic, 1> simplex_constrain(
  * The simplex transform is defined through a centered
  * stick-breaking process.
  *
- * @tparam T type of elements in the vector
+ * @tparam ColVec type of the vector
  * @param y Free vector input of dimensionality K - 1.
  * @param lp Log probability reference to increment.
  * @return Simplex of dimensionality K.
  */
-template <typename T>
-Eigen::Matrix<T, Eigen::Dynamic, 1> simplex_constrain(
-    const Eigen::Matrix<T, Eigen::Dynamic, 1>& y, T& lp) {
+template <typename ColVec, require_eigen_col_vector_t<ColVec>* = nullptr>
+auto simplex_constrain(const ColVec& y, value_type_t<ColVec>& lp) {
   using Eigen::Dynamic;
   using Eigen::Matrix;
   using std::log;
-
-  using size_type = index_type_t<Matrix<T, Dynamic, 1>>;
+  using T = value_type_t<ColVec>;
 
   int Km1 = y.size();  // K = Km1 + 1
   Matrix<T, Dynamic, 1> x(Km1 + 1);
   T stick_len(1.0);
-  for (size_type k = 0; k < Km1; ++k) {
+  for (Eigen::Index k = 0; k < Km1; ++k) {
     double eq_share = -log(Km1 - k);  // = logit(1.0/(Km1 + 1 - k));
-    T adj_y_k(y(k) + eq_share);
-    T z_k(inv_logit(adj_y_k));
-    x(k) = stick_len * z_k;
+    T adj_y_k = y.coeff(k) + eq_share;
+    T z_k = inv_logit(adj_y_k);
+    x.coeffRef(k) = stick_len * z_k;
     lp += log(stick_len);
     lp -= log1p_exp(-adj_y_k);
     lp -= log1p_exp(adj_y_k);
-    stick_len -= x(k);  // equivalently *= (1 - z_k);
+    stick_len -= x.coeff(k);  // equivalently *= (1 - z_k);
   }
-  x(Km1) = stick_len;  // no Jacobian contrib for last dim
+  x.coeffRef(Km1) = stick_len;  // no Jacobian contrib for last dim
   return x;
 }
 

stan/math/prim/fun/simplex_free.hpp

@@ -19,26 +19,25 @@ namespace math {
  * <p>The simplex transform is defined through a centered
  * stick-breaking process.
  *
- * @tparam T type of elements in the simplex
+ * @tparam ColVec type of the simplex (must be a column vector)
  * @param x Simplex of dimensionality K.
  * @return Free vector of dimensionality (K-1) that transforms to
  * the simplex.
  * @throw std::domain_error if x is not a valid simplex
  */
-template <typename T>
-Eigen::Matrix<T, Eigen::Dynamic, 1> simplex_free(
-    const Eigen::Matrix<T, Eigen::Dynamic, 1>& x) {
+template <typename ColVec, require_eigen_col_vector_t<ColVec>* = nullptr>
+auto simplex_free(const ColVec& x) {
   using std::log;
-  using size_type = index_type_t<Eigen::Matrix<T, Eigen::Dynamic, 1>>;
+  using T = value_type_t<ColVec>;
 
   check_simplex("stan::math::simplex_free", "Simplex variable", x);
   int Km1 = x.size() - 1;
   Eigen::Matrix<T, Eigen::Dynamic, 1> y(Km1);
-  T stick_len(x(Km1));
-  for (size_type k = Km1; --k >= 0;) {
-    stick_len += x(k);
-    T z_k(x(k) / stick_len);
-    y(k) = logit(z_k) + log(Km1 - k);
+  T stick_len = x.coeff(Km1);
+  for (Eigen::Index k = Km1; --k >= 0;) {
+    stick_len += x.coeff(k);
+    T z_k = x.coeff(k) / stick_len;
+    y.coeffRef(k) = logit(z_k) + log(Km1 - k);
     // note: log(Km1 - k) = logit(1.0 / (Km1 + 1 - k));
   }
   return y;

stan/math/prim/fun/singular_values.hpp

@@ -1,6 +1,7 @@
 #ifndef STAN_MATH_PRIM_FUN_SINGULAR_VALUES_HPP
 #define STAN_MATH_PRIM_FUN_SINGULAR_VALUES_HPP
 
+#include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
 
 namespace stan {
@@ -12,18 +13,19 @@ namespace math {
  * <p>See the documentation for <code>svd()</code> for
  * information on the singular values.
  *
- * @tparam T type of elements in the matrix
+ * @tparam EigMat type of the matrix
  * @param m Specified matrix.
  * @return Singular values of the matrix.
  */
-template <typename T>
-Eigen::Matrix<T, Eigen::Dynamic, 1> singular_values(
-    const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>& m) {
+template <typename EigMat, require_eigen_matrix_t<EigMat>* = nullptr>
+Eigen::Matrix<value_type_t<EigMat>, Eigen::Dynamic, 1> singular_values(
+    const EigMat& m) {
   if (m.size() == 0) {
     return {};
   }
 
-  return Eigen::JacobiSVD<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic> >(m)
+  return Eigen::JacobiSVD<Eigen::Matrix<value_type_t<EigMat>, Eigen::Dynamic,
+                                        Eigen::Dynamic> >(m)
       .singularValues();
 }