diff --git a/CHANGELOG.rst b/CHANGELOG.rst index cd7024db..a1092e0d 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -9,6 +9,8 @@ The format is based on `Keep a Changelog ` Unreleased ---------- +- Modified `ep.colorbar()` to work with GeoDataFrame plots (@nkorinek, #557) + 0.9.2 ----- diff --git a/earthpy/plot.py b/earthpy/plot.py index c594f4b4..2963dfa1 100644 --- a/earthpy/plot.py +++ b/earthpy/plot.py @@ -16,7 +16,7 @@ import earthpy.spatial as es -def colorbar(mapobj, size="3%", pad=0.09): +def colorbar(mapobj, size="3%", pad=0.09, return_cax=False): """Adjust colorbar height to match the matplotlib axis height. NOTE: This function requires matplotlib v 3.0.1 or greater or v 2.9 or @@ -31,6 +31,9 @@ def colorbar(mapobj, size="3%", pad=0.09): The percent width of the colorbar relative to the plot. pad : int (default = 0.09) The space between the plot and the color bar. + return_cax : bool (default = False) + Boolean to return the cax created before it is applied to the axis. + Can be used to help adjust colorbars on non-array plots. Returns ------- @@ -70,7 +73,10 @@ def colorbar(mapobj, size="3%", pad=0.09): fig = ax.figure divider = make_axes_locatable(ax) cax = divider.append_axes("right", size=size, pad=pad) - return fig.colorbar(mapobj, cax=cax) + if return_cax: + return cax + else: + return fig.colorbar(mapobj, cax=cax) def _plot_image( diff --git a/earthpy/tests/test_plot.py b/earthpy/tests/test_plot.py index c4af5002..0dc05bcc 100644 --- a/earthpy/tests/test_plot.py +++ b/earthpy/tests/test_plot.py @@ -21,3 +21,12 @@ def test_colorbar_raises_value_error(): with pytest.raises(AttributeError, match="requires a matplotlib"): ep.colorbar(list()) plt.close() + + +def test_colorbar_returns_cax(basic_geometry_gdf): + """Test that colorbar works with geodataframes when cax is returned.""" + f, ax = plt.subplots(figsize=(5, 5)) + cb_ax = ep.colorbar(ax, return_cax=True) + im = basic_geometry_gdf.plot(ax=ax, legend=True, cax=cb_ax) + assert ax.get_position().height == im.axes.get_position().height + plt.close(f)