get_spd_field.py

Go to the documentation of this file.

import pygrib 
import os, math
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap

##\ingroup graphics
def get_spd_field(fn,lvl,typeOfLevel,outdir,outmap,fhr,var):
    print "GET ARGS: ", locals()
    tmp    = locals()
    alist  = {}
    for i in tmp:
      if tmp[i] is not None and i != 'fn':
        print i , tmp[i]
        if i == 'level':
           alist[i] = int(tmp[i])
        else:
           alist[i] = tmp[i]
    print "ALIST : ", alist
    grbs=pygrib.open(fn)
    plt.figure(figsize=(15,12))
    lev   = int(lvl)
    if var == 'wind':
         grbu=grbs.select(typeOfLevel='isobaricInhPa', shortName='u', level=lev)[0]
         grbv=grbs.select(typeOfLevel='isobaricInhPa', shortName='v', level=lev)[0]
         clevs=np.arange(30,120,5)
    elif var == '10wind':
         grbu=grbs.select(typeOfLevel='heightAboveGround', shortName='10u', level=lev)[0]
         grbv=grbs.select(typeOfLevel='heightAboveGround', shortName='10v', level=lev)[0]
         clevs=np.arange(5,35,1)
    
    datau=grbu.values
    datav=grbv.values
    atime=grbu.analDate.strftime('Analysis: %Y%m%d %H UTC')
    altime=grbu.analDate.strftime('%Y%m%d%H')
    vtime=grbu.validDate.strftime('Valid: %Y%m%d %H UTC')
    date=os.getenv('START_TIME')     # @Y@m@d@H
    grid  = outmap

    lat,lon = grbu.latlons()
    data=np.sqrt(datau**2+datav**2)

    if grid == 'glob' or grid == 'G3':
     m =Basemap(projection='cyl',lat_ts=10,llcrnrlon=lon.min(), \
      urcrnrlon=lon.max(),llcrnrlat=lat.min(),urcrnrlat=lat.max(), \
      resolution='c')
    elif grid == 'nh_mill':
     m =Basemap(projection='mill',lat_ts=10,llcrnrlon=lon.min(), \
      urcrnrlon=lon.max(),llcrnrlat=0,urcrnrlat=90, \
      resolution='c')
    elif grid == 'sh_mill':
     m =Basemap(projection='mill',lat_ts=10,llcrnrlon=lon.min(), \
      urcrnrlon=lon.max(),llcrnrlat=-90,urcrnrlat=0, \
      resolution='c')
    elif grid == 'trop_mill':
     m =Basemap(projection='mill',lat_ts=10,llcrnrlon=lon.min(), \
      urcrnrlon=lon.max(),llcrnrlat=-23,urcrnrlat=23, \
      resolution='c')
    elif grid == 'nps':
     m =Basemap(projection='npstere',boundinglat=-0.268, \
      lon_0=-105,resolution='c')
    elif grid == 'sps':
     m =Basemap(projection='npstere',boundinglat=-60, \
      lon_0=105,resolution='c')
    elif grid == 'G218':
     m=Basemap(projection='lcc',lat_ts=10,llcrnrlon=226.514, \
      urcrnrlon=-49.420,llcrnrlat=12.190,urcrnrlat=57.328, \
      resolution='c',lat_1=25,lat_2=25,lon_0=265,rsphere=6371200)
    elif grid == 'G104':   
     lon_w,lon_e,lat_s,lat_n =[-147, -64, 10.8, 89.9]
     lon_0 = lon_w + (lon_e - lon_w) / 2.
     ref = lat_s if abs(lat_s) > abs(lat_n) else lat_n
     lat_0 = math.copysign(90., ref)
     proj = 'npstere' if lat_0 > 0 else 'spstere'
     prj = Basemap(projection=proj, lon_0=lon_0, lat_0=lat_0,
                          boundinglat=0, resolution='c')
     lons = [lon_w, lon_e, lon_w, lon_e, lon_0, lon_0]
     lats = [lat_s, lat_s, lat_n, lat_n, lat_s, lat_n]
     x, y = prj(lons, lats)
     ll_lon, ll_lat = prj(min(x), min(y), inverse=True)
     ur_lon, ur_lat = prj(max(x), max(y), inverse=True)
     m=Basemap(projection='stere', lat_0=lat_0, lon_0=lon_0,
                           llcrnrlon=ll_lon, llcrnrlat=ll_lat,
                           urcrnrlon=ur_lon, urcrnrlat=ur_lat, resolution='c')
   
    x, y = m(lon,lat)
    
    # Draw political boundaries
    m.drawcoastlines()
    m.drawstates()
    m.drawcountries()
    m.drawmapboundary()
#    m.drawlsmask(ocean_color='0.8',land_color='white')
    m.drawlsmask(ocean_color='white',land_color='white')

    # Draw lat/lon grid
    parallels = np.arange(-90.,90,10.)
    m.drawparallels(parallels,labels=[1,0,0,0],fontsize=10)
    meridians = np.arange(0.,360.,20.)
    m.drawmeridians(meridians,labels=[0,0,0,1],fontsize=10)

#    clevs=np.arange(30,100,5)
    cs  = m.contourf(x,y,data,clevs,cmap=plt.cm.rainbow)

    # Determine how to label the vertical coordinate based on the contents of the grib file, or other
    #       information.'''
    if  typeOfLevel == 'isobaricInhPa':
             vert_unit = 'hPa'
    elif typeOfLevel == 'hybrid':
             vert_unit = 'sigma'
    elif typeOfLevel == 'heightAboveGround':
             vert_unit = 'm'
    else:
             vert_unit ='m'

    # Plot the wind field in barbs. Masks are applied here to 
    # thin the quantity of barbs plotted for calrity.
    maskarray = np.ones(datau.shape)
    if grid == 'G3':
             maskarray[1::7,1::14] = 0
    elif grid == 'G104':
             maskarray[1::3,1::6] = 0
    else:
             maskarray[1::20,1::40] = 0

    mu = np.ma.masked_array(datau, mask=maskarray)
    mv = np.ma.masked_array(datav, mask=maskarray)

    barb_control = dict(height=0.6, width=0.3, emptybarb=0.07)
    m.barbs(x, y, mu, mv, barbcolor='k',pivot='middle',sizes=barb_control, linewidth=0.4, length=5)



    #Add a colorbar and title, and then show the plot.
    plt.colorbar(cs,orientation='horizontal',shrink=0.8)
    plt.title('%s \nFcst Hr: %s' % (atime, fhr) , loc='left')

    plt.title('wind magnitude (m s**-1,shaded)', loc = 'right')
    plt.title('%s %s' % (lvl,vert_unit))
    plt.xlabel('%s' % (vtime), fontsize=18, labelpad=40)

    # Output file name definition
    filename=''.join(['wind','%s' % (lvl),vert_unit,'_',grid,'_',altime,'_f',fhr,'.png'])
     

    # Check to see if output location exists, if not, make it.
    if not os.path.exists(outdir):
           os.makedirs(outdir)

    plt.savefig('/'.join([outdir, filename]))
    plt.close()