From 649da80c91ef68b2de2571e0fe14485dc652365c Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Tue, 29 Oct 2024 10:03:23 -0400
Subject: [PATCH] prune nonlinear factors in HybridNonlinearISAM

---
 gtsam/hybrid/HybridNonlinearISAM.cpp | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/gtsam/hybrid/HybridNonlinearISAM.cpp b/gtsam/hybrid/HybridNonlinearISAM.cpp
index 503afaa72..37a556096 100644
--- a/gtsam/hybrid/HybridNonlinearISAM.cpp
+++ b/gtsam/hybrid/HybridNonlinearISAM.cpp
@@ -16,6 +16,7 @@
  */
 
 #include <gtsam/hybrid/HybridGaussianFactorGraph.h>
+#include <gtsam/hybrid/HybridNonlinearFactor.h>
 #include <gtsam/hybrid/HybridNonlinearISAM.h>
 #include <gtsam/inference/Ordering.h>
 
@@ -39,7 +40,6 @@ void HybridNonlinearISAM::update(const HybridNonlinearFactorGraph& newFactors,
   if (newFactors.size() > 0) {
     // Reorder and relinearize every reorderInterval updates
     if (reorderInterval_ > 0 && ++reorderCounter_ >= reorderInterval_) {
-      // TODO(Varun) Re-linearization doesn't take into account pruning
       reorderRelinearize();
       reorderCounter_ = 0;
     }
@@ -65,8 +65,22 @@ void HybridNonlinearISAM::reorderRelinearize() {
     // Obtain the new linearization point
     const Values newLinPoint = estimate();
 
+    auto discreteProbs = *(isam_.roots().at(0)->conditional()->asDiscrete());
+
     isam_.clear();
 
+    // Prune nonlinear factors based on discrete conditional probabilities
+    HybridNonlinearFactorGraph pruned_factors;
+    for (auto&& factor : factors_) {
+      if (auto nf = std::dynamic_pointer_cast<HybridNonlinearFactor>(factor)) {
+        pruned_factors.push_back(nf->prune(discreteProbs));
+      } else {
+        pruned_factors.push_back(factor);
+      }
+    }
+    factors_ = pruned_factors;
+    factors_.print("OG factors");
+
     // Just recreate the whole BayesTree
     // TODO: allow for constrained ordering here
     // TODO: decouple re-linearization and reordering to avoid