Deserializer for sum_to_zero matrix

src/stan/io/deserializer.hpp

@@ -592,11 +592,10 @@ class deserializer {
    * @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 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,6 +605,32 @@ 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), incrementing the specified reference with the
+   * log absolute Jacobian determinant (no adjustment, in this case).
+   *
+   * <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 The reference to the variable holding the log
+   * @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
@@ -1229,6 +1254,17 @@ class deserializer {
     return stan::math::sum_to_zero_free(this->read<Ret>(size));
   }
 
+  /**
+   * Read a serialized sum_to_zero matrix and unconstrain it
+   *
+   * @tparam Ret Type of output
+   * @return Unconstrained matrix
+   */
+  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
    *

src/test/unit/io/deserializer_free_test.cpp

@@ -272,6 +272,15 @@ TEST(deserializer_vector, read_free_sum_to_zero) {
                     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(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,49 @@ 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 i = 0; i < phi.rows(); ++i) {
+    for (int j = 0; j < phi.cols(); ++j) {
+      EXPECT_FLOAT_EQ(reference(i, j), phi(i, j));
+    }
+  }
+}
+
+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 i = 0; i < phi.rows(); ++i) {
+    for (int j = 0; j < phi.cols(); ++j) {
+      EXPECT_FLOAT_EQ(reference(i, j), phi(i, j));
+    }
+  }
+  EXPECT_FLOAT_EQ(lp_ref, lp);
+}
+
 // ordered
 
 TEST(deserializer_vector, ordered_constrain) {

src/test/unit/io/deserializer_var_test.cpp

@@ -640,6 +640,54 @@ 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 i = 0; i < phi.rows(); ++i) {
+    for (int j = 0; j < phi.cols(); ++j) {
+      EXPECT_FLOAT_EQ(reference(i, j).val(), phi(i, j).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 i = 0; i < phi.rows(); ++i) {
+    for (int j = 0; j < phi.cols(); ++j) {
+      EXPECT_FLOAT_EQ(reference(i, j).val(), phi(i, j).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>(