Hybrid Multifrontal

gtsam/hybrid/HybridBayesTree.cpp

@@ -73,6 +73,8 @@ struct HybridAssignmentData {
   GaussianBayesTree::sharedNode parentClique_;
   // The gaussian bayes tree that will be recursively created.
   GaussianBayesTree* gaussianbayesTree_;
+  // Flag indicating if all the nodes are valid. Used in optimize().
+  bool valid_;
 
   /**
    * @brief Construct a new Hybrid Assignment Data object.
@@ -83,10 +85,13 @@ struct HybridAssignmentData {
    */
   HybridAssignmentData(const DiscreteValues& assignment,
                        const GaussianBayesTree::sharedNode& parentClique,
-                       GaussianBayesTree* gbt)
+                       GaussianBayesTree* gbt, bool valid = true)
       : assignment_(assignment),
         parentClique_(parentClique),
-        gaussianbayesTree_(gbt) {}
+        gaussianbayesTree_(gbt),
+        valid_(valid) {}
+
+  bool isValid() const { return valid_; }
 
   /**
    * @brief A function used during tree traversal that operates on each node
@@ -101,6 +106,7 @@ struct HybridAssignmentData {
       HybridAssignmentData& parentData) {
     // Extract the gaussian conditional from the Hybrid clique
     HybridConditional::shared_ptr hybrid_conditional = node->conditional();
+
     GaussianConditional::shared_ptr conditional;
     if (hybrid_conditional->isHybrid()) {
       conditional = (*hybrid_conditional->asMixture())(parentData.assignment_);
@@ -111,15 +117,21 @@ struct HybridAssignmentData {
       conditional = boost::make_shared<GaussianConditional>();
     }
 
-    // Create the GaussianClique for the current node
-    auto clique = boost::make_shared<GaussianBayesTree::Node>(conditional);
-    // Add the current clique to the GaussianBayesTree.
-    parentData.gaussianbayesTree_->addClique(clique, parentData.parentClique_);
+    GaussianBayesTree::sharedNode clique;
+    if (conditional) {
+      // Create the GaussianClique for the current node
+      clique = boost::make_shared<GaussianBayesTree::Node>(conditional);
+      // Add the current clique to the GaussianBayesTree.
+      parentData.gaussianbayesTree_->addClique(clique,
+                                               parentData.parentClique_);
+    } else {
+      parentData.valid_ = false;
+    }
 
     // Create new HybridAssignmentData where the current node is the parent
     // This will be passed down to the children nodes
     HybridAssignmentData data(parentData.assignment_, clique,
-                              parentData.gaussianbayesTree_);
+                              parentData.gaussianbayesTree_, parentData.valid_);
     return data;
   }
 };
@@ -138,6 +150,9 @@ VectorValues HybridBayesTree::optimize(const DiscreteValues& assignment) const {
         visitorPost);
   }
 
+  if (!rootData.isValid()) {
+    return VectorValues();
+  }
   VectorValues result = gbt.optimize();
 
   // Return the optimized bayes net result.
@@ -151,6 +166,8 @@ void HybridBayesTree::prune(const size_t maxNrLeaves) {
 
   DecisionTreeFactor prunedDecisionTree = decisionTree->prune(maxNrLeaves);
   decisionTree->root_ = prunedDecisionTree.root_;
+  // this->print();
+  // decisionTree->print("", DefaultKeyFormatter);
 
   /// Helper struct for pruning the hybrid bayes tree.
   struct HybridPrunerData {

gtsam/hybrid/HybridBayesTree.h

@@ -50,9 +50,12 @@ class GTSAM_EXPORT HybridBayesTreeClique
   typedef boost::shared_ptr<This> shared_ptr;
   typedef boost::weak_ptr<This> weak_ptr;
   HybridBayesTreeClique() {}
-  virtual ~HybridBayesTreeClique() {}
   HybridBayesTreeClique(const boost::shared_ptr<HybridConditional>& conditional)
       : Base(conditional) {}
+  ///< Copy constructor
+  HybridBayesTreeClique(const HybridBayesTreeClique& clique) : Base(clique) {}
+
+  virtual ~HybridBayesTreeClique() {}
 };
 
 /* ************************************************************************* */

gtsam/hybrid/HybridGaussianFactorGraph.cpp

@@ -546,6 +546,24 @@ HybridGaussianFactorGraph::continuousDelta(
   return delta_tree;
 }
 
+/* ************************************************************************ */
+DecisionTree<Key, VectorValues::shared_ptr>
+HybridGaussianFactorGraph::continuousDelta(
+    const DiscreteKeys &discrete_keys,
+    const boost::shared_ptr<BayesTreeType> &continuousBayesTree,
+    const std::vector<DiscreteValues> &assignments) const {
+  // Create a decision tree of all the different VectorValues
+  std::vector<VectorValues::shared_ptr> vector_values;
+  for (const DiscreteValues &assignment : assignments) {
+    VectorValues values = continuousBayesTree->optimize(assignment);
+    vector_values.push_back(boost::make_shared<VectorValues>(values));
+  }
+  DecisionTree<Key, VectorValues::shared_ptr> delta_tree(discrete_keys,
+                                                         vector_values);
+
+  return delta_tree;
+}
+
 /* ************************************************************************ */
 AlgebraicDecisionTree<Key> HybridGaussianFactorGraph::continuousProbPrimes(
     const DiscreteKeys &orig_discrete_keys,
@@ -584,6 +602,67 @@ AlgebraicDecisionTree<Key> HybridGaussianFactorGraph::continuousProbPrimes(
   return probPrimeTree;
 }
 
+/* ************************************************************************ */
+AlgebraicDecisionTree<Key> HybridGaussianFactorGraph::continuousProbPrimes(
+    const DiscreteKeys &orig_discrete_keys,
+    const boost::shared_ptr<BayesTreeType> &continuousBayesTree) const {
+  // Generate all possible assignments.
+  const std::vector<DiscreteValues> assignments =
+      DiscreteValues::CartesianProduct(orig_discrete_keys);
+
+  // Save a copy of the original discrete key ordering
+  DiscreteKeys discrete_keys(orig_discrete_keys);
+  // Reverse discrete keys order for correct tree construction
+  std::reverse(discrete_keys.begin(), discrete_keys.end());
+
+  // Create a decision tree of all the different VectorValues
+  DecisionTree<Key, VectorValues::shared_ptr> delta_tree =
+      this->continuousDelta(discrete_keys, continuousBayesTree, assignments);
+
+  // Get the probPrime tree with the correct leaf probabilities
+  std::vector<double> probPrimes;
+  for (const DiscreteValues &assignment : assignments) {
+    VectorValues delta = *delta_tree(assignment);
+
+    // If VectorValues is empty, it means this is a pruned branch.
+    // Set thr probPrime to 0.0.
+    if (delta.size() == 0) {
+      probPrimes.push_back(0.0);
+      continue;
+    }
+
+    // Compute the error given the delta and the assignment.
+    double error = this->error(delta, assignment);
+    probPrimes.push_back(exp(-error));
+  }
+
+  AlgebraicDecisionTree<Key> probPrimeTree(discrete_keys, probPrimes);
+  return probPrimeTree;
+}
+
+/* ************************************************************************ */
+std::pair<Ordering, Ordering>
+HybridGaussianFactorGraph::separateContinuousDiscreteOrdering(
+    const Ordering &ordering) const {
+  KeySet all_continuous_keys = this->continuousKeys();
+  KeySet all_discrete_keys = this->discreteKeys();
+  Ordering continuous_ordering, discrete_ordering;
+
+  for (auto &&key : ordering) {
+    if (std::find(all_continuous_keys.begin(), all_continuous_keys.end(),
+                  key) != all_continuous_keys.end()) {
+      continuous_ordering.push_back(key);
+    } else if (std::find(all_discrete_keys.begin(), all_discrete_keys.end(),
+                         key) != all_discrete_keys.end()) {
+      discrete_ordering.push_back(key);
+    } else {
+      throw std::runtime_error("Key in ordering not present in factors.");
+    }
+  }
+
+  return std::make_pair(continuous_ordering, discrete_ordering);
+}
+
 /* ************************************************************************ */
 boost::shared_ptr<HybridGaussianFactorGraph::BayesNetType>
 HybridGaussianFactorGraph::eliminateHybridSequential(
@@ -640,25 +719,99 @@ boost::shared_ptr<HybridGaussianFactorGraph::BayesNetType>
 HybridGaussianFactorGraph::eliminateSequential(
     const Ordering &ordering, const Eliminate &function,
     OptionalVariableIndex variableIndex) const {
-  KeySet all_continuous_keys = this->continuousKeys();
-  KeySet all_discrete_keys = this->discreteKeys();
+  // Segregate the continuous and the discrete keys
   Ordering continuous_ordering, discrete_ordering;
+  std::tie(continuous_ordering, discrete_ordering) =
+      this->separateContinuousDiscreteOrdering(ordering);
+
+  return this->eliminateHybridSequential(continuous_ordering, discrete_ordering,
+                                         function, variableIndex);
+}
+
+/* ************************************************************************ */
+boost::shared_ptr<HybridGaussianFactorGraph::BayesTreeType>
+HybridGaussianFactorGraph::eliminateHybridMultifrontal(
+    const boost::optional<Ordering> continuous,
+    const boost::optional<Ordering> discrete, const Eliminate &function,
+    OptionalVariableIndex variableIndex) const {
+  Ordering continuous_ordering =
+      continuous ? *continuous : Ordering(this->continuousKeys());
+  Ordering discrete_ordering =
+      discrete ? *discrete : Ordering(this->discreteKeys());
+
+  // Eliminate continuous
+  HybridBayesTree::shared_ptr bayesTree;
+  HybridGaussianFactorGraph::shared_ptr discreteGraph;
+  std::tie(bayesTree, discreteGraph) =
+      BaseEliminateable::eliminatePartialMultifrontal(continuous_ordering,
+                                                      function, variableIndex);
+
+  // Get the last continuous conditional which will have all the discrete
+  Key last_continuous_key =
+      continuous_ordering.at(continuous_ordering.size() - 1);
+  auto last_conditional = (*bayesTree)[last_continuous_key]->conditional();
+  DiscreteKeys discrete_keys = last_conditional->discreteKeys();
+
+  // If not discrete variables, return the eliminated bayes net.
+  if (discrete_keys.size() == 0) {
+    return bayesTree;
+  }
+
+  AlgebraicDecisionTree<Key> probPrimeTree =
+      this->continuousProbPrimes(discrete_keys, bayesTree);
+
+  discreteGraph->add(DecisionTreeFactor(discrete_keys, probPrimeTree));
+
+  auto updatedBayesTree =
+      discreteGraph->BaseEliminateable::eliminateMultifrontal(discrete_ordering,
+                                                              function);
+
+  auto discrete_clique = (*updatedBayesTree)[discrete_ordering.at(0)];
+
+  std::set<HybridBayesTreeClique::shared_ptr> clique_set;
+  for (auto node : bayesTree->nodes()) {
+    clique_set.insert(node.second);
+  }
+
+  // Set the root of the bayes tree as the discrete clique
+  for (auto clique : clique_set) {
+    if (clique->conditional()->parents() ==
+        discrete_clique->conditional()->frontals()) {
+      updatedBayesTree->addClique(clique, discrete_clique);
 
-  // Segregate the continuous and the discrete keys
-  for (auto &&key : ordering) {
-    if (std::find(all_continuous_keys.begin(), all_continuous_keys.end(),
-                  key) != all_continuous_keys.end()) {
-      continuous_ordering.push_back(key);
-    } else if (std::find(all_discrete_keys.begin(), all_discrete_keys.end(),
-                         key) != all_discrete_keys.end()) {
-      discrete_ordering.push_back(key);
     } else {
-      throw std::runtime_error("Key in ordering not present in factors.");
+      // Remove the clique from the children of the parents since it will get
+      // added again in addClique.
+      auto clique_it = std::find(clique->parent()->children.begin(),
+                                 clique->parent()->children.end(), clique);
+      clique->parent()->children.erase(clique_it);
+      updatedBayesTree->addClique(clique, clique->parent());
     }
   }
+  return updatedBayesTree;
+}
+
+/* ************************************************************************ */
+boost::shared_ptr<HybridGaussianFactorGraph::BayesTreeType>
+HybridGaussianFactorGraph::eliminateMultifrontal(
+    OptionalOrderingType orderingType, const Eliminate &function,
+    OptionalVariableIndex variableIndex) const {
+  return BaseEliminateable::eliminateMultifrontal(orderingType, function,
+                                                  variableIndex);
+}
+
+/* ************************************************************************ */
+boost::shared_ptr<HybridGaussianFactorGraph::BayesTreeType>
+HybridGaussianFactorGraph::eliminateMultifrontal(
+    const Ordering &ordering, const Eliminate &function,
+    OptionalVariableIndex variableIndex) const {
+  // Segregate the continuous and the discrete keys
+  Ordering continuous_ordering, discrete_ordering;
+  std::tie(continuous_ordering, discrete_ordering) =
+      this->separateContinuousDiscreteOrdering(ordering);
 
-  return this->eliminateHybridSequential(continuous_ordering,
-                                         discrete_ordering);
+  return this->eliminateHybridMultifrontal(
+      continuous_ordering, discrete_ordering, function, variableIndex);
 }
 
 }  // namespace gtsam

gtsam/hybrid/HybridGaussianFactorGraph.h

@@ -231,6 +231,10 @@ class GTSAM_EXPORT HybridGaussianFactorGraph
       const DiscreteKeys& discrete_keys,
       const boost::shared_ptr<BayesNetType>& continuousBayesNet,
       const std::vector<DiscreteValues>& assignments) const;
+  DecisionTree<Key, VectorValues::shared_ptr> continuousDelta(
+      const DiscreteKeys& discrete_keys,
+      const boost::shared_ptr<BayesTreeType>& continuousBayesTree,
+      const std::vector<DiscreteValues>& assignments) const;
 
   /**
    * @brief Compute the unnormalized probabilities of the continuous variables
@@ -244,6 +248,12 @@ class GTSAM_EXPORT HybridGaussianFactorGraph
   AlgebraicDecisionTree<Key> continuousProbPrimes(
       const DiscreteKeys& discrete_keys,
       const boost::shared_ptr<BayesNetType>& continuousBayesNet) const;
+  AlgebraicDecisionTree<Key> continuousProbPrimes(
+      const DiscreteKeys& discrete_keys,
+      const boost::shared_ptr<BayesTreeType>& continuousBayesTree) const;
+
+  std::pair<Ordering, Ordering> separateContinuousDiscreteOrdering(
+      const Ordering& ordering) const;
 
   /**
    * @brief Custom elimination function which computes the correct
@@ -269,6 +279,22 @@ class GTSAM_EXPORT HybridGaussianFactorGraph
       const Eliminate& function = EliminationTraitsType::DefaultEliminate,
       OptionalVariableIndex variableIndex = boost::none) const;
 
+  boost::shared_ptr<BayesTreeType> eliminateHybridMultifrontal(
+      const boost::optional<Ordering> continuous = boost::none,
+      const boost::optional<Ordering> discrete = boost::none,
+      const Eliminate& function = EliminationTraitsType::DefaultEliminate,
+      OptionalVariableIndex variableIndex = boost::none) const;
+
+  boost::shared_ptr<BayesTreeType> eliminateMultifrontal(
+      OptionalOrderingType orderingType = boost::none,
+      const Eliminate& function = EliminationTraitsType::DefaultEliminate,
+      OptionalVariableIndex variableIndex = boost::none) const;
+
+  boost::shared_ptr<BayesTreeType> eliminateMultifrontal(
+      const Ordering& ordering,
+      const Eliminate& function = EliminationTraitsType::DefaultEliminate,
+      OptionalVariableIndex variableIndex = boost::none) const;
+
   /**
    * @brief Return a Colamd constrained ordering where the discrete keys are
    * eliminated after the continuous keys.

gtsam/hybrid/tests/testHybridBayesTree.cpp

@@ -136,6 +136,13 @@ TEST(HybridBayesTree, Optimize) {
     dfg.push_back(
         boost::dynamic_pointer_cast<DecisionTreeFactor>(factor->inner()));
   }
+
+  // Add the probabilities for each branch
+  DiscreteKeys discrete_keys = {{M(0), 2}, {M(1), 2}, {M(2), 2}};
+  vector<double> probs = {0.012519475, 0.041280228, 0.075018647, 0.081663656,
+                          0.037152205, 0.12248971,  0.07349729,  0.08};
+  AlgebraicDecisionTree<Key> potentials(discrete_keys, probs);
+  dfg.emplace_shared<DecisionTreeFactor>(discrete_keys, probs);
 
   DiscreteValues expectedMPE = dfg.optimize();
   VectorValues expectedValues = hybridBayesNet->optimize(expectedMPE);