Add sum_to_zero constraint, free, and check

stan/math/prim/constraint.hpp

@@ -37,6 +37,8 @@
 #include <stan/math/prim/constraint/stochastic_column_free.hpp>
 #include <stan/math/prim/constraint/stochastic_row_constrain.hpp>
 #include <stan/math/prim/constraint/stochastic_row_free.hpp>
+#include <stan/math/prim/constraint/sum_to_zero_constrain.hpp>
+#include <stan/math/prim/constraint/sum_to_zero_free.hpp>
 #include <stan/math/prim/constraint/ub_constrain.hpp>
 #include <stan/math/prim/constraint/ub_free.hpp>
 #include <stan/math/prim/constraint/unit_vector_constrain.hpp>

stan/math/prim/constraint/sum_to_zero_constrain.hpp

@@ -0,0 +1,111 @@
+#ifndef STAN_MATH_PRIM_CONSTRAINT_SUM_TO_ZERO_CONSTRAIN_HPP
+#define STAN_MATH_PRIM_CONSTRAINT_SUM_TO_ZERO_CONSTRAIN_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/fun/sum.hpp>
+#include <stan/math/prim/functor/apply_vector_unary.hpp>
+#include <cmath>
+
+namespace stan {
+namespace math {
+
+/**
+ * Return a vector with sum zero corresponding to the specified
+ * free vector.
+ *
+ * The sum-to-zero transform is defined such that the first K-1
+ * elements are unconstrained and the last element is the negative
+ * sum of those elements.
+ *
+ * @tparam Vec type of the vector
+ * @param y Free vector input of dimensionality K - 1.
+ * @return Zero-sum vector of dimensionality K.
+ */
+template <typename Vec, 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) {
+  const auto Km1 = y.size();
+  if (unlikely(Km1 == 0)) {
+    return plain_type_t<Vec>(Eigen::VectorXd{{0}});
+  }
+  plain_type_t<Vec> x(Km1 + 1);
+  // copy the first Km1 elements
+  auto&& y_ref = to_ref(y);
+  x.head(Km1) = y_ref;
+  // set the last element to -sum(y)
+  x.coeffRef(Km1) = -sum(y_ref);
+  return x;
+}
+
+/**
+ * Return a vector with sum zero corresponding to the specified
+ * free vector.
+ *
+ * The sum-to-zero transform is defined such that the first K-1
+ * elements are unconstrained and the last element is the negative
+ * sum of those elements. This is a linear transform, with no
+ * Jacobian.
+ *
+ * @tparam Vec type of the vector
+ * @param y Free vector input of dimensionality K - 1.
+ * @param lp unused
+ * @return Zero-sum vector of dimensionality K.
+ */
+template <typename Vec, 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,
+                                               value_type_t<Vec>& lp) {
+  return sum_to_zero_constrain(y);
+}
+
+/**
+ * Return a vector with sum zero corresponding to the specified
+ * free vector.
+ *
+ * The sum-to-zero transform is defined such that the first K-1
+ * elements are unconstrained and the last element is the negative
+ * sum of those elements. 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
+ * @param[in] y free vector
+ * @param[in, out] lp unused
+ * @return Zero-sum vector of dimensionality one greater than `y`
+ */
+template <bool Jacobian, typename Vec, require_not_std_vector_t<Vec>* = nullptr>
+inline plain_type_t<Vec> sum_to_zero_constrain(const Vec& y,
+                                               return_type_t<Vec>& lp) {
+  return sum_to_zero_constrain(y);
+}
+
+/**
+ * Return a vector with sum zero corresponding to the specified
+ * free vector.
+ *
+ * The sum-to-zero transform is defined such that the first K-1
+ * elements are unconstrained and the last element is the negative
+ * sum of those elements. This is a linear transform, with no
+ * Jacobian.
+ *
+ * @tparam Jacobian unused
+ * @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
+ * @param[in, out] lp unused
+ * @return Zero-sum vectors of dimensionality one greater than `y`
+ */
+template <bool Jacobian, typename T, require_std_vector_t<T>* = nullptr>
+inline auto sum_to_zero_constrain(const T& y, return_type_t<T>& lp) {
+  return apply_vector_unary<T>::apply(
+      y, [](auto&& v) { return sum_to_zero_constrain(v); });
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/prim/constraint/sum_to_zero_free.hpp

@@ -0,0 +1,51 @@
+#ifndef STAN_MATH_PRIM_CONSTRAINT_SUM_TO_ZERO_FREE_HPP
+#define STAN_MATH_PRIM_CONSTRAINT_SUM_TO_ZERO_FREE_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/fun/to_ref.hpp>
+#include <stan/math/prim/functor/apply_vector_unary.hpp>
+#include <cmath>
+
+namespace stan {
+namespace math {
+
+/**
+ * Return an unconstrained vector.
+ *
+ * The sum-to-zero transform is defined such that the first K-1
+ * elements are unconstrained and the last element is the negative
+ * sum of those elements.
+ *
+ * @tparam ColVec a column vector type
+ * @param x Vector of length K.
+ * @return Free vector of length (K-1).
+ * @throw std::domain_error if x 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);
+  check_sum_to_zero("stan::math::sum_to_zero_free", "sum_to_zero variable",
+                    x_ref);
+
+  return x_ref.head(x_ref.size() - 1);
+}
+
+/**
+ * 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.
+ */
+template <typename T, require_std_vector_t<T>* = nullptr>
+auto sum_to_zero_free(const T& x) {
+  return apply_vector_unary<T>::apply(
+      x, [](auto&& v) { return sum_to_zero_free(v); });
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/prim/err.hpp

@@ -41,6 +41,7 @@
 #include <stan/math/prim/err/check_std_vector_index.hpp>
 #include <stan/math/prim/err/check_stochastic_column.hpp>
 #include <stan/math/prim/err/check_stochastic_row.hpp>
+#include <stan/math/prim/err/check_sum_to_zero.hpp>
 #include <stan/math/prim/err/check_symmetric.hpp>
 #include <stan/math/prim/err/check_unit_vector.hpp>
 #include <stan/math/prim/err/check_vector.hpp>

stan/math/prim/err/check_sum_to_zero.hpp

@@ -0,0 +1,72 @@
+#ifndef STAN_MATH_PRIM_ERR_CHECK_SUM_TO_ZERO_HPP
+#define STAN_MATH_PRIM_ERR_CHECK_SUM_TO_ZERO_HPP
+
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/err/constraint_tolerance.hpp>
+#include <stan/math/prim/err/make_iter_name.hpp>
+#include <stan/math/prim/err/throw_domain_error.hpp>
+#include <stan/math/prim/fun/to_ref.hpp>
+#include <stan/math/prim/fun/value_of_rec.hpp>
+#include <sstream>
+#include <string>
+
+namespace stan {
+namespace math {
+
+/**
+ * Throw an exception if the specified vector does not sum to 0.
+ * This function tests that the sum is within the tolerance specified by
+ * `CONSTRAINT_TOLERANCE`.
+ * This function only accepts Eigen vectors, statically
+ * typed vectors, not general matrices with 1 column.
+ * @tparam T A type inheriting from `Eigen::EigenBase`
+ * @param function Function name (for error messages)
+ * @param name Variable name (for error messages)
+ * @param theta Vector to test
+ * @throw `std::invalid_argument` if `theta` is a 0-vector
+ * @throw `std::domain_error` if the vector does not sum to zero
+ */
+template <typename T, require_matrix_t<T>* = nullptr>
+void check_sum_to_zero(const char* function, const char* name, const T& theta) {
+  using std::fabs;
+  // the size-zero case is technically a valid sum-to-zero vector,
+  // but it cannot be unconstrained to anything
+  check_nonzero_size(function, name, theta);
+  auto&& theta_ref = to_ref(value_of_rec(theta));
+  if (unlikely(!(fabs(theta_ref.sum()) <= CONSTRAINT_TOLERANCE))) {
+    [&]() STAN_COLD_PATH {
+      std::stringstream msg;
+      scalar_type_t<T> sum = theta_ref.sum();
+      msg << "does not sum to zero.";
+      msg.precision(10);
+      msg << " sum(" << name << ") = " << sum << ", but should be ";
+      std::string msg_str(msg.str());
+      throw_domain_error(function, name, 0.0, msg_str.c_str());
+    }();
+  }
+}
+
+/**
+ * Throw an exception if any vector in a standard vector does not sum to 0.
+ * This function tests that the sum is within the tolerance specified by
+ * `CONSTRAINT_TOLERANCE`.
+ * @tparam T A standard vector with inner type inheriting from
+ * `Eigen::EigenBase`
+ * @param function Function name (for error messages)
+ * @param name Variable name (for error messages)
+ * @param theta Vector to test.
+ * @throw `std::invalid_argument` if `theta` is a 0-vector
+ * @throw `std::domain_error` if the vector does not sum to zero
+ */
+template <typename T, require_std_vector_t<T>* = nullptr>
+void check_sum_to_zero(const char* function, const char* name, const T& theta) {
+  for (size_t i = 0; i < theta.size(); ++i) {
+    check_sum_to_zero(function, internal::make_iter_name(name, i).c_str(),
+                      theta[i]);
+  }
+}
+
+}  // namespace math
+}  // namespace stan
+#endif

stan/math/rev/constraint.hpp

@@ -17,6 +17,7 @@
 #include <stan/math/rev/constraint/simplex_constrain.hpp>
 #include <stan/math/rev/constraint/stochastic_column_constrain.hpp>
 #include <stan/math/rev/constraint/stochastic_row_constrain.hpp>
+#include <stan/math/rev/constraint/sum_to_zero_constrain.hpp>
 #include <stan/math/rev/constraint/unit_vector_constrain.hpp>
 #include <stan/math/rev/constraint/ub_constrain.hpp>
 

stan/math/rev/constraint/sum_to_zero_constrain.hpp

@@ -0,0 +1,70 @@
+#ifndef STAN_MATH_REV_CONSTRAINT_SUM_TO_ZERO_CONSTRAIN_HPP
+#define STAN_MATH_REV_CONSTRAINT_SUM_TO_ZERO_CONSTRAIN_HPP
+
+#include <stan/math/rev/meta.hpp>
+#include <stan/math/rev/core/reverse_pass_callback.hpp>
+#include <stan/math/rev/core/arena_matrix.hpp>
+#include <stan/math/rev/fun/value_of.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <cmath>
+#include <tuple>
+#include <vector>
+
+namespace stan {
+namespace math {
+
+/**
+ * Return a vector with sum zero corresponding to the specified
+ * free vector.
+ *
+ * The sum-to-zero transform is defined such that the first K-1
+ * elements are unconstrained and the last element is the negative
+ * sum of those elements.
+ *
+ * @tparam T type of the vector
+ * @param y Free vector input of dimensionality K - 1.
+ * @return Zero-sum vector of dimensionality K.
+ */
+template <typename T, require_rev_col_vector_t<T>* = nullptr>
+inline auto sum_to_zero_constrain(const T& y) {
+  using ret_type = plain_type_t<T>;
+
+  const auto N = y.size();
+  if (unlikely(N == 0)) {
+    return arena_t<ret_type>(Eigen::VectorXd{{0}});
+  }
+  Eigen::VectorXd x_val = Eigen::VectorXd::Zero(N + 1);
+  auto arena_y = to_arena(y);
+  double x_sum = -sum(arena_y.val());
+  x_val.head(N) = arena_y.val();
+  x_val(N) = x_sum;
+  arena_t<ret_type> arena_x = x_val;
+  reverse_pass_callback([arena_y, arena_x, x_sum, N]() mutable {
+    arena_y.adj().array() -= arena_x.adj_op()(N);
+    arena_y.adj() += arena_x.adj_op().head(N);
+  });
+  return arena_x;
+}
+
+/**
+ * Return a vector with sum zero corresponding to the specified
+ * free vector.
+ *
+ * The sum-to-zero transform is defined such that the first K-1
+ * elements are unconstrained and the last element is the negative
+ * sum of those elements. This is a linear transform, with no
+ * Jacobian.
+ *
+ * @tparam Vec type of the vector
+ * @param y Free vector input of dimensionality K - 1.
+ * @param lp unused
+ * @return Zero-sum vector of dimensionality K.
+ */
+template <typename T, require_rev_col_vector_t<T>* = nullptr>
+inline auto sum_to_zero_constrain(const T& y, scalar_type_t<T>& lp) {
+  return sum_to_zero_constrain(y);
+}
+
+}  // namespace math
+}  // namespace stan
+#endif

test/unit/math/mix/constraint/sum_to_zero_constrain_test.cpp

@@ -0,0 +1,58 @@
+#include <test/unit/math/test_ad.hpp>
+
+namespace sum_to_zero_constrain_test {
+template <typename T>
+T g1(const T& x) {
+  stan::scalar_type_t<T> lp = 0;
+  return stan::math::sum_to_zero_constrain<false>(x, lp);
+}
+template <typename T>
+T g2(const T& x) {
+  stan::scalar_type_t<T> lp = 0;
+  return stan::math::sum_to_zero_constrain<true>(x, lp);
+}
+template <typename T>
+typename stan::scalar_type<T>::type g3(const T& x) {
+  stan::scalar_type_t<T> lp = 0;
+  stan::math::sum_to_zero_constrain<true>(x, lp);
+  return lp;
+}
+
+template <typename T>
+void expect_sum_to_zero_transform(const T& x) {
+  auto f1 = [](const auto& x) { return g1(x); };
+  auto f2 = [](const auto& x) { return g2(x); };
+  auto f3 = [](const auto& x) { return g3(x); };
+  stan::test::expect_ad(f1, x);
+  stan::test::expect_ad_matvar(f1, x);
+  stan::test::expect_ad(f2, x);
+  stan::test::expect_ad_matvar(f2, x);
+  stan::test::expect_ad(f3, x);
+  stan::test::expect_ad_matvar(f3, x);
+}
+}  // namespace sum_to_zero_constrain_test
+
+TEST(MathMixMatFun, sum_to_zeroTransform) {
+  Eigen::VectorXd v0(0);
+  sum_to_zero_constrain_test::expect_sum_to_zero_transform(v0);
+
+  Eigen::VectorXd v1(1);
+  v1 << 1;
+  sum_to_zero_constrain_test::expect_sum_to_zero_transform(v1);
+
+  Eigen::VectorXd v2(2);
+  v2 << 3, -1;
+  sum_to_zero_constrain_test::expect_sum_to_zero_transform(v2);
+
+  Eigen::VectorXd v3(3);
+  v3 << 2, 3, -1;
+  sum_to_zero_constrain_test::expect_sum_to_zero_transform(v3);
+
+  Eigen::VectorXd v4(4);
+  v4 << 2, -1, 0, -1.1;
+  sum_to_zero_constrain_test::expect_sum_to_zero_transform(v4);
+
+  Eigen::VectorXd v5(5);
+  v5 << 1, -3, 2, 0, -1;
+  sum_to_zero_constrain_test::expect_sum_to_zero_transform(v5);
+}