From 35e7acbf16548a286d3677717f8dfc1b9398d1f8 Mon Sep 17 00:00:00 2001
From: Frank Dellaert <dellaert@gmail.com>
Date: Mon, 27 Jan 2025 13:05:20 -0500
Subject: [PATCH] Print factors

---
 gtsam/discrete/DiscreteJunctionTree.cpp | 41 +++++++---
 gtsam/discrete/DiscreteJunctionTree.h   | 99 +++++++++++++++----------
 2 files changed, 90 insertions(+), 50 deletions(-)

diff --git a/gtsam/discrete/DiscreteJunctionTree.cpp b/gtsam/discrete/DiscreteJunctionTree.cpp
index dc24860eb..b4657ec8d 100644
--- a/gtsam/discrete/DiscreteJunctionTree.cpp
+++ b/gtsam/discrete/DiscreteJunctionTree.cpp
@@ -16,19 +16,42 @@
  * @author Richard Roberts
  */
 
-#include <gtsam/inference/JunctionTree-inst.h>
-#include <gtsam/discrete/DiscreteJunctionTree.h>
 #include <gtsam/discrete/DiscreteEliminationTree.h>
+#include <gtsam/discrete/DiscreteJunctionTree.h>
+#include <gtsam/inference/JunctionTree-inst.h>
 
 namespace gtsam {
 
-  // Instantiate base classes
-  template class EliminatableClusterTree<DiscreteBayesTree, DiscreteFactorGraph>;
-  template class JunctionTree<DiscreteBayesTree, DiscreteFactorGraph>;
+// Instantiate base classes
+template class EliminatableClusterTree<DiscreteBayesTree, DiscreteFactorGraph>;
+template class JunctionTree<DiscreteBayesTree, DiscreteFactorGraph>;
 
-  /* ************************************************************************* */
-  DiscreteJunctionTree::DiscreteJunctionTree(
-    const DiscreteEliminationTree& eliminationTree) :
-  Base(eliminationTree) {}
+/* ************************************************************************* */
+DiscreteJunctionTree::DiscreteJunctionTree(
+    const DiscreteEliminationTree& eliminationTree)
+    : Base(eliminationTree) {}
+/* ************************************************************************* */
+namespace {
+struct PrintForestVisitorPre {
+  const KeyFormatter& formatter;
+  PrintForestVisitorPre(const KeyFormatter& formatter) : formatter(formatter) {}
+  std::string operator()(
+      const std::shared_ptr<DiscreteJunctionTree::Cluster>& node,
+      const std::string& parentString) {
+    // Print the current node
+    node->print(parentString + "-", formatter);
+    node->factors.print(parentString + "-", formatter);
+    std::cout << std::endl;
+    // Increment the indentation
+    return parentString + "| ";
+  }
+};
+}  // namespace
 
+void DiscreteJunctionTree::print(const std::string& s,
+                                 const KeyFormatter& keyFormatter) const {
+  PrintForestVisitorPre visitor(keyFormatter);
+  treeTraversal::DepthFirstForest(*this, s, visitor);
 }
+
+}  // namespace gtsam
diff --git a/gtsam/discrete/DiscreteJunctionTree.h b/gtsam/discrete/DiscreteJunctionTree.h
index f6171c672..4b9241036 100644
--- a/gtsam/discrete/DiscreteJunctionTree.h
+++ b/gtsam/discrete/DiscreteJunctionTree.h
@@ -18,54 +18,71 @@
 
 #pragma once
 
-#include <gtsam/discrete/DiscreteFactorGraph.h>
 #include <gtsam/discrete/DiscreteBayesTree.h>
+#include <gtsam/discrete/DiscreteFactorGraph.h>
 #include <gtsam/inference/JunctionTree.h>
 
 namespace gtsam {
 
-  // Forward declarations
-  class DiscreteEliminationTree;
+// Forward declarations
+class DiscreteEliminationTree;
+
+/**
+ * An EliminatableClusterTree, i.e., a set of variable clusters with factors,
+ * arranged in a tree, with the additional property that it represents the
+ * clique tree associated with a Bayes net.
+ *
+ * In GTSAM a junction tree is an intermediate data structure in multifrontal
+ * variable elimination.  Each node is a cluster of factors, along with a
+ * clique of variables that are eliminated all at once. In detail, every node k
+ * represents a clique (maximal fully connected subset) of an associated chordal
+ * graph, such as a chordal Bayes net resulting from elimination.
+ *
+ * The difference with the BayesTree is that a JunctionTree stores factors,
+ * whereas a BayesTree stores conditionals, that are the product of eliminating
+ * the factors in the corresponding JunctionTree cliques.
+ *
+ * The tree structure and elimination method are exactly analogous to the
+ * EliminationTree, except that in the JunctionTree, at each node multiple
+ * variables are eliminated at a time.
+ *
+ * \ingroup Multifrontal
+ * @ingroup discrete
+ * \nosubgrouping
+ */
+class GTSAM_EXPORT DiscreteJunctionTree
+    : public JunctionTree<DiscreteBayesTree, DiscreteFactorGraph> {
+ public:
+  typedef JunctionTree<DiscreteBayesTree, DiscreteFactorGraph>
+      Base;                                  ///< Base class
+  typedef DiscreteJunctionTree This;         ///< This class
+  typedef std::shared_ptr<This> shared_ptr;  ///< Shared pointer to this class
+
+  /// @name Constructors
+  /// @{
 
   /**
-   * An EliminatableClusterTree, i.e., a set of variable clusters with factors, arranged in a tree,
-   * with the additional property that it represents the clique tree associated with a Bayes net.
-   *
-   * In GTSAM a junction tree is an intermediate data structure in multifrontal
-   * variable elimination.  Each node is a cluster of factors, along with a
-   * clique of variables that are eliminated all at once. In detail, every node k represents
-   * a clique (maximal fully connected subset) of an associated chordal graph, such as a
-   * chordal Bayes net resulting from elimination.
-   *
-   * The difference with the BayesTree is that a JunctionTree stores factors, whereas a
-   * BayesTree stores conditionals, that are the product of eliminating the factors in the
-   * corresponding JunctionTree cliques.
-   *
-   * The tree structure and elimination method are exactly analogous to the EliminationTree,
-   * except that in the JunctionTree, at each node multiple variables are eliminated at a time.
-   *
-   * \ingroup Multifrontal
-   * @ingroup discrete
-   * \nosubgrouping
+   * Build the elimination tree of a factor graph using precomputed column
+   * structure.
+   * @param factorGraph The factor graph for which to build the elimination tree
+   * @param structure The set of factors involving each variable.  If this is
+   * not precomputed, you can call the Create(const FactorGraph<DERIVEDFACTOR>&)
+   * named constructor instead.
+   * @return The elimination tree
    */
-  class GTSAM_EXPORT DiscreteJunctionTree :
-    public JunctionTree<DiscreteBayesTree, DiscreteFactorGraph> {
-  public:
-    typedef JunctionTree<DiscreteBayesTree, DiscreteFactorGraph> Base; ///< Base class
-    typedef DiscreteJunctionTree This; ///< This class
-    typedef std::shared_ptr<This> shared_ptr; ///< Shared pointer to this class
+  DiscreteJunctionTree(const DiscreteEliminationTree& eliminationTree);
 
-    /**
-    * Build the elimination tree of a factor graph using precomputed column structure.
-    * @param factorGraph The factor graph for which to build the elimination tree
-    * @param structure The set of factors involving each variable.  If this is not
-    * precomputed, you can call the Create(const FactorGraph<DERIVEDFACTOR>&)
-    * named constructor instead.
-    * @return The elimination tree
-    */
-    DiscreteJunctionTree(const DiscreteEliminationTree& eliminationTree);
-  };
+  /// @}
+  /// @name Testable
+  /// @{
 
-  /// typedef for wrapper:
-  using DiscreteCluster = DiscreteJunctionTree::Cluster;
-}
+  /** Print the tree to cout */
+  void print(const std::string& name = "DiscreteJunctionTree: ",
+             const KeyFormatter& formatter = DefaultKeyFormatter) const;
+
+  /// @}
+};
+
+/// typedef for wrapper:
+using DiscreteCluster = DiscreteJunctionTree::Cluster;
+}  // namespace gtsam