def main():
plot_utils.apply_plot_params(width_cm=20, height_cm=20, font_size=10)
high_hles_years = [1993, 1995, 1998]
low_hles_years = [1997, 2001, 2006]
data_path = "/BIG1/skynet1_rech1/diro/sample_obsdata/eraint/eraint_uvslp_years_198111_201102_NDJmean_ts.nc"
with xr.open_dataset(data_path) as ds:
print(ds)
u = get_composit_for_name(ds, "u10", high_years_list=high_hles_years, low_years_list=low_hles_years)
v = get_composit_for_name(ds, "v10", high_years_list=high_hles_years, low_years_list=low_hles_years)
msl = get_composit_for_name(ds, "msl", high_years_list=high_hles_years, low_years_list=low_hles_years)
lons = ds["longitude"].values
lats = ds["latitude"].values
print(lats)
print(msl.shape)
print(lons.shape, lats.shape)
lons2d, lats2d = np.meshgrid(lons, lats)
fig = plt.figure()
map = Basemap(llcrnrlon=-130, llcrnrlat=22, urcrnrlon=-28,
urcrnrlat=65, projection='lcc', lat_1=33, lat_2=45,
lon_0=-95, resolution='i', area_thresh=10000)
clevs = np.arange(-11.5, 12, 1)
cmap = cm.get_cmap("bwr", len(clevs) - 1)
bn = BoundaryNorm(clevs, len(clevs) - 1)
x, y = map(lons2d, lats2d)
im = map.contourf(x, y, msl / 100, levels=clevs, norm=bn, cmap=cmap) # convert to mb (i.e hpa)
map.colorbar(im)
stride = 2
ux, vy = map.rotate_vector(u, v, lons2d, lats2d)
qk = map.quiver(x[::stride, ::stride], y[::stride, ::stride], ux[::stride, ::stride], vy[::stride, ::stride],
scale=10, width=0.01, units="inches")
plt.quiverkey(qk, 0.5, -0.1, 2, "2 m/s", coordinates="axes")
map.drawcoastlines(linewidth=0.5)
map.drawcountries()
map.drawstates()
#plt.show()
fig.savefig("hles_wind_compoosits.png", bbox_inches="tight", dpi=300)
def test_cylindrical(self):
# Cylindrical case
B = Basemap()
u, v, lat, lon = self.make_array()
ru, rv = B.rotate_vector(u, v, lon, lat)
# Check that the vectors are identical.
assert_almost_equal(ru, u)
assert_almost_equal(rv, v)
def test_npstere(self):
# NP Stereographic case
B = Basemap(projection='npstere', boundinglat=50., lon_0=0.)
u, v, lat, lon = self.make_array()
v = np.ones((len(lat), len(lon)))
ru, rv = B.rotate_vector(u, v, lon, lat)
assert_almost_equal(ru[2, :], [1, -1, -1, 1], 6)
assert_almost_equal(rv[2, :], [1, 1, -1, -1], 6)
def test_cylindrical(self):
# Cylindrical case
B = Basemap()
u,v,lat,lon=self.make_array()
ru, rv = B.rotate_vector(u,v, lon, lat)
# Check that the vectors are identical.
assert_almost_equal(ru, u)
assert_almost_equal(rv, v)
def test_npstere(self):
# NP Stereographic case
B=Basemap(projection='npstere', boundinglat=50., lon_0=0.)
u,v,lat,lon=self.make_array()
v = np.ones((len(lat), len(lon)))
ru, rv = B.rotate_vector(u,v, lon, lat)
assert_almost_equal(ru[2, :],[1,-1,-1,1], 6)
assert_almost_equal(rv[2, :],[1,1,-1,-1], 6)
def test():
plt.figure()
b = Basemap()
u = np.array([1, ])
v = np.array([1, ])
lon, lat = np.array([-90, ]), np.array([45, ])
xx, yy = b(lon, lat)
print(xx.shape)
b.quiver(xx, yy, u, v, color="r")
urot, vrot = b.rotate_vector(u, v, lon, lat)
b.quiver(xx, yy, urot, vrot, color="g")
b.drawcoastlines()
# Plot the same in rotpole projection
HL_LABEL = "CRCM5_HL"
NEMO_LABEL = "CRCM5_NEMO"
sim_label_to_path = OrderedDict(
[(HL_LABEL, "/RESCUE/skynet3_rech1/huziy/CNRCWP/C5/2016/2-year-runs/coupled-GL+stfl_oneway/Samples"),
(NEMO_LABEL, "/HOME/huziy/skynet3_rech1/CNRCWP/C5/2016/2-year-runs/coupled-GL+stfl/Samples")]
)
# get a coord file ...
coord_file = ""
found_coord_file = False
for mdir in os.listdir(sim_label_to_path[HL_LABEL]):
mdir_path = os.path.join(sim_label_to_path[HL_LABEL], mdir)
if not os.path.isdir(mdir_path):
continue
for fn in os.listdir(mdir_path):
print(fn)
if fn[:2] not in ["pm", "dm", "pp", "dp"]:
continue
coord_file = os.path.join(mdir_path, fn)
found_coord_file = True
if found_coord_file:
break
bmp, lons, lats = nemo_hl_util.get_basemap_obj_and_coords_from_rpn_file(path=coord_file)
plt.figure()
urot, vrot = bmp.rotate_vector(u, v, lon, lat)
xx, yy = bmp(lon, lat)
bmp.quiver(xx, yy, urot, vrot, color="b")
bmp.drawcoastlines()
plt.show()
def test_nan(self):
B = Basemap()
u,v,lat,lon=self.make_array()
# Set one element to 0, so that the vector magnitude is 0.
u[1,1] = 0.
ru, rv = B.rotate_vector(u,v, lon, lat)
assert not np.isnan(ru).any()
assert_almost_equal(u, ru)
assert_almost_equal(v, rv)
def test_nan(self):
B = Basemap()
u, v, lat, lon = self.make_array()
# Set one element to 0, so that the vector magnitude is 0.
u[1, 1] = 0.
ru, rv = B.rotate_vector(u, v, lon, lat)
assert not np.isnan(ru).any()
assert_almost_equal(u, ru)
assert_almost_equal(v, rv)
def plot(fig, num, heading, lat, lon, u, v):
ax = fig.add_subplot(str(num))
ax.set_title(heading, pad=20)
# make South pole Stereographic projection
m = Basemap(projection='spstere',
boundinglat=-30,
lon_0=180.,
resolution='h',
round=True)
# may not be required
#ax.title(heading)
# draw parallel longitude lines
m.drawmeridians(np.arange(-180., 181., 30.), labels=[1, 1, 1, 1])
# draw parallel latitude lines
m.drawparallels(np.arange(-90., -40., 10.), labels=[1, 0, 0, 1])
# add the observations to the map
m.fillcontinents(alpha=0.42)
skip = 8
urot, vrot, x, y = m.rotate_vector(u, v, lon, lat, returnxy=True)
c_m = np.hypot(x, y)
c_m_k = np.ones(x[::skip, ::skip].shape)
m.quiver(x, y, urot, vrot, c_m, color='k', scale=80)
print(x.shape, y.shape, urot.shape, vrot.shape)
#print(x[::2,::2].shape,y[::2].shape,urot[::2].shape,vrot[::2].shape)
#m.quiver(x[::skip,::skip], y[::skip,::skip], urot[::skip,::skip], vrot[::skip,::skip], c_m_k, color='k', scale=80)
#m.quiver(x, y, urot, vrot, c_m, c='scale=20)
#m.quiver(lon, lat, u, v, latlon=True, scale=10)
#q2 = m.quiver(x, y, u, v, latlon=True, scale=30, scale_units='inches')
#m.streamplot(lon,lat,u,v)
#ax.set_aspect('equal')
plt.colorbar(orientation='horizontal', ax=ax)
lats1 = -90.+dellat*np.arange(nlats)
lons1 = -180.+dellon*np.arange(nlons)
lons, lats = np.meshgrid(lons1, lats1)
# plot vectors in geographical (lat/lon) coordinates.
# north polar projection.
m = Basemap(lon_0=-135,boundinglat=25,round=True,
resolution='c',area_thresh=10000.,projection='npstere')
# create a figure, add an axes.
fig=plt.figure(figsize=(8,8))
ax = fig.add_axes([0.1,0.1,0.7,0.7])
# rotate wind vectors to map projection coordinates.
# (also compute native map projections coordinates of lat/lon grid)
# only do Northern Hemisphere.
urot,vrot,x,y = m.rotate_vector(u[36:,:],v[36:,:],lons[36:,:],lats[36:,:],returnxy=True)
# plot filled contours over map.
cs = m.contourf(x,y,p[36:,:],15,cmap=plt.cm.jet)
# plot wind vectors over map.
Q = m.quiver(x,y,urot,vrot) #or specify, e.g., width=0.003, scale=400)
qk = plt.quiverkey(Q, 0.95, 1.05, 25, '25 m/s', labelpos='W')
m.colorbar(pad='12%') # draw colorbar
m.drawcoastlines()
m.drawcountries()
# draw parallels
delat = 20.
circles = np.arange(0.,90.+delat,delat).tolist()+\
np.arange(-delat,-90.-delat,-delat).tolist()
m.drawparallels(circles,labels=[1,1,1,1])
# draw meridians
delon = 45.
def make_plev_map(filename=None,
level=None,
cfld=None,
cscl=None,
coff=None,
colormap=None,
cmin=None,
cmax=None,
cbunits=None,
nc=None,
z1fld=None,
z1levs=None,
z1off=None,
z1scl=None,
z2fld=None,
z2levs=None,
z2off=None,
z2scl=None,
ufld=None,
vfld=None,
titletext=None,
filepat=None):
wrf=xr.open_dataset(filename,engine='pynio')
run=filename.split('/')[-1]
init=datetime.strptime(run.split('_')[0],'%Y%m%d%H')
fhr=run.split('.')[1]
valid=init+timedelta(hours=int(fhr))
level=float(level)*100.
terr=wrf['HGT_P0_L1_GLC0'].values
hgt=wrf['HGT_P0_L100_GLC0'].sel(lv_ISBL0=level,method='nearest').values
lon=wrf['gridlon_0'].values
lat=wrf['gridlat_0'].values
ug=wrf[ufld].sel(lv_ISBL0=level,method='nearest').values
vg=wrf[vfld].sel(lv_ISBL0=level,method='nearest').values
u=ug*np.cos(wrf['gridrot_0'])-vg*np.sin(wrf['gridrot_0'])
v=ug*np.sin(wrf['gridrot_0'])+vg*np.cos(wrf['gridrot_0'])
if cfld=='SPD':
c=cscl*np.sqrt(u**2+v**2)+coff
elif cfld=='AVO':
f=2.*7.29e-5*np.sin(lat* np.pi / 180.)
dudx,dudy=np.gradient(u,3000.,edge_order=2)
dvdx,dvdy=np.gradient(v,3000.,edge_order=2)
c=cscl*(dvdx-dudy+f)+coff
else:
c=cscl*wrf[cfld].sel(lv_ISBL0=level).values+coff
z1=(z1scl*wrf[z1fld].sel(lv_ISBL0=level).values)+z1off
z2=(z2scl*wrf[z2fld].sel(lv_ISBL0=level).values)+z2off
fields=titletext+str(int(level/100.))+' hPa'
c=np.ma.masked_where(terr > hgt,c)
u=np.ma.masked_where(terr > hgt,u)
v=np.ma.masked_where(terr > hgt,v)
z1=np.ma.masked_where(terr > hgt,z1)
z2=np.ma.masked_where(terr > hgt,z2)
mask=np.ones_like(terr)
mask=np.ma.masked_where(terr < hgt+50,mask)
m=Basemap(projection='lcc',width=3000*550,height=3000*375,
resolution='i',lat_1=-32.8,lat_2=-32.8,lat_0=-32.8,lon_0=-67.0)
x,y=m(lon,lat)
N=19.
font0 = FontProperties()
font0.set_family('monospace')
my_dpi=100
fig, ax = plt.subplots(figsize=(11.0, 8.5))
m.drawcoastlines()
m.drawcountries(linewidth=1.0)
m.drawstates(color=(0.5,0.5,0.5),linewidth=0.5)
#m.drawparallels(np.arange(-80.,81.,1.))
#m.drawmeridians(np.arange(-180.,181.,1.))
C=m.pcolormesh(x,y,c,vmin=cmin,vmax=cmax,ax=ax,cmap=colormap)
plt.title('Initialized '+init.isoformat()+' UTC\n'
'F'+fhr+' Valid '+valid.isoformat()+' UTC',loc='right',fontdict={'family': 'monospace'})#plt.colorbar(orientation='horizontal',shrink=0.5,pad=0.0)
plt.title('University of Illinois 3 km WRF Forecast\n'+
fields,loc='left',fontdict={'family': 'sans-serif'})#plt.colorbar(orientation='horizontal',shrink=0.5,pad=0.0)
for key in landmark_dict.keys():
kx,ky=m(landmark_dict[key][0],landmark_dict[key][1])
plt.text(kx,ky,key,fontsize=8,fontweight='light',
ha='center',va='center',color='b')
CS = m.contour(x,y,z1,z1levs,colors='k',lw=2,ax=ax)
cl=plt.clabel(CS, fontsize=9, inline=1, fmt='%1.0f',fontproperties=font0)
CS2 = m.contour(x,y,z2,z2levs,colors='g',lw=2)
for cl in CS2.collections:
cl.set_dashes([(0, (2.0, 2.0))])
print(level)
print(np.min(c))
print(np.max(c))
cl2=plt.clabel(CS2, fontsize=9, inline=1, fmt='%1.0f',fontproperties=font0,ax=ax)
Urot, Vrot = m.rotate_vector(u,v,lon,lat)
m.barbs(x[::20,::20],y[::20,::20],Urot[::20,::20],Vrot[::20,::20],
barb_increments=dict(half=2.5, full=5., flag=25),
length=5,flagcolor='none',barbcolor='k',lw=0.5,ax=ax)
m.contour(x,y,terr,[500.,1500.],colors=('blue','red'),alpha=0.5)
m.pcolormesh(x,y,mask,cmap=cm.Gray5,vmin=0.,vmax=1.)
cax = fig.add_axes([0.2, 0.12, 0.4, 0.02])
cb=plt.colorbar(C, cax=cax, orientation='horizontal')
cb.set_label(cbunits, labelpad=-10, x=1.1)
ram = cStringIO.StringIO()
plt.savefig(ram, format='png',dpi=my_dpi, bbox_inches='tight')
ram.seek(0)
im = Image.open(ram)
im2 = im.convert('RGB')
im2.save( run+'_'+filepat+'_'+str(int(level/100.))+'_'+fhr+'.png' , format='PNG')
resolution='c')
#m = Basemap(projection='ortho',lat_0=lats.min(),lon_0=lons.min(), resolution='l')
lon, lat = np.meshgrid(lons, lats)
XX, YY = m(lon, lat)
cmap = plt.cm.Reds
#cmap = plt.cm.RdBu_r
#norm = mpl.colors.Normalize(vmin=5, vmax=10)
clevs = np.arange(990, 1010, 1)
cs1 = m.contour(XX, YY, p, clevs, linewidths=1.0, colors='k', animated=True)
#cs2 = m.contourf(XX,YY,p,clevs,cmap=cmap,animated=True)
urot, vrot, x, y = m.rotate_vector(u[:, :],
v[:, :],
lons[:],
lats[:],
returnxy=True)
Q = m.quiver(lon, lat, u, v,
latlon=True) #or specify, e.g., width=0.003, scale=400)
qk = plt.quiverkey(Q, 0.95, 1.05, 25, '25 m/s', labelpos='W')
###########################################################################################################
#ugrid,newlons = shiftgrid(180.,u,lons,start=False)
#vgrid,newlons = shiftgrid(180.,v,lons,start=False)
# transform vectors to projection grid.
#uproj,vproj,xx,yy = m.transform_vector(ugrid,vgrid,newlons,latitudes,31,31,returnxy=True,masked=True)
#Q = m.quiver(XX,YY,u,v,scale=9000)
# make quiver key.
#qk = plt.quiverkey(Q, 0.1, 0.1, 20, '20 m/s', labelpos='W')
#m.barbs(XX,YY,u,v, length=7, color='red')
fontsize=7,
linewidth=0.5)
levels = np.arange(-10, 10, 0.1)
cmap1 = plt.cm.coolwarm
# Plot anything you want underneath as a colour map, e.g. q or precip
mymapf = plt.contourf(lonall2,
latall2,
divg,
levels,
cmap=cmap1,
extend='both')
# This is important for plotting vectors but I can't remember why...
xx, yy = m(*np.meshgrid(lonx2, latx2))
u_out, v_out = m.rotate_vector(qu, qv, lonx2, latx2)
# Plotting vectors
# xx[::5,::5], yy[::5,::5], u_out[::5,::5], v_out[::5,::5] - this part is the spacing, this means 1 in every 5 vectors will be plotted (otherwise it can be too many to see). Note that more vectors need to be plotted for lower-res data e.g. NCEP
Q = m.quiver(
xx[::4, ::4],
yy[::4, ::4],
u_out[::4, ::4],
v_out[::4, ::4],
width=0.0025,
scale=0.85,
color='k',
pivot='mid'
) # width and scale are the size and length of vectors - play around with them to get it right
plt.quiverkey(
Q,
m.drawparallels(np.arange(0.,81.,20.),linewidth=0.4,color='gray')
m.drawmeridians(np.arange(0.,360.,60.),linewidth=0.4,color='gray',labels=[True,True,True,True])
psi_p1,lon1=addcyclic(psi_p,lon) #This might raise an error when you use basemap 1.2.0 on python3.7
lonsn, latsn = np.meshgrid(lon1, lat[0:160])
x_cyc,y_cyc=m(lonsn,latsn)
#Plot Pertubation stream-function
lev=np.arange(-24,28,4)
cf=m.contourf(x_cyc,y_cyc,psi_p1[0:160,:]/1e6,levels=lev,cmap='RdBu_r',extend='both')
c=m.contour(x_cyc,y_cyc,psi_p1[0:160,:]/1e6,colors='black',linewidths=0.5,levels=lev)
#plot T-N Flux vector on a Polar Lambert Azimuthal Projection map
lons, lats = np.meshgrid(lon, lat[0:160])
# the vector must be rotated to fit in the projection
px_r, py_r, x_r, y_r = m.rotate_vector(px[0:160,:], py[0:160,:], lons, lats, returnxy=True)
step=10
step90_60=20
Q=m.quiver(x_r[0:59:step,::step90_60],y_r[0:59:step,::step90_60],px_r[0:59:step,::step90_60],py_r[0:59:step,::step90_60],pivot='mid',width=0.0025,scale=500,headwidth=3)
qk = ax.quiverkey(Q, 0.76, 0.8, 25, r'25 m$^2$/s$^2$', labelpos='E',coordinates='figure')
m.quiver(x_r[60::step,::step],y_r[60::step,::step],px_r[60:160:step,::step],py_r[60:160:step,::step],pivot='mid',width=0.0025,scale=500,headwidth=3)
ax1 = fig.add_axes([0.25, 0.05, 0.55, 0.02])
cb=plt.colorbar(cf, cax=ax1, shrink=0.0, orientation='horizontal')
cb.ax.tick_params(labelsize='small')
ax1.set_xlabel('Stream function anomalies $\mathcal{\psi}$$^\prime$(10$^6$ m$^2$/s)')
ax.text(0.005, 1.045, 'ECMWF ERA-Interim [T$_L$255L60 Cy31r2@4D-Var]\nWave Activity Flux (${Takaya & Nakamura,2001}$)',color='blue',fontsize=7,transform=ax.transAxes)
ax.text(0.83, 1.063, 'January,1981',color='red',fontsize=10,transform=ax.transAxes)
def main():
plot_utils.apply_plot_params(width_cm=20, height_cm=20, font_size=10)
high_hles_years = [1993, 1995, 1998]
low_hles_years = [1997, 2001, 2006]
data_path = "/BIG1/skynet1_rech1/diro/sample_obsdata/eraint/eraint_uvslp_years_198111_201102_NDJmean_ts.nc"
with xr.open_dataset(data_path) as ds:
print(ds)
u = get_composit_for_name(ds,
"u10",
high_years_list=high_hles_years,
low_years_list=low_hles_years)
v = get_composit_for_name(ds,
"v10",
high_years_list=high_hles_years,
low_years_list=low_hles_years)
msl = get_composit_for_name(ds,
"msl",
high_years_list=high_hles_years,
low_years_list=low_hles_years)
lons = ds["longitude"].values
lats = ds["latitude"].values
print(lats)
print(msl.shape)
print(lons.shape, lats.shape)
lons2d, lats2d = np.meshgrid(lons, lats)
fig = plt.figure()
map = Basemap(llcrnrlon=-130,
llcrnrlat=22,
urcrnrlon=-28,
urcrnrlat=65,
projection='lcc',
lat_1=33,
lat_2=45,
lon_0=-95,
resolution='i',
area_thresh=10000)
clevs = np.arange(-11.5, 12, 1)
cmap = cm.get_cmap("bwr", len(clevs) - 1)
bn = BoundaryNorm(clevs, len(clevs) - 1)
x, y = map(lons2d, lats2d)
im = map.contourf(x, y, msl / 100, levels=clevs, norm=bn,
cmap=cmap) # convert to mb (i.e hpa)
map.colorbar(im)
stride = 2
ux, vy = map.rotate_vector(u, v, lons2d, lats2d)
qk = map.quiver(x[::stride, ::stride],
y[::stride, ::stride],
ux[::stride, ::stride],
vy[::stride, ::stride],
scale=10,
width=0.01,
units="inches")
plt.quiverkey(qk, 0.5, -0.1, 2, "2 m/s", coordinates="axes")
map.drawcoastlines(linewidth=0.5)
map.drawcountries()
map.drawstates()
#plt.show()
fig.savefig("hles_wind_compoosits.png", bbox_inches="tight", dpi=300)
print 'Plotting ' + ut.isoformat()
sys.stdout.flush()
bm.drawcoastlines()
bm.fillcontinents(color=land_color, lake_color=water_color)
# maybe parallels and meridians should be configurable (?)
bm.drawparallels(np.arange(-80., 81., 20.), labels=[1, 1, 0, 0])
bm.drawmeridians(np.arange(-180., 181., 20.), labels=[0, 0, 0, 1])
bm.drawmapboundary(fill_color=water_color)
if nx == None and ny == None:
# vectors remain on input grid
u, v, x, y = bm.rotate_vector(Ys[i] * Ms[i],
Xs[i] * Ms[i],
Lon,
Lat,
returnxy=True)
else:
# Basemap.transform_vector() does not generate results consistent
# with Basemap.rotate_vector(). We re-implment transform_vector()
# here, but call a different 2d interpolator.
uin, vin, xin, yin = bm.rotate_vector(Ys[i] * Ms[i],
Xs[i] * Ms[i],
Lon,
Lat,
returnxy=True)
longs, lats, x, y = bm.makegrid(nx, ny, returnxy=True)
u = sInterp.griddata((xin.flatten(), yin.flatten()),
uin.flatten(), (x, y),
def test():
plt.figure()
b = Basemap()
u = np.array([
1,
])
v = np.array([
1,
])
lon, lat = np.array([
-90,
]), np.array([
45,
])
xx, yy = b(lon, lat)
print(xx.shape)
b.quiver(xx, yy, u, v, color="r")
urot, vrot = b.rotate_vector(u, v, lon, lat)
b.quiver(xx, yy, urot, vrot, color="g")
b.drawcoastlines()
# Plot the same in rotpole projection
HL_LABEL = "CRCM5_HL"
NEMO_LABEL = "CRCM5_NEMO"
sim_label_to_path = OrderedDict([
(HL_LABEL,
"/RESCUE/skynet3_rech1/huziy/CNRCWP/C5/2016/2-year-runs/coupled-GL+stfl_oneway/Samples"
),
(NEMO_LABEL,
"/HOME/huziy/skynet3_rech1/CNRCWP/C5/2016/2-year-runs/coupled-GL+stfl/Samples"
)
])
# get a coord file ...
coord_file = ""
found_coord_file = False
for mdir in os.listdir(sim_label_to_path[HL_LABEL]):
mdir_path = os.path.join(sim_label_to_path[HL_LABEL], mdir)
if not os.path.isdir(mdir_path):
continue
for fn in os.listdir(mdir_path):
print(fn)
if fn[:2] not in ["pm", "dm", "pp", "dp"]:
continue
coord_file = os.path.join(mdir_path, fn)
found_coord_file = True
if found_coord_file:
break
bmp, lons, lats = nemo_hl_util.get_basemap_obj_and_coords_from_rpn_file(
path=coord_file)
plt.figure()
urot, vrot = bmp.rotate_vector(u, v, lon, lat)
xx, yy = bmp(lon, lat)
bmp.quiver(xx, yy, urot, vrot, color="b")
bmp.drawcoastlines()
plt.show()
def plot_flow_vectors_basemap(lons, lats, uu, vv, flow_speed, subregion:list=None, grid_shape=None, ax:Axes=None,
streamplot=False, draw_colorbar=True):
from mpl_toolkits.basemap import Basemap
if grid_shape is None:
grid_shape = (300, 300)
if subregion is None:
subregion = [0, 1, 0, 1]
if ax is None:
fig = plt.figure()
nx, ny = lons.shape
b = Basemap(lon_0=180,
llcrnrlon=lons[35, 35],
llcrnrlat=lats[35, 35],
urcrnrlon=lons[nx // 2, ny // 2],
urcrnrlat=lats[nx // 2, ny // 2],
resolution="i", area_thresh=2000)
# im = b.pcolormesh(xx, yy, flow_speed)
# b.colorbar(im)
stride = 6
uu1, vv1 = b.rotate_vector(uu, vv, lons, lats)
lons_g, lats_g, xx_g, yy_g = b.makegrid(*grid_shape, returnxy=True)
nx, ny = lons_g.shape
i_start, i_end = int(nx * subregion[0]), int(nx * subregion[1])
j_start, j_end = int(ny * subregion[2]), int(ny * subregion[3])
xt, yt, zt = lat_lon.lon_lat_to_cartesian(lons_g.flatten(), lats_g.flatten())
xs, ys, zs = lat_lon.lon_lat_to_cartesian(lons.flatten(), lats.flatten())
ktree = KDTree(data=list(zip(xs, ys, zs)))
dists, inds = ktree.query(list(zip(xt, yt, zt)))
uu_to_plot = uu1.flatten()[inds].reshape(lons_g.shape)
vv_to_plot = vv1.flatten()[inds].reshape(lons_g.shape)
flow_speed_to_plot = flow_speed.flatten()[inds].reshape(lons_g.shape)
clevs = [0, 0.01, 0.02, 0.04, 0.06, 0.08, 0.12, 0.16]
ncolors = len(clevs) - 1
norm = BoundaryNorm(clevs, ncolors)
cmap = cm.get_cmap("gist_ncar_r", ncolors)
im = b.pcolormesh(xx_g, yy_g, flow_speed_to_plot, alpha=0.5, cmap=cmap, norm=norm)
if not streamplot:
b.quiver(xx_g[i_start:i_end:stride, j_start:j_end:stride], yy_g[i_start:i_end:stride, j_start:j_end:stride],
uu_to_plot[i_start:i_end:stride, j_start:j_end:stride], vv_to_plot[i_start:i_end:stride, j_start:j_end:stride],
headlength=2, headaxislength=2, headwidth=4, units="inches", color="k")
else:
b.streamplot(xx_g, yy_g, uu_to_plot, vv_to_plot, linewidth=0.4, density=3, arrowstyle="fancy", arrowsize=0.4, ax=ax, color="k")
# im = b.contourf(xx_g, yy_g, flow_speed_to_plot, levels=clevs, alpha=0.5, cmap=cmap, norm=norm)
cb = b.colorbar(im, location="bottom")
cb.ax.set_visible(draw_colorbar)
b.drawcoastlines(linewidth=0.3, ax=ax)
if ax is None:
fig.savefig("nemo/circ_annual_mean_basemap.png", bbox_inches="tight", dpi=300)
plt.close(fig)
return im
dates = [dates[tidx],]
tidx_offset=tidx
idxs =arange(len(x))
np.random.shuffle(idxs)
idxs = idxs[:20000]
os.system('mkdir -p jpgs/vel')
if True:
if True:
fig=figure()
fig.subplots_adjust(left=0.0,right=1.0,bottom=0.0,top=1.0)
pc=tripcolor(x,y,nv,vel,cmap=cmap,rasterized=True)
del hvel,zcor,uvel,vvel
urot,vrot = proj.rotate_vector(u[idxs],v[idxs],lon[idxs],lat[idxs])
del u,v
#proj.quiver(x[idxs],y[idxs],u[idxs],v[idxs],pivot='mid',alpha=0.5)
#urot=u[idxs]; vrot=v[idxs]
if (uselim):
upper_scale=vmax
else:
upper_scale=10.0 # deep transports
upper_scale=0.1 # surface transports
urot = ma.masked_where(vel[idxs]<0.05*upper_scale,urot)
vrot = ma.masked_where(vel[idxs]<0.05*upper_scale,vrot)
#quiverkws={'scale_units':'x','scale':10./10000.} # for deep transports
quiverkws={'scale_units':'x','scale':upper_scale/10000.}
proj.quiver(x[idxs],y[idxs],urot,vrot,pivot='mid',alpha=0.5,**quiverkws)
# test vector
def geoplot(data=None, lon=None, lat=None, **kw):
'''Show 2D data in a lon-lat plane.
Parameters
-----------
data: array of shape (n_lat, n_lon), or [u_array, v_array]-like for (u,v)
data or None(default) when only plotting the basemap.
lon: n_lon length vector or None(default).
lat: n_lat length vector or None(default).
kw: dict parameters related to basemap or plot functions.
Basemap related parameters:
----------------------------
basemap_kw: dict parameter in the initialization of a Basemap.
proj or projection: map projection name (default='moll')
popular projections: 'ortho', 'np'(='nplaea'), 'sp'(='splaea')
and other projections given from basemap.
lon_0: map center longitude (None as default).
lat_0: map center latitude (None as default).
lonlatcorner: (llcrnrlon, urcrnrlon, llcrnrlat, urcrnrlat).
boundinglat: latitude at the map out boundary (None as default).
basemap_round or round: True(default) or False.
fill_continents: bool value (False as default).
continents_kw: dict parameter used in the Basemap.fillcontinents method.
continents_color: color of continents ('0.33' as default).
lake_color: color of lakes ('none' as default).
ocean_color: color of ocean (None as default).
coastlines_kw: dict parameter used in the Basemap.drawcoastlines method.
coastlines_color: color of coast lines ('0.33' as default).
gridOn: bool value (True as default).
gridLabelOn: bool value (False as default).
parallels_kw: dict parameters used in the Basemap.drawparallels method.
parallels: parallels to be drawn (None as default).
parallels_color: color of parallels ('0.5' as default).
parallels_labels:[0,0,0,0] or [1,0,0,0].
meridians_kw: dict parameters used in the Basemap.drawmeridians method.
meridians: meridians to be drawn (None as default).
meridians_color: color of meridians (parallels_color as default).
meridians_labels: [0,0,0,0] or [0,0,0,1].
lonlatbox: None or (lon_start, lon_end, lat_start, lat_end).
lonlatbox_kw: dict parameters in the plot of lon-lat box.
lonlatbox_color:
General plot parameters:
-------------------------
ax: axis object, default is plt.gca()
plot_type: a string of plot type from ('pcolor', 'pcolormesh', 'imshow', 'contourf', 'contour', 'quiver', 'scatter') or None(default).
cmap: pyplot colormap.
clim: a tuple of colormap limit.
levels: sequence, int or None (default=None)
plot_kw: dict parameters used in the plot functions.
Pcolor/Pcolormesh related parameters:
-------------------------------
rasterized: bool (default is True).
Imshow related parameters
---------------------------
origin: 'lower' or 'upper'.
extent: horizontal range.
interpolation: 'nearest' (default) or 'bilinear' or 'cubic'.
Contourf related parameters:
-------------------------------
extend: 'both'(default).
Contour related parameters:
----------------------------
label_contour: False(default) or True.
Whether to label contours or not.
colors: contour color (default is 'gray').
Quiver plot related parameters:
--------------------------------
stride: stride along lon and lat.
stride_lon: stride along lon.
stride_lat: stride along lat.
quiver_scale: quiver scale.
quiver_color: quiver color.
sparse_polar_grids: bool, default is True.
quiver_kw: dict parameters used in the plt.quiver function.
hide_qkey: bool value, whether to show the quiverkey plot.
qkey_X: X parameter in the plt.quiverkey function (default is 0.85).
qkey_Y: Y parameter in the plt.quiverkey function (default is 1.02).
qkey_U: U parameter in the plt.quiverkey function (default is 2).
qkey_label: label parameter in the plt.quiverkey function.
qkey_labelpos: labelpos parameter in the plt.quiverkey function.
qkey_kw: dict parameters used in the plt.quiverkey function.
Scatter related parameters:
------------------------------
scatter_data: None(default) or (lonvec, latvec).
Hatch plot related parameters:
----------------------------------
hatches: ['///'] is default.
Colorbar related parameters:
-------------------------------
hide_cbar: bool value, whether to show the colorbar.
cbar_type: 'vertical'(shorten as 'v') or 'horizontal' (shorten as 'h').
cbar_extend: extend parameter in the plt.colorbar function.
'neither' as default here.
cbar_size: default '2.5%' for vertical colorbar,
'5%' for horizontal colorbar.
cbar_pad: default 0.1 for vertical colorbar,
0.4 for horizontal colorbar.
cbar_kw: dict parameters used in the plt.colorbar function.
units: str
long_name: str
Returns
--------
basemap object if only basemap is plotted.
plot object if data is shown.
'''
# target axis
ax = kw.pop('ax', None)
if ax is not None:
plt.sca(ax)
if isinstance(data, xr.DataArray):
data_array = data.copy()
data = data_array.values
if np.any(np.isnan(data)):
data = ma.masked_invalid(data)
if lon is None:
try:
lonname = [
s for s in data_array.dims
if s in ('lon', 'longitude', 'X')
][0]
except IndexError:
lonname = [s for s in data_array.dims if 'lon' in s][0]
lon = data_array[lonname]
if lat is None:
try:
latname = [
s for s in data_array.dims if s in ('lat', 'latitude', 'Y')
][0]
except IndexError:
latname = [s for s in data_array.dims if 'lat' in s][0]
lat = data_array[latname]
# guess data name
try:
data_name = data_array.attrs['long_name']
except KeyError:
try:
data_name = data_array.name
except AttributeError:
data_name = ''
# guess data units
try:
data_units = data_array.attrs['units']
except KeyError:
data_units = ''
# copy the original data
elif data is not None:
try:
data = data.copy()
except:
try:
# data has two components(u,v)
data = (data[0].copy(), data[1].copy())
except:
pass
if lon is not None:
lon = lon.copy()
if lat is not None:
lat = lat.copy()
# #### basemap parameters
# basemap kw parameter
basemap_kw = kw.pop('basemap_kw', {})
# projection
# proj = kw.pop('proj', 'cyl')
# proj = kw.pop('proj', 'moll')
proj = kw.pop('proj', 'hammer')
proj = kw.pop('projection',
proj) # projection overrides the proj parameter
# short names for nplaea and splaea projections
if proj in ('npolar', 'polar', 'np'):
proj = 'nplaea'
elif proj in ('spolar', 'sp'):
proj = 'splaea'
# lon_0
lon_0 = kw.pop('lon_0', None)
if lon_0 is None:
if lon is not None:
if np.isclose(np.abs(lon[-1] - lon[0]), 360):
lon_0 = (lon[0] + lon[-2]) / 2.0
else:
lon_0 = (lon[0] + lon[-1]) / 2.0
else:
# dummy = np.linspace(0, 360, data.shape[1]+1)[0:-1]
# lon_0 = ( dummy[0] + dummy[-1] )/2.0
lon_0 = 180
# elif proj in ('moll', 'cyl', 'hammer', 'robin'):
# lon_0 = 180
# elif proj in ('ortho','npstere', 'nplaea', 'npaeqd', 'spstere', 'splaea', 'spaeqd'):
# lon_0 = 0
else: # lon_0 is specified
if lon is not None and proj in ('moll', 'cyl', 'hammer'):
# correct the lon_0 so that it is at the edge of a grid box
lon_0_data = (lon[0] + lon[-1]) / 2.0
d_lon = lon[1] - lon[0]
d_lon_0 = lon_0 - lon_0_data
lon_0 = float(int(d_lon_0 / d_lon)) * d_lon + lon_0_data
# lat_0
lat_0 = kw.pop('lat_0', None)
if lat_0 is None:
if lat is not None:
lat_0 = (lat[0] + lat[-1]) / 2.0
elif proj in ('ortho', ):
lat_0 = 45
# lonlatcorner = (llcrnrlon, urcrnrlon, llcrnrlat, urcrnrlat)
lonlatcorner = kw.pop('lonlatcorner', None)
if lonlatcorner is not None:
llcrnrlon = lonlatcorner[0]
urcrnrlon = lonlatcorner[1]
llcrnrlat = lonlatcorner[2]
urcrnrlat = lonlatcorner[3]
else:
llcrnrlon = None
urcrnrlon = None
llcrnrlat = None
urcrnrlat = None
llcrnrlon = basemap_kw.pop('llcrnrlon', llcrnrlon)
urcrnrlon = basemap_kw.pop('urcrnrlon', urcrnrlon)
llcrnrlat = basemap_kw.pop('llcrnrlat', llcrnrlat)
urcrnrlat = basemap_kw.pop('urcrnrlat', urcrnrlat)
if llcrnrlon is None and urcrnrlon is None and llcrnrlat is None and urcrnrlat is None:
lonlatcorner = None
else:
lonlatcorner = (llcrnrlon, urcrnrlon, llcrnrlat, urcrnrlat)
# boundinglat
boundinglat = kw.pop('boundinglat', None)
if boundinglat is None:
if proj in ('npstere', 'nplaea', 'npaeqd'):
boundinglat = 30
elif proj in ('spstere', 'splaea', 'spaeqd'):
boundinglat = -30
# basemap round: True or False
if proj in ('npstere', 'nplaea', 'npaeqd', 'spstere', 'splaea', 'spaeqd'):
basemap_round = kw.pop('basemap_round', True)
else:
basemap_round = kw.pop('basemap_round', False)
basemap_round = kw.pop('round', basemap_round)
# base map
proj = basemap_kw.pop('projection', proj)
lon_0 = basemap_kw.pop('lon_0', lon_0)
lat_0 = basemap_kw.pop('lat_0', lat_0)
boundinglat = basemap_kw.pop('boundinglat', boundinglat)
basemap_round = basemap_kw.pop('round', basemap_round)
m = Basemap(projection=proj,
lon_0=lon_0,
lat_0=lat_0,
boundinglat=boundinglat,
round=basemap_round,
llcrnrlon=llcrnrlon,
urcrnrlon=urcrnrlon,
llcrnrlat=llcrnrlat,
urcrnrlat=urcrnrlat,
**basemap_kw)
# fill continents or plot coast lines
fill_continents = kw.pop('fill_continents', False)
if fill_continents:
# use Basemap.fillcontinents method
continents_kw = kw.pop('continents_kw', {})
continents_color = kw.pop('continents_color', '0.33')
continents_color = continents_kw.pop('color', continents_color)
lake_color = kw.pop('lake_color', 'none')
lake_color = continents_kw.pop('lake_color', lake_color)
m.fillcontinents(color=continents_color,
lake_color=lake_color,
**continents_kw)
# else:
draw_coastlines = kw.pop('draw_coastlines', not fill_continents)
if draw_coastlines:
# use Basemap.drawcoastlines method
coastlines_kw = kw.pop('coastlines_kw', {})
coastlines_color = kw.pop('coastlines_color', '0.33')
coastlines_color = coastlines_kw.pop('color', coastlines_color)
coastlines_lw = kw.pop('coastlines_lw', 0.5)
coastlines_lw = coastlines_kw.pop('lw', coastlines_lw)
m.drawcoastlines(color=coastlines_color,
linewidth=coastlines_lw,
**coastlines_kw)
ocean_color = kw.pop('ocean_color', None)
if ocean_color is not None:
m.drawmapboundary(fill_color=ocean_color)
# parallels
gridOn = kw.pop('gridOn', True)
gridLabelOn = kw.pop('gridLabelOn', False)
parallels_kw = kw.pop('parallels_kw', {})
# parallels = kw.pop('parallels', np.arange(-90,91,30))
parallels = kw.pop('parallels', None)
parallels = parallels_kw.pop('parallels', parallels)
parallels_color = kw.pop('parallels_color', '0.5')
parallels_color = parallels_kw.pop('color', parallels_color)
parallels_lw = kw.pop('parallels_lw', 0.5)
parallels_lw = parallels_kw.pop('lw', parallels_lw)
parallels_labels = kw.pop('parallels_labels', None)
parallels_labels = parallels_kw.pop('labels', parallels_labels)
if parallels_labels is None:
if gridLabelOn:
parallels_labels = [1, 0, 0, 0]
else:
parallels_labels = [0, 0, 0, 0]
if parallels is not None:
m.drawparallels(parallels,
color=parallels_color,
labels=parallels_labels,
linewidth=parallels_lw,
**parallels_kw)
elif gridOn:
m.drawparallels(np.arange(-90, 91, 30),
color=parallels_color,
linewidth=parallels_lw,
labels=parallels_labels,
**parallels_kw)
# meridians
meridians_kw = kw.pop('meridians_kw', {})
# meridians = kw.pop('meridians', np.arange(-180, 360, 30))
meridians = kw.pop('meridians', None)
meridians = meridians_kw.pop('meridians', meridians)
meridians_color = kw.pop('meridians_color', parallels_color)
meridians_color = meridians_kw.pop('color', meridians_color)
meridians_lw = kw.pop('meridians_lw', parallels_lw)
meridians_lw = meridians_kw.pop('lw', meridians_lw)
meridians_labels = kw.pop('meridians_labels', None)
meridians_labels = meridians_kw.pop('labels', meridians_labels)
if meridians_labels is None:
if gridLabelOn:
if proj in ('npstere', 'nplaea', 'npaeqd', 'spstere', 'splaea',
'spaeqd'):
meridians_labels = [1, 1, 0, 0]
elif proj in ('cyl', ):
meridians_labels = [0, 0, 0, 1]
else:
meridians_labels = [0, 0, 0, 0]
print('Meridian are not labeled.')
else:
meridians_labels = [0, 0, 0, 0]
if meridians is not None:
m.drawmeridians(meridians,
color=meridians_color,
labels=meridians_labels,
linewidth=meridians_lw,
**meridians_kw)
elif gridOn:
m.drawmeridians(np.arange(0, 360, 30),
color=meridians_color,
labels=meridians_labels,
linewidth=meridians_lw,
**meridians_kw)
# lonlatbox
lonlatbox = kw.pop('lonlatbox', None)
if lonlatbox is not None:
lonlon = np.array([
np.linspace(lonlatbox[0], lonlatbox[1],
100), lonlatbox[1] * np.ones(100),
np.linspace(lonlatbox[1], lonlatbox[0], 100),
lonlatbox[0] * np.ones(100)
]).ravel()
latlat = np.array([
lonlatbox[2] * np.ones(100),
np.linspace(lonlatbox[2], lonlatbox[3], 100),
lonlatbox[3] * np.ones(100),
np.linspace(lonlatbox[3], lonlatbox[2], 100)
]).ravel()
lonlatbox_kw = kw.pop('lonlatbox_kw', {})
lonlatbox_color = kw.pop('lonlatbox_color', 'k')
lonlatbox_color = lonlatbox_kw.pop('color', lonlatbox_color)
m.plot(lonlon,
latlat,
latlon=True,
color=lonlatbox_color,
**lonlatbox_kw)
# scatter
scatter_data = kw.pop('scatter_data', None)
if scatter_data is not None:
# L = data.astype('bool')
# marker_color = kw.pop('marker_color', 'k')
# plot_obj = m.scatter(X[L], Y[L], color=marker_color, **kw)
lonvec, latvec = scatter_data
plot_obj = m.scatter(lonvec, latvec, **kw)
# #### stop here and return the map object if data is None
if data is None:
return m
# data prepare
input_data_have_two_components = isinstance(data, tuple) \
or isinstance(data, list)
if input_data_have_two_components:
# input data is (u,v) or [u, v] where u, v are ndarray and two components of a vector
assert len(
data) == 2, 'quiver data must contain only two componets u and v'
u = data[0].squeeze()
v = data[1].squeeze()
assert u.ndim == 2, 'u component data must be two dimensional'
assert v.ndim == 2, 'v component data must be two dimensional'
data = np.sqrt(u**2 + v**2) # calculate wind speed
else: # input data is a ndarray
data = data.squeeze()
assert data.ndim == 2, 'Input data must be two dimensional!'
# lon
if lon is None:
# lon = np.linspace(0, 360, data.shape[1]+1)[0:-1]
# lon_edge = np.hstack((
# lon[0]*2 - lon[1],
# lon,
# lon[-1]*2 - lon[-2]))
# lon_edge = ( lon_edge[:-1] + lon_edge[1:] )/2.0
lon_edge = np.linspace(lon_0 - 180, lon_0 + 180, data.shape[1] + 1)
lon = (lon_edge[:-1] + lon_edge[1:]) / 2.0
else: # lon is specified
if np.isclose(np.abs(lon[-1] - lon[0]), 360):
# first and last longitude point to the same location: remove the last longitude
lon = lon[:-1]
data = data[:, :-1]
if input_data_have_two_components:
u = u[:, :-1]
v = v[:, :-1]
if (not np.isclose(lon_0, (lon[0] + lon[-1]) / 2.0)
and proj in ('moll', 'cyl', 'hammer', 'robin')):
# lon_0 not at the center of lon, need to shift grid
lon_west_end = lon_0 - 180 + (
lon[1] - lon[0]) / 2.0 # longitude of west end
# make sure the longitude of west end within the lon
if lon_west_end < lon[0]:
lon_west_end += 360
elif lon_west_end > lon[-1]:
lon_west_end -= 360
data, lon_shift = shiftgrid(lon_west_end, data, lon, start=True)
if input_data_have_two_components:
u, lon_shift = shiftgrid(lon_west_end, u, lon, start=True)
v, lon_shift = shiftgrid(lon_west_end, v, lon, start=True)
lon = lon_shift
if lon[0] < -180:
lon += 360
elif lon[-1] >= 540:
lon -= 360
lon_hstack = np.hstack(
(2 * lon[0] - lon[1], lon, 2 * lon[-1] - lon[-2]))
lon_edge = (lon_hstack[:-1] + lon_hstack[1:]) / 2.0
# lat
if lat is None:
lat_edge = np.linspace(-90, 90, data.shape[0] + 1)
lat = (lat_edge[:-1] + lat_edge[1:]) / 2.0
else:
lat_hstack = np.hstack(
(2 * lat[0] - lat[1], lat, 2 * lat[-1] - lat[-2]))
lat_edge = (lat_hstack[:-1] + lat_hstack[1:]) / 2.0
lat_edge[lat_edge > 90] = 90
lat_edge[lat_edge < -90] = -90
Lon, Lat = np.meshgrid(lon, lat)
X, Y = m(Lon, Lat)
Lon_edge, Lat_edge = np.meshgrid(lon_edge, lat_edge)
X_edge, Y_edge = m(Lon_edge, Lat_edge)
# ###### plot parameters
# plot_type
plot_type = kw.pop('plot_type', None)
if plot_type is None:
if input_data_have_two_components:
plot_type = 'quiver'
# elif ( proj in ('cyl',)
# and lonlatcorner is None
# ):
# plot_type = 'imshow'
# print('plot_type **** imshow **** is used.')
elif proj in ('nplaea', 'splaea', 'ortho'):
# pcolormesh has a problem for these projections
plot_type = 'pcolor'
print('plot_type **** pcolor **** is used.')
else:
plot_type = 'pcolormesh'
print('plot_type **** pcolormesh **** is used.')
# cmap
cmap = kw.pop('cmap', None)
if cmap is None:
zz_max = data.max()
zz_min = data.min()
if zz_min >= 0:
try:
cmap = plt.get_cmap('viridis')
except:
cmap = plt.get_cmap('OrRd')
elif zz_max <= 0:
try:
cmap = plt.get_cmap('viridis')
except:
cmap = plt.get_cmap('Blues_r')
else:
cmap = plt.get_cmap('RdBu_r')
elif isinstance(cmap, str):
cmap = plt.get_cmap(cmap)
# clim parameters
clim = kw.pop('clim', None)
robust = kw.pop('robust', False)
if clim is None:
if isinstance(data, np.ma.core.MaskedArray):
data1d = data.compressed()
else:
data1d = data.ravel()
notNaNs = np.logical_not(np.isnan(data1d))
data1d = data1d[notNaNs]
if robust:
a = np.percentile(data1d, 2)
b = np.percentile(data1d, 98)
else:
a = data1d.min()
b = data1d.max()
if a * b < 0:
b = max(abs(a), abs(b))
a = -b
clim = a, b
# levels
levels = kw.pop('levels', None)
if levels is None:
if plot_type in ('contour', 'contourf', 'contourf+'):
a, b = clim
levels = np.linspace(a, b, 11)
elif isinstance(levels, int):
if plot_type in ('contour', 'contourf', 'contourf+'):
a, b = clim
levels = np.linspace(a, b, levels)
elif plot_type in ('pcolor', 'pcolormesh', 'imshow'):
cmap = plt.get_cmap(cmap.name, levels - 1)
else: # levels is a sequence
if plot_type in ('pcolor', 'pcolormesh', 'imshow'):
cmap = plt.get_cmap(cmap.name, len(levels) - 1)
clim = min(levels), max(levels)
# colorbar parameters
if plot_type in ('pcolor', 'pcolormesh', 'contourf', 'contourf+',
'imshow'):
cbar_type = kw.pop('cbar_type', 'vertical')
cbar_kw = kw.pop('cbar_kw', {})
cbar_extend = kw.pop('cbar_extend', 'neither')
cbar_extend = cbar_kw.pop('extend', cbar_extend)
hide_cbar = kw.pop('hide_cbar', False)
if cbar_type in ('v', 'vertical'):
cbar_size = kw.pop('cbar_size', '2.5%')
cbar_size = cbar_kw.pop('size', cbar_size)
cbar_pad = kw.pop('cbar_pad', 0.1)
cbar_pad = cbar_kw.pop('pad', cbar_pad)
cbar_position = 'right'
cbar_orientation = 'vertical'
elif cbar_type in ('h', 'horizontal'):
# cbar = hcolorbar(units=units)
cbar_size = kw.pop('cbar_size', '5%')
cbar_size = cbar_kw.pop('size', cbar_size)
cbar_pad = kw.pop('cbar_pad', 0.4)
cbar_pad = cbar_kw.pop('pad', cbar_pad)
cbar_position = 'bottom'
cbar_orientation = 'horizontal'
# units in colorbar
units = kw.pop('units', None)
if units is None:
try:
units = data_units # input data is a DataArray
except:
units = ''
# long_name in colorbar
long_name = kw.pop('long_name', None)
if long_name is None:
try:
long_name = data_name # if input data is a DataArray
if long_name is None:
long_name = ''
except:
long_name = ''
# ###### plot
# pcolor
if plot_type in ('pcolor', ):
rasterized = kw.pop('rasterized', True)
plot_obj = m.pcolor(X_edge,
Y_edge,
data,
cmap=cmap,
rasterized=rasterized,
**kw)
# pcolormesh
elif plot_type in ('pcolormesh', ):
rasterized = kw.pop('rasterized', True)
plot_obj = m.pcolormesh(X_edge,
Y_edge,
data,
cmap=cmap,
rasterized=rasterized,
**kw)
# imshow
elif plot_type in ('imshow', ):
if Y_edge[-1, 0] > Y_edge[0, 0]:
origin = kw.pop('origin', 'lower')
else:
origin = kw.pop('origin', 'upper')
extent = kw.pop(
'extent',
[X_edge[0, 0], X_edge[0, -1], Y_edge[0, 0], Y_edge[-1, 0]])
interpolation = kw.pop('interpolation', 'nearest')
plot_obj = m.imshow(data,
origin=origin,
cmap=cmap,
extent=extent,
interpolation=interpolation,
**kw)
# contourf
elif plot_type in ('contourf', ):
if proj in ('ortho','npstere', 'nplaea', 'npaeqd', 'spstere',
'splaea', 'spaeqd')\
and np.isclose(np.abs(lon_edge[-1]-lon_edge[0]), 360):
data, lon = addcyclic(data, lon)
Lon, Lat = np.meshgrid(lon, lat)
X, Y = m(Lon, Lat)
extend = kw.pop('extend', 'both')
plot_obj = m.contourf(X,
Y,
data,
extend=extend,
cmap=cmap,
levels=levels,
**kw)
# contour
elif plot_type in ('contour', ):
if proj in ('ortho','npstere', 'nplaea', 'npaeqd', 'spstere',
'splaea', 'spaeqd')\
and np.isclose(np.abs(lon_edge[-1]-lon_edge[0]), 360):
data, lon = addcyclic(data, lon)
Lon, Lat = np.meshgrid(lon, lat)
X, Y = m(Lon, Lat)
colors = kw.pop('colors', 'k')
if colors is not None:
cmap = None
alpha = kw.pop('alpha', 0.5)
plot_obj = m.contour(X,
Y,
data,
cmap=cmap,
colors=colors,
levels=levels,
alpha=alpha,
**kw)
label_contour = kw.pop('label_contour', False)
if label_contour:
plt.clabel(plot_obj, plot_obj.levels[::2], fmt='%.2G')
# contourf + contour
elif plot_type in ('contourf+', ):
if proj in ('ortho','npstere', 'nplaea', 'npaeqd', 'spstere',
'splaea', 'spaeqd')\
and np.isclose(np.abs(lon_edge[-1]-lon_edge[0]), 360):
data, lon = addcyclic(data, lon)
Lon, Lat = np.meshgrid(lon, lat)
X, Y = m(Lon, Lat)
extend = kw.pop('extend', 'both')
linewidths = kw.pop('linewidths', 1)
plot_obj = m.contourf(X,
Y,
data,
extend=extend,
cmap=cmap,
levels=levels,
**kw)
colors = kw.pop('colors', 'k')
if colors is not None:
cmap = None
alpha = kw.pop('alpha', 0.5)
m.contour(X,
Y,
data,
cmap=cmap,
colors=colors,
alpha=alpha,
levels=levels,
linewidths=linewidths,
**kw)
# quiverplot
elif plot_type in ('quiver', ):
nGrids = kw.pop('nGrids', 50)
nx = kw.pop('nx', nGrids)
ny = kw.pop('ny', nGrids)
stride = kw.pop('stride', 1)
stride_lon = kw.pop('stride_lon', stride)
stride_lat = kw.pop('stride_lat', stride)
print(stride, stride_lon, stride_lat)
if (stride != 1) or (stride_lon != 1) or (
stride_lat != 1
): # compatible with old api, with stride, stride_lon, stride_lat controling quiver grids. To be obsolete. Use nGrids, nx, ny instead
print('stride used')
lon_ = lon[::stride_lon] # subset of lon
lat_ = lat[::stride_lat]
u_ = u[::stride_lat, ::stride_lon]
v_ = v[::stride_lat, ::stride_lon]
# sparse polar area
sparse_polar_grids = kw.pop('sparse_polar_grids', True)
if sparse_polar_grids:
msk = np.empty(u_.shape) * np.nan
for i in range(lat_.size):
step = int(1. / np.cos(lat_[i] * np.pi / 180))
msk[i, 0::step] = 0
u_ += msk
v_ += msk
Lon_, Lat_ = np.meshgrid(lon_, lat_)
u_rot, v_rot, X_, Y_ = m.rotate_vector(u_,
v_,
Lon_,
Lat_,
returnxy=True)
else: # use nGrids, nx, ny to control the quiver grids
if lon.max() > 180:
u_, lon_ = shiftgrid(180., u, lon, start=False)
v_, lon_ = shiftgrid(180., v, lon, start=False)
else:
u_ = u
v_ = v
lon_ = lon
u_rot, v_rot, X_, Y_ = m.transform_vector(u_,
v_,
lon_,
lat,
nx,
ny,
returnxy=True)
quiver_color = kw.pop('quiver_color', 'g')
quiver_scale = kw.pop('quiver_scale', None)
hide_qkey = kw.pop('hide_qkey', False)
qkey_kw = kw.pop('qkey_kw', {})
qkey_X = kw.pop('qkey_X', 0.85)
qkey_X = qkey_kw.pop('X', qkey_X)
qkey_Y = kw.pop('qkey_Y', 1.02)
qkey_Y = qkey_kw.pop('Y', qkey_Y)
qkey_U = kw.pop('qkey_U', 2)
qkey_U = qkey_kw.pop('U', qkey_U)
qkey_label = kw.pop('qkey_label', '{:g} '.format(qkey_U) + units)
qkey_label = qkey_kw.pop('label', qkey_label)
qkey_labelpos = kw.pop('qkey_labelpos', 'W')
qkey_labelpos = qkey_kw.pop('labelpos', qkey_labelpos)
plot_obj = m.quiver(X_,
Y_,
u_rot,
v_rot,
color=quiver_color,
scale=quiver_scale,
**kw)
if not hide_qkey:
# quiverkey plot
plt.quiverkey(plot_obj,
qkey_X,
qkey_Y,
qkey_U,
label=qkey_label,
labelpos=qkey_labelpos,
**qkey_kw)
# hatch plot
elif plot_type in ('hatch', 'hatches'):
hatches = kw.pop('hatches', ['///'])
plot_obj = m.contourf(X,
Y,
data,
colors='none',
hatches=hatches,
extend='both',
**kw)
else:
print(
'Please choose a right plot_type from ("pcolor", "contourf", "contour")!'
)
# set clim
if plot_type in ('pcolor', 'pcolormesh', 'imshow'):
plt.clim(clim)
# plot colorbar
if plot_type in ('pcolor', 'pcolormesh', 'contourf', 'contourf+',
'imshow'):
ax_current = plt.gca()
divider = make_axes_locatable(ax_current)
cax = divider.append_axes(cbar_position, size=cbar_size, pad=cbar_pad)
cbar = plt.colorbar(plot_obj,
cax=cax,
extend=cbar_extend,
orientation=cbar_orientation,
**cbar_kw)
if cbar_type in ('v', 'vertical'):
# put the units on the top of the vertical colorbar
cbar.ax.xaxis.set_label_position('top')
cbar.ax.set_xlabel(units)
cbar.ax.set_ylabel(long_name)
elif cbar_type in ('h', 'horizontal'):
# cbar.ax.yaxis.set_label_position('right')
# cbar.ax.set_ylabel(units, rotation=0, ha='left', va='center')
if long_name == '' or units == '':
cbar.ax.set_xlabel('{}{}'.format(long_name, units))
else:
cbar.ax.set_xlabel('{} [{}]'.format(long_name, units))
# remove the colorbar to avoid repeated colorbars
if hide_cbar:
cbar.remove()
# set back the main axes as the current axes
plt.sca(ax_current)
return plot_obj
# Create a map for plotting
bm = Basemap(projection='tmerc',
lat_0=90.0,
lon_0=-100.0,
lat_ts=40.0,
llcrnrlon=-121,
llcrnrlat=24,
urcrnrlon=-65,
urcrnrlat=46,
resolution='l')
# Get U,V components from wind speed and direction
u, v = metpy.get_wind_components(speed, wdir)
# Rotate the vectors to be properly aligned in the map projection
u, v = bm.rotate_vector(u, v, lon, lat)
# Generate grid of x,y positions
lon_grid, lat_grid, x_grid, y_grid = bm.makegrid(130, 60, returnxy=True)
# Transform the obs to basemap space for gridding
obx, oby = bm(lon, lat)
# Perform analysis of height obs using Cressman weights
heights_oban = grid_data(height, x_grid, y_grid, obx, oby,
obans[which_oban][0], obans[which_oban][1])
heights_oban = maskoceans(lon_grid, lat_grid, heights_oban)
# Map plotting
contours = np.arange(5000., 5800., 60.0)
vmin=np.min(speed); vmax=np.max(speed) print(speed.max()) ; print(np.mean(speed)); print(np.std(speed)) #UN = uvel/speed ; VN = vvel/speed #print fh.variables['hi'].units #lon_0=-60.; lat_0=90. m = Basemap(projection='npstere',boundinglat=55,lon_0=340,resolution='l', round=False) xi, yi = m(lons, lats) print(xi[0, :]) print(xi.min(), xi.max()) urot, vrot = m.rotate_vector(uvel, vvel, lons, lats) #yy = np.arange(0, yi.shape[0], 10) #xx = np.arange(0, xi.shape[1], 10) #points = np.meshgrid(yy, xx) plt.figure(figsize=(12,12)) ax = plt.gca() #levs=np.arange(0.01,0.2,0.005) cs = m.pcolor(xi,yi,urot,vmin=None, vmax=None) stride = 10
urcrnrlon=urlon,
resolution='h')
else:
map = Basemap(projection='cyl',
llcrnrlat=lllat,
urcrnrlat=lllat,
llcrnrlon=lllon,
urcrnrlon=lllon,
resolution='h')
x, y = map(lons[:, :], lats[:, :])
ur = sample.variables['U10'][0]
vr = sample.variables['V10'][0]
ve = vr * cosalpha + ur * sinalpha
urot, vrot = map.rotate_vector(ue, ve, lons, lats)
# do the plotting
map.barbs(x,
y,
urot,
vrot,
color='blue',
label='Rotated to latlon and then to map - CORRECT')
# bad barbs
map.barbs(x, y, ue, ve, color='red', label='Rotated to latlon - insufficient')
plt.savefig('map')
#---- Open the real data and do the processing ----#
def main():
# HL_LABEL = "CRCM5_HL"
# NEMO_LABEL = "CRCM5_NEMO"
#
#
# sim_label_to_path = OrderedDict(
# [(HL_LABEL, "/RESCUE/skynet3_rech1/huziy/CNRCWP/C5/2016/2-year-runs/coupled-GL+stfl_oneway/Samples"),
# (NEMO_LABEL, "/HOME/huziy/skynet3_rech1/CNRCWP/C5/2016/2-year-runs/coupled-GL+stfl/Samples")]
# )
#
#
# # get a coord file ... (use pm* files, since they contain NEM1 variable)
# # Should be NEMO_LABEL, since the hostetler case does not calculate NEM? vars
# coord_file = ""
# found_coord_file = False
# for mdir in os.listdir(sim_label_to_path[NEMO_LABEL]):
#
# mdir_path = os.path.join(sim_label_to_path[NEMO_LABEL], mdir)
# if not os.path.isdir(mdir_path):
# continue
#
# for fn in os.listdir(mdir_path):
#
# if fn[:2] not in ["pm", ]:
# continue
#
# if fn[-9:-1] == "0" * 8:
# continue
#
# coord_file = os.path.join(mdir_path, fn)
# found_coord_file = True
#
# if found_coord_file:
# break
#
#
#
# bmp, lons, lats = nemo_hl_util.get_basemap_obj_and_coords_from_rpn_file(path=coord_file)
# xx, yy = bmp(lons, lats)
u = np.array([0.0, ])
v = np.array([0.5, ])
lon = -84
lat = 45
lon = np.array([lon, ])
lat = np.array([lat, ])
b = Basemap(lon_0=0)
urot, vrot = b.rotate_vector(u, v, lon, lat)
xx, yy = b(lon, lat)
b.quiver(xx, yy, urot, vrot, color="r")
b.quiver(xx, yy, u, v, color="g")
b.drawcoastlines()
plt.show()
def geoplot(data=None, lon=None, lat=None, **kw):
'''Show 2D data in a lon-lat plane.
Parameters
-----------
data: array of shape (n_lat, n_lon), or [u_array, v_array]-like for (u,v)
data or None(default) when only plotting the basemap.
lon: n_lon length vector or None(default).
lat: n_lat length vector or None(default).
kw: dict parameters related to basemap or plot functions.
Basemap related parameters:
----------------------------
basemap_kw: dict parameter in the initialization of a Basemap.
proj or projection: map projection name (default='moll')
popular projections: 'ortho', 'np'(='nplaea'), 'sp'(='splaea')
and other projections given from basemap.
lon_0: map center longitude (None as default).
lat_0: map center latitude (None as default).
lonlatcorner: (llcrnrlon, urcrnrlon, llcrnrlat, urcrnrlat).
boundinglat: latitude at the map out boundary (None as default).
basemap_round or round: True(default) or False.
fill_continents: bool value (False as default).
continents_kw: dict parameter used in the Basemap.fillcontinents method.
continents_color: color of continents ('0.5' as default).
lake_color: color of lakes ('none' as default).
coastlines_kw: dict parameter used in the Basemap.drawcoastlines method.
coastlines_color: color of coast lines ('0.66' as default).
parallels_kw: dict parameters used in the Basemap.drawparallels method.
parallels: parallels to be drawn (None as default).
parallels_color: color of parallels ('0.75' as default).
parallels_labels:[0,0,0,0] or [1,0,0,0].
meridians_kw: dict parameters used in the Basemap.drawmeridians method.
meridians: meridians to be drawn (None as default).
meridians_color: color of meridians ('0.75' as default).
meridians_labels: [0,0,0,0] or [0,0,0,1].
lonlatbox: None or (lon_start, lon_end, lat_start, lat_end).
lonlatbox_kw: dict parameters in the plot of lon-lat box.
lonlatbox_color:
General plot parameters:
-------------------------
ax: axis object, default is plt.gca()
plot_type: a string of plot type from ('pcolor', 'pcolormesh', 'imshow', 'contourf', 'contour', 'quiver', 'scatter') or None(default).
cmap: pyplot colormap.
clim: a tuple of colormap limit.
levels: sequence, int or None (default=None)
plot_kw: dict parameters used in the plot functions.
Pcolor/Pcolormesh related parameters:
-------------------------------
rasterized: bool (default is True).
Imshow related parameters
---------------------------
origin: 'lower' or 'upper'.
extent: horizontal range.
interpolation: 'nearest' (default) or 'bilinear' or 'cubic'.
Contourf related parameters:
-------------------------------
extend: 'both'(default).
Contour related parameters:
----------------------------
label_contour: False(default) or True.
Whether to label contours or not.
colors: contour color (default is 'gray').
Quiver plot related parameters:
--------------------------------
stride: stride along lon and lat.
stride_lon: stride along lon.
stride_lat: stride along lat.
quiver_scale: quiver scale.
quiver_color: quiver color.
quiver_kw: dict parameters used in the plt.quiver function.
hide_qkey: bool value, whether to show the quiverkey plot.
qkey_X: X parameter in the plt.quiverkey function (default is 0.85).
qkey_Y: Y parameter in the plt.quiverkey function (default is 1.02).
qkey_U: U parameter in the plt.quiverkey function (default is 2).
qkey_label: label parameter in the plt.quiverkey function.
qkey_labelpos: labelpos parameter in the plt.quiverkey function.
qkey_kw: dict parameters used in the plt.quiverkey function.
Scatter related parameters:
------------------------------
scatter_data: None(default) or (lonvec, latvec).
Hatch plot related parameters:
----------------------------------
hatches: ['///'] is default.
Colorbar related parameters:
-------------------------------
hide_cbar: bool value, whether to show the colorbar.
cbar_type: 'vertical'(shorten as 'v') or 'horizontal' (shorten as 'h').
cbar_extend: extend parameter in the plt.colorbar function.
'neither' as default here.
cbar_size: default '2.5%' for vertical colorbar,
'5%' for horizontal colorbar.
cbar_pad: default 0.1 for vertical colorbar,
0.4 for horizontal colorbar.
cbar_kw: dict parameters used in the plt.colorbar function.
units: str
long_name: str
Returns
--------
basemap object if only basemap is plotted.
plot object if data is shown.
'''
# target axis
ax = kw.pop('ax', None)
if ax is not None:
plt.sca(ax)
if isinstance(data, xr.DataArray):
data_array = data.copy()
data = data_array.values
if np.any(np.isnan(data)):
data = ma.masked_invalid(data)
if lon is None:
try:
lonname = [s for s in data_array.dims
if s in ('lon', 'longitude', 'X')][0]
except IndexError:
lonname = [s for s in data_array.dims
if 'lon' in s][0]
lon = data_array[lonname]
if lat is None:
try:
latname = [s for s in data_array.dims
if s in ('lat', 'latitude', 'Y')][0]
except IndexError:
latname = [s for s in data_array.dims
if 'lat' in s][0]
lat = data_array[latname]
# guess data name
try:
data_name = data_array.attrs['long_name']
except KeyError:
try:
data_name = data_array.name
except AttributeError:
data_name = ''
# guess data units
try:
data_units = data_array.attrs['units']
except KeyError:
data_units = ''
# copy the original data
if data is not None and hasattr(data, 'copy'):
data = data.copy()
if lon is not None:
lon = lon.copy()
if lat is not None:
lat = lat.copy()
# #### basemap parameters
# basemap kw parameter
basemap_kw = kw.pop('basemap_kw', {})
# projection
# proj = kw.pop('proj', 'cyl')
proj = kw.pop('proj', 'moll')
proj = kw.pop('projection', proj) # projection overrides the proj parameter
# short names for nplaea and splaea projections
if proj in ('npolar', 'polar', 'np'):
proj = 'nplaea'
elif proj in ('spolar', 'sp'):
proj = 'splaea'
# lon_0
lon_0 = kw.pop('lon_0', None)
if lon_0 is None:
if lon is not None:
if np.isclose( np.abs(lon[-1] - lon[0]), 360 ):
lon_0 = (lon[0] + lon[-2])/2.0
else:
lon_0 = (lon[0] + lon[-1])/2.0
else:
# dummy = np.linspace(0, 360, data.shape[1]+1)[0:-1]
# lon_0 = ( dummy[0] + dummy[-1] )/2.0
lon_0 = 180
# elif proj in ('moll', 'cyl', 'hammer', 'robin'):
# lon_0 = 180
# elif proj in ('ortho','npstere', 'nplaea', 'npaeqd', 'spstere', 'splaea', 'spaeqd'):
# lon_0 = 0
else: # lon_0 is specified
if lon is not None and proj in ('moll', 'cyl'):
# correct the lon_0 so that it is at the edge of a grid box
lon_0_data = (lon[0] + lon[-1])/2.0
d_lon = lon[1] - lon[0]
d_lon_0 = lon_0 - lon_0_data
lon_0 = float(int(d_lon_0 / d_lon)) * d_lon + lon_0_data
# lat_0
lat_0 = kw.pop('lat_0', None)
if lat_0 is None:
if lat is not None:
lat_0 = ( lat[0] + lat[-1] )/2.0
elif proj in ('ortho',):
lat_0 = 45
# lonlatcorner = (llcrnrlon, urcrnrlon, llcrnrlat, urcrnrlat)
lonlatcorner = kw.pop('lonlatcorner', None)
if lonlatcorner is not None:
llcrnrlon = lonlatcorner[0]
urcrnrlon = lonlatcorner[1]
llcrnrlat = lonlatcorner[2]
urcrnrlat = lonlatcorner[3]
else:
llcrnrlon = None
urcrnrlon = None
llcrnrlat = None
urcrnrlat = None
llcrnrlon = basemap_kw.pop('llcrnrlon', llcrnrlon)
urcrnrlon = basemap_kw.pop('urcrnrlon', urcrnrlon)
llcrnrlat = basemap_kw.pop('llcrnrlat', llcrnrlat)
urcrnrlat = basemap_kw.pop('urcrnrlat', urcrnrlat)
if llcrnrlon is None and urcrnrlon is None and llcrnrlat is None and urcrnrlat is None:
lonlatcorner = None
else:
lonlatcorner = (llcrnrlon, urcrnrlon, llcrnrlat, urcrnrlat)
# boundinglat
boundinglat = kw.pop('boundinglat', None)
if boundinglat is None:
if proj in ('npstere', 'nplaea', 'npaeqd'):
boundinglat = 30
elif proj in ('spstere', 'splaea', 'spaeqd'):
boundinglat = -30
# basemap round: True or False
if proj in ('npstere', 'nplaea', 'npaeqd', 'spstere', 'splaea', 'spaeqd'):
basemap_round = kw.pop('basemap_round', True)
else:
basemap_round = kw.pop('basemap_round', False)
basemap_round = kw.pop('round', basemap_round)
# base map
proj = basemap_kw.pop('projection', proj)
lon_0 = basemap_kw.pop('lon_0', lon_0)
lat_0 = basemap_kw.pop('lat_0', lat_0)
boundinglat = basemap_kw.pop('boundinglat', boundinglat)
basemap_round = basemap_kw.pop('round', basemap_round)
m = Basemap(projection=proj, lon_0=lon_0, lat_0=lat_0, boundinglat=boundinglat,
round=basemap_round,
llcrnrlon=llcrnrlon, urcrnrlon=urcrnrlon,
llcrnrlat=llcrnrlat, urcrnrlat=urcrnrlat,
**basemap_kw)
# fill continents or plot coast lines
fill_continents = kw.pop('fill_continents', False)
if fill_continents:
# use Basemap.fillcontinents method
continents_kw = kw.pop('continents_kw', {})
continents_color = kw.pop('continents_color', '0.5')
continents_color = continents_kw.pop('color', continents_color)
lake_color = kw.pop('lake_color', 'none')
lake_color = continents_kw.pop('lake_color', lake_color)
m.fillcontinents(color=continents_color, lake_color=lake_color,
**continents_kw)
else:
# use Basemap.drawcoastlines method
coastlines_kw = kw.pop('coastlines_kw', {})
coastlines_color = kw.pop('coastlines_color', '0.66')
coastlines_color = coastlines_kw.pop('color', coastlines_color)
m.drawcoastlines(color=coastlines_color, linewidth=0.5)
# parallels
gridon = kw.pop('gridon', False)
parallels_kw = kw.pop('parallels_kw', {})
# parallels = kw.pop('parallels', np.arange(-90,91,30))
parallels = kw.pop('parallels', None)
parallels = parallels_kw.pop('parallels', parallels)
parallels_color = kw.pop('parallels_color', '0.75')
parallels_color = parallels_kw.pop('color', parallels_color)
if proj in ('cyl',):
parallels_labels = kw.pop('parallels_labels', [1,0,0,0])
else:
parallels_labels = kw.pop('parallels_labels', [0,0,0,0])
parallels_labels = parallels_kw.pop('labels', parallels_labels)
if parallels is not None:
m.drawparallels(parallels, color=parallels_color,
labels=parallels_labels, linewidth=1.0, **parallels_kw)
elif gridon:
m.drawparallels(np.arange(-90, 91, 30), color=parallels_color,
labels=parallels_labels, linewidth=1.0, **parallels_kw)
# meridians
meridians_kw = kw.pop('meridians_kw', {})
# meridians = kw.pop('meridians', np.arange(-180, 360, 30))
meridians = kw.pop('meridians', None)
meridians = meridians_kw.pop('meridians', meridians)
meridians_color = kw.pop('meridians_color', parallels_color)
meridians_color = meridians_kw.pop('color', meridians_color)
if proj in ('cyl',):
meridians_labels = kw.pop('meridians_labels', [0,0,0,1])
# elif proj in ('npstere', 'nplaea', 'npaeqd', 'spstere', 'splaea', 'spaeqd'):
# meridians_labels = kw.pop('meridians_labels', [1,0,0,0])
else:
meridians_labels = kw.pop('meridians_labels', [0,0,0,0])
meridians_labels = meridians_kw.pop('labels', meridians_labels)
if meridians is not None:
m.drawmeridians(meridians, color=meridians_color,
labels=meridians_labels, linewidth=1.0, **meridians_kw)
elif gridon:
m.drawmeridians(np.arange(-180, 360, 30), color=meridians_color,
labels=meridians_labels, linewidth=1.0, **meridians_kw)
# lonlatbox
lonlatbox = kw.pop('lonlatbox', None)
if lonlatbox is not None:
lonlon = np.array([
np.linspace(lonlatbox[0], lonlatbox[1], 100),
lonlatbox[1]*np.ones(100),
np.linspace(lonlatbox[1], lonlatbox[0], 100),
lonlatbox[0]*np.ones(100)
]).ravel()
latlat = np.array([
lonlatbox[2]*np.ones(100),
np.linspace(lonlatbox[2], lonlatbox[3], 100),
lonlatbox[3]*np.ones(100),
np.linspace(lonlatbox[3], lonlatbox[2], 100)
]).ravel()
lonlatbox_kw = kw.pop('lonlatbox_kw', {})
lonlatbox_color = kw.pop('lonlatbox_color', 'k')
lonlatbox_color = lonlatbox_kw.pop('color', lonlatbox_color)
m.plot(lonlon, latlat, latlon=True, color=lonlatbox_color, **lonlatbox_kw)
# scatter
scatter_data = kw.pop('scatter_data', None)
if scatter_data is not None:
# L = data.astype('bool')
# marker_color = kw.pop('marker_color', 'k')
# plot_obj = m.scatter(X[L], Y[L], color=marker_color, **kw)
lonvec, latvec = scatter_data
plot_obj = m.scatter(lonvec, latvec, **kw)
# #### stop here and return the map object if data is None
if data is None:
return m
# data prepare
input_data_have_two_components = isinstance(data, tuple) \
or isinstance(data, list)
if input_data_have_two_components:
# input data is (u,v) or [u, v] where u, v are ndarray and two components of a vector
assert len(data) == 2,'quiver data must contain only two componets u and v'
u = data[0].squeeze()
v = data[1].squeeze()
assert u.ndim == 2, 'u component data must be two dimensional'
assert v.ndim == 2, 'v component data must be two dimensional'
data = np.sqrt( u**2 + v**2 ) # calculate wind speed
else:# input data is a ndarray
data = data.squeeze()
assert data.ndim == 2, 'Input data must be two dimensional!'
# lon
if lon is None:
# lon = np.linspace(0, 360, data.shape[1]+1)[0:-1]
# lon_edge = np.hstack((
# lon[0]*2 - lon[1],
# lon,
# lon[-1]*2 - lon[-2]))
# lon_edge = ( lon_edge[:-1] + lon_edge[1:] )/2.0
lon_edge = np.linspace(lon_0-180, lon_0+180, data.shape[1]+1)
lon = (lon_edge[:-1] + lon_edge[1:])/2.0
else:# lon is specified
if np.isclose( np.abs(lon[-1] - lon[0]), 360 ):
# first and last longitude point to the same location: remove the last longitude
lon = lon[:-1]
data = data[:, :-1]
if input_data_have_two_components:
u = u[:, :-1]
v = v[:, :-1]
if (not np.isclose(lon_0, (lon[0] + lon[-1])/2.0)
and proj in ('moll', 'cyl', 'hammer', 'robin')
):
# lon_0 not at the center of lon, need to shift grid
lon_west_end = lon_0 - 180 + (lon[1] - lon[0])/2.0 # longitude of west end
# make sure the longitude of west end within the lon
if lon_west_end < lon[0]:
lon_west_end += 360
elif lon_west_end > lon[-1]:
lon_west_end -= 360
data, lon_shift = shiftgrid(lon_west_end, data, lon, start=True)
if input_data_have_two_components:
u, lon_shift = shiftgrid(lon_west_end, u, lon, start=True)
v, lon_shift = shiftgrid(lon_west_end, v, lon, start=True)
lon = lon_shift
if lon[0]<-180:
lon += 360
elif lon[-1]>=540:
lon -= 360
lon_hstack = np.hstack((2*lon[0] - lon[1], lon, 2*lon[-1] - lon[-2]))
lon_edge = (lon_hstack[:-1] + lon_hstack[1:])/2.0
# lat
if lat is None:
lat_edge = np.linspace(-90, 90, data.shape[0]+1)
lat = (lat_edge[:-1] + lat_edge[1:])/2.0
else:
lat_hstack = np.hstack((2*lat[0] - lat[1], lat, 2*lat[-1] - lat[-2]))
lat_edge = (lat_hstack[:-1] + lat_hstack[1:])/2.0
lat_edge[lat_edge>90] = 90
lat_edge[lat_edge<-90] = -90
Lon, Lat = np.meshgrid(lon, lat)
X, Y = m(Lon, Lat)
Lon_edge, Lat_edge = np.meshgrid(lon_edge, lat_edge)
X_edge, Y_edge = m(Lon_edge, Lat_edge)
# ###### plot parameters
# plot_type
plot_type = kw.pop('plot_type', None)
if plot_type is None:
if input_data_have_two_components:
plot_type = 'quiver'
# elif ( proj in ('cyl',)
# and lonlatcorner is None
# ):
# plot_type = 'imshow'
# print('plot_type **** imshow **** is used.')
elif proj in ('nplaea', 'splaea', 'ortho'):
# pcolormesh has a problem for these projections
plot_type = 'pcolor'
print('plot_type **** pcolor **** is used.')
else:
plot_type = 'pcolormesh'
print ('plot_type **** pcolormesh **** is used.')
# cmap
cmap = kw.pop('cmap', None)
if cmap is None:
zz_max = data.max()
zz_min = data.min()
if zz_min >=0:
try:
cmap = plt.get_cmap('viridis')
except:
cmap = plt.get_cmap('OrRd')
elif zz_max<=0:
try:
cmap = plt.get_cmap('viridis')
except:
cmap = plt.get_cmap('Blues_r')
else:
cmap = plt.get_cmap('RdBu_r')
elif isinstance(cmap, str):
cmap = plt.get_cmap(cmap)
# clim parameters
clim = kw.pop('clim', None)
robust = kw.pop('robust', False)
if clim is None:
if isinstance(data,np.ma.core.MaskedArray):
data1d = data.compressed()
else:
data1d = data.ravel()
notNaNs = np.logical_not(np.isnan(data1d))
data1d = data1d[notNaNs]
if robust:
a = np.percentile(data1d,2)
b = np.percentile(data1d,98)
else:
a = data1d.min()
b = data1d.max()
if a * b < 0:
b = max(abs(a), abs(b))
a = -b
clim = a, b
# levels
levels = kw.pop('levels', None)
if levels is None:
if plot_type in ('contour', 'contourf', 'contourf+'):
a, b = clim
levels = np.linspace(a, b, 11)
elif isinstance(levels, int):
if plot_type in ('contour', 'contourf', 'contourf+'):
a, b = clim
levels = np.linspace(a, b, levels)
elif plot_type in ('pcolor', 'pcolormesh', 'imshow'):
cmap = plt.get_cmap(cmap.name, levels-1)
else: # levels is a sequence
if plot_type in ('pcolor', 'pcolormesh', 'imshow'):
cmap = plt.get_cmap(cmap.name, len(levels)-1)
clim = min(levels), max(levels)
# colorbar parameters
if plot_type in ('pcolor', 'pcolormesh', 'contourf', 'contourf+',
'imshow'):
cbar_type = kw.pop('cbar_type', 'vertical')
cbar_kw = kw.pop('cbar_kw', {})
cbar_extend = kw.pop('cbar_extend', 'neither')
cbar_extend = cbar_kw.pop('extend', cbar_extend)
hide_cbar = kw.pop('hide_cbar', False)
if cbar_type in ('v', 'vertical'):
cbar_size = kw.pop('cbar_size', '2.5%')
cbar_size = cbar_kw.pop('size', cbar_size)
cbar_pad = kw.pop('cbar_pad', 0.1)
cbar_pad = cbar_kw.pop('pad', cbar_pad)
cbar_position = 'right'
cbar_orientation = 'vertical'
elif cbar_type in ('h', 'horizontal'):
# cbar = hcolorbar(units=units)
cbar_size = kw.pop('cbar_size', '5%')
cbar_size = cbar_kw.pop('size', cbar_size)
cbar_pad = kw.pop('cbar_pad', 0.4)
cbar_pad = cbar_kw.pop('pad', cbar_pad)
cbar_position = 'bottom'
cbar_orientation = 'horizontal'
# units in colorbar
units = kw.pop('units', None)
if units is None:
try:
units = data_units # input data is a DataArray
except:
units = ''
# long_name in colorbar
long_name = kw.pop('long_name', None)
if long_name is None:
try:
long_name = data_name # if input data is a DataArray
if long_name is None:
long_name = ''
except:
long_name = ''
# ###### plot
# pcolor
if plot_type in ('pcolor',):
rasterized = kw.pop('rasterized', True)
plot_obj = m.pcolor(X_edge, Y_edge, data, cmap=cmap,
rasterized=rasterized, **kw)
# pcolormesh
elif plot_type in ('pcolormesh',):
rasterized = kw.pop('rasterized', True)
plot_obj = m.pcolormesh(X_edge, Y_edge, data, cmap=cmap,
rasterized=rasterized, **kw)
# imshow
elif plot_type in ('imshow',):
if Y_edge[-1,0] > Y_edge[0,0]:
origin = kw.pop('origin', 'lower')
else:
origin = kw.pop('origin', 'upper')
extent = kw.pop('extent', [X_edge[0,0], X_edge[0,-1], Y_edge[0,0], Y_edge[-1,0]])
interpolation = kw.pop('interpolation', 'nearest')
plot_obj = m.imshow(data, origin=origin, cmap=cmap, extent=extent,
interpolation=interpolation, **kw)
# contourf
elif plot_type in ('contourf',):
if proj in ('ortho','npstere', 'nplaea', 'npaeqd', 'spstere',
'splaea', 'spaeqd')\
and np.isclose(np.abs(lon_edge[-1]-lon_edge[0]), 360):
data, lon = addcyclic(data, lon)
Lon, Lat = np.meshgrid(lon,lat)
X, Y = m(Lon, Lat)
extend = kw.pop('extend', 'both')
plot_obj = m.contourf(X, Y, data, extend=extend, cmap=cmap,
levels=levels, **kw)
# contour
elif plot_type in ('contour',):
if proj in ('ortho','npstere', 'nplaea', 'npaeqd', 'spstere',
'splaea', 'spaeqd')\
and np.isclose(np.abs(lon_edge[-1]-lon_edge[0]), 360):
data, lon = addcyclic(data, lon)
Lon, Lat = np.meshgrid(lon,lat)
X, Y = m(Lon, Lat)
colors = kw.pop('colors', 'k')
if colors is not None:
cmap = None
alpha = kw.pop('alpha', 0.5)
plot_obj = m.contour(X, Y, data, cmap=cmap, colors=colors,
levels=levels, alpha=alpha, **kw)
label_contour = kw.pop('label_contour', False)
if label_contour:
plt.clabel(plot_obj,plot_obj.levels[::2],fmt='%.2G')
# contourf + contour
elif plot_type in ('contourf+',):
if proj in ('ortho','npstere', 'nplaea', 'npaeqd', 'spstere',
'splaea', 'spaeqd')\
and np.isclose(np.abs(lon_edge[-1]-lon_edge[0]), 360):
data, lon = addcyclic(data, lon)
Lon, Lat = np.meshgrid(lon,lat)
X, Y = m(Lon, Lat)
extend = kw.pop('extend', 'both')
linewidths = kw.pop('linewidths', 1)
plot_obj = m.contourf(X, Y, data, extend=extend, cmap=cmap,
levels=levels, **kw)
colors = kw.pop('colors', 'k')
if colors is not None:
cmap = None
alpha = kw.pop('alpha', 0.5)
m.contour(X, Y, data, cmap=cmap, colors=colors, alpha=alpha,
levels=levels, linewidths=linewidths, **kw)
# quiverplot
elif plot_type in ('quiver',):
stride = kw.pop('stride', 1)
stride_lon = kw.pop('stride_lon', stride)
stride_lat = kw.pop('stride_lat', stride)
lon_ = lon[::stride_lon] # subset of lon
lat_ = lat[::stride_lat]
u_ = u[::stride_lat, ::stride_lon]
v_ = v[::stride_lat, ::stride_lon]
Lon_, Lat_ = np.meshgrid(lon_, lat_)
u_rot, v_rot, X_, Y_ = m.rotate_vector(
u_, v_, Lon_, Lat_, returnxy=True
)
quiver_color = kw.pop('quiver_color', 'g')
quiver_scale = kw.pop('quiver_scale', None)
hide_qkey = kw.pop('hide_qkey', False)
qkey_kw = kw.pop('qkey_kw', {})
qkey_X = kw.pop('qkey_X', 0.85)
qkey_X = qkey_kw.pop('X', qkey_X)
qkey_Y = kw.pop('qkey_Y', 1.02)
qkey_Y = qkey_kw.pop('Y', qkey_Y)
qkey_U = kw.pop('qkey_U', 2)
qkey_U = qkey_kw.pop('U', qkey_U)
qkey_label = kw.pop('qkey_label', '{:g} '.format(qkey_U) + units)
qkey_label = qkey_kw.pop('label', qkey_label)
qkey_labelpos = kw.pop('qkey_labelpos', 'W')
qkey_labelpos = qkey_kw.pop('labelpos', qkey_labelpos)
plot_obj = m.quiver(X_, Y_, u_rot, v_rot, color=quiver_color,
scale=quiver_scale, **kw)
if not hide_qkey:
# quiverkey plot
plt.quiverkey(plot_obj, qkey_X, qkey_Y, qkey_U,
label=qkey_label, labelpos=qkey_labelpos, **qkey_kw)
# hatch plot
elif plot_type in ('hatch', 'hatches'):
hatches = kw.pop('hatches', ['///'])
plot_obj = m.contourf(X, Y, data, colors='none', hatches=hatches,
extend='both', **kw)
else:
print('Please choose a right plot_type from ("pcolor", "contourf", "contour")!')
# set clim
if plot_type in ('pcolor', 'pcolormesh', 'imshow'):
plt.clim(clim)
# plot colorbar
if plot_type in ('pcolor', 'pcolormesh', 'contourf', 'contourf+',
'imshow'):
ax_current = plt.gca()
divider = make_axes_locatable(ax_current)
cax = divider.append_axes(cbar_position, size=cbar_size, pad=cbar_pad)
cbar = plt.colorbar(plot_obj, cax=cax, extend=cbar_extend,
orientation=cbar_orientation, **cbar_kw)
if cbar_type in ('v', 'vertical'):
# put the units on the top of the vertical colorbar
cbar.ax.xaxis.set_label_position('top')
cbar.ax.set_xlabel(units)
cbar.ax.set_ylabel(long_name)
elif cbar_type in ('h', 'horizontal'):
# cbar.ax.yaxis.set_label_position('right')
# cbar.ax.set_ylabel(units, rotation=0, ha='left', va='center')
if long_name == '' or units =='':
cbar.ax.set_xlabel('{}{}'.format(long_name, units))
else:
cbar.ax.set_xlabel('{} [{}]'.format(long_name, units))
# remove the colorbar to avoid repeated colorbars
if hide_cbar:
cbar.remove()
# set back the main axes as the current axes
plt.sca(ax_current)
return plot_obj
def _make_EICS_plots(dtime=None,
vplot_sized=False,
contour_den=8,
s_loc=False,
quiver_scale=30):
"""
@Parameter: dtime input as a string
@Parameter: s_loc input as a bool, which means the locations of the virtual stations.
"""
dtype = 'EICS'
if not os.path.exists(CONFIG['plots_dir']):
os.makedirs(CONFIG['plots_dir'])
dtime_range = [dtime, dtime]
pathformat_prefix = dtype + '/%Y/%m/'
pathformat_unzipped = pathformat_prefix + '%d/' + dtype + '%Y%m%d_%H%M%S.dat'
filename_unzipped = dailynames(file_format=pathformat_unzipped,
trange=dtime_range,
res=10)
out_files_unzipped = [
CONFIG['local_data_dir'] + rf_res for rf_res in filename_unzipped
]
Data_Days_time = read_data_files(out_files=out_files_unzipped,
dtype=dtype,
out_type='df')
J_comp = Data_Days_time['Jy']
Jc_max, Jc_min = J_comp.max(), J_comp.min()
Jcm_abs = max(abs(Jc_max), abs(Jc_min))
contour_density = np.linspace(-Jcm_abs, Jcm_abs, num=contour_den)
tp = dtime
datetime_tp = tp[0:4] + tp[5:7] + tp[8:10] + '_' + tp[11:13] + tp[
14:16] + tp[17:19]
lon = Data_Days_time['longitude']
lat = Data_Days_time['latitude']
Jx = Data_Days_time['Jx'] # Note: positive is Northward
Jy = Data_Days_time['Jy'] # Note: positive is Eastward
# plot 1:
# plot map ground (North hemisphere)
fig1 = plt.figure(figsize=(8, 8))
ax1 = plt.gca()
m = Basemap(projection='lcc',
resolution='c',
width=8E6,
height=8E6,
lat_0=60,
lon_0=-100)
# draw coastlines, country boundaries, fill continents.
m.drawcoastlines(linewidth=0.25)
m.drawcountries(linewidth=0.25)
m.fillcontinents(color='None', lake_color='None')
# draw the edge of the map projection region (the projection limb)
m.drawmapboundary(fill_color=None)
# m.drawgreatcircle(-100,0,0,90)
m.drawlsmask()
# m.bluemarble()
m.shadedrelief()
# draw parallels and meridians.
# label parallels on right and top
# meridians on bottom and left
parallels = np.arange(0., 81, 10.)
# labels = [left,right,top,bottom]
m.drawparallels(parallels, labels=[False, True, True, False])
meridians = np.arange(10., 351., 20.)
m.drawmeridians(meridians, labels=[True, False, False, True])
date_nightshade = datetime.strptime(dtime, '%Y-%m-%d/%H:%M:%S')
m.nightshade(date=date_nightshade)
draw_map(m)
# plot vector field:
lon = lon.to_numpy()
lat = lat.to_numpy()
Jx = Jx.to_numpy() # Note: positive is Northward
Jy = Jy.to_numpy() # Note: positive is Eastward
Jx_uni = Jx / np.sqrt(Jx**2 + Jy**2)
Jy_uni = Jy / np.sqrt(Jx**2 + Jy**2)
n = -2
color = np.sqrt(((Jx_uni - n) / 2)**2 + ((Jy_uni - n) / 2)**2)
if vplot_sized == False:
qv = m.quiver(
lon,
lat,
Jx_uni,
Jy_uni,
color,
headlength=7,
latlon=True,
cmap='GnBu') # autumn_r #, color=cm(norm(o)))#, cmap = 'jet')
plt.colorbar()
else:
Jy_rot, Jx_rot, x, y = m.rotate_vector(Jy, Jx, lon, lat, returnxy=True)
qv = m.quiver(lon,
lat,
Jy_rot,
Jx_rot,
headlength=7,
latlon=True,
scale_units='dots',
scale=quiver_scale) # , transform='lcc')
qk = ax1.quiverkey(qv,
0.3,
-0.1,
100,
r'$100 \ mA/m$',
labelpos='E',
coordinates='data') # figure
plt.title(label='EICS ' + tp, fontsize=20, color="black", pad=20)
plt.tight_layout()
plt.savefig(CONFIG['plots_dir'] + 'EICS' + '_vector_' +
date_nightshade.strftime('%Y%m%d%H%M%S') + '.jpeg')
plt.show()
# plot 2: contour plot
# plot map ground (North hemisphere)
fig2 = plt.figure(figsize=(8, 8))
ax2 = plt.gca()
m = Basemap(projection='lcc',
resolution='c',
width=8E6,
height=8E6,
lat_0=60,
lon_0=-100)
# draw coastlines, country boundaries, fill continents.
m.drawcoastlines(linewidth=0.25)
m.drawcountries(linewidth=0.25)
m.fillcontinents(color='None', lake_color='None')
# draw the edge of the map projection region (the projection limb)
m.drawmapboundary(fill_color=None)
m.drawlsmask()
m.shadedrelief()
# draw parallels and meridians.
# label parallels on right and top
# meridians on bottom and left
parallels = np.arange(0., 81, 10.)
m.drawparallels(parallels, labels=[False, True, True, False])
meridians = np.arange(10., 351., 20.)
m.drawmeridians(meridians, labels=[True, False, False, True])
date_nightshade = datetime.strptime(dtime, '%Y-%m-%d/%H:%M:%S')
# m.nightshade(date=date_nightshade, alpha = 0.0)
delta = 0.25
lons_dd, lats_dd, tau, dec = daynight_terminator(date_nightshade, delta,
m.lonmin, m.lonmax)
xy = [lons_dd, lats_dd]
xy = np.array(xy)
xb, yb = xy[0], xy[1]
m.plot(xb, yb, marker=None, color='m',
latlon=True) # for dawn-dusk circle line
# Plot the noon-midnight line.
n_interval = len(lons_dd)
ni_half = int(np.floor(len(lons_dd) / 2))
ni_otherhalf = n_interval - ni_half
noon_midnight = noon_midnight_meridian(dtime, delta)
m.plot(noon_midnight['lons_noon'],
noon_midnight['lats_noon'],
marker=None,
color='deepskyblue',
latlon=True) # noon semi-circle
m.plot(noon_midnight['lons_midnight'],
noon_midnight['lats_midnight'],
marker=None,
color='k',
latlon=True) # midnight semi-circle
draw_map(m)
Jy_log = Jy / np.abs(Jy) * np.log10(np.abs(Jy))
norm_cb = CenteredNorm()
# norm_cb = NoNorm()
# norm_cb = CenteredNorm(vmin=Jy.min(), vcenter=0, vmax=Jy.max())
# use Jy for the contour map, not Jy_rot.
ctrf = m.contourf(lon,
lat,
Jy,
contour_density,
latlon=True,
tri=True,
cmap='jet_r',
norm=norm_cb)
##ctrf = m.contourf(lon, lat, Jy, contour_density, latlon=True, tri=True, cmap='jet_r', norm=norm_cb)
# -------------
if s_loc:
m.scatter(lon, lat, latlon=True, marker='*', c='black')
# -------------
cb = m.colorbar(matplotlib.cm.ScalarMappable(norm=norm_cb, cmap='jet_r'),
pad='15%')
cb.set_label(r'$\mathit{J}_y \ (mA/m)$')
ax_cb = cb.ax
text = ax_cb.yaxis.label
font_cb = matplotlib.font_manager.FontProperties(family='times new roman',
style='italic',
size=20)
text.set_font_properties(font_cb)
plt.title(label='EICS ' + tp, fontsize=20, color="black", pad=20)
plt.tight_layout()
plt.savefig(CONFIG['plots_dir'] + 'EICS' + '_contour_' +
date_nightshade.strftime('%Y%m%d%H%M%S') + '.jpeg')
plt.show()
print('EICS plots completed!')
return
fill=False,
linewidth=3,
linestyle='-')
plt.gca().add_patch(poly)
##Robs
#lonband = [5, 15,25,35,45,55,70,70,55,45,35,25,14,5]
#latband = [79,79,79,79,79,79,79,83,83,83,83,83,83,83]
#xband,yband = m(lonband,latband)
#xyband = zip(xband,yband)
#poly_band = Polygon( xyband, edgecolor='k', alpha=1, fill=False, linewidth=3)
#plt.gca().add_patch(poly_band)
#plot velocities
x, y = m(coarse_lons, coarse_lats)
ur, vr = Basemap.rotate_vector(m, u_coarse, v_coarse, coarse_lons,
coarse_lats)
Q = plt.quiver(x, y, ur, vr, units='width', scale=6, width=.002)
qk = plt.quiverkey(Q,
0.9,
0.1,
.1,
r'$10 \frac{cm}{s}$',
labelpos='E',
coordinates='axes',
fontproperties={'size': 16},
labelcolor='w',
color='w')
fig.tight_layout()
outname = 'icecon' + dt
fig.savefig(outpath + outname, bbox_inches='tight')
# x1,y1 = m(70.313, -3.157)
# x2,y2 = m(70.313, 1.754)
# x3,y3 = m(75.234, 1.754)
# x4,y4 = m(75.234, -3.154)
# poly = Polygon([(x1,y1),(x2,y2),(x3,y3),(x4,y4)],facecolor='none',edgecolor='black',linewidth=5)
# plt.gca().add_patch(poly)
# lon_Gan = 73.155
# lat_Gan = -0.694
# x_Gan,y_Gan = m(lon_Gan, lat_Gan)
# m.plot(x_Gan, y_Gan, 'ko', markersize=18)
x_bulk,y_bulk = m(lats_bulk, lons_bulk)
x_winds,y_winds = m(lats_winds, lons_winds)
urot,vrot,x_winds,y_winds = m.rotate_vector(u_850[::4,::4],
v_850[::4,::4],
lons_winds[::4,::4],lats_winds[::4,::4],returnxy=True)
x_winds2,y_winds2 = m(lats_winds, lons_winds)
urot2,vrot2,x_winds2,y_winds2 = m.rotate_vector(u_850[::1,::1],
v_850[::1,::1],
lons_winds[::1,::1],lats_winds[::1,::1],returnxy=True)
s = m.pcolormesh(x_winds2,y_winds2,RH_850[::1,::1],cmap=my_RH_cmap)
plt.clim([20,100])
#cbar=fig3.colorbar(s,shrink=0.5)
cbar=fig1.colorbar(s)
cbar.set_label('Relative Humidity (%)', fontsize=30)
cbar.ax.tick_params(labelsize=18)
B = plt.barbs(x_winds,y_winds,urot,vrot, length=8, flagcolor='k',barbcolor=['k', 'k'])
def __init__(self, radial, nBEAR=72, nRNGE=40):
"""
Constructor
Parametros
----------
radial: Objeto de la clase Radial
Parametros por defecto
----------------------
nBEAR: Número de direcciones
nRNGE: Número de distancias
"""
self.nBEAR, self.nRNGE = nBEAR, nRNGE
origen_lat, origen_lon = radial.Origin
# Escogemos una proyección. Tmercator está bien. La idea es trabajar en un plano:
#m = Basemap(llcrnrlon=-11.0, llcrnrlat=41.8, urcrnrlon=-8, urcrnrlat=44.5, resolution='h', projection='tmerc', lon_0=-8, lat_0=45)
m = Basemap(llcrnrlon=-11.0,
llcrnrlat=41.8,
urcrnrlon=-8,
urcrnrlat=44.5,
resolution='h',
projection='tmerc',
lon_0=origen_lon,
lat_0=origen_lat)
# Necesito las coordenadas del origen y su proyección:
origen_x, origen_y = m(origen_lon, origen_lat)
# Coordenadas polares de los puntos:
RangeResolutionKMeters = radial.RangeResolutionKMeters
AntennaBearing = radial.AntennaBearing
AngularResolution = radial.AngularResolution
# Radios:
RNGE = np.arange(nRNGE) * RangeResolutionKMeters * 1000
# Ángulos:
BEAR = np.arange(nBEAR) * AngularResolution + AntennaBearing
# BEAR = np.sort(BEAR%360)*deg2rad
BEAR = np.sort(BEAR % 360)
# Generamos la lista de vectores unitarios en las direcciones:
X, Y = m.rotate_vector(np.sin(BEAR * deg2rad), np.cos(BEAR * deg2rad),
np.repeat(origen_lon, len(BEAR)),
np.repeat(origen_lat, len(BEAR)))
X = np.array([RNGE * x + origen_x for x in X])
Y = np.array([RNGE * y + origen_y for y in Y])
# ... y las coordenadas esféricas reconstruidas:
longitud, latitud = m(X, Y, inverse=True)
# Preparamos las variables para guardar (las queremos en km no en m y referidas al origen de coordenadas):
X -= origen_x
Y -= origen_y
X /= 1000
Y /= 1000
RNGE /= 1000
# Guardamos las coordenadas proyectadas para trabajar en el plano:
self.X = xr.DataArray(X,
dims={
'BEAR': nBEAR,
'RNGE': nRNGE
},
coords={
'BEAR': BEAR,
'RNGE': RNGE
})
self.Y = xr.DataArray(Y,
dims={
'BEAR': nBEAR,
'RNGE': nRNGE
},
coords={
'BEAR': BEAR,
'RNGE': RNGE
})
# ... que se guardan como xr.DataArray para su uso futuro en la definición de las variables:
self.longitud = xr.DataArray(longitud,
dims={
'BEAR': nBEAR,
'RNGE': nRNGE
},
coords={
'BEAR': BEAR,
'RNGE': RNGE
})
self.latitud = xr.DataArray(latitud,
dims={
'BEAR': nBEAR,
'RNGE': nRNGE
},
coords={
'BEAR': BEAR,
'RNGE': RNGE
})
# ... y las coordenadas polares de los puntos:
self.RNGE, self.BEAR = RNGE, BEAR
dates = [dates[tidx],]
tidx_offset=tidx
idxs =arange(len(x))
np.random.shuffle(idxs)
idxs = idxs[:2000]
os.system('mkdir -p jpgs/vel')
if True:
if True:
fig=figure()
fig.subplots_adjust(left=0.0,right=1.0,bottom=0.0,top=1.0)
pc=tripcolor(x,y,nv,vel,cmap=cmap,rasterized=True)
del hvel,zcor,uvel,vvel
urot,vrot = proj.rotate_vector(u[idxs],v[idxs],lon[idxs],lat[idxs])
del u,v
#proj.quiver(x[idxs],y[idxs],u[idxs],v[idxs],pivot='mid',alpha=0.5)
#urot=u[idxs]; vrot=v[idxs]
if (uselim):
upper_scale=vmax
else:
upper_scale=10.0 # deep transports
upper_scale=0.1 # surface transports
urot = ma.masked_where(vel[idxs]<0.05*upper_scale,urot)
vrot = ma.masked_where(vel[idxs]<0.05*upper_scale,vrot)
#quiverkws={'scale_units':'x','scale':10./10000.} # for deep transports
quiverkws={'scale_units':'x','scale':upper_scale/10000.,'width':0.001}
proj.quiver(x[idxs],y[idxs],urot,vrot,pivot='mid',alpha=0.5,**quiverkws)
# test vector
num_days = size(files) #int(size(files)/int(num_days_lag)+1)
print num_days
if size(files)>0:
drift_days = ma.masked_all((num_days,2, nx, ny))
drift_days_int_mask = ma.masked_all((num_days, 2, nx, ny))
curl_days = ma.masked_all((num_days, nx, ny))
for x in xrange(0, size(files)):
print x
day = '%02d' % x
f = Dataset(files[x], 'r')
u = f.variables['zonal_motion'][0]/(60.*60.*24.*num_days_lag)
v = f.variables['meridional_motion'][0]/(60.*60.*24.*num_days_lag)
q = f.variables['quality_flag'][0]
#ROTATE VECTORS TO X/Y GRID on NSIDC GRID FOR CURL CALC (SO IT IS FLAT ALONG THE BOTTOM)
u_r,v_r = m.rotate_vector(u,v,lon,lat)
#grid data onto 100km grid - also matches wind forcing fields
drift_days[x, 0], drift_days[x, 1], q_int = BGF.interp_uvCSAT(u_r, v_r, q, xpts, ypts, xpts2m, ypts2m)
curl_days[x] = BGF.calc_curl_sq_2d_xy_gradient(drift_days[x, 0], drift_days[x, 1], dx_res)
drift_days_mask_count = ma.count_masked(drift_days[:, 0], axis=0)
days_in_month=15
curl_months[y, mon] = ma.masked_where(drift_days_mask_count>days_in_month, ma.mean(curl_days, axis=0))
drift_month[y, mon, 0] = ma.masked_where(drift_days_mask_count>days_in_month, ma.mean(drift_days[:, 0] , axis=0))
drift_month[y, mon, 1] = ma.masked_where(drift_days_mask_count>days_in_month, ma.mean(drift_days[:, 1] , axis=0))
curl_months.dump(outpath+str(start_year)+'-'+str(end_year)+'-curl_data_months'+grid_str+'.txt')
plot_curl=0
lon = data.variables['longitude']
lat = data.variables['latitude']
u = data.variables['u10'][0, :, :]
v = data.variables['v10'][0, :, :]
spd = np.sqrt(u**2 + v**2)
#draw a map backgroud
m = Basemap(projection='cyl',llcrnrlat=min(lat)-1,urcrnrlat=max(lat)+1,\
resolution='l',llcrnrlon=min(lon)-1,urcrnrlon=max(lon)+1)
#set figure attributes
fig = plt.figure(figsize=(8, 10))
ax = fig.add_axes([0.1, 0.1, 0.8, 0.8])
#add lon, lat, u, v as quiver
uproj, vproj, xx, yy = m.rotate_vector(u, v, lon, lat, returnxy=True)
Q = m.quiver(xx, yy, uproj, vproj, spd, cmap=cm.jet, headlength=7)
#draw coastlines, state and country boundaries, edge of map
m.drawcoastlines()
#create and draw meridians and parallels grid lines
m.drawparallels(np.arange(-90., 90., 10.), labels=[1, 0, 0, 0], fontsize=10)
m.drawmeridians(np.arange(-180., 180., 10.), labels=[0, 0, 0, 1], fontsize=10)
#draw shape
m.readshapefile('CHN_adm1', 'CHN_adm1')
#save
plt.savefig('uv_ecmwf.pdf')
#fig.clf(); lh = llc.pcol(xg,yg,yg); plt.show()
#sys.exit()
#d3=rdmds('diag3Dm',iter)
i = 0
for i in range(1):
uu = np.mean(f[:, 0, :, :], axis=0)
vv = np.mean(f[:, 1, :, :], axis=0)
# uu = np.mean(d3[:,2,1,:,:],axis=0)
# vv = np.mean(d3[:,3,1,:,:],axis=0)
# uu = np.ma.masked_array(uu,hf[0,:,:]<=0.0)
# vv = np.ma.masked_array(vv,hf[0,:,:]<=0.0)
spd = np.sqrt(uu * uu + vv * vv)
u, v = Basemap.rotate_vector(m, uu * cs - vv * sn, uu * sn + vv * cs, xc,
yc)
fig.clf()
lh = llc.pcol(xg, yg, sq(wspd[i, :, :] * hf[0, :, :]), m)
m.drawcoastlines(color='k')
di = 1
q = m.quiver(xc[::di, ::di],
yc[::di, ::di],
u[::di, ::di],
v[::di, ::di],
latlon=True,
scale=200,
pivot='mid',
edgecolor='none',
headaxislength=5)
# pivot = 'mid',units = 'x', edgecolor = 'none',headaxislength = 5)
nice_cmap = plt.get_cmap('RdYlGn')
#nice_cmap= plt.get_cmap(mymap)
#clevs=[-15,-10,-5,-0,5,10,15,20,25,30,35,40,45,50,55,60,70]
clevs = [-9, -6, -3, -0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39]
skip = (slice(None, None, 5), slice(None, None, 5))
# #cs = (m.contourf(x,y,data,leves=clevs,cmap=cmap,norm=norml,extend='both'))
# kw = dict( color='black', alpha=0.75)
# m.quiver(x,y,g_uu,g_vv,width=0.0003, scale=500)
# plt.show()
cs = (m.contourf(x, y, data, clevs, cmap=nice_cmap, extended='both'))
urot, vrot, xx, yy = m.rotate_vector(g_uu[:, :],
g_vv[:, :],
lon[:],
lat[:],
returnxy=True)
Q = m.quiver(xx[skip],
yy[skip],
urot[skip],
vrot[skip],
units='xy',
color='b',
headwidth=4,
headlength=5)
#qk = plt.quiverkey(Q, 0.95, 1.05, 25, '25 m/s', labelpos='W')
#ua, va, xv, yv = m.transform_vector(g_uu,g_vv,lon1,lat1,lon1.shape[0],lon1.shape[0],returnxy=True)
#m.barbs(xv[skip],yv[skip],ua[skip],va[skip],cmap=nice_cmap,length=5,sizes=dict(emptybarb=0.25, spacing=0.05, height=0.5),barbcolor='b',flagcolor='r')
def main():
start_year = 1979
end_year = 1981
HL_LABEL = "CRCM5_HL"
NEMO_LABEL = "CRCM5_NEMO"
dx = 0.1
dy = 0.1
file_prefix = "pm"
PR_level = -1
PR_level_type = level_kinds.ARBITRARY
tprecip_vname = "PR"
sprecip_vname = "SN"
TT_level = 1
TT_level_type = level_kinds.HYBRID
sim_label_to_path = OrderedDict(
[(HL_LABEL, "/RESCUE/skynet3_rech1/huziy/CNRCWP/C5/2016/2-year-runs/coupled-GL+stfl_oneway/Samples"),
(NEMO_LABEL, "/HOME/huziy/skynet3_rech1/CNRCWP/C5/2016/2-year-runs/coupled-GL+stfl/Samples")]
)
# get a coord file ... (use pm* files, since they contain NEM1 variable)
# Should be NEMO_LABEL, since the hostetler case does not calculate NEM? vars
coord_file = ""
found_coord_file = False
for mdir in os.listdir(sim_label_to_path[NEMO_LABEL]):
mdir_path = os.path.join(sim_label_to_path[NEMO_LABEL], mdir)
if not os.path.isdir(mdir_path):
continue
for fn in os.listdir(mdir_path):
if fn[:2] not in ["pm", ]:
continue
if fn[-9:-1] == "0" * 8:
continue
coord_file = os.path.join(mdir_path, fn)
found_coord_file = True
if found_coord_file:
break
bmp, lons, lats = nemo_hl_util.get_basemap_obj_and_coords_from_rpn_file(path=coord_file)
xx, yy = bmp(lons, lats)
stride = 3
#
rot_path = "/HOME/huziy/skynet3_rech1/CNRCWP/C5/2016/2-year-runs/coupled-GL+stfl/rotated_wind_CRCM5_NEMO.nc"
fig = plt.figure()
with Dataset(rot_path) as ds:
ncvars = ds.variables
uu_rot, vv_rot = ncvars["UU"][10, ...], ncvars["VV"][10, ...]
plt.title("rotated in the file")
bmp.quiver(xx[::stride, ::stride], yy[::stride, ::stride], uu_rot[::stride, ::stride], vv_rot[::stride, ::stride], scale=1000, color="r")
lons[lons > 180] -= 360
not_rot_path = "/HOME/huziy/skynet3_rech1/CNRCWP/C5/2016/2-year-runs/coupled-GL+stfl/not_rotated_wind_CRCM5_NEMO.nc"
with Dataset(not_rot_path) as ds:
ncvars = ds.variables
uu, vv = ncvars["UU"][10, ...], ncvars["VV"][10, ...]
uu_rot1, vv_rot1 = rotate_vecs_from_geo_to_rotpole(uu, vv, lons, lats, bmp=bmp)
plt.title("not rotated in the file")
bmp.quiver(xx[::stride, ::stride], yy[::stride, ::stride], uu_rot1[::stride, ::stride], vv_rot1[::stride, ::stride],
scale=1000)
bmp.drawcoastlines()
plt.figure()
b = Basemap(lon_0=0)
xx1, yy1 = b(lons, lats)
uu_rot2, vv_rot2 = b.rotate_vector(uu, vv, lons, lats)
b.quiver(xx1[::stride, ::stride], yy1[::stride, ::stride], uu_rot2[::stride, ::stride], vv_rot2[::stride, ::stride],
scale=1000)
b.drawcoastlines()
plt.show()
class PlotSkyPatch:
"""
Class to plot a close-up look of a region of interest (ROI) in the sky.
To use this class you need to install the Basemap package:
https://matplotlib.org/basemap/users/installing.html
.. code-block:: python
from astrotools.skymap import PlotSkyPatch
patch = PlotSkyPatch(lon0, lat0, r_roi, title='My Skypatch')
mappable = patch.plot_crs("/path/to/cosmic_rays.CosmicRaysSets.npz", set_idx=0)
patch.mark_roi()
patch.plot_grid()
patch.colorbar(mappable)
patch.savefig("/tmp/test-skypatch.png")
"""
def __init__(self, lon_roi, lat_roi, r_roi, ax=None, title=None, **kwargs):
"""
:param lon_roi: Longitude of center of ROI in radians (0..2*pi)
:param lat_roi: Latitude of center of ROI in radians (0..2*pi)
:param r_roi: Radius of ROI to be plotted (in radians)
:param ax: Matplotlib axes in case you want to plot on certain axes
:param title: Optional title of plot (plotted in upper left corner)
:param kwargs: keywords passed to matplotlib.figure()
"""
from mpl_toolkits.basemap import Basemap # pylint: disable=import-error,no-name-in-module
import matplotlib as mpl
with_latex_style = {
"text.usetex": True,
"font.family": "serif",
"axes.labelsize": 30,
"font.size": 30,
"legend.fontsize": 30,
"xtick.labelsize": 26,
"ytick.labelsize": 26,
"legend.fancybox": False,
"lines.linewidth": 3.0,
"patch.linewidth": 3.0
}
mpl.rcParams.update(with_latex_style)
assert (isinstance(lon_roi, (float, int))) and (isinstance(lat_roi, (float, int))) and \
(isinstance(r_roi, (float, int))), "Keywords 'lon_roi', 'lat_roi' and 'r_roi' have to be floats or ints!"
self.vec_0 = coord.ang2vec(lon_roi, lat_roi)
self.lon_0 = np.rad2deg(lon_roi)
self.lat_0 = np.rad2deg(lat_roi)
self.r_roi = r_roi
self.scale = 5500000 * (r_roi / 0.3)
self.fig = None
self.ax = ax
if ax is None:
kwargs.setdefault('figsize', [8, 8])
self.fig = plt.figure(**kwargs)
self.ax = plt.axes()
self.title = title
if title is not None:
self.text(0.02, 0.98, title, verticalalignment='top', fontsize=36)
self.m = Basemap(width=self.scale,
height=self.scale,
resolution='l',
projection='stere',
celestial=True,
lat_0=self.lat_0,
lon_0=-360 -
self.lon_0 if self.lon_0 < 0 else -self.lon_0,
ax=ax)
def plot_crs(self, crs, set_idx=0, zorder=0, cmap='viridis', **kwargs):
"""
Plot cosmic ray events in the sky.
:param crs: Either cosmic_rays.CosmicRaysBase or cosmic_rays.CosmicRaysSets object (or path) or dict object
:param set_idx: In case of CosmicRaysSets object, chose the respective set index
:param zorder: Usual matplotlib zorder keyword (order of plotting)
:param cmap: Matplotlib colormap object or string
"""
if isinstance(crs, str):
from astrotools import cosmic_rays
try:
crs = cosmic_rays.CosmicRaysBase(crs)
except AttributeError:
crs = cosmic_rays.CosmicRaysSets(crs)
if hasattr(crs, 'type') and (crs.type == "CosmicRaysSet"):
crs = crs[set_idx]
if 'log10e' in crs.keys():
log10e = crs['log10e']
assert np.all(
log10e < 25
), "Input energies ('log10e' key) are too high for being plotted"
kwargs.setdefault('s', 10**(log10e - 18.))
kwargs.setdefault('c', log10e)
kwargs.setdefault('lw', 0)
return self.scatter(crs['lon'],
crs['lat'],
zorder=zorder,
cmap=cmap,
**kwargs)
def plot(self, lons, lats, **kwargs):
""" Replaces matplotlib.pyplot.plot() function """
kwargs.setdefault('rasterized', True)
x, y = self.m(np.rad2deg(lons), np.rad2deg(lats))
return self.m.plot(x, y, **kwargs)
def scatter(self, lons, lats, **kwargs):
""" Replaces matplotlib.pyplot.scatter() function """
kwargs.setdefault('rasterized', True)
x, y = self.m(np.rad2deg(lons), np.rad2deg(lats))
return self.m.scatter(x, y, **kwargs)
def tissot(self, lon, lat, radius, npts=1000, **kwargs):
""" Replaces the Basemap tissot() function (plot circles) """
kwargs.setdefault('fill', False)
kwargs.setdefault('lw', 1)
kwargs.setdefault('color', 'grey')
return self.m.tissot(np.rad2deg(lon), np.rad2deg(lat),
np.rad2deg(radius), npts, **kwargs)
def mark_roi(self, alpha=0.4, **kwargs):
"""
Marks the ROI by a circle ans shades cosmic rays outside the ROI.
:param kwargs: Passed to Basemaps tissot() function
"""
from matplotlib import path, collections
kwargs.setdefault('lw', 2)
kwargs.setdefault('zorder', 3)
try:
t = self.tissot(np.deg2rad(self.lon_0), np.deg2rad(self.lat_0),
self.r_roi, **kwargs)
xyb = np.array([[0., 0.], [1., 0.], [1., 1.], [0., 1.], [0., 0.]
]) * self.scale
p = path.Path(np.concatenate([xyb, t.get_xy()[::-1]]))
p.codes = np.ones(len(p.vertices), dtype=p.code_type) * p.LINETO
p.codes[0] = path.Path.MOVETO
p.codes[4] = path.Path.CLOSEPOLY
p.codes[5] = path.Path.MOVETO
p.codes[-1] = path.Path.CLOSEPOLY
col = collections.PathCollection([p],
facecolor='white',
alpha=alpha,
zorder=1)
self.ax.add_collection(col)
except ValueError:
print(
"Warning: Could not plot ROI circle due to undefined inverse geodesic!"
)
self.mark_roi_center()
def mark_roi_center(self, **kwargs):
"""
Mark the ROI center
:param kwargs: keywords for matplotlib.pyplot.plot() function
"""
kwargs.setdefault('marker', '+')
kwargs.setdefault('markersize', 20)
kwargs.setdefault('color', 'k')
kwargs.setdefault('lw', 2)
x, y = self.m(self.lon_0, self.lat_0)
self.m.plot((x), (y), **kwargs)
def plot_grid(self,
meridians=None,
parallels=None,
mer_labels=None,
par_labels=None):
""" Plot the longitude and latitude grid in the skypatch """
if meridians is None:
meridians = np.arange(-180, 181,
60) if abs(self.lat_0) > 60 else np.arange(
-180, 181, 20)
if parallels is None:
parallels = np.arange(-90, 91,
15) if abs(self.lat_0) > 60 else np.arange(
-90, 91, 20)
self.m.drawmeridians(meridians,
labels=[False, False, True, False]
if mer_labels is None else mer_labels)
self.m.drawparallels(parallels,
labels=[True, True, False, False]
if par_labels is None else par_labels)
def plot_thrust(self, n, t, **kwargs):
"""
Visualize the thrust observables in the ROI.
:param n: Thrust axis as given by astrotools.obs.thrust()[1]
:param t: Thrust values as returned by astrotools.obs.thrust()[0]
:param kwargs: Keywords passed to matplotlib.pyplot.plot() for axis visualization
"""
kwargs.setdefault('c', 'red')
linestyle_may = kwargs.pop('linestyle', 'solid')
alpha_may = kwargs.pop('alpha', 0.5)
lon, lat = coord.vec2ang(n[0])
# fill thrust array (unit vector phi runs in negative lon direction)
e_phi = coord.sph_unit_vectors(lon, lat)[1]
sign = np.sign(e_phi[2] - n[1][2])
phi_major = sign * coord.angle(e_phi, n[1])[0]
phi_minor = sign * coord.angle(e_phi, n[2])[0]
if np.abs(phi_major - phi_minor) < 0.99 * np.pi / 2.:
phi_minor = 2 * np.pi - phi_minor
t23_ratio = t[1] / t[2]
# mark the principal axes n3
u = np.array(np.cos(phi_minor))
v = -1. * np.array(np.sin(phi_minor))
urot, vrot, x, y = self.m.rotate_vector(u,
v,
np.rad2deg(lon),
np.rad2deg(lat),
returnxy=True)
_phi = np.arctan2(vrot, urot)
s = self.r_roi * (t[1] / 0.15) * self.scale / t23_ratio
self.m.plot([x - np.cos(_phi) * s, x + np.cos(_phi) * s],
[y - np.sin(_phi) * s, y + np.sin(_phi) * s],
linestyle='dashed',
alpha=0.5,
**kwargs)
# mark the principal axes n2
u = np.array(np.cos(phi_major))
v = -1. * np.array(np.sin(phi_major))
urot, vrot, x, y = self.m.rotate_vector(u,
v,
np.rad2deg(lon),
np.rad2deg(lat),
returnxy=True)
_phi = np.arctan2(vrot, urot)
s = self.r_roi * (t[1] / 0.15) * self.scale
self.m.plot([x - np.cos(_phi) * s, x + np.cos(_phi) * s],
[y - np.sin(_phi) * s, y + np.sin(_phi) * s],
linestyle=linestyle_may,
alpha=alpha_may,
**kwargs)
# mark the center point
self.m.plot((x), (y), 'o', color=kwargs.pop('c'), markersize=10)
def colorbar(self,
mappable,
cblabel='Energy [eV]',
labelsize=12,
ticks=None,
**kwargs):
"""
Adds a colorbar to a mappable in matplotlib. Replaces matplotlib colorbar() function.
Use e.g:
patch = PlotSkyPatch(...)
mappable = patch.plot_crs(crs)
patch.colorbar(mappable)
:param mappable: Mappable in matplotlib.
:param clabel: Label for the colorbar
:param ticks: Ticks for the colorbar (either array-like or integer for number of ticks)
:param kwargs: Keywords passed to matplotlib colorbar() function
"""
# add a colorbar
try:
kwargs.setdefault('location', 'bottom')
cb = self.m.colorbar(mappable, **kwargs)
if (ticks is None) or isinstance(ticks, (int, float)):
vmin, vmax = mappable.get_clim()
vmin, vmax = smart_round(vmin), smart_round(vmax)
n_ticks = float(3) if ticks is None else float(ticks)
step = smart_round((vmax - vmin) / n_ticks, order=1)
ticks = np.arange(vmin, vmax, step)
ticks = ticks[(ticks >= vmin) & (ticks <= vmax)]
cb.set_ticks(ticks)
cb.set_label(cblabel, fontsize=3 * labelsize)
t = ['$10^{%.1f}$' % (f) for f in ticks]
cb.ax.set_xticklabels(t, fontsize=int(max(0.8 * 3 * labelsize, 1)))
except KeyError:
print("Can not plot colorbar on axis.")
def text(self, x, y, s, **kwargs):
""" Substitudes matplotlib.pyplot.text() function """
kwargs.setdefault('transform', self.ax.transAxes)
self.ax.text(x, y, s, **kwargs)
def savefig(self, path, **kwargs):
""" Substitudes matplotlib savefig() function """
kwargs.setdefault('dpi', 150)
kwargs.setdefault('bbox_inches', 'tight')
self.fig.savefig(path, **kwargs)
fig=plt.figure(figsize=(9, 3))
map = Basemap(projection='sinu',
lat_0=0, lon_0=0)
lons = np.linspace(-180, 180, 10)
lats = np.linspace(-90, 90, 10)
lons, lats = np.meshgrid(lons, lats)
v10 = np.ones((lons.shape)) * 15
u10 = np.zeros((lons.shape))
u10_rot, v10_rot, x, y = map.rotate_vector(u10, v10, lons, lats, returnxy=True)
ax = fig.add_subplot(121)
ax.set_title('Without rotation')
map.drawmapboundary(fill_color='aqua')
map.fillcontinents(color='#cc9955', lake_color='aqua', zorder = 0)
map.drawcoastlines(color = '0.15')
map.barbs(x, y, u10, v10,
pivot='middle', barbcolor='#333333')
ax = fig.add_subplot(122)
ax.set_title('Rotated vectors')
fig=plt.figure(figsize=(9, 3))
map = Basemap(projection='sinu',
lat_0=0, lon_0=0)
lons = np.linspace(-180, 180, 10)
lats = np.linspace(-90, 90, 10)
lons, lats = np.meshgrid(lons, lats)
v10 = np.ones((lons.shape)) * 15
u10 = np.zeros((lons.shape))
u10_rot, v10_rot, x, y = map.rotate_vector(u10, v10, lons, lats, returnxy=True)
ax = fig.add_subplot(121)
ax.set_title('Without rotation')
map.drawmapboundary(fill_color='aqua')
map.fillcontinents(color='#cc9955', lake_color='aqua', zorder = 0)
map.drawcoastlines(color = '0.15')
map.barbs(x, y, u10, v10,
pivot='middle', barbcolor='#333333')
ax = fig.add_subplot(122)
ax.set_title('Rotated vectors')
if size(files) > 0:
drift_days = ma.masked_all((num_days, 2, nx, ny))
drift_days_int_mask = ma.masked_all((num_days, 2, nx, ny))
curl_days = ma.masked_all((num_days, nx, ny))
for x in xrange(0, size(files)):
print x
day = '%02d' % x
f = Dataset(files[x], 'r')
u = f.variables['zonal_motion'][0] / (60. * 60. * 24. *
num_days_lag)
v = f.variables['meridional_motion'][0] / (60. * 60. * 24. *
num_days_lag)
q = f.variables['quality_flag'][0]
#ROTATE VECTORS TO X/Y GRID on NSIDC GRID FOR CURL CALC (SO IT IS FLAT ALONG THE BOTTOM)
u_r, v_r = m.rotate_vector(u, v, lon, lat)
#grid data onto 100km grid - also matches wind forcing fields
drift_days[x, 0], drift_days[x, 1], q_int = BGF.interp_uvCSAT(
u_r, v_r, q, xpts, ypts, xpts2m, ypts2m)
curl_days[x] = BGF.calc_curl_sq_2d_xy_gradient(
drift_days[x, 0], drift_days[x, 1], dx_res)
drift_days_mask_count = ma.count_masked(drift_days[:, 0], axis=0)
days_in_month = 15
curl_months[y, mon] = ma.masked_where(
drift_days_mask_count > days_in_month, ma.mean(curl_days, axis=0))
drift_month[y, mon,
0] = ma.masked_where(drift_days_mask_count > days_in_month,
ma.mean(drift_days[:, 0], axis=0))
drift_month[y, mon,
1] = ma.masked_where(drift_days_mask_count > days_in_month,
latg=nvargrd['lat_rho'][:]
long=nvargrd['lon_rho'][:]
nlat,nlon=long.shape
cang=nvargrd['tide_Cangle'][:]
cangr=pl.zeros_like(cang)
u0=pl.ones_like (long)
v0=pl.zeros_like(latg)
ur=pl.zeros_like (long)
vr=pl.zeros_like(latg)
ur,vr=proj.rotate_vector(u0,v0,long,latg)
alpha=pl.arccos(ur)*180./pi
for id in range(len(cang)):
cangr[id,:]=cang[id,:]-alpha
cangr[cangr<0.] =cangr[cangr<0.] +360.
cangr[cangr>360.]=cangr[cangr>360.]-360.
nvargrd['tide_Cangle'][:]=cangr
#dimensions:
# two = 2 ;
# eta_rho = 300 ;
# xi_rho = 240 ;
71964,60.7,-135.1,5310.0,250.0,36.5
'''
# Extracts data from IDV output and masks out stations outside of North America
lat,lon,height,wdir,speed = np.loadtxt(StringIO(data), delimiter=',',
unpack=True, usecols=(1,2,3,4,5))
# Create a map for plotting
bm = Basemap(projection='tmerc', lat_0=90.0, lon_0=-100.0, lat_ts=40.0,
llcrnrlon=-121, llcrnrlat=24, urcrnrlon=-65, urcrnrlat=46,
resolution='l')
# Get U,V components from wind speed and direction
u,v = metpy.get_wind_components(speed, wdir)
# Rotate the vectors to be properly aligned in the map projection
u,v = bm.rotate_vector(u, v, lon, lat)
# Generate grid of x,y positions
lon_grid, lat_grid, x_grid, y_grid = bm.makegrid(130, 60, returnxy=True)
# Transform the obs to basemap space for gridding
obx,oby = bm(lon, lat)
# Perform analysis of height obs using Cressman weights
heights_oban = grid_data(height, x_grid, y_grid, obx, oby,
obans[which_oban][0], obans[which_oban][1])
heights_oban = maskoceans(lon_grid, lat_grid, heights_oban)
# Map plotting
contours = np.arange(5000., 5800., 60.0)
def plot_flow_vectors_basemap(lons,
lats,
uu,
vv,
flow_speed,
subregion: list = None,
grid_shape=None,
ax: Axes = None,
streamplot=False,
draw_colorbar=True):
from mpl_toolkits.basemap import Basemap
if grid_shape is None:
grid_shape = (300, 300)
if subregion is None:
subregion = [0, 1, 0, 1]
if ax is None:
fig = plt.figure()
nx, ny = lons.shape
b = Basemap(lon_0=180,
llcrnrlon=lons[35, 35],
llcrnrlat=lats[35, 35],
urcrnrlon=lons[nx // 2, ny // 2],
urcrnrlat=lats[nx // 2, ny // 2],
resolution="i",
area_thresh=2000)
# im = b.pcolormesh(xx, yy, flow_speed)
# b.colorbar(im)
stride = 6
uu1, vv1 = b.rotate_vector(uu, vv, lons, lats)
lons_g, lats_g, xx_g, yy_g = b.makegrid(*grid_shape, returnxy=True)
nx, ny = lons_g.shape
i_start, i_end = int(nx * subregion[0]), int(nx * subregion[1])
j_start, j_end = int(ny * subregion[2]), int(ny * subregion[3])
xt, yt, zt = lat_lon.lon_lat_to_cartesian(lons_g.flatten(),
lats_g.flatten())
xs, ys, zs = lat_lon.lon_lat_to_cartesian(lons.flatten(), lats.flatten())
ktree = KDTree(data=list(zip(xs, ys, zs)))
dists, inds = ktree.query(list(zip(xt, yt, zt)))
uu_to_plot = uu1.flatten()[inds].reshape(lons_g.shape)
vv_to_plot = vv1.flatten()[inds].reshape(lons_g.shape)
flow_speed_to_plot = flow_speed.flatten()[inds].reshape(lons_g.shape)
clevs = [0, 0.01, 0.02, 0.04, 0.06, 0.08, 0.12, 0.16]
ncolors = len(clevs) - 1
norm = BoundaryNorm(clevs, ncolors)
cmap = cm.get_cmap("gist_ncar_r", ncolors)
im = b.pcolormesh(xx_g,
yy_g,
flow_speed_to_plot,
alpha=0.5,
cmap=cmap,
norm=norm)
if not streamplot:
b.quiver(xx_g[i_start:i_end:stride, j_start:j_end:stride],
yy_g[i_start:i_end:stride, j_start:j_end:stride],
uu_to_plot[i_start:i_end:stride, j_start:j_end:stride],
vv_to_plot[i_start:i_end:stride, j_start:j_end:stride],
headlength=2,
headaxislength=2,
headwidth=4,
units="inches",
color="k")
else:
b.streamplot(xx_g,
yy_g,
uu_to_plot,
vv_to_plot,
linewidth=0.4,
density=3,
arrowstyle="fancy",
arrowsize=0.4,
ax=ax,
color="k")
# im = b.contourf(xx_g, yy_g, flow_speed_to_plot, levels=clevs, alpha=0.5, cmap=cmap, norm=norm)
cb = b.colorbar(im, location="bottom")
cb.ax.set_visible(draw_colorbar)
b.drawcoastlines(linewidth=0.3, ax=ax)
if ax is None:
fig.savefig("nemo/circ_annual_mean_basemap.png",
bbox_inches="tight",
dpi=300)
plt.close(fig)
return im
