From 2d7690dbb7c2468a86df5b63ae585a81c10d71cb Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Thu, 23 Jan 2025 20:32:49 -0500
Subject: [PATCH] update addConditionals to only use factor graph keys and
 remove an extra loop

---
 gtsam/hybrid/HybridSmoother.cpp | 33 ++++++++++++++-------------------
 gtsam/hybrid/HybridSmoother.h   |  2 +-
 2 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/gtsam/hybrid/HybridSmoother.cpp b/gtsam/hybrid/HybridSmoother.cpp
index ca3e27252..3e06d98ae 100644
--- a/gtsam/hybrid/HybridSmoother.cpp
+++ b/gtsam/hybrid/HybridSmoother.cpp
@@ -68,8 +68,7 @@ void HybridSmoother::update(HybridGaussianFactorGraph graph,
   }
 
   // Add the necessary conditionals from the previous timestep(s).
-  std::tie(graph, hybridBayesNet_) =
-      addConditionals(graph, hybridBayesNet_, ordering);
+  std::tie(graph, hybridBayesNet_) = addConditionals(graph, hybridBayesNet_);
 
   // Eliminate.
   HybridBayesNet bayesNetFragment = *graph.eliminateSequential(ordering);
@@ -88,10 +87,11 @@ void HybridSmoother::update(HybridGaussianFactorGraph graph,
 /* ************************************************************************* */
 std::pair<HybridGaussianFactorGraph, HybridBayesNet>
 HybridSmoother::addConditionals(const HybridGaussianFactorGraph &originalGraph,
-                                const HybridBayesNet &originalHybridBayesNet,
-                                const Ordering &ordering) const {
+                                const HybridBayesNet &hybridBayesNet) const {
   HybridGaussianFactorGraph graph(originalGraph);
-  HybridBayesNet hybridBayesNet(originalHybridBayesNet);
+  HybridBayesNet updatedHybridBayesNet(hybridBayesNet);
+
+  KeySet factorKeys = graph.keys();
 
   // If hybridBayesNet is not empty,
   // it means we have conditionals to add to the factor graph.
@@ -99,10 +99,6 @@ HybridSmoother::addConditionals(const HybridGaussianFactorGraph &originalGraph,
     // We add all relevant hybrid conditionals on the last continuous variable
     // in the previous `hybridBayesNet` to the graph
 
-    // Conditionals to remove from the bayes net
-    // since the conditional will be updated.
-    std::vector<HybridConditional::shared_ptr> conditionals_to_erase;
-
     // New conditionals to add to the graph
     gtsam::HybridBayesNet newConditionals;
 
@@ -112,25 +108,24 @@ HybridSmoother::addConditionals(const HybridGaussianFactorGraph &originalGraph,
       auto conditional = hybridBayesNet.at(i);
 
       for (auto &key : conditional->frontals()) {
-        if (std::find(ordering.begin(), ordering.end(), key) !=
-            ordering.end()) {
+        if (std::find(factorKeys.begin(), factorKeys.end(), key) !=
+            factorKeys.end()) {
           newConditionals.push_back(conditional);
-          conditionals_to_erase.push_back(conditional);
+
+          // Remove the conditional from the updated Bayes net
+          auto it = find(updatedHybridBayesNet.begin(),
+                         updatedHybridBayesNet.end(), conditional);
+          updatedHybridBayesNet.erase(it);
 
           break;
         }
       }
     }
-    // Remove conditionals at the end so we don't affect the order in the
-    // original bayes net.
-    for (auto &&conditional : conditionals_to_erase) {
-      auto it = find(hybridBayesNet.begin(), hybridBayesNet.end(), conditional);
-      hybridBayesNet.erase(it);
-    }
 
     graph.push_back(newConditionals);
   }
-  return {graph, hybridBayesNet};
+
+  return {graph, updatedHybridBayesNet};
 }
 
 /* ************************************************************************* */
diff --git a/gtsam/hybrid/HybridSmoother.h b/gtsam/hybrid/HybridSmoother.h
index 66edf86d6..4669b1d8f 100644
--- a/gtsam/hybrid/HybridSmoother.h
+++ b/gtsam/hybrid/HybridSmoother.h
@@ -66,7 +66,7 @@ class GTSAM_EXPORT HybridSmoother {
    */
   std::pair<HybridGaussianFactorGraph, HybridBayesNet> addConditionals(
       const HybridGaussianFactorGraph& graph,
-      const HybridBayesNet& hybridBayesNet, const Ordering& ordering) const;
+      const HybridBayesNet& hybridBayesNet) const;
 
   /**
    * @brief Get the hybrid Gaussian conditional from