Search code examples
pythonmatplotlibpython-xarraycolorbar

How to adjust Matplotlib colorbar range in xarray plot?

I have a plot that looks like this

Pseudocolor maps

I cannot understand how to manually change or set the range of data values for the colorbar. I would like to experiment with ranges based on the data values shown in the plots and change the colorbar to (-4,4). I see that plt.clim, vmin and vmax are functions to possibly use.

Here is my code:

import cdsapi
import xarray as xr
import matplotlib.pyplot as plt
import numpy as np
import cartopy.crs as ccrs
# Also requires cfgrib library.

c = cdsapi.Client()

url = c.retrieve(
    'reanalysis-era5-single-levels-monthly-means',
    {
        'product_type': 'monthly_averaged_reanalysis',
        'format': 'grib',
        'variable': ['100m_u_component_of_wind','100m_v_component_of_wind'],
        'year': ['2006','2007','2008','2009','2010','2011','2012','2013','2014','2015','2016','2017','2018','2019','2020','2021'],
        'month': ['01','02','03','04','05','06','07','08','09','10','11','12'],
        'time': '00:00',
        'grid': [0.25, 0.25],
        'area': [70.00, -180.00, -40.00, 180.00],
    },
    "C:\\Users\\U321103\\.spyder-py3\\ERA5_MAPPING\\100m_wind_U_V.grib")
path = "C:\\Users\\U321103\\.spyder-py3\\ERA5_MAPPING\\100m_wind_U_V.grib"
ds = xr.load_dataset(path, engine='cfgrib')

wind_abs = np.sqrt(ds.u100**2 + ds.v100**2)
monthly_means = wind_abs.mean(dim='time')
wind_abs_clim = wind_abs.sel(time=slice('2006-01','2020-12')).groupby('time.month').mean(dim='time') # select averaging period

wind_abs_anom = ((wind_abs.groupby('time.month') / wind_abs_clim))-1 #deviation from climo

fg = wind_abs_anom.sel(time=slice('2021-01',None)).groupby('time.month').mean(dim='time').plot(col='month',
                        col_wrap=3,transform=ccrs.PlateCarree(),
                        cbar_kwargs={'orientation':'horizontal','shrink':0.6, 'aspect':40,'label':'Percent Deviation'},robust=False,subplot_kws={'projection': ccrs.Mercator()})

fg.map(lambda: plt.gca().coastlines())
Solution

  • I was able to reproduce your figure and found that I could add vmin and vmax as shown below. For some reason that meant I also had to specify the colormap, otherwise I ended up with viridis. But the code below works for me (with a bit of refactoring as I got it working — the only material change here is in the plotting section at the bottom).

    First, loading the data:

    import cdsapi

c = cdsapi.Client()
params = {
    'product_type': 'monthly_averaged_reanalysis',
    'format': 'grib',
    'variable': ['100m_u_component_of_wind', '100m_v_component_of_wind'],
    'year': [f'{n}' for n in range(2006, 2022)],
    'month': [f'{n:02d}' for n in range(1, 13)],
    'time': '00:00',
    'grid': [0.25, 0.25],
    'area': [70.00, -180.00, -40.00, 180.00],
}
path = '100m_wind_U_V.grib'
url = c.retrieve('reanalysis-era5-single-levels-monthly-means',
                 params,
                 path,
                )

    Then there's the data pipeline:

    import xarray as xr
import numpy as np
# Also need cfgrib library.

ds = xr.load_dataset(path, engine='cfgrib')
wind_abs = np.sqrt(ds.u100**2 + ds.v100**2)
monthly_means = wind_abs.mean(dim='time')
wind_abs_clim = (wind_abs.sel(time=slice('2006-01','2020-12'))
                         .groupby('time.month')
                         .mean(dim='time'))
wind_abs_anom = ((wind_abs.groupby('time.month') / wind_abs_clim)) - 1

    Finally the plotting:

    import cartopy.crs as ccrs
import matplotlib.pyplot as plt

cbar_kwargs = {'orientation':'horizontal', 'shrink':0.6, 'aspect':40, 'label':'Percent Deviation'}
subplot_kws = {'projection': ccrs.Mercator()}
fg = (wind_abs_anom.sel(time=slice('2021-01', None))
                   .groupby('time.month')
                   .mean(dim='time')
                   .plot(col='month',
                         col_wrap=3,
                         transform=ccrs.PlateCarree(),
                         cmap='RdBu_r', vmin=-3, vmax=3,  # <-- New bit.
                         cbar_kwargs=cbar_kwargs,
                         robust=False,
                         subplot_kws=subplot_kws
                        ))
fg.map(lambda: plt.gca().coastlines())

    Sometimes I'll use a percentile to control the values for vmin and vmax automatically, like max_ = np.percentile(data, 99), then vmin=-max_, vmax=max_. This deals nicely with outliers that stretch the colormap, but it requires you to be able to calculate those values before making the plot.

    If you want to start having more control over the plot, it might be a good idea to stop using the xarray plotting interface and use matplotlib and cartopy directly. Here's what that might look like (replacing all of the plotting code above):

    import cartopy.crs as ccrs
import matplotlib.pyplot as plt

sel = wind_abs_anom.sel(time=slice('2021-01', None))

left, *_, right = wind_abs_anom.longitude
top, *_, bottom = wind_abs_anom.latitude  # Min and max latitude.
extent = [left, right, bottom, top]

fig, axs = plt.subplots(nrows=2, ncols=3,
                        figsize=(15, 6),
                        subplot_kw={'projection': ccrs.PlateCarree()},
                       )

for ax, (month, group) in zip(axs.flat, sel.groupby('time.month')):
    mean = group.mean(dim='time')
    im = ax.imshow(mean,
                   transform=ccrs.PlateCarree(),
                   extent=extent,
                   cmap='RdBu_r', vmin=-3, vmax=3)
    ax.set_title(f'month = {month}')
    ax.coastlines()

cbar_ax = fig.add_axes([0.2, 0.0, 0.6, 0.05])  # Left, bottom, width, height.
cbar = fig.colorbar(im, cax=cbar_ax, extend='both', orientation='horizontal')
cbar.set_label('Percent deviation')
    
plt.show()

    For some reason, when I try to use ccra.Mercator() for the map, the data gets distorted; maybe you can figure that bit out.