borglab · dellaert · Jan 8, 2023 · Jan 6, 2023 · Jan 7, 2023 · Jan 7, 2023
diff --git a/gtsam/hybrid/GaussianMixture.cpp b/gtsam/hybrid/GaussianMixture.cpp
@@ -59,6 +59,8 @@ GaussianMixture::GaussianMixture(
                       Conditionals(discreteParents, conditionals)) {}
 
 /* *******************************************************************************/
+// TODO(dellaert): This is copy/paste: GaussianMixture should be derived from
+// GaussianMixtureFactor, no?
 GaussianFactorGraphTree GaussianMixture::add(
     const GaussianFactorGraphTree &sum) const {
   using Y = GraphAndConstant;

diff --git a/gtsam/hybrid/HybridBayesNet.cpp b/gtsam/hybrid/HybridBayesNet.cpp
@@ -141,8 +141,8 @@ std::function<double(const Assignment<Key> &, double)> prunerFunc(
 
 /* ************************************************************************* */
 void HybridBayesNet::updateDiscreteConditionals(
-    const DecisionTreeFactor::shared_ptr &prunedDecisionTree) {
-  KeyVector prunedTreeKeys = prunedDecisionTree->keys();
+    const DecisionTreeFactor &prunedDecisionTree) {
+  KeyVector prunedTreeKeys = prunedDecisionTree.keys();
 
   // Loop with index since we need it later.
   for (size_t i = 0; i < this->size(); i++) {
@@ -154,7 +154,7 @@ void HybridBayesNet::updateDiscreteConditionals(
       auto discreteTree =
           boost::dynamic_pointer_cast<DecisionTreeFactor::ADT>(discrete);
       DecisionTreeFactor::ADT prunedDiscreteTree =
-          discreteTree->apply(prunerFunc(*prunedDecisionTree, *conditional));
+          discreteTree->apply(prunerFunc(prunedDecisionTree, *conditional));
 
       // Create the new (hybrid) conditional
       KeyVector frontals(discrete->frontals().begin(),
@@ -173,9 +173,7 @@ void HybridBayesNet::updateDiscreteConditionals(
 HybridBayesNet HybridBayesNet::prune(size_t maxNrLeaves) {
   // Get the decision tree of only the discrete keys
   auto discreteConditionals = this->discreteConditionals();
-  const DecisionTreeFactor::shared_ptr decisionTree =
-      boost::make_shared<DecisionTreeFactor>(
-          discreteConditionals->prune(maxNrLeaves));
+  const auto decisionTree = discreteConditionals->prune(maxNrLeaves);
 
   this->updateDiscreteConditionals(decisionTree);
 
@@ -194,7 +192,7 @@ HybridBayesNet HybridBayesNet::prune(size_t maxNrLeaves) {
     if (auto gm = conditional->asMixture()) {
       // Make a copy of the Gaussian mixture and prune it!
       auto prunedGaussianMixture = boost::make_shared<GaussianMixture>(*gm);
-      prunedGaussianMixture->prune(*decisionTree);  // imperative :-(
+      prunedGaussianMixture->prune(decisionTree);  // imperative :-(
 
       // Type-erase and add to the pruned Bayes Net fragment.
       prunedBayesNetFragment.push_back(prunedGaussianMixture);

diff --git a/gtsam/hybrid/HybridBayesNet.h b/gtsam/hybrid/HybridBayesNet.h
@@ -51,33 +51,51 @@ class GTSAM_EXPORT HybridBayesNet : public BayesNet<HybridConditional> {
   /// @{
 
   /// GTSAM-style printing
-  void print(
-      const std::string &s = "",
-      const KeyFormatter &formatter = DefaultKeyFormatter) const override;
+  void print(const std::string &s = "", const KeyFormatter &formatter =
+                                            DefaultKeyFormatter) const override;
 
   /// GTSAM-style equals
-  bool equals(const This& fg, double tol = 1e-9) const;
-  
+  bool equals(const This &fg, double tol = 1e-9) const;
+
   /// @}
   /// @name Standard Interface
   /// @{
 
-  /// Add HybridConditional to Bayes Net
-  using Base::emplace_shared;
+  /**
+   * @brief Add a hybrid conditional using a shared_ptr.
+   *
+   * This is the "native" push back, as this class stores hybrid conditionals.
+   */
+  void push_back(boost::shared_ptr<HybridConditional> conditional) {
+    factors_.push_back(conditional);
+  }
 
-  /// Add a conditional directly using a pointer.
+  /**
+   * Preferred: add a conditional directly using a pointer.
+   *
+   * Examples:
+   *   hbn.emplace_back(new GaussianMixture(...)));
+   *   hbn.emplace_back(new GaussianConditional(...)));
+   *   hbn.emplace_back(new DiscreteConditional(...)));
+   */
   template <class Conditional>
   void emplace_back(Conditional *conditional) {
     factors_.push_back(boost::make_shared<HybridConditional>(
         boost::shared_ptr<Conditional>(conditional)));
   }
 
-  /// Add a conditional directly using a shared_ptr.
-  void push_back(boost::shared_ptr<HybridConditional> conditional) {
-    factors_.push_back(conditional);
-  }
-
-  /// Add a conditional directly using implicit conversion.
+  /**
+   * Add a conditional using a shared_ptr, using implicit conversion to
+   * a HybridConditional.
+   *
+   * This is useful when you create a conditional shared pointer as you need it
+   * somewhere else.
+   *
+   * Example:
+   *   auto shared_ptr_to_a_conditional =
+   *     boost::make_shared<GaussianMixture>(...);
+   *  hbn.push_back(shared_ptr_to_a_conditional);
+   */
   void push_back(HybridConditional &&conditional) {
     factors_.push_back(
         boost::make_shared<HybridConditional>(std::move(conditional)));
@@ -214,8 +232,7 @@ class GTSAM_EXPORT HybridBayesNet : public BayesNet<HybridConditional> {
    *
    * @param prunedDecisionTree
    */
-  void updateDiscreteConditionals(
-      const DecisionTreeFactor::shared_ptr &prunedDecisionTree);
+  void updateDiscreteConditionals(const DecisionTreeFactor &prunedDecisionTree);
 
   /** Serialization function */
   friend class boost::serialization::access;

diff --git a/gtsam/hybrid/HybridBayesTree.h b/gtsam/hybrid/HybridBayesTree.h
@@ -132,18 +132,15 @@ struct traits<HybridBayesTree> : public Testable<HybridBayesTree> {};
  * This object stores parent keys in our base type factor so that
  * eliminating those parent keys will pull this subtree into the
  * elimination.
- * This does special stuff for the hybrid case.
  *
- * @tparam CLIQUE
+ * This is a template instantiation for hybrid Bayes tree cliques, storing both
+ * the regular keys *and* discrete keys in the HybridConditional.
  */
-template <class CLIQUE>
-class BayesTreeOrphanWrapper<
-    CLIQUE, typename std::enable_if<
-                boost::is_same<CLIQUE, HybridBayesTreeClique>::value> >
-    : public CLIQUE::ConditionalType {
+template <>
+class BayesTreeOrphanWrapper<HybridBayesTreeClique> : public HybridConditional {
  public:
-  typedef CLIQUE CliqueType;
-  typedef typename CLIQUE::ConditionalType Base;
+  typedef HybridBayesTreeClique CliqueType;
+  typedef HybridConditional Base;
 
   boost::shared_ptr<CliqueType> clique;
 

diff --git a/gtsam/hybrid/HybridDiscreteFactor.cpp b/gtsam/hybrid/HybridDiscreteFactor.cpp
diff --git a/gtsam/hybrid/HybridDiscreteFactor.h b/gtsam/hybrid/HybridDiscreteFactor.h
diff --git a/gtsam/hybrid/HybridFactor.h b/gtsam/hybrid/HybridFactor.h
@@ -67,11 +67,10 @@ DiscreteKeys CollectDiscreteKeys(const DiscreteKeys &key1,
                                  const DiscreteKeys &key2);
 
 /**
- * Base class for hybrid probabilistic factors
+ * Base class for *truly* hybrid probabilistic factors
  *
  * Examples:
- *  - HybridGaussianFactor
- *  - HybridDiscreteFactor
+ *  - MixtureFactor
  *  - GaussianMixtureFactor
  *  - GaussianMixture
  *

diff --git a/gtsam/hybrid/HybridFactorGraph.cpp b/gtsam/hybrid/HybridFactorGraph.cpp
@@ -0,0 +1,78 @@
+/* ----------------------------------------------------------------------------
+
+ * GTSAM Copyright 2010-2022, Georgia Tech Research Corporation,
+ * Atlanta, Georgia 30332-0415
+ * All Rights Reserved
+ * Authors: Frank Dellaert, et al. (see THANKS for the full author list)
+
+ * See LICENSE for the license information
+
+ * -------------------------------------------------------------------------- */
+
+/**
+ * @file   HybridFactorGraph.cpp
+ * @brief  Factor graph with utilities for hybrid factors.
+ * @author Varun Agrawal
+ * @author Frank Dellaert
+ * @date   January, 2023
+ */
+
+#include <gtsam/discrete/DecisionTreeFactor.h>
+#include <gtsam/hybrid/HybridFactorGraph.h>
+
+#include <boost/format.hpp>
+
+namespace gtsam {
+
+/* ************************************************************************* */
+DiscreteKeys HybridFactorGraph::discreteKeys() const {
+  DiscreteKeys keys;
+  for (auto& factor : factors_) {
+    if (auto p = boost::dynamic_pointer_cast<DecisionTreeFactor>(factor)) {
+      for (const DiscreteKey& key : p->discreteKeys()) {
+        keys.push_back(key);
+      }
+    }
+    if (auto p = boost::dynamic_pointer_cast<HybridFactor>(factor)) {
+      for (const DiscreteKey& key : p->discreteKeys()) {
+        keys.push_back(key);
+      }
+    }
+  }
+  return keys;
+}
+
+/* ************************************************************************* */
+KeySet HybridFactorGraph::discreteKeySet() const {
+  KeySet keys;
+  for (const DiscreteKey& k : discreteKeys()) {
+    keys.insert(k.first);
+  }
+  return keys;
+}
+
+/* ************************************************************************* */
+std::unordered_map<Key, DiscreteKey> HybridFactorGraph::discreteKeyMap() const {
+  std::unordered_map<Key, DiscreteKey> result;
+  for (const DiscreteKey& k : discreteKeys()) {
+    result[k.first] = k;
+  }
+  return result;
+}
+
+/* ************************************************************************* */
+const KeySet HybridFactorGraph::continuousKeySet() const {
+  KeySet keys;
+  for (auto& factor : factors_) {
+    if (auto p = boost::dynamic_pointer_cast<HybridFactor>(factor)) {
+      for (const Key& key : p->continuousKeys()) {
+        keys.insert(key);
+      }
+    }
+  }
+  return keys;
+}
+
+/* ************************************************************************* */
+
+}  // namespace gtsam