From 8c38e45c83c6bbeb2ef5e7c7d73807b2eed6831f Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Wed, 28 Jun 2023 12:19:41 -0400
Subject: [PATCH 1/4] enumerate all assignments for computing probabilities to
 prune

---
 gtsam/discrete/DecisionTreeFactor.cpp         | 11 ++--
 .../tests/testGaussianMixtureFactor.cpp       |  4 +-
 .../tests/testHybridNonlinearFactorGraph.cpp  | 58 +++++++++----------
 gtsam/hybrid/tests/testMixtureFactor.cpp      |  4 +-
 4 files changed, 39 insertions(+), 38 deletions(-)

diff --git a/gtsam/discrete/DecisionTreeFactor.cpp b/gtsam/discrete/DecisionTreeFactor.cpp
index 4aca10a284..6f20907e14 100644
--- a/gtsam/discrete/DecisionTreeFactor.cpp
+++ b/gtsam/discrete/DecisionTreeFactor.cpp
@@ -306,11 +306,12 @@ namespace gtsam {
 
     // Get the probabilities in the decision tree so we can threshold.
     std::vector<double> probabilities;
-    this->visitLeaf([&](const Leaf& leaf) {
-      const size_t nrAssignments = leaf.nrAssignments();
-      double prob = leaf.constant();
-      probabilities.insert(probabilities.end(), nrAssignments, prob);
-    });
+    // NOTE(Varun) this is potentially slow due to the cartesian product
+    auto allValues = DiscreteValues::CartesianProduct(this->discreteKeys());
+    for (auto&& val : allValues) {
+      double prob = (*this)(val);
+      probabilities.push_back(prob);
+    }
 
     // The number of probabilities can be lower than max_leaves
     if (probabilities.size() <= N) {
diff --git a/gtsam/hybrid/tests/testGaussianMixtureFactor.cpp b/gtsam/hybrid/tests/testGaussianMixtureFactor.cpp
index 5207e9372a..75ba5a0594 100644
--- a/gtsam/hybrid/tests/testGaussianMixtureFactor.cpp
+++ b/gtsam/hybrid/tests/testGaussianMixtureFactor.cpp
@@ -108,7 +108,7 @@ TEST(GaussianMixtureFactor, Printing) {
   std::string expected =
       R"(Hybrid [x1 x2; 1]{
  Choice(1) 
- 0 Leaf [1]:
+ 0 Leaf :
   A[x1] = [
 	0;
 	0
@@ -120,7 +120,7 @@ TEST(GaussianMixtureFactor, Printing) {
   b = [ 0 0 ]
   No noise model
 
- 1 Leaf [1]:
+ 1 Leaf :
   A[x1] = [
 	0;
 	0
diff --git a/gtsam/hybrid/tests/testHybridNonlinearFactorGraph.cpp b/gtsam/hybrid/tests/testHybridNonlinearFactorGraph.cpp
index 0a621c42d2..29136c62ad 100644
--- a/gtsam/hybrid/tests/testHybridNonlinearFactorGraph.cpp
+++ b/gtsam/hybrid/tests/testHybridNonlinearFactorGraph.cpp
@@ -493,7 +493,7 @@ factor 0:
 factor 1: 
 Hybrid [x0 x1; m0]{
  Choice(m0) 
- 0 Leaf [1]:
+ 0 Leaf :
   A[x0] = [
 	-1
 ]
@@ -503,7 +503,7 @@ Hybrid [x0 x1; m0]{
   b = [ -1 ]
   No noise model
 
- 1 Leaf [1]:
+ 1 Leaf :
   A[x0] = [
 	-1
 ]
@@ -517,7 +517,7 @@ Hybrid [x0 x1; m0]{
 factor 2: 
 Hybrid [x1 x2; m1]{
  Choice(m1) 
- 0 Leaf [1]:
+ 0 Leaf :
   A[x1] = [
 	-1
 ]
@@ -527,7 +527,7 @@ Hybrid [x1 x2; m1]{
   b = [ -1 ]
   No noise model
 
- 1 Leaf [1]:
+ 1 Leaf :
   A[x1] = [
 	-1
 ]
@@ -551,16 +551,16 @@ factor 4:
   b = [ -10 ]
   No noise model
 factor 5:  P( m0 ):
- Leaf [2] 0.5
+ Leaf  0.5
 
 factor 6:  P( m1 | m0 ):
  Choice(m1) 
  0 Choice(m0) 
- 0 0 Leaf [1]0.33333333
- 0 1 Leaf [1] 0.6
+ 0 0 Leaf 0.33333333
+ 0 1 Leaf  0.6
  1 Choice(m0) 
- 1 0 Leaf [1]0.66666667
- 1 1 Leaf [1] 0.4
+ 1 0 Leaf 0.66666667
+ 1 1 Leaf  0.4
 
 )";
 #else
@@ -575,7 +575,7 @@ factor 0:
 factor 1: 
 Hybrid [x0 x1; m0]{
  Choice(m0) 
- 0 Leaf [1]:
+ 0 Leaf :
   A[x0] = [
 	-1
 ]
@@ -585,7 +585,7 @@ Hybrid [x0 x1; m0]{
   b = [ -1 ]
   No noise model
 
- 1 Leaf [1]:
+ 1 Leaf :
   A[x0] = [
 	-1
 ]
@@ -599,7 +599,7 @@ Hybrid [x0 x1; m0]{
 factor 2: 
 Hybrid [x1 x2; m1]{
  Choice(m1) 
- 0 Leaf [1]:
+ 0 Leaf :
   A[x1] = [
 	-1
 ]
@@ -609,7 +609,7 @@ Hybrid [x1 x2; m1]{
   b = [ -1 ]
   No noise model
 
- 1 Leaf [1]:
+ 1 Leaf :
   A[x1] = [
 	-1
 ]
@@ -634,17 +634,17 @@ factor 4:
   No noise model
 factor 5:  P( m0 ):
  Choice(m0) 
- 0 Leaf [1] 0.5
- 1 Leaf [1] 0.5
+ 0 Leaf  0.5
+ 1 Leaf  0.5
 
 factor 6:  P( m1 | m0 ):
  Choice(m1) 
  0 Choice(m0) 
- 0 0 Leaf [1]0.33333333
- 0 1 Leaf [1] 0.6
+ 0 0 Leaf 0.33333333
+ 0 1 Leaf  0.6
  1 Choice(m0) 
- 1 0 Leaf [1]0.66666667
- 1 1 Leaf [1] 0.4
+ 1 0 Leaf 0.66666667
+ 1 1 Leaf  0.4
 
 )";
 #endif
@@ -657,13 +657,13 @@ size: 3
 conditional 0: Hybrid  P( x0 | x1 m0)
  Discrete Keys = (m0, 2), 
  Choice(m0) 
- 0 Leaf [1] p(x0 | x1)
+ 0 Leaf  p(x0 | x1)
   R = [ 10.0499 ]
   S[x1] = [ -0.0995037 ]
   d = [ -9.85087 ]
   No noise model
 
- 1 Leaf [1] p(x0 | x1)
+ 1 Leaf  p(x0 | x1)
   R = [ 10.0499 ]
   S[x1] = [ -0.0995037 ]
   d = [ -9.95037 ]
@@ -673,26 +673,26 @@ conditional 1: Hybrid  P( x1 | x2 m0 m1)
  Discrete Keys = (m0, 2), (m1, 2), 
  Choice(m1) 
  0 Choice(m0) 
- 0 0 Leaf [1] p(x1 | x2)
+ 0 0 Leaf  p(x1 | x2)
   R = [ 10.099 ]
   S[x2] = [ -0.0990196 ]
   d = [ -9.99901 ]
   No noise model
 
- 0 1 Leaf [1] p(x1 | x2)
+ 0 1 Leaf  p(x1 | x2)
   R = [ 10.099 ]
   S[x2] = [ -0.0990196 ]
   d = [ -9.90098 ]
   No noise model
 
  1 Choice(m0) 
- 1 0 Leaf [1] p(x1 | x2)
+ 1 0 Leaf  p(x1 | x2)
   R = [ 10.099 ]
   S[x2] = [ -0.0990196 ]
   d = [ -10.098 ]
   No noise model
 
- 1 1 Leaf [1] p(x1 | x2)
+ 1 1 Leaf  p(x1 | x2)
   R = [ 10.099 ]
   S[x2] = [ -0.0990196 ]
   d = [ -10 ]
@@ -702,14 +702,14 @@ conditional 2: Hybrid  P( x2 | m0 m1)
  Discrete Keys = (m0, 2), (m1, 2), 
  Choice(m1) 
  0 Choice(m0) 
- 0 0 Leaf [1] p(x2)
+ 0 0 Leaf  p(x2)
   R = [ 10.0494 ]
   d = [ -10.1489 ]
   mean: 1 elements
   x2: -1.0099
   No noise model
 
- 0 1 Leaf [1] p(x2)
+ 0 1 Leaf  p(x2)
   R = [ 10.0494 ]
   d = [ -10.1479 ]
   mean: 1 elements
@@ -717,14 +717,14 @@ conditional 2: Hybrid  P( x2 | m0 m1)
   No noise model
 
  1 Choice(m0) 
- 1 0 Leaf [1] p(x2)
+ 1 0 Leaf  p(x2)
   R = [ 10.0494 ]
   d = [ -10.0504 ]
   mean: 1 elements
   x2: -1.0001
   No noise model
 
- 1 1 Leaf [1] p(x2)
+ 1 1 Leaf  p(x2)
   R = [ 10.0494 ]
   d = [ -10.0494 ]
   mean: 1 elements
diff --git a/gtsam/hybrid/tests/testMixtureFactor.cpp b/gtsam/hybrid/tests/testMixtureFactor.cpp
index 03fdccff26..0827e4eedb 100644
--- a/gtsam/hybrid/tests/testMixtureFactor.cpp
+++ b/gtsam/hybrid/tests/testMixtureFactor.cpp
@@ -63,8 +63,8 @@ TEST(MixtureFactor, Printing) {
       R"(Hybrid [x1 x2; 1]
 MixtureFactor
  Choice(1) 
- 0 Leaf [1] Nonlinear factor on 2 keys
- 1 Leaf [1] Nonlinear factor on 2 keys
+ 0 Leaf  Nonlinear factor on 2 keys
+ 1 Leaf  Nonlinear factor on 2 keys
 )";
   EXPECT(assert_print_equal(expected, mixtureFactor));
 }

From 647d3c0744171ffa78f585a39a88bfa4be4d2002 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Wed, 28 Jun 2023 13:26:07 -0400
Subject: [PATCH 2/4] remove nrAssignments from the DecisionTree

---
 gtsam/discrete/DecisionTree-inl.h             |  47 ++------
 gtsam/discrete/DecisionTree.h                 |  36 ------
 gtsam/discrete/tests/testDecisionTree.cpp     |  85 -------------
 .../tests/testDiscreteFactorGraph.cpp         | 113 ------------------
 4 files changed, 8 insertions(+), 273 deletions(-)

diff --git a/gtsam/discrete/DecisionTree-inl.h b/gtsam/discrete/DecisionTree-inl.h
index 156177d039..10532284cf 100644
--- a/gtsam/discrete/DecisionTree-inl.h
+++ b/gtsam/discrete/DecisionTree-inl.h
@@ -53,26 +53,17 @@ namespace gtsam {
     /** constant stored in this leaf */
     Y constant_;
 
-    /** The number of assignments contained within this leaf.
-     * Particularly useful when leaves have been pruned.
-     */
-    size_t nrAssignments_;
-
     /// Default constructor for serialization.
     Leaf() {}
 
     /// Constructor from constant
-    Leaf(const Y& constant, size_t nrAssignments = 1)
-        : constant_(constant), nrAssignments_(nrAssignments) {}
+    Leaf(const Y& constant) : constant_(constant) {}
 
     /// Return the constant
     const Y& constant() const {
       return constant_;
     }
 
-    /// Return the number of assignments contained within this leaf.
-    size_t nrAssignments() const { return nrAssignments_; }
-
     /// Leaf-Leaf equality
     bool sameLeaf(const Leaf& q) const override {
       return constant_ == q.constant_;
@@ -93,8 +84,7 @@ namespace gtsam {
     /// print
     void print(const std::string& s, const LabelFormatter& labelFormatter,
                const ValueFormatter& valueFormatter) const override {
-      std::cout << s << " Leaf [" << nrAssignments() << "]"
-                << valueFormatter(constant_) << std::endl;
+      std::cout << s << " Leaf " << valueFormatter(constant_) << std::endl;
     }
 
     /** Write graphviz format to stream `os`. */
@@ -114,14 +104,14 @@ namespace gtsam {
 
     /** apply unary operator */
     NodePtr apply(const Unary& op) const override {
-      NodePtr f(new Leaf(op(constant_), nrAssignments_));
+      NodePtr f(new Leaf(op(constant_)));
       return f;
     }
 
     /// Apply unary operator with assignment
     NodePtr apply(const UnaryAssignment& op,
                   const Assignment<L>& assignment) const override {
-      NodePtr f(new Leaf(op(assignment, constant_), nrAssignments_));
+      NodePtr f(new Leaf(op(assignment, constant_)));
       return f;
     }
 
@@ -137,9 +127,7 @@ namespace gtsam {
     // Applying binary operator to two leaves results in a leaf
     NodePtr apply_g_op_fL(const Leaf& fL, const Binary& op) const override {
       // fL op gL
-      // TODO(Varun) nrAssignments setting is not correct.
-      // Depending on f and g, the nrAssignments can be different. This is a bug!
-      NodePtr h(new Leaf(op(fL.constant_, constant_), fL.nrAssignments()));
+      NodePtr h(new Leaf(op(fL.constant_, constant_)));
       return h;
     }
 
@@ -150,7 +138,7 @@ namespace gtsam {
 
     /** choose a branch, create new memory ! */
     NodePtr choose(const L& label, size_t index) const override {
-      return NodePtr(new Leaf(constant(), nrAssignments()));
+      return NodePtr(new Leaf(constant()));
     }
 
     bool isLeaf() const override { return true; }
@@ -165,7 +153,6 @@ namespace gtsam {
     void serialize(ARCHIVE& ar, const unsigned int /*version*/) {
       ar & BOOST_SERIALIZATION_BASE_OBJECT_NVP(Base);
       ar& BOOST_SERIALIZATION_NVP(constant_);
-      ar& BOOST_SERIALIZATION_NVP(nrAssignments_);
     }
 #endif
   };  // Leaf
@@ -235,16 +222,8 @@ namespace gtsam {
           assert(f->branches().size() > 0);
           NodePtr f0 = f->branches_[0];
 
-          // Compute total number of assignments
-          size_t nrAssignments = 0;
-          for (auto branch : f->branches()) {
-            if (auto leaf = std::dynamic_pointer_cast<const Leaf>(branch)) {
-              nrAssignments += leaf->nrAssignments();
-            }
-          }
           NodePtr newLeaf(
-              new Leaf(std::dynamic_pointer_cast<const Leaf>(f0)->constant(),
-                       nrAssignments));
+              new Leaf(std::dynamic_pointer_cast<const Leaf>(f0)->constant()));
           return newLeaf;
         }
 #endif
@@ -730,7 +709,7 @@ namespace gtsam {
     // If leaf, apply unary conversion "op" and create a unique leaf.
     using MXLeaf = typename DecisionTree<M, X>::Leaf;
     if (auto leaf = std::dynamic_pointer_cast<const MXLeaf>(f)) {
-      return NodePtr(new Leaf(Y_of_X(leaf->constant()), leaf->nrAssignments()));
+      return NodePtr(new Leaf(Y_of_X(leaf->constant())));
     }
 
     // Check if Choice
@@ -877,16 +856,6 @@ namespace gtsam {
     return total;
   }
 
-  /****************************************************************************/
-  template <typename L, typename Y>
-  size_t DecisionTree<L, Y>::nrAssignments() const {
-    size_t n = 0;
-    this->visitLeaf([&n](const DecisionTree<L, Y>::Leaf& leaf) {
-      n += leaf.nrAssignments();
-    });
-    return n;
-  }
-
   /****************************************************************************/
   // fold is just done with a visit
   template <typename L, typename Y>
diff --git a/gtsam/discrete/DecisionTree.h b/gtsam/discrete/DecisionTree.h
index 9a8eac65e7..733ab1ad1f 100644
--- a/gtsam/discrete/DecisionTree.h
+++ b/gtsam/discrete/DecisionTree.h
@@ -307,42 +307,6 @@ namespace gtsam {
     /// Return the number of leaves in the tree.
     size_t nrLeaves() const;
 
-    /**
-     * @brief This is a convenience function which returns the total number of
-     * leaf assignments in the decision tree.
-     * This function is not used for anymajor operations within the discrete
-     * factor graph framework.
-     *
-     * Leaf assignments represent the cardinality of each leaf node, e.g. in a
-     * binary tree each leaf has 2 assignments. This includes counts removed
-     * from implicit pruning hence, it will always be >= nrLeaves().
-     *
-     * E.g. we have a decision tree as below, where each node has 2 branches:
-     *
-     * Choice(m1)
-     * 0 Choice(m0)
-     * 0 0 Leaf 0.0
-     * 0 1 Leaf 0.0
-     * 1 Choice(m0)
-     * 1 0 Leaf 1.0
-     * 1 1 Leaf 2.0
-     *
-     * In the unpruned form, the tree will have 4 assignments, 2 for each key,
-     * and 4 leaves.
-     *
-     * In the pruned form, the number of assignments is still 4 but the number
-     * of leaves is now 3, as below:
-     *
-     * Choice(m1)
-     * 0 Leaf 0.0
-     * 1 Choice(m0)
-     * 1 0 Leaf 1.0
-     * 1 1 Leaf 2.0
-     *
-     * @return size_t
-     */
-    size_t nrAssignments() const;
-
     /**
      * @brief Fold a binary function over the tree, returning accumulator.
      *
diff --git a/gtsam/discrete/tests/testDecisionTree.cpp b/gtsam/discrete/tests/testDecisionTree.cpp
index 653360fb7c..cdbad8c1c0 100644
--- a/gtsam/discrete/tests/testDecisionTree.cpp
+++ b/gtsam/discrete/tests/testDecisionTree.cpp
@@ -328,59 +328,6 @@ TEST(DecisionTree, Containers) {
   StringContainerTree converted(stringIntTree, container_of_int);
 }
 
-/* ************************************************************************** */
-// Test nrAssignments.
-TEST(DecisionTree, NrAssignments) {
-  const std::pair<string, size_t> A("A", 2), B("B", 2), C("C", 2);
-  DT tree({A, B, C}, "1 1 1 1 1 1 1 1");
-
-  EXPECT_LONGS_EQUAL(8, tree.nrAssignments());
-
-#ifdef GTSAM_DT_MERGING
-  EXPECT(tree.root_->isLeaf());
-  auto leaf = std::dynamic_pointer_cast<const DT::Leaf>(tree.root_);
-  EXPECT_LONGS_EQUAL(8, leaf->nrAssignments());
-#endif
-
-  DT tree2({C, B, A}, "1 1 1 2 3 4 5 5");
-  /* The tree is
-    Choice(C) 
-    0 Choice(B) 
-    0 0 Leaf 1
-    0 1 Choice(A) 
-    0 1 0 Leaf 1
-    0 1 1 Leaf 2
-    1 Choice(B) 
-    1 0 Choice(A) 
-    1 0 0 Leaf 3
-    1 0 1 Leaf 4
-    1 1 Leaf 5
-  */
-
-  EXPECT_LONGS_EQUAL(8, tree2.nrAssignments());
-
-  auto root = std::dynamic_pointer_cast<const DT::Choice>(tree2.root_);
-  CHECK(root);
-  auto choice0 = std::dynamic_pointer_cast<const DT::Choice>(root->branches()[0]);
-  CHECK(choice0);
-
-#ifdef GTSAM_DT_MERGING
-  EXPECT(choice0->branches()[0]->isLeaf());
-  auto choice00 = std::dynamic_pointer_cast<const DT::Leaf>(choice0->branches()[0]);
-  CHECK(choice00);
-  EXPECT_LONGS_EQUAL(2, choice00->nrAssignments());
-
-  auto choice1 = std::dynamic_pointer_cast<const DT::Choice>(root->branches()[1]);
-  CHECK(choice1);
-  auto choice10 = std::dynamic_pointer_cast<const DT::Choice>(choice1->branches()[0]);
-  CHECK(choice10);
-  auto choice11 = std::dynamic_pointer_cast<const DT::Leaf>(choice1->branches()[1]);
-  CHECK(choice11);
-  EXPECT(choice11->isLeaf());
-  EXPECT_LONGS_EQUAL(2, choice11->nrAssignments());
-#endif
-}
-
 /* ************************************************************************** */
 // Test visit.
 TEST(DecisionTree, visit) {
@@ -593,38 +540,6 @@ TEST(DecisionTree, ApplyWithAssignment) {
 #endif
 }
 
-/* ************************************************************************** */
-// Test number of assignments.
-TEST(DecisionTree, NrAssignments2) {
-  using gtsam::symbol_shorthand::M;
-
-  std::vector<double> probs = {0, 0, 1, 2};
-
-  /* Create the decision tree
-    Choice(m1)
-    0 Leaf 0.000000
-    1 Choice(m0)
-    1 0 Leaf 1.000000
-    1 1 Leaf 2.000000
-  */
-  DiscreteKeys keys{{M(1), 2}, {M(0), 2}};
-  DecisionTree<Key, double> dt1(keys, probs);
-  EXPECT_LONGS_EQUAL(4, dt1.nrAssignments());
-
-  /* Create the DecisionTree
-    Choice(m1)
-    0 Choice(m0)
-    0 0 Leaf 0.000000
-    0 1 Leaf 1.000000
-    1 Choice(m0)
-    1 0 Leaf 0.000000
-    1 1 Leaf 2.000000
-  */
-  DiscreteKeys keys2{{M(0), 2}, {M(1), 2}};
-  DecisionTree<Key, double> dt2(keys2, probs);
-  EXPECT_LONGS_EQUAL(4, dt2.nrAssignments());
-}
-
 /* ************************************************************************* */
 int main() {
   TestResult tr;
diff --git a/gtsam/discrete/tests/testDiscreteFactorGraph.cpp b/gtsam/discrete/tests/testDiscreteFactorGraph.cpp
index 6e86215957..341eb63e38 100644
--- a/gtsam/discrete/tests/testDiscreteFactorGraph.cpp
+++ b/gtsam/discrete/tests/testDiscreteFactorGraph.cpp
@@ -349,119 +349,6 @@ TEST(DiscreteFactorGraph, markdown) {
   EXPECT_DOUBLES_EQUAL(0.3, graph[0]->operator()(values), 1e-9);
 }
 
-TEST(DiscreteFactorGraph, NrAssignments) {
-#ifdef GTSAM_DT_MERGING
-  string expected_dfg = R"(
-size: 2
-factor 0:  f[ (m0,2), (m1,2), (m2,2), ]
- Choice(m2) 
- 0 Choice(m1) 
- 0 0 Leaf [2]   0
- 0 1 Choice(m0) 
- 0 1 0 Leaf [1]0.27527634
- 0 1 1 Leaf [1]   0
- 1 Choice(m1) 
- 1 0 Leaf [2]   0
- 1 1 Choice(m0) 
- 1 1 0 Leaf [1]0.44944733
- 1 1 1 Leaf [1]0.27527634
-factor 1:  f[ (m0,2), (m1,2), (m2,2), (m3,2), ]
- Choice(m3) 
- 0 Choice(m2) 
- 0 0 Choice(m1) 
- 0 0 0 Leaf [2]   1
- 0 0 1 Leaf [2]0.015366387
- 0 1 Choice(m1) 
- 0 1 0 Leaf [2]   1
- 0 1 1 Choice(m0) 
- 0 1 1 0 Leaf [1]   1
- 0 1 1 1 Leaf [1]0.015365663
- 1 Choice(m2) 
- 1 0 Choice(m1) 
- 1 0 0 Leaf [2]   1
- 1 0 1 Choice(m0) 
- 1 0 1 0 Leaf [1]0.0094115739
- 1 0 1 1 Leaf [1]0.0094115652
- 1 1 Choice(m1) 
- 1 1 0 Leaf [2]   1
- 1 1 1 Choice(m0) 
- 1 1 1 0 Leaf [1]   1
- 1 1 1 1 Leaf [1]0.009321081
-)";
-#else
-  string expected_dfg = R"(
-size: 2
-factor 0:  f[ (m0,2), (m1,2), (m2,2), ]
- Choice(m2) 
- 0 Choice(m1) 
- 0 0 Choice(m0) 
- 0 0 0 Leaf [1]   0
- 0 0 1 Leaf [1]   0
- 0 1 Choice(m0) 
- 0 1 0 Leaf [1]0.27527634
- 0 1 1 Leaf [1]0.44944733
- 1 Choice(m1) 
- 1 0 Choice(m0) 
- 1 0 0 Leaf [1]   0
- 1 0 1 Leaf [1]   0
- 1 1 Choice(m0) 
- 1 1 0 Leaf [1]   0
- 1 1 1 Leaf [1]0.27527634
-factor 1:  f[ (m0,2), (m1,2), (m2,2), (m3,2), ]
- Choice(m3) 
- 0 Choice(m2) 
- 0 0 Choice(m1) 
- 0 0 0 Choice(m0) 
- 0 0 0 0 Leaf [1]   1
- 0 0 0 1 Leaf [1]   1
- 0 0 1 Choice(m0) 
- 0 0 1 0 Leaf [1]0.015366387
- 0 0 1 1 Leaf [1]0.015366387
- 0 1 Choice(m1) 
- 0 1 0 Choice(m0) 
- 0 1 0 0 Leaf [1]   1
- 0 1 0 1 Leaf [1]   1
- 0 1 1 Choice(m0) 
- 0 1 1 0 Leaf [1]   1
- 0 1 1 1 Leaf [1]0.015365663
- 1 Choice(m2) 
- 1 0 Choice(m1) 
- 1 0 0 Choice(m0) 
- 1 0 0 0 Leaf [1]   1
- 1 0 0 1 Leaf [1]   1
- 1 0 1 Choice(m0) 
- 1 0 1 0 Leaf [1]0.0094115739
- 1 0 1 1 Leaf [1]0.0094115652
- 1 1 Choice(m1) 
- 1 1 0 Choice(m0) 
- 1 1 0 0 Leaf [1]   1
- 1 1 0 1 Leaf [1]   1
- 1 1 1 Choice(m0) 
- 1 1 1 0 Leaf [1]   1
- 1 1 1 1 Leaf [1]0.009321081
-)";
-#endif
-
-  DiscreteKeys d0{{M(0), 2}, {M(1), 2}, {M(2), 2}};
-  std::vector<double> p0 = {0, 0, 0.17054468, 0.27845056, 0, 0, 0, 0.17054468};
-  AlgebraicDecisionTree<Key> dt(d0, p0);
-  //TODO(Varun) Passing ADT to DiscreteConditional causes nrAssignments to get messed up
-  // Issue seems to be in DecisionTreeFactor.cpp L104
-  DiscreteConditional f0(3, DecisionTreeFactor(d0, dt));
-
-  DiscreteKeys d1{{M(0), 2}, {M(1), 2}, {M(2), 2}, {M(3), 2}};
-  std::vector<double> p1 = {
-      1, 1, 1, 1, 0.015366387, 0.0094115739, 1,           1,
-      1, 1, 1, 1, 0.015366387, 0.0094115652, 0.015365663, 0.009321081};
-  DecisionTreeFactor f1(d1, p1);
-
-  DiscreteFactorGraph dfg;
-  dfg.add(f0);
-  dfg.add(f1);
-
-  EXPECT(assert_print_equal(expected_dfg, dfg));
-}
-
 /* ************************************************************************* */
 int main() {
 TestResult tr;

From 2db08281c65d3f47feeb76bc467f9649f3c29a6e Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Mon, 10 Jul 2023 19:39:28 -0400
Subject: [PATCH 3/4] Revert "remove nrAssignments from the DecisionTree"

This reverts commit 647d3c0744171ffa78f585a39a88bfa4be4d2002.
---
 gtsam/discrete/DecisionTree-inl.h             |  47 ++++++--
 gtsam/discrete/DecisionTree.h                 |  36 ++++++
 gtsam/discrete/tests/testDecisionTree.cpp     |  85 +++++++++++++
 .../tests/testDiscreteFactorGraph.cpp         | 113 ++++++++++++++++++
 4 files changed, 273 insertions(+), 8 deletions(-)

diff --git a/gtsam/discrete/DecisionTree-inl.h b/gtsam/discrete/DecisionTree-inl.h
index 10532284cf..156177d039 100644
--- a/gtsam/discrete/DecisionTree-inl.h
+++ b/gtsam/discrete/DecisionTree-inl.h
@@ -53,17 +53,26 @@ namespace gtsam {
     /** constant stored in this leaf */
     Y constant_;
 
+    /** The number of assignments contained within this leaf.
+     * Particularly useful when leaves have been pruned.
+     */
+    size_t nrAssignments_;
+
     /// Default constructor for serialization.
     Leaf() {}
 
     /// Constructor from constant
-    Leaf(const Y& constant) : constant_(constant) {}
+    Leaf(const Y& constant, size_t nrAssignments = 1)
+        : constant_(constant), nrAssignments_(nrAssignments) {}
 
     /// Return the constant
     const Y& constant() const {
       return constant_;
     }
 
+    /// Return the number of assignments contained within this leaf.
+    size_t nrAssignments() const { return nrAssignments_; }
+
     /// Leaf-Leaf equality
     bool sameLeaf(const Leaf& q) const override {
       return constant_ == q.constant_;
@@ -84,7 +93,8 @@ namespace gtsam {
     /// print
     void print(const std::string& s, const LabelFormatter& labelFormatter,
                const ValueFormatter& valueFormatter) const override {
-      std::cout << s << " Leaf " << valueFormatter(constant_) << std::endl;
+      std::cout << s << " Leaf [" << nrAssignments() << "]"
+                << valueFormatter(constant_) << std::endl;
     }
 
     /** Write graphviz format to stream `os`. */
@@ -104,14 +114,14 @@ namespace gtsam {
 
     /** apply unary operator */
     NodePtr apply(const Unary& op) const override {
-      NodePtr f(new Leaf(op(constant_)));
+      NodePtr f(new Leaf(op(constant_), nrAssignments_));
       return f;
     }
 
     /// Apply unary operator with assignment
     NodePtr apply(const UnaryAssignment& op,
                   const Assignment<L>& assignment) const override {
-      NodePtr f(new Leaf(op(assignment, constant_)));
+      NodePtr f(new Leaf(op(assignment, constant_), nrAssignments_));
       return f;
     }
 
@@ -127,7 +137,9 @@ namespace gtsam {
     // Applying binary operator to two leaves results in a leaf
     NodePtr apply_g_op_fL(const Leaf& fL, const Binary& op) const override {
       // fL op gL
-      NodePtr h(new Leaf(op(fL.constant_, constant_)));
+      // TODO(Varun) nrAssignments setting is not correct.
+      // Depending on f and g, the nrAssignments can be different. This is a bug!
+      NodePtr h(new Leaf(op(fL.constant_, constant_), fL.nrAssignments()));
       return h;
     }
 
@@ -138,7 +150,7 @@ namespace gtsam {
 
     /** choose a branch, create new memory ! */
     NodePtr choose(const L& label, size_t index) const override {
-      return NodePtr(new Leaf(constant()));
+      return NodePtr(new Leaf(constant(), nrAssignments()));
     }
 
     bool isLeaf() const override { return true; }
@@ -153,6 +165,7 @@ namespace gtsam {
     void serialize(ARCHIVE& ar, const unsigned int /*version*/) {
       ar & BOOST_SERIALIZATION_BASE_OBJECT_NVP(Base);
       ar& BOOST_SERIALIZATION_NVP(constant_);
+      ar& BOOST_SERIALIZATION_NVP(nrAssignments_);
     }
 #endif
   };  // Leaf
@@ -222,8 +235,16 @@ namespace gtsam {
           assert(f->branches().size() > 0);
           NodePtr f0 = f->branches_[0];
 
+          // Compute total number of assignments
+          size_t nrAssignments = 0;
+          for (auto branch : f->branches()) {
+            if (auto leaf = std::dynamic_pointer_cast<const Leaf>(branch)) {
+              nrAssignments += leaf->nrAssignments();
+            }
+          }
           NodePtr newLeaf(
-              new Leaf(std::dynamic_pointer_cast<const Leaf>(f0)->constant()));
+              new Leaf(std::dynamic_pointer_cast<const Leaf>(f0)->constant(),
+                       nrAssignments));
           return newLeaf;
         }
 #endif
@@ -709,7 +730,7 @@ namespace gtsam {
     // If leaf, apply unary conversion "op" and create a unique leaf.
     using MXLeaf = typename DecisionTree<M, X>::Leaf;
     if (auto leaf = std::dynamic_pointer_cast<const MXLeaf>(f)) {
-      return NodePtr(new Leaf(Y_of_X(leaf->constant())));
+      return NodePtr(new Leaf(Y_of_X(leaf->constant()), leaf->nrAssignments()));
     }
 
     // Check if Choice
@@ -856,6 +877,16 @@ namespace gtsam {
     return total;
   }
 
+  /****************************************************************************/
+  template <typename L, typename Y>
+  size_t DecisionTree<L, Y>::nrAssignments() const {
+    size_t n = 0;
+    this->visitLeaf([&n](const DecisionTree<L, Y>::Leaf& leaf) {
+      n += leaf.nrAssignments();
+    });
+    return n;
+  }
+
   /****************************************************************************/
   // fold is just done with a visit
   template <typename L, typename Y>
diff --git a/gtsam/discrete/DecisionTree.h b/gtsam/discrete/DecisionTree.h
index 733ab1ad1f..9a8eac65e7 100644
--- a/gtsam/discrete/DecisionTree.h
+++ b/gtsam/discrete/DecisionTree.h
@@ -307,6 +307,42 @@ namespace gtsam {
     /// Return the number of leaves in the tree.
     size_t nrLeaves() const;
 
+    /**
+     * @brief This is a convenience function which returns the total number of
+     * leaf assignments in the decision tree.
+     * This function is not used for anymajor operations within the discrete
+     * factor graph framework.
+     *
+     * Leaf assignments represent the cardinality of each leaf node, e.g. in a
+     * binary tree each leaf has 2 assignments. This includes counts removed
+     * from implicit pruning hence, it will always be >= nrLeaves().
+     *
+     * E.g. we have a decision tree as below, where each node has 2 branches:
+     *
+     * Choice(m1)
+     * 0 Choice(m0)
+     * 0 0 Leaf 0.0
+     * 0 1 Leaf 0.0
+     * 1 Choice(m0)
+     * 1 0 Leaf 1.0
+     * 1 1 Leaf 2.0
+     *
+     * In the unpruned form, the tree will have 4 assignments, 2 for each key,
+     * and 4 leaves.
+     *
+     * In the pruned form, the number of assignments is still 4 but the number
+     * of leaves is now 3, as below:
+     *
+     * Choice(m1)
+     * 0 Leaf 0.0
+     * 1 Choice(m0)
+     * 1 0 Leaf 1.0
+     * 1 1 Leaf 2.0
+     *
+     * @return size_t
+     */
+    size_t nrAssignments() const;
+
     /**
      * @brief Fold a binary function over the tree, returning accumulator.
      *
diff --git a/gtsam/discrete/tests/testDecisionTree.cpp b/gtsam/discrete/tests/testDecisionTree.cpp
index cdbad8c1c0..653360fb7c 100644
--- a/gtsam/discrete/tests/testDecisionTree.cpp
+++ b/gtsam/discrete/tests/testDecisionTree.cpp
@@ -328,6 +328,59 @@ TEST(DecisionTree, Containers) {
   StringContainerTree converted(stringIntTree, container_of_int);
 }
 
+/* ************************************************************************** */
+// Test nrAssignments.
+TEST(DecisionTree, NrAssignments) {
+  const std::pair<string, size_t> A("A", 2), B("B", 2), C("C", 2);
+  DT tree({A, B, C}, "1 1 1 1 1 1 1 1");
+
+  EXPECT_LONGS_EQUAL(8, tree.nrAssignments());
+
+#ifdef GTSAM_DT_MERGING
+  EXPECT(tree.root_->isLeaf());
+  auto leaf = std::dynamic_pointer_cast<const DT::Leaf>(tree.root_);
+  EXPECT_LONGS_EQUAL(8, leaf->nrAssignments());
+#endif
+
+  DT tree2({C, B, A}, "1 1 1 2 3 4 5 5");
+  /* The tree is
+    Choice(C) 
+    0 Choice(B) 
+    0 0 Leaf 1
+    0 1 Choice(A) 
+    0 1 0 Leaf 1
+    0 1 1 Leaf 2
+    1 Choice(B) 
+    1 0 Choice(A) 
+    1 0 0 Leaf 3
+    1 0 1 Leaf 4
+    1 1 Leaf 5
+  */
+
+  EXPECT_LONGS_EQUAL(8, tree2.nrAssignments());
+
+  auto root = std::dynamic_pointer_cast<const DT::Choice>(tree2.root_);
+  CHECK(root);
+  auto choice0 = std::dynamic_pointer_cast<const DT::Choice>(root->branches()[0]);
+  CHECK(choice0);
+
+#ifdef GTSAM_DT_MERGING
+  EXPECT(choice0->branches()[0]->isLeaf());
+  auto choice00 = std::dynamic_pointer_cast<const DT::Leaf>(choice0->branches()[0]);
+  CHECK(choice00);
+  EXPECT_LONGS_EQUAL(2, choice00->nrAssignments());
+
+  auto choice1 = std::dynamic_pointer_cast<const DT::Choice>(root->branches()[1]);
+  CHECK(choice1);
+  auto choice10 = std::dynamic_pointer_cast<const DT::Choice>(choice1->branches()[0]);
+  CHECK(choice10);
+  auto choice11 = std::dynamic_pointer_cast<const DT::Leaf>(choice1->branches()[1]);
+  CHECK(choice11);
+  EXPECT(choice11->isLeaf());
+  EXPECT_LONGS_EQUAL(2, choice11->nrAssignments());
+#endif
+}
+
 /* ************************************************************************** */
 // Test visit.
 TEST(DecisionTree, visit) {
@@ -540,6 +593,38 @@ TEST(DecisionTree, ApplyWithAssignment) {
 #endif
 }
 
+/* ************************************************************************** */
+// Test number of assignments.
+TEST(DecisionTree, NrAssignments2) {
+  using gtsam::symbol_shorthand::M;
+
+  std::vector<double> probs = {0, 0, 1, 2};
+
+  /* Create the decision tree
+    Choice(m1)
+    0 Leaf 0.000000
+    1 Choice(m0)
+    1 0 Leaf 1.000000
+    1 1 Leaf 2.000000
+  */
+  DiscreteKeys keys{{M(1), 2}, {M(0), 2}};
+  DecisionTree<Key, double> dt1(keys, probs);
+  EXPECT_LONGS_EQUAL(4, dt1.nrAssignments());
+
+  /* Create the DecisionTree
+    Choice(m1)
+    0 Choice(m0)
+    0 0 Leaf 0.000000
+    0 1 Leaf 1.000000
+    1 Choice(m0)
+    1 0 Leaf 0.000000
+    1 1 Leaf 2.000000
+  */
+  DiscreteKeys keys2{{M(0), 2}, {M(1), 2}};
+  DecisionTree<Key, double> dt2(keys2, probs);
+  EXPECT_LONGS_EQUAL(4, dt2.nrAssignments());
+}
+
 /* ************************************************************************* */
 int main() {
   TestResult tr;
diff --git a/gtsam/discrete/tests/testDiscreteFactorGraph.cpp b/gtsam/discrete/tests/testDiscreteFactorGraph.cpp
index 341eb63e38..6e86215957 100644
--- a/gtsam/discrete/tests/testDiscreteFactorGraph.cpp
+++ b/gtsam/discrete/tests/testDiscreteFactorGraph.cpp
@@ -349,6 +349,119 @@ TEST(DiscreteFactorGraph, markdown) {
   EXPECT_DOUBLES_EQUAL(0.3, graph[0]->operator()(values), 1e-9);
 }
 
+TEST(DiscreteFactorGraph, NrAssignments) {
+#ifdef GTSAM_DT_MERGING
+  string expected_dfg = R"(
+size: 2
+factor 0:  f[ (m0,2), (m1,2), (m2,2), ]
+ Choice(m2) 
+ 0 Choice(m1) 
+ 0 0 Leaf [2]   0
+ 0 1 Choice(m0) 
+ 0 1 0 Leaf [1]0.27527634
+ 0 1 1 Leaf [1]   0
+ 1 Choice(m1) 
+ 1 0 Leaf [2]   0
+ 1 1 Choice(m0) 
+ 1 1 0 Leaf [1]0.44944733
+ 1 1 1 Leaf [1]0.27527634
+factor 1:  f[ (m0,2), (m1,2), (m2,2), (m3,2), ]
+ Choice(m3) 
+ 0 Choice(m2) 
+ 0 0 Choice(m1) 
+ 0 0 0 Leaf [2]   1
+ 0 0 1 Leaf [2]0.015366387
+ 0 1 Choice(m1) 
+ 0 1 0 Leaf [2]   1
+ 0 1 1 Choice(m0) 
+ 0 1 1 0 Leaf [1]   1
+ 0 1 1 1 Leaf [1]0.015365663
+ 1 Choice(m2) 
+ 1 0 Choice(m1) 
+ 1 0 0 Leaf [2]   1
+ 1 0 1 Choice(m0) 
+ 1 0 1 0 Leaf [1]0.0094115739
+ 1 0 1 1 Leaf [1]0.0094115652
+ 1 1 Choice(m1) 
+ 1 1 0 Leaf [2]   1
+ 1 1 1 Choice(m0) 
+ 1 1 1 0 Leaf [1]   1
+ 1 1 1 1 Leaf [1]0.009321081
+)";
+#else
+  string expected_dfg = R"(
+size: 2
+factor 0:  f[ (m0,2), (m1,2), (m2,2), ]
+ Choice(m2) 
+ 0 Choice(m1) 
+ 0 0 Choice(m0) 
+ 0 0 0 Leaf [1]   0
+ 0 0 1 Leaf [1]   0
+ 0 1 Choice(m0) 
+ 0 1 0 Leaf [1]0.27527634
+ 0 1 1 Leaf [1]0.44944733
+ 1 Choice(m1) 
+ 1 0 Choice(m0) 
+ 1 0 0 Leaf [1]   0
+ 1 0 1 Leaf [1]   0
+ 1 1 Choice(m0) 
+ 1 1 0 Leaf [1]   0
+ 1 1 1 Leaf [1]0.27527634
+factor 1:  f[ (m0,2), (m1,2), (m2,2), (m3,2), ]
+ Choice(m3) 
+ 0 Choice(m2) 
+ 0 0 Choice(m1) 
+ 0 0 0 Choice(m0) 
+ 0 0 0 0 Leaf [1]   1
+ 0 0 0 1 Leaf [1]   1
+ 0 0 1 Choice(m0) 
+ 0 0 1 0 Leaf [1]0.015366387
+ 0 0 1 1 Leaf [1]0.015366387
+ 0 1 Choice(m1) 
+ 0 1 0 Choice(m0) 
+ 0 1 0 0 Leaf [1]   1
+ 0 1 0 1 Leaf [1]   1
+ 0 1 1 Choice(m0) 
+ 0 1 1 0 Leaf [1]   1
+ 0 1 1 1 Leaf [1]0.015365663
+ 1 Choice(m2) 
+ 1 0 Choice(m1) 
+ 1 0 0 Choice(m0) 
+ 1 0 0 0 Leaf [1]   1
+ 1 0 0 1 Leaf [1]   1
+ 1 0 1 Choice(m0) 
+ 1 0 1 0 Leaf [1]0.0094115739
+ 1 0 1 1 Leaf [1]0.0094115652
+ 1 1 Choice(m1) 
+ 1 1 0 Choice(m0) 
+ 1 1 0 0 Leaf [1]   1
+ 1 1 0 1 Leaf [1]   1
+ 1 1 1 Choice(m0) 
+ 1 1 1 0 Leaf [1]   1
+ 1 1 1 1 Leaf [1]0.009321081
+)";
+#endif
+
+  DiscreteKeys d0{{M(0), 2}, {M(1), 2}, {M(2), 2}};
+  std::vector<double> p0 = {0, 0, 0.17054468, 0.27845056, 0, 0, 0, 0.17054468};
+  AlgebraicDecisionTree<Key> dt(d0, p0);
+  //TODO(Varun) Passing ADT to DiscreteConditional causes nrAssignments to get messed up
+  // Issue seems to be in DecisionTreeFactor.cpp L104
+  DiscreteConditional f0(3, DecisionTreeFactor(d0, dt));
+
+  DiscreteKeys d1{{M(0), 2}, {M(1), 2}, {M(2), 2}, {M(3), 2}};
+  std::vector<double> p1 = {
+      1, 1, 1, 1, 0.015366387, 0.0094115739, 1,           1,
+      1, 1, 1, 1, 0.015366387, 0.0094115652, 0.015365663, 0.009321081};
+  DecisionTreeFactor f1(d1, p1);
+
+  DiscreteFactorGraph dfg;
+  dfg.add(f0);
+  dfg.add(f1);
+
+  EXPECT(assert_print_equal(expected_dfg, dfg));
+}
+
 /* ************************************************************************* */
 int main() {
 TestResult tr;

From b7deefd744b8a5cb544d863eb1e6f35da56f7bff Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Mon, 10 Jul 2023 19:39:36 -0400
Subject: [PATCH 4/4] Revert "enumerate all assignments for computing
 probabilities to prune"

This reverts commit 8c38e45c83c6bbeb2ef5e7c7d73807b2eed6831f.
---
 gtsam/discrete/DecisionTreeFactor.cpp         | 11 ++--
 .../tests/testGaussianMixtureFactor.cpp       |  4 +-
 .../tests/testHybridNonlinearFactorGraph.cpp  | 58 +++++++++----------
 gtsam/hybrid/tests/testMixtureFactor.cpp      |  4 +-
 4 files changed, 38 insertions(+), 39 deletions(-)

diff --git a/gtsam/discrete/DecisionTreeFactor.cpp b/gtsam/discrete/DecisionTreeFactor.cpp
index 6f20907e14..4aca10a284 100644
--- a/gtsam/discrete/DecisionTreeFactor.cpp
+++ b/gtsam/discrete/DecisionTreeFactor.cpp
@@ -306,12 +306,11 @@ namespace gtsam {
 
     // Get the probabilities in the decision tree so we can threshold.
     std::vector<double> probabilities;
-    // NOTE(Varun) this is potentially slow due to the cartesian product
-    auto allValues = DiscreteValues::CartesianProduct(this->discreteKeys());
-    for (auto&& val : allValues) {
-      double prob = (*this)(val);
-      probabilities.push_back(prob);
-    }
+    this->visitLeaf([&](const Leaf& leaf) {
+      const size_t nrAssignments = leaf.nrAssignments();
+      double prob = leaf.constant();
+      probabilities.insert(probabilities.end(), nrAssignments, prob);
+    });
 
     // The number of probabilities can be lower than max_leaves
     if (probabilities.size() <= N) {
diff --git a/gtsam/hybrid/tests/testGaussianMixtureFactor.cpp b/gtsam/hybrid/tests/testGaussianMixtureFactor.cpp
index 75ba5a0594..5207e9372a 100644
--- a/gtsam/hybrid/tests/testGaussianMixtureFactor.cpp
+++ b/gtsam/hybrid/tests/testGaussianMixtureFactor.cpp
@@ -108,7 +108,7 @@ TEST(GaussianMixtureFactor, Printing) {
   std::string expected =
       R"(Hybrid [x1 x2; 1]{
  Choice(1) 
- 0 Leaf :
+ 0 Leaf [1]:
   A[x1] = [
 	0;
 	0
@@ -120,7 +120,7 @@ TEST(GaussianMixtureFactor, Printing) {
   b = [ 0 0 ]
   No noise model
 
- 1 Leaf :
+ 1 Leaf [1]:
   A[x1] = [
 	0;
 	0
diff --git a/gtsam/hybrid/tests/testHybridNonlinearFactorGraph.cpp b/gtsam/hybrid/tests/testHybridNonlinearFactorGraph.cpp
index 29136c62ad..0a621c42d2 100644
--- a/gtsam/hybrid/tests/testHybridNonlinearFactorGraph.cpp
+++ b/gtsam/hybrid/tests/testHybridNonlinearFactorGraph.cpp
@@ -493,7 +493,7 @@ factor 0:
 factor 1: 
 Hybrid [x0 x1; m0]{
  Choice(m0) 
- 0 Leaf :
+ 0 Leaf [1]:
   A[x0] = [
 	-1
 ]
@@ -503,7 +503,7 @@ Hybrid [x0 x1; m0]{
   b = [ -1 ]
   No noise model
 
- 1 Leaf :
+ 1 Leaf [1]:
   A[x0] = [
 	-1
 ]
@@ -517,7 +517,7 @@ Hybrid [x0 x1; m0]{
 factor 2: 
 Hybrid [x1 x2; m1]{
  Choice(m1) 
- 0 Leaf :
+ 0 Leaf [1]:
   A[x1] = [
 	-1
 ]
@@ -527,7 +527,7 @@ Hybrid [x1 x2; m1]{
   b = [ -1 ]
   No noise model
 
- 1 Leaf :
+ 1 Leaf [1]:
   A[x1] = [
 	-1
 ]
@@ -551,16 +551,16 @@ factor 4:
   b = [ -10 ]
   No noise model
 factor 5:  P( m0 ):
- Leaf  0.5
+ Leaf [2] 0.5
 
 factor 6:  P( m1 | m0 ):
  Choice(m1) 
  0 Choice(m0) 
- 0 0 Leaf 0.33333333
- 0 1 Leaf  0.6
+ 0 0 Leaf [1]0.33333333
+ 0 1 Leaf [1] 0.6
  1 Choice(m0) 
- 1 0 Leaf 0.66666667
- 1 1 Leaf  0.4
+ 1 0 Leaf [1]0.66666667
+ 1 1 Leaf [1] 0.4
 
 )";
 #else
@@ -575,7 +575,7 @@ factor 0:
 factor 1: 
 Hybrid [x0 x1; m0]{
  Choice(m0) 
- 0 Leaf :
+ 0 Leaf [1]:
   A[x0] = [
 	-1
 ]
@@ -585,7 +585,7 @@ Hybrid [x0 x1; m0]{
   b = [ -1 ]
   No noise model
 
- 1 Leaf :
+ 1 Leaf [1]:
   A[x0] = [
 	-1
 ]
@@ -599,7 +599,7 @@ Hybrid [x0 x1; m0]{
 factor 2: 
 Hybrid [x1 x2; m1]{
  Choice(m1) 
- 0 Leaf :
+ 0 Leaf [1]:
   A[x1] = [
 	-1
 ]
@@ -609,7 +609,7 @@ Hybrid [x1 x2; m1]{
   b = [ -1 ]
   No noise model
 
- 1 Leaf :
+ 1 Leaf [1]:
   A[x1] = [
 	-1
 ]
@@ -634,17 +634,17 @@ factor 4:
   No noise model
 factor 5:  P( m0 ):
  Choice(m0) 
- 0 Leaf  0.5
- 1 Leaf  0.5
+ 0 Leaf [1] 0.5
+ 1 Leaf [1] 0.5
 
 factor 6:  P( m1 | m0 ):
  Choice(m1) 
  0 Choice(m0) 
- 0 0 Leaf 0.33333333
- 0 1 Leaf  0.6
+ 0 0 Leaf [1]0.33333333
+ 0 1 Leaf [1] 0.6
  1 Choice(m0) 
- 1 0 Leaf 0.66666667
- 1 1 Leaf  0.4
+ 1 0 Leaf [1]0.66666667
+ 1 1 Leaf [1] 0.4
 
 )";
 #endif
@@ -657,13 +657,13 @@ size: 3
 conditional 0: Hybrid  P( x0 | x1 m0)
  Discrete Keys = (m0, 2), 
  Choice(m0) 
- 0 Leaf  p(x0 | x1)
+ 0 Leaf [1] p(x0 | x1)
   R = [ 10.0499 ]
   S[x1] = [ -0.0995037 ]
   d = [ -9.85087 ]
   No noise model
 
- 1 Leaf  p(x0 | x1)
+ 1 Leaf [1] p(x0 | x1)
   R = [ 10.0499 ]
   S[x1] = [ -0.0995037 ]
   d = [ -9.95037 ]
@@ -673,26 +673,26 @@ conditional 1: Hybrid  P( x1 | x2 m0 m1)
  Discrete Keys = (m0, 2), (m1, 2), 
  Choice(m1) 
  0 Choice(m0) 
- 0 0 Leaf  p(x1 | x2)
+ 0 0 Leaf [1] p(x1 | x2)
   R = [ 10.099 ]
   S[x2] = [ -0.0990196 ]
   d = [ -9.99901 ]
   No noise model
 
- 0 1 Leaf  p(x1 | x2)
+ 0 1 Leaf [1] p(x1 | x2)
   R = [ 10.099 ]
   S[x2] = [ -0.0990196 ]
   d = [ -9.90098 ]
   No noise model
 
  1 Choice(m0) 
- 1 0 Leaf  p(x1 | x2)
+ 1 0 Leaf [1] p(x1 | x2)
   R = [ 10.099 ]
   S[x2] = [ -0.0990196 ]
   d = [ -10.098 ]
   No noise model
 
- 1 1 Leaf  p(x1 | x2)
+ 1 1 Leaf [1] p(x1 | x2)
   R = [ 10.099 ]
   S[x2] = [ -0.0990196 ]
   d = [ -10 ]
@@ -702,14 +702,14 @@ conditional 2: Hybrid  P( x2 | m0 m1)
  Discrete Keys = (m0, 2), (m1, 2), 
  Choice(m1) 
  0 Choice(m0) 
- 0 0 Leaf  p(x2)
+ 0 0 Leaf [1] p(x2)
   R = [ 10.0494 ]
   d = [ -10.1489 ]
   mean: 1 elements
   x2: -1.0099
   No noise model
 
- 0 1 Leaf  p(x2)
+ 0 1 Leaf [1] p(x2)
   R = [ 10.0494 ]
   d = [ -10.1479 ]
   mean: 1 elements
@@ -717,14 +717,14 @@ conditional 2: Hybrid  P( x2 | m0 m1)
   No noise model
 
  1 Choice(m0) 
- 1 0 Leaf  p(x2)
+ 1 0 Leaf [1] p(x2)
   R = [ 10.0494 ]
   d = [ -10.0504 ]
   mean: 1 elements
   x2: -1.0001
   No noise model
 
- 1 1 Leaf  p(x2)
+ 1 1 Leaf [1] p(x2)
   R = [ 10.0494 ]
   d = [ -10.0494 ]
   mean: 1 elements
diff --git a/gtsam/hybrid/tests/testMixtureFactor.cpp b/gtsam/hybrid/tests/testMixtureFactor.cpp
index 0827e4eedb..03fdccff26 100644
--- a/gtsam/hybrid/tests/testMixtureFactor.cpp
+++ b/gtsam/hybrid/tests/testMixtureFactor.cpp
@@ -63,8 +63,8 @@ TEST(MixtureFactor, Printing) {
       R"(Hybrid [x1 x2; 1]
 MixtureFactor
  Choice(1) 
- 0 Leaf  Nonlinear factor on 2 keys
- 1 Leaf  Nonlinear factor on 2 keys
+ 0 Leaf [1] Nonlinear factor on 2 keys
+ 1 Leaf [1] Nonlinear factor on 2 keys
 )";
   EXPECT(assert_print_equal(expected, mixtureFactor));
 }