Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 45 additions & 27 deletions stan/math/prim/constraint/cholesky_corr_constrain.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So this will allow lp to be var and return_type_t<EigVec> to be double or var.

Do we ever want to increment the jacobian when y just has a scalar type of double?

lp += 0.5 * log1m(sum_sqs);

What are the use types you want to pass into this? Is it just the return_type_t<EigVec> = double and lp = var case? Since the data is considered constant in that case shouldn't we just call the one that does not increment the jacobian?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

return_type_t<EigVec> = double and lp = var

Yes, that's the relevant one that comes up during code generation. Based on my conversation with @bob-carpenter, I think we still want to let the user do it, even if it's weird. The implementation in stanc3#1494 gives the user both a _jacobian function and a _constrain function, and I think it's reasonable to view these as similar to the difference between doing a target+= with a _lpdf versus a _lupdf.

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;
Expand Down Expand Up @@ -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); });
}

/**
Expand All @@ -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
Expand Down
9 changes: 5 additions & 4 deletions stan/math/prim/constraint/cholesky_corr_free.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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));
}
}
Expand All @@ -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); });
}
Expand Down
78 changes: 49 additions & 29 deletions stan/math/prim/constraint/cholesky_factor_constrain.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand All @@ -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);
});
}

/**
Expand All @@ -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
Expand Down
2 changes: 1 addition & 1 deletion stan/math/prim/constraint/cholesky_factor_free.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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); });
}
Expand Down
2 changes: 1 addition & 1 deletion stan/math/prim/constraint/corr_constrain.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion stan/math/prim/constraint/corr_free.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
85 changes: 53 additions & 32 deletions stan/math/prim/constraint/corr_matrix_constrain.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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); });
}

/**
Expand All @@ -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
Expand Down
Loading