allow for boxes!

gtsam/inference/BayesNet-inst.h

@@ -44,7 +44,7 @@ void BayesNet<CONDITIONAL>::dot(std::ostream& os,
   // Create nodes for each variable in the graph
   for (Key key : this->keys()) {
     auto position = writer.variablePos(key);
-    writer.DrawVariable(key, keyFormatter, position, &os);
+    writer.drawVariable(key, keyFormatter, position, &os);
   }
   os << "\n";
 

gtsam/inference/DotWriter.cpp

@@ -39,15 +39,18 @@ void DotWriter::digraphPreamble(ostream* os) const {
       << "\";\n\n";
 }
 
-void DotWriter::DrawVariable(Key key, const KeyFormatter& keyFormatter,
+void DotWriter::drawVariable(Key key, const KeyFormatter& keyFormatter,
                              const boost::optional<Vector2>& position,
-                             ostream* os) {
+                             ostream* os) const {
   // Label the node with the label from the KeyFormatter
   *os << "  var" << keyFormatter(key) << "[label=\"" << keyFormatter(key)
       << "\"";
   if (position) {
     *os << ", pos=\"" << position->x() << "," << position->y() << "!\"";
   }
+  if (boxes.count(key)) {
+    *os << ", shape=box";
+  }
   *os << "];\n";
 }
 

gtsam/inference/DotWriter.h

@@ -24,6 +24,7 @@
 
 #include <iosfwd>
 #include <map>
+#include <set>
 
 namespace gtsam {
 
@@ -43,7 +44,7 @@ struct GTSAM_EXPORT DotWriter {
    * Variable positions can be optionally specified and will be included in the
    * dor file with a "!' sign, so "neato" can use it to render them.
    */
-  std::map<gtsam::Key, Vector2> variablePositions;
+  std::map<Key, Vector2> variablePositions;
 
   /**
    * The position hints allow one to use symbol character and index to specify
@@ -52,6 +53,9 @@ struct GTSAM_EXPORT DotWriter {
    */
   std::map<char, double> positionHints;
 
+  /** A set of keys that will be displayed as a box */
+  std::set<Key> boxes;
+
   explicit DotWriter(double figureWidthInches = 5,
                      double figureHeightInches = 5,
                      bool plotFactorPoints = true,
@@ -69,9 +73,9 @@ struct GTSAM_EXPORT DotWriter {
   void digraphPreamble(std::ostream* os) const;
 
   /// Create a variable dot fragment.
-  static void DrawVariable(Key key, const KeyFormatter& keyFormatter,
+  void drawVariable(Key key, const KeyFormatter& keyFormatter,
                     const boost::optional<Vector2>& position,
-                           std::ostream* os);
+                    std::ostream* os) const;
 
   /// Create factor dot.
   static void DrawFactor(size_t i, const boost::optional<Vector2>& position,

gtsam/inference/FactorGraph-inst.h

@@ -135,7 +135,7 @@ void FactorGraph<FACTOR>::dot(std::ostream& os,
 
   // Create nodes for each variable in the graph
   for (Key key : keys()) {
-    writer.DrawVariable(key, keyFormatter, boost::none, &os);
+    writer.drawVariable(key, keyFormatter, boost::none, &os);
   }
   os << "\n";
 

gtsam/inference/inference.i

@@ -130,6 +130,7 @@ class DotWriter {
 
   std::map<gtsam::Key, gtsam::Vector2> variablePositions;
   std::map<char, double> positionHints;
+  std::set<Key> boxes;
 };
 
 #include <gtsam/inference/VariableIndex.h>

gtsam/nonlinear/NonlinearFactorGraph.cpp

@@ -111,7 +111,7 @@ void NonlinearFactorGraph::dot(std::ostream& os, const Values& values,
   // Create nodes for each variable in the graph
   for (Key key : keys) {
     auto position = writer.variablePos(values, min, key);
-    writer.DrawVariable(key, keyFormatter, position, &os);
+    writer.drawVariable(key, keyFormatter, position, &os);
   }
   os << "\n";
 

gtsam/symbolic/tests/testSymbolicBayesNet.cpp

@@ -91,18 +91,21 @@ TEST(SymbolicBayesNet, Dot) {
   DotWriter writer;
   writer.positionHints.emplace('a', 2);
   writer.positionHints.emplace('x', 1);
+  writer.boxes.emplace(A(1));
+  writer.boxes.emplace(A(2));
   
   auto position = writer.variablePos(A(1));
   CHECK(position);
   EXPECT(assert_equal(Vector2(1, 2), *position, 1e-5));
 
   string actual = bn.dot(DefaultKeyFormatter, writer);
+  bn.saveGraph("bn.dot", DefaultKeyFormatter, writer);
   EXPECT(actual ==
          "digraph {\n"
          "  size=\"5,5\";\n"
          "\n"
-         "  vara1[label=\"a1\", pos=\"1,2!\"];\n"
+         "  vara1[label=\"a1\", pos=\"1,2!\", shape=box];\n"
-         "  vara2[label=\"a2\", pos=\"2,2!\"];\n"
+         "  vara2[label=\"a2\", pos=\"2,2!\", shape=box];\n"
          "  varx1[label=\"x1\", pos=\"1,1!\"];\n"
          "  varx2[label=\"x2\", pos=\"2,1!\"];\n"
          "  varx3[label=\"x3\", pos=\"3,1!\"];\n"