Python matplotlib example plot curvilinear data#
Software requirements:
Python 3
numpy
xarray
matplotlib
cartopy
Run the curvilinear grid example script:
python curvilinear_grid_wrapped_lines_corrected.py
Script curvilinear_grid_wrapped_lines_corrected.py:
#!/usr/bin/env python
# coding: utf-8
# DKRZ tutorial
#
# Curvilinear grids
#
# Curvilinear grids have 2-dimensional longitude lon(y,x) and latitude lat(y,x)
# coordinate arrays, where x and y are just the indexes.
# These kind of grids often makes it difficult to plot the data on its native
# grid without remapping to a regular lonlat grid. For example, the
# Matplotlib's contour and contourf functions wrap the contour lines around
# most projections which looks streaky.
#
# In this example we use the function z_masked_overlap() from htonchia at
# https://github.com/SciTools/cartopy/issues/1421 to fix the wrapping problem.
#
# The ncdump output of the input dataset:
#
# netcdf MPI-ESM-LR_tos_20060101 {
# dimensions:
# x = 256 ;
# y = 220 ;
# nv4 = 4 ;
# time = UNLIMITED ; // (1 currently)
# nb2 = 2 ;
# variables:
# float lon(y, x) ;
# lon:long_name = "longitude coordinate" ;
# lon:units = "degrees_east" ;
# lon:standard_name = "longitude" ;
# lon:_CoordinateAxisType = "Lon" ;
# lon:bounds = "lon_bnds" ;
# float lon_bnds(y, x, nv4) ;
# float lat(y, x) ;
# lat:long_name = "latitude coordinate" ;
# lat:units = "degrees_north" ;
# lat:standard_name = "latitude" ;
# lat:_CoordinateAxisType = "Lat" ;
# lat:bounds = "lat_bnds" ;
# float lat_bnds(y, x, nv4) ;
# double time(time) ;
# time:bounds = "time_bnds" ;
# time:units = "days since 1850-01-01 00:00:00" ;
# time:calendar = "proleptic_gregorian" ;
# double time_bnds(time, nb2) ;
# time_bnds:units = "days since 1850-01-01 00:00:00" ;
# time_bnds:calendar = "proleptic_gregorian" ;
# float tos(time, y, x) ;
# tos:long_name = "Sea Surface Temperature" ;
# tos:standard_name = "sea_surface_temperature" ;
# tos:units = "K" ;
# tos:coordinates = "lon lat" ;
# tos:_FillValue = 1.e+20f ;
#
#
# See also https://en.wikipedia.org/wiki/Curvilinear_coordinates
#
# 06.12.2022 copyright DKRZ, kmf
#
import os
import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
# Function z_masked_overlap() to prevent the line wrapping in contour plots
# The function z_masked_overlap was taken from
# https://github.com/SciTools/cartopy/issues/1421 from htonchia.
def z_masked_overlap(axe, X, Y, Z, source_projection=None):
"""
for data in projection axe.projection
find and mask the overlaps (more 1/2 the axe.projection range)
X, Y either the coordinates in axe.projection or longitudes latitudes
Z the data
operation one of 'pcorlor', 'pcolormesh', 'countour', 'countourf'
if source_projection is a geodetic CRS data is in geodetic coordinates
and should first be projected in axe.projection
X, Y are 2D same dimension as Z for contour and contourf
same dimension as Z or with an extra row and column for pcolor
and pcolormesh
return ptx, pty, Z
"""
if not hasattr(axe, 'projection'):
return X, Y, Z
if not isinstance(axe.projection, ccrs.Projection):
return X, Y, Z
if len(X.shape) != 2 or len(Y.shape) != 2:
return X, Y, Z
if (source_projection is not None and
isinstance(source_projection, ccrs.Geodetic)):
transformed_pts = axe.projection.transform_points(
source_projection, X, Y)
ptx, pty = transformed_pts[..., 0], transformed_pts[..., 1]
else:
ptx, pty = X, Y
with np.errstate(invalid='ignore'):
# diagonals have one less row and one less columns
diagonal0_lengths = np.hypot(
ptx[1:, 1:] - ptx[:-1, :-1],
pty[1:, 1:] - pty[:-1, :-1]
)
diagonal1_lengths = np.hypot(
ptx[1:, :-1] - ptx[:-1, 1:],
pty[1:, :-1] - pty[:-1, 1:]
)
to_mask = (
(diagonal0_lengths > (
abs(axe.projection.x_limits[1]
- axe.projection.x_limits[0])) / 2) |
np.isnan(diagonal0_lengths) |
(diagonal1_lengths > (
abs(axe.projection.x_limits[1]
- axe.projection.x_limits[0])) / 2) |
np.isnan(diagonal1_lengths)
)
# TODO check if we need to do something about surrounding vertices
# add one extra colum and row for contour and contourf
if (to_mask.shape[0] == Z.shape[0] - 1 and
to_mask.shape[1] == Z.shape[1] - 1):
to_mask_extended = np.zeros(Z.shape, dtype=bool)
to_mask_extended[:-1, :-1] = to_mask
to_mask_extended[-1, :] = to_mask_extended[-2, :]
to_mask_extended[:, -1] = to_mask_extended[:, -2]
to_mask = to_mask_extended
if np.any(to_mask):
Z_mask = getattr(Z, 'mask', None)
to_mask = to_mask if Z_mask is None else to_mask | Z_mask
Z = np.ma.masked_where(to_mask, Z)
return ptx, pty, Z
def main():
#-- Open curvilinear dataset
fname = os.environ['HOME']+'/data/MPIOM/MPI-ESM-LR_tos_20060101.nc'
ds = xr.open_dataset(fname)
var = ds.tos.isel(time=0)
lat = ds.lat
lon = ds.lon
#-- The coordinate arrays lat and lon are bounds, so the variable var should be
#-- the value inside those bounds. Therefore, the last value from the var array
#-- has to be removed.
var = var[:-1, :-1]
#-- Colormap
cmap = 'RdBu_r'
#-- Set vmin, vmax, and levels
vmin = 260.
vmax = 300.
levels = 20
#-- Create the figure
projection = ccrs.PlateCarree()
fig = plt.figure(figsize=(14, 8), constrained_layout=False)
#-- Define space for the four plots
gs = fig.add_gridspec(nrows=2, ncols=2, hspace=0, wspace=0.05)
#-- Generate the subplots
ax1 = plt.subplot(gs[0, 0], projection=projection)
ax2 = plt.subplot(gs[1, 0], projection=projection)
ax3 = plt.subplot(gs[0, 1], projection=projection)
ax4 = plt.subplot(gs[1, 1], projection=projection)
#-- Upper left plot
ax1.set_title('curvilinear grid - wrapped lines', fontsize=16)
ax1.set_global()
ax1.add_feature(cfeature.LAND, color='gray')
ax1.add_feature(cfeature.OCEAN, color='gray')
ax1.add_feature(cfeature.COASTLINE)
ax1.gridlines(draw_labels=True)
plot = ax1.contour(lon[:-1,:-1], lat[:-1,:-1], var,
cmap=cmap,
levels=levels,
vmin=vmin,
vmax=vmax,
linewidths=1.0)
plt.colorbar(plot, ax=ax1, shrink=0.6, pad=0.09)
#-- Compute the corrected data to prevent line wrapping. Use the dataset
#-- coordinates and variable.
X, Y, maskedZ = z_masked_overlap(ax1,
lon.data,
lat.data,
ds.tos.isel(time=0),
source_projection=ccrs.Geodetic())
#-- Lower left plot
ax2.set_global()
ax2.add_feature(cfeature.LAND, color='gray')
ax2.add_feature(cfeature.OCEAN, color='gray')
ax2.add_feature(cfeature.COASTLINE)
ax2.gridlines(draw_labels=True)
ax2.set_title('curvilinear grid - corrected', fontsize=16)
plot = ax2.contour(X, Y, maskedZ,
cmap=cmap,
levels=levels,
vmin=vmin,
vmax=vmax,
linewidths=1.0)
plt.colorbar(plot, ax=ax2, shrink=0.6, pad=0.09)
#-- Check if it works the same for contour fill.
#-- Display the incorrect and the corrected contour fill plot
#--
#-- Upper right plot
ax3.set_title('curvilinear grid - wrapped contours', fontsize=16)
ax3.set_global()
ax3.add_feature(cfeature.LAND, color='gray')
ax3.add_feature(cfeature.OCEAN, color='gray')
ax3.add_feature(cfeature.COASTLINE)
ax3.gridlines(draw_labels=True)
plot = ax3.contourf(lon[:-1,:-1], lat[:-1,:-1], var,
cmap=cmap,
levels=levels,
vmin=vmin,
vmax=vmax)
plt.colorbar(plot, ax=ax3, shrink=0.6, pad=0.09)
#-- Use the already computed X, Y, maskedZ from above.
#--
#-- lower right plot
ax4.set_global()
ax4.add_feature(cfeature.LAND, color='gray')
ax4.add_feature(cfeature.OCEAN, color='gray')
ax4.add_feature(cfeature.COASTLINE)
ax4.gridlines(draw_labels=True)
ax4.set_title('curvilinear grid - corrected', fontsize=16)
plot = ax4.contourf(X, Y, maskedZ,
cmap=cmap,
levels=levels,
vmin=vmin,
vmax=vmax)
plt.colorbar(plot, ax=ax4, shrink=0.6, pad=0.09)
plt.savefig('plot_curvilinear_contour_2x2_plots.png', bbox_inches='tight', dpi=100)
if __name__ == '__main__':
main()
Plot result:
