Add plotting functionality (#11)

* Initial work towards plotting

* Add default colors for plotting

* Implement Graph.plot() method

* Plotting works

* Add matplotlib to tests requirements

* Add matplotlib to Travis-CI and readthedocs requirements

* Attempt to fix failing test on Travis-CI

* Add parameters to Graph.plot()

.travis.yml

@@ -3,7 +3,7 @@ python:
   - "3.7"
 install:
   - pip install .
-  - pip install flake8 pylint coveralls
+  - pip install flake8 pylint coveralls matplotlib
 script:
 - flake8 graphslam/ && pylint graphslam/ && flake8 tests/ && pylint tests/ && coverage run --source graphslam setup.py test && coverage report -m
 after_success:

docs/requirements.txt

@@ -1,4 +1,5 @@
 sphinx==2.1.2
 graphslam
+matplotlib
 numpy
 scipy

graphslam/edge/base_edge.py

@@ -160,3 +160,14 @@ def to_g2o(self):
 
         """
         raise NotImplementedError
+
+    def plot(self, color=''):
+        """Plot the edge.
+
+        Parameters
+        ----------
+        color : str
+            The color that will be used to plot the edge
+
+        """
+        raise NotImplementedError

graphslam/edge/edge_odometry.py

@@ -7,9 +7,16 @@
 
 import numpy as np
 
+try:
+    import matplotlib.pyplot as plt
+except ImportError:  # pragma: no cover
+    plt = None
+
 from .base_edge import BaseEdge
 
+from ..pose.r2 import PoseR2
 from ..pose.se2 import PoseSE2
+from ..pose.r3 import PoseR3
 from ..pose.se3 import PoseSE3
 
 
@@ -85,3 +92,26 @@ def to_g2o(self):
             return "EDGE_SE3:QUAT {} {} {} {} {} {} {} {} {} ".format(self.vertex_ids[0], self.vertex_ids[1], self.estimate[0], self.estimate[1], self.estimate[2], self.estimate[3], self.estimate[4], self.estimate[5], self.estimate[6]) + " ".join([str(x) for x in self.information[np.triu_indices(6, 0)]]) + "\n"
 
         raise NotImplementedError
+
+    def plot(self, color='b'):
+        """Plot the edge.
+
+        Parameters
+        ----------
+        color : str
+            The color that will be used to plot the edge
+
+        """
+        if plt is None:  # pragma: no cover
+            raise NotImplementedError
+
+        if isinstance(self.vertices[0].pose, (PoseR2, PoseSE2)):
+            xy = np.array([v.pose.position for v in self.vertices])
+            plt.plot(xy[:, 0], xy[:, 1], color=color)
+
+        elif isinstance(self.vertices[0].pose, (PoseR3, PoseSE3)):
+            xyz = np.array([v.pose.position for v in self.vertices])
+            plt.plot(xyz[:, 0], xyz[:, 1], xyz[:, 2], color=color)
+
+        else:
+            raise NotImplementedError

graphslam/graph.py

@@ -83,6 +83,12 @@
 from scipy.sparse import SparseEfficiencyWarning, lil_matrix
 from scipy.sparse.linalg import spsolve
 
+try:
+    import matplotlib.pyplot as plt
+    from mpl_toolkits.mplot3d import Axes3D  # noqa pylint: disable=unused-import
+except ImportError:  # pragma: no cover
+    plt = None
+
 
 warnings.simplefilter("ignore", SparseEfficiencyWarning)
 warnings.filterwarnings("ignore", category=SparseEfficiencyWarning)
@@ -292,3 +298,33 @@ def to_g2o(self, outfile):
 
             for e in self._edges:
                 f.write(e.to_g2o())
+
+    def plot(self, vertex_color='r', vertex_marker='o', vertex_markersize=3, edge_color='b'):
+        """Plot the graph.
+
+        Parameters
+        ----------
+        vertex_color : str
+            The color that will be used to plot the vertices
+        vertex_marker : str
+            The marker that will be used to plot the vertices
+        vertex_markersize : int
+            The size of the plotted vertices
+        edge_color : str
+            The color that will be used to plot the edges
+
+        """
+        if plt is None:  # pragma: no cover
+            raise NotImplementedError
+
+        fig = plt.figure()
+        if len(self._vertices[0].pose.position) == 3:
+            fig.add_subplot(111, projection='3d')
+
+        for e in self._edges:
+            e.plot(edge_color)
+
+        for v in self._vertices:
+            v.plot(vertex_color, vertex_marker, vertex_markersize)
+
+        plt.show()

graphslam/vertex.py

@@ -4,6 +4,13 @@
 
 """
 
+try:
+    import matplotlib.pyplot as plt
+except ImportError:  # pragma: no cover
+    plt = None
+
+from .pose.r2 import PoseR2
+from .pose.r3 import PoseR3
 from .pose.se2 import PoseSE2
 from .pose.se3 import PoseSE3
 
@@ -52,3 +59,30 @@ def to_g2o(self):
             return "VERTEX_SE3:QUAT {} {} {} {} {} {} {} {}\n".format(self.id, self.pose[0], self.pose[1], self.pose[2], self.pose[3], self.pose[4], self.pose[5], self.pose[6])
 
         raise NotImplementedError
+
+    def plot(self, color='r', marker='o', markersize=3):
+        """Plot the vertex.
+
+        Parameters
+        ----------
+        color : str
+            The color that will be used to plot the vertex
+        marker : str
+            The marker that will be used to plot the vertex
+        markersize : int
+            The size of the plotted vertex
+
+        """
+        if plt is None:  # pragma: no cover
+            raise NotImplementedError
+
+        if isinstance(self.pose, (PoseR2, PoseSE2)):
+            x, y = self.pose.position
+            plt.plot(x, y, color=color, marker=marker, markersize=markersize)
+
+        elif isinstance(self.pose, (PoseR3, PoseSE3)):
+            x, y, z = self.pose.position
+            plt.plot([x], [y], [z], markerfacecolor=color, markeredgecolor=color, marker=marker, markersize=markersize)
+
+        else:
+            raise NotImplementedError

setup.py

@@ -22,7 +22,7 @@
     author_email='jefflirion@users.noreply.github.com',
     packages=['graphslam', 'graphslam.pose', 'graphslam.edge'],
     install_requires=['numpy', 'scipy'],
-    # tests_require=[],
+    tests_require=['matplotlib'],
     classifiers=['License :: OSI Approved :: MIT License',
                  'Operating System :: OS Independent',
                  'Programming Language :: Python :: 3'],

tests/test_base_edge.py

@@ -116,6 +116,17 @@ def test_to_g2o(self):
         with self.assertRaises(NotImplementedError):
             _ = e.to_g2o()
 
+    def test_plot(self):
+        """Test that the ``plot`` method is not implemented.
+
+        """
+        p = PoseSE2([0, 0], 0)
+        v = Vertex(0, p)
+        e = BaseEdge(0, 1, 0, [v])
+
+        with self.assertRaises(NotImplementedError):
+            e.plot()
+
 
 if __name__ == '__main__':
     unittest.main()

tests/test_edge_odometry.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2020 Jeff Irion and contributors
+
+"""Unit tests for the graph.py module.
+
+"""
+
+
+import unittest
+
+from graphslam.vertex import Vertex
+from graphslam.edge.edge_odometry import EdgeOdometry
+from graphslam.pose.r2 import PoseR2
+from graphslam.pose.r3 import PoseR3
+from graphslam.pose.se2 import PoseSE2
+from graphslam.pose.se3 import PoseSE3
+
+
+class TestEdgeOdometry(unittest.TestCase):
+    """Tests for the ``EdgeOdometry`` class.
+
+    """
+
+    def test_plot(self):
+        """Test that the ``plot`` method is not implemented.
+
+        """
+        v_none = Vertex(0, None)
+        v_r2 = Vertex(1, PoseR2([1, 2]))
+        v_se2 = Vertex(2, PoseSE2([1, 2], 3))
+        v_r3 = Vertex(3, PoseR3([1, 2, 3]))
+        v_se3 = Vertex(4, PoseSE3([1, 2, 3], [0.5, 0.5, 0.5, 0.5]))
+
+        with self.assertRaises(NotImplementedError):
+            e = EdgeOdometry(0, 1, 0, [v_none, v_none])
+            e.plot()
+
+        for v in [v_r2, v_se2, v_r3, v_se3]:
+            e = EdgeOdometry(0, 1, 0, [v, v])
+            e.plot()
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_graph.py

@@ -89,6 +89,14 @@ def test_to_g2o(self):
                 with self.assertRaises(NotImplementedError):
                     self.g.to_g2o("test.g2o")
 
+    def test_plot(self):
+        """Test that the ``plot`` method does not raise an exception.
+
+        """
+        # avoid showing the plots
+        with patch("graphslam.graph.plt.show"):
+            self.g.plot()
+
 
 class TestGraphR3(TestGraphR2):
     r"""Tests for the ``Graph`` class with :math:`\mathbb{R}^3` poses.

tests/test_vertex.py

@@ -8,9 +8,14 @@
 import unittest
 
 import numpy as np
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D  # noqa pylint: disable=unused-import
 
 from graphslam.vertex import Vertex
+from graphslam.pose.r2 import PoseR2
+from graphslam.pose.r3 import PoseR3
 from graphslam.pose.se2 import PoseSE2
+from graphslam.pose.se3 import PoseSE3
 
 
 class TestVertex(unittest.TestCase):
@@ -27,6 +32,25 @@ def test_constructor(self):
         self.assertEqual(v.id, 1)
         self.assertAlmostEqual(np.linalg.norm(v.pose.to_array() - np.array([1., 2., 3.])), 0.)
 
+    def test_plot(self):
+        """Test that a ``Vertex`` can be plotted.
+
+        """
+        v_none = Vertex(0, None)
+        v_r2 = Vertex(1, PoseR2([1, 2]))
+        v_se2 = Vertex(2, PoseSE2([1, 2], 3))
+        v_r3 = Vertex(3, PoseR3([1, 2, 3]))
+        v_se3 = Vertex(4, PoseSE3([1, 2, 3], [0.5, 0.5, 0.5, 0.5]))
+
+        with self.assertRaises(NotImplementedError):
+            v_none.plot()
+
+        for v in [v_r2, v_se2, v_r3, v_se3]:
+            fig = plt.figure()
+            if len(v.pose.position) == 3:
+                fig.add_subplot(111, projection='3d')
+            v.plot()
+
 
 if __name__ == '__main__':
     unittest.main()