Merge pull request #3151 from stan-dev/broaden-constraint-lp-template

Broaden the template bound on the lp argument to constraints

Author: WardBrian

stan/math/prim/constraint/cholesky_corr_constrain.hpp

@@ -43,9 +43,11 @@ cholesky_corr_constrain(const EigVec& y, int K) {
 }
 
 // FIXME to match above after debugged
-template <typename EigVec, require_eigen_vector_t<EigVec>* = nullptr>
+template <typename EigVec, typename Lp,
+          require_eigen_vector_t<EigVec>* = nullptr,
+          require_convertible_t<return_type_t<EigVec>, Lp>* = nullptr>
 inline Eigen::Matrix<value_type_t<EigVec>, Eigen::Dynamic, Eigen::Dynamic>
-cholesky_corr_constrain(const EigVec& y, int K, return_type_t<EigVec>& lp) {
+cholesky_corr_constrain(const EigVec& y, int K, Lp& lp) {
   using Eigen::Dynamic;
   using Eigen::Matrix;
   using std::sqrt;
@@ -75,28 +77,39 @@ cholesky_corr_constrain(const EigVec& y, int K, return_type_t<EigVec>& lp) {
 }
 
 /**
- * Return The cholesky of a `KxK` correlation matrix. If the `Jacobian`
- * parameter is `true`, the log density accumulator is incremented with the log
- * absolute Jacobian determinant of the transform.  All of the transforms are
- * specified with their Jacobians in the *Stan Reference Manual* chapter
- * Constraint Transforms.
- * @tparam Jacobian if `true`, increment log density accumulator with log
- * absolute Jacobian determinant of constraining transform
- * @tparam T A type inheriting from `Eigen::DenseBase` or a `var_value` with
- *  inner type inheriting from `Eigen::DenseBase` with compile time dynamic rows
- *  and 1 column
+ * Return The cholesky of a `KxK` correlation matrix.
+ * This overload handles looping over the elements of a standard vector.
+ * @tparam T A standard vector with inner type inheriting from
+ * `Eigen::DenseBase` or a `var_value` with inner type inheriting from
+ * `Eigen::DenseBase` with compile time dynamic rows and 1 column
+ * @param y Linearly Serialized vector of size `(K * (K - 1))/2` holding the
+ *  column major order elements of the lower triangurlar
+ * @param K The size of the matrix to return
+ */
+template <typename T, require_std_vector_t<T>* = nullptr>
+inline auto cholesky_corr_constrain(const T& y, int K) {
+  return apply_vector_unary<T>::apply(
+      y, [K](auto&& v) { return cholesky_corr_constrain(v, K); });
+}
+
+/**
+ * Return The cholesky of a `KxK` correlation matrix.
+ * This overload handles looping over the elements of a standard vector.
+ * @tparam T A standard vector with inner type inheriting from
+ * `Eigen::DenseBase` or a `var_value` with inner type inheriting from
+ * `Eigen::DenseBase` with compile time dynamic rows and 1 column
+ * @tparam Lp Scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param y Linearly Serialized vector of size `(K * (K - 1))/2` holding the
  *  column major order elements of the lower triangurlar
  * @param K The size of the matrix to return
  * @param[in,out] lp log density accumulator
  */
-template <bool Jacobian, typename T, require_not_std_vector_t<T>* = nullptr>
-inline auto cholesky_corr_constrain(const T& y, int K, return_type_t<T>& lp) {
-  if (Jacobian) {
-    return cholesky_corr_constrain(y, K, lp);
-  } else {
-    return cholesky_corr_constrain(y, K);
-  }
+template <typename T, typename Lp, require_std_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
+inline auto cholesky_corr_constrain(const T& y, int K, Lp& lp) {
+  return apply_vector_unary<T>::apply(
+      y, [&lp, K](auto&& v) { return cholesky_corr_constrain(v, K, lp); });
 }
 
 /**
@@ -107,19 +120,24 @@ inline auto cholesky_corr_constrain(const T& y, int K, return_type_t<T>& lp) {
  * Constraint Transforms.
  * @tparam Jacobian if `true`, increment log density accumulator with log
  * absolute Jacobian determinant of constraining transform
- * @tparam T A standard vector with inner type inheriting from
- * `Eigen::DenseBase` or a `var_value` with inner type inheriting from
- * `Eigen::DenseBase` with compile time dynamic rows and 1 column
+ * @tparam T A type inheriting from `Eigen::DenseBase` or a `var_value` with
+ *  inner type inheriting from `Eigen::DenseBase` with compile time dynamic rows
+ *  and 1 column, or a standard vector thereof
+ * @tparam Lp A scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param y Linearly Serialized vector of size `(K * (K - 1))/2` holding the
  *  column major order elements of the lower triangurlar
  * @param K The size of the matrix to return
  * @param[in,out] lp log density accumulator
  */
-template <bool Jacobian, typename T, require_std_vector_t<T>* = nullptr>
-inline auto cholesky_corr_constrain(const T& y, int K, return_type_t<T>& lp) {
-  return apply_vector_unary<T>::apply(y, [&lp, K](auto&& v) {
-    return cholesky_corr_constrain<Jacobian>(v, K, lp);
-  });
+template <bool Jacobian, typename T, typename Lp,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
+inline auto cholesky_corr_constrain(const T& y, int K, Lp& lp) {
+  if constexpr (Jacobian) {
+    return cholesky_corr_constrain(y, K, lp);
+  } else {
+    return cholesky_corr_constrain(y, K);
+  }
 }
 
 }  // namespace math

stan/math/prim/constraint/cholesky_corr_free.hpp

@@ -11,9 +11,10 @@ namespace stan {
 namespace math {
 
 template <typename T, require_eigen_t<T>* = nullptr>
-auto cholesky_corr_free(const T& x) {
+inline auto cholesky_corr_free(const T& x) {
   using Eigen::Dynamic;
   using Eigen::Matrix;
+  using std::sqrt;
 
   check_square("cholesky_corr_free", "x", x);
   // should validate lower-triangular, unit lengths
@@ -24,9 +25,9 @@ auto cholesky_corr_free(const T& x) {
   int k = 0;
   for (int i = 1; i < x.rows(); ++i) {
     z.coeffRef(k++) = corr_free(x_ref.coeff(i, 0));
-    double sum_sqs = square(x_ref.coeff(i, 0));
+    auto sum_sqs = square(x_ref.coeff(i, 0));
     for (int j = 1; j < i; ++j) {
-      z.coeffRef(k++) = corr_free(x_ref.coeff(i, j) / std::sqrt(1.0 - sum_sqs));
+      z.coeffRef(k++) = corr_free(x_ref.coeff(i, j) / sqrt(1.0 - sum_sqs));
       sum_sqs += square(x_ref.coeff(i, j));
     }
   }
@@ -41,7 +42,7 @@ auto cholesky_corr_free(const T& x) {
  * @param x The standard vector to untransform.
  */
 template <typename T, require_std_vector_t<T>* = nullptr>
-auto cholesky_corr_free(const T& x) {
+inline auto cholesky_corr_free(const T& x) {
   return apply_vector_unary<T>::apply(
       x, [](auto&& v) { return cholesky_corr_free(v); });
 }

stan/math/prim/constraint/cholesky_factor_constrain.hpp

@@ -70,9 +70,10 @@ cholesky_factor_constrain(const T& x, int M, int N) {
  * determinant
  * @return Cholesky factor
  */
-template <typename T, require_eigen_vector_t<T>* = nullptr>
+template <typename T, typename Lp, require_eigen_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
 inline Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
-cholesky_factor_constrain(const T& x, int M, int N, return_type_t<T>& lp) {
+cholesky_factor_constrain(const T& x, int M, int N, Lp& lp) {
   check_size_match("cholesky_factor_constrain", "x.size()", x.size(),
                    "((N * (N + 1)) / 2 + (M - N) * N)",
                    ((N * (N + 1)) / 2 + (M - N) * N));
@@ -88,31 +89,46 @@ cholesky_factor_constrain(const T& x, int M, int N, return_type_t<T>& lp) {
 /**
  * Return the Cholesky factor of the specified size read from the specified
  * vector. A total of (N choose 2) + N + N * (M - N) free parameters are
- * required to read an M by N Cholesky factor. If the `Jacobian` parameter is
- * `true`, the log density accumulator is incremented with the log absolute
- * Jacobian determinant of the transform.  All of the transforms are specified
- * with their Jacobians in the *Stan Reference Manual* chapter Constraint
- * Transforms.
+ * required to read an M by N Cholesky factor.
+ * This overload handles looping over the elements of a standard vector.
  *
- * @tparam Jacobian if `true`, increment log density accumulator with log
- * absolute Jacobian determinant of constraining transform
- * @tparam T A type inheriting from `Eigen::DenseBase` or a `var_value` with
- *  inner type inheriting from `Eigen::DenseBase` with compile time dynamic rows
- *  and 1 column
+ * @tparam T A standard vector with inner type inheriting from
+ * `Eigen::DenseBase` or a `var_value` with inner type inheriting from
+ * `Eigen::DenseBase` with compile time dynamic rows and 1 column
+ * @param x Vector of unconstrained values
+ * @param M number of rows
+ * @param N number of columns
+ * @return Cholesky factor
+ */
+template <typename T, require_std_vector_t<T>* = nullptr>
+inline auto cholesky_factor_constrain(const T& x, int M, int N) {
+  return apply_vector_unary<T>::apply(
+      x, [M, N](auto&& v) { return cholesky_factor_constrain(v, M, N); });
+}
+
+/**
+ * Return the Cholesky factor of the specified size read from the specified
+ * vector. A total of (N choose 2) + N + N * (M - N) free parameters are
+ * required to read an M by N Cholesky factor.
+ * This overload handles looping over the elements of a standard vector.
+ *
+ * @tparam T A standard vector with inner type inheriting from
+ * `Eigen::DenseBase` or a `var_value` with inner type inheriting from
+ * `Eigen::DenseBase` with compile time dynamic rows and 1 column
+ * @tparam Lp Scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param x Vector of unconstrained values
  * @param M number of rows
  * @param N number of columns
  * @param[in,out] lp log density accumulator
  * @return Cholesky factor
  */
-template <bool Jacobian, typename T, require_not_std_vector_t<T>* = nullptr>
-inline auto cholesky_factor_constrain(const T& x, int M, int N,
-                                      return_type_t<T>& lp) {
-  if (Jacobian) {
-    return cholesky_factor_constrain(x, M, N, lp);
-  } else {
-    return cholesky_factor_constrain(x, M, N);
-  }
+template <typename T, typename Lp, require_std_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
+inline auto cholesky_factor_constrain(const T& x, int M, int N, Lp& lp) {
+  return apply_vector_unary<T>::apply(x, [&lp, M, N](auto&& v) {
+    return cholesky_factor_constrain(v, M, N, lp);
+  });
 }
 
 /**
@@ -126,21 +142,25 @@ inline auto cholesky_factor_constrain(const T& x, int M, int N,
  *
  * @tparam Jacobian if `true`, increment log density accumulator with log
  * absolute Jacobian determinant of constraining transform
- * @tparam T A standard vector with inner type inheriting from
- * `Eigen::DenseBase` or a `var_value` with inner type inheriting from
- * `Eigen::DenseBase` with compile time dynamic rows and 1 column
+ * @tparam T A type inheriting from `Eigen::DenseBase` or a `var_value` with
+ *  inner type inheriting from `Eigen::DenseBase` with compile time dynamic rows
+ *  and 1 column
+ * @tparam Lp A scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param x Vector of unconstrained values
  * @param M number of rows
  * @param N number of columns
  * @param[in,out] lp log density accumulator
  * @return Cholesky factor
  */
-template <bool Jacobian, typename T, require_std_vector_t<T>* = nullptr>
-inline auto cholesky_factor_constrain(const T& x, int M, int N,
-                                      return_type_t<T>& lp) {
-  return apply_vector_unary<T>::apply(x, [&lp, M, N](auto&& v) {
-    return cholesky_factor_constrain<Jacobian>(v, M, N, lp);
-  });
+template <bool Jacobian, typename T, typename Lp,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
+inline auto cholesky_factor_constrain(const T& x, int M, int N, Lp& lp) {
+  if constexpr (Jacobian) {
+    return cholesky_factor_constrain(x, M, N, lp);
+  } else {
+    return cholesky_factor_constrain(x, M, N);
+  }
 }
 
 }  // namespace math

stan/math/prim/constraint/cholesky_factor_free.hpp

@@ -56,7 +56,7 @@ Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, 1> cholesky_factor_free(
  * @param x The standard vector to untransform.
  */
 template <typename T, require_std_vector_t<T>* = nullptr>
-auto cholesky_factor_free(const T& x) {
+inline auto cholesky_factor_free(const T& x) {
   return apply_vector_unary<T>::apply(
       x, [](auto&& v) { return cholesky_factor_free(v); });
 }

stan/math/prim/constraint/corr_constrain.hpp

@@ -66,7 +66,7 @@ inline auto corr_constrain(const T_x& x, T_lp& lp) {
  */
 template <bool Jacobian, typename T_x, typename T_lp>
 inline auto corr_constrain(const T_x& x, T_lp& lp) {
-  if (Jacobian) {
+  if constexpr (Jacobian) {
     return corr_constrain(x, lp);
   } else {
     return corr_constrain(x);

stan/math/prim/constraint/corr_free.hpp

@@ -25,7 +25,7 @@ namespace math {
  * @return free scalar that transforms to the specified input
  */
 template <typename T>
-inline T corr_free(const T& y) {
+inline plain_type_t<T> corr_free(const T& y) {
   check_bounded("lub_free", "Correlation variable", y, -1.0, 1.0);
   return atanh(y);
 }

stan/math/prim/constraint/corr_matrix_constrain.hpp

@@ -61,13 +61,16 @@ corr_matrix_constrain(const T& x, Eigen::Index k) {
  *
  * @tparam T type of the vector (must be derived from \c Eigen::MatrixBase and
  * have one compile-time dimension equal to 1)
+ * @tparam Lp A scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param x Vector of unconstrained partial correlations.
  * @param k Dimensionality of returned correlation matrix.
  * @param lp Log probability reference to increment.
  */
-template <typename T, require_eigen_col_vector_t<T>* = nullptr>
+template <typename T, typename Lp, require_eigen_col_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
 inline Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
-corr_matrix_constrain(const T& x, Eigen::Index k, return_type_t<T>& lp) {
+corr_matrix_constrain(const T& x, Eigen::Index k, Lp& lp) {
   Eigen::Index k_choose_2 = (k * (k - 1)) / 2;
   check_size_match("cov_matrix_constrain", "x.size()", x.size(), "k_choose_2",
                    k_choose_2);
@@ -79,28 +82,41 @@ corr_matrix_constrain(const T& x, Eigen::Index k, return_type_t<T>& lp) {
  * the specified vector of unconstrained values. The input vector must be of
  * length \f${k \choose 2} = \frac{k(k-1)}{2}\f$.  The values in the input
  * vector represent unconstrained (partial) correlations among the dimensions.
- * If the `Jacobian` parameter is `true`, the log density accumulator is
- * incremented with the log absolute Jacobian determinant of the transform.  All
- * of the transforms are specified with their Jacobians in the *Stan Reference
- * Manual* chapter Constraint Transforms.
+ * This overload handles looping over the elements of a standard vector.
  *
- * @tparam Jacobian if `true`, increment log density accumulator with log
- * absolute Jacobian determinant of constraining transform
- * @tparam T A type inheriting from `Eigen::DenseBase` or a `var_value` with
- *  inner type inheriting from `Eigen::DenseBase` with compile time dynamic rows
- *  and 1 column
- * @param x Vector of unconstrained partial correlations
- * @param k Dimensionality of returned correlation matrix
- * @param[in,out] lp log density accumulator
+ * @tparam T A standard vector with inner type inheriting from
+ * `Eigen::DenseBase` or a `var_value` with inner type inheriting from
+ * `Eigen::DenseBase` with compile time dynamic rows and 1 column
+ * @param y Vector of unconstrained partial correlations
+ * @param K Dimensionality of returned correlation matrix
  */
-template <bool Jacobian, typename T, require_not_std_vector_t<T>* = nullptr>
-inline auto corr_matrix_constrain(const T& x, Eigen::Index k,
-                                  return_type_t<T>& lp) {
-  if (Jacobian) {
-    return corr_matrix_constrain(x, k, lp);
-  } else {
-    return corr_matrix_constrain(x, k);
-  }
+template <typename T, require_std_vector_t<T>* = nullptr>
+inline auto corr_matrix_constrain(const T& y, int K) {
+  return apply_vector_unary<T>::apply(
+      y, [K](auto&& v) { return corr_matrix_constrain(v, K); });
+}
+
+/**
+ * Return the correlation matrix of the specified dimensionality derived from
+ * the specified vector of unconstrained values. The input vector must be of
+ * length \f${k \choose 2} = \frac{k(k-1)}{2}\f$.  The values in the input
+ * vector represent unconstrained (partial) correlations among the dimensions.
+ * This overload handles looping over the elements of a standard vector.
+ *
+ * @tparam T A standard vector with inner type inheriting from
+ * `Eigen::DenseBase` or a `var_value` with inner type inheriting from
+ * `Eigen::DenseBase` with compile time dynamic rows and 1 column
+ * @tparam Lp Scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
+ * @param y Vector of unconstrained partial correlations
+ * @param K Dimensionality of returned correlation matrix
+ * @param[in, out] lp log density accumulator o
+ */
+template <typename T, typename Lp, require_std_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
+inline auto corr_matrix_constrain(const T& y, int K, Lp& lp) {
+  return apply_vector_unary<T>::apply(
+      y, [&lp, K](auto&& v) { return corr_matrix_constrain(v, K, lp); });
 }
 
 /**
@@ -115,18 +131,23 @@ inline auto corr_matrix_constrain(const T& x, Eigen::Index k,
  *
  * @tparam Jacobian if `true`, increment log density accumulator with log
  * absolute Jacobian determinant of constraining transform
- * @tparam T A standard vector with inner type inheriting from
- * `Eigen::DenseBase` or a `var_value` with inner type inheriting from
- * `Eigen::DenseBase` with compile time dynamic rows and 1 column
- * @param y Vector of unconstrained partial correlations
- * @param K Dimensionality of returned correlation matrix
+ * @tparam T A type inheriting from `Eigen::DenseBase` or a `var_value` with
+ *  inner type inheriting from `Eigen::DenseBase` with compile time dynamic rows
+ *  and 1 column or standard vector thereof
+ * @tparam Lp A scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
+ * @param x Vector of unconstrained partial correlations
+ * @param k Dimensionality of returned correlation matrix
  * @param[in,out] lp log density accumulator
  */
-template <bool Jacobian, typename T, require_std_vector_t<T>* = nullptr>
-inline auto corr_matrix_constrain(const T& y, int K, return_type_t<T>& lp) {
-  return apply_vector_unary<T>::apply(y, [&lp, K](auto&& v) {
-    return corr_matrix_constrain<Jacobian>(v, K, lp);
-  });
+template <bool Jacobian, typename T, typename Lp,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
+inline auto corr_matrix_constrain(const T& x, Eigen::Index k, Lp& lp) {
+  if constexpr (Jacobian) {
+    return corr_matrix_constrain(x, k, lp);
+  } else {
+    return corr_matrix_constrain(x, k);
+  }
 }
 
 }  // namespace math