Improved BayesTree pruning

gtsam/discrete/Assignment.h

@@ -11,15 +11,17 @@
 
 /**
  * @file    Assignment.h
- * @brief    An assignment from labels to a discrete value index (size_t)
+ * @brief   An assignment from labels to a discrete value index (size_t)
  * @author  Frank Dellaert
  * @date    Feb 5, 2012
  */
 
 #pragma once
 
+#include <functional>
 #include <iostream>
 #include <map>
+#include <sstream>
 #include <utility>
 #include <vector>
 
@@ -32,13 +34,30 @@ namespace gtsam {
  */
 template <class L>
 class Assignment : public std::map<L, size_t> {
+  /**
+   * @brief Default method used by `labelFormatter` or `valueFormatter` when
+   * printing.
+   *
+   * @param x The value passed to format.
+   * @return std::string
+   */
+  static std::string DefaultFormatter(const L& x) {
+    std::stringstream ss;
+    ss << x;
+    return ss.str();
+  }
+
  public:
   using std::map<L, size_t>::operator=;
 
-  void print(const std::string& s = "Assignment: ") const {
+  void print(const std::string& s = "Assignment: ",
+             const std::function<std::string(L)>& labelFormatter =
+                 &DefaultFormatter) const {
     std::cout << s << ": ";
-    for (const typename Assignment::value_type& keyValue : *this)
-      std::cout << "(" << keyValue.first << ", " << keyValue.second << ")";
+    for (const typename Assignment::value_type& keyValue : *this) {
+      std::cout << "(" << labelFormatter(keyValue.first) << ", "
+                << keyValue.second << ")";
+    }
     std::cout << std::endl;
   }
 

gtsam/hybrid/GaussianMixture.cpp

@@ -119,11 +119,12 @@ void GaussianMixture::print(const std::string &s,
       "", [&](Key k) { return formatter(k); },
       [&](const GaussianConditional::shared_ptr &gf) -> std::string {
         RedirectCout rd;
-        if (gf && !gf->empty())
+        if (gf && !gf->empty()) {
           gf->print("", formatter);
-        else
-          return {"nullptr"};
-        return rd.str();
+          return rd.str();
+        } else {
+          return "nullptr";
+        }
       });
 }
 

gtsam/hybrid/HybridBayesNet.cpp

@@ -31,8 +31,32 @@ static std::set<DiscreteKey> DiscreteKeysAsSet(const DiscreteKeys &dkeys) {
 }
 
 /* ************************************************************************* */
-HybridBayesNet HybridBayesNet::prune(
-    const DecisionTreeFactor::shared_ptr &discreteFactor) const {
+DecisionTreeFactor::shared_ptr HybridBayesNet::discreteConditionals() const {
+  AlgebraicDecisionTree<Key> decisionTree;
+
+  // The canonical decision tree factor which will get the discrete conditionals
+  // added to it.
+  DecisionTreeFactor dtFactor;
+
+  for (size_t i = 0; i < this->size(); i++) {
+    HybridConditional::shared_ptr conditional = this->at(i);
+    if (conditional->isDiscrete()) {
+      // Convert to a DecisionTreeFactor and add it to the main factor.
+      DecisionTreeFactor f(*conditional->asDiscreteConditional());
+      dtFactor = dtFactor * f;
+    }
+  }
+  return boost::make_shared<DecisionTreeFactor>(dtFactor);
+}
+
+/* ************************************************************************* */
+HybridBayesNet HybridBayesNet::prune(size_t maxNrLeaves) const {
+  // Get the decision tree of only the discrete keys
+  auto discreteConditionals = this->discreteConditionals();
+  const DecisionTreeFactor::shared_ptr discreteFactor =
+      boost::make_shared<DecisionTreeFactor>(
+          discreteConditionals->prune(maxNrLeaves));
+
   /* To Prune, we visitWith every leaf in the GaussianMixture.
    * For each leaf, using the assignment we can check the discrete decision tree
    * for 0.0 probability, then just set the leaf to a nullptr.

gtsam/hybrid/HybridBayesNet.h

@@ -73,6 +73,8 @@ class GTSAM_EXPORT HybridBayesNet : public BayesNet<HybridConditional> {
         HybridConditional(boost::make_shared<DiscreteConditional>(key, table)));
   }
 
+  using Base::push_back;
+
   /// Get a specific Gaussian mixture by index `i`.
   GaussianMixture::shared_ptr atMixture(size_t i) const;
 
@@ -109,9 +111,17 @@ class GTSAM_EXPORT HybridBayesNet : public BayesNet<HybridConditional> {
    */
   VectorValues optimize(const DiscreteValues &assignment) const;
 
-  /// Prune the Hybrid Bayes Net given the discrete decision tree.
-  HybridBayesNet prune(
-      const DecisionTreeFactor::shared_ptr &discreteFactor) const;
+ protected:
+  /**
+   * @brief Get all the discrete conditionals as a decision tree factor.
+   *
+   * @return DecisionTreeFactor::shared_ptr
+   */
+  DecisionTreeFactor::shared_ptr discreteConditionals() const;
+
+ public:
+  /// Prune the Hybrid Bayes Net such that we have at most maxNrLeaves leaves.
+  HybridBayesNet prune(size_t maxNrLeaves) const;
 
   /// @}
 

gtsam/hybrid/HybridBayesTree.cpp

@@ -89,12 +89,12 @@ struct HybridAssignmentData {
         gaussianbayesTree_(gbt) {}
 
   /**
-   * @brief A function used during tree traversal that operators on each node
+   * @brief A function used during tree traversal that operates on each node
    * before visiting the node's children.
    *
    * @param node The current node being visited.
    * @param parentData The HybridAssignmentData from the parent node.
-   * @return HybridAssignmentData
+   * @return HybridAssignmentData which is passed to the children.
    */
   static HybridAssignmentData AssignmentPreOrderVisitor(
       const HybridBayesTree::sharedNode& node,
@@ -144,4 +144,61 @@ VectorValues HybridBayesTree::optimize(const DiscreteValues& assignment) const {
   return result;
 }
 
+/* ************************************************************************* */
+void HybridBayesTree::prune(const size_t maxNrLeaves) {
+  auto decisionTree = boost::dynamic_pointer_cast<DecisionTreeFactor>(
+      this->roots_.at(0)->conditional()->inner());
+
+  DecisionTreeFactor prunedDiscreteFactor = decisionTree->prune(maxNrLeaves);
+  decisionTree->root_ = prunedDiscreteFactor.root_;
+
+  /// Helper struct for pruning the hybrid bayes tree.
+  struct HybridPrunerData {
+    /// The discrete decision tree after pruning.
+    DecisionTreeFactor prunedDiscreteFactor;
+    HybridPrunerData(const DecisionTreeFactor& prunedDiscreteFactor,
+                     const HybridBayesTree::sharedNode& parentClique)
+        : prunedDiscreteFactor(prunedDiscreteFactor) {}
+
+    /**
+     * @brief A function used during tree traversal that operates on each node
+     * before visiting the node's children.
+     *
+     * @param node The current node being visited.
+     * @param parentData The data from the parent node.
+     * @return HybridPrunerData which is passed to the children.
+     */
+    static HybridPrunerData AssignmentPreOrderVisitor(
+        const HybridBayesTree::sharedNode& clique,
+        HybridPrunerData& parentData) {
+      // Get the conditional
+      HybridConditional::shared_ptr conditional = clique->conditional();
+
+      // If conditional is hybrid, we prune it.
+      if (conditional->isHybrid()) {
+        auto gaussianMixture = conditional->asMixture();
+
+        // Check if the number of discrete keys match,
+        // else we get an assignment error.
+        // TODO(Varun) Update prune method to handle assignment subset?
+        if (gaussianMixture->discreteKeys() ==
+            parentData.prunedDiscreteFactor.discreteKeys()) {
+          gaussianMixture->prune(parentData.prunedDiscreteFactor);
+        }
+      }
+      return parentData;
+    }
+  };
+
+  HybridPrunerData rootData(prunedDiscreteFactor, 0);
+  {
+    treeTraversal::no_op visitorPost;
+    // Limits OpenMP threads since we're mixing TBB and OpenMP
+    TbbOpenMPMixedScope threadLimiter;
+    treeTraversal::DepthFirstForestParallel(
+        *this, rootData, HybridPrunerData::AssignmentPreOrderVisitor,
+        visitorPost);
+  }
+}
+
 }  // namespace gtsam

gtsam/hybrid/HybridBayesTree.h

@@ -88,6 +88,13 @@ class GTSAM_EXPORT HybridBayesTree : public BayesTree<HybridBayesTreeClique> {
    */
   VectorValues optimize(const DiscreteValues& assignment) const;
 
+  /**
+   * @brief Prune the underlying Bayes tree.
+   *
+   * @param maxNumberLeaves The max number of leaf nodes to keep.
+   */
+  void prune(const size_t maxNumberLeaves);
+
   /// @}
 
  private:

gtsam/hybrid/HybridConditional.h

@@ -34,8 +34,6 @@
 
 namespace gtsam {
 
-class HybridGaussianFactorGraph;
-
 /**
  * Hybrid Conditional Density
  *

gtsam/hybrid/HybridFactorGraph.h

@@ -135,6 +135,28 @@ class HybridFactorGraph : public FactorGraph<HybridFactor> {
       push_hybrid(p);
     }
   }
+
+  /// Get all the discrete keys in the factor graph.
+  const KeySet discreteKeys() const {
+    KeySet discrete_keys;
+    for (auto& factor : factors_) {
+      for (const DiscreteKey& k : factor->discreteKeys()) {
+        discrete_keys.insert(k.first);
+      }
+    }
+    return discrete_keys;
+  }
+
+  /// Get all the continuous keys in the factor graph.
+  const KeySet continuousKeys() const {
+    KeySet keys;
+    for (auto& factor : factors_) {
+      for (const Key& key : factor->continuousKeys()) {
+        keys.insert(key);
+      }
+    }
+    return keys;
+  }
 };
 
 }  // namespace gtsam

gtsam/hybrid/HybridGaussianFactorGraph.cpp

@@ -404,31 +404,9 @@ void HybridGaussianFactorGraph::add(DecisionTreeFactor::shared_ptr factor) {
   FactorGraph::add(boost::make_shared<HybridDiscreteFactor>(factor));
 }
 
-/* ************************************************************************ */
-const KeySet HybridGaussianFactorGraph::getDiscreteKeys() const {
-  KeySet discrete_keys;
-  for (auto &factor : factors_) {
-    for (const DiscreteKey &k : factor->discreteKeys()) {
-      discrete_keys.insert(k.first);
-    }
-  }
-  return discrete_keys;
-}
-
-/* ************************************************************************ */
-const KeySet HybridGaussianFactorGraph::getContinuousKeys() const {
-  KeySet keys;
-  for (auto &factor : factors_) {
-    for (const Key &key : factor->continuousKeys()) {
-      keys.insert(key);
-    }
-  }
-  return keys;
-}
-
 /* ************************************************************************ */
 const Ordering HybridGaussianFactorGraph::getHybridOrdering() const {
-  KeySet discrete_keys = getDiscreteKeys();
+  KeySet discrete_keys = discreteKeys();
   for (auto &factor : factors_) {
     for (const DiscreteKey &k : factor->discreteKeys()) {
       discrete_keys.insert(k.first);

gtsam/hybrid/HybridGaussianFactorGraph.h

@@ -161,12 +161,6 @@ class GTSAM_EXPORT HybridGaussianFactorGraph
     }
   }
 
-  /// Get all the discrete keys in the factor graph.
-  const KeySet getDiscreteKeys() const;
-
-  /// Get all the continuous keys in the factor graph.
-  const KeySet getContinuousKeys() const;
-
   /**
    * @brief Return a Colamd constrained ordering where the discrete keys are
    * eliminated after the continuous keys.