import copy
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
[docs]def pcol( x, y, data, projection=None, vmin=None, vmax=None, **kwargs):
"""
Plots 2D scalar fields on the MITgcm cubed sphere grid with pcolormesh.
Parameters
----------
x : array_like
'xg', that is, x coordinate of the points one half grid cell to the
left and bottom, that is vorticity points for tracers, etc.
y : array_like
'yg', that is, y coordinate of same points
data : array_like
scalar field at tracer points
projection : Basemap instance, optional
used to transform if present.
Unfortunatly, cylindrical and conic maps are limited to
the [-180 180] range.
projection = 'sphere' results in a 3D visualization on the sphere
without any specific projection. Good for debugging.
Example
-------
>>> from mpl_toolkits.basemap import Basemap
>>> import MITgcmutils as mit
>>> import matplotlib.pyplot as plt
>>> from sq import sq
>>>
>>> x=mit.rdmds('XG'); y=mit.rdmds('YG'); e=mit.rdmds('Eta',np.Inf)
>>> fig = plt.figure();
>>> mp = Basemap(projection='moll',lon_0 = 0.,
>>> resolution = 'l', area_thresh = 1000.)
>>> plt.clf()
>>> h = mit.cs.pcol(x,y,sq(e), projection = mp)
>>> mp.fillcontinents(color = 'grey')
>>> mp.drawmapboundary()
>>> mp.drawmeridians(np.arange(0, 360, 30))
>>> mp.drawparallels(np.arange(-90, 90, 30))
>>> plt.show()
"""
# pcol first divides the 2D cs-field(6*n,n) into six faces. Then for
# each face, an extra row and colum is added from the neighboring faces in
# order to fool pcolor into drawing the entire field and not just
# (n-1,m-1) data points. There are two corner points that have no explicit
# coordinates so that they have to be found by
# interpolation/averaging. Then each face is divided into 4 tiles,
# assuming cs-geometry, and each tile is plotted individually in
# order to avoid problems due to ambigous longitude values (the jump
# between -180 and 180, or 360 and 0 degrees). As long as the poles
# are at the centers of the north and south faces and the first tile is
# symmetric about its center this should work.
# get the figure handle
fig=plt.gcf()
mapit = 0
if projection!=None:
mp = projection
if mp=='sphere': mapit=-1
else: mapit = 1
# convert to [-180 180[ representation
x = np.where(x>180,x-360.,x)
ny,nx = data.shape
# determine range for color range
cax = [data.min(),data.max()]
if cax[1]-cax[0]==0: cax = [cax[0]-1,cax[1]+1]
if vmin!=None: cax[0]=vmin
if vmax!=None: cax[1]=vmax
if mapit == -1:
# set up 3D plot
if len(fig.axes)>0:
# if present, remove and replace the last axis of fig
geom=fig.axes[-1].get_geometry()
plt.delaxes(fig.axes[-1])
else:
# otherwise use full figure
geom = ((1,1,1))
ax = fig.add_subplot(geom[0],geom[1],geom[2],projection = '3d',
facecolor='None')
# define color range
tmp = data - data.min()
N = tmp/tmp.max()
# use this colormap
colmap = copy.copy(cm.jet)
colmap.set_bad('w',1.0)
mycolmap = colmap(N) #cm.jet(N)
ph=np.array([])
jc=x.shape[0]//2
xxf=np.empty((jc+1,jc+1,4))
yyf=xxf
ffld=np.empty((jc,jc,4))
xff=[]
yff=[]
fldf=[]
for k in range(0,6):
ix = np.arange(0,ny) + k*ny
xff.append(x[0:ny,ix])
yff.append(y[0:ny,ix])
fldf.append(data[0:ny,ix])
# find the missing corners by interpolation (one in the North Atlantic)
xfodd = (xff[0][-1,0]+xff[2][-1,0]+xff[4][-1,0])/3.
yfodd = (yff[0][-1,0]+yff[2][-1,0]+yff[4][-1,0])/3.
# and one south of Australia
xfeven= (xff[1][0,-1]+xff[3][0,-1]+xff[5][0,-1])/3.
yfeven= (yff[1][0,-1]+yff[3][0,-1]+yff[5][0,-1])/3.
# loop over tiles
for k in range(0,6):
kodd = 2*(k//2)
kodd2 = kodd
if kodd==4: kodd2=kodd-6
keven = 2*(k//2)
keven2 = keven
if keven==4: keven2=keven-6
fld = fldf[k]
if np.mod(k+1,2):
xf = np.vstack( [ np.column_stack( [xff[k],xff[1+kodd][:,0]] ),
np.flipud(np.append(xff[2+kodd2][:,0],xfodd))] )
yf = np.vstack( [ np.column_stack( [yff[k],yff[1+kodd][:,0]] ),
np.flipud(np.append(yff[2+kodd2][:,0],yfodd))] )
else:
xf = np.column_stack( [np.vstack( [xff[k],xff[2+keven2][0,:]] ),
np.flipud(np.append(xff[3+keven2][0,:],
xfeven))] )
yf = np.column_stack( [np.vstack( [yff[k],yff[2+keven2][0,:]] ),
np.flipud(np.append(yff[3+keven2][0,:],
yfeven))] )
if mapit==-1:
ix = np.arange(0,ny) + k*ny
# no projection at all (projection argument is 'sphere'),
# just convert to cartesian coordinates and plot a 3D sphere
deg2rad=np.pi/180.
xcart,ycart,zcart = sph2cart( xf*deg2rad, yf*deg2rad )
ax.plot_surface(xcart,ycart,zcart,rstride=1,cstride=1,
facecolors=mycolmap[0:ny,ix],
linewidth=2,shade=False)
ph = np.append(ph, ax)
else:
# divide all faces into 4 because potential problems arise at
# the centers
for kf in range(0,4):
if kf==0: i0,i1,j0,j1 = 0, jc+1, 0, jc+1
elif kf==1: i0,i1,j0,j1 = 0, jc+1,jc,2*jc+1
elif kf==2: i0,i1,j0,j1 = jc,2*jc+1, 0, jc+1
elif kf==3: i0,i1,j0,j1 = jc,2*jc+1,jc,2*jc+1
xx = xf[i0:i1,j0:j1]
yy = yf[i0:i1,j0:j1]
ff = fld[i0:i1-1,j0:j1-1]
if np.median(xx) < 0:
xx = np.where(xx>=180,xx-360.,xx)
else:
xx = np.where(xx<=-180,xx+360.,xx)
# if provided use projection
if mapit==1: xx,yy = mp(xx,yy)
# now finally plot 4x6 tiles
ph = np.append(ph, plt.pcolormesh(xx, yy, ff,
vmin=cax[0], vmax=cax[1],
**kwargs))
if mapit == -1:
# ax.axis('image')
ax.set_axis_off()
# ax.set_visible=False
# add a reasonable colormap
m = cm.ScalarMappable(cmap=colmap)
m.set_array(data)
plt.colorbar(m)
elif mapit == 0:
ax = fig.axes[-1]
ax.axis('image')
plt.grid(True)
return ph
def sph2cart(azim_sph_coord, elev_sph_coord):
r = np.cos(elev_sph_coord)
x = -r * np.sin(azim_sph_coord)
y = r * np.cos(azim_sph_coord)
z = np.sin(elev_sph_coord)
return x, y, z