From 1778190af481e2e90eb2fdd2ff0781af371d60b6 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Mon, 3 Feb 2020 08:31:04 -0700 Subject: [PATCH] Switch COMPASS from basemap to cartopy Fix cartopy issues with tricontourf --- .../interpolate_time_varying_forcing.py | 61 +++++++++++-------- .../ocean/jigsaw_to_MPAS/coastal_tools.py | 46 ++++++++------ 2 files changed, 63 insertions(+), 44 deletions(-) diff --git a/testing_and_setup/compass/ocean/hurricane/interpolate_time_varying_forcing.py b/testing_and_setup/compass/ocean/hurricane/interpolate_time_varying_forcing.py index ee1f0e7914..b3e0cef560 100644 --- a/testing_and_setup/compass/ocean/hurricane/interpolate_time_varying_forcing.py +++ b/testing_and_setup/compass/ocean/hurricane/interpolate_time_varying_forcing.py @@ -1,4 +1,4 @@ -import netCDF4 +import netCDF4 import matplotlib.pyplot as plt import numpy as np import glob @@ -8,16 +8,20 @@ import yaml import subprocess import argparse -from mpl_toolkits.basemap import Basemap +import cartopy +import cartopy.crs as ccrs +import cartopy.feature as cfeature from scipy import interpolate plt.switch_backend('agg') +cartopy.config['pre_existing_data_dir'] = \ + os.getenv('CARTOPY_DIR', cartopy.config.get('pre_existing_data_dir')) ################################################################################################## ################################################################################################## def interpolate_data_to_grid(grid_file,data_file,var): - # Open files + # Open files data_nc = netCDF4.Dataset(data_file,'r') grid_nc = netCDF4.Dataset(grid_file,'r') @@ -84,9 +88,9 @@ def write_to_file(filename,data,var,xtime): data_nc.createDimension('StrLen',64) data_nc.createDimension('Time',None) - # Create time variable + # Create time variable time = data_nc.createVariable('xtime','S1',('Time','StrLen')) - time[:,:] = netCDF4.stringtochar(xtime) + time[:,:] = netCDF4.stringtochar(xtime) # Set variables data_var = data_nc.createVariable(var,np.float64,('Time','nCells')) @@ -98,28 +102,31 @@ def write_to_file(filename,data,var,xtime): def plot_interp_data(lon_data,lat_data,data,lon_grid,lat_grid,interp_data,var_label,var_abrev,time): - levels = np.linspace(np.amin(data),np.amax(data),10) # Plot data fig = plt.figure() - ax0 = fig.add_subplot(2,1,1) - cf = ax0.contourf(lon_data,lat_data,data,levels=levels) - m = Basemap(projection='cyl',llcrnrlat=-90,urcrnrlat=90,\ - llcrnrlon=0,urcrnrlon=360,resolution='c') - m.fillcontinents(color='tan',lake_color='lightblue') - m.drawcoastlines() - ax0.set_title('data '+time.strip()) + levels = np.linspace(np.amin(data),np.amax(data),100) + ax0 = fig.add_subplot(2, 1, 1, projection=ccrs.PlateCarree()) + cf = ax0.contourf(lon_data, lat_data, data, levels=levels, + transform=ccrs.PlateCarree()) + ax0.set_extent([0, 359.9, -90, 90], crs=ccrs.PlateCarree()) + ax0.add_feature(cfeature.LAND, zorder=100) + ax0.add_feature(cfeature.LAKES, alpha=0.5, zorder=101) + ax0.add_feature(cfeature.COASTLINE, zorder=101) + ax0.set_title('data '+time.strip().decode()) cbar = fig.colorbar(cf,ax=ax0) cbar.set_label(var_label) # Plot interpolated data - ax1 = fig.add_subplot(2,1,2) - cf = ax1.tricontourf(lon_grid,lat_grid,interp_data,levels=levels) - m = Basemap(projection='cyl',llcrnrlat=-90,urcrnrlat=90,\ - llcrnrlon=0,urcrnrlon=360,resolution='c') - m.fillcontinents(color='tan',lake_color='lightblue') - m.drawcoastlines() - ax1.set_title('interpolated data '+time.strip()) + ax1 = fig.add_subplot(2, 1, 2, projection=ccrs.PlateCarree()) + levels = np.linspace(np.amin(interp_data),np.amax(interp_data),100) + cf = ax1.tricontourf(lon_grid,lat_grid,interp_data,levels=levels, + transform=ccrs.PlateCarree()) + ax1.set_extent([0, 359.9, -90, 90], crs=ccrs.PlateCarree()) + ax1.add_feature(cfeature.LAND, zorder=100) + ax1.add_feature(cfeature.LAKES, alpha=0.5, zorder=101) + ax1.add_feature(cfeature.COASTLINE, zorder=101) + ax1.set_title('interpolated data '+time.strip().decode()) cbar = fig.colorbar(cf,ax=ax1) cbar.set_label(var_label) @@ -136,9 +143,9 @@ def plot_interp_data(lon_data,lat_data,data,lon_grid,lat_grid,interp_data,var_la parser = argparse.ArgumentParser() parser.add_argument('--plot',action='store_true') args = parser.parse_args() - + nplot = 10 - + # Files to interpolate to/from grid_file = './mesh.nc' wind_file = './wnd10m.nc' @@ -148,22 +155,22 @@ def plot_interp_data(lon_data,lat_data,data,lon_grid,lat_grid,interp_data,var_la # Interpolation of u and v velocities lon_grid,lat_grid,u_interp,lon_data,lat_data,u_data,xtime = interpolate_data_to_grid(grid_file,wind_file,'U_GRD_L103') lon_grid,lat_grid,v_interp,lon_data,lat_data,v_data,xtime = interpolate_data_to_grid(grid_file,wind_file,'V_GRD_L103') - + # Calculate and plot velocity magnitude if args.plot: for i in range(u_data.shape[0]): if i % nplot == 0: print('Plotting vel: '+str(i)) - + data = np.sqrt(np.square(u_data[i,:,:]) + np.square(v_data[i,:,:])) interp_data = np.sqrt(np.square(u_interp[i,:]) + np.square(v_interp[i,:])) - + plot_interp_data(lon_data,lat_data,data,lon_grid,lat_grid,interp_data,'velocity magnitude','vel',xtime[i]) # Interpolation of atmospheric pressure lon_grid,lat_grid,p_interp,lon_data,lat_data,p_data,xtime = interpolate_data_to_grid(grid_file,pres_file,'PRMSL_L101') - + # Plot atmopheric pressure if args.plot: for i in range(p_data.shape[0]): @@ -172,7 +179,7 @@ def plot_interp_data(lon_data,lat_data,data,lon_grid,lat_grid,interp_data,var_la print('Plotting pres: '+str(i)) plot_interp_data(lon_data,lat_data,p_data[i,:,:],lon_grid,lat_grid,p_interp[i,:],'atmospheric pressure','pres',xtime[i]) - + # Write to NetCDF file subprocess.call(['rm',forcing_file]) write_to_file(forcing_file,u_interp,'windSpeedU',xtime) diff --git a/testing_and_setup/compass/ocean/jigsaw_to_MPAS/coastal_tools.py b/testing_and_setup/compass/ocean/jigsaw_to_MPAS/coastal_tools.py index 5280aaafd4..12a1b62d4f 100755 --- a/testing_and_setup/compass/ocean/jigsaw_to_MPAS/coastal_tools.py +++ b/testing_and_setup/compass/ocean/jigsaw_to_MPAS/coastal_tools.py @@ -16,13 +16,18 @@ import matplotlib.pyplot as plt from scipy import spatial, io import timeit -from mpl_toolkits.basemap import Basemap +import os +import cartopy +import cartopy.crs as ccrs +import cartopy.feature as cfeature from rasterio.features import rasterize from affine import Affine import shapely.geometry from geometric_features.plot import subdivide_geom import jigsaw_to_MPAS.mesh_definition_tools as mdt plt.switch_backend('agg') +cartopy.config['pre_existing_data_dir'] = \ + os.getenv('CARTOPY_DIR', cartopy.config.get('pre_existing_data_dir')) # Constants km = 1000.0 @@ -354,9 +359,10 @@ def create_background_mesh(grd_box, ddeg, mesh_type, dx_min, dx_max, # {{{ plt.plot(lat_grd, cell_width_lat) plt.savefig('bckgrnd_grid_cell_width_vs_lat' + str(call) + '.png') - plt.figure() - plt.contourf(lon_grd, lat_grd, cell_width) - plot_coarse_coast(plot_box) + fig = plt.figure() + ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree()) + plt.contourf(lon_grd, lat_grd, cell_width, transform=ccrs.PlateCarree()) + plot_coarse_coast(ax, plot_box) plt.colorbar() plt.savefig( 'bckgnd_grid_cell_width' + @@ -447,14 +453,17 @@ def extract_coastlines(nc_file, nc_vars, region_box, z_contour=0, n_longest=10, lon, lat, bathy_data, plot_box) # Plot bathymetry data, coastlines and region boxes - plt.figure() + + fig = plt.figure() + ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree()) levels = np.linspace(np.amin(z_plot), np.amax(z_plot), 100) ds = 100 # Downsample dsx = np.arange(0, lon_plot.size, ds) # bathy data dsy = np.arange(0, lat_plot.size, ds) # to speed up dsxy = np.ix_(dsy, dsx) # plotting - plt.contourf(lon_plot[dsx], lat_plot[dsy], z_plot[dsxy], levels=levels) - plot_coarse_coast(plot_box) + plt.contourf(lon_plot[dsx], lat_plot[dsy], z_plot[dsxy], levels=levels, + transform=ccrs.PlateCarree()) + plot_coarse_coast(ax, plot_box) plt.plot(coastlines[:, 0], coastlines[:, 1], color='white') for box in region_box["include"]: plot_region_box(box, 'b') @@ -531,11 +540,13 @@ def distance_to_coast(coastlines, lon_grd, lat_grd, origin, nn_search, smooth_wi lon_grd, lat_grd, D, plot_box) # Plot distance to coast - plt.figure() + fig = plt.figure() + ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree()) D_plot = D_plot / km levels = np.linspace(np.amin(D_plot), np.amax(D_plot), 10) - plt.contourf(lon_plot, lat_plot, D_plot, levels=levels) - plot_coarse_coast(plot_box) + plt.contourf(lon_plot, lat_plot, D_plot, levels=levels, + transform=ccrs.PlateCarree()) + plot_coarse_coast(ax, plot_box) plt.plot(coastlines[:, 0], coastlines[:, 1], color='white') plt.grid( xdata=lon_plot, @@ -784,11 +795,13 @@ def compute_cell_width(D, cell_width, lon, lat, dx_min, trans_start, trans_width lon_grd, lat_grd, cell_width / km, plot_box) # Plot cell width - plt.figure() + fig = plt.figure() + ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree()) levels = np.linspace(np.amin(cell_width_plot), np.amax(cell_width_plot), 100) - plt.contourf(lon_plot, lat_plot, cell_width_plot, levels=levels) - plot_coarse_coast(plot_box) + plt.contourf(lon_plot, lat_plot, cell_width_plot, levels=levels, + transform=ccrs.PlateCarree()) + plot_coarse_coast(ax, plot_box) plt.plot(coastlines[:, 0], coastlines[:, 1], color='white') if restrict_box: for box in restrict_box["include"]: @@ -1073,11 +1086,10 @@ def flag_wrap(box): ############################################################## -def plot_coarse_coast(plot_box): +def plot_coarse_coast(ax, plot_box): - m = Basemap(projection='cyl', llcrnrlat=plot_box[2], urcrnrlat=plot_box[3], - llcrnrlon=plot_box[0], urcrnrlon=plot_box[1], resolution='c') - m.drawcoastlines() + ax.set_extent(plot_box, crs=ccrs.PlateCarree()) + ax.add_feature(cfeature.COASTLINE) ##############################################################