Generic dot export with DotWriter

gtsam/discrete/tests/testDiscreteFactorGraph.cpp

@@ -359,6 +359,31 @@ cout << unicorns;
 }
 #endif
 
+/* ************************************************************************* */
+TEST(DiscreteFactorGraph, Dot) {
+  // Declare a bunch of keys
+  DiscreteKey C(0, 2), A(1, 2), B(2, 2);
+
+  // Create Factor graph
+  DiscreteFactorGraph graph;
+  graph.add(C & A, "0.2 0.8 0.3 0.7");
+  graph.add(C & B, "0.1 0.9 0.4 0.6");
+
+  string actual = graph.dot();
+  string expected =
+      "graph {\n"
+      "  size=\"5,5\";\n"
+      "\n"
+      "  var0[label=\"0\"];\n"
+      "  var1[label=\"1\"];\n"
+      "  var2[label=\"2\"];\n"
+      "\n"
+      "  var0--var1;\n"
+      "  var0--var2;\n"
+      "}\n";
+  EXPECT(actual == expected);
+}
+
 /* ************************************************************************* */
 int main() {
 TestResult tr;

gtsam/inference/DotWriter.cpp

@@ -0,0 +1,93 @@
+/* ----------------------------------------------------------------------------
+
+ * GTSAM Copyright 2010-2021, 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   DotWriter.cpp
+ * @brief  Graphviz formatting for factor graphs.
+ * @author Frank Dellaert
+ * @date December, 2021
+ */
+
+#include <gtsam/base/Vector.h>
+#include <gtsam/inference/DotWriter.h>
+
+#include <ostream>
+
+using namespace std;
+
+namespace gtsam {
+
+void DotWriter::writePreamble(ostream* os) const {
+  *os << "graph {\n";
+  *os << "  size=\"" << figureWidthInches << "," << figureHeightInches
+      << "\";\n\n";
+}
+
+void DotWriter::DrawVariable(Key key, const KeyFormatter& keyFormatter,
+                             const boost::optional<Vector2>& position,
+                             ostream* os) {
+  // Label the node with the label from the KeyFormatter
+  *os << "  var" << key << "[label=\"" << keyFormatter(key) << "\"";
+  if (position) {
+    *os << ", pos=\"" << position->x() << "," << position->y() << "!\"";
+  }
+  *os << "];\n";
+}
+
+void DotWriter::DrawFactor(size_t i, const boost::optional<Vector2>& position,
+                           ostream* os) {
+  *os << "  factor" << i << "[label=\"\", shape=point";
+  if (position) {
+    *os << ", pos=\"" << position->x() << "," << position->y() << "!\"";
+  }
+  *os << "];\n";
+}
+
+void DotWriter::ConnectVariables(Key key1, Key key2, ostream* os) {
+  *os << "  var" << key1 << "--"
+      << "var" << key2 << ";\n";
+}
+
+void DotWriter::ConnectVariableFactor(Key key, size_t i, ostream* os) {
+  *os << "  var" << key << "--"
+      << "factor" << i << ";\n";
+}
+
+void DotWriter::ProcessFactor(size_t i, const KeyVector& keys,
+                              const boost::optional<Vector2>& position,
+                              ostream* os) const {
+  if (plotFactorPoints) {
+    if (binaryEdges && keys.size() == 2) {
+      ConnectVariables(keys[0], keys[1], os);
+    } else {
+      // Create dot for the factor.
+      DrawFactor(i, position, os);
+
+      // Make factor-variable connections
+      if (connectKeysToFactor) {
+        for (Key key : keys) {
+          ConnectVariableFactor(key, i, os);
+        }
+      }
+    }
+  } else {
+    // just connect variables in a clique
+    for (Key key1 : keys) {
+      for (Key key2 : keys) {
+        if (key2 > key1) {
+          ConnectVariables(key1, key2, os);
+        }
+      }
+    }
+  }
+}
+
+}  // namespace gtsam

gtsam/inference/DotWriter.h

@@ -0,0 +1,69 @@
+/* ----------------------------------------------------------------------------
+
+ * GTSAM Copyright 2010-2021, 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 DotWriter.h
+ * @brief Graphviz formatter
+ * @author Frank Dellaert
+ * @date December, 2021
+ */
+
+#pragma once
+
+#include <gtsam/base/FastVector.h>
+#include <gtsam/base/Vector.h>
+#include <gtsam/inference/Key.h>
+
+#include <iosfwd>
+
+namespace gtsam {
+
+/// Graphviz formatter.
+struct GTSAM_EXPORT DotWriter {
+  double figureWidthInches;   ///< The figure width on paper in inches
+  double figureHeightInches;  ///< The figure height on paper in inches
+  bool plotFactorPoints;  ///< Plots each factor as a dot between the variables
+  bool connectKeysToFactor;  ///< Draw a line from each key within a factor to
+                             ///< the dot of the factor
+  bool binaryEdges;          ///< just use non-dotted edges for binary factors
+
+  DotWriter()
+      : figureWidthInches(5),
+        figureHeightInches(5),
+        plotFactorPoints(true),
+        connectKeysToFactor(true),
+        binaryEdges(true) {}
+
+  /// Write out preamble, including size.
+  void writePreamble(std::ostream* os) const;
+
+  /// Create a variable dot fragment.
+  static void DrawVariable(Key key, const KeyFormatter& keyFormatter,
+                           const boost::optional<Vector2>& position,
+                           std::ostream* os);
+
+  /// Create factor dot.
+  static void DrawFactor(size_t i, const boost::optional<Vector2>& position,
+                         std::ostream* os);
+
+  /// Connect two variables.
+  static void ConnectVariables(Key key1, Key key2, std::ostream* os);
+
+  /// Connect variable and factor.
+  static void ConnectVariableFactor(Key key, size_t i, std::ostream* os);
+
+  /// Draw a single factor, specified by its index i and its variable keys.
+  void ProcessFactor(size_t i, const KeyVector& keys,
+                     const boost::optional<Vector2>& position,
+                     std::ostream* os) const;
+};
+
+}  // namespace gtsam

gtsam/inference/FactorGraph-inst.h

@@ -26,6 +26,7 @@
 #include <stdio.h>
 #include <algorithm>
 #include <iostream>  // for cout :-(
+#include <fstream>
 #include <sstream>
 #include <string>
 
@@ -125,4 +126,48 @@ FactorIndices FactorGraph<FACTOR>::add_factors(const CONTAINER& factors,
   return newFactorIndices;
 }
 
+/* ************************************************************************* */
+template <class FACTOR>
+void FactorGraph<FACTOR>::dot(std::ostream& os, const DotWriter& writer,
+                              const KeyFormatter& keyFormatter) const {
+  writer.writePreamble(&os);
+
+  // Create nodes for each variable in the graph
+  for (Key key : keys()) {
+    writer.DrawVariable(key, keyFormatter, boost::none, &os);
+  }
+  os << "\n";
+
+  // Create factors and variable connections
+  for (size_t i = 0; i < size(); ++i) {
+    const auto& factor = at(i);
+    if (factor) {
+      const KeyVector& factorKeys = factor->keys();
+      writer.ProcessFactor(i, factorKeys, boost::none, &os);
+    }
+  }
+
+  os << "}\n";
+  std::flush(os);
+}
+
+/* ************************************************************************* */
+template <class FACTOR>
+std::string FactorGraph<FACTOR>::dot(const DotWriter& writer,
+                                     const KeyFormatter& keyFormatter) const {
+  std::stringstream ss;
+  dot(ss, writer, keyFormatter);
+  return ss.str();
+}
+
+/* ************************************************************************* */
+template <class FACTOR>
+void FactorGraph<FACTOR>::saveGraph(const std::string& filename,
+                                    const DotWriter& writer,
+                                    const KeyFormatter& keyFormatter) const {
+  std::ofstream of(filename.c_str());
+  dot(of, writer, keyFormatter);
+  of.close();
+}
+
 }  // namespace gtsam

gtsam/inference/FactorGraph.h

@@ -22,9 +22,10 @@
 
 #pragma once
 
+#include <gtsam/inference/DotWriter.h>
+#include <gtsam/inference/Key.h>
 #include <gtsam/base/FastVector.h>
 #include <gtsam/base/Testable.h>
-#include <gtsam/inference/Key.h>
 
 #include <Eigen/Core>  // for Eigen::aligned_allocator
 
@@ -36,6 +37,7 @@
 #include <string>
 #include <type_traits>
 #include <utility>
+#include <iosfwd>
 
 namespace gtsam {
 /// Define collection type:
@@ -371,6 +373,23 @@ class FactorGraph {
     return factors_.erase(first, last);
   }
 
+  /// @}
+  /// @name Graph Display
+  /// @{
+
+  /// Output to graphviz format, stream version.
+  void dot(std::ostream& os, const DotWriter& writer = DotWriter(),
+           const KeyFormatter& keyFormatter = DefaultKeyFormatter) const;
+
+  /// Output to graphviz format string.
+  std::string dot(const DotWriter& writer = DotWriter(),
+                  const KeyFormatter& keyFormatter = DefaultKeyFormatter) const;
+
+  /// output to file with graphviz format.
+  void saveGraph(const std::string& filename,
+                 const DotWriter& writer = DotWriter(),
+                 const KeyFormatter& keyFormatter = DefaultKeyFormatter) const;
+
   /// @}
   /// @name Advanced Interface
   /// @{