diff --git a/gtsam/discrete/discrete.i b/gtsam/discrete/discrete.i index e2310f4344..1bee25379f 100644 --- a/gtsam/discrete/discrete.i +++ b/gtsam/discrete/discrete.i @@ -222,6 +222,12 @@ class DotWriter { DotWriter(double figureWidthInches = 5, double figureHeightInches = 5, bool plotFactorPoints = true, bool connectKeysToFactor = true, bool binaryEdges = true); + + double figureWidthInches; + double figureHeightInches; + bool plotFactorPoints; + bool connectKeysToFactor; + bool binaryEdges; }; #include diff --git a/gtsam/nonlinear/nonlinear.i b/gtsam/nonlinear/nonlinear.i index 84c4939f49..b6ab086c45 100644 --- a/gtsam/nonlinear/nonlinear.i +++ b/gtsam/nonlinear/nonlinear.i @@ -133,6 +133,18 @@ class Ordering { void serialize() const; }; +#include +class GraphvizFormatting : gtsam::DotWriter { + GraphvizFormatting(); + + enum Axis { X, Y, Z, NEGX, NEGY, NEGZ }; + Axis paperHorizontalAxis; + Axis paperVerticalAxis; + + double scale; + bool mergeSimilarFactors; +}; + #include class NonlinearFactorGraph { NonlinearFactorGraph(); @@ -195,10 +207,13 @@ class NonlinearFactorGraph { string dot( const gtsam::Values& values, - const gtsam::KeyFormatter& keyFormatter = gtsam::DefaultKeyFormatter); + const gtsam::KeyFormatter& keyFormatter = gtsam::DefaultKeyFormatter, + const GraphvizFormatting& writer = GraphvizFormatting()); void saveGraph(const string& s, const gtsam::Values& values, const gtsam::KeyFormatter& keyFormatter = - gtsam::DefaultKeyFormatter) const; + gtsam::DefaultKeyFormatter, + const GraphvizFormatting& writer = + GraphvizFormatting()) const; }; #include diff --git a/python/gtsam/tests/test_GraphvizFormatting.py b/python/gtsam/tests/test_GraphvizFormatting.py new file mode 100644 index 0000000000..bd675b988f --- /dev/null +++ b/python/gtsam/tests/test_GraphvizFormatting.py @@ -0,0 +1,139 @@ +""" +GTSAM Copyright 2010-2021, Georgia Tech Research Corporation, +Atlanta, Georgia 30332-0415 +All Rights Reserved + +See LICENSE for the license information + +Unit tests for Graphviz formatting of NonlinearFactorGraph. +Author: Frank Dellaert & senselessDev +""" + +# pylint: disable=no-member, invalid-name + +import unittest +import textwrap + +import numpy as np + +import gtsam +from gtsam.utils.test_case import GtsamTestCase + + +class TestGraphvizFormatting(GtsamTestCase): + """Tests for saving NonlinearFactorGraph to GraphViz format.""" + + def setUp(self): + self.graph = gtsam.NonlinearFactorGraph() + + odometry = gtsam.Pose2(2.0, 0.0, 0.0) + odometryNoise = gtsam.noiseModel.Diagonal.Sigmas( + np.array([0.2, 0.2, 0.1])) + self.graph.add(gtsam.BetweenFactorPose2(0, 1, odometry, odometryNoise)) + self.graph.add(gtsam.BetweenFactorPose2(1, 2, odometry, odometryNoise)) + + self.values = gtsam.Values() + self.values.insert_pose2(0, gtsam.Pose2(0., 0., 0.)) + self.values.insert_pose2(1, gtsam.Pose2(2., 0., 0.)) + self.values.insert_pose2(2, gtsam.Pose2(4., 0., 0.)) + + def test_default(self): + """Test with default GraphvizFormatting""" + expected_result = """\ + graph { + size="5,5"; + + var0[label="0", pos="0,0!"]; + var1[label="1", pos="0,2!"]; + var2[label="2", pos="0,4!"]; + + factor0[label="", shape=point]; + var0--factor0; + var1--factor0; + factor1[label="", shape=point]; + var1--factor1; + var2--factor1; + } + """ + + self.assertEqual(self.graph.dot(self.values), + textwrap.dedent(expected_result)) + + def test_swapped_axes(self): + """Test with user-defined GraphvizFormatting swapping x and y""" + expected_result = """\ + graph { + size="5,5"; + + var0[label="0", pos="0,0!"]; + var1[label="1", pos="2,0!"]; + var2[label="2", pos="4,0!"]; + + factor0[label="", shape=point]; + var0--factor0; + var1--factor0; + factor1[label="", shape=point]; + var1--factor1; + var2--factor1; + } + """ + + graphviz_formatting = gtsam.GraphvizFormatting() + graphviz_formatting.paperHorizontalAxis = gtsam.GraphvizFormatting.Axis.X + graphviz_formatting.paperVerticalAxis = gtsam.GraphvizFormatting.Axis.Y + self.assertEqual(self.graph.dot(self.values, + writer=graphviz_formatting), + textwrap.dedent(expected_result)) + + def test_factor_points(self): + """Test with user-defined GraphvizFormatting without factor points""" + expected_result = """\ + graph { + size="5,5"; + + var0[label="0", pos="0,0!"]; + var1[label="1", pos="0,2!"]; + var2[label="2", pos="0,4!"]; + + var0--var1; + var1--var2; + } + """ + + graphviz_formatting = gtsam.GraphvizFormatting() + graphviz_formatting.plotFactorPoints = False + + self.assertEqual(self.graph.dot(self.values, + writer=graphviz_formatting), + textwrap.dedent(expected_result)) + + def test_width_height(self): + """Test with user-defined GraphvizFormatting for width and height""" + expected_result = """\ + graph { + size="20,10"; + + var0[label="0", pos="0,0!"]; + var1[label="1", pos="0,2!"]; + var2[label="2", pos="0,4!"]; + + factor0[label="", shape=point]; + var0--factor0; + var1--factor0; + factor1[label="", shape=point]; + var1--factor1; + var2--factor1; + } + """ + + graphviz_formatting = gtsam.GraphvizFormatting() + graphviz_formatting.figureWidthInches = 20 + graphviz_formatting.figureHeightInches = 10 + + self.assertEqual(self.graph.dot(self.values, + writer=graphviz_formatting), + textwrap.dedent(expected_result)) + + +if __name__ == "__main__": + unittest.main()