From 4fc2387a6307e7395b50e0a07f82b28853aa51f4 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Tue, 7 Jan 2025 15:15:16 -0500
Subject: [PATCH] fix relinearization in HybridNonlinearISAM

---
 gtsam/hybrid/HybridNonlinearISAM.cpp | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/gtsam/hybrid/HybridNonlinearISAM.cpp b/gtsam/hybrid/HybridNonlinearISAM.cpp
index 29e467d86..3b4856dfb 100644
--- a/gtsam/hybrid/HybridNonlinearISAM.cpp
+++ b/gtsam/hybrid/HybridNonlinearISAM.cpp
@@ -15,6 +15,7 @@
  * @author Varun Agrawal
  */
 
+#include <gtsam/discrete/TableDistribution.h>
 #include <gtsam/hybrid/HybridGaussianFactorGraph.h>
 #include <gtsam/hybrid/HybridNonlinearFactor.h>
 #include <gtsam/hybrid/HybridNonlinearISAM.h>
@@ -65,7 +66,14 @@ void HybridNonlinearISAM::reorderRelinearize() {
     // Obtain the new linearization point
     const Values newLinPoint = estimate();
 
-    auto discreteProbs = *(isam_.roots().at(0)->conditional()->asDiscrete());
+    DiscreteConditional::shared_ptr discreteProbabilities;
+
+    auto discreteRoot = isam_.roots().at(0)->conditional();
+    if (discreteRoot->asDiscrete<TableDistribution>()) {
+      discreteProbabilities = discreteRoot->asDiscrete<TableDistribution>();
+    } else {
+      discreteProbabilities = discreteRoot->asDiscrete();
+    }
 
     isam_.clear();
 
@@ -73,7 +81,7 @@ void HybridNonlinearISAM::reorderRelinearize() {
     HybridNonlinearFactorGraph pruned_factors;
     for (auto&& factor : factors_) {
       if (auto nf = std::dynamic_pointer_cast<HybridNonlinearFactor>(factor)) {
-        pruned_factors.push_back(nf->prune(discreteProbs));
+        pruned_factors.push_back(nf->prune(*discreteProbabilities));
       } else {
         pruned_factors.push_back(factor);
       }