From 5c4038c7c02ed49958fe12336cfef7a9c76bf039 Mon Sep 17 00:00:00 2001
From: Frank Dellaert <dellaert@gmail.com>
Date: Sun, 2 Jan 2022 18:08:09 -0500
Subject: [PATCH] Fixed dot to have right arguments

---
 gtsam/discrete/DecisionTreeFactor.cpp | 25 +++++++++++++++++++++++++
 gtsam/discrete/DecisionTreeFactor.h   | 14 ++++++++++++++
 gtsam/discrete/discrete.i             |  4 +++-
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/gtsam/discrete/DecisionTreeFactor.cpp b/gtsam/discrete/DecisionTreeFactor.cpp
index 7aed00c57..75018cf92 100644
--- a/gtsam/discrete/DecisionTreeFactor.cpp
+++ b/gtsam/discrete/DecisionTreeFactor.cpp
@@ -153,6 +153,31 @@ namespace gtsam {
     return result;
   }
 
+  /* ************************************************************************* */
+  static std::string valueFormatter(const double& v) {
+    return (boost::format("%4.2g") % v).str();
+  }
+
+  /** output to graphviz format, stream version */
+  void DecisionTreeFactor::dot(std::ostream& os,
+                               const KeyFormatter& keyFormatter,
+                               bool showZero) const {
+    Potentials::dot(os, keyFormatter, valueFormatter, showZero);
+  }
+
+  /** output to graphviz format, open a file */
+  void DecisionTreeFactor::dot(const std::string& name,
+                               const KeyFormatter& keyFormatter,
+                               bool showZero) const {
+    Potentials::dot(name, keyFormatter, valueFormatter, showZero);
+  }
+
+  /** output to graphviz format string */
+  std::string DecisionTreeFactor::dot(const KeyFormatter& keyFormatter,
+                                      bool showZero) const {
+    return Potentials::dot(keyFormatter, valueFormatter, showZero);
+  }
+
   /* ************************************************************************* */
   std::string DecisionTreeFactor::markdown(
       const KeyFormatter& keyFormatter) const {
diff --git a/gtsam/discrete/DecisionTreeFactor.h b/gtsam/discrete/DecisionTreeFactor.h
index f90af56dd..46509db82 100644
--- a/gtsam/discrete/DecisionTreeFactor.h
+++ b/gtsam/discrete/DecisionTreeFactor.h
@@ -178,6 +178,20 @@ namespace gtsam {
     /// @name Wrapper support
     /// @{
     
+    /** output to graphviz format, stream version */
+    void dot(std::ostream& os,
+             const KeyFormatter& keyFormatter = DefaultKeyFormatter,
+             bool showZero = true) const;
+
+    /** output to graphviz format, open a file */
+    void dot(const std::string& name,
+             const KeyFormatter& keyFormatter = DefaultKeyFormatter,
+             bool showZero = true) const;
+
+    /** output to graphviz format string */
+    std::string dot(const KeyFormatter& keyFormatter = DefaultKeyFormatter,
+                    bool showZero = true) const;
+
     /// Render as markdown table.
     std::string markdown(
         const KeyFormatter& keyFormatter = DefaultKeyFormatter) const override;
diff --git a/gtsam/discrete/discrete.i b/gtsam/discrete/discrete.i
index 36caccfc8..5bd4a2913 100644
--- a/gtsam/discrete/discrete.i
+++ b/gtsam/discrete/discrete.i
@@ -46,7 +46,9 @@ virtual class DecisionTreeFactor : gtsam::DiscreteFactor {
              const gtsam::KeyFormatter& keyFormatter =
                  gtsam::DefaultKeyFormatter) const;
   bool equals(const gtsam::DecisionTreeFactor& other, double tol = 1e-9) const;
-  string dot(bool showZero = false) const;
+  string dot(
+      const gtsam::KeyFormatter& keyFormatter = gtsam::DefaultKeyFormatter,
+      bool showZero = false) const;
   std::vector<std::pair<DiscreteValues, double>> enumerate() const;
   string markdown(const gtsam::KeyFormatter& keyFormatter =
                       gtsam::DefaultKeyFormatter) const;