Merge pull request #186 from jorisvandenbossche/plot-colorbar

ENH: add plot points with colormap

geopandas/plotting.py

@@ -49,9 +49,10 @@ def plot_multilinestring(ax, geom, color='red', linewidth=1.0):
             plot_linestring(ax, line, color=color, linewidth=linewidth)
 
 
-def plot_point(ax, pt, marker='o', markersize=2):
+def plot_point(ax, pt, marker='o', markersize=2, color="black"):
     """ Plot a single Point geometry """
-    ax.plot(pt.x, pt.y, marker=marker, markersize=markersize, linewidth=0)
+    ax.plot(pt.x, pt.y, marker=marker, markersize=markersize, linewidth=0,
+            color=color)
 
 
 def gencolor(N, colormap='Set1'):
@@ -126,7 +127,7 @@ def plot_series(s, colormap='Set1', axes=None, linewidth=1.0, figsize=None, **co
         elif geom.type == 'LineString' or geom.type == 'MultiLineString':
             plot_multilinestring(ax, geom, color=next(color), linewidth=linewidth)
         elif geom.type == 'Point':
-            plot_point(ax, geom)
+            plot_point(ax, geom, color=next(color))
     plt.draw()
     return ax
 
@@ -244,9 +245,8 @@ def plot_dataframe(s, column=None, colormap=None, linewidth=1.0,
                 plot_multipolygon(ax, geom, facecolor=cmap.to_rgba(value), linewidth=linewidth, **color_kwds)
             elif geom.type == 'LineString' or geom.type == 'MultiLineString':
                 plot_multilinestring(ax, geom, color=cmap.to_rgba(value), linewidth=linewidth)
-            # TODO: color point geometries
             elif geom.type == 'Point':
-                plot_point(ax, geom)
+                plot_point(ax, geom, color=cmap.to_rgba(value))
         if legend:
             if categorical:
                 patches = []
@@ -337,6 +337,7 @@ def norm_cmap(values, cmap, normalize, cm, vmin=None, vmax=None):
         -------
         n_cmap
             mapping of normalized values to colormap (cmap)
+
     """
 
     mn = vmin or min(values)

tests/test_plotting.py

@@ -1,4 +1,4 @@
-from __future__ import absolute_import
+from __future__ import absolute_import, division
 
 import numpy as np
 import os
@@ -7,7 +7,7 @@ import tempfile
 
 import matplotlib
 matplotlib.use('Agg', warn=False)
-from matplotlib.pyplot import Artist, savefig, clf, cm
+from matplotlib.pyplot import Artist, savefig, clf, cm, get_cmap
 from matplotlib.testing.noseclasses import ImageComparisonFailure
 from matplotlib.testing.compare import compare_images
 from numpy import cos, sin, pi
@@ -107,6 +107,62 @@ class PlotTests(unittest.TestCase):
         self._compare_images(ax=ax, filename=filename)
 
 
+class TestPointPlotting(unittest.TestCase):
+
+    def setUp(self):
+
+        self.N = 10
+        self.points = GeoSeries(Point(i, i) for i in range(self.N))
+        values = np.arange(self.N)
+        self.df = GeoDataFrame({'geometry': self.points, 'values': values})
+
+    def test_default_colors(self):
+
+        ## without specifying values -> max 9 different colors
+
+        # GeoSeries
+        ax = self.points.plot()
+        cmap = get_cmap('Set1', 9)
+        expected_colors = cmap(list(range(9))*2)
+        _check_colors(ax.get_lines(), expected_colors)
+
+        # GeoDataFrame -> uses 'jet' instead of 'Set1'
+        ax = self.df.plot()
+        cmap = get_cmap('jet', 9)
+        expected_colors = cmap(list(range(9))*2)
+        _check_colors(ax.get_lines(), expected_colors)
+
+        ## with specifying values
+
+        ax = self.df.plot(column='values')
+        cmap = get_cmap('jet')
+        expected_colors = cmap(np.arange(self.N)/(self.N-1))
+
+        _check_colors(ax.get_lines(), expected_colors)
+
+    def test_colormap(self):
+
+        ## without specifying values -> max 9 different colors
+
+        # GeoSeries
+        ax = self.points.plot(colormap='RdYlGn')
+        cmap = get_cmap('RdYlGn', 9)
+        expected_colors = cmap(list(range(9))*2)
+        _check_colors(ax.get_lines(), expected_colors)
+
+        # GeoDataFrame -> same as GeoSeries in this case
+        ax = self.df.plot(colormap='RdYlGn')
+        _check_colors(ax.get_lines(), expected_colors)
+
+        ## with specifying values
+
+        ax = self.df.plot(column='values', colormap='RdYlGn')
+        cmap = get_cmap('RdYlGn')
+        expected_colors = cmap(np.arange(self.N)/(self.N-1))
+
+        _check_colors(ax.get_lines(), expected_colors)
+
+
 class TestPySALPlotting(unittest.TestCase):
 
     @classmethod
@@ -128,5 +184,15 @@ class TestPySALPlotting(unittest.TestCase):
         self.assertEqual(labels, expected)
 
 
+def _check_colors(collection, expected_colors):
+
+    import matplotlib.colors as colors
+    conv = colors.colorConverter
+
+    for patch, color in zip(collection, expected_colors):
+        result = patch.get_color()
+        assert conv.to_rgba(result) == conv.to_rgba(color)
+
+
 if __name__ == '__main__':
     unittest.main()