enumerate all assignments for computing probabilities to prune

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) {

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

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

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));
 }