plots.py 文件源码

def basicsurfmap(file,varstr,outputdir):
  """
  Creates plots for each timestep of a file
  for the specified variable and outputs them
  to a specified directory
  """
  import matplotlib
  matplotlib.use('Agg')
  import numpy as np
  import sys
  from netCDF4 import Dataset
  import matplotlib.pyplot as plt
  from mpl_toolkits.basemap import Basemap
  # read in file and vars
  f = Dataset(file,'r')
  wrfeta = f.variables['ZNU'][0][:]
  times = f.variables['Times'][:]
  wrflats = f.variables['XLAT'][0][:]
  wrflons = f.variables['XLONG'][0][:]
  var = f.variables[varstr][:]
  # four corners of domain
  print wrflats.shape
  print wrflons.shape
  print wrflons[0].shape
  wrflat_s = wrflats[0,len(wrflats)-1]
  wrflat_n = wrflats[len(wrflats)-1,len(wrflons[0])-1]
  wrflon_w = wrflons[0,0]
  wrflon_e = wrflons[len(wrflats)-1,len(wrflons[0])-1]

  z = 0 # assuming lowest level of model

  # set up map
  map = Basemap(projection='merc',llcrnrlon=wrflon_w,urcrnrlon=wrflon_e,llcrnrlat=wrflat_s,urcrnrlat=wrflat_n,resolution='i')
  map.drawstates()
  map.drawcounties()
  map.drawcoastlines()
  x,y = map(wrflons,wrflats)

  # loop through times
  for t in range(len(times)):
    timestr = ''.join(times[t,:])
    map.drawstates()
    map.drawcounties()
    map.drawcoastlines()
    plt1 = map.pcolormesh(x,y,var[t,z,:,:],vmin=np.amin(var),vmax=np.amax(var))
    colorbar = map.colorbar(plt1,"right", size="5%",pad="2%")
    colorbar.set_label(f.variables[varstr].description+' '+f.variables[varstr].units)   
    plt.title('WRF output valid: '+timestr)
    plt.savefig(outputdir+'/%03d_' % (t) +timestr+'_'+varstr+'.png')
    plt.clf()