Consider a new plotting artist for improving colormap legends (#894)

* Consider a new plotting artist for improving colormap legends

* test: add a unit test for verifying legend height

* plotting: manage 'cax' argument without 'ax' one

* test: add a missing 'abs(.)' function to cax height test

* doc: complete 'mapping.rst' with details about the choropleth legends

* Update mapping.rst

Make explanation of colorbar example more concise

doc/source/mapping.rst

@@ -52,6 +52,32 @@ Choropleth Maps
     world.plot(column='gdp_per_cap');
 
 
+Creating a legend
+~~~~~~~~~~~~~~~~~
+
+When plotting a map, one can enable a legend using the ``legend`` argument:
+
+.. ipython:: python
+
+    # Plot population estimates with an accurate legend
+    import matplotlib.pyplot as plt
+    fig, ax = plt.subplots(1, 1)
+    @savefig world_pop_est.png
+    world.plot(column='pop_est', ax=ax, legend=True)
+
+However, the default appearance of the legend and plot axes may not be desirable. One can define the plot axes (with ``ax``) and the legend axes (with ``cax``) and then pass those in to the ``plot`` call. The following example uses ``mpl_toolkits`` to vertically align the plot axes and the legend axes:
+
+.. ipython:: python
+
+    # Plot population estimates with an accurate legend
+    from mpl_toolkits.axes_grid1 import make_axes_locatable
+    fig, ax = plt.subplots(1, 1)
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="5%", pad=0.1)
+    @savefig world_pop_est_fixed_legend_height.png
+    world.plot(column='pop_est', ax=ax, legend=True, cax=cax)
+
+
 Choosing colors
 ~~~~~~~~~~~~~~~~
 

geopandas/plotting.py

@@ -314,7 +314,7 @@ def plot_series(s, cmap=None, color=None, ax=None, figsize=None, **style_kwds):
     return ax
 
 
-def plot_dataframe(df, column=None, cmap=None, color=None, ax=None,
+def plot_dataframe(df, column=None, cmap=None, color=None, ax=None, cax=None,
                    categorical=False, legend=False, scheme=None, k=5,
                    vmin=None, vmax=None, markersize=None, figsize=None,
                    legend_kwds=None, classification_kwds=None, **style_kwds):
@@ -342,6 +342,8 @@ def plot_dataframe(df, column=None, cmap=None, color=None, ax=None,
         If specified, all objects will be colored uniformly.
     ax : matplotlib.pyplot.Artist (default None)
         axes on which to draw the plot
+    cax : matplotlib.pyplot Artist (default None)
+        axes on which to draw the legend in case of color map.
     categorical : bool (default False)
         If False, cmap will reflect numerical values of the
         column being plotted.  For non-numerical columns, this
@@ -406,6 +408,8 @@ def plot_dataframe(df, column=None, cmap=None, color=None, ax=None,
     import matplotlib.pyplot as plt
 
     if ax is None:
+        if cax is not None:
+            raise ValueError("'ax' can not be None if 'cax' is not.")
         fig, ax = plt.subplots(figsize=figsize)
     ax.set_aspect('equal')
 
@@ -487,6 +491,11 @@ def plot_dataframe(df, column=None, cmap=None, color=None, ax=None,
         plot_linestring_collection(ax, lines, values[line_idx],
                                    vmin=mn, vmax=mx, cmap=cmap, **style_kwds)
 
+    if cax is not None:
+        cbar_kwargs = {"cax": cax}
+    else:
+        cbar_kwargs = {"ax": ax}
+
     # plot all Points in the same collection
     points = df.geometry[point_idx]
     if not points.empty:
@@ -518,7 +527,7 @@ def plot_dataframe(df, column=None, cmap=None, color=None, ax=None,
             ax.legend(patches, categories, **legend_kwds)
         else:
             n_cmap.set_array([])
-            ax.get_figure().colorbar(n_cmap, ax=ax)
+            ax.get_figure().colorbar(n_cmap, **cbar_kwargs)
 
     plt.draw()
     return ax

geopandas/tests/test_plotting.py

@@ -437,6 +437,45 @@ class TestMapclassifyPlotting:
             self.df.plot(column='gdp_md_est', scheme=scheme, k=3,
                          cmap='OrRd', legend=True)
 
+    def test_cax_legend_passing(self):
+        """Pass a 'cax' argument to 'df.plot(.)', that is valid only if 'ax' is
+        passed as well (if not, a new figure is created ad hoc, and 'cax' is
+        ignored)
+        """
+        ax = plt.axes()
+        from mpl_toolkits.axes_grid1 import make_axes_locatable
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes('right', size='5%', pad=0.1)
+        with pytest.raises(ValueError):
+            ax = self.df.plot(
+                column='pop_est', cmap='OrRd', legend=True, cax=cax
+            )
+
+    def test_cax_legend_height(self):
+        """Pass a cax argument to 'df.plot(.)', the legend location must be
+        aligned with those of main plot
+        """
+        # base case
+        with warnings.catch_warnings(record=True) as _:  # don't print warning
+            ax = self.df.plot(
+                column='pop_est', cmap='OrRd', legend=True
+            )
+        plot_height = ax.get_figure().get_axes()[0].get_position().height
+        legend_height = ax.get_figure().get_axes()[1].get_position().height
+        assert abs(plot_height - legend_height) >= 1e-6
+        # fix heights with cax argument
+        ax2 = plt.axes()
+        from mpl_toolkits.axes_grid1 import make_axes_locatable
+        divider = make_axes_locatable(ax2)
+        cax = divider.append_axes('right', size='5%', pad=0.1)
+        with warnings.catch_warnings(record=True) as _:
+            ax2 = self.df.plot(
+                column='pop_est', cmap='OrRd', legend=True, cax=cax, ax=ax2
+            )
+        plot_height = ax2.get_figure().get_axes()[0].get_position().height
+        legend_height = ax2.get_figure().get_axes()[1].get_position().height
+        assert abs(plot_height - legend_height) < 1e-6
+
 
 class TestPlotCollections: