From c0f6cd247b61cb6dc2c8971e041ab3a7374b61eb Mon Sep 17 00:00:00 2001 From: Frank Dellaert Date: Thu, 27 Jan 2022 12:53:27 -0500 Subject: [PATCH] allow for boxes! --- gtsam/inference/BayesNet-inst.h | 2 +- gtsam/inference/DotWriter.cpp | 7 +++++-- gtsam/inference/DotWriter.h | 12 ++++++++---- gtsam/inference/FactorGraph-inst.h | 2 +- gtsam/inference/inference.i | 1 + gtsam/nonlinear/NonlinearFactorGraph.cpp | 2 +- gtsam/symbolic/tests/testSymbolicBayesNet.cpp | 7 +++++-- 7 files changed, 22 insertions(+), 11 deletions(-) diff --git a/gtsam/inference/BayesNet-inst.h b/gtsam/inference/BayesNet-inst.h index bd90f4e4b..c201475c5 100644 --- a/gtsam/inference/BayesNet-inst.h +++ b/gtsam/inference/BayesNet-inst.h @@ -44,7 +44,7 @@ void BayesNet::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"; diff --git a/gtsam/inference/DotWriter.cpp b/gtsam/inference/DotWriter.cpp index 9220bd168..a6a33bc74 100644 --- a/gtsam/inference/DotWriter.cpp +++ b/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& 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"; } diff --git a/gtsam/inference/DotWriter.h b/gtsam/inference/DotWriter.h index a00cf4b07..13683e338 100644 --- a/gtsam/inference/DotWriter.h +++ b/gtsam/inference/DotWriter.h @@ -24,6 +24,7 @@ #include #include +#include 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 variablePositions; + std::map variablePositions; /** * The position hints allow one to use symbol character and index to specify @@ -52,6 +53,9 @@ struct GTSAM_EXPORT DotWriter { */ std::map positionHints; + /** A set of keys that will be displayed as a box */ + std::set 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, - const boost::optional& position, - std::ostream* os); + void drawVariable(Key key, const KeyFormatter& keyFormatter, + const boost::optional& position, + std::ostream* os) const; /// Create factor dot. static void DrawFactor(size_t i, const boost::optional& position, diff --git a/gtsam/inference/FactorGraph-inst.h b/gtsam/inference/FactorGraph-inst.h index 2034fdcb6..3d85be49e 100644 --- a/gtsam/inference/FactorGraph-inst.h +++ b/gtsam/inference/FactorGraph-inst.h @@ -135,7 +135,7 @@ void FactorGraph::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"; diff --git a/gtsam/inference/inference.i b/gtsam/inference/inference.i index 30f51ea23..9dd5b9812 100644 --- a/gtsam/inference/inference.i +++ b/gtsam/inference/inference.i @@ -130,6 +130,7 @@ class DotWriter { std::map variablePositions; std::map positionHints; + std::set boxes; }; #include diff --git a/gtsam/nonlinear/NonlinearFactorGraph.cpp b/gtsam/nonlinear/NonlinearFactorGraph.cpp index c03caed75..dfa54f26f 100644 --- a/gtsam/nonlinear/NonlinearFactorGraph.cpp +++ b/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"; diff --git a/gtsam/symbolic/tests/testSymbolicBayesNet.cpp b/gtsam/symbolic/tests/testSymbolicBayesNet.cpp index f9cd07d23..2e13be10e 100644 --- a/gtsam/symbolic/tests/testSymbolicBayesNet.cpp +++ b/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" - " vara2[label=\"a2\", pos=\"2,2!\"];\n" + " vara1[label=\"a1\", pos=\"1,2!\", shape=box];\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"