Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Create Jacobian template parameter overload for constrain functions #2559

Merged
merged 8 commits into from
Aug 27, 2021
32 changes: 28 additions & 4 deletions stan/math/prim/fun/cholesky_corr_constrain.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,14 @@ namespace stan {
namespace math {

template <typename EigVec, require_eigen_col_vector_t<EigVec>* = nullptr>
Eigen::Matrix<value_type_t<EigVec>, Eigen::Dynamic, Eigen::Dynamic>
inline Eigen::Matrix<value_type_t<EigVec>, Eigen::Dynamic, Eigen::Dynamic>
cholesky_corr_constrain(const EigVec& y, int K) {
using Eigen::Dynamic;
using Eigen::Matrix;
using std::sqrt;
using T_scalar = value_type_t<EigVec>;
int k_choose_2 = (K * (K - 1)) / 2;
check_size_match("cholesky_corr_constrain", "y.size()", y.size(),
check_size_match("cholesky_corr_constrain", "constrain size", y.size(),
"k_choose_2", k_choose_2);
Matrix<T_scalar, Dynamic, 1> z = corr_constrain(y);
Matrix<T_scalar, Dynamic, Dynamic> x(K, K);
Expand All @@ -44,8 +44,8 @@ cholesky_corr_constrain(const EigVec& y, int K) {

// FIXME to match above after debugged
template <typename EigVec, require_eigen_vector_t<EigVec>* = nullptr>
Eigen::Matrix<value_type_t<EigVec>, Eigen::Dynamic, Eigen::Dynamic>
cholesky_corr_constrain(const EigVec& y, int K, value_type_t<EigVec>& lp) {
inline Eigen::Matrix<value_type_t<EigVec>, Eigen::Dynamic, Eigen::Dynamic>
cholesky_corr_constrain(const EigVec& y, int K, return_type_t<EigVec>& lp) {
using Eigen::Dynamic;
using Eigen::Matrix;
using std::sqrt;
Expand Down Expand Up @@ -74,6 +74,30 @@ cholesky_corr_constrain(const EigVec& y, int K, value_type_t<EigVec>& lp) {
return x;
}

/**
* 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
* @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>
inline auto cholesky_corr_constrain(const T& y, int K, return_type_t<T>& lp) {
if (Jacobian) {
Copy link
Contributor

Choose a reason for hiding this comment

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

[question]
I'm unclear on how this function gets resolved since it seems to call itself inside the if (Jacobian) block. Is this only resolved because the nested call doesn't provide a template parameter?

[comment]
I opened an issue to refactor the implementation: #2561

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Yes since there's no default template value for Jacobian it diverts to the signatures without the Jacobian template parameter

return cholesky_corr_constrain(y, K, lp);
} else {
return cholesky_corr_constrain(y, K);
}
}
} // namespace math
} // namespace stan
#endif
47 changes: 39 additions & 8 deletions stan/math/prim/fun/cholesky_factor_constrain.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -26,14 +26,14 @@ namespace math {
* @return Cholesky factor
*/
template <typename T, require_eigen_col_vector_t<T>* = nullptr>
Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
inline Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
cholesky_factor_constrain(const T& x, int M, int N) {
using std::exp;
using T_scalar = value_type_t<T>;
check_greater_or_equal("cholesky_factor_constrain",
"num rows (must be greater or equal to num cols)", M,
N);
check_size_match("cholesky_factor_constrain", "x.size()", x.size(),
check_size_match("cholesky_factor_constrain", "constrain size", x.size(),
"((N * (N + 1)) / 2 + (M - N) * N)",
((N * (N + 1)) / 2 + (M - N) * N));
Eigen::Matrix<T_scalar, Eigen::Dynamic, Eigen::Dynamic> y(M, N);
Expand All @@ -58,21 +58,22 @@ cholesky_factor_constrain(const T& x, int M, int N) {
/**
* Return the Cholesky factor of the specified size read from the
* specified vector and increment the specified log probability
* reference with the log Jacobian adjustment of the transform. A total
* of (N choose 2) + N + N * (M - N) free parameters are required to read
* an M by N Cholesky factor.
* reference with the log absolute Jacobian determinant adjustment of the
* transform. A total of (N choose 2) + N + N * (M - N) free parameters are
* required to read an M by N Cholesky factor.
*
* @tparam T type of the vector (must be derived from \c Eigen::MatrixBase and
* have one compile-time dimension equal to 1)
* @param x Vector of unconstrained values
* @param M number of rows
* @param N number of columns
* @param lp Log probability that is incremented with the log Jacobian
* @param lp Log probability that is incremented with the log absolute Jacobian
* determinant
* @return Cholesky factor
*/
template <typename T, require_eigen_vector_t<T>* = nullptr>
Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
cholesky_factor_constrain(const T& x, int M, int N, value_type_t<T>& lp) {
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) {
check_size_match("cholesky_factor_constrain", "x.size()", x.size(),
Copy link
Contributor

Choose a reason for hiding this comment

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

[optional]
Ideally the error messages will be phrased in a way that a Stan user can understand them from the Stan program. So maybe "constrained size" would be better than "x.size()". I don't think this can fail for Stan programs because we manage the sizes, and it's not part of this issue, so I'm marking "optional".

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Yep just fixed these up to `constrained size'.

"((N * (N + 1)) / 2 + (M - N) * N)",
((N * (N + 1)) / 2 + (M - N) * N));
Expand All @@ -85,6 +86,36 @@ cholesky_factor_constrain(const T& x, int M, int N, value_type_t<T>& lp) {
return cholesky_factor_constrain(x_ref, 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. 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
* @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>
inline auto cholesky_factor_constrain(const T& x, int M, int N,
return_type_t<T>& lp) {
if (Jacobian) {
Copy link
Contributor

Choose a reason for hiding this comment

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

[optional]
I generally prefer a ternary operator when appropriate, so this would be

return Jacobian ? cholesky_factor_constrain(x, M, N, lp) : cholesky_factor_constrain(X, M, N);

if it fits or

return Jacobian
  ? cholesky_factor_constrain(x, M, N, lp)
  : cholesky_factor_constrain(X, M, N);

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

imo I like the if (Jacobian) for these just because it future proofs us a bit so that if we ever move to C++17 we can pretty much just do a copy replace for if (Jacobian) for if constexpr (Jacobian)

Copy link
Contributor

Choose a reason for hiding this comment

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

Interesting. Does the ternary operator not get reduced at compile time if the condition is constexpr? Obviously the template parameters are all constexpr!

return cholesky_factor_constrain(x, M, N, lp);
} else {
return cholesky_factor_constrain(x, M, N);
}
}

} // namespace math
} // namespace stan

Expand Down
26 changes: 25 additions & 1 deletion stan/math/prim/fun/corr_constrain.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ namespace math {
* @return tanh transform
*/
template <typename T>
inline auto corr_constrain(const T& x) {
inline plain_type_t<T> corr_constrain(const T& x) {
return tanh(x);
}

Expand All @@ -49,6 +49,30 @@ inline auto corr_constrain(const T_x& x, T_lp& lp) {
return tanh_x;
}

/**
* Return the result of transforming the specified scalar or container of values
Copy link
Contributor

Choose a reason for hiding this comment

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

The doc for all of these needs to mention that the Jacobian is incremented with the log absolute Jacobian determinant of the transform if the template parameter Jacobian is set to true. The goal is that the first sentence of the doc should give as thorough and precise an explanation of the function's arguments and returns as possible. If you don't think it fits in one sentence, add a second sentence on the Jacobian.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Yes I appended the doc you mention below with all of these so should be good now

Copy link
Contributor

Choose a reason for hiding this comment

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

Looks good. Thanks.

* to have a valid correlation value between -1 and 1 (inclusive). 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_x Type of scalar or container
* @tparam T_lp type of target log density accumulator
* @param[in] x value or container
* @param[in,out] lp log density accumulator
*/
template <bool Jacobian, typename T_x, typename T_lp>
inline auto corr_constrain(const T_x& x, T_lp& lp) {
if (Jacobian) {
return corr_constrain(x, lp);
} else {
return corr_constrain(x);
}
}

} // namespace math
} // namespace stan
#endif
35 changes: 32 additions & 3 deletions stan/math/prim/fun/corr_matrix_constrain.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ namespace math {
* matrix.
*/
template <typename T, require_eigen_col_vector_t<T>* = nullptr>
Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
inline Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
corr_matrix_constrain(const T& x, Eigen::Index k) {
Eigen::Index k_choose_2 = (k * (k - 1)) / 2;
check_size_match("cov_matrix_constrain", "x.size()", x.size(), "k_choose_2",
Expand Down Expand Up @@ -66,14 +66,43 @@ corr_matrix_constrain(const T& x, Eigen::Index k) {
* @param lp Log probability reference to increment.
*/
template <typename T, require_eigen_col_vector_t<T>* = nullptr>
Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
corr_matrix_constrain(const T& x, Eigen::Index k, value_type_t<T>& lp) {
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) {
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);
return read_corr_matrix(corr_constrain(x, lp), k, lp);
}

/**
* 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.
* 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
* @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>
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);
}
}

} // namespace math
} // namespace stan

Expand Down
36 changes: 31 additions & 5 deletions stan/math/prim/fun/cov_matrix_constrain.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ namespace math {
* @throws std::invalid_argument if (x.size() != K + (K choose 2)).
*/
template <typename T, require_eigen_col_vector_t<T>* = nullptr>
Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
inline Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
cov_matrix_constrain(const T& x, Eigen::Index K) {
using Eigen::Dynamic;
using Eigen::Matrix;
Expand All @@ -53,8 +53,6 @@ cov_matrix_constrain(const T& x, Eigen::Index K) {
* by K resulting from transforming the specified finite vector of
* size K plus (K choose 2).
*
* <p>See <code>cov_matrix_free()</code> for the inverse transform.
*
* @tparam T type of the vector (must be derived from \c Eigen::MatrixBase and
* have one compile-time dimension equal to 1)
* @param x The vector to convert to a covariance matrix.
Expand All @@ -63,8 +61,8 @@ cov_matrix_constrain(const T& x, Eigen::Index K) {
* @throws std::domain_error if (x.size() != K + (K choose 2)).
*/
template <typename T, require_eigen_col_vector_t<T>* = nullptr>
Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
cov_matrix_constrain(const T& x, Eigen::Index K, value_type_t<T>& lp) {
inline Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
cov_matrix_constrain(const T& x, Eigen::Index K, return_type_t<T>& lp) {
using Eigen::Dynamic;
using Eigen::Matrix;
using std::exp;
Expand All @@ -88,6 +86,34 @@ cov_matrix_constrain(const T& x, Eigen::Index K, value_type_t<T>& lp) {
return multiply_lower_tri_self_transpose(L);
}

/**
* Return the symmetric, positive-definite matrix of dimensions K by K resulting
* from transforming the specified finite vector of size K plus (K choose 2). 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.
*
Copy link
Contributor

Choose a reason for hiding this comment

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

Good doc on function, but needs to mention Jacobian and its template parameter either in this first sentence or the second one.

* @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
Copy link
Contributor

Choose a reason for hiding this comment

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

Same comments as on other doc:

  • incremented -> increment
  • only capitalize sentences
  • only put periods at the end of sentences

* inner type inheriting from `Eigen::DenseBase` with compile time dynamic rows
* and 1 column
* @param x The vector to convert to a covariance matrix
* @param K The dimensions of the resulting covariance matrix
* @param[in, out] lp log density accumulator
* @throws std::domain_error if (x.size() != K + (K choose 2)).
*/
template <bool Jacobian, typename T>
inline auto cov_matrix_constrain(const T& x, Eigen::Index K,
return_type_t<T>& lp) {
if (Jacobian) {
return cov_matrix_constrain(x, K, lp);
} else {
return cov_matrix_constrain(x, K);
}
}

} // namespace math
} // namespace stan

Expand Down
46 changes: 33 additions & 13 deletions stan/math/prim/fun/cov_matrix_constrain_lkj.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ namespace math {
* correlations and deviations.
*/
template <typename T, require_eigen_vector_t<T>* = nullptr>
Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
inline Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
cov_matrix_constrain_lkj(const T& x, size_t k) {
size_t k_choose_2 = (k * (k - 1)) / 2;
const auto& x_ref = to_ref(x);
Expand All @@ -45,16 +45,6 @@ cov_matrix_constrain_lkj(const T& x, size_t k) {
* values and increment the specified log probability reference
* with the log absolute Jacobian determinant.
*
* <p>The transform is defined as for
* <code>cov_matrix_constrain(Matrix, size_t)</code>.
*
* <p>The log absolute Jacobian determinant is derived by
* composing the log absolute Jacobian determinant for the
* underlying correlation matrix as defined in
* <code>cov_matrix_constrain(Matrix, size_t, T&)</code> with
* the Jacobian of the transform of the correlation matrix
* into a covariance matrix by scaling by standard deviations.
*
* @tparam T type of the vector (must be derived from \c Eigen::MatrixBase and
* have one compile-time dimension equal to 1)
* @param x Input vector of unconstrained partial correlations and
Expand All @@ -65,14 +55,44 @@ cov_matrix_constrain_lkj(const T& x, size_t k) {
* correlations and deviations.
*/
template <typename T, require_eigen_vector_t<T>* = nullptr>
Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
cov_matrix_constrain_lkj(const T& x, size_t k, value_type_t<T>& lp) {
inline Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
cov_matrix_constrain_lkj(const T& x, size_t k, return_type_t<T>& lp) {
size_t k_choose_2 = (k * (k - 1)) / 2;
const auto& x_ref = x;
return read_cov_matrix(corr_constrain(x_ref.head(k_choose_2)),
positive_constrain(x_ref.tail(k)), lp);
}

/**
* Return the covariance matrix of the specified dimensionality derived from
* constraining the specified vector of unconstrained values. 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 rows or
* columns equal to 1
* @param x Input vector of unconstrained partial correlations and
* standard deviations
* @param k Dimensionality of returned covariance matrix
* @param[in, out] lp log density accumulator
* @return Covariance matrix derived from the unconstrained partial
* correlations and deviations.
*/
template <bool Jacobian, typename T>
inline auto cov_matrix_constrain_lkj(const T& x, size_t k,
return_type_t<T>& lp) {
if (Jacobian) {
return cov_matrix_constrain_lkj(x, k, lp);
} else {
return cov_matrix_constrain_lkj(x, k);
}
}

} // namespace math
} // namespace stan

Expand Down
4 changes: 3 additions & 1 deletion stan/math/prim/fun/identity_constrain.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,14 @@ namespace math {
* transform to the input. This promotes the input to the least upper
* bound of the input types.
*
* @tparam Jacobian if `true`, increment log density accumulator with log
* absolute Jacobian determinant of constraining transform
* @tparam T type of value to promote
* @tparam Types Other types.
* @param[in] x object
* @return transformed input
*/
template <typename T, typename... Types,
template <bool Jacobian = false, typename T, typename... Types,
Copy link
Contributor

Choose a reason for hiding this comment

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

[question]
Can type inference work here when the first template parameter gets a default? If so, it seems like we can eventually save a function definition in #2561 by using this for all the other constraints.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

With a value of false yes we could take all of current ones with lp and give them the bool template parameter. Though imo I think this could actually end up with more / confusing code as we have specializations for several of the functions with lp as a parameter. I kind of prefer the scheme here where we have a sort of higher level function that will defer to the impls based on the template parameter value.

require_all_not_var_matrix_t<T, Types...>* = nullptr>
inline auto identity_constrain(T&& x, Types&&... /* args */) {
return promote_scalar_t<return_type_t<T, Types...>, T>(x);
Expand Down
Loading