Correct placement of colorbar relative to geo axes (cartopy)

Question

Using Cartopy, I would like to have full control of where my colorbar goes. Usually I do this by getting the current axes position as basis and then create new axes for the colorbar. This works well for standard matplotlib axes but not when using Cartopy and geo_axes, because this will distort the axes.

So, my question is: how do I get the exact position of my geo_axes?

Here is a code example based on the Cartopy docs http://scitools.org.uk/cartopy/docs/latest/matplotlib/advanced_plotting.html:

import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import os
from netCDF4 import Dataset as netcdf_dataset
from cartopy import config

def main():
    fname = os.path.join(config["repo_data_dir"],
                     'netcdf', 'HadISST1_SST_update.nc'
                     )

    dataset = netcdf_dataset(fname)

    sst = dataset.variables['sst'][0, :, :]
    lats = dataset.variables['lat'][:]
    lons = dataset.variables['lon'][:]

    #my preferred way of creating plots (even if it is only one plot)
    ef, ax = plt.subplots(1,1,figsize=(10,5),subplot_kw={'projection': ccrs.PlateCarree()})
    ef.subplots_adjust(hspace=0,wspace=0,top=0.925,left=0.1)

    #get size and extent of axes:
    axpos = ax.get_position()
    pos_x = axpos.x0+axpos.width + 0.01# + 0.25*axpos.width
    pos_y = axpos.y0
    cax_width = 0.04
    cax_height = axpos.height
    #create new axes where the colorbar should go.
    #it should be next to the original axes and have the same height!
    pos_cax = ef.add_axes([pos_x,pos_y,cax_width,cax_height])

    im = ax.contourf(lons, lats, sst, 60, transform=ccrs.PlateCarree())

    ax.coastlines()

    plt.colorbar(im, cax=pos_cax)

    ax.coastlines(resolution='110m')
    ax.gridlines()
    ax.set_extent([-20, 60, 33, 63])

    #when using this line the positioning of the colorbar is correct,
    #but the image gets distorted.
    #when omitting this line, the positioning of the colorbar is wrong,
    #but the image is well represented (not distorted).
    ax.set_aspect('auto', adjustable=None)

    plt.savefig('sst_aspect.png')
    plt.close()



if __name__ == '__main__': main()

Resulting Figure, when using "set_aspect":

Resulting Figure, when omitting "set_aspect":

Basically, I'd like to obtain the first figure (correctly placed colorbar) but without using the "set_aspect". I guess this should be possible with some transformations, but I didn't find a solution so far.

Thanks!

Resilience

Answer 1

Great question! Thanks for the code, and pictures, it makes the problem a lot easier to understand as well as making it easier to quickly iterate on possible solutions.

The problem here is essentially a matplotlib one. Cartopy calls ax.set_aspect('equal') as this is part of the the Cartesian units of a projection's definition.

Matplotlib's equal aspect ratio functionality resizes the axes to match the x and y limits, rather than changing the limits to fit to the axes rectangle. It is for this reason that the axes does not fill the space allocated to it on the figure. If you interactively resize the figure you will see that the amount of space that the axes occupies varies depending on the aspect that you resize your figure to.

The simplest way of identifying the location of an axes is with the ax.get_position() method you have already been using. However, as we now know, this "position" changes with the size of the figure. One solution therefore is to re-calculate the position of the colorbar each time the figure is resized.

The matplotlib event machinery has a "resize_event" which is triggered each time a figure is resized. If we use this machinery for your colorbar, our event might look something like:

def resize_colobar(event):
    # Tell matplotlib to re-draw everything, so that we can get
    # the correct location from get_position.
    plt.draw()

    posn = ax.get_position()
    colorbar_ax.set_position([posn.x0 + posn.width + 0.01, posn.y0,
                             0.04, axpos.height])

fig.canvas.mpl_connect('resize_event', resize_colobar)

So if we relate this back to cartopy, and your original question, it is now possible to resize the colorbar based on the position of the geo-axes. The full code to do this might look like:

import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import os
from netCDF4 import Dataset as netcdf_dataset
from cartopy import config


fname = os.path.join(config["repo_data_dir"],
                 'netcdf', 'HadISST1_SST_update.nc'
                 )
dataset = netcdf_dataset(fname)
sst = dataset.variables['sst'][0, :, :]
lats = dataset.variables['lat'][:]
lons = dataset.variables['lon'][:]

fig, ax = plt.subplots(1, 1, figsize=(10,5),
                       subplot_kw={'projection': ccrs.PlateCarree()})

# Add the colorbar axes anywhere in the figure. Its position will be
# re-calculated at each figure resize. 
cbar_ax = fig.add_axes([0, 0, 0.1, 0.1])

fig.subplots_adjust(hspace=0, wspace=0, top=0.925, left=0.1)

sst_contour = ax.contourf(lons, lats, sst, 60, transform=ccrs.PlateCarree())


def resize_colobar(event):
    plt.draw()

    posn = ax.get_position()
    cbar_ax.set_position([posn.x0 + posn.width + 0.01, posn.y0,
                          0.04, posn.height])

fig.canvas.mpl_connect('resize_event', resize_colobar)

ax.coastlines()

plt.colorbar(sst_contour, cax=cbar_ax)


ax.gridlines()
ax.set_extent([-20, 60, 33, 63])

plt.show()

Bugloss

Great answer! Would never have found this solution on my own! – Resilience

Maybe a short side-question if I may: Is it possible to put the definition of resize_colorbar somewhere else and pass ax and cbar_ax as arguments (so that this function becomes reusable)? – Resilience

Update: while plt.show() delivers the correct result, plt.savefig('sst.png') does not (colorbar stays at the original position at [0,0,0.1,0.1]) . Tried to change the backend to matplotlib.use('Agg'), but this does not help. Any idea how I could get it to work with savefig? – Resilience

Question1: How about a function which returns the "resize_colorbar" function. ax and cbar_ax can be the input to that function... – Bugloss

Question2: Untested, but try just calling resize_colorbar(None) before the save. I suspect there is no resize_event being triggered when saving. – Bugloss

No luck regarding Q1, but more importantly, solution for Q2 works!! Thanks! – Resilience

@pelson: There is a typo in the upper code example axpos.height istead of posn.height. It can not be edited because it is less than 6 characters... – Lucilelucilia

Here is the solution for Q1. Shouldn't this all be built into cartopy? Such that ax.colorbar() on a GeoAxes makes a cbar_ax with the right extensions and takes care of the updates? – Berga

Answer 2

Bearing in mind that mpl_toolkits.axes_grid1 is not be the best-tested part of matplotlib, we can use it's functionality to achieve what you want.

We can use the Example given in the mpl_toolkits documentation, but the axes_class needs to be set explicitly, it has to be set as axes_class=plt.Axes, else it attempts to create a GeoAxes as colorbar

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

def sample_data_3d(shape):
    """Returns `lons`, `lats`, and fake `data`

    adapted from:
    http://scitools.org.uk/cartopy/docs/v0.15/examples/axes_grid_basic.html
    """


    nlons, nlats = shape
    lats = np.linspace(-np.pi / 2, np.pi / 2, nlats)
    lons = np.linspace(0, 2 * np.pi, nlons)
    lons, lats = np.meshgrid(lons, lats)
    wave = 0.75 * (np.sin(2 * lats) ** 8) * np.cos(4 * lons)
    mean = 0.5 * np.cos(2 * lats) * ((np.sin(2 * lats)) ** 2 + 2)

    lats = np.rad2deg(lats)
    lons = np.rad2deg(lons)
    data = wave + mean

    return lons, lats, data


# get data
lons, lats, data = sample_data_3d((180, 90))


# set up the plot
proj = ccrs.PlateCarree()

f, ax = plt.subplots(1, 1, subplot_kw=dict(projection=proj))
h = ax.pcolormesh(lons, lats, data, transform=proj, cmap='RdBu')
ax.coastlines()

# following https://matplotlib.org/2.0.2/mpl_toolkits/axes_grid/users/overview.html#colorbar-whose-height-or-width-in-sync-with-the-master-axes
# we need to set axes_class=plt.Axes, else it attempts to create
# a GeoAxes as colorbar

divider = make_axes_locatable(ax)
ax_cb = divider.new_horizontal(size="5%", pad=0.1, axes_class=plt.Axes)


f.add_axes(ax_cb)
plt.colorbar(h, cax=ax_cb)

Also note the cartopy example that uses AxesGrid from mpl_toolkits.axes_grid1.

Lucilelucilia

The colorbar in this solutions is correctly saved in the png file. – Rambling

Answer 3

Building on the immensely helpful answer by pelson above, if you don't care about resizing with plt.show but having properly sized and placed colorbars when using plt.savefig, I just developed this code snippet, which hopefully will be useful to others:

import wrf
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import cartopy.crs as ccrs

fig = plt.figure()
ax = plt.subplot(cart_proj) # cart_proj is from wrf.get_cartopy
ax.set_xlim(cart_xlim) # cart_xlim is from wrf.cartopy_xlim
ax.set_ylim(cart_ylim) # cart_ylim is from wrf.cartopy_ylim
ax.coastlines()
data_crs = ccrs.PlateCarree()

## Draw the contour plot (assume lons, lats, var are defined previously)
extend = 'max'
cmap = mpl.rainbow
bounds = np.arange(min_val, max_val, int_val) # preset these values to something meaningful
norm = mpl.colors.BoundaryNorm(bounds, cmap.N, extend=extend)
plt.contourf(wrf.to_np(lons), wrf.to_np(lats), wrf.to_np(var), bounds, cmap=cmap, norm=norm, extend=extend, transform=data_crs, transform_first=(ax,True))

## Create colorbar axes (temporarily) anywhere
cax = fig.add_axes([0,0,0.1,0.1])

## Find the location of the main plot axes
posn = ax.get_position()

## Where do you want the colorbar? 'bottom' or 'right' as examples
cbar_loc = 'bottom'
cbar_lab = 'Variable Name [units]'

## Adjust the positioning and orientation of the colorbar, and draw it
if cbar_loc == 'bottom':
   cax.set_position([posn.x0, posn.y0-0.09, posn.width, 0.05])
   plt.colorbar(cax=cax, orientation='horizontal', label=cbar_lab)
elif cbar_loc == 'right':
   cax.set_position([posn.x0+posn.width+0.05, posn.y0, 0.04, posn.height])
   plt.colorbar(cax=cax, orientation='vertical', label=cbar_lab)

## Save the figure (fname is set to whatever file path for the figure)
plt.savefig(fname)

Not only will the colorbar be correctly positioned and exactly match the width or height of the GeoAxes frame (depending on where you want to position it), but the colorbar will also inherit the correct extend properties in this way from plt.contourf, if both it and norm are defined with the extend attribute. I have tested this using WRF domains of various sizes and aspect ratios, and it seems to be robust.

Quietly