From 5da56c139310a9701ea330ed95fce1e7de1fb135 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Wed, 21 Dec 2022 20:19:51 +0530
Subject: [PATCH 01/10] add choose method to HybridBayesTree

---
 gtsam/hybrid/HybridBayesTree.cpp           | 14 +++++-
 gtsam/hybrid/HybridBayesTree.h             | 10 +++++
 gtsam/hybrid/tests/testHybridBayesTree.cpp | 51 ++++++++++++++++++++++
 3 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/gtsam/hybrid/HybridBayesTree.cpp b/gtsam/hybrid/HybridBayesTree.cpp
index b706fb745..4ab344a1d 100644
--- a/gtsam/hybrid/HybridBayesTree.cpp
+++ b/gtsam/hybrid/HybridBayesTree.cpp
@@ -138,7 +138,8 @@ struct HybridAssignmentData {
 
 /* *************************************************************************
  */
-VectorValues HybridBayesTree::optimize(const DiscreteValues& assignment) const {
+GaussianBayesTree HybridBayesTree::choose(
+    const DiscreteValues& assignment) const {
   GaussianBayesTree gbt;
   HybridAssignmentData rootData(assignment, 0, &gbt);
   {
@@ -151,6 +152,17 @@ VectorValues HybridBayesTree::optimize(const DiscreteValues& assignment) const {
   }
 
   if (!rootData.isValid()) {
+    return GaussianBayesTree();
+  }
+  return gbt;
+}
+
+/* *************************************************************************
+ */
+VectorValues HybridBayesTree::optimize(const DiscreteValues& assignment) const {
+  GaussianBayesTree gbt = this->choose(assignment);
+  // If empty GaussianBayesTree, means a clique is pruned hence invalid
+  if (gbt.size() == 0) {
     return VectorValues();
   }
   VectorValues result = gbt.optimize();
diff --git a/gtsam/hybrid/HybridBayesTree.h b/gtsam/hybrid/HybridBayesTree.h
index 2d01aab76..628a453a6 100644
--- a/gtsam/hybrid/HybridBayesTree.h
+++ b/gtsam/hybrid/HybridBayesTree.h
@@ -24,6 +24,7 @@
 #include <gtsam/inference/BayesTree.h>
 #include <gtsam/inference/BayesTreeCliqueBase.h>
 #include <gtsam/inference/Conditional.h>
+#include <gtsam/linear/GaussianBayesTree.h>
 
 #include <string>
 
@@ -76,6 +77,15 @@ class GTSAM_EXPORT HybridBayesTree : public BayesTree<HybridBayesTreeClique> {
   /** Check equality */
   bool equals(const This& other, double tol = 1e-9) const;
 
+  /**
+   * @brief Get the Gaussian Bayes Tree which corresponds to a specific discrete
+   * value assignment.
+   *
+   * @param assignment The discrete value assignment for the discrete keys.
+   * @return GaussianBayesTree
+   */
+  GaussianBayesTree choose(const DiscreteValues& assignment) const;
+
   /**
    * @brief Optimize the hybrid Bayes tree by computing the MPE for the current
    * set of discrete variables and using it to compute the best continuous
diff --git a/gtsam/hybrid/tests/testHybridBayesTree.cpp b/gtsam/hybrid/tests/testHybridBayesTree.cpp
index 3992aa023..b4d049210 100644
--- a/gtsam/hybrid/tests/testHybridBayesTree.cpp
+++ b/gtsam/hybrid/tests/testHybridBayesTree.cpp
@@ -169,6 +169,57 @@ TEST(HybridBayesTree, Optimize) {
   EXPECT(assert_equal(expectedValues, delta.continuous()));
 }
 
+/* ****************************************************************************/
+// Test for choosing a GaussianBayesTree from a HybridBayesTree.
+TEST(HybridBayesTree, Choose) {
+  Switching s(4);
+
+  HybridGaussianISAM isam;
+  HybridGaussianFactorGraph graph1;
+
+  // Add the 3 hybrid factors, x1-x2, x2-x3, x3-x4
+  for (size_t i = 1; i < 4; i++) {
+    graph1.push_back(s.linearizedFactorGraph.at(i));
+  }
+
+  // Add the Gaussian factors, 1 prior on X(0),
+  // 3 measurements on X(2), X(3), X(4)
+  graph1.push_back(s.linearizedFactorGraph.at(0));
+  for (size_t i = 4; i <= 6; i++) {
+    graph1.push_back(s.linearizedFactorGraph.at(i));
+  }
+
+  // Add the discrete factors
+  for (size_t i = 7; i <= 9; i++) {
+    graph1.push_back(s.linearizedFactorGraph.at(i));
+  }
+
+  isam.update(graph1);
+
+  DiscreteValues assignment;
+  assignment[M(0)] = 1;
+  assignment[M(1)] = 1;
+  assignment[M(2)] = 1;
+
+  GaussianBayesTree gbt = isam.choose(assignment);
+
+  Ordering ordering;
+  ordering += X(0);
+  ordering += X(1);
+  ordering += X(2);
+  ordering += X(3);
+  ordering += M(0);
+  ordering += M(1);
+  ordering += M(2);
+  
+  //TODO(Varun) get segfault if ordering not provided
+  auto bayesTree = s.linearizedFactorGraph.eliminateMultifrontal(ordering);
+  
+  auto expected_gbt = bayesTree->choose(assignment);
+
+  EXPECT(assert_equal(expected_gbt, gbt));
+}
+
 /* ****************************************************************************/
 // Test HybridBayesTree serialization.
 TEST(HybridBayesTree, Serialization) {

From a27979e84ba4be26ec9754915477cb02d5df014a Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Wed, 21 Dec 2022 20:20:56 +0530
Subject: [PATCH 02/10] minor refactoring

---
 gtsam/hybrid/HybridGaussianFactorGraph.cpp |  7 +++++--
 gtsam/hybrid/HybridGaussianFactorGraph.h   | 11 +++++------
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/gtsam/hybrid/HybridGaussianFactorGraph.cpp b/gtsam/hybrid/HybridGaussianFactorGraph.cpp
index 39c3c2965..66beb3e4c 100644
--- a/gtsam/hybrid/HybridGaussianFactorGraph.cpp
+++ b/gtsam/hybrid/HybridGaussianFactorGraph.cpp
@@ -521,6 +521,7 @@ double HybridGaussianFactorGraph::probPrime(
     const VectorValues &continuousValues,
     const DiscreteValues &discreteValues) const {
   double error = this->error(continuousValues, discreteValues);
+  // NOTE: The 0.5 term is handled by each factor
   return std::exp(-error);
 }
 
@@ -528,8 +529,10 @@ double HybridGaussianFactorGraph::probPrime(
 AlgebraicDecisionTree<Key> HybridGaussianFactorGraph::probPrime(
     const VectorValues &continuousValues) const {
   AlgebraicDecisionTree<Key> error_tree = this->error(continuousValues);
-  AlgebraicDecisionTree<Key> prob_tree =
-      error_tree.apply([](double error) { return exp(-error); });
+  AlgebraicDecisionTree<Key> prob_tree = error_tree.apply([](double error) {
+    // NOTE: The 0.5 term is handled by each factor
+    return exp(-error);
+  });
   return prob_tree;
 }
 
diff --git a/gtsam/hybrid/HybridGaussianFactorGraph.h b/gtsam/hybrid/HybridGaussianFactorGraph.h
index b3bf8c0f5..02ebea74a 100644
--- a/gtsam/hybrid/HybridGaussianFactorGraph.h
+++ b/gtsam/hybrid/HybridGaussianFactorGraph.h
@@ -12,7 +12,7 @@
 /**
  * @file   HybridGaussianFactorGraph.h
  * @brief  Linearized Hybrid factor graph that uses type erasure
- * @author Fan Jiang
+ * @author Fan Jiang, Varun Agrawal
  * @date   Mar 11, 2022
  */
 
@@ -270,10 +270,9 @@ class GTSAM_EXPORT HybridGaussianFactorGraph
     const std::vector<DiscreteValues> assignments =
         DiscreteValues::CartesianProduct(discrete_keys);
 
-    // Save a copy of the original discrete key ordering
-    DiscreteKeys reversed_discrete_keys(discrete_keys);
-    // Reverse discrete keys order for correct tree construction
-    std::reverse(reversed_discrete_keys.begin(), reversed_discrete_keys.end());
+    // Get reversed discrete key ordering for correct tree construction
+    DiscreteKeys reversed_discrete_keys(discrete_keys.rbegin(),
+                                        discrete_keys.rend());
 
     // Create a decision tree of all the different VectorValues
     DecisionTree<Key, VectorValues::shared_ptr> delta_tree =
@@ -286,7 +285,7 @@ class GTSAM_EXPORT HybridGaussianFactorGraph
       VectorValues delta = *delta_tree(assignment);
 
       // If VectorValues is empty, it means this is a pruned branch.
-      // Set thr probPrime to 0.0.
+      // Set the probPrime to 0.0.
       if (delta.size() == 0) {
         probPrimes.push_back(0.0);
         continue;

From 80000b7e1b7a7a556b68402c3affce723207b2ef Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Wed, 21 Dec 2022 20:23:25 +0530
Subject: [PATCH 03/10] fix bug in prunerFunc due to missing keys in assignment

---
 gtsam/hybrid/HybridBayesNet.cpp | 52 +++++++++++++++++++++++++--------
 1 file changed, 40 insertions(+), 12 deletions(-)

diff --git a/gtsam/hybrid/HybridBayesNet.cpp b/gtsam/hybrid/HybridBayesNet.cpp
index f5c11c6e1..04636f74e 100644
--- a/gtsam/hybrid/HybridBayesNet.cpp
+++ b/gtsam/hybrid/HybridBayesNet.cpp
@@ -47,19 +47,21 @@ DecisionTreeFactor::shared_ptr HybridBayesNet::discreteConditionals() const {
 /**
  * @brief Helper function to get the pruner functional.
  *
- * @param decisionTree The probability decision tree of only discrete keys.
- * @return std::function<GaussianConditional::shared_ptr(
- * const Assignment<Key> &, const GaussianConditional::shared_ptr &)>
+ * @param prunedDecisionTree  The prob. decision tree of only discrete keys.
+ * @param conditional Conditional to prune. Used to get full assignment.
+ * @return std::function<double(const Assignment<Key> &, double)>
  */
 std::function<double(const Assignment<Key> &, double)> prunerFunc(
-    const DecisionTreeFactor &decisionTree,
+    const DecisionTreeFactor &prunedDecisionTree,
     const HybridConditional &conditional) {
   // Get the discrete keys as sets for the decision tree
   // and the Gaussian mixture.
-  auto decisionTreeKeySet = DiscreteKeysAsSet(decisionTree.discreteKeys());
-  auto conditionalKeySet = DiscreteKeysAsSet(conditional.discreteKeys());
+  std::set<DiscreteKey> decisionTreeKeySet =
+      DiscreteKeysAsSet(prunedDecisionTree.discreteKeys());
+  std::set<DiscreteKey> conditionalKeySet =
+      DiscreteKeysAsSet(conditional.discreteKeys());
 
-  auto pruner = [decisionTree, decisionTreeKeySet, conditionalKeySet](
+  auto pruner = [prunedDecisionTree, decisionTreeKeySet, conditionalKeySet](
                     const Assignment<Key> &choices,
                     double probability) -> double {
     // typecast so we can use this to get probability value
@@ -67,17 +69,44 @@ std::function<double(const Assignment<Key> &, double)> prunerFunc(
     // Case where the Gaussian mixture has the same
     // discrete keys as the decision tree.
     if (conditionalKeySet == decisionTreeKeySet) {
-      if (decisionTree(values) == 0) {
+      if (prunedDecisionTree(values) == 0) {
         return 0.0;
       } else {
         return probability;
       }
     } else {
+      // Due to branch merging (aka pruning) in DecisionTree, it is possible we
+      // get a `values` which doesn't have the full set of keys.
+      std::set<Key> valuesKeys;
+      for (auto kvp : values) {
+        valuesKeys.insert(kvp.first);
+      }
+      std::set<Key> conditionalKeys;
+      for (auto kvp : conditionalKeySet) {
+        conditionalKeys.insert(kvp.first);
+      }
+      // If true, then values is missing some keys
+      if (conditionalKeys != valuesKeys) {
+        // Get the keys present in conditionalKeys but not in valuesKeys
+        std::vector<Key> missing_keys;
+        std::set_difference(conditionalKeys.begin(), conditionalKeys.end(),
+                            valuesKeys.begin(), valuesKeys.end(),
+                            std::back_inserter(missing_keys));
+        // Insert missing keys with a default assignment.
+        for (auto missing_key : missing_keys) {
+          values[missing_key] = 0;
+        }
+      }
+
+      // Now we generate the full assignment by enumerating
+      // over all keys in the prunedDecisionTree.
+      // First we find the differing keys
       std::vector<DiscreteKey> set_diff;
       std::set_difference(decisionTreeKeySet.begin(), decisionTreeKeySet.end(),
                           conditionalKeySet.begin(), conditionalKeySet.end(),
                           std::back_inserter(set_diff));
 
+      // Now enumerate over all assignments of the differing keys
       const std::vector<DiscreteValues> assignments =
           DiscreteValues::CartesianProduct(set_diff);
       for (const DiscreteValues &assignment : assignments) {
@@ -86,7 +115,7 @@ std::function<double(const Assignment<Key> &, double)> prunerFunc(
 
         // If any one of the sub-branches are non-zero,
         // we need this probability.
-        if (decisionTree(augmented_values) > 0.0) {
+        if (prunedDecisionTree(augmented_values) > 0.0) {
           return probability;
         }
       }
@@ -107,10 +136,7 @@ void HybridBayesNet::updateDiscreteConditionals(
   for (size_t i = 0; i < this->size(); i++) {
     HybridConditional::shared_ptr conditional = this->at(i);
     if (conditional->isDiscrete()) {
-      // std::cout << demangle(typeid(conditional).name()) << std::endl;
       auto discrete = conditional->asDiscreteConditional();
-      KeyVector frontals(discrete->frontals().begin(),
-                         discrete->frontals().end());
 
       // Apply prunerFunc to the underlying AlgebraicDecisionTree
       auto discreteTree =
@@ -119,6 +145,8 @@ void HybridBayesNet::updateDiscreteConditionals(
           discreteTree->apply(prunerFunc(*prunedDecisionTree, *conditional));
 
       // Create the new (hybrid) conditional
+      KeyVector frontals(discrete->frontals().begin(),
+                         discrete->frontals().end());
       auto prunedDiscrete = boost::make_shared<DiscreteLookupTable>(
           frontals.size(), conditional->discreteKeys(), prunedDiscreteTree);
       conditional = boost::make_shared<HybridConditional>(prunedDiscrete);

From 0bbd279bbb52691551f6d5d73dbf6776ecedae81 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Thu, 22 Dec 2022 07:19:47 +0530
Subject: [PATCH 04/10] improved comments

---
 gtsam/hybrid/HybridGaussianFactorGraph.cpp | 1 +
 gtsam/hybrid/HybridGaussianFactorGraph.h   | 6 ++----
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/gtsam/hybrid/HybridGaussianFactorGraph.cpp b/gtsam/hybrid/HybridGaussianFactorGraph.cpp
index 66beb3e4c..a04f5776d 100644
--- a/gtsam/hybrid/HybridGaussianFactorGraph.cpp
+++ b/gtsam/hybrid/HybridGaussianFactorGraph.cpp
@@ -261,6 +261,7 @@ hybridElimination(const HybridGaussianFactorGraph &factors,
       if (!factor) {
         return 0.0;  // If nullptr, return 0.0 probability
       } else {
+        // This is the probability q(μ) at the MLE point.
         double error =
             0.5 * std::abs(factor->augmentedInformation().determinant());
         return std::exp(-error);
diff --git a/gtsam/hybrid/HybridGaussianFactorGraph.h b/gtsam/hybrid/HybridGaussianFactorGraph.h
index 02ebea74a..c448400f0 100644
--- a/gtsam/hybrid/HybridGaussianFactorGraph.h
+++ b/gtsam/hybrid/HybridGaussianFactorGraph.h
@@ -252,8 +252,8 @@ class GTSAM_EXPORT HybridGaussianFactorGraph
   }
 
   /**
-   * @brief Compute the unnormalized probabilities of the continuous variables
-   * for each of the modes.
+   * @brief Compute the unnormalized probabilities P(X | M, Z)
+   * of the continuous variables for each of the mode sequences.
    *
    * @tparam BAYES  Template on the type of Bayes graph, either a bayes net or a
    * bayes tree.
@@ -304,8 +304,6 @@ class GTSAM_EXPORT HybridGaussianFactorGraph
   std::pair<Ordering, Ordering> separateContinuousDiscreteOrdering(
       const Ordering& ordering) const;
 
-  
-
   /**
    * @brief Return a Colamd constrained ordering where the discrete keys are
    * eliminated after the continuous keys.

From b68530dcbb9ba8f2694ca3ccac55d5fc84a8a2a9 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Thu, 22 Dec 2022 15:44:25 +0530
Subject: [PATCH 05/10] cleanup and better assertions

---
 gtsam/hybrid/tests/testHybridEstimation.cpp | 81 +++++++++++++--------
 1 file changed, 52 insertions(+), 29 deletions(-)

diff --git a/gtsam/hybrid/tests/testHybridEstimation.cpp b/gtsam/hybrid/tests/testHybridEstimation.cpp
index e4a45bf4d..cc27f2fa2 100644
--- a/gtsam/hybrid/tests/testHybridEstimation.cpp
+++ b/gtsam/hybrid/tests/testHybridEstimation.cpp
@@ -72,25 +72,45 @@ Ordering getOrdering(HybridGaussianFactorGraph& factors,
 }
 
 TEST(HybridEstimation, Full) {
-  size_t K = 3;
-  std::vector<double> measurements = {0, 1, 2};
+  size_t K = 6;
+  std::vector<double> measurements = {0, 1, 2, 2, 2, 3};
   // Ground truth discrete seq
-  std::vector<size_t> discrete_seq = {1, 1, 0};
+  std::vector<size_t> discrete_seq = {1, 1, 0, 0, 1};
   // Switching example of robot moving in 1D
   // with given measurements and equal mode priors.
   Switching switching(K, 1.0, 0.1, measurements, "1/1 1/1");
   HybridGaussianFactorGraph graph = switching.linearizedFactorGraph;
 
   Ordering hybridOrdering;
-  hybridOrdering += X(0);
-  hybridOrdering += X(1);
-  hybridOrdering += X(2);
-  hybridOrdering += M(0);
-  hybridOrdering += M(1);
+  for (size_t k = 0; k < K; k++) {
+    hybridOrdering += X(k);
+  }
+  for (size_t k = 0; k < K - 1; k++) {
+    hybridOrdering += M(k);
+  }
+
   HybridBayesNet::shared_ptr bayesNet =
       graph.eliminateSequential(hybridOrdering);
 
-  EXPECT_LONGS_EQUAL(5, bayesNet->size());
+  bayesNet->print();
+  EXPECT_LONGS_EQUAL(2 * K - 1, bayesNet->size());
+
+  HybridValues delta = bayesNet->optimize();
+
+  Values initial = switching.linearizationPoint;
+  Values result = initial.retract(delta.continuous());
+
+  DiscreteValues expected_discrete;
+  for (size_t k = 0; k < K - 1; k++) {
+    expected_discrete[M(k)] = discrete_seq[k];
+  }
+  EXPECT(assert_equal(expected_discrete, delta.discrete()));
+
+  Values expected_continuous;
+  for (size_t k = 0; k < K; k++) {
+    expected_continuous.insert(X(k), measurements[k]);
+  }
+  EXPECT(assert_equal(expected_continuous, result));
 }
 
 /****************************************************************************/
@@ -102,8 +122,8 @@ TEST(HybridEstimation, Incremental) {
   // Ground truth discrete seq
   std::vector<size_t> discrete_seq = {1, 1, 0, 0, 0, 1, 1, 1, 1, 0,
                                       1, 1, 1, 0, 0, 1, 1, 0, 0, 0};
-  // Switching example of robot moving in 1D with given measurements and equal
-  // mode priors.
+  // Switching example of robot moving in 1D
+  // with given measurements and equal mode priors.
   Switching switching(K, 1.0, 0.1, measurements, "1/1 1/1");
   HybridSmoother smoother;
   HybridNonlinearFactorGraph graph;
@@ -209,13 +229,16 @@ std::vector<size_t> getDiscreteSequence(size_t x) {
 }
 
 /**
- * @brief Helper method to get the tree of unnormalized probabilities
- * as per the new elimination scheme.
+ * @brief Helper method to get the tree of
+ * unnormalized probabilities as per the elimination scheme.
+ *
+ * Used as a helper to compute q(\mu | M, Z) which is used by
+ * both P(X | M, Z) and P(M | Z).
  *
  * @param graph The HybridGaussianFactorGraph to eliminate.
  * @return AlgebraicDecisionTree<Key>
  */
-AlgebraicDecisionTree<Key> probPrimeTree(
+AlgebraicDecisionTree<Key> getProbPrimeTree(
     const HybridGaussianFactorGraph& graph) {
   HybridBayesNet::shared_ptr bayesNet;
   HybridGaussianFactorGraph::shared_ptr remainingGraph;
@@ -239,20 +262,19 @@ AlgebraicDecisionTree<Key> probPrimeTree(
   DecisionTree<Key, VectorValues::shared_ptr> delta_tree(discrete_keys,
                                                          vector_values);
 
+  // Get the probPrime tree with the correct leaf probabilities
   std::vector<double> probPrimes;
   for (const DiscreteValues& assignment : assignments) {
-    double error = 0.0;
     VectorValues delta = *delta_tree(assignment);
-    for (auto factor : graph) {
-      if (factor->isHybrid()) {
-        auto f = boost::static_pointer_cast<GaussianMixtureFactor>(factor);
-        error += f->error(delta, assignment);
 
-      } else if (factor->isContinuous()) {
-        auto f = boost::static_pointer_cast<HybridGaussianFactor>(factor);
-        error += f->inner()->error(delta);
-      }
+    // If VectorValues is empty, it means this is a pruned branch.
+    // Set the probPrime to 0.0.
+    if (delta.size() == 0) {
+      probPrimes.push_back(0.0);
+      continue;
     }
+
+    double error = graph.error(delta, assignment);
     probPrimes.push_back(exp(-error));
   }
   AlgebraicDecisionTree<Key> probPrimeTree(discrete_keys, probPrimes);
@@ -261,7 +283,8 @@ AlgebraicDecisionTree<Key> probPrimeTree(
 
 /****************************************************************************/
 /**
- * Test for correctness of different branches of the P'(Continuous | Discrete).
+ * Test for correctness of different branches of the P'(Continuous |
+ Discrete).
  * The values should match those of P'(Continuous) for each discrete mode.
  */
 TEST(HybridEstimation, Probability) {
@@ -287,8 +310,8 @@ TEST(HybridEstimation, Probability) {
     expected_prob_primes.push_back(prob_prime);
   }
 
-  // Switching example of robot moving in 1D with given measurements and equal
-  // mode priors.
+  // Switching example of robot moving in 1D with
+  // given measurements and equal mode priors.
   Switching switching(K, between_sigma, measurement_sigma, measurements,
                       "1/1 1/1");
   auto graph = switching.linearizedFactorGraph;
@@ -358,7 +381,7 @@ TEST(HybridEstimation, ProbabilityMultifrontal) {
   Ordering ordering = getOrdering(graph, HybridGaussianFactorGraph());
 
   // Get the tree of unnormalized probabilities for each mode sequence.
-  AlgebraicDecisionTree<Key> expected_probPrimeTree = probPrimeTree(graph);
+  AlgebraicDecisionTree<Key> expected_probPrimeTree = getProbPrimeTree(graph);
 
   // Eliminate continuous
   Ordering continuous_ordering(graph.continuousKeys());
@@ -412,8 +435,8 @@ TEST(HybridEstimation, ProbabilityMultifrontal) {
       discreteBayesTree->addClique(clique, discrete_clique);
 
     } else {
-      // Remove the clique from the children of the parents since it will get
-      // added again in addClique.
+      // Remove the clique from the children of the parents since
+      // it will get added again in addClique.
       auto clique_it = std::find(clique->parent()->children.begin(),
                                  clique->parent()->children.end(), clique);
       clique->parent()->children.erase(clique_it);

From 8272854378627e82c597dd542bd3a8a73c67c9cf Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Thu, 22 Dec 2022 21:01:43 +0530
Subject: [PATCH 06/10] refactor tests a bit

---
 gtsam/hybrid/tests/testHybridEstimation.cpp | 25 ++++++++++++++-------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/gtsam/hybrid/tests/testHybridEstimation.cpp b/gtsam/hybrid/tests/testHybridEstimation.cpp
index cc27f2fa2..3774ad700 100644
--- a/gtsam/hybrid/tests/testHybridEstimation.cpp
+++ b/gtsam/hybrid/tests/testHybridEstimation.cpp
@@ -92,7 +92,6 @@ TEST(HybridEstimation, Full) {
   HybridBayesNet::shared_ptr bayesNet =
       graph.eliminateSequential(hybridOrdering);
 
-  bayesNet->print();
   EXPECT_LONGS_EQUAL(2 * K - 1, bayesNet->size());
 
   HybridValues delta = bayesNet->optimize();
@@ -315,10 +314,23 @@ TEST(HybridEstimation, Probability) {
   Switching switching(K, between_sigma, measurement_sigma, measurements,
                       "1/1 1/1");
   auto graph = switching.linearizedFactorGraph;
-  Ordering ordering = getOrdering(graph, HybridGaussianFactorGraph());
 
-  HybridBayesNet::shared_ptr bayesNet = graph.eliminateSequential(ordering);
-  auto discreteConditional = bayesNet->atDiscrete(bayesNet->size() - 3);
+  // Continuous elimination
+  Ordering continuous_ordering(graph.continuousKeys());
+  HybridBayesNet::shared_ptr bayesNet;
+  HybridGaussianFactorGraph::shared_ptr discreteGraph;
+  std::tie(bayesNet, discreteGraph) =
+      graph.eliminatePartialSequential(continuous_ordering);
+
+  // Discrete elimination
+  Ordering discrete_ordering(graph.discreteKeys());
+  auto discreteBayesNet = discreteGraph->eliminateSequential(discrete_ordering);
+
+  // Add the discrete conditionals to make it a full bayes net.
+  for (auto discrete_conditional : *discreteBayesNet) {
+    bayesNet->add(discrete_conditional);
+  }
+  auto discreteConditional = discreteBayesNet->atDiscrete(0);
 
   // Test if the probPrimeTree matches the probability of
   // the individual factor graphs
@@ -413,11 +425,8 @@ TEST(HybridEstimation, ProbabilityMultifrontal) {
                          probPrimeTree(discrete_assignment), 1e-8);
   }
 
-  discreteGraph->add(DecisionTreeFactor(discrete_keys, probPrimeTree));
-
   Ordering discrete(graph.discreteKeys());
-  auto discreteBayesTree =
-      discreteGraph->BaseEliminateable::eliminateMultifrontal(discrete);
+  auto discreteBayesTree = discreteGraph->eliminateMultifrontal(discrete);
 
   EXPECT_LONGS_EQUAL(1, discreteBayesTree->size());
   // DiscreteBayesTree should have only 1 clique

From 0fb67995c68fa300c336d7df4f4a61c77e42bfc7 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Fri, 23 Dec 2022 00:12:51 +0530
Subject: [PATCH 07/10] added comments and minor refactor

---
 gtsam/hybrid/HybridBayesNet.cpp | 6 +++---
 gtsam/hybrid/HybridSmoother.cpp | 3 +--
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/gtsam/hybrid/HybridBayesNet.cpp b/gtsam/hybrid/HybridBayesNet.cpp
index 04636f74e..dffbb366a 100644
--- a/gtsam/hybrid/HybridBayesNet.cpp
+++ b/gtsam/hybrid/HybridBayesNet.cpp
@@ -242,7 +242,7 @@ GaussianBayesNet HybridBayesNet::choose(
 
 /* ************************************************************************* */
 HybridValues HybridBayesNet::optimize() const {
-  // Solve for the MPE
+  // Collect all the discrete factors to compute MPE
   DiscreteBayesNet discrete_bn;
   for (auto &&conditional : *this) {
     if (conditional->isDiscrete()) {
@@ -250,11 +250,11 @@ HybridValues HybridBayesNet::optimize() const {
     }
   }
 
+  // Solve for the MPE
   DiscreteValues mpe = DiscreteFactorGraph(discrete_bn).optimize();
 
   // Given the MPE, compute the optimal continuous values.
-  GaussianBayesNet gbn = this->choose(mpe);
-  return HybridValues(mpe, gbn.optimize());
+  return HybridValues(mpe, this->optimize(mpe));
 }
 
 /* ************************************************************************* */
diff --git a/gtsam/hybrid/HybridSmoother.cpp b/gtsam/hybrid/HybridSmoother.cpp
index 07a7a4e77..ef77a2413 100644
--- a/gtsam/hybrid/HybridSmoother.cpp
+++ b/gtsam/hybrid/HybridSmoother.cpp
@@ -100,8 +100,7 @@ HybridSmoother::addConditionals(const HybridGaussianFactorGraph &originalGraph,
 /* ************************************************************************* */
 GaussianMixture::shared_ptr HybridSmoother::gaussianMixture(
     size_t index) const {
-  return boost::dynamic_pointer_cast<GaussianMixture>(
-      hybridBayesNet_.at(index));
+  return hybridBayesNet_.atMixture(index);
 }
 
 /* ************************************************************************* */

From 5ea63be8c52ca4d8a49a7cba191cdcaa0d64e3cc Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Fri, 30 Dec 2022 14:23:40 +0530
Subject: [PATCH 08/10] fixes based on previous PR

---
 gtsam/hybrid/HybridBayesNet.cpp              |  2 -
 gtsam/hybrid/HybridValues.h                  |  2 +-
 python/gtsam/tests/test_HybridFactorGraph.py | 78 ++++++++++----------
 3 files changed, 41 insertions(+), 41 deletions(-)

diff --git a/gtsam/hybrid/HybridBayesNet.cpp b/gtsam/hybrid/HybridBayesNet.cpp
index 2e47793ef..fd78986df 100644
--- a/gtsam/hybrid/HybridBayesNet.cpp
+++ b/gtsam/hybrid/HybridBayesNet.cpp
@@ -138,8 +138,6 @@ void HybridBayesNet::updateDiscreteConditionals(
     HybridConditional::shared_ptr conditional = this->at(i);
     if (conditional->isDiscrete()) {
       auto discrete = conditional->asDiscrete();
-      KeyVector frontals(discrete->frontals().begin(),
-                         discrete->frontals().end());
 
       // Apply prunerFunc to the underlying AlgebraicDecisionTree
       auto discreteTree =
diff --git a/gtsam/hybrid/HybridValues.h b/gtsam/hybrid/HybridValues.h
index 80c942a83..adeee5963 100644
--- a/gtsam/hybrid/HybridValues.h
+++ b/gtsam/hybrid/HybridValues.h
@@ -104,7 +104,7 @@ class GTSAM_EXPORT HybridValues {
    * @param j The index with which the value will be associated. */
   void insert(Key j, const Vector& value) { continuous_.insert(j, value); }
 
-  // TODO(Shangjie)- update() and insert_or_assign() , similar to Values.h
+  // TODO(Shangjie)- insert_or_assign() , similar to Values.h
 
   /**
    * Read/write access to the discrete value with key \c j, throws
diff --git a/python/gtsam/tests/test_HybridFactorGraph.py b/python/gtsam/tests/test_HybridFactorGraph.py
index 2ebc87971..40016c8bf 100644
--- a/python/gtsam/tests/test_HybridFactorGraph.py
+++ b/python/gtsam/tests/test_HybridFactorGraph.py
@@ -11,30 +11,20 @@ Author: Fan Jiang
 # pylint: disable=invalid-name, no-name-in-module, no-member
 
 import unittest
-import math
 
 import numpy as np
 from gtsam.symbol_shorthand import C, M, X, Z
 from gtsam.utils.test_case import GtsamTestCase
 
 import gtsam
-from gtsam import (
-    DecisionTreeFactor,
-    DiscreteConditional,
-    DiscreteKeys,
-    GaussianConditional,
-    GaussianMixture,
-    GaussianMixtureFactor,
-    HybridGaussianFactorGraph,
-    JacobianFactor,
-    Ordering,
-    noiseModel,
-)
+from gtsam import (DiscreteConditional, DiscreteKeys, GaussianConditional,
+                   GaussianMixture, GaussianMixtureFactor,
+                   HybridGaussianFactorGraph, JacobianFactor, Ordering,
+                   noiseModel)
 
 
 class TestHybridGaussianFactorGraph(GtsamTestCase):
     """Unit tests for HybridGaussianFactorGraph."""
-
     def test_create(self):
         """Test construction of hybrid factor graph."""
         model = noiseModel.Unit.Create(3)
@@ -52,8 +42,8 @@ class TestHybridGaussianFactorGraph(GtsamTestCase):
         hfg.push_back(gmf)
 
         hbn = hfg.eliminateSequential(
-            Ordering.ColamdConstrainedLastHybridGaussianFactorGraph(hfg, [C(0)])
-        )
+            Ordering.ColamdConstrainedLastHybridGaussianFactorGraph(
+                hfg, [C(0)]))
 
         self.assertEqual(hbn.size(), 2)
 
@@ -84,36 +74,39 @@ class TestHybridGaussianFactorGraph(GtsamTestCase):
         hfg.push_back(dtf)
 
         hbn = hfg.eliminateSequential(
-            Ordering.ColamdConstrainedLastHybridGaussianFactorGraph(hfg, [C(0)])
-        )
+            Ordering.ColamdConstrainedLastHybridGaussianFactorGraph(
+                hfg, [C(0)]))
 
         hv = hbn.optimize()
         self.assertEqual(hv.atDiscrete(C(0)), 1)
 
     @staticmethod
-    def tiny(num_measurements: int = 1):
-        """Create a tiny two variable hybrid model."""
+    def tiny(num_measurements: int = 1) -> gtsam.HybridBayesNet:
+        """
+        Create a tiny two variable hybrid model which represents
+        the generative probability P(z, x, n) = P(z | x, n)P(x)P(n).
+        """
         # Create hybrid Bayes net.
         bayesNet = gtsam.HybridBayesNet()
 
         # Create mode key: 0 is low-noise, 1 is high-noise.
-        modeKey = M(0)
-        mode = (modeKey, 2)
+        mode = (M(0), 2)
 
         # Create Gaussian mixture Z(0) = X(0) + noise for each measurement.
         I = np.eye(1)
         keys = DiscreteKeys()
         keys.push_back(mode)
         for i in range(num_measurements):
-            conditional0 = GaussianConditional.FromMeanAndStddev(
-                Z(i), I, X(0), [0], sigma=0.5
-            )
-            conditional1 = GaussianConditional.FromMeanAndStddev(
-                Z(i), I, X(0), [0], sigma=3
-            )
-            bayesNet.emplaceMixture(
-                [Z(i)], [X(0)], keys, [conditional0, conditional1]
-            )
+            conditional0 = GaussianConditional.FromMeanAndStddev(Z(i),
+                                                                 I,
+                                                                 X(0), [0],
+                                                                 sigma=0.5)
+            conditional1 = GaussianConditional.FromMeanAndStddev(Z(i),
+                                                                 I,
+                                                                 X(0), [0],
+                                                                 sigma=3)
+            bayesNet.emplaceMixture([Z(i)], [X(0)], keys,
+                                    [conditional0, conditional1])
 
         # Create prior on X(0).
         prior_on_x0 = GaussianConditional.FromMeanAndStddev(X(0), [5.0], 5.0)
@@ -142,14 +135,22 @@ class TestHybridGaussianFactorGraph(GtsamTestCase):
 
         self.assertEqual(fg.size(), 3)
 
-    def test_tiny2(self):
-        """Test a tiny two variable hybrid model, with 2 measurements."""
-        # Create the Bayes net and sample from it.
+    def test_ratio(self):
+        """
+        Given a tiny two variable hybrid model, with 2 measurements,
+        test the ratio of the bayes net model representing P(z, x, n)=P(z|x, n)P(x)P(n)
+        and the factor graph P(x, n | z)=P(x | n, z)P(n|z),
+        both of which represent the same posterior.
+        """
+        # Create the Bayes net representing the generative model P(z, x, n)=P(z|x, n)P(x)P(n)
         bayesNet = self.tiny(num_measurements=2)
-        sample = bayesNet.sample()
+        # Sample from the Bayes net.
+        sample: gtsam.HybridValues = bayesNet.sample()
         # print(sample)
 
         # Create a factor graph from the Bayes net with sampled measurements.
+        # The factor graph is `P(x)P(n) ϕ(x, n; z1) ϕ(x, n; z2)`
+        # and thus represents the same joint probability as the Bayes net.
         fg = HybridGaussianFactorGraph()
         for i in range(2):
             conditional = bayesNet.atMixture(i)
@@ -161,13 +162,14 @@ class TestHybridGaussianFactorGraph(GtsamTestCase):
         fg.push_back(bayesNet.atDiscrete(3))
 
         self.assertEqual(fg.size(), 4)
-        # Calculate ratio  between Bayes net probability and the factor graph:
+
+        # Calculate ratio between Bayes net probability and the factor graph:
         continuousValues = gtsam.VectorValues()
         continuousValues.insert(X(0), sample.at(X(0)))
         discreteValues = sample.discrete()
         expected_ratio = bayesNet.evaluate(sample) / fg.probPrime(
-            continuousValues, discreteValues
-        )
+            continuousValues, discreteValues)
+        #TODO(Varun) This should be 1. Adding the normalizing factor should fix fg.probPrime
         print(expected_ratio)
 
         # TODO(dellaert): Change the mode to 0 and calculate the ratio again.

From 9b92d153fc60c11715474f1be5c440d9145a599d Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Fri, 30 Dec 2022 14:34:11 +0530
Subject: [PATCH 09/10] improve tests

---
 gtsam/hybrid/tests/testHybridBayesNet.cpp   | 2 ++
 gtsam/hybrid/tests/testHybridEstimation.cpp | 3 +--
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/gtsam/hybrid/tests/testHybridBayesNet.cpp b/gtsam/hybrid/tests/testHybridBayesNet.cpp
index 0170ce423..43cee6f74 100644
--- a/gtsam/hybrid/tests/testHybridBayesNet.cpp
+++ b/gtsam/hybrid/tests/testHybridBayesNet.cpp
@@ -188,12 +188,14 @@ TEST(HybridBayesNet, Optimize) {
 
   HybridValues delta = hybridBayesNet->optimize();
 
+  //TODO(Varun) The expectedAssignment should be 111, not 101
   DiscreteValues expectedAssignment;
   expectedAssignment[M(0)] = 1;
   expectedAssignment[M(1)] = 0;
   expectedAssignment[M(2)] = 1;
   EXPECT(assert_equal(expectedAssignment, delta.discrete()));
 
+  //TODO(Varun) This should be all -Vector1::Ones()
   VectorValues expectedValues;
   expectedValues.insert(X(0), -0.999904 * Vector1::Ones());
   expectedValues.insert(X(1), -0.99029 * Vector1::Ones());
diff --git a/gtsam/hybrid/tests/testHybridEstimation.cpp b/gtsam/hybrid/tests/testHybridEstimation.cpp
index 5429c2266..867760e29 100644
--- a/gtsam/hybrid/tests/testHybridEstimation.cpp
+++ b/gtsam/hybrid/tests/testHybridEstimation.cpp
@@ -282,8 +282,7 @@ AlgebraicDecisionTree<Key> getProbPrimeTree(
 
 /****************************************************************************/
 /**
- * Test for correctness of different branches of the P'(Continuous |
- Discrete).
+ * Test for correctness of different branches of the P'(Continuous | Discrete).
  * The values should match those of P'(Continuous) for each discrete mode.
  */
 TEST(HybridEstimation, Probability) {

From 6f8a23fe3463f8efc415857bcac525582437aa0f Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Fri, 30 Dec 2022 17:32:31 +0530
Subject: [PATCH 10/10] minor fixes

---
 gtsam/hybrid/HybridBayesNet.cpp             |  1 -
 gtsam/hybrid/HybridGaussianFactorGraph.cpp  | 16 +++++++---------
 gtsam/hybrid/tests/testHybridEstimation.cpp |  2 +-
 3 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/gtsam/hybrid/HybridBayesNet.cpp b/gtsam/hybrid/HybridBayesNet.cpp
index fd78986df..8e01c0c76 100644
--- a/gtsam/hybrid/HybridBayesNet.cpp
+++ b/gtsam/hybrid/HybridBayesNet.cpp
@@ -128,7 +128,6 @@ std::function<double(const Assignment<Key> &, double)> prunerFunc(
 }
 
 /* ************************************************************************* */
-// TODO(dellaert): what is this non-const method used for? Abolish it?
 void HybridBayesNet::updateDiscreteConditionals(
     const DecisionTreeFactor::shared_ptr &prunedDecisionTree) {
   KeyVector prunedTreeKeys = prunedDecisionTree->keys();
diff --git a/gtsam/hybrid/HybridGaussianFactorGraph.cpp b/gtsam/hybrid/HybridGaussianFactorGraph.cpp
index 9d010d2a1..aac37bc24 100644
--- a/gtsam/hybrid/HybridGaussianFactorGraph.cpp
+++ b/gtsam/hybrid/HybridGaussianFactorGraph.cpp
@@ -397,18 +397,16 @@ EliminateHybrid(const HybridGaussianFactorGraph &factors,
   if (discrete_only) {
     // Case 1: we are only dealing with discrete
     return discreteElimination(factors, frontalKeys);
-  } else {
+  } else if (mapFromKeyToDiscreteKey.empty()) {
     // Case 2: we are only dealing with continuous
-    if (mapFromKeyToDiscreteKey.empty()) {
-      return continuousElimination(factors, frontalKeys);
-    } else {
-      // Case 3: We are now in the hybrid land!
+    return continuousElimination(factors, frontalKeys);
+  } else {
+    // Case 3: We are now in the hybrid land!
 #ifdef HYBRID_TIMING
-      tictoc_reset_();
+    tictoc_reset_();
 #endif
-      return hybridElimination(factors, frontalKeys, continuousSeparator,
-                               discreteSeparatorSet);
-    }
+    return hybridElimination(factors, frontalKeys, continuousSeparator,
+                             discreteSeparatorSet);
   }
 }
 
diff --git a/gtsam/hybrid/tests/testHybridEstimation.cpp b/gtsam/hybrid/tests/testHybridEstimation.cpp
index 867760e29..927f5c047 100644
--- a/gtsam/hybrid/tests/testHybridEstimation.cpp
+++ b/gtsam/hybrid/tests/testHybridEstimation.cpp
@@ -425,7 +425,7 @@ static HybridNonlinearFactorGraph createHybridNonlinearFactorGraph() {
 }
 
 /*********************************************************************************
-  // Create a hybrid nonlinear factor graph f(x0, x1, m0; z0, z1)
+  // Create a hybrid linear factor graph f(x0, x1, m0; z0, z1)
  ********************************************************************************/
 static HybridGaussianFactorGraph::shared_ptr createHybridGaussianFactorGraph() {
   HybridNonlinearFactorGraph nfg = createHybridNonlinearFactorGraph();