enable set_map_layout for AxesGrid (#116)

* enable set_map_layout for AxesGrid

* changelog

* Apply suggestions from code review

Author: mathause

CHANGELOG.md

@@ -10,7 +10,9 @@
   | ----------- | ---- | ------- |
   | matplotlib* | 3.6  | 3.7     |
 
+### Enhancements
 
+- Enable passing `AxesGrid` (from `mpl_toolkits.axes_grid1`) to `set_map_layout` ([#116](https://github.com/mathause/mplotutils/pull/116)).
 
 ## v0.5.0 (27.03.2024)
 

README.md

@@ -11,6 +11,18 @@ Without mplotutils         |  With mplotutils
 
 The code to create the examples can be found in [docs/example.py](docs/example.py).
 
+
+**axes_grid**
+
+Matplotlib's [axes_grid](https://matplotlib.org/stable/users/explain/toolkits/axes_grid.html) can also display data with a fixed aspect ratio. However, the size of the figure will not be correct. mplotutils (from version 0.6) can also help with this
+
+| Axes grid - without mplotutils     | Axes grid - with mplotutils |
+| :--------------------------------: | :-------------------------: |
+| <img src="docs/example_axes_grid_no_mpu.png" alt="Without mplotutils" width="250"/> | <img src="docs/example_axes_grid_mpu.png" alt="With mplotutils" width="250"/> |
+
+The code to create the example can be found in [docs/example_axes_grid.py](docs/example_axes_grid.py).
+
+
 ## Installation
 
 See [docs/installation.md](docs/installation.md).

docs/example_axes_grid.py

@@ -0,0 +1,74 @@
+import cartopy.crs as ccrs
+import matplotlib.pyplot as plt
+from cartopy.mpl.geoaxes import GeoAxes
+from mpl_toolkits.axes_grid1 import AxesGrid
+
+import mplotutils as mpu
+
+
+def plot_map_axes_grid():
+
+    # create example data
+    da = mpu.sample_dataarray(30, 30)
+
+    axes_class = (GeoAxes, {"projection": ccrs.Robinson()})
+
+    f = plt.figure()
+
+    axgr = AxesGrid(
+        f,
+        111,
+        axes_class=axes_class,
+        nrows_ncols=(1, 2),
+        axes_pad=0.5 / 2.54,
+        cbar_mode="single",
+        cbar_location="bottom",
+        cbar_size="10%",
+        # see https://github.com/matplotlib/matplotlib/issues/28343
+        cbar_pad=0.0 / 2.54,
+    )
+
+    for ax in axgr.axes_all:
+        mappable = da.plot(ax=ax, transform=ccrs.PlateCarree(), add_colorbar=False)
+        ax.coastlines()
+
+    cbax = axgr.cbar_axes[0]
+
+    cbax.colorbar(mappable)
+
+    f.subplots_adjust(left=0.05, right=0.95)  # , bottom=0, top=1)
+
+    return f, axgr
+
+
+def plot_map_no_mpu():
+    """plot 2 x 2 global maps _not_ using mplotutils"""
+
+    f, __ = plot_map_axes_grid()
+
+    print(f.get_size_inches())  # * 2.54)
+
+    f.suptitle("AxesGrid - without mplotutils")
+
+
+def plot_map_mpu():
+    """plot 2 x 2 global maps using mplotutils"""
+
+    f, axgr = plot_map_axes_grid()
+
+    # ensure the figure has the correct size
+    mpu.set_map_layout(axgr, width=6.4 * 2.54)
+
+    print(f.get_size_inches())  # * 2.54)
+
+    f.suptitle("AxesGrid - with mplotutils")
+
+
+if __name__ == "__main__":
+    opt = {"dpi": 200, "facecolor": "0.9", "transparent": False}
+
+    plot_map_no_mpu()
+    plt.savefig("example_axes_grid_no_mpu.png", **opt)
+
+    plot_map_mpu()
+    plt.savefig("example_axes_grid_mpu.png", **opt)

docs/example_axes_grid_mpu.png

@@ -1,19 +1,23 @@
+import warnings
+
 import matplotlib.pyplot as plt
 import numpy as np
+from mpl_toolkits.axes_grid1 import AxesGrid
 
 from mplotutils._deprecate import _deprecate_positional_args
+from mplotutils.mpl import _get_renderer
 
 
 @_deprecate_positional_args("0.3")
-def set_map_layout(axes, width=17.0, *, nrow=None, ncol=None):
+def set_map_layout(obj=None, width=17.0, *, nrow=None, ncol=None, axes=None):
     """set figure height, given width, taking axes' aspect ratio into account
 
     Needs to be called after all plotting is done.
 
     Parameters
     ----------
-    axes : (Geo)Axes | iterable of (Geo)Axes
-        Array with all axes of the figure.
+    obj : (Geo)Axes | iterable of (Geo)Axes | AxesGrid
+        Array with all axes of the figure or mpl_toolkits.axes_grid1 AxesGrid.
     width : float, default: 17
         Width of the full figure in cm.
     nrow : integer, default: None
@@ -28,6 +32,26 @@ def set_map_layout(axes, width=17.0, *, nrow=None, ncol=None):
     Only works if all the axes have the same aspect ratio.
     """
 
+    if axes is not None and obj is not None:
+        raise TypeError("Cannot pass 'obj' and 'axes'")
+
+    if axes is not None:
+        warnings.warn("The 'axes' keyword has been renamed to 'obj'", FutureWarning)
+        obj = axes
+
+    if obj is None:
+        raise TypeError(
+            "set_map_layout() missing 1 required positional argument: 'obj'"
+        )
+
+    if isinstance(obj, AxesGrid):
+        _set_map_layout_axes_grid(obj, width, nrow, ncol)
+    else:
+        _set_map_layout_axes(obj, width, nrow, ncol)
+
+
+def _set_map_layout_axes(axes, width, nrow, ncol):
+
     if (nrow is None and ncol is not None) or (nrow is not None and ncol is None):
         raise ValueError("Must set none or both of 'nrow' and 'ncol'")
 
@@ -69,3 +93,51 @@ def set_map_layout(axes, width=17.0, *, nrow=None, ncol=None):
 
     f.set_figwidth(width / 2.54)
     f.set_figheight(height / 2.54)
+
+
+def _set_map_layout_axes_grid(axgr, width, nrow, ncol):
+
+    if nrow is not None or ncol is not None:
+        raise TypeError("Cannot pass 'nrow' or 'ncol' for and 'AxesGrid'")
+
+        # assumes the first of the axes is representative for all
+    ax = axgr.axes_all[0]
+
+    f = ax.get_figure()
+
+    # getting the correct data ratio of geoaxes requires draw
+    f.canvas.draw()
+
+    bottom = f.subplotpars.bottom
+    top = f.subplotpars.top
+    left = f.subplotpars.left
+    right = f.subplotpars.right
+
+    width_fraction = right - left
+    height_fraction = top - bottom
+
+    inner_width = width_fraction * width
+
+    renderer = _get_renderer(f)
+
+    # divider get_*_sizes contains the relative and absolute sizes of all plot elements
+    # (subplots, colorbars, pad & colorbar pad)
+
+    divider = axgr.get_divider()
+
+    vertical_sizes = divider.get_vertical_sizes(renderer)
+    vs_rel = vertical_sizes[:, 0].sum()
+    vs_abs = vertical_sizes[:, 1].sum() * 2.54
+
+    horizontal_sizes = divider.get_horizontal_sizes(renderer)
+    hs_rel = horizontal_sizes[:, 0].sum()
+    hs_abs = horizontal_sizes[:, 1].sum() * 2.54
+
+    inner_height = (inner_width - hs_abs) / hs_rel * vs_rel + vs_abs
+
+    if inner_height <= 0:
+        raise ValueError("Not enough space on figure")
+
+    height = inner_height / height_fraction
+
+    f.set_size_inches(width / 2.54, height / 2.54)

docs/example_axes_grid_no_mpu.png

@@ -50,3 +50,13 @@ def _set_backend(backend):
         matplotlib.use(backend)
     except ImportError:
         pytest.skip(backend)
+
+
+def get_rtol(f):
+    # macosx is only exact up to 1 / dpi
+
+    if plt.get_backend().lower() != "macosx":
+        rtol = 1e-07
+    else:
+        rtol = 1 / f.get_dpi()
+    return rtol

mplotutils/map_layout.py

@@ -4,17 +4,25 @@
 
 from mplotutils import set_map_layout
 
-from . import figure_context, subplots_context
+from . import figure_context, get_rtol, subplots_context
 
 
-def get_rtol(f):
-    # macosx is only exact up to 1 / dpi
+def test_set_map_layout_deprecated_kwarg():
 
-    if plt.get_backend().lower() != "macosx":
-        rtol = 1e-07
-    else:
-        rtol = 1 / f.get_dpi()
-    return rtol
+    with subplots_context() as (__, ax):
+        with pytest.warns(
+            FutureWarning, match="The 'axes' keyword has been renamed to 'obj'"
+        ):
+            set_map_layout(axes=ax)
+
+    with pytest.raises(
+        TypeError,
+        match=r"set_map_layout\(\) missing 1 required positional argument: 'obj'",
+    ):
+        set_map_layout()
+
+    with pytest.raises(TypeError, match="Cannot pass 'obj' and 'axes'"):
+        set_map_layout(object, axes=object)
 
 
 def test_set_map_layout_default_width():
@@ -195,10 +203,10 @@ def test_set_map_layout_two_axes_horz():
 
 def test_set_map_layout_nrow_ncol_only_one_raises():
     with pytest.raises(ValueError, match="Must set none or both of 'nrow' and 'ncol'"):
-        set_map_layout(None, width=17.0, nrow=1, ncol=None)
+        set_map_layout(object, width=17.0, nrow=1, ncol=None)
 
     with pytest.raises(ValueError, match="Must set none or both of 'nrow' and 'ncol'"):
-        set_map_layout(None, width=17.0, nrow=None, ncol=1)
+        set_map_layout(object, width=17.0, nrow=None, ncol=1)
 
 
 def test_set_map_layout_cartopy_2_2():