Prim implementation and test

WardBrian · WardBrian · commit fb463df6f578 · 2025-03-14T10:03:55.000-04:00
diff --git a/stan/math/prim/constraint/sum_to_zero_constrain.hpp b/stan/math/prim/constraint/sum_to_zero_constrain.hpp
@@ -61,47 +61,82 @@ inline plain_type_t<Vec> sum_to_zero_constrain(const Vec& y) {
 }
 
 /**
- * Return a vector with sum zero corresponding to the specified
- * free vector.
+ * Return a matrix that sums to zero over both the rows
+ * and columns corresponding to the free matrix x.
  *
- * The sum-to-zero transform is defined using a modified version of the
- * the inverse of the isometric log ratio transform (ILR).
- * See:
- * Egozcue, Juan Jose; Pawlowsky-Glahn, Vera; Mateu-Figueras, Gloria;
- * Barcelo-Vidal, Carles (2003), "Isometric logratio transformations for
- * compositional data analysis", Mathematical Geology, 35 (3): 279–300,
- * doi:10.1023/A:1023818214614, S2CID 122844634
+ * This is a linear transform, with no Jacobian.
  *
- * This implementation is closer to the description of the same using "pivot
- * coordinates" in
- * Filzmoser, P., Hron, K., Templ, M. (2018). Geometrical Properties of
- * Compositional Data. In: Applied Compositional Data Analysis. Springer Series
- * in Statistics. Springer, Cham. https://doi.org/10.1007/978-3-319-96422-5_3
+ * @tparam Mat type of the matrix
+ * @param x Free matrix input of dimensionality (N - 1, M - 1).
+ * @return Zero-sum matrix of dimensionality (N, M).
+ */
+template <typename Mat, require_eigen_matrix_dynamic_t<Mat>* = nullptr,
+          require_not_st_var<Mat>* = nullptr>
+inline plain_type_t<Mat> sum_to_zero_constrain(const Mat& x) {
+  const auto N = x.rows();
+  const auto M = x.cols();
+  const auto s = std::max(N, M);
+
+  plain_type_t<Mat> Z = Eigen::MatrixXd::Zero(N + 1, M + 1);
+  if (unlikely(N == 0 || M == 0)) {
+    return Z;
+  }
+  auto&& x_ref = to_ref(x);
+
+  Eigen::VectorXd beta = Eigen::VectorXd::Zero(N);
+
+  for (int j = M - 1; j >= 0; --j) {
+    value_type_t<Mat> ax_previous(0);
+
+    double a_j = 1.0 / std::sqrt((j + 1.0) * (j + 2.0));
+    double b_j = (j + 1.0) * a_j;
+
+    for (int i = N - 1; i >= 0; --i) {
+      double a_i = 1.0 / std::sqrt((i + 1.0) * (i + 2.0));
+      double b_i = (i + 1.0) * a_i;
+
+      auto b_i_x = b_i * x_ref(i, j) - ax_previous;
+
+      Z(i, j) = (b_j * b_i_x) - beta(i);
+      beta(i) += a_j * b_i_x;
+
+      Z(N, j) -= Z(i, j);
+      Z(i, M) -= Z(i, j);
+
+      ax_previous += a_i * x_ref(i, j);
+    }
+    Z(N, M) -= Z(N, j);
+  }
+
+  return Z;
+}
+
+/**
+ * Return a vector or matrix with sum zero corresponding to the specified
+ * free input.
  *
  * This is a linear transform, with no Jacobian.
  *
- * @tparam Vec type of the vector
+ * @tparam T type of the input, either a vector or a matrix
  * @tparam Lp unused
- * @param y Free vector input of dimensionality K - 1.
+ * @param y Free vector or matrix
  * @param lp unused
- * @return Zero-sum vector of dimensionality K.
+ * @return Zero-sum vector or matrix which is one larger in each dimension
  */
-template <typename Vec, typename Lp, require_eigen_col_vector_t<Vec>* = nullptr,
-          require_not_st_var<Vec>* = nullptr>
-inline plain_type_t<Vec> sum_to_zero_constrain(const Vec& y, Lp& lp) {
+template <typename T, typename Lp, require_not_st_var<T>* = nullptr>
+inline plain_type_t<T> sum_to_zero_constrain(const T& y, Lp& lp) {
   return sum_to_zero_constrain(y);
 }
 
 /**
- * Return a vector with sum zero corresponding to the specified
- * free vector.
+ * Return a vector or matrix with sum zero corresponding to the specified
+ * free input.
  * This overload handles looping over the elements of a standard vector.
  *
- * @tparam Vec 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[in] y free vector
- * @return Zero-sum vectors of dimensionality one greater than `y`
+ * @tparam T A standard vector with inner type that is either a vector or a
+ * matrix
+ * @param[in] y free vector or matrix
+ * @return Zero-sum vectors or matrices which are one larger in each dimension
  */
 template <typename T, require_std_vector_t<T>* = nullptr>
 inline auto sum_to_zero_constrain(const T& y) {
@@ -114,13 +149,12 @@ inline auto sum_to_zero_constrain(const T& y) {
  * free vector.
  * This overload handles looping over the elements of a standard vector.
  *
- * @tparam Vec 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 standard vector with inner type that is either a vector or a
+ * matrix
  * @tparam Lp unused
- * @param[in] y free vector
+ * @param[in] y free vector or matrix
  * @param[in, out] lp unused
- * @return Zero-sum vectors of dimensionality one greater than `y`
+ * @return Zero-sum vectors or matrices which are one larger in each dimension
  */
 template <typename T, typename Lp, require_std_vector_t<T>* = nullptr,
           require_convertible_t<return_type_t<T>, Lp>* = nullptr>
@@ -130,36 +164,19 @@ inline auto sum_to_zero_constrain(const T& y, Lp& lp) {
 }
 
 /**
- * Return a vector with sum zero corresponding to the specified
- * free vector.
- *
- * The sum-to-zero transform is defined using a modified version of
- * the inverse of the isometric log ratio transform (ILR).
- * See:
- * Egozcue, Juan Jose; Pawlowsky-Glahn, Vera; Mateu-Figueras, Gloria;
- * Barcelo-Vidal, Carles (2003), "Isometric logratio transformations for
- * compositional data analysis", Mathematical Geology, 35 (3): 279–300,
- * doi:10.1023/A:1023818214614, S2CID 122844634
- *
- * This implementation is closer to the description of the same using "pivot
- * coordinates" in
- * Filzmoser, P., Hron, K., Templ, M. (2018). Geometrical Properties of
- * Compositional Data. In: Applied Compositional Data Analysis. Springer Series
- * in Statistics. Springer, Cham. https://doi.org/10.1007/978-3-319-96422-5_3
- *
+ * Return a vector or matrix with sum zero corresponding to the specified
+ * free input.
  * This is a linear transform, with no Jacobian.
  *
  * @tparam Jacobian unused
- * @tparam Vec 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 T type of the input
  * @tparam Lp unused
- * @param[in] y free vector
+ * @param[in] y free vector or matrix
  * @param[in, out] lp unused
- * @return Zero-sum vector of dimensionality one greater than `y`
+ * @return Zero-sum vector or matrix which is one larger in each dimension
  */
-template <bool Jacobian, typename Vec, typename Lp>
-inline plain_type_t<Vec> sum_to_zero_constrain(const Vec& y, Lp& lp) {
+template <bool Jacobian, typename T, typename Lp>
+inline plain_type_t<T> sum_to_zero_constrain(const T& y, Lp& lp) {
   return sum_to_zero_constrain(y);
 }
 
diff --git a/stan/math/prim/constraint/sum_to_zero_free.hpp b/stan/math/prim/constraint/sum_to_zero_free.hpp
@@ -62,10 +62,56 @@ inline plain_type_t<Vec> sum_to_zero_free(const Vec& z) {
 }
 
 /**
- * Overload of `sum_to_zero_free()` to untransform each Eigen vector
+ * Return an unconstrained matrix.
+ *
+ * @tparam Mat a column vector type
+ * @param z Matrix of size (N, M)
+ * @return Free matrix of length (N - 1, M - 1)
+ * @throw std::domain_error if z does not sum to zero
+ */
+template <typename Mat, require_eigen_matrix_dynamic_t<Mat>* = nullptr>
+inline plain_type_t<Mat> sum_to_zero_free(const Mat& z) {
+  const auto& z_ref = to_ref(z);
+  check_sum_to_zero("stan::math::sum_to_zero_free", "sum_to_zero variable",
+                    z_ref);
+
+  const auto N = z_ref.rows() - 1;
+  const auto M = z_ref.cols() - 1;
+  const auto s = std::max(N, M);
+
+  plain_type_t<Mat> x = Eigen::MatrixXd::Zero(N, M);
+  if (unlikely(N == 0 || M == 0)) {
+    return x;
+  }
+
+  Eigen::Matrix<value_type_t<Mat>, -1, 1> beta = Eigen::VectorXd::Zero(N);
+
+  for (int j = M - 1; j >= 0; --j) {
+    value_type_t<Mat> ax_previous(0);
+
+    double a_j = 1.0 / std::sqrt((j + 1.0) * (j + 2.0));
+    double b_j = (j + 1.0) * a_j;
+
+    for (int i = N - 1; i >= 0; --i) {
+      double a_i = 1.0 / std::sqrt((i + 1.0) * (i + 2.0));
+      double b_i = (i + 1.0) * a_i;
+
+      auto alpha_plus_beta = z_ref(i, j) + beta(i);
+
+      x(i, j) = (alpha_plus_beta + b_j * ax_previous) / (b_j * b_i);
+      beta(i) += a_j * (b_i * x(i, j) - ax_previous);
+      ax_previous += a_i * x(i, j);
+    }
+  }
+
+  return x;
+}
+
+/**
+ * Overload of `sum_to_zero_free()` to untransform each Eigen type
  * in a standard vector.
  * @tparam T A standard vector with with a `value_type` which inherits from
- *  `Eigen::MatrixBase` with compile time rows or columns equal to 1.
+ *  `Eigen::MatrixBase`
  * @param z The standard vector to untransform.
  */
 template <typename T, require_std_vector_t<T>* = nullptr>
diff --git a/test/unit/math/prim/constraint/sum_to_zero_transform_test.cpp b/test/unit/math/prim/constraint/sum_to_zero_transform_test.cpp
@@ -51,7 +51,7 @@ TEST(prob_transform, sum_to_zero_match) {
 
   EXPECT_EQ(4, y.size());
   EXPECT_EQ(4, y2.size());
-  for (int i = 0; i < x.size(); ++i)
+  for (int i = 0; i < y.size(); ++i)
     EXPECT_FLOAT_EQ(y[i], y2[i]);
 }
 
@@ -62,3 +62,62 @@ TEST(prob_transform, sum_to_zero_f_exception) {
   y << 0.5, -0.55;
   EXPECT_THROW(stan::math::sum_to_zero_free(y), std::domain_error);
 }
+
+TEST(prob_transform, sum_to_zero_mat_rt0) {
+  double lp = 0;
+  auto x = Eigen::MatrixXd::Zero(4, 3);
+  std::vector<Eigen::MatrixXd> x_vec{x, x, x};
+  std::vector<Eigen::MatrixXd> y_vec
+      = stan::math::sum_to_zero_constrain<false>(x_vec, lp);
+  EXPECT_NO_THROW(stan::math::check_sum_to_zero("checkSumToZero", "y", y_vec));
+
+  for (auto&& y_i : y_vec) {
+    EXPECT_MATRIX_FLOAT_EQ(Eigen::MatrixXd::Zero(5, 4), y_i);
+  }
+  std::vector<Eigen::MatrixXd> xrt = stan::math::sum_to_zero_free(y_vec);
+  EXPECT_EQ(x.rows() + 1, y_vec[2].rows());
+  EXPECT_EQ(x.cols() + 1, y_vec[2].cols());
+  for (auto&& x_i : xrt) {
+    EXPECT_MATRIX_NEAR(x, x_i, 1E-10);
+  }
+}
+
+TEST(prob_transform, sum_to_zero_mat_rt) {
+  Eigen::MatrixXd x(3, 4);
+  x << 1.0, -1.0, 2.0, 1.1, -1.1, 2.1, 1.2, -1.2, 2.2, 1.3, -1.3, 2.3;
+
+  Eigen::MatrixXd y = stan::math::sum_to_zero_constrain(x);
+
+  EXPECT_EQ(x.rows() + 1, y.rows());
+  EXPECT_EQ(x.cols() + 1, y.cols());
+  EXPECT_FLOAT_EQ(y.sum(), 0.0);
+  EXPECT_MATRIX_NEAR(y.rowwise().sum(), Eigen::VectorXd::Zero(4), 1E-15);
+  EXPECT_MATRIX_NEAR(y.colwise().sum(), Eigen::RowVectorXd::Zero(5), 1E-15);
+  EXPECT_NO_THROW(stan::math::check_sum_to_zero("checkSumToZero", "y", y));
+
+  Eigen::MatrixXd xrt = stan::math::sum_to_zero_free(y);
+
+  EXPECT_EQ(x.size(), xrt.size());
+  EXPECT_MATRIX_FLOAT_EQ(x, xrt);
+}
+
+TEST(prob_transform, sum_to_zero_mat_match) {
+  Eigen::MatrixXd x(3, 4);
+  x << 1.0, -1.0, 2.0, 1.1, -1.1, 2.1, 1.2, -1.2, 2.2, 1.3, -1.3, 2.3;
+
+  double lp = 0;
+  Eigen::MatrixXd y = stan::math::sum_to_zero_constrain(x);
+  Eigen::MatrixXd y2 = stan::math::sum_to_zero_constrain(x, lp);
+
+  EXPECT_EQ(y.size(), y2.size());
+  EXPECT_MATRIX_FLOAT_EQ(y, y2);
+  EXPECT_EQ(lp, 0);
+}
+
+TEST(prob_transform, sum_to_zero_mat_f_exception) {
+  Eigen::MatrixXd x(3, 4);
+  x << 1.0, -1.0, 2.0, 1.1, -1.1, 2.1, 1.2, -1.2, 2.2, 1.3, -1.3, 2.3;
+  auto y = stan::math::sum_to_zero_constrain(x);
+  y(0, 0) += 0.1;
+  EXPECT_THROW(stan::math::sum_to_zero_free(y), std::domain_error);
+}
