Python matplotlib example: unstructured ICON triangles plot with vectors (python 3)#

Software requirements:

  • Python 3

  • Numpy

  • matplotlib

  • cartopy

  • xarray

Script:

'''
DKRZ matplotlib script:  DKRZ_example_ICON_triangles_uas_vas_quiver.py

Description:    Python script using matplotlib
                - ICON model data
                - variable 'tas', 'uas', 'vas'
                - use cell vertices
                - add colorbar
                - use gadm36_DEU shapefile to draw state lines
                - draw location of Marburg

Data:           Cells:  1239

31.05.21 kmf
'''
import time, os
import xarray as xr

import numpy as np
import numpy.ma as ma

import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.tri as tri
from   matplotlib.collections import PolyCollection

import cartopy.crs as ccrs
import cartopy.feature as cfeature
from   cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
import cartopy.io.shapereader as shpreader


def main():

    t1 = time.time()    #-- retrieve start time
    print('')

    #-- set color map
    colormap = 'RdYlBu_r'

    #-- set title string
    title = 'ICON temperature data'

    #--  define path, file and variable name
    dir_name  = '/Users/k204045/data/ICON/SummerSchool/'
    file_name = 'icon_atmo.nc'

    fname = os.path.join(dir_name,file_name)

    varName = 'tas'
    uName   = 'uas'
    vName   = 'vas'

    #-- open dataset
    ds = xr.open_dataset(fname)

    #-- get missing_value
    try:
        missing_value = ds[varName].encoding['missing_value']
    except:
        print('Warning: no missing_value/_FillValue defined!')

    #-- get variable
    var = ds[varName][0,:].values
    var = var - 273.15

    u = ds[uName][0,:].values
    v = ds[vName][0,:].values

    vec_units = ds[uName].units

    #-- get coordinates and convert radians to degrees
    clon = np.rad2deg(ds.clon.values)
    clat = np.rad2deg(ds.clat.values)
    clon_vertices = np.rad2deg(ds.clon_bnds.values)
    clat_vertices = np.rad2deg(ds.clat_bnds.values)

    ncells, nv = clon_vertices.shape[0], clon_vertices.shape[1]

    #-- set contour levels, labels
    varMin, varMax, varInt = 0., 6., 0.5
    levels = np.arange(varMin, varMax+varInt, varInt)
    nlevs  = levels.size
    labels = ['{:.2f}'.format(x) for x in levels]

    #-- print information to stdout
    print('')
    print('Cells:            %6d ' % clon.size)
    print('Variable min/max: %6.2f ' % np.nanmin(var)+'/'+' %.2f' % np.nanmax(var))
    print('Contour  min/max: %6.2f ' % varMin+'/'+' %.2f' % varMax)
    print('')

    #-- set projection
    projection = ccrs.PlateCarree()

    #-- create figure and axes instances; we need subplots for plot and colorbar
    fig, ax = plt.subplots(figsize=(12,11), subplot_kw=dict(projection=projection))

    bg = [0.8, 0.8, 0.8, 1.0]

    ax.set_facecolor(bg)

    ax.set_extent([7.95, 10.05, 49.95, 52.05], crs=ccrs.PlateCarree())

    #-- plot land areas at last to get rid of the contour lines at land
    ax.coastlines(linewidth=0.5)

    gl = ax.gridlines(draw_labels=True, linewidth=0.5, color='dimgray',
                      alpha=0.4, zorder=3)
    gl.xlabel_style = {'size':7}
    gl.ylabel_style = {'size':7}
    gl.top_labels   = False
    gl.right_labels = False
    gl.xformatter = LONGITUDE_FORMATTER
    gl.yformatter = LATITUDE_FORMATTER

    #-- plot the title string
    plt.title(title, fontsize=16)

    #-- define color map
    cmap     = plt.get_cmap(colormap, nlevs)            #-- read the color map
    cmaplist = [i for i in range(cmap.N)]               #-- color bar indices
    ncol     = len(cmaplist)                            #-- number of colors
    colors   = np.ndarray([ncells,4], np.float32)       #-- assign color array for triangles

    #-- preset all colors elements to black
    colors[:,0] = 0.8
    colors[:,1] = 0.8
    colors[:,2] = 0.8
    colors[:,3] = 1.

    print('levels:      ',levels)
    print('nlevs:       %3d' %nlevs)
    print('ncol:        %3d' %ncol)
    print('')

    #-- set color index of all cells in between levels
    for m in range(0,ncol-1):
        vind = []
        for i in range(0,ncells-2, 1):
            if (var[i] >= levels[m] and var[i] < levels[m+1]):
                colors[i,:] = cmap(cmaplist[m])
                vind.append(i)

        print('set colors: finished level %3d' % m , ' -- %5d ' % len(vind) , ' polygons considered')
        del vind

    colors[np.where(var < varMin),:]  = cmap(cmaplist[0])
    colors[np.where(var >= varMax),:] = cmap(cmaplist[ncol-1])

    #-- create the triangles
    clon_vertices = np.where(clon_vertices < -180., clon_vertices + 360., clon_vertices)
    clon_vertices = np.where(clon_vertices >  180., clon_vertices - 360., clon_vertices)

    triangles = np.zeros((ncells, nv, 2), np.float32)

    for i in range(0, ncells, 1):
        triangles[i,:,0] = np.array(clon_vertices[i,:])
        triangles[i,:,1] = np.array(clat_vertices[i,:])

    print('')
    print('--> triangles done')

    #-- create polygon/triangle collection
    coll = PolyCollection(triangles, array=None,
                          fc=colors,
                          edgecolors='none',
                          linewidth=0.05,
                          transform=ccrs.Geodetic())
    ax.add_collection(coll)

    print('--> polygon collection done')

    #-- vectors
    vplot = ax.quiver(clon, clat, u, v,
                      width=0.003,
                      headwidth=4.,
                      headlength=5.,
                      units='xy',
                      color='black',
                      transform=ccrs.PlateCarree())

    vec_ref = ax.quiverkey(vplot, 0.92, 0.84, 2,
                           r'$2 \frac{m}{s}$',
                           labelpos='E',
                           labelsep=0.05,
                           coordinates='figure')

    #-- add States of Germany from shapefile
    shpname = '/Users/k204045/data/Shapefiles/gadm36_DEU_shp/gadm36_DEU_1.shp'
    adm1_shapes = list(shpreader.Reader(shpname).geometries())
    ax.add_geometries(adm1_shapes,
                      ccrs.PlateCarree(),
                      linewidth=1,
                      edgecolor='black',
                      facecolor='none',
                      alpha=0.5)

    #-- mark location of Marburg
    ax.plot(8.7667, 50.8167,
            marker='.',
            color='white',
            markersize=12,
            markeredgecolor='black',
            transform=ccrs.Geodetic())

    ax.text(8.8, 50.8167, 'Marburg',
            fontsize=16,
            fontweight='bold',
            horizontalalignment='center',
            verticalalignment='top',
            transform=ccrs.Geodetic())

    #-- add a color bar
    cax  = fig.add_axes([0.91, 0.165, 0.02, 0.65], autoscalex_on=True) #-- x,y,w,h

    bounds = [0., 0.5, 0.75, 1., 1.5, 2., 3., 4., 5., 6.]
    norm = mpl.colors.BoundaryNorm(bounds, cmap.N, extend='both')

    cbar = fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap),
                        cax=cax,
                        orientation='vertical',
                        ticks=levels,
                        boundaries=levels,
                        format='%3.1f')

    #plt.setp(cbar.ax.get_xticklabels()[::2], visible=False)
    cbar.set_label('[deg C]')

    #-- add text 'missing = black'
    ax.annotate('gray = missing',
                xy=(.83, .08),
                xycoords='figure fraction',
                horizontalalignment='right',
                verticalalignment='top',
                fontsize=6)

    #-- maximize and save the PNG file
    plt.savefig('plot_ICON_tas_triangles_uas_vas_quiver.png', bbox_inches='tight',dpi=100)

    #-- get wallclock time
    t2 = time.time()
    print('Wallclock time:  %0.3f seconds' % (t2-t1))
    print('')

if __name__ == '__main__':
    main()

Result:

Matplotlib unstructured ICON triangles with vectors py3 w400