Merge pull request #3342 from stan-dev/serde/sum-to-zero-matrix

Author: WardBrian

src/stan/io/deserializer.hpp

@@ -583,20 +583,19 @@ class deserializer {
 
   /**
    * Return the next zero sum vector of the specified size (using one fewer
-   * unconstrained scalars), incrementing the specified reference with the
-   * log absolute Jacobian determinant (no adjustment, in this case).
+   * unconstrained scalars). This particular transformation is linear, not
+   * requiring a Jacobian adjustment.
    *
-   * <p>See <code>stan::math::sum_to_zero_constrain(Eigen::Matrix,T&)</code>.
+   * <p>See <code>stan::math::sum_to_zero_constrain(Eigen::Matrix, T&)</code>.
    *
    * @tparam Ret The type to return.
    * @tparam Jacobian Whether to increment the log of the absolute Jacobian
    * determinant of the transform.
    * @tparam LP Type of log probability.
-   * @tparam Sizes A parameter pack of integral types.
-   * @param lp The reference to the variable holding the log
+   * @param lp (Unused) The reference to the variable holding the log density.
    * @param size Number of cells in zero sum vector to generate.
    * @return The next zero sum of the specified size.
-   * @throws std::invalid_argument if number of dimensions (`k`) is zero
+   * @throws std::invalid_argument if `size` is zero
    */
   template <typename Ret, bool Jacobian, typename LP,
             require_not_std_vector_t<Ret>* = nullptr>
@@ -606,10 +605,37 @@ class deserializer {
         this->read<Ret>(size - 1), lp);
   }
 
+  /**
+   * Return the next zero sum matrix of the specified dimensions (requires
+   * (N - 1) * (M - 1) unconstrained scalars).
+   * This particular transformation is linear, not requiring a Jacobian
+   * adjustment.
+   *
+   * <p>See <code>stan::math::sum_to_zero_constrain(Eigen::Matrix, T&)</code>.
+   *
+   * @tparam Ret The type to return.
+   * @tparam Jacobian Whether to increment the log of the absolute Jacobian
+   * determinant of the transform.
+   * @tparam LP Type of log probability.
+   * @param lp (Unused) The reference to the variable holding the log density.
+   * @param N Number of rows in zero sum matrix to generate.
+   * @param M Number of columns in zero sum matrix to generate.
+   * @return The next zero sum of the specified size.
+   * @throws std::invalid_argument either `N` or `M` is zero
+   */
+  template <typename Ret, bool Jacobian, typename LP,
+            require_matrix_t<Ret>* = nullptr>
+  inline auto read_constrain_sum_to_zero(LP& lp, size_t N, size_t M) {
+    stan::math::check_positive("read_sum_to_zero", "N", N);
+    stan::math::check_positive("read_sum_to_zero", "M", M);
+    return stan::math::sum_to_zero_constrain<Jacobian>(
+        this->read<conditional_var_val_t<Ret, matrix_t>>(N - 1, M - 1), lp);
+  }
+
   /**
    * Return the next zero sum vector of the specified size (using one fewer
-   * unconstrained scalars), incrementing the specified reference with the
-   * log absolute Jacobian determinant (no adjustment, in this case).
+   * unconstrained scalars). This particular transformation is linear, not
+   * requiring a Jacobian adjustment.
    *
    * <p>See <code>stan::math::sum_to_zero_constrain(Eigen::Matrix,T&)</code>.
    *
@@ -1219,18 +1245,32 @@ class deserializer {
   }
 
   /**
-   * Read a serialized sum_to_zero vector and unconstrain it
+   * Read a serialized sum_to_zero vector and unconstrain it.
    *
    * @tparam Ret Type of output
-   * @return Unconstrained vector
+   * @param size Vector size
+   * @return Unconstrained vector of length (size - 1)
    */
   template <typename Ret, require_not_std_vector_t<Ret>* = nullptr>
   inline auto read_free_sum_to_zero(size_t size) {
     return stan::math::sum_to_zero_free(this->read<Ret>(size));
   }
 
   /**
-   * Read serialized zero-sum vectors and unconstrain them
+   * Read a serialized sum_to_zero matrix and unconstrain it.
+   *
+   * @tparam Ret Type of output
+   * @param N Rows of matrix
+   * @param M Cols of matrix
+   * @return Unconstrained matrix of size (N-1) x (M-1)
+   */
+  template <typename Ret, require_matrix_t<Ret>* = nullptr>
+  inline auto read_free_sum_to_zero(size_t N, size_t M) {
+    return stan::math::sum_to_zero_free(this->read<Ret>(N, M));
+  }
+
+  /**
+   * Read serialized zero-sum vectors and unconstrain them.
    *
    * @tparam Ret Type of output
    * @tparam Sizes Types of dimensions of output

src/test/unit/io/deserializer_free_test.cpp

@@ -265,13 +265,24 @@ struct SumToZeroConstrain {
 
 TEST(deserializer_vector, read_free_sum_to_zero) {
   using stan::test::deserializer_test;
+  deserializer_test<Eigen::VectorXd, SumToZeroConstrain>(std::make_tuple(1));
   deserializer_test<Eigen::VectorXd, SumToZeroConstrain>(std::make_tuple(4));
   deserializer_test<std::vector<Eigen::VectorXd>, SumToZeroConstrain>(
       std::make_tuple(2, 4));
   deserializer_test<std::vector<std::vector<Eigen::VectorXd>>,
                     SumToZeroConstrain>(std::make_tuple(3, 2, 4));
 }
 
+TEST(deserializer_matrix, read_free_sum_to_zero) {
+  using stan::test::deserializer_test;
+  deserializer_test<Eigen::MatrixXd, SumToZeroConstrain>(std::make_tuple(1, 1));
+  deserializer_test<Eigen::MatrixXd, SumToZeroConstrain>(std::make_tuple(4, 5));
+  deserializer_test<std::vector<Eigen::MatrixXd>, SumToZeroConstrain>(
+      std::make_tuple(2, 4, 5));
+  deserializer_test<std::vector<std::vector<Eigen::MatrixXd>>,
+                    SumToZeroConstrain>(std::make_tuple(3, 2, 4, 2));
+}
+
 // ordered
 template <typename Ret>
 struct OrderedConstrain {

src/test/unit/io/deserializer_stdvector_test.cpp

@@ -175,6 +175,50 @@ TEST(deserializer_array, sum_to_zero) {
   }
 }
 
+TEST(deserializer_array, sum_to_zero_matrix) {
+  std::vector<int> theta_i;
+  std::vector<double> theta;
+  for (size_t i = 0; i < 100U; ++i)
+    theta.push_back(static_cast<double>(i));
+
+  stan::io::deserializer<double> deserializer1(theta, theta_i);
+  stan::io::deserializer<double> deserializer2(theta, theta_i);
+
+  // no jac
+  {
+    double lp_ref = 0.0;
+    double lp = 0.0;
+    auto y
+        = deserializer1
+              .read_constrain_sum_to_zero<std::vector<Eigen::MatrixXd>, false>(
+                  lp, 4, 3, 2);
+    for (size_t i = 0; i < 4; ++i) {
+      stan::test::expect_near_rel(
+          "test_std_vector_deserializer", y[i],
+          deserializer2.read_constrain_sum_to_zero<Eigen::MatrixXd, false>(
+              lp_ref, 3, 2));
+    }
+    EXPECT_FLOAT_EQ(lp_ref, lp);
+  }
+
+  // jac
+  {
+    double lp_ref = 0.0;
+    double lp = 0.0;
+    auto y
+        = deserializer1
+              .read_constrain_sum_to_zero<std::vector<Eigen::MatrixXd>, true>(
+                  lp, 4, 3, 2);
+    for (size_t i = 0; i < 4; ++i) {
+      stan::test::expect_near_rel(
+          "test_std_vector_deserializer", y[i],
+          deserializer2.read_constrain_sum_to_zero<Eigen::MatrixXd, true>(
+              lp_ref, 3, 2));
+    }
+    EXPECT_FLOAT_EQ(lp_ref, lp);
+  }
+}
+
 // ordered
 
 TEST(deserializer_array, ordered) {

src/test/unit/io/deserializer_test.cpp

@@ -714,6 +714,63 @@ TEST(deserializer_vector, sum_to_zero_jacobian) {
   EXPECT_FLOAT_EQ(lp_ref, lp);
 }
 
+TEST(deserializer_matrix, sum_to_zero_constrain) {
+  std::vector<int> theta_i;
+  std::vector<double> theta;
+  theta.push_back(3.0);
+  theta.push_back(-1.0);
+  theta.push_back(-2.0);
+  theta.push_back(0.0);
+  stan::io::deserializer<double> deserializer(theta, theta_i);
+  double lp = 0.0;
+  Eigen::MatrixXd reference
+      = stan::math::sum_to_zero_constrain(stan::math::to_matrix(theta, 2, 2));
+  Eigen::MatrixXd phi(
+      deserializer.read_constrain_sum_to_zero<Eigen::MatrixXd, false>(lp, 3,
+                                                                      3));
+  for (int n = 0; n < phi.size(); ++n) {
+    EXPECT_FLOAT_EQ(reference(n), phi(n));
+  }
+}
+
+TEST(deserializer_matrix, sum_to_zero_constrain_empty) {
+  std::vector<int> theta_i;
+  std::vector<double> theta;
+
+  stan::io::deserializer<double> deserializer(theta, theta_i);
+  double lp = 0.0;
+
+  Eigen::MatrixXd phi(
+      deserializer.read_constrain_sum_to_zero<Eigen::MatrixXd, false>(lp, 1,
+                                                                      1));
+
+  EXPECT_EQ(1, phi.rows());
+  EXPECT_EQ(1, phi.cols());
+  for (int n = 0; n < phi.size(); ++n) {
+    EXPECT_FLOAT_EQ(0.0, phi(n));
+  }
+}
+
+TEST(deserializer_matrix, sum_to_zero_jacobian) {
+  std::vector<int> theta_i;
+  std::vector<double> theta;
+  theta.push_back(3.0);
+  theta.push_back(-1.0);
+  theta.push_back(-2.0);
+  theta.push_back(0.0);
+  stan::io::deserializer<double> deserializer(theta, theta_i);
+  double lp = 0.0;
+  double lp_ref = 0.0;
+  Eigen::MatrixXd reference = stan::math::sum_to_zero_constrain(
+      stan::math::to_matrix(theta, 2, 2), lp_ref);
+  Eigen::MatrixXd phi(
+      deserializer.read_constrain_sum_to_zero<Eigen::MatrixXd, true>(lp, 3, 3));
+  for (int n = 0; n < phi.size(); ++n) {
+    EXPECT_FLOAT_EQ(reference(n), phi(n));
+  }
+  EXPECT_FLOAT_EQ(lp_ref, lp);
+}
+
 // ordered
 
 TEST(deserializer_vector, ordered_constrain) {

src/test/unit/io/deserializer_var_test.cpp

@@ -640,6 +640,50 @@ TEST(deserializer_vector, sum_to_zero_jacobian) {
   }
   EXPECT_FLOAT_EQ(lp_ref.val(), lp.val());
 }
+
+TEST(deserializer_matrix, sum_to_zero_constrain) {
+  std::vector<int> theta_i;
+  std::vector<stan::math::var> theta;
+  theta.push_back(3.0);
+  theta.push_back(-1.0);
+  theta.push_back(-2.0);
+  theta.push_back(0.0);
+  stan::io::deserializer<stan::math::var> deserializer(theta, theta_i);
+  stan::math::var lp = 0;
+  Eigen::Matrix<stan::math::var, Eigen::Dynamic, Eigen::Dynamic> reference
+      = stan::math::sum_to_zero_constrain(stan::math::to_matrix(theta, 2, 2));
+  Eigen::Matrix<stan::math::var, Eigen::Dynamic, Eigen::Dynamic> phi(
+      deserializer.read_constrain_sum_to_zero<
+          Eigen::Matrix<stan::math::var, Eigen::Dynamic, Eigen::Dynamic>,
+          false>(lp, 3, 3));
+  for (int n = 0; n < phi.size(); ++n) {
+    EXPECT_FLOAT_EQ(reference(n).val(), phi(n).val());
+  }
+}
+
+TEST(deserializer_matrix, sum_to_zero_jacobian) {
+  std::vector<int> theta_i;
+  std::vector<stan::math::var> theta;
+  theta.push_back(3.0);
+  theta.push_back(-1.0);
+  theta.push_back(-2.0);
+  theta.push_back(0.0);
+  stan::io::deserializer<stan::math::var> deserializer(theta, theta_i);
+  stan::math::var lp = 0.0;
+  stan::math::var lp_ref = 0.0;
+  Eigen::Matrix<stan::math::var, Eigen::Dynamic, Eigen::Dynamic> reference
+      = stan::math::sum_to_zero_constrain(stan::math::to_matrix(theta, 2, 2),
+                                          lp_ref);
+  Eigen::Matrix<stan::math::var, Eigen::Dynamic, Eigen::Dynamic> phi(
+      deserializer.read_constrain_sum_to_zero<
+          Eigen::Matrix<stan::math::var, Eigen::Dynamic, Eigen::Dynamic>, true>(
+          lp, 3, 3));
+  for (int n = 0; n < phi.size(); ++n) {
+    EXPECT_FLOAT_EQ(reference(n).val(), phi(n).val());
+  }
+  EXPECT_FLOAT_EQ(lp_ref.val(), lp.val());
+}
+
 // ordered
 
 TEST(deserializer_vector, ordered_constrain) {

src/test/unit/io/deserializer_varmat_test.cpp

@@ -363,6 +363,42 @@ TEST(deserializer, read_constrain_sum_to_zero_jacobian) {
   EXPECT_FLOAT_EQ(lp_ref.val(), lp.val());
 }
 
+TEST(deserializer_var_matrix, read_constrain_sum_to_zero_constrain) {
+  std::vector<int> theta_i;
+  std::vector<stan::math::var> theta;
+  theta.push_back(-2.0);
+  theta.push_back(3.0);
+  theta.push_back(-1.0);
+  theta.push_back(0.0);
+  stan::io::deserializer<stan::math::var> deserializer(theta, theta_i);
+  stan::math::var lp = 0.0;
+  auto reference
+      = stan::math::sum_to_zero_constrain(stan::math::to_matrix(theta, 2, 2));
+  auto y
+      = deserializer.read_constrain_sum_to_zero<var_matrix_t, false>(lp, 3, 3);
+  EXPECT_TRUE((std::is_same<var_matrix_t, decltype(y)>::value));
+  stan::test::expect_near_rel("deserializer tests", reference.val(), y.val());
+}
+
+TEST(deserializer_var_matrix, read_constrain_sum_to_zero_jacobian) {
+  std::vector<int> theta_i;
+  std::vector<stan::math::var> theta;
+  theta.push_back(-2.0);
+  theta.push_back(3.0);
+  theta.push_back(-1.0);
+  theta.push_back(0.0);
+  stan::io::deserializer<stan::math::var> deserializer(theta, theta_i);
+  stan::math::var lp_ref = 0.0;
+  stan::math::var lp = 0.0;
+  auto reference = stan::math::sum_to_zero_constrain(
+      stan::math::to_matrix(theta, 2, 2), lp_ref);
+  auto y
+      = deserializer.read_constrain_sum_to_zero<var_matrix_t, false>(lp, 3, 3);
+  EXPECT_TRUE((std::is_same<var_matrix_t, decltype(y)>::value));
+  stan::test::expect_near_rel("deserializer tests", reference.val(), y.val());
+  EXPECT_FLOAT_EQ(lp_ref.val(), lp.val());
+}
+
 // ordered
 
 TEST(deserializer, read_constrain_ordered_constrain) {

src/test/unit/io/serializer_test.cpp

@@ -522,6 +522,15 @@ TEST(serializer_vectorized, write_free_sum_to_zero) {
                   SumToZeroConstrain>(std::make_tuple(3, 2, 4));
 }
 
+TEST(serializer_vectorized, write_free_sum_to_zero_mat) {
+  using stan::test::serializer_test;
+  serializer_test<Eigen::MatrixXd, SumToZeroConstrain>(std::make_tuple(4, 5));
+  serializer_test<std::vector<Eigen::MatrixXd>, SumToZeroConstrain>(
+      std::make_tuple(2, 4, 3));
+  serializer_test<std::vector<std::vector<Eigen::MatrixXd>>,
+                  SumToZeroConstrain>(std::make_tuple(3, 2, 4, 2));
+}
+
 // ordered
 
 template <typename Ret>
@@ -692,7 +701,7 @@ struct StochasticCol {
     };
   }
 };
-TEST(deserializer_vector, read_stochastic_column_matrix) {
+TEST(serializer_vectorized, read_stochastic_column_matrix) {
   using stan::test::serializer_test;
   serializer_test<Eigen::MatrixXd, StochasticCol>(std::make_tuple(3, 3));
   serializer_test<std::vector<Eigen::MatrixXd>, StochasticCol>(
@@ -716,7 +725,7 @@ struct StochasticRow {
     };
   }
 };
-TEST(deserializer_vector, read_stochastic_row_matrix) {
+TEST(serializer_vectorized, read_stochastic_row_matrix) {
   using stan::test::serializer_test;
   serializer_test<Eigen::MatrixXd, StochasticRow>(std::make_tuple(3, 3));
   serializer_test<std::vector<Eigen::MatrixXd>, StochasticRow>(