From d18569be62c00d48777db4f3b6230d52f3a508d7 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Wed, 1 Jan 2025 21:53:07 -0500
Subject: [PATCH] fix testGaussianMixture

---
 gtsam/hybrid/tests/testGaussianMixture.cpp | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/gtsam/hybrid/tests/testGaussianMixture.cpp b/gtsam/hybrid/tests/testGaussianMixture.cpp
index d5137ca38d..266b05c95a 100644
--- a/gtsam/hybrid/tests/testGaussianMixture.cpp
+++ b/gtsam/hybrid/tests/testGaussianMixture.cpp
@@ -20,6 +20,7 @@
 #include <gtsam/discrete/DecisionTreeFactor.h>
 #include <gtsam/discrete/DiscreteConditional.h>
 #include <gtsam/discrete/DiscreteKey.h>
+#include <gtsam/discrete/DiscreteTableConditional.h>
 #include <gtsam/hybrid/HybridBayesNet.h>
 #include <gtsam/hybrid/HybridGaussianConditional.h>
 #include <gtsam/hybrid/HybridGaussianFactorGraph.h>
@@ -79,8 +80,8 @@ TEST(GaussianMixture, GaussianMixtureModel) {
   double midway = mu1 - mu0;
   auto eliminationResult =
       gmm.toFactorGraph({{Z(0), Vector1(midway)}}).eliminateSequential();
-  auto pMid = eliminationResult->at(0)->asDiscrete();
-  EXPECT(assert_equal(DiscreteConditional(m, "60/40"), *pMid));
+  auto pMid = eliminationResult->at(0)->asDiscrete<DiscreteTableConditional>();
+  EXPECT(assert_equal(DiscreteTableConditional(m, "60/40"), *pMid));
 
   // Everywhere else, the result should be a sigmoid.
   for (const double shift : {-4, -2, 0, 2, 4}) {
@@ -90,7 +91,8 @@ TEST(GaussianMixture, GaussianMixtureModel) {
     // Workflow 1: convert HBN to HFG and solve
     auto eliminationResult1 =
         gmm.toFactorGraph({{Z(0), Vector1(z)}}).eliminateSequential();
-    auto posterior1 = *eliminationResult1->at(0)->asDiscrete();
+    auto posterior1 =
+        *eliminationResult1->at(0)->asDiscrete<DiscreteTableConditional>();
     EXPECT_DOUBLES_EQUAL(expected, posterior1(m1Assignment), 1e-8);
 
     // Workflow 2: directly specify HFG and solve
@@ -99,7 +101,8 @@ TEST(GaussianMixture, GaussianMixtureModel) {
         m, std::vector{Gaussian(mu0, sigma, z), Gaussian(mu1, sigma, z)});
     hfg1.push_back(mixing);
     auto eliminationResult2 = hfg1.eliminateSequential();
-    auto posterior2 = *eliminationResult2->at(0)->asDiscrete();
+    auto posterior2 =
+        *eliminationResult2->at(0)->asDiscrete<DiscreteTableConditional>();
     EXPECT_DOUBLES_EQUAL(expected, posterior2(m1Assignment), 1e-8);
   }
 }
@@ -138,8 +141,9 @@ TEST(GaussianMixture, GaussianMixtureModel2) {
   EXPECT(assert_equal(expectedDiscretePosterior,
                       eliminationResultMax->discretePosterior(vv)));
 
-  auto pMax = *eliminationResultMax->at(0)->asDiscrete();
-  EXPECT(assert_equal(DiscreteConditional(m, "42/58"), pMax, 1e-4));
+  auto pMax =
+      *eliminationResultMax->at(0)->asDiscrete<DiscreteTableConditional>();
+  EXPECT(assert_equal(DiscreteTableConditional(m, "42/58"), pMax, 1e-4));
 
   // Everywhere else, the result should be a bell curve like function.
   for (const double shift : {-4, -2, 0, 2, 4}) {
@@ -149,7 +153,8 @@ TEST(GaussianMixture, GaussianMixtureModel2) {
     // Workflow 1: convert HBN to HFG and solve
     auto eliminationResult1 =
         gmm.toFactorGraph({{Z(0), Vector1(z)}}).eliminateSequential();
-    auto posterior1 = *eliminationResult1->at(0)->asDiscrete();
+    auto posterior1 =
+        *eliminationResult1->at(0)->asDiscrete<DiscreteTableConditional>();
     EXPECT_DOUBLES_EQUAL(expected, posterior1(m1Assignment), 1e-8);
 
     // Workflow 2: directly specify HFG and solve
@@ -158,7 +163,8 @@ TEST(GaussianMixture, GaussianMixtureModel2) {
         m, std::vector{Gaussian(mu0, sigma0, z), Gaussian(mu1, sigma1, z)});
     hfg.push_back(mixing);
     auto eliminationResult2 = hfg.eliminateSequential();
-    auto posterior2 = *eliminationResult2->at(0)->asDiscrete();
+    auto posterior2 =
+        *eliminationResult2->at(0)->asDiscrete<DiscreteTableConditional>();
     EXPECT_DOUBLES_EQUAL(expected, posterior2(m1Assignment), 1e-8);
   }
 }