Hybrid Bayes Net sampling

gtsam/hybrid/HybridBayesNet.cpp

@@ -20,6 +20,9 @@
 #include <gtsam/hybrid/HybridBayesNet.h>
 #include <gtsam/hybrid/HybridValues.h>
 
+// In Wrappers we have no access to this so have a default ready
+static std::mt19937_64 kRandomNumberGenerator(42);
+
 namespace gtsam {
 
 /* ************************************************************************* */
@@ -232,6 +235,43 @@ VectorValues HybridBayesNet::optimize(const DiscreteValues &assignment) const {
   return gbn.optimize();
 }
 
+/* ************************************************************************* */
+HybridValues HybridBayesNet::sample(VectorValues given, std::mt19937_64 *rng,
+                                    SharedDiagonal model) const {
+  DiscreteBayesNet dbn;
+  for (size_t idx = 0; idx < size(); idx++) {
+    if (factors_.at(idx)->isDiscrete()) {
+      // If factor at `idx` is discrete-only, we add to the discrete bayes net.
+      dbn.push_back(this->atDiscrete(idx));
+    }
+  }
+  // Sample a discrete assignment.
+  DiscreteValues assignment = dbn.sample();
+  // Select the continuous bayes net corresponding to the assignment.
+  GaussianBayesNet gbn = this->choose(assignment);
+  // Sample from the gaussian bayes net.
+  VectorValues sample = gbn.sample(given, rng, model);
+  return HybridValues(assignment, sample);
+}
+
+/* ************************************************************************* */
+HybridValues HybridBayesNet::sample(std::mt19937_64 *rng,
+                                    SharedDiagonal model) const {
+  VectorValues given;
+  return sample(given, rng, model);
+}
+
+/* ************************************************************************* */
+HybridValues HybridBayesNet::sample(VectorValues given,
+                                    SharedDiagonal model) const {
+  return sample(given, &kRandomNumberGenerator, model);
+}
+
+/* ************************************************************************* */
+HybridValues HybridBayesNet::sample(SharedDiagonal model) const {
+  return sample(&kRandomNumberGenerator, model);
+}
+
 /* ************************************************************************* */
 double HybridBayesNet::error(const VectorValues &continuousValues,
                              const DiscreteValues &discreteValues) const {

gtsam/hybrid/HybridBayesNet.h

@@ -120,7 +120,53 @@ class GTSAM_EXPORT HybridBayesNet : public BayesNet<HybridConditional> {
    */
   DecisionTreeFactor::shared_ptr discreteConditionals() const;
 
- public:
+  /**
+   * @brief Sample from an incomplete BayesNet, given missing variables.
+   *
+   * Example:
+   *   std::mt19937_64 rng(42);
+   *   VectorValues given = ...;
+   *   auto sample = bn.sample(given, &rng);
+   *
+   * @param given Values of missing variables.
+   * @param rng The pseudo-random number generator.
+   * @param model Optional diagonal noise model to use in sampling.
+   * @return HybridValues
+   */
+  HybridValues sample(VectorValues given, std::mt19937_64 *rng,
+                      SharedDiagonal model = nullptr) const;
+
+  /**
+   * @brief Sample using ancestral sampling.
+   *
+   * Example:
+   *   std::mt19937_64 rng(42);
+   *   auto sample = bn.sample(&rng);
+   *
+   * @param rng The pseudo-random number generator.
+   * @param model Optional diagonal noise model to use in sampling.
+   * @return HybridValues
+   */
+  HybridValues sample(std::mt19937_64 *rng,
+                      SharedDiagonal model = nullptr) const;
+
+  /**
+   * @brief Sample from an incomplete BayesNet, use default rng.
+   *
+   * @param given Values of missing variables.
+   * @param model Optional diagonal noise model to use in sampling.
+   * @return HybridValues
+   */
+  HybridValues sample(VectorValues given, SharedDiagonal model = nullptr) const;
+
+  /**
+   * @brief Sample using ancestral sampling, use default rng.
+   *
+   * @param model Optional diagonal noise model to use in sampling.
+   * @return HybridValues
+   */
+  HybridValues sample(SharedDiagonal model = nullptr) const;
+
   /// Prune the Hybrid Bayes Net such that we have at most maxNrLeaves leaves.
   HybridBayesNet prune(size_t maxNrLeaves);
 

gtsam/hybrid/tests/testHybridBayesNet.cpp

@@ -316,6 +316,70 @@ TEST(HybridBayesNet, Serialization) {
   EXPECT(equalsBinary<HybridBayesNet>(hbn));
 }
 
+/* ****************************************************************************/
+// Test HybridBayesNet sampling.
+TEST(HybridBayesNet, Sampling) {
+  HybridNonlinearFactorGraph nfg;
+
+  auto noise_model = noiseModel::Diagonal::Sigmas(Vector1(1.0));
+  auto zero_motion =
+      boost::make_shared<BetweenFactor<double>>(X(0), X(1), 0, noise_model);
+  auto one_motion =
+      boost::make_shared<BetweenFactor<double>>(X(0), X(1), 1, noise_model);
+  std::vector<NonlinearFactor::shared_ptr> factors = {zero_motion, one_motion};
+  nfg.emplace_nonlinear<PriorFactor<double>>(X(0), 0.0, noise_model);
+  nfg.emplace_hybrid<MixtureFactor>(
+      KeyVector{X(0), X(1)}, DiscreteKeys{DiscreteKey(M(0), 2)}, factors);
+
+  DiscreteKey mode(M(0), 2);
+  auto discrete_prior = boost::make_shared<DiscreteDistribution>(mode, "1/1");
+  nfg.push_discrete(discrete_prior);
+
+  Values initial;
+  double z0 = 0.0, z1 = 1.0;
+  initial.insert<double>(X(0), z0);
+  initial.insert<double>(X(1), z1);
+
+  // Create the factor graph from the nonlinear factor graph.
+  HybridGaussianFactorGraph::shared_ptr fg = nfg.linearize(initial);
+  // Eliminate into BN
+  Ordering ordering = fg->getHybridOrdering();
+  HybridBayesNet::shared_ptr bn = fg->eliminateSequential(ordering);
+
+  // Set up sampling
+  std::mt19937_64 gen(11);
+
+  // Initialize containers for computing the mean values.
+  vector<double> discrete_samples;
+  VectorValues average_continuous;
+
+  size_t num_samples = 1000;
+  for (size_t i = 0; i < num_samples; i++) {
+    // Sample
+    HybridValues sample = bn->sample(&gen, noise_model);
+
+    discrete_samples.push_back(sample.discrete()[M(0)]);
+
+    if (i == 0) {
+      average_continuous.insert(sample.continuous());
+    } else {
+      average_continuous += sample.continuous();
+    }
+  }
+  double discrete_sum =
+      std::accumulate(discrete_samples.begin(), discrete_samples.end(),
+                      decltype(discrete_samples)::value_type(0));
+
+  // regression for specific RNG seed
+  EXPECT_DOUBLES_EQUAL(0.477, discrete_sum / num_samples, 1e-9);
+
+  VectorValues expected;
+  expected.insert({X(0), Vector1(-0.0131207162712)});
+  expected.insert({X(1), Vector1(-0.499026377568)});
+  // regression for specific RNG seed
+  EXPECT(assert_equal(expected, average_continuous.scale(1.0 / num_samples)));
+}
+
 /* ************************************************************************* */
 int main() {
   TestResult tr;

gtsam/linear/GaussianBayesNet.cpp

@@ -59,27 +59,30 @@ namespace gtsam {
   }
 
   /* ************************************************************************ */
-  VectorValues GaussianBayesNet::sample(std::mt19937_64* rng) const {
+  VectorValues GaussianBayesNet::sample(std::mt19937_64* rng,
+                                        const SharedDiagonal& model) const {
     VectorValues result;  // no missing variables -> create an empty vector
-    return sample(result, rng);
+    return sample(result, rng, model);
   }
 
   VectorValues GaussianBayesNet::sample(VectorValues result,
-                                        std::mt19937_64* rng) const {
+                                        std::mt19937_64* rng,
+                                        const SharedDiagonal& model) const {
     // sample each node in reverse topological sort order (parents first)
     for (auto cg : boost::adaptors::reverse(*this)) {
-      const VectorValues sampled = cg->sample(result, rng);
+      const VectorValues sampled = cg->sample(result, rng, model);
       result.insert(sampled);
     }
     return result;
   }
 
   /* ************************************************************************ */
-  VectorValues GaussianBayesNet::sample() const {
+  VectorValues GaussianBayesNet::sample(const SharedDiagonal& model) const {
     return sample(&kRandomNumberGenerator);
   }
 
-  VectorValues GaussianBayesNet::sample(VectorValues given) const {
+  VectorValues GaussianBayesNet::sample(VectorValues given,
+                                        const SharedDiagonal& model) const {
     return sample(given, &kRandomNumberGenerator);
   }
 

gtsam/linear/GaussianBayesNet.h

@@ -101,7 +101,8 @@ namespace gtsam {
      *   std::mt19937_64 rng(42);
      *   auto sample = gbn.sample(&rng);
      */
-    VectorValues sample(std::mt19937_64* rng) const;
+    VectorValues sample(std::mt19937_64* rng,
+                        const SharedDiagonal& model = nullptr) const;
 
     /**
      * Sample from an incomplete BayesNet, given missing variables
@@ -110,13 +111,15 @@ namespace gtsam {
      *   VectorValues given = ...;
      *   auto sample = gbn.sample(given, &rng);
      */
-    VectorValues sample(VectorValues given, std::mt19937_64* rng) const;
+    VectorValues sample(VectorValues given, std::mt19937_64* rng,
+                        const SharedDiagonal& model = nullptr) const;
 
     /// Sample using ancestral sampling, use default rng
-    VectorValues sample() const;
+    VectorValues sample(const SharedDiagonal& model = nullptr) const;
 
     /// Sample from an incomplete BayesNet, use default rng
-    VectorValues sample(VectorValues given) const;
+    VectorValues sample(VectorValues given,
+                        const SharedDiagonal& model = nullptr) const;
 
     /**
      * Return ordering corresponding to a topological sort.

gtsam/linear/GaussianConditional.cpp

@@ -293,39 +293,48 @@ double GaussianConditional::logDeterminant() const {
 
   /* ************************************************************************ */
   VectorValues GaussianConditional::sample(const VectorValues& parentsValues,
-                                           std::mt19937_64* rng) const {
+                                           std::mt19937_64* rng,
+                                           const SharedDiagonal& model) const {
     if (nrFrontals() != 1) {
       throw std::invalid_argument(
           "GaussianConditional::sample can only be called on single variable "
           "conditionals");
     }
-    if (!model_) {
+
+    VectorValues solution = solve(parentsValues);
+    Key key = firstFrontalKey();
+
+    Vector sigmas;
+    if (model_) {
+      sigmas = model_->sigmas();
+    } else if (model) {
+      sigmas = model->sigmas();
+    } else {
       throw std::invalid_argument(
           "GaussianConditional::sample can only be called if a diagonal noise "
           "model was specified at construction.");
     }
-    VectorValues solution = solve(parentsValues);
-    Key key = firstFrontalKey();
-    const Vector& sigmas = model_->sigmas();
     solution[key] += Sampler::sampleDiagonal(sigmas, rng);
     return solution;
   }
 
-  VectorValues GaussianConditional::sample(std::mt19937_64* rng) const {
+  VectorValues GaussianConditional::sample(std::mt19937_64* rng,
+                                           const SharedDiagonal& model) const {
     if (nrParents() != 0)
       throw std::invalid_argument(
           "sample() can only be invoked on no-parent prior");
     VectorValues values;
-    return sample(values);
+    return sample(values, rng, model);
   }
 
   /* ************************************************************************ */
-  VectorValues GaussianConditional::sample() const {
-    return sample(&kRandomNumberGenerator);
+  VectorValues GaussianConditional::sample(const SharedDiagonal& model) const {
+    return sample(&kRandomNumberGenerator, model);
   }
 
-  VectorValues GaussianConditional::sample(const VectorValues& given) const {
-    return sample(given, &kRandomNumberGenerator);
+  VectorValues GaussianConditional::sample(const VectorValues& given,
+                                           const SharedDiagonal& model) const {
+    return sample(given, &kRandomNumberGenerator, model);
   }
 
   /* ************************************************************************ */

gtsam/linear/GaussianConditional.h

@@ -188,7 +188,8 @@ namespace gtsam {
      *   std::mt19937_64 rng(42);
      *   auto sample = gbn.sample(&rng);
      */
-    VectorValues sample(std::mt19937_64* rng) const;
+    VectorValues sample(std::mt19937_64* rng,
+                        const SharedDiagonal& model = nullptr) const;
 
     /**
      * Sample from conditional, given missing variables
@@ -198,13 +199,15 @@ namespace gtsam {
      *   auto sample = gbn.sample(given, &rng);
      */
     VectorValues sample(const VectorValues& parentsValues,
-                        std::mt19937_64* rng) const;
+                        std::mt19937_64* rng,
+                        const SharedDiagonal& model = nullptr) const;
 
     /// Sample, use default rng
-    VectorValues sample() const;
+    VectorValues sample(const SharedDiagonal& model = nullptr) const;
 
     /// Sample with given values, use default rng
-    VectorValues sample(const VectorValues& parentsValues) const;
+    VectorValues sample(const VectorValues& parentsValues,
+                        const SharedDiagonal& model = nullptr) const;
 
     /// @}