plot: split contour and contourf.

Author: Benjamin Chrétien

src/roboptim/core/visualization/plotter.py

@@ -6,10 +6,10 @@
 import numpy as np
 
 class PlotStyle2D(object):
-    Contour, PColorMesh = range(2)
+    Contour, Contourf, PColorMesh = range(3)
 
 class PlotStyle3D(object):
-    Contour, Wireframe, Triangle = range(3)
+    Contour, Contourf, Wireframe, Triangle = range(4)
 
 class Plotter(object):
     """
@@ -77,20 +77,24 @@ def plot(self, f, plot_style=PlotStyle2D.PColorMesh, cmap=None, *args, **kwargs)
 
         if cmap is not None:
             if self.isstr(cmap):
-                cmap = cm.get_cmap(cmap)
+                kwargs['cmap'] = cm.get_cmap(cmap)
+            else:
+                kwargs['cmap'] = cmap
 
         Z = self.compute_z(f,X,Y)
 
         # Contour plotting
         if plot_style == PlotStyle2D.Contour:
-            self.ax.contourf(X, Y, Z, 1, cmap=cmap, *args, **kwargs)
-            if 'levels' in kwargs:
-                self.ax.contour(X, Y, Z, max(self.x_res, self.y_res),
-                                *args, **kwargs)
+            cs = self.ax.contour(X, Y, Z, max(self.x_res, self.y_res),
+                                 *args, **kwargs)
+            if 'label' in kwargs and kwargs['label'] is True:
+                plt.clabel(cs, *args, **kwargs)
+        # Filled contour plotting
+        elif plot_style == PlotStyle2D.Contourf:
+            self.ax.contourf(X, Y, Z, 1, *args, **kwargs)
         # Color mesh plotting
         elif plot_style == PlotStyle2D.PColorMesh:
-            mesh = self.ax.pcolormesh(X, Y, Z, cmap=cmap,
-                                      *args, **kwargs)
+            mesh = self.ax.pcolormesh(X, Y, Z, *args, **kwargs)
             plt.colorbar(mesh)
         else:
             return
@@ -117,7 +121,9 @@ def plot(self, f, plot_style=PlotStyle3D.Triangle, cmap=None,
 
         if cmap is not None:
             if self.isstr(cmap):
-                cmap = cm.get_cmap(cmap)
+                kwargs['cmap'] = cm.get_cmap(cmap)
+            else:
+                kwargs['cmap'] = cmap
 
         Z = self.compute_z(f,X,Y)
 
@@ -127,12 +133,13 @@ def plot(self, f, plot_style=PlotStyle3D.Triangle, cmap=None,
         # Triangle plotting
         elif plot_style == PlotStyle3D.Triangle:
             self.ax.plot_trisurf(X.flatten(), Y.flatten(), Z.flatten(),
-                                 cmap=cmap, *args, **kwargs)
+                                 *args, **kwargs)
+        # Filled contour plotting
+        elif plot_style == PlotStyle3D.Contourf:
+            self.ax.contourf(X, Y, Z, 1, *args, **kwargs)
         # Contour plotting
         elif plot_style == PlotStyle3D.Contour:
-            self.ax.contourf(X, Y, Z, 1, cmap=cmap, *args, **kwargs)
-            if 'levels' in kwargs:
-                self.ax.contour(X, Y, Z, max(self.x_res, self.y_res),
-                                *args, **kwargs)
+            self.ax.contour(X, Y, Z, max(self.x_res, self.y_res),
+                            *args, **kwargs)
         else:
             return

tests/plot.py

@@ -26,7 +26,8 @@ def test_2d_plot(self):
         plotter = Plotter2D([-10,10],[-10,10])
         plotter.x_res = 10
         plotter.y_res = 10
-        plotter.plot(f, plot_style=PlotStyle2D.Contour)
+        plotter.plot(f, plot_style=PlotStyle2D.Contour, label=True)
+        plotter.plot(f, plot_style=PlotStyle2D.Contourf)
         plotter.plot(f, plot_style=PlotStyle2D.PColorMesh)
         plotter.add_marker([0], [0], color="black", marker="x",
                            markersize=10, markeredgewidth=2)
@@ -38,6 +39,7 @@ def test_3d_plot(self):
         plotter.x_res = 10
         plotter.y_res = 10
         plotter.plot(f, plot_style=PlotStyle3D.Contour)
+        plotter.plot(f, plot_style=PlotStyle3D.Contourf)
         plotter.plot(f, plot_style=PlotStyle3D.Wireframe)
         plotter.plot(f, plot_style=PlotStyle3D.Triangle)
         plotter.add_marker([0], [0], color="black", marker="x",