Adding json serialization to random effects container

Author: andrewherren

debug/api_debug.cpp

@@ -280,7 +280,7 @@ void RunAPI() {
   }
 
   // Run the MCMC sampler
-  int num_mcmc_samples = 100;
+  int num_mcmc_samples = 10;
   for (int i = num_gfr_samples; i < num_gfr_samples + num_mcmc_samples; i++) {
     if (i == 0) {
       global_variance = global_variance_init;
@@ -313,17 +313,16 @@ void RunAPI() {
   std::vector<double> rfx_predictions(n*num_samples);
   rfx_container.Predict(rfx_dataset, label_mapper, rfx_predictions);
 
-  // // Write model to a file
-  // std::string filename = "model.json";
-  // forest_samples.SaveToJsonFile(filename);
+  // Write model to a file
+  std::string filename = "model.json";
+  forest_samples.SaveToJsonFile(filename);
 
-  // // Read and parse json from file
-  // ForestContainer forest_samples_parsed = ForestContainer(num_trees, output_dimension, is_leaf_constant);
-  // forest_samples_parsed.LoadFromJsonFile(filename);
+  // Read and parse json from file
+  ForestContainer forest_samples_parsed = ForestContainer(num_trees, output_dimension, is_leaf_constant);
+  forest_samples_parsed.LoadFromJsonFile(filename);
 
-  // // Make sure we can predict from both the original and parsed forest containers
-  // std::vector<double> pred_orig = forest_samples.Predict(dataset);
-  // std::vector<double> pred_parsed = forest_samples_parsed.Predict(dataset);
+  // Make sure we can predict from both the original (above) and parsed forest containers
+  std::vector<double> pred_parsed = forest_samples_parsed.Predict(dataset);
 }
 
 } // namespace StochTree

include/stochtree/random_effects.h

@@ -14,6 +14,7 @@
 #include <stochtree/normal_sampler.h>
 #include <stochtree/partition_tracker.h>
 #include <stochtree/prior.h>
+#include <nlohmann/json.hpp>
 #include <Eigen/Dense>
 
 #include <cmath>
@@ -247,6 +248,11 @@ class RandomEffectsContainer {
     num_groups_ = num_groups;
     num_samples_ = 0;
   }
+  RandomEffectsContainer() {
+    num_components_ = 0;
+    num_groups_ = 0;
+    num_samples_ = 0;
+  }
   ~RandomEffectsContainer() {}
   void AddSample(MultivariateRegressionRandomEffectsModel& model);
   void Predict(RandomEffectsDataset& dataset, LabelMapper& label_mapper, std::vector<double>& output);
@@ -257,6 +263,8 @@ class RandomEffectsContainer {
   std::vector<double>& GetAlpha() {return alpha_;}
   std::vector<double>& GetXi() {return xi_;}
   std::vector<double>& GetSigma() {return sigma_xi_;}
+  nlohmann::json to_json();
+  void from_json(const nlohmann::json& rfx_container_json);
  private:
   int num_samples_;
   int num_components_;

src/random_effects.cpp

@@ -199,4 +199,72 @@ void RandomEffectsContainer::Predict(RandomEffectsDataset& dataset, LabelMapper&
   }
 }
 
+nlohmann::json RandomEffectsContainer::to_json() {
+  json result_obj;
+  // Store the non-array fields in json
+  result_obj.emplace("num_samples", num_samples_);
+  result_obj.emplace("num_components", num_components_);
+  result_obj.emplace("num_groups", num_groups_);
+
+  // Store some meta-level information about the containers
+  int beta_size = num_groups_*num_components_*num_samples_;
+  int alpha_size = num_components_*num_samples_;
+  result_obj.emplace("beta_size", beta_size);
+  result_obj.emplace("alpha_size", alpha_size);
+
+  // Initialize a map with names of the node vectors and empty json arrays
+  std::map<std::string, json> tree_array_map;
+  tree_array_map.emplace(std::pair("beta", json::array()));
+  tree_array_map.emplace(std::pair("xi", json::array()));
+  tree_array_map.emplace(std::pair("alpha", json::array()));
+  tree_array_map.emplace(std::pair("sigma_xi", json::array()));
+
+  // Unpack beta and xi into json arrays
+  for (int i = 0; i < beta_size; i++) {
+    tree_array_map["beta"].emplace_back(beta_.at(i));
+    tree_array_map["xi"].emplace_back(xi_.at(i));
+  }
+
+  // Unpack alpha and sigma into json arrays
+  for (int i = 0; i < alpha_size; i++) {
+    tree_array_map["alpha"].emplace_back(alpha_.at(i));
+    tree_array_map["sigma_xi"].emplace_back(sigma_xi_.at(i));
+  }
+
+  // Unpack the map into the reference JSON object
+  for (auto& pair : tree_array_map) {
+    result_obj.emplace(pair);
+  }
+
+return result_obj;
+}
+
+void RandomEffectsContainer::from_json(const nlohmann::json& rfx_container_json) {
+  int beta_size = rfx_container_json.at("beta_size");
+  int alpha_size = rfx_container_json.at("alpha_size");
+
+  // Clear all internal arrays
+  beta_.clear();
+  xi_.clear();
+  alpha_.clear();
+  sigma_xi_.clear();
+
+  // Unpack internal counts
+  this->num_samples_ = rfx_container_json.at("num_samples");
+  this->num_components_ = rfx_container_json.at("num_components");
+  this->num_groups_ = rfx_container_json.at("num_groups");
+  
+  // Unpack beta and xi
+  for (int i = 0; i < beta_size; i++) {
+    beta_.push_back(rfx_container_json.at("beta").at(i));
+    xi_.push_back(rfx_container_json.at("xi").at(i));
+  }
+  
+  // Unpack alpha and sigma_xi
+  for (int i = 0; i < alpha_size; i++) {
+    alpha_.push_back(rfx_container_json.at("alpha").at(i));
+    sigma_xi_.push_back(rfx_container_json.at("sigma_xi").at(i));
+  }
+}
+
 }  // namespace StochTree

test/test_random_effects.cpp

@@ -160,3 +160,73 @@ TEST(RandomEffects, Predict) {
     ASSERT_EQ(output[i], output_expected[i]);
   }
 }
+
+TEST(RandomEffects, Serialization) {
+  // Load test data
+  StochTree::TestUtils::TestDataset test_dataset;
+  test_dataset = StochTree::TestUtils::LoadSmallDatasetUnivariateBasis();
+  std::vector<StochTree::FeatureType> feature_types(test_dataset.x_cols, StochTree::FeatureType::kNumeric);
+
+  // Construct dataset
+  int n = test_dataset.n;
+  StochTree::RandomEffectsDataset dataset = StochTree::RandomEffectsDataset();
+  dataset.AddBasis(test_dataset.rfx_basis.data(), test_dataset.n, test_dataset.rfx_basis_cols, test_dataset.row_major);
+  dataset.AddGroupLabels(test_dataset.rfx_groups);
+  
+  // Construct tracker, model state, and container
+  StochTree::RandomEffectsTracker tracker = StochTree::RandomEffectsTracker(test_dataset.rfx_groups);
+  StochTree::MultivariateRegressionRandomEffectsModel model = StochTree::MultivariateRegressionRandomEffectsModel(test_dataset.rfx_basis_cols, test_dataset.rfx_num_groups);
+  StochTree::RandomEffectsContainer container = StochTree::RandomEffectsContainer(test_dataset.rfx_basis_cols, test_dataset.rfx_num_groups);
+  StochTree::LabelMapper label_mapper = StochTree::LabelMapper(tracker.GetLabelMap());
+
+  // Set the values of alpha, xi and sigma in the model state (rather than simulating)
+  Eigen::VectorXd alpha(test_dataset.rfx_basis_cols);
+  Eigen::MatrixXd xi(test_dataset.rfx_basis_cols, test_dataset.rfx_num_groups);
+  Eigen::MatrixXd sigma(test_dataset.rfx_basis_cols, test_dataset.rfx_basis_cols);
+  alpha << 1.5;
+  xi << 2, 4;
+  Eigen::VectorXd xi0 = xi(Eigen::all, 0);
+  Eigen::VectorXd xi1 = xi(Eigen::all, 1);
+  sigma << 1;
+  model.SetWorkingParameter(alpha);
+  model.SetGroupParameter(xi0, 0);
+  model.SetGroupParameter(xi1, 1);
+  model.SetGroupParameterCovariance(sigma);
+  
+  // Push to the container
+  container.AddSample(model);
+
+  // Change values and push a second "sample" to the container
+  alpha << 2.0;
+  xi << 1, 3;
+  xi0 = xi(Eigen::all, 0);
+  xi1 = xi(Eigen::all, 1);
+  sigma << 1;
+  model.SetWorkingParameter(alpha);
+  model.SetGroupParameter(xi0, 0);
+  model.SetGroupParameter(xi1, 1);
+  model.SetGroupParameterCovariance(sigma);
+  container.AddSample(model);
+
+  // Json round trip
+  nlohmann::json container_json = container.to_json();
+  StochTree::RandomEffectsContainer container_deserialized = StochTree::RandomEffectsContainer();
+  container_deserialized.from_json(container_json);
+
+  // Check data in the container
+  std::vector<double> alpha_original = container.GetAlpha();
+  std::vector<double> alpha_deserialized = container_deserialized.GetAlpha();
+  for (int i = 0; i < alpha_deserialized.size(); i++) {
+    ASSERT_EQ(alpha_original[i], alpha_deserialized[i]);
+  }
+  std::vector<double> xi_original = container.GetXi();
+  std::vector<double> xi_deserialized = container_deserialized.GetXi();
+  for (int i = 0; i < xi_deserialized.size(); i++) {
+    ASSERT_EQ(xi_original[i], xi_deserialized[i]);
+  }
+  std::vector<double> beta_original = container.GetBeta();
+  std::vector<double> beta_deserialized = container_deserialized.GetBeta();
+  for (int i = 0; i < beta_deserialized.size(); i++) {
+    ASSERT_EQ(beta_original[i], beta_deserialized[i]);
+  }
+}