Efficient Discrete Elimination

gtsam/discrete/DecisionTreeFactor.cpp

@@ -82,6 +82,22 @@ namespace gtsam {
     ADT::print("", formatter);
   }
 
+  /* ************************************************************************ */
+  DecisionTreeFactor DecisionTreeFactor::apply(ADT::Unary op) const {
+    // apply operand
+    ADT result = ADT::apply(op);
+    // Make a new factor
+    return DecisionTreeFactor(discreteKeys(), result);
+  }
+
+  /* ************************************************************************ */
+  DecisionTreeFactor DecisionTreeFactor::apply(ADT::UnaryAssignment op) const {
+    // apply operand
+    ADT result = ADT::apply(op);
+    // Make a new factor
+    return DecisionTreeFactor(discreteKeys(), result);
+  }
+
   /* ************************************************************************ */
   DecisionTreeFactor DecisionTreeFactor::apply(const DecisionTreeFactor& f,
                                               ADT::Binary op) const {
@@ -101,14 +117,6 @@ namespace gtsam {
     return DecisionTreeFactor(keys, result);
   }
 
-  /* ************************************************************************ */
-  DecisionTreeFactor DecisionTreeFactor::apply(ADT::UnaryAssignment op) const {
-    // apply operand
-    ADT result = ADT::apply(op);
-    // Make a new factor
-    return DecisionTreeFactor(discreteKeys(), result);
-  }
-
   /* ************************************************************************ */
   DecisionTreeFactor::shared_ptr DecisionTreeFactor::combine(
       size_t nrFrontals, ADT::Binary op) const {
@@ -188,10 +196,45 @@ namespace gtsam {
 
   /* ************************************************************************ */
   std::vector<double> DecisionTreeFactor::probabilities() const {
+    // Set of all keys
+    std::set<Key> allKeys(keys().begin(), keys().end());
+
     std::vector<double> probs;
-    for (auto&& [key, value] : enumerate()) {
-      probs.push_back(value);
-    }
+
+    /* An operation that takes each leaf probability, and computes the
+     * nrAssignments by checking the difference between the keys in the factor
+     * and the keys in the assignment.
+     * The nrAssignments is then used to append
+     * the correct number of leaf probability values to the `probs` vector
+     * defined above.
+     */
+    auto op = [&](const Assignment<Key>& a, double p) {
+      // Get all the keys in the current assignment
+      std::set<Key> assignment_keys;
+      for (auto&& [k, _] : a) {
+        assignment_keys.insert(k);
+      }
+
+      // Find the keys missing in the assignment
+      std::vector<Key> diff;
+      std::set_difference(allKeys.begin(), allKeys.end(),
+                          assignment_keys.begin(), assignment_keys.end(),
+                          std::back_inserter(diff));
+
+      // Compute the total number of assignments in the (pruned) subtree
+      size_t nrAssignments = 1;
+      for (auto&& k : diff) {
+        nrAssignments *= cardinalities_.at(k);
+      }
+      // Add p `nrAssignments` times to the probs vector.
+      probs.insert(probs.end(), nrAssignments, p);
+
+      return p;
+    };
+
+    // Go through the tree
+    this->apply(op);
+
     return probs;
   }
 
@@ -305,11 +348,7 @@ namespace gtsam {
     const size_t N = maxNrAssignments;
 
     // 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
-    for (auto&& [assignment, prob] : this->enumerate()) {
-      probabilities.push_back(prob);
-    }
+    std::vector<double> probabilities = this->probabilities();
 
     // The number of probabilities can be lower than max_leaves
     if (probabilities.size() <= N) {

gtsam/discrete/DecisionTreeFactor.h

@@ -186,6 +186,13 @@ namespace gtsam {
      * Apply unary operator (*this) "op" f
      * @param op a unary operator that operates on AlgebraicDecisionTree
      */
+    DecisionTreeFactor apply(ADT::Unary op) const;
+
+    /**
+     * Apply unary operator (*this) "op" f
+     * @param op a unary operator that operates on AlgebraicDecisionTree. Takes
+     * both the assignment and the value.
+     */
     DecisionTreeFactor apply(ADT::UnaryAssignment op) const;
 
     /**

gtsam/hybrid/GaussianMixture.cpp

@@ -286,8 +286,6 @@ GaussianMixture::prunerFunc(const DecisionTreeFactor &discreteProbs) {
 
 /* *******************************************************************************/
 void GaussianMixture::prune(const DecisionTreeFactor &discreteProbs) {
-  auto discreteProbsKeySet = DiscreteKeysAsSet(discreteProbs.discreteKeys());
-  auto gmKeySet = DiscreteKeysAsSet(this->discreteKeys());
   // Functional which loops over all assignments and create a set of
   // GaussianConditionals
   auto pruner = prunerFunc(discreteProbs);

gtsam/hybrid/HybridBayesNet.cpp

@@ -129,7 +129,6 @@ std::function<double(const Assignment<Key> &, double)> prunerFunc(
 DecisionTreeFactor HybridBayesNet::pruneDiscreteConditionals(
     size_t maxNrLeaves) {
   // Get the joint distribution of only the discrete keys
-  gttic_(HybridBayesNet_PruneDiscreteConditionals);
   // The joint discrete probability.
   DiscreteConditional discreteProbs;
 
@@ -147,12 +146,11 @@ DecisionTreeFactor HybridBayesNet::pruneDiscreteConditionals(
       discrete_factor_idxs.push_back(i);
     }
   }
+
   const DecisionTreeFactor prunedDiscreteProbs =
       discreteProbs.prune(maxNrLeaves);
-  gttoc_(HybridBayesNet_PruneDiscreteConditionals);
 
   // Eliminate joint probability back into conditionals
-  gttic_(HybridBayesNet_UpdateDiscreteConditionals);
   DiscreteFactorGraph dfg{prunedDiscreteProbs};
   DiscreteBayesNet::shared_ptr dbn = dfg.eliminateSequential(discrete_frontals);
 
@@ -161,7 +159,6 @@ DecisionTreeFactor HybridBayesNet::pruneDiscreteConditionals(
     size_t idx = discrete_factor_idxs.at(i);
     this->at(idx) = std::make_shared<HybridConditional>(dbn->at(i));
   }
-  gttoc_(HybridBayesNet_UpdateDiscreteConditionals);
 
   return prunedDiscreteProbs;
 }
@@ -180,7 +177,6 @@ HybridBayesNet HybridBayesNet::prune(size_t maxNrLeaves) {
 
   HybridBayesNet prunedBayesNetFragment;
 
-  gttic_(HybridBayesNet_PruneMixtures);
   // Go through all the conditionals in the
   // Bayes Net and prune them as per prunedDiscreteProbs.
   for (auto &&conditional : *this) {
@@ -197,7 +193,6 @@ HybridBayesNet HybridBayesNet::prune(size_t maxNrLeaves) {
       prunedBayesNetFragment.push_back(conditional);
     }
   }
-  gttoc_(HybridBayesNet_PruneMixtures);
 
   return prunedBayesNetFragment;
 }

gtsam/hybrid/HybridGaussianFactorGraph.cpp

@@ -96,7 +96,6 @@ static GaussianFactorGraphTree addGaussian(
 // TODO(dellaert): it's probably more efficient to first collect the discrete
 // keys, and then loop over all assignments to populate a vector.
 GaussianFactorGraphTree HybridGaussianFactorGraph::assembleGraphTree() const {
-  gttic_(assembleGraphTree);
 
   GaussianFactorGraphTree result;
 
@@ -129,8 +128,6 @@ GaussianFactorGraphTree HybridGaussianFactorGraph::assembleGraphTree() const {
     }
   }
 
-  gttoc_(assembleGraphTree);
-
   return result;
 }
 

gtsam/hybrid/tests/testHybridNonlinearFactorGraph.cpp

@@ -420,7 +420,7 @@ TEST(HybridFactorGraph, Full_Elimination) {
     DiscreteFactorGraph discrete_fg;
     // TODO(Varun) Make this a function of HybridGaussianFactorGraph?
     for (auto& factor : (*remainingFactorGraph_partial)) {
-      auto df = dynamic_pointer_cast<DecisionTreeFactor>(factor);
+      auto df = dynamic_pointer_cast<DiscreteFactor>(factor);
       assert(df);
       discrete_fg.push_back(df);
     }