#!/usr/bin/env python import traceback import sys import sqlite3 from datetime import datetime as dt import numpy as n import matplotlib matplotlib.use('agg') import matplotlib.pyplot as plt #from wrf.core import Constants #from wrf.var.temp import calc_temp #from wrf.plot.matplotlib.defaults import (get_basemap, get_default_map_opts, # get_null_opts) #from wrf.plot.matplotlib.helper import (add_plot_info_text, plot_map, # plot_contourf) #__all__ = ["plot_2d"] def get_basemap(wrfnc): #TODO: handle multiple projections lat2d = wrfnc.variables["XLAT"][0,:,:] lon2d = wrfnc.variables["XLONG"][0,:,:] ny = len(wrfnc.dimensions["south_north"]) nx = len(wrfnc.dimensions["west_east"]) nz = len(wrfnc.dimensions["bottom_top"]) dx = wrfnc.DX dy = wrfnc.DY center_lat = wrfnc.CEN_LAT center_lon = wrfnc.CEN_LON true_lat1 = wrfnc.TRUELAT1 true_lat2 = wrfnc.TRUELAT2 basemap = Basemap(projection="lcc", lat_0=center_lat, lon_0=center_lon, lat_1=true_lat1, lat_2=true_lat2, llcrnrlon=lon2d[0,0], llcrnrlat=lat2d[0,0], urcrnrlon=lon2d[ny-1, nx-1], urcrnrlat=lat2d[ny-1, nx-1], resolution='i') return basemap def get_default_map_opts(): landcolor = (204/255.0, 204/255.0, 153/255.0) oceancolor = (102/255.0, 204/255.0, 255/255.0) return MapOptions( coastargs = {"linewidth":1.0, "linestyle":'solid', "color":'k'}, countryargs = {"linewidth":0.5, "linestyle":'solid', "color":'k'}, stateargs = {"linewidth":0.5, "linestyle":'solid', "color":'k'}, mapboundaryargs = {"color":'k', "linewidth":1.0, "fill_color":oceancolor}, continentfillargs = {"color":landcolor, "lake_color":oceancolor, "zorder":0}) def get_null_opts(): return FilledContourOptions(fcontourargs={}, colorbarargs={"location" : "bottom", "size" : "5%", "pad" : "2%", "extend" : "both"}) def add_plot_info_text(ax, top_left_text="", top_right_text="", bot_left_text="", bot_right_text=""): plt.ioff() if top_left_text != "": plt.text(0.005,.995,top_left_text, bbox=dict(facecolor="white"), horizontalalignment="left", verticalalignment="top", transform = ax.transAxes, fontsize=10) if top_right_text != "": plt.text(.995,.995,top_right_text, bbox=dict(facecolor="white"), horizontalalignment="right", verticalalignment="top", transform = ax.transAxes, fontsize=10) if bot_left_text != "": plt.text(0.005,0.005,bot_left_text, bbox=dict(facecolor="white"), horizontalalignment="left", verticalalignment="bottom", transform = ax.transAxes, fontsize=10) if bot_right_text != "": plt.text(.995,0.005,bot_right_text, bbox=dict(facecolor="white"), horizontalalignment="right", verticalalignment="bottom", transform = ax.transAxes, fontsize=10) def plot_map(basemap, mapoptions): plt.ioff() if mapoptions.mapboundaryargs is not None: basemap.drawmapboundary(**mapoptions.mapboundaryargs) if mapoptions.continentfillargs is not None: basemap.fillcontinents(**mapoptions.continentfillargs) if mapoptions.coastargs is not None: basemap.drawcoastlines(**mapoptions.coastargs) if mapoptions.countryargs is not None: basemap.drawcountries(**mapoptions.countryargs) if mapoptions.stateargs is not None: basemap.drawstates(**mapoptions.stateargs) if mapoptions.countyargs is not None: basemap.drawcounties(**mapoptions.countyargs) if mapoptions.riverargs is not None: basemap.drawrivers(**mapoptions.riverargs) def plot_contourf(x,y,data,basemap, contourfoptions): plt.ioff() cs1 = None if contourfoptions.contourargs is not None: cs1 = basemap.contour(x,y,data,**contourfoptions.contourargs) cs2 = None if contourfoptions.fcontourargs is not None: cs2 = basemap.contourf(x,y,data,**contourfoptions.fcontourargs) cb = None if contourfoptions.colorbarargs is not None: cb = basemap.colorbar(cs2, **contourfoptions.colorbarargs) if contourfoptions.labelargs is not None: plt.clabel(cs1, **contourfoptions.labelargs) return cs1, cs2, cb def get_null_opts(): return FilledContourOptions(fcontourargs={}, colorbarargs={"location" : "bottom", "size" : "5%", "pad" : "2%", "extend" : "both"}) class MapOptions(object): def __init__(self, coastargs = None, countyargs = None, countryargs = None, riverargs = None, stateargs = None, mapboundaryargs = None, continentfillargs = None): self.coastargs = coastargs self.countyargs = countyargs self.countryargs = countryargs self.riverargs = riverargs self.stateargs = stateargs self.mapboundaryargs = mapboundaryargs self.continentfillargs = continentfillargs def plot_2d(wrfnc, varname=None, outfile=None, title=None, map_opts=None, plot_opts=None, top_left_info="", top_right_info="", bot_left_info="", bot_right_info="", wks_type="png", var=None, time_in=0): try: plt.ioff() print "generating %s.%s" % (outfile, wks_type) if var is not None: field = var elif varname is not None: field = wrfnc.variables[varname][time_in,:,:] lat2d = wrfnc.variables["XLAT"][time_in,:,:] lon2d = wrfnc.variables["XLONG"][time_in,:,:] times = wrfnc.variables["Times"][time_in,:] model_time = "".join(times) start_date = dt.strptime(model_time, "%Y-%m-%d_%H:%M:%S") ny = len(wrfnc.dimensions["south_north"]) nx = len(wrfnc.dimensions["west_east"]) nz = len(wrfnc.dimensions["bottom_top"]) fig = plt.figure(figsize=(8,8), dpi=200) ax = fig.add_axes([0.1,0.1,0.8,0.8]) bm = get_basemap(wrfnc) if map_opts is None: map_opts = get_default_map_opts() if plot_opts is None: plot_opts = get_null_opts() x,y = bm(lon2d, lat2d) plot_map(bm,map_opts) plt.xticks(rotation=70) tplot = plot_contourf(x,y,field,bm,plot_opts) add_plot_info_text(ax, top_left_info, top_right_info, bot_left_info, bot_right_info) ax.set_title(title,fontdict={"fontsize" : 20}) plt.savefig("%s.%s" % (outfile, wks_type)) plt.clf() plt.close(fig) except: # print the stack trace since it will be lost when used in a # multiprocessing worker. print traceback.format_exc() raise finally: sys.stdout.flush() def main(): parser = argparse.ArgumentParser(description="Generate meteorological " "plots for a specific data file") parser.add_argument("-v", "--var", required=True, help="variable name") parser.add_argument("-f", "--filename", required=True, help="WRF file to plot") parser.add_argument("-o", "--outdir", default=".", required=False, help="output directory for images") parser.add_argument("-l", "--levels", required=False, type=float, nargs="+", default=None, help=("the start, end, and increment for contour levels" " as a list of items with spaces between them" "example: 1 10 2 ")) parser.add_argument("-c", "--customlevels", required=False, type=float, nargs="+", default=None, help=("a list of space delimited contour levels" "example: 1 2 3 4 5 19 28 200 ")) args = parser.parse_args() if not os.path.exists(args.filename): raise RuntimeError ("%s not found" % args.filename) if not os.path.exists(args.outdir): os.makedirs(args.outdir) basename = os.path.basename(args.filename) wrfnc = NetCDF(args.filename, mode='r') outfile = os.path.join(args.outdir, domain, args.var, "%s.%s" % (basename,args.var)) if not os.path.exists(os.path.dirname(outfile)): os.makedirs(os.path.dirname(outfile)) if args.levels is not None or args.customlevels is not None: plot_opts = get_null_opts() if args.levels is not None: if len(args.levels) < 2 or len(args.levels) > 3: raise RuntimeError("levels argument is invalid") plot_opts.fcontourargs["levels"] = [x for x in n.arange(args.levels[0], args.levels[1], args.levels[2])] plot_opts.fcontourargs["extend"] = "both" elif args.customlevels is not None: plot_opts.fcontourargs["levels"] = args.customlevels plot_opts.fcontourargs["extend"] = "both" else: plot_opts = None plot_2d(wrfnc, args.var, outfile, "%s"%args.var, plot_opts = plot_opts) if __name__ == "__main__": main()