Merge pull request #2373 from stan-dev/feature/lb_constrain-matvar

Adds lb/ub/lub_constrain specializations

Author: bbbales2

stan/math/prim/fun/identity_constrain.hpp

@@ -8,37 +8,18 @@ namespace math {
 
 /**
  * Returns the result of applying the identity constraint
- * transform to the input.
+ * transform to the input. This promotes the input to the least upper
+ * bound of the input types.
  *
- * <p>This method is effectively a no-op and is mainly useful as a
- * placeholder in auto-generated code.
- *
- * @tparam T Any type.
- * @param[in] x object
- * @return transformed input
- */
-template <typename T>
-inline auto identity_constrain(T&& x) {
-  return std::forward<T>(x);
-}
-
-/**
- * Returns the result of applying the identity constraint
- * transform to the input and increments the log probability
- * reference with the log absolute Jacobian determinant.
- *
- * <p>This method is effectively a no-op and mainly useful as a
- * placeholder in auto-generated code.
- *
- * @tparam T Any type
- * @tparam S type of log probability
+ * @tparam T type of value to promote
+ * @tparam Types Other types.
  * @param[in] x object
- * @param[in] lp log density reference
  * @return transformed input
  */
-template <typename T, typename S>
-inline auto identity_constrain(T&& x, S& lp) {
-  return std::forward<T>(x);
+template <typename T, typename... Types,
+          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);
 }
 
 }  // namespace math

stan/math/prim/fun/identity_free.hpp

@@ -8,18 +8,19 @@ namespace math {
 
 /**
  * Returns the result of applying the inverse of the identity
- * constraint transform to the input.
+ * constraint transform to the input. This promotes the input to the least upper
+ * bound of the input types.
  *
- * <p>This function is a no-op and mainly useful as a placeholder
- * in auto-generated code.
  *
- * @tparam T type of value
- * @param[in] y value
+ * @tparam T type of value to promote
+ * @tparam Types Other types.
+ * @param[in] x value to promote
  * @return value
  */
-template <typename T>
-inline auto identity_free(T&& y) {
-  return std::forward<T>(y);
+template <typename T, typename... Types,
+          require_all_not_var_matrix_t<T, Types...>* = nullptr>
+inline auto identity_free(T&& x, Types&&... /* args */) {
+  return promote_scalar_t<return_type_t<T, Types...>, T>(x);
 }
 
 }  // namespace math

stan/math/prim/fun/lb_constrain.hpp

@@ -3,8 +3,12 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/constants.hpp>
+#include <stan/math/prim/fun/identity_constrain.hpp>
+#include <stan/math/prim/fun/identity_free.hpp>
 #include <stan/math/prim/fun/add.hpp>
 #include <stan/math/prim/fun/exp.hpp>
+#include <stan/math/prim/fun/eval.hpp>
 #include <stan/math/prim/fun/sum.hpp>
 #include <stan/math/prim/fun/value_of.hpp>
 #include <cmath>
@@ -22,17 +26,20 @@ namespace math {
  *
  * <p>where \f$L\f$ is the constant lower bound.
  *
- * @tparam T type of Matrix
- * @tparam L type of lower bound
- * @param[in] x Unconstrained Matrix input
+ * @tparam T Scalar.
+ * @tparam L Scalar.
+ * @param[in] x Unconstrained input
  * @param[in] lb Lower bound
  * @return Constrained matrix
  */
-template <typename T, typename L>
+template <typename T, typename L, require_all_stan_scalar_t<T, L>* = nullptr,
+          require_all_not_st_var<T, L>* = nullptr>
 inline auto lb_constrain(const T& x, const L& lb) {
-  auto& lb_ref = to_ref(lb);
-  check_finite("lb_constrain", "lb", value_of(lb_ref));
-  return eval(add(exp(x), lb_ref));
+  if (unlikely(value_of_rec(lb) == NEGATIVE_INFTY)) {
+    return identity_constrain(x, lb);
+  } else {
+    return add(exp(x), lb);
+  }
 }
 
 /**
@@ -41,21 +48,177 @@ inline auto lb_constrain(const T& x, const L& lb) {
  * reference with the log absolute Jacobian determinant of the
  * transform.
  *
- * @tparam T Type of Matrix
- * @tparam L type of lower bound
- * @tparam S type of log probability
- * @param[in] x unconstrained Matrix input
+ * @tparam T Scalar.
+ * @tparam L Scalar.
+ * @param[in] x unconstrained input
  * @param[in] lb lower bound on output
  * @param[in,out] lp reference to log probability to increment
  * @return lower-bound constrained value corresponding to inputs
  */
-template <typename T, typename L>
+template <typename T, typename L, require_all_stan_scalar_t<T, L>* = nullptr,
+          require_all_not_st_var<T, L>* = nullptr>
+inline auto lb_constrain(const T& x, const L& lb, return_type_t<T, L>& lp) {
+  if (value_of_rec(lb) == NEGATIVE_INFTY) {
+    return identity_constrain(x, lb);
+  } else {
+    lp += x;
+    return add(exp(x), lb);
+  }
+}
+
+/**
+ * Specialization of `lb_constrain` to apply a scalar lower bound elementwise
+ *  to each input.
+ *
+ * @tparam T A type inheriting from `EigenBase`.
+ * @tparam L Scalar.
+ * @param[in] x unconstrained input
+ * @param[in] lb lower bound on output
+ * @return lower-bound constrained value corresponding to inputs
+ */
+template <typename T, typename L, require_eigen_t<T>* = nullptr,
+          require_stan_scalar_t<L>* = nullptr,
+          require_all_not_st_var<T, L>* = nullptr>
+inline auto lb_constrain(T&& x, L&& lb) {
+  return eval(x.unaryExpr([lb](auto&& x) { return lb_constrain(x, lb); }));
+}
+
+/**
+ * Specialization of `lb_constrain` to apply a scalar lower bound elementwise
+ *  to each input.
+ *
+ * @tparam T A type inheriting from `EigenBase`.
+ * @tparam L Scalar.
+ * @param[in] x unconstrained input
+ * @param[in] lb lower bound on output
+ * @param[in,out] lp reference to log probability to increment
+ * @return lower-bound constrained value corresponding to inputs
+ */
+template <typename T, typename L, require_eigen_t<T>* = nullptr,
+          require_stan_scalar_t<L>* = nullptr,
+          require_all_not_st_var<T, L>* = nullptr>
+inline auto lb_constrain(const T& x, const L& lb, return_type_t<T, L>& lp) {
+  return eval(
+      x.unaryExpr([lb, &lp](auto&& xx) { return lb_constrain(xx, lb, lp); }));
+}
+
+/**
+ * Specialization of `lb_constrain` to apply a matrix of lower bounds
+ * elementwise to each input element.
+ *
+ * @tparam T A type inheriting from `EigenBase`.
+ * @tparam L A type inheriting from `EigenBase`.
+ * @param[in] x unconstrained input
+ * @param[in] lb lower bound on output
+ * @return lower-bound constrained value corresponding to inputs
+ */
+template <typename T, typename L, require_all_eigen_t<T, L>* = nullptr,
+          require_all_not_st_var<T, L>* = nullptr>
+inline auto lb_constrain(T&& x, L&& lb) {
+  check_matching_dims("lb_constrain", "x", x, "lb", lb);
+  return eval(x.binaryExpr(
+      lb, [](auto&& x, auto&& lb) { return lb_constrain(x, lb); }));
+}
+
+/**
+ * Specialization of `lb_constrain` to apply a matrix of lower bounds
+ * elementwise to each input element.
+ *
+ * @tparam T A type inheriting from `EigenBase`.
+ * @tparam L A type inheriting from `EigenBase`.
+ * @param[in] x unconstrained input
+ * @param[in] lb lower bound on output
+ * @param[in,out] lp reference to log probability to increment
+ * @return lower-bound constrained value corresponding to inputs
+ */
+template <typename T, typename L, require_all_eigen_t<T, L>* = nullptr,
+          require_all_not_st_var<T, L>* = nullptr>
 inline auto lb_constrain(const T& x, const L& lb, return_type_t<T, L>& lp) {
-  auto& x_ref = to_ref(x);
-  auto& lb_ref = to_ref(lb);
-  check_finite("lb_constrain", "lb", value_of(lb_ref));
-  lp += sum(x_ref);
-  return eval(add(exp(x_ref), lb_ref));
+  check_matching_dims("lb_constrain", "x", x, "lb", lb);
+  return eval(x.binaryExpr(
+      lb, [&lp](auto&& xx, auto&& lbb) { return lb_constrain(xx, lbb, lp); }));
+}
+
+/**
+ * Specialization of `lb_constrain` to apply a container of lower bounds
+ * elementwise to each input element.
+ *
+ * @tparam T A Any type with a Scalar `scalar_type`.
+ * @tparam L A type inheriting from `EigenBase` or a scalar.
+ * @param[in] x unconstrained input
+ * @param[in] lb lower bound on output
+ * @return lower-bound constrained value corresponding to inputs
+ */
+template <typename T, typename L, require_not_std_vector_t<L>* = nullptr>
+inline auto lb_constrain(const std::vector<T>& x, const L& lb) {
+  std::vector<plain_type_t<decltype(lb_constrain(x[0], lb))>> ret(x.size());
+  for (size_t i = 0; i < x.size(); ++i) {
+    ret[i] = lb_constrain(x[i], lb);
+  }
+  return ret;
+}
+
+/**
+ * Specialization of `lb_constrain` to apply a container of lower bounds
+ * elementwise to each input element.
+ *
+ * @tparam T A Any type with a Scalar `scalar_type`.
+ * @tparam L A type inheriting from `EigenBase` or a standard vector.
+ * @param[in] x unconstrained input
+ * @param[in] lb lower bound on output
+ * @param[in,out] lp reference to log probability to increment
+ * @return lower-bound constrained value corresponding to inputs
+ */
+template <typename T, typename L, require_not_std_vector_t<L>* = nullptr>
+inline auto lb_constrain(const std::vector<T>& x, const L& lb,
+                         return_type_t<T, L>& lp) {
+  std::vector<plain_type_t<decltype(lb_constrain(x[0], lb))>> ret(x.size());
+  for (size_t i = 0; i < x.size(); ++i) {
+    ret[i] = lb_constrain(x[i], lb, lp);
+  }
+  return ret;
+}
+
+/**
+ * Specialization of `lb_constrain` to apply a container of lower bounds
+ * elementwise to each input element.
+ *
+ * @tparam T A Any type with a Scalar `scalar_type`.
+ * @tparam L A type inheriting from `EigenBase` or a standard vector.
+ * @param[in] x unconstrained input
+ * @param[in] lb lower bound on output
+ * @return lower-bound constrained value corresponding to inputs
+ */
+template <typename T, typename L>
+inline auto lb_constrain(const std::vector<T>& x, const std::vector<L>& lb) {
+  check_matching_dims("lb_constrain", "x", x, "lb", lb);
+  std::vector<plain_type_t<decltype(lb_constrain(x[0], lb[0]))>> ret(x.size());
+  for (size_t i = 0; i < x.size(); ++i) {
+    ret[i] = lb_constrain(x[i], lb[i]);
+  }
+  return ret;
+}
+
+/**
+ * Specialization of `lb_constrain` to apply a container of lower bounds
+ * elementwise to each input element.
+ *
+ * @tparam T A Any type with a Scalar `scalar_type`.
+ * @tparam L A type inheriting from `EigenBase` or a standard vector.
+ * @param[in] x unconstrained input
+ * @param[in] lb lower bound on output
+ * @param[in,out] lp reference to log probability to increment
+ * @return lower-bound constrained value corresponding to inputs
+ */
+template <typename T, typename L>
+inline auto lb_constrain(const std::vector<T>& x, const std::vector<L>& lb,
+                         return_type_t<T, L>& lp) {
+  check_matching_dims("lb_constrain", "x", x, "lb", lb);
+  std::vector<plain_type_t<decltype(lb_constrain(x[0], lb[0]))>> ret(x.size());
+  for (size_t i = 0; i < x.size(); ++i) {
+    ret[i] = lb_constrain(x[i], lb[i], lp);
+  }
+  return ret;
 }
 
 }  // namespace math