From fff828f5998a9e3e83336cc35565a060745e5c41 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Tue, 21 Jan 2025 15:31:01 -0500
Subject: [PATCH] move unit test for pruning

---
 gtsam/hybrid/tests/testHybridBayesNet.cpp | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/gtsam/hybrid/tests/testHybridBayesNet.cpp b/gtsam/hybrid/tests/testHybridBayesNet.cpp
index 327b5b3d0..6a108b941 100644
--- a/gtsam/hybrid/tests/testHybridBayesNet.cpp
+++ b/gtsam/hybrid/tests/testHybridBayesNet.cpp
@@ -400,6 +400,27 @@ TEST(HybridBayesNet, Error) {
   EXPECT_DOUBLES_EQUAL(density, discretePosterior(discrete_values), 1e-6);
 }
 
+/* ****************************************************************************/
+// Test Bayes net pruning
+TEST(HybridBayesNet, Prune) {
+  Switching s(3);
+
+  HybridBayesNet::shared_ptr posterior =
+      s.linearizedFactorGraph().eliminateSequential();
+  EXPECT_LONGS_EQUAL(7, posterior->size());
+
+  // Call Max-Product to get MAP
+  HybridValues delta = posterior->optimize();
+
+  // Prune the Bayes net
+  auto prunedBayesNet = posterior->prune(2);
+
+  // Test if Max-Product gives the same result as unpruned version
+  HybridValues pruned_delta = prunedBayesNet.optimize();
+  EXPECT(assert_equal(delta.discrete(), pruned_delta.discrete()));
+  EXPECT(assert_equal(delta.continuous(), pruned_delta.continuous()));
+}
+
 /* ****************************************************************************/
 // Test Bayes net error and log-probability after pruning
 TEST(HybridBayesNet, ErrorAfterPruning) {