From 539bf708294f26896edd0cfe45d35f6e12964271 Mon Sep 17 00:00:00 2001
From: Frank Dellaert <dellaert@gmail.com>
Date: Sun, 6 Feb 2022 17:56:49 -0500
Subject: [PATCH] Sampling from GaussianBayesNet

---
 gtsam/linear/GaussianBayesNet.cpp           | 42 +++++++++++++--------
 gtsam/linear/GaussianBayesNet.h             | 13 +++++--
 gtsam/linear/tests/testGaussianBayesNet.cpp | 19 ++++++++++
 3 files changed, 56 insertions(+), 18 deletions(-)

diff --git a/gtsam/linear/GaussianBayesNet.cpp b/gtsam/linear/GaussianBayesNet.cpp
index 8fd4f2c26..87b1c6cb4 100644
--- a/gtsam/linear/GaussianBayesNet.cpp
+++ b/gtsam/linear/GaussianBayesNet.cpp
@@ -37,28 +37,40 @@ namespace gtsam {
     return Base::equals(bn, tol);
   }
 
-  /* ************************************************************************* */
-  VectorValues GaussianBayesNet::optimize() const
-  {
-    VectorValues soln; // no missing variables -> just create an empty vector
-    return optimize(soln);
+  /* ************************************************************************ */
+  VectorValues GaussianBayesNet::optimize() const {
+    VectorValues solution;  // no missing variables -> create an empty vector
+    return optimize(solution);
   }
 
-  /* ************************************************************************* */
-  VectorValues GaussianBayesNet::optimize(
-      const VectorValues& solutionForMissing) const {
-    VectorValues soln(solutionForMissing); // possibly empty
+  VectorValues GaussianBayesNet::optimize(VectorValues solution) const {
     // (R*x)./sigmas = y by solving x=inv(R)*(y.*sigmas)
-    /** solve each node in turn in topological sort order (parents first)*/
-    for (auto cg: boost::adaptors::reverse(*this)) {
+    // solve each node in reverse topological sort order (parents first)
+    for (auto cg : boost::adaptors::reverse(*this)) {
       // i^th part of R*x=y, x=inv(R)*y
-      // (Rii*xi + R_i*x(i+1:))./si = yi <-> xi = inv(Rii)*(yi.*si - R_i*x(i+1:))
-      soln.insert(cg->solve(soln));
+      // (Rii*xi + R_i*x(i+1:))./si = yi =>
+      // xi = inv(Rii)*(yi.*si - R_i*x(i+1:))
+      solution.insert(cg->solve(solution));
     }
-    return soln;
+    return solution;
   }
 
-  /* ************************************************************************* */
+  /* ************************************************************************ */
+  VectorValues GaussianBayesNet::sample() const {
+    VectorValues result;  // no missing variables -> create an empty vector
+    return sample(result);
+  }
+
+  VectorValues GaussianBayesNet::sample(VectorValues result) const {
+    // sample each node in reverse topological sort order (parents first)
+    for (auto cg : boost::adaptors::reverse(*this)) {
+      const VectorValues sampled = cg->sample(result);
+      result.insert(sampled);
+    }
+    return result;
+  }
+
+  /* ************************************************************************ */
   VectorValues GaussianBayesNet::optimizeGradientSearch() const
   {
     gttic(GaussianBayesTree_optimizeGradientSearch);
diff --git a/gtsam/linear/GaussianBayesNet.h b/gtsam/linear/GaussianBayesNet.h
index 6d906d65e..18da3ed8a 100644
--- a/gtsam/linear/GaussianBayesNet.h
+++ b/gtsam/linear/GaussianBayesNet.h
@@ -88,11 +88,18 @@ namespace gtsam {
     /// @name Standard Interface
     /// @{
 
-    /// Solve the GaussianBayesNet, i.e. return \f$ x = R^{-1}*d \f$, by back-substitution
+    /// Solve the GaussianBayesNet, i.e. return \f$ x = R^{-1}*d \f$, by
+    /// back-substitution
     VectorValues optimize() const;
 
-    /// Version of optimize for incomplete BayesNet, needs solution for missing variables
-    VectorValues optimize(const VectorValues& solutionForMissing) const;
+    /// Version of optimize for incomplete BayesNet, given missing variables
+    VectorValues optimize(const VectorValues given) const;
+
+    /// Sample using ancestral sampling
+    VectorValues sample() const;
+
+    /// Sample from an incomplete BayesNet, given missing variables
+    VectorValues sample(VectorValues given) const;
 
     /**
      * Return ordering corresponding to a topological sort.
diff --git a/gtsam/linear/tests/testGaussianBayesNet.cpp b/gtsam/linear/tests/testGaussianBayesNet.cpp
index f62da15dd..d9c6ee93d 100644
--- a/gtsam/linear/tests/testGaussianBayesNet.cpp
+++ b/gtsam/linear/tests/testGaussianBayesNet.cpp
@@ -16,10 +16,12 @@
  */
 
 #include <gtsam/linear/GaussianBayesNet.h>
+#include <gtsam/linear/GaussianDensity.h>
 #include <gtsam/linear/JacobianFactor.h>
 #include <gtsam/linear/GaussianFactorGraph.h>
 #include <gtsam/base/Testable.h>
 #include <gtsam/base/numericalDerivative.h>
+#include <gtsam/inference/Symbol.h>
 
 #include <CppUnitLite/TestHarness.h>
 #include <boost/tuple/tuple.hpp>
@@ -35,6 +37,7 @@ using namespace boost::assign;
 using namespace std::placeholders;
 using namespace std;
 using namespace gtsam;
+using symbol_shorthand::X;
 
 static const Key _x_ = 11, _y_ = 22, _z_ = 33;
 
@@ -138,6 +141,22 @@ TEST( GaussianBayesNet, optimize3 )
   EXPECT(assert_equal(expected, actual));
 }
 
+/* ************************************************************************* */
+TEST(GaussianBayesNet, sample) {
+  GaussianBayesNet gbn;
+  Matrix A1 = (Matrix(2, 2) << 1., 2., 3., 4.).finished();
+  const Vector2 mean(20, 40), b(10, 10);
+  const double sigma = 0.01;
+
+  gbn.add(GaussianConditional::FromMeanAndStddev(X(0), A1, X(1), b, sigma));
+  gbn.add(GaussianDensity::FromMeanAndStddev(X(1), mean, sigma));
+
+  auto actual = gbn.sample();
+  EXPECT_LONGS_EQUAL(2, actual.size());
+  EXPECT(assert_equal(mean, actual[X(1)], 50 * sigma));
+  EXPECT(assert_equal(A1 * mean + b, actual[X(0)], 50 * sigma));
+}
+
 /* ************************************************************************* */
 TEST(GaussianBayesNet, ordering)
 {