Free impl

Author: WardBrian

stan/math/prim/constraint/sum_to_zero_free.hpp

@@ -19,30 +19,47 @@ namespace math {
  * sum of those elements.
  *
  * @tparam ColVec a column vector type
- * @param x Vector of length K.
+ * @param z Vector of length K.
  * @return Free vector of length (K-1).
- * @throw std::domain_error if x does not sum to zero
+ * @throw std::domain_error if z does not sum to zero
  */
 template <typename Vec, require_eigen_vector_t<Vec>* = nullptr>
-inline plain_type_t<Vec> sum_to_zero_free(const Vec& x) {
-  const auto& x_ref = to_ref(x);
+inline plain_type_t<Vec> sum_to_zero_free(const Vec& z) {
+  const auto& z_ref = to_ref(z);
   check_sum_to_zero("stan::math::sum_to_zero_free", "sum_to_zero variable",
-                    x_ref);
+                    z_ref);
 
-  return x_ref.head(x_ref.size() - 1);
+  const auto N = z.size() - 1;
+
+  plain_type_t<Vec> y = Eigen::VectorXd::Zero(N);
+  if (unlikely(N == 0)) {
+    return y;
+  }
+
+  y.coeffRef(N - 1) = -z_ref(N) * sqrt(N * (N + 1)) / N;
+  typename plain_type_t<Vec>::Scalar total(0);
+
+  for (int i = N - 2; i >= 0; --i) {
+    double n = i + 1;
+    auto w = y(i + 1) / sqrt((n + 1) * (n + 2));
+    total += w;
+    y.coeffRef(i) = (total - z_ref(i + 1)) * sqrt(n * (n + 1)) / n;
+  }
+
+  return y;
 }
 
 /**
  * Overload of `sum_to_zero_free()` to untransform each Eigen vector
  * 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.
- * @param x The standard vector to untransform.
+ * @param z The standard vector to untransform.
  */
 template <typename T, require_std_vector_t<T>* = nullptr>
-auto sum_to_zero_free(const T& x) {
+auto sum_to_zero_free(const T& z) {
   return apply_vector_unary<T>::apply(
-      x, [](auto&& v) { return sum_to_zero_free(v); });
+      z, [](auto&& v) { return sum_to_zero_free(v); });
 }
 
 }  // namespace math

stan/math/rev/constraint/sum_to_zero_constrain.hpp

@@ -6,6 +6,7 @@
 #include <stan/math/rev/core/arena_matrix.hpp>
 #include <stan/math/rev/fun/value_of.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/constraint/sum_to_zero_constrain.hpp>
 #include <cmath>
 #include <tuple>
 #include <vector>

test/unit/math/prim/constraint/sum_to_zero_transform_test.cpp

@@ -11,6 +11,8 @@ TEST(prob_transform, sum_to_zero_rt0) {
   std::vector<Matrix<double, Dynamic, 1>> x_vec{x, x, x};
   std::vector<Matrix<double, Dynamic, 1>> 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::VectorXd::Zero(5), y_i);
   }