Add simplex row and column matrix constraints

stan/math/prim/fun.hpp

@@ -315,8 +315,10 @@
 #include <stan/math/prim/fun/select.hpp>
 #include <stan/math/prim/fun/sign.hpp>
 #include <stan/math/prim/fun/signbit.hpp>
+#include <stan/math/prim/fun/simplex_column_constrain.hpp>
 #include <stan/math/prim/fun/simplex_constrain.hpp>
 #include <stan/math/prim/fun/simplex_free.hpp>
+#include <stan/math/prim/fun/simplex_row_constrain.hpp>
 #include <stan/math/prim/fun/sin.hpp>
 #include <stan/math/prim/fun/singular_values.hpp>
 #include <stan/math/prim/fun/sinh.hpp>

stan/math/prim/fun/simplex_column_constrain.hpp

@@ -0,0 +1,115 @@
+#ifndef STAN_MATH_PRIM_FUN_SIMPLEX_COLUMN_CONSTRAIN_HPP
+#define STAN_MATH_PRIM_FUN_SIMPLEX_COLUMN_CONSTRAIN_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/fun/inv_logit.hpp>
+#include <stan/math/prim/fun/log.hpp>
+#include <stan/math/prim/fun/log1p_exp.hpp>
+#include <stan/math/prim/fun/logit.hpp>
+#include <stan/math/prim/fun/simplex_constrain.hpp>
+#include <cmath>
+
+namespace stan {
+namespace math {
+
+/**
+ * Return a matrix with columns as simplex vectors.
+ * A simplex is a vector containing values greater than or equal
+ * to 0 that sum to 1.  A vector with (K-1) unconstrained values
+ * will produce a simplex of size K.
+ *
+ * The transform is based on a centered stick-breaking process.
+ *
+ * @tparam Mat type of the Matrix
+ * @param y Free Matrix input of dimensionality (K - 1, M).
+ * @return Matrix with simplex columns of dimensionality (K, M).
+ */
+template <typename Mat, require_eigen_matrix_dynamic_t<Mat>* = nullptr,
+          require_not_st_var<Mat>* = nullptr>
+inline auto simplex_column_constrain(const Mat& y) {
+  // cut & paste simplex_column_constrain(Eigen::Matrix, T) w/o Jacobian
+  auto&& y_ref = to_ref(y);
+  const Eigen::Index M = y_ref.cols();
+  plain_type_t<Mat> ret(y_ref.rows() + 1, M);
+  for (Eigen::Index i = 0; i < M; ++i) {
+    ret.col(i) = simplex_constrain(y_ref.col(i));
+  }
+  return ret;
+}
+
+/**
+ * Return a matrix with columns as simplex vectors
+ * and increment the specified log probability reference with
+ * the log absolute Jacobian determinant of the transform.
+ *
+ * The simplex transform is defined through a centered
+ * stick-breaking process.
+ *
+ * @tparam Mat type of the Matrix
+ * @param y Free Matrix input of dimensionality (K - 1, M).
+ * @param lp Log probability reference to increment.
+ * @return Matrix with simplex columns of dimensionality (K, M).
+ */
+template <typename Mat, require_eigen_matrix_dynamic_t<Mat>* = nullptr,
+          require_not_st_var<Mat>* = nullptr>
+inline auto simplex_column_constrain(const Mat& y, value_type_t<Mat>& lp) {
+  auto&& y_ref = to_ref(y);
+  const Eigen::Index M = y_ref.cols();
+  plain_type_t<Mat> ret(y_ref.rows() + 1, M);
+  for (Eigen::Index i = 0; i < M; ++i) {
+    ret.col(i) = simplex_constrain(y_ref.col(i), lp);
+  }
+  return ret;
+}
+
+/**
+ * Return a matrix with columns as simplex vectors. 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 Mat type of the Matrix
+ * @param y Free Matrix input of dimensionality (K - 1, M).
+ * @param[in, out] lp log density accumulator
+ * @return Matrix with simplex columns of dimensionality (K, M).
+ */
+template <bool Jacobian, typename Mat, require_not_std_vector_t<Mat>* = nullptr>
+auto simplex_column_constrain(const Mat& y, return_type_t<Mat>& lp) {
+  if (Jacobian) {
+    return simplex_column_constrain(y, lp);
+  } else {
+    return simplex_column_constrain(y);
+  }
+}
+
+/**
+ * Return a standard vector of matrices with columns as simplex vectors. 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 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 dynamic columns
+ * @param[in] y free vector
+ * @param[in, out] lp log density accumulator
+ * @return Standard vector containing matrices with simplex columns of
+ * dimensionality (K, M).
+ */
+template <bool Jacobian, typename T, require_std_vector_t<T>* = nullptr>
+inline auto simplex_column_constrain(const T& y, return_type_t<T>& lp) {
+  return apply_vector_unary<T>::apply(
+      y, [&lp](auto&& v) { return simplex_column_constrain<Jacobian>(v, lp); });
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/prim/fun/simplex_constrain.hpp

@@ -24,7 +24,7 @@ namespace math {
  * @param y Free vector input of dimensionality K - 1.
  * @return Simplex of dimensionality K.
  */
-template <typename Vec, require_eigen_col_vector_t<Vec>* = nullptr,
+template <typename Vec, require_eigen_vector_t<Vec>* = nullptr,
           require_not_st_var<Vec>* = nullptr>
 inline auto simplex_constrain(const Vec& y) {
   // cut & paste simplex_constrain(Eigen::Matrix, T) w/o Jacobian
@@ -56,7 +56,7 @@ inline auto simplex_constrain(const Vec& y) {
  * @param lp Log probability reference to increment.
  * @return Simplex of dimensionality K.
  */
-template <typename Vec, require_eigen_col_vector_t<Vec>* = nullptr,
+template <typename Vec, require_eigen_vector_t<Vec>* = nullptr,
           require_not_st_var<Vec>* = nullptr>
 inline auto simplex_constrain(const Vec& y, value_type_t<Vec>& lp) {
   using Eigen::Dynamic;

stan/math/prim/fun/simplex_row_constrain.hpp

@@ -0,0 +1,132 @@
+#ifndef STAN_MATH_PRIM_FUN_SIMPLEX_ROW_CONSTRAIN_HPP
+#define STAN_MATH_PRIM_FUN_SIMPLEX_ROW_CONSTRAIN_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/fun/inv_logit.hpp>
+#include <stan/math/prim/fun/log.hpp>
+#include <stan/math/prim/fun/log1p_exp.hpp>
+#include <stan/math/prim/fun/logit.hpp>
+#include <stan/math/prim/fun/simplex_constrain.hpp>
+#include <cmath>
+
+namespace stan {
+namespace math {
+
+/**
+ * Return the simplex corresponding to the specified free vector.
+ * A simplex is a vector containing values greater than or equal
+ * to 0 that sum to 1.  A vector with (K-1) unconstrained values
+ * will produce a simplex of size K.
+ *
+ * The transform is based on a centered stick-breaking process.
+ *
+ * @tparam Mat type of the Matrix
+ * @param y Free Matrix input of dimensionality (N, K - 1).
+ * @return Matrix with simplexes along the rows of dimensionality (N, K).
+ */
+template <typename Mat, require_eigen_matrix_dynamic_t<Mat>* = nullptr,
+          require_not_st_var<Mat>* = nullptr>
+inline auto simplex_row_constrain(const Mat& y) {
+  auto&& y_ref = to_ref(y);
+  const Eigen::Index N = y_ref.rows();
+  int Km1 = y_ref.cols();
+  plain_type_t<Mat> x(N, Km1 + 1);
+  using eigen_arr = Eigen::Array<scalar_type_t<Mat>, -1, 1>;
+  eigen_arr stick_len = eigen_arr::Constant(N, 1.0);
+  for (Eigen::Index k = 0; k < Km1; ++k) {
+    auto z_k = inv_logit(y_ref.array().col(k) - log(Km1 - k));
+    x.array().col(k) = stick_len * z_k;
+    stick_len -= x.array().col(k);
+  }
+  x.array().col(Km1) = stick_len;
+  return x;
+}
+
+/**
+ * Return a matrix with simplex rows corresponding to the specified free matrix
+ * and increment the specified log probability reference with
+ * the log absolute Jacobian determinant of the transform.
+ *
+ * The simplex transform is defined through a centered
+ * stick-breaking process.
+ *
+ * @tparam Mat type of the matrix
+ * @param y Free matrix input of dimensionality (N, K - 1).
+ * @param lp Log probability reference to increment.
+ * @return Matrix with simplexes along the rows of dimensionality (N, K).
+ */
+template <typename Mat, require_eigen_matrix_dynamic_t<Mat>* = nullptr,
+          require_not_st_var<Mat>* = nullptr>
+inline auto simplex_row_constrain(const Mat& y, value_type_t<Mat>& lp) {
+  auto&& y_ref = to_ref(y);
+  const Eigen::Index N = y_ref.rows();
+  Eigen::Index Km1 = y_ref.cols();
+  plain_type_t<Mat> x(N, Km1 + 1);
+  Eigen::Array<scalar_type_t<Mat>, -1, 1> stick_len
+      = Eigen::Array<scalar_type_t<Mat>, -1, 1>::Constant(N, 1.0);
+  for (Eigen::Index k = 0; k < Km1; ++k) {
+    auto eq_share = -log(Km1 - k);  // = logit(1.0/(Km1 + 1 - k));
+    auto adj_y_k = (y_ref.array().col(k) + eq_share).eval();
+    auto z_k = inv_logit(adj_y_k);
+    x.array().col(k) = stick_len * z_k;
+    lp += sum(log(stick_len));
+    lp -= sum(log1p_exp(-adj_y_k));
+    lp -= sum(log1p_exp(adj_y_k));
+    stick_len -= x.array().col(k);  // equivalently *= (1 - z_k);
+  }
+  x.col(Km1).array() = stick_len;
+  return x;
+}
+
+/**
+ * Return a matrix with simplex rows corresponding to the specified free 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 Mat A type inheriting from `Eigen::DenseBase` or a `var_value` with
+ *  inner type inheriting from `Eigen::DenseBase` with compile time dynamic rows
+ *  and dynamic columns
+ * @param[in] y free matrix
+ * @param[in, out] lp log density accumulator
+ * @return Matrix with simplexes along the rows of dimensionality (N, K).
+ */
+template <bool Jacobian, typename Mat, require_not_std_vector_t<Mat>* = nullptr>
+inline auto simplex_row_constrain(const Mat& y, return_type_t<Mat>& lp) {
+  if (Jacobian) {
+    return simplex_row_constrain(y, lp);
+  } else {
+    return simplex_row_constrain(y);
+  }
+}
+
+/**
+ * Return the simplex corresponding to the specified free vector. 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 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 dynamic columns
+ * @param[in] y free vector with matrices of size (N, K - 1)
+ * @param[in, out] lp log density accumulator
+ * @return vector of matrices with simplex rows of dimensionality (N, K)
+ */
+template <bool Jacobian, typename T, require_std_vector_t<T>* = nullptr>
+inline auto simplex_row_constrain(const T& y, return_type_t<T>& lp) {
+  return apply_vector_unary<T>::apply(
+      y, [&lp](auto&& v) { return simplex_row_constrain<Jacobian>(v, lp); });
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/rev/core/var.hpp

@@ -371,10 +371,23 @@ class var_value<T, internal::require_matrix_var_value<T>> {
    * @param other the value to assign
    * @return this
    */
-  template <typename S, require_assignable_t<value_type, S>* = nullptr,
+  template <typename S,
+            require_assignable_t<vari_type,
+                                 typename var_value<S>::vari_type>* = nullptr,
             require_all_plain_type_t<T, S>* = nullptr>
   var_value(const var_value<S>& other) : vi_(other.vi_) {}
 
+  template <typename S,
+            require_not_assignable_t<
+                vari_type, typename var_value<S>::vari_type>* = nullptr,
+            require_constructible_t<vari_type, S>* = nullptr,
+            require_all_plain_type_t<T, S>* = nullptr>
+  var_value(const var_value<S>& other) : vi_(new vari_type(other.vi_->val_)) {
+    reverse_pass_callback([this_vi = this->vi_, other_vi = other.vi_]() {
+      other_vi->adj_ += this_vi->adj_;
+    });
+  }
+
   /**
    * Construct a `var_value` with a plain type
    *  from another `var_value` containing an expression.

stan/math/rev/fun.hpp

@@ -160,7 +160,9 @@
 #include <stan/math/rev/fun/rows_dot_product.hpp>
 #include <stan/math/rev/fun/rows_dot_self.hpp>
 #include <stan/math/rev/fun/sd.hpp>
+#include <stan/math/rev/fun/simplex_column_constrain.hpp>
 #include <stan/math/rev/fun/simplex_constrain.hpp>
+#include <stan/math/rev/fun/simplex_row_constrain.hpp>
 #include <stan/math/rev/fun/sin.hpp>
 #include <stan/math/rev/fun/singular_values.hpp>
 #include <stan/math/rev/fun/svd.hpp>