Merge pull request #1271 from borglab/feature/nonlinear-incremental

gtsam/hybrid/GaussianMixture.cpp

@@ -119,11 +119,36 @@ void GaussianMixture::print(const std::string &s,
       "", [&](Key k) { return formatter(k); },
       [&](const GaussianConditional::shared_ptr &gf) -> std::string {
         RedirectCout rd;
-        if (!gf->empty())
+        if (gf && !gf->empty())
           gf->print("", formatter);
         else
           return {"nullptr"};
         return rd.str();
       });
 }
+
+/* *******************************************************************************/
+void GaussianMixture::prune(const DecisionTreeFactor &decisionTree) {
+  // Functional which loops over all assignments and create a set of
+  // GaussianConditionals
+  auto pruner = [&decisionTree](
+                    const Assignment<Key> &choices,
+                    const GaussianConditional::shared_ptr &conditional)
+      -> GaussianConditional::shared_ptr {
+    // typecast so we can use this to get probability value
+    DiscreteValues values(choices);
+
+    if (decisionTree(values) == 0.0) {
+      // empty aka null pointer
+      boost::shared_ptr<GaussianConditional> null;
+      return null;
+    } else {
+      return conditional;
+    }
+  };
+
+  auto pruned_conditionals = conditionals_.apply(pruner);
+  conditionals_.root_ = pruned_conditionals.root_;
+}
+
 }  // namespace gtsam

gtsam/hybrid/GaussianMixture.h

@@ -21,6 +21,7 @@
 
 #include <gtsam/discrete/DecisionTree-inl.h>
 #include <gtsam/discrete/DecisionTree.h>
+#include <gtsam/discrete/DecisionTreeFactor.h>
 #include <gtsam/discrete/DiscreteKey.h>
 #include <gtsam/hybrid/HybridFactor.h>
 #include <gtsam/inference/Conditional.h>
@@ -121,7 +122,7 @@ class GTSAM_EXPORT GaussianMixture
   /// Test equality with base HybridFactor
   bool equals(const HybridFactor &lf, double tol = 1e-9) const override;
 
-  /* print utility */
+  /// Print utility
   void print(
       const std::string &s = "GaussianMixture\n",
       const KeyFormatter &formatter = DefaultKeyFormatter) const override;
@@ -131,6 +132,15 @@ class GTSAM_EXPORT GaussianMixture
   /// Getter for the underlying Conditionals DecisionTree
   const Conditionals &conditionals();
 
+  /**
+   * @brief Prune the decision tree of Gaussian factors as per the discrete
+   * `decisionTree`.
+   *
+   * @param decisionTree A pruned decision tree of discrete keys where the
+   * leaves are probabilities.
+   */
+  void prune(const DecisionTreeFactor &decisionTree);
+
   /**
    * @brief Merge the Gaussian Factor Graphs in `this` and `sum` while
    * maintaining the decision tree structure.

gtsam/hybrid/HybridGaussianISAM.cpp

@@ -41,6 +41,7 @@ HybridGaussianISAM::HybridGaussianISAM(const HybridBayesTree& bayesTree)
 void HybridGaussianISAM::updateInternal(
     const HybridGaussianFactorGraph& newFactors,
     HybridBayesTree::Cliques* orphans,
+    const boost::optional<Ordering>& ordering,
     const HybridBayesTree::Eliminate& function) {
   // Remove the contaminated part of the Bayes tree
   BayesNetType bn;
@@ -78,12 +79,19 @@ void HybridGaussianISAM::updateInternal(
 
   // Get an ordering where the new keys are eliminated last
   const VariableIndex index(factors);
-  const Ordering ordering = Ordering::ColamdConstrainedLast(
+  Ordering elimination_ordering;
-      index, KeyVector(newKeysDiscreteLast.begin(), newKeysDiscreteLast.end()),
+  if (ordering) {
+    elimination_ordering = *ordering;
+  } else {
+    elimination_ordering = Ordering::ColamdConstrainedLast(
+        index,
+        KeyVector(newKeysDiscreteLast.begin(), newKeysDiscreteLast.end()),
         true);
+  }
 
   // eliminate all factors (top, added, orphans) into a new Bayes tree
-  auto bayesTree = factors.eliminateMultifrontal(ordering, function, index);
+  HybridBayesTree::shared_ptr bayesTree =
+      factors.eliminateMultifrontal(elimination_ordering, function, index);
 
   // Re-add into Bayes tree data structures
   this->roots_.insert(this->roots_.end(), bayesTree->roots().begin(),
@@ -93,9 +101,45 @@ void HybridGaussianISAM::updateInternal(
 
 /* ************************************************************************* */
 void HybridGaussianISAM::update(const HybridGaussianFactorGraph& newFactors,
+                                const boost::optional<Ordering>& ordering,
                                 const HybridBayesTree::Eliminate& function) {
   Cliques orphans;
-  this->updateInternal(newFactors, &orphans, function);
+  this->updateInternal(newFactors, &orphans, ordering, function);
+}
+
+void HybridGaussianISAM::prune(const Key& root, const size_t maxNrLeaves) {
+  auto decisionTree = boost::dynamic_pointer_cast<DecisionTreeFactor>(
+      this->clique(root)->conditional()->inner());
+  DecisionTreeFactor prunedDiscreteFactor = decisionTree->prune(maxNrLeaves);
+  decisionTree->root_ = prunedDiscreteFactor.root_;
+
+  std::vector<gtsam::Key> prunedKeys;
+  for (auto&& clique : nodes()) {
+    // The cliques can be repeated for each frontal so we record it in
+    // prunedKeys and check if we have already pruned a particular clique.
+    if (std::find(prunedKeys.begin(), prunedKeys.end(), clique.first) !=
+        prunedKeys.end()) {
+      continue;
+    }
+
+    // Add all the keys of the current clique to be pruned to prunedKeys
+    for (auto&& key : clique.second->conditional()->frontals()) {
+      prunedKeys.push_back(key);
+    }
+
+    // Convert parents() to a KeyVector for comparison
+    KeyVector parents;
+    for (auto&& parent : clique.second->conditional()->parents()) {
+      parents.push_back(parent);
+    }
+
+    if (parents == decisionTree->keys()) {
+      auto gaussianMixture = boost::dynamic_pointer_cast<GaussianMixture>(
+          clique.second->conditional()->inner());
+
+      gaussianMixture->prune(prunedDiscreteFactor);
+    }
+  }
 }
 
 }  // namespace gtsam

gtsam/hybrid/HybridGaussianISAM.h

@@ -48,6 +48,7 @@ class GTSAM_EXPORT HybridGaussianISAM : public ISAM<HybridBayesTree> {
   void updateInternal(
       const HybridGaussianFactorGraph& newFactors,
       HybridBayesTree::Cliques* orphans,
+      const boost::optional<Ordering>& ordering = boost::none,
       const HybridBayesTree::Eliminate& function =
           HybridBayesTree::EliminationTraitsType::DefaultEliminate);
 
@@ -59,8 +60,17 @@ class GTSAM_EXPORT HybridGaussianISAM : public ISAM<HybridBayesTree> {
    * @param function Elimination function.
    */
   void update(const HybridGaussianFactorGraph& newFactors,
+              const boost::optional<Ordering>& ordering = boost::none,
               const HybridBayesTree::Eliminate& function =
                   HybridBayesTree::EliminationTraitsType::DefaultEliminate);
+
+  /**
+   * @brief 
+   * 
+   * @param root The root key in the discrete conditional decision tree.
+   * @param maxNumberLeaves 
+   */
+  void prune(const Key& root, const size_t maxNumberLeaves);
 };
 
 /// traits

gtsam/hybrid/tests/testHybridIncremental.cpp

@@ -18,7 +18,9 @@
 
 #include <gtsam/discrete/DiscreteBayesNet.h>
 #include <gtsam/discrete/DiscreteDistribution.h>
+#include <gtsam/discrete/DiscreteFactorGraph.h>
 #include <gtsam/geometry/Pose2.h>
+#include <gtsam/hybrid/HybridConditional.h>
 #include <gtsam/hybrid/HybridGaussianISAM.h>
 #include <gtsam/linear/GaussianBayesNet.h>
 #include <gtsam/linear/GaussianFactorGraph.h>
@@ -52,7 +54,7 @@ TEST(HybridGaussianElimination, IncrementalElimination) {
   // Create initial factor graph
   //  *        *      *
   //  |        |      |
-  // *- X1  -*-  X2 -*- X3
+  //  X1  -*-  X2 -*- X3
   //   \*-M1-*/
   graph1.push_back(switching.linearizedFactorGraph.at(0));  // P(X1)
   graph1.push_back(switching.linearizedFactorGraph.at(1));  // P(X1, X2 | M1)
@@ -108,7 +110,7 @@ TEST(HybridGaussianElimination, IncrementalInference) {
   // Create initial factor graph
   //    *        *        *
   //    |        |        |
-  // *- X1  -*-  X2  -*-  X3
+  //    X1  -*-  X2  -*-  X3
   //         |        |
   //      *-M1 - * - M2
   graph1.push_back(switching.linearizedFactorGraph.at(0));  // P(X1)
@@ -116,15 +118,13 @@ TEST(HybridGaussianElimination, IncrementalInference) {
   graph1.push_back(switching.linearizedFactorGraph.at(3));  // P(X2)
   graph1.push_back(switching.linearizedFactorGraph.at(5));  // P(M1)
 
-  //TODO(Varun) we cannot enforce ordering
-  // // Create ordering.
-  // Ordering ordering1;
-  // ordering1 += X(1);
-  // ordering1 += X(2);
-
   // Run update step
   isam.update(graph1);
 
+  auto discreteConditional_m1 =
+      isam[M(1)]->conditional()->asDiscreteConditional();
+  EXPECT(discreteConditional_m1->keys() == KeyVector({M(1)}));
+
   /********************************************************/
   // New factor graph for incremental update.
   HybridGaussianFactorGraph graph2;
@@ -133,14 +133,7 @@ TEST(HybridGaussianElimination, IncrementalInference) {
   graph2.push_back(switching.linearizedFactorGraph.at(4));  // P(X3)
   graph2.push_back(switching.linearizedFactorGraph.at(6));  // P(M1, M2)
 
-  //TODO(Varun) we cannot enforce ordering
-  // // Create ordering.
-  // Ordering ordering2;
-  // ordering2 += X(2);
-  // ordering2 += X(3);
-
   isam.update(graph2);
-  GTSAM_PRINT(isam);
 
   /********************************************************/
   // Run batch elimination so we can compare results.
@@ -150,269 +143,255 @@ TEST(HybridGaussianElimination, IncrementalInference) {
   ordering += X(3);
 
   // Now we calculate the actual factors using full elimination
-  HybridBayesNet::shared_ptr expectedHybridBayesNet;
+  HybridBayesTree::shared_ptr expectedHybridBayesTree;
   HybridGaussianFactorGraph::shared_ptr expectedRemainingGraph;
-  std::tie(expectedHybridBayesNet, expectedRemainingGraph) =
+  std::tie(expectedHybridBayesTree, expectedRemainingGraph) =
-      switching.linearizedFactorGraph.eliminatePartialSequential(ordering);
+      switching.linearizedFactorGraph.eliminatePartialMultifrontal(ordering);
 
   // The densities on X(1) should be the same
   auto x1_conditional =
       dynamic_pointer_cast<GaussianMixture>(isam[X(1)]->conditional()->inner());
-  EXPECT(
+  auto actual_x1_conditional = dynamic_pointer_cast<GaussianMixture>(
-      assert_equal(*x1_conditional, *(expectedHybridBayesNet->atGaussian(0))));
+      (*expectedHybridBayesTree)[X(1)]->conditional()->inner());
+  EXPECT(assert_equal(*x1_conditional, *actual_x1_conditional));
 
   // The densities on X(2) should be the same
   auto x2_conditional =
       dynamic_pointer_cast<GaussianMixture>(isam[X(2)]->conditional()->inner());
-  EXPECT(
+  auto actual_x2_conditional = dynamic_pointer_cast<GaussianMixture>(
-      assert_equal(*x2_conditional, *(expectedHybridBayesNet->atGaussian(1))));
+      (*expectedHybridBayesTree)[X(2)]->conditional()->inner());
+  EXPECT(assert_equal(*x2_conditional, *actual_x2_conditional));
 
-  // // The densities on X(3) should be the same
+  // The densities on X(3) should be the same
-  // auto x3_conditional =
+  auto x3_conditional =
-  //     dynamic_pointer_cast<GaussianMixture>(isam[X(3)]->conditional()->inner());
+      dynamic_pointer_cast<GaussianMixture>(isam[X(3)]->conditional()->inner());
-  // EXPECT(
+  auto actual_x3_conditional = dynamic_pointer_cast<GaussianMixture>(
-  //     assert_equal(*x3_conditional, *(expectedHybridBayesNet->atGaussian(2))));
+      (*expectedHybridBayesTree)[X(2)]->conditional()->inner());
+  EXPECT(assert_equal(*x3_conditional, *actual_x3_conditional));
 
-  GTSAM_PRINT(*expectedHybridBayesNet);
+  // We only perform manual continuous elimination for 0,0.
-
+  // The other discrete probabilities on M(2) are calculated the same way
-  // we only do the manual continuous elimination for 0,0
+  Ordering discrete_ordering;
-  // the other discrete probabilities on M(2) are calculated the same way
+  discrete_ordering += M(1);
-  auto m00_prob = [&]() {
+  discrete_ordering += M(2);
-    GaussianFactorGraph gf;
+  HybridBayesTree::shared_ptr discreteBayesTree =
-    // gf.add(switching.linearizedFactorGraph.gaussianGraph().at(3));
+      expectedRemainingGraph->eliminateMultifrontal(discrete_ordering);
 
   DiscreteValues m00;
   m00[M(1)] = 0, m00[M(2)] = 0;
-    // auto dcMixture =
+  DiscreteConditional decisionTree =
-    //     dynamic_pointer_cast<DCGaussianMixtureFactor>(graph2.dcGraph().at(0));
+      *(*discreteBayesTree)[M(2)]->conditional()->asDiscreteConditional();
-    // gf.add(dcMixture->factors()(m00));
+  double m00_prob = decisionTree(m00);
-    // auto x2_mixed =
-    //     boost::dynamic_pointer_cast<GaussianMixture>(hybridBayesNet.at(1));
-    // gf.add(x2_mixed->factors()(m00));
-    auto result_gf = gf.eliminateSequential();
-    return gf.probPrime(result_gf->optimize());
-  }();
 
-  /// Test if the probability values are as expected with regression tests.
+  auto discreteConditional = isam[M(2)]->conditional()->asDiscreteConditional();
-//   DiscreteValues assignment;
-//   EXPECT(assert_equal(m00_prob, 0.60656, 1e-5));
-//   assignment[M(1)] = 0;
-//   assignment[M(2)] = 0;
-//   EXPECT(assert_equal(m00_prob, (*discreteFactor)(assignment), 1e-5));
-//   assignment[M(1)] = 1;
-//   assignment[M(2)] = 0;
-//   EXPECT(assert_equal(0.612477, (*discreteFactor)(assignment), 1e-5));
-//   assignment[M(1)] = 0;
-//   assignment[M(2)] = 1;
-//   EXPECT(assert_equal(0.999952, (*discreteFactor)(assignment), 1e-5));
-//   assignment[M(1)] = 1;
-//   assignment[M(2)] = 1;
-//   EXPECT(assert_equal(1.0, (*discreteFactor)(assignment), 1e-5));
 
-//   DiscreteFactorGraph dfg;
+  // Test if the probability values are as expected with regression tests.
-//   dfg.add(*discreteFactor);
+  DiscreteValues assignment;
-//   dfg.add(discreteFactor_m1);
+  EXPECT(assert_equal(m00_prob, 0.0619233, 1e-5));
-//   dfg.add_factors(switching.linearizedFactorGraph.discreteGraph());
+  assignment[M(1)] = 0;
+  assignment[M(2)] = 0;
+  EXPECT(assert_equal(m00_prob, (*discreteConditional)(assignment), 1e-5));
+  assignment[M(1)] = 1;
+  assignment[M(2)] = 0;
+  EXPECT(assert_equal(0.183743, (*discreteConditional)(assignment), 1e-5));
+  assignment[M(1)] = 0;
+  assignment[M(2)] = 1;
+  EXPECT(assert_equal(0.204159, (*discreteConditional)(assignment), 1e-5));
+  assignment[M(1)] = 1;
+  assignment[M(2)] = 1;
+  EXPECT(assert_equal(0.2, (*discreteConditional)(assignment), 1e-5));
 
-//   // Check if the chordal graph generated from incremental elimination
+  // Check if the clique conditional generated from incremental elimination
-//   matches
+  // matches that of batch elimination.
-//   // that of batch elimination.
+  auto expectedChordal = expectedRemainingGraph->eliminateMultifrontal();
-//   auto chordal = dfg.eliminateSequential();
+  auto expectedConditional = dynamic_pointer_cast<DecisionTreeFactor>(
-//   auto expectedChordal =
+      (*expectedChordal)[M(2)]->conditional()->inner());
-//       expectedRemainingGraph->discreteGraph().eliminateSequential();
+  auto actualConditional = dynamic_pointer_cast<DecisionTreeFactor>(
-
+      isam[M(2)]->conditional()->inner());
-//   EXPECT(assert_equal(*expectedChordal, *chordal, 1e-6));
+  EXPECT(assert_equal(*actualConditional, *expectedConditional, 1e-6));
 }
 
 /* ****************************************************************************/
-// // Test if we can approximately do the inference
+// Test if we can approximately do the inference
-// TEST(DCGaussianElimination, Approx_inference) {
+TEST(HybridGaussianElimination, Approx_inference) {
-//   Switching switching(4);
+  Switching switching(4);
-//   IncrementalHybrid incrementalHybrid;
+  HybridGaussianISAM incrementalHybrid;
-//   HybridGaussianFactorGraph graph1;
+  HybridGaussianFactorGraph graph1;
 
-//   // Add the 3 DC factors, x1-x2, x2-x3, x3-x4
+  // Add the 3 hybrid factors, x1-x2, x2-x3, x3-x4
-//   for (size_t i = 0; i < 3; i++) {
+  for (size_t i = 1; i < 4; i++) {
-//     graph1.push_back(switching.linearizedFactorGraph.dcGraph().at(i));
+    graph1.push_back(switching.linearizedFactorGraph.at(i));
-//   }
+  }
 
-//   // Add the Gaussian factors, 1 prior on X(1), 4 measurements
+  // Add the Gaussian factors, 1 prior on X(1),
-//   for (size_t i = 0; i <= 4; i++) {
+  // 3 measurements on X(2), X(3), X(4)
-// graph1.push_back(switching.linearizedFactorGraph.gaussianGraph().at(i));
+  graph1.push_back(switching.linearizedFactorGraph.at(0));
-//   }
+  for (size_t i = 4; i <= 7; i++) {
+    graph1.push_back(switching.linearizedFactorGraph.at(i));
+  }
 
-//   // Create ordering.
+  // Create ordering.
-//   Ordering ordering;
+  Ordering ordering;
-//   for (size_t j = 1; j <= 4; j++) {
+  for (size_t j = 1; j <= 4; j++) {
-//     ordering += X(j);
+    ordering += X(j);
-//   }
+  }
 
-//   // Now we calculate the actual factors using full elimination
+  // Now we calculate the actual factors using full elimination
-//   HybridBayesNet::shared_ptr unprunedHybridBayesNet;
+  HybridBayesTree::shared_ptr unprunedHybridBayesTree;
-//   HybridGaussianFactorGraph::shared_ptr unprunedRemainingGraph;
+  HybridGaussianFactorGraph::shared_ptr unprunedRemainingGraph;
-//   std::tie(unprunedHybridBayesNet, unprunedRemainingGraph) =
+  std::tie(unprunedHybridBayesTree, unprunedRemainingGraph) =
-// switching.linearizedFactorGraph.eliminatePartialSequential(ordering);
+      switching.linearizedFactorGraph.eliminatePartialMultifrontal(ordering);
 
-//   size_t maxComponents = 5;
+  size_t maxNrLeaves = 5;
-//   incrementalHybrid.update(graph1, ordering, maxComponents);
+  incrementalHybrid.update(graph1);
 
-//   /*
+  incrementalHybrid.prune(M(3), maxNrLeaves);
-//   unpruned factor is:
-//       Choice(m3)
-//       0 Choice(m2)
-//       0 0 Choice(m1)
-//       0 0 0 Leaf 0.2248 -
-//       0 0 1 Leaf 0.3715 -
-//       0 1 Choice(m1)
-//       0 1 0 Leaf 0.3742 *
-//       0 1 1 Leaf 0.6125 *
-//       1 Choice(m2)
-//       1 0 Choice(m1)
-//       1 0 0 Leaf 0.3706 -
-//       1 0 1 Leaf 0.6124 *
-//       1 1 Choice(m1)
-//       1 1 0 Leaf 0.611 *
-//       1 1 1 Leaf    1 *
 
-//   pruned factors is:
+  /*
-//       Choice(m3)
+  unpruned factor is:
-//       0 Choice(m2)
+    Choice(m3)
-//       0 0 Leaf    0
+    0 Choice(m2)
-//       0 1 Choice(m1)
+    0 0 Choice(m1)
-//       0 1 0 Leaf 0.3742
+    0 0 0 Leaf 0.11267528
-//       0 1 1 Leaf 0.6125
+    0 0 1 Leaf 0.18576102
-//       1 Choice(m2)
+    0 1 Choice(m1)
-//       1 0 Choice(m1)
+    0 1 0 Leaf 0.18754662
-//       1 0 0 Leaf    0
+    0 1 1 Leaf 0.30623871
-//       1 0 1 Leaf 0.6124
+    1 Choice(m2)
-//       1 1 Choice(m1)
+    1 0 Choice(m1)
-//       1 1 0 Leaf 0.611
+    1 0 0 Leaf 0.18576102
-//       1 1 1 Leaf    1
+    1 0 1 Leaf 0.30622428
-//    */
+    1 1 Choice(m1)
+    1 1 0 Leaf 0.30623871
+    1 1 1 Leaf  0.5
 
-//   // Test that the remaining factor graph has one
+  pruned factors is:
-//   // DecisionTreeFactor on {M3, M2, M1}.
+    Choice(m3)
-//   auto remainingFactorGraph = incrementalHybrid.remainingFactorGraph();
+    0 Choice(m2)
-//   EXPECT_LONGS_EQUAL(1, remainingFactorGraph.size());
+    0 0 Leaf    0
+    0 1 Choice(m1)
+    0 1 0 Leaf 0.18754662
+    0 1 1 Leaf 0.30623871
+    1 Choice(m2)
+    1 0 Choice(m1)
+    1 0 0 Leaf    0
+    1 0 1 Leaf 0.30622428
+    1 1 Choice(m1)
+    1 1 0 Leaf 0.30623871
+    1 1 1 Leaf  0.5
+  */
 
-//   auto discreteFactor_m1 = *dynamic_pointer_cast<DecisionTreeFactor>(
+  auto discreteConditional_m1 = *dynamic_pointer_cast<DiscreteConditional>(
-//       incrementalHybrid.remainingDiscreteGraph().at(0));
+      incrementalHybrid[M(1)]->conditional()->inner());
-//   EXPECT(discreteFactor_m1.keys() == KeyVector({M(3), M(2), M(1)}));
+  EXPECT(discreteConditional_m1.keys() == KeyVector({M(1), M(2), M(3)}));
 
-//   // Get the number of elements which are equal to 0.
+  // Get the number of elements which are greater than 0.
-//   auto count = [](const double &value, int count) {
+  auto count = [](const double &value, int count) {
-//     return value > 0 ? count + 1 : count;
+    return value > 0 ? count + 1 : count;
-//   };
+  };
-//   // Check that the number of leaves after pruning is 5.
+  // Check that the number of leaves after pruning is 5.
-//   EXPECT_LONGS_EQUAL(5, discreteFactor_m1.fold(count, 0));
+  EXPECT_LONGS_EQUAL(5, discreteConditional_m1.fold(count, 0));
 
-//   /* Expected hybrid Bayes net
+  // Check that the hybrid nodes of the bayes net match those of the pre-pruning
-//    * factor 0:  [x1 | x2 m1 ], 2 components
+  // bayes net, at the same positions.
-//    * factor 1:  [x2 | x3 m2 m1 ], 4 components
+  auto &unprunedLastDensity = *dynamic_pointer_cast<GaussianMixture>(
-//    * factor 2:  [x3 | x4 m3 m2 m1 ], 8 components
+      unprunedHybridBayesTree->clique(X(4))->conditional()->inner());
-//    * factor 3:  [x4 | m3 m2 m1 ], 8 components
+  auto &lastDensity = *dynamic_pointer_cast<GaussianMixture>(
-//    */
+      incrementalHybrid[X(4)]->conditional()->inner());
-//   auto hybridBayesNet = incrementalHybrid.hybridBayesNet();
 
-//   // Check if we have a bayes net with 4 hybrid nodes,
+  std::vector<std::pair<DiscreteValues, double>> assignments =
-//   // each with 2, 4, 5 (pruned), and 5 (pruned) leaves respetively.
+      discreteConditional_m1.enumerate();
-//   EXPECT_LONGS_EQUAL(4, hybridBayesNet.size());
+  // Loop over all assignments and check the pruned components
-//   EXPECT_LONGS_EQUAL(2, hybridBayesNet.atGaussian(0)->nrComponents());
+  for (auto &&av : assignments) {
-//   EXPECT_LONGS_EQUAL(4, hybridBayesNet.atGaussian(1)->nrComponents());
+    const DiscreteValues &assignment = av.first;
-//   EXPECT_LONGS_EQUAL(5, hybridBayesNet.atGaussian(2)->nrComponents());
+    const double value = av.second;
-//   EXPECT_LONGS_EQUAL(5, hybridBayesNet.atGaussian(3)->nrComponents());
 
-//   // Check that the hybrid nodes of the bayes net match those of the bayes
+    if (value == 0.0) {
-//   net
+      EXPECT(lastDensity(assignment) == nullptr);
-//   // before pruning, at the same positions.
+    } else {
-//   auto &lastDensity = *(hybridBayesNet.atGaussian(3));
+      CHECK(lastDensity(assignment));
-//   auto &unprunedLastDensity = *(unprunedHybridBayesNet->atGaussian(3));
+      EXPECT(assert_equal(*unprunedLastDensity(assignment),
-//   std::vector<std::pair<DiscreteValues, double>> assignments =
+                          *lastDensity(assignment)));
-//       discreteFactor_m1.enumerate();
+    }
-//   // Loop over all assignments and check the pruned components
+  }
-//   for (auto &&av : assignments) {
+}
-//     const DiscreteValues &assignment = av.first;
-//     const double value = av.second;
-
-//     if (value == 0.0) {
-//       EXPECT(lastDensity(assignment) == nullptr);
-//     } else {
-//       CHECK(lastDensity(assignment));
-//       EXPECT(assert_equal(*unprunedLastDensity(assignment),
-//                           *lastDensity(assignment)));
-//     }
-//   }
-// }
 
 /* ****************************************************************************/
-// // Test approximate inference with an additional pruning step.
+// Test approximate inference with an additional pruning step.
-// TEST(DCGaussianElimination, Incremental_approximate) {
+TEST(HybridGaussianElimination, Incremental_approximate) {
-//   Switching switching(5);
+  Switching switching(5);
-//   IncrementalHybrid incrementalHybrid;
+  HybridGaussianISAM incrementalHybrid;
-//   HybridGaussianFactorGraph graph1;
+  HybridGaussianFactorGraph graph1;
 
-//   // Add the 3 DC factors, x1-x2, x2-x3, x3-x4
+  /***** Run Round 1 *****/
-//   for (size_t i = 0; i < 3; i++) {
+  // Add the 3 hybrid factors, x1-x2, x2-x3, x3-x4
-//     graph1.push_back(switching.linearizedFactorGraph.dcGraph().at(i));
+  for (size_t i = 1; i < 4; i++) {
-//   }
+    graph1.push_back(switching.linearizedFactorGraph.at(i));
+  }
 
-//   // Add the Gaussian factors, 1 prior on X(1), 4 measurements
+  // Add the Gaussian factors, 1 prior on X(1),
-//   for (size_t i = 0; i <= 4; i++) {
+  // 3 measurements on X(2), X(3), X(4)
-// graph1.push_back(switching.linearizedFactorGraph.gaussianGraph().at(i));
+  graph1.push_back(switching.linearizedFactorGraph.at(0));
-//   }
+  for (size_t i = 5; i <= 7; i++) {
+    graph1.push_back(switching.linearizedFactorGraph.at(i));
+  }
 
-//   // Create ordering.
+  // Create ordering.
-//   Ordering ordering;
+  Ordering ordering;
-//   for (size_t j = 1; j <= 4; j++) {
+  for (size_t j = 1; j <= 4; j++) {
-//     ordering += X(j);
+    ordering += X(j);
-//   }
+  }
 
-//   // Run update with pruning
+  // Run update with pruning
-//   size_t maxComponents = 5;
+  size_t maxComponents = 5;
-//   incrementalHybrid.update(graph1, ordering, maxComponents);
+  incrementalHybrid.update(graph1);
+  incrementalHybrid.prune(M(3), maxComponents);
 
-//   // Check if we have a bayes net with 4 hybrid nodes,
+  // Check if we have a bayes tree with 4 hybrid nodes,
-//   // each with 2, 4, 8, and 5 (pruned) leaves respetively.
+  // each with 2, 4, 8, and 5 (pruned) leaves respetively.
-//   auto actualBayesNet1 = incrementalHybrid.hybridBayesNet();
+  EXPECT_LONGS_EQUAL(4, incrementalHybrid.size());
-//   CHECK_EQUAL(4, actualBayesNet1.size());
+  EXPECT_LONGS_EQUAL(
-//   EXPECT_LONGS_EQUAL(2, actualBayesNet1.atGaussian(0)->nrComponents());
+      2, incrementalHybrid[X(1)]->conditional()->asMixture()->nrComponents());
-//   EXPECT_LONGS_EQUAL(4, actualBayesNet1.atGaussian(1)->nrComponents());
+  EXPECT_LONGS_EQUAL(
-//   EXPECT_LONGS_EQUAL(8, actualBayesNet1.atGaussian(2)->nrComponents());
+      4, incrementalHybrid[X(2)]->conditional()->asMixture()->nrComponents());
-//   EXPECT_LONGS_EQUAL(5, actualBayesNet1.atGaussian(3)->nrComponents());
+  EXPECT_LONGS_EQUAL(
+      5, incrementalHybrid[X(3)]->conditional()->asMixture()->nrComponents());
+  EXPECT_LONGS_EQUAL(
+      5, incrementalHybrid[X(4)]->conditional()->asMixture()->nrComponents());
 
-//   /***** Run Round 2 *****/
+  /***** Run Round 2 *****/
-//   HybridGaussianFactorGraph graph2;
+  HybridGaussianFactorGraph graph2;
-//   graph2.push_back(switching.linearizedFactorGraph.dcGraph().at(3));
+  graph2.push_back(switching.linearizedFactorGraph.at(4));
-//   graph2.push_back(switching.linearizedFactorGraph.gaussianGraph().at(5));
+  graph2.push_back(switching.linearizedFactorGraph.at(8));
 
-//   Ordering ordering2;
+  Ordering ordering2;
-//   ordering2 += X(4);
+  ordering2 += X(4);
-//   ordering2 += X(5);
+  ordering2 += X(5);
 
-//   // Run update with pruning a second time.
+  // Run update with pruning a second time.
-//   incrementalHybrid.update(graph2, ordering2, maxComponents);
+  incrementalHybrid.update(graph2);
+  incrementalHybrid.prune(M(4), maxComponents);
 
-//   // Check if we have a bayes net with 2 hybrid nodes,
+  // Check if we have a bayes tree with pruned hybrid nodes,
-//   // each with 10 (pruned), and 5 (pruned) leaves respetively.
+  // with 5 (pruned) leaves.
-//   auto actualBayesNet = incrementalHybrid.hybridBayesNet();
+  CHECK_EQUAL(5, incrementalHybrid.size());
-//   CHECK_EQUAL(2, actualBayesNet.size());
+  EXPECT_LONGS_EQUAL(
-//   EXPECT_LONGS_EQUAL(10, actualBayesNet.atGaussian(0)->nrComponents());
+      5, incrementalHybrid[X(4)]->conditional()->asMixture()->nrComponents());
-//   EXPECT_LONGS_EQUAL(5, actualBayesNet.atGaussian(1)->nrComponents());
+  EXPECT_LONGS_EQUAL(
-// }
+      5, incrementalHybrid[X(5)]->conditional()->asMixture()->nrComponents());
+}
 
 /* ************************************************************************/
-// // Test for figuring out the optimal ordering to ensure we get
+// Test for figuring out the optimal ordering to ensure we get
-// // a discrete graph after elimination.
+// a discrete graph after elimination.
-// TEST(IncrementalHybrid, NonTrivial) {
+TEST(HybridGaussianISAM, NonTrivial) {
-//   // This is a GTSAM-only test for running inference on a single legged
+  // This is a GTSAM-only test for running inference on a single legged robot.
-//   robot.
+  // The leg links are represented by the chain X-Y-Z-W, where X is the base and
-//   // The leg links are represented by the chain X-Y-Z-W, where X is the base
+  // W is the foot.
-//   and
+  // We use BetweenFactor<Pose2> as constraints between each of the poses.
-//   // W is the foot.
-//   // We use BetweenFactor<Pose2> as constraints between each of the poses.
 
-//   /*************** Run Round 1 ***************/
+  /*************** Run Round 1 ***************/
-//   NonlinearHybridFactorGraph fg;
+  HybridNonlinearFactorGraph fg;
 
   //   // Add a prior on pose x1 at the origin.
   //   // A prior factor consists of a mean  and
@@ -444,9 +423,9 @@ TEST(HybridGaussianElimination, IncrementalInference) {
   //   initial.insert(W(0), Pose2(0.0, 3.0, 0.0));
 
   //   HybridGaussianFactorGraph gfg = fg.linearize(initial);
-//   fg = NonlinearHybridFactorGraph();
+  //   fg = HybridNonlinearFactorGraph();
 
-//   IncrementalHybrid inc;
+  //   HybridGaussianISAM inc;
 
   //   // Regular ordering going up the chain.
   //   Ordering ordering;
@@ -469,8 +448,8 @@ TEST(HybridGaussianElimination, IncrementalInference) {
   //   Pose2 odometry(1.0, 0.0, 0.0);
   //   KeyVector contKeys = {W(0), W(1)};
   //   auto noise_model = noiseModel::Isotropic::Sigma(3, 1.0);
-//   auto still = boost::make_shared<PlanarMotionModel>(W(0), W(1), Pose2(0, 0,
+  //   auto still = boost::make_shared<PlanarMotionModel>(W(0), W(1), Pose2(0,
-//   0),
+  //   0, 0),
   //                                                      noise_model),
   //        moving = boost::make_shared<PlanarMotionModel>(W(0), W(1), odometry,
   //                                                       noise_model);
@@ -480,7 +459,8 @@ TEST(HybridGaussianElimination, IncrementalInference) {
   //   fg.push_back(dcFactor);
 
   //   // Add equivalent of ImuFactor
-//   fg.emplace_nonlinear<BetweenFactor<Pose2>>(X(0), X(1), Pose2(1.0, 0.0, 0),
+  //   fg.emplace_nonlinear<BetweenFactor<Pose2>>(X(0), X(1), Pose2(1.0, 0.0,
+  //   0),
   //                                              poseNoise);
   //   // PoseFactors-like at k=1
   //   fg.emplace_nonlinear<BetweenFactor<Pose2>>(X(1), Y(1), Pose2(0, 1, 0),
@@ -510,11 +490,12 @@ TEST(HybridGaussianElimination, IncrementalInference) {
   //   ordering += X(1);
 
   //   gfg = fg.linearize(initial);
-//   fg = NonlinearHybridFactorGraph();
+  //   fg = HybridNonlinearFactorGraph();
 
   //   // Update without pruning
   //   // The result is a HybridBayesNet with 1 discrete variable M(1).
-//   // P(X | measurements) = P(W0|Z0, W1, M1) P(Z0|Y0, W1, M1) P(Y0|X0, W1, M1)
+  //   // P(X | measurements) = P(W0|Z0, W1, M1) P(Z0|Y0, W1, M1) P(Y0|X0, W1,
+  //   M1)
   //   //                       P(X0 | X1, W1, M1) P(W1|Z1, X1, M1) P(Z1|Y1, X1,
   //   M1)
   //   //                       P(Y1 | X1, M1)P(X1 | M1)P(M1)
@@ -535,7 +516,8 @@ TEST(HybridGaussianElimination, IncrementalInference) {
   //   fg.push_back(dcFactor);
 
   //   // Add equivalent of ImuFactor
-//   fg.emplace_nonlinear<BetweenFactor<Pose2>>(X(1), X(2), Pose2(1.0, 0.0, 0),
+  //   fg.emplace_nonlinear<BetweenFactor<Pose2>>(X(1), X(2), Pose2(1.0, 0.0,
+  //   0),
   //                                              poseNoise);
   //   // PoseFactors-like at k=1
   //   fg.emplace_nonlinear<BetweenFactor<Pose2>>(X(2), Y(2), Pose2(0, 1, 0),
@@ -562,13 +544,13 @@ TEST(HybridGaussianElimination, IncrementalInference) {
   //   ordering += X(2);
 
   //   gfg = fg.linearize(initial);
-//   fg = NonlinearHybridFactorGraph();
+  //   fg = HybridNonlinearFactorGraph();
 
   //   // Now we prune!
-//   // P(X | measurements) = P(W0|Z0, W1, M1) P(Z0|Y0, W1, M1) P(Y0|X0, W1, M1)
+  //   // P(X | measurements) = P(W0|Z0, W1, M1) P(Z0|Y0, W1, M1) P(Y0|X0, W1,
-//   P(X0 | X1, W1, M1)
+  //   M1) P(X0 | X1, W1, M1)
-//   //                       P(W1|W2, Z1, X1, M1, M2) P(Z1| W2, Y1, X1, M1, M2)
+  //   //                       P(W1|W2, Z1, X1, M1, M2) P(Z1| W2, Y1, X1, M1,
-//   P(Y1 | W2, X1, M1, M2)
+  //   M2) P(Y1 | W2, X1, M1, M2)
   //   //                       P(X1 | W2, X2, M1, M2) P(W2|Z2, X2, M1, M2)
   //   P(Z2|Y2, X2, M1, M2)
   //   //                       P(Y2 | X2, M1, M2)P(X2 | M1, M2) P(M1, M2)
@@ -590,7 +572,8 @@ TEST(HybridGaussianElimination, IncrementalInference) {
   //   fg.push_back(dcFactor);
 
   //   // Add equivalent of ImuFactor
-//   fg.emplace_nonlinear<BetweenFactor<Pose2>>(X(2), X(3), Pose2(1.0, 0.0, 0),
+  //   fg.emplace_nonlinear<BetweenFactor<Pose2>>(X(2), X(3), Pose2(1.0, 0.0,
+  //   0),
   //                                              poseNoise);
   //   // PoseFactors-like at k=3
   //   fg.emplace_nonlinear<BetweenFactor<Pose2>>(X(3), Y(3), Pose2(0, 1, 0),
@@ -617,7 +600,7 @@ TEST(HybridGaussianElimination, IncrementalInference) {
   //   ordering += X(3);
 
   //   gfg = fg.linearize(initial);
-//   fg = NonlinearHybridFactorGraph();
+  //   fg = HybridNonlinearFactorGraph();
 
   //   // Keep pruning!
   //   inc.update(gfg, ordering, 3);
@@ -634,13 +617,13 @@ TEST(HybridGaussianElimination, IncrementalInference) {
   //   expected_assignment[M(3)] = 1;
   //   EXPECT(assert_equal(expected_assignment, optimal_assignment));
 
-//   // Test if pruning worked correctly by checking that we only have 3 leaves
+  //   // Test if pruning worked correctly by checking that we only have 3
-//   in
+  //   leaves in
   //   // the last node.
   //   auto lastConditional = boost::dynamic_pointer_cast<GaussianMixture>(
   //       inc.hybridBayesNet().at(inc.hybridBayesNet().size() - 1));
   //   EXPECT_LONGS_EQUAL(3, lastConditional->nrComponents());
-// }
+}
 
 /* ************************************************************************* */
 int main() {