GMTB_graphics.py

Go to the documentation of this file.

#!/usr/bin/env python
##\file GMTB_graphics.py

import matplotlib as mpl 
mpl.use('Agg')
import pygrib
from mpl_toolkits.basemap import Basemap
import numpy as np
import matplotlib.pyplot as plt
import datetime
import os,sys,math
import maps, regrib
import get_grib_field as ggf
import get_spd_field as gsf

##\ingroup graphics
# This function sets up a variable table to plot (map_confs) using key/value pairs.
#\param lev The value of \e level
#\param pvar The value of \e shortName
#\param barbs Flag for overlaying wind barbs
#\param hgt   Flag for overlaying height contour
#\param vc The value for \e typeOfLevel
#\param acc_field Flag for if \e stepTypeInternal is accum 
def set_def_maps():
    map_confs = {
        '200_wind':   {'lev': '200', 'pvar':  'wind', 'barbs':  True,  'hgt':  False, 'vc': 'isobaricInhPa'},
        '200_temp':   {'lev': '200', 'pvar':     't', 'barbs':  True,  'hgt':  False, 'vc': 'isobaricInhPa'},
        '500_hgt':    {'lev': '500', 'pvar':    'gh', 'barbs':  True,  'hgt':  True,  'vc': 'isobaricInhPa'},
        '500_temp':   {'lev': '500', 'pvar':     't', 'barbs':  True,  'hgt':  False, 'vc': 'isobaricInhPa'},
        '500_vort':   {'lev': '500', 'pvar':  'absv', 'barbs': False,  'hgt':  False, 'vc': 'isobaricInhPa'},
        '700_temp':   {'lev': '700', 'pvar':     't', 'barbs':  True,  'hgt':  False, 'vc': 'isobaricInhPa'},
        '700_vvel':   {'lev': '700', 'pvar':     'w', 'barbs': False,  'hgt':  False, 'vc': 'isobaricInhPa'},
        '850_hgt':    {'lev': '850', 'pvar':    'gh', 'barbs':  True,  'hgt':  True,  'vc': 'isobaricInhPa'},
        '850_temp':   {'lev': '850', 'pvar':     't', 'barbs':  True,  'hgt':  False, 'vc': 'isobaricInhPa'},
        '850_rh':     {'lev': '850', 'pvar':     'r', 'barbs':  True,  'hgt':  False, 'vc': 'isobaricInhPa'},
        '2m_tmp':     {'lev':   '2', 'pvar':    '2t', 'barbs': False,  'hgt':  False, 'vc': 'heightAboveGround'},
        '2m_dpt':     {'lev':   '2', 'pvar':    '2d', 'barbs': False,  'hgt':  False, 'vc': 'heightAboveGround'},
        '10m_u':      {'lev':  '10', 'pvar':   '10u', 'barbs': False,  'hgt':  False, 'vc': 'heightAboveGround'},
        '10m_v':      {'lev':  '10', 'pvar':   '10v', 'barbs': False,  'hgt':  False, 'vc': 'heightAboveGround'},
        '10m_wind':   {'lev':  '10', 'pvar':'10wind', 'barbs':  True,  'hgt':  False, 'vc': 'heightAboveGround'},
        'prmsl':      {'lev':   '0', 'pvar': 'prmsl', 'pline': True, 'barbs':  True,  'hgt':  False, 'vc': 'meanSea'},
        'pwat':       {'lev':   '0', 'pvar':  'pwat', 'pline': False, 'barbs':  False,  'hgt':  False, 'vc': 'unknown'},
        'prate':      {'lev':   '0', 'pvar':  'prate','pline': False, 'barbs':  False,  'hgt':  False, 'vc': 'surface', 'acc_field': True},
        'cprat':      {'lev':   '0', 'pvar':  'cprat','pline': False, 'barbs':  False,  'hgt':  False, 'vc': 'surface', 'acc_field': True},
        '6h_acpcp':   {'lev':   '0', 'pvar': 'acpcp', 'barbs': False,  'hgt':  False, 'vc': 'surface', 'acc_field': True},
        '6h_apcp':    {'lev':   '0', 'pvar':    'tp', 'barbs': False,  'hgt':  False, 'vc': 'surface', 'acc_field': True}
                }

    return map_confs


##\defgroup graphics graphics
def main():
 
    ## Get the cycle information from env vars
    moad_dataroot=os.getenv('MOAD_DATAROOT')
    rotdir=os.getenv('ROTDIR')
    date=os.getenv('START_TIME')     # @Y@m@d@H
    fhr=os.getenv('FCST_TIME') # 3 digit fcst hour
    grid_list=os.getenv('GRID_VX_LIST')
    vx_grid=grid_list.split(" ")
    filename=''.join(['pgrbq', fhr, '.gfs.', date, '.grib2'])
    gribin='/'.join([rotdir, filename])
    user_def = os.getenv('USER_MAPS',False)

    ## Set the output file location (template set in maps.py).
    outdir='/'.join([moad_dataroot, 'figprd'])


    # Change projection if needed.
    for grid in vx_grid: 
        outmap = grid
        print grid
        
        # Wind rotation is a user choice, not necessarily a predefined choice, so is included here.
        wind_rot = {'G104': 'grid', 'G218': 'grid', 'G3': 'grid'}

        # Assumes that NCEP pre-defined grids will be defined with a leading "G"  
        if grid[0] == 'G':
           newgrib = ''.join(['pgrbq', fhr, '.gfs.', date,'_', grid, '.grib2'])
           newgrib = '/'.join([outdir,newgrib])
           if not os.path.exists(outdir):
             os.makedirs(outdir)
           vx = regrib.new_grid(grid,gribin,newgrib,wind_rot.get(grid,None))
           vx.regrid()
           grib = newgrib
        else:
           grib = gribin

        plot_wind=True
#z        print gribin, outdir
    
    
        # Set forecast hour
        pvars=['gh', 't']
        plevels=[ 850, 700, 500, 200 ]
        coord = [ "isobaricInhPa", "isobaricInhPa", "isobaricInhPa", "isobaricInhPa" ]
        # Variables include: t, r, u, v, w, gh, absv, clwmr,soilw, 2t, 2d, pwat,
        # cwat,press,q,pt,tp, acpcp
        if user_def:  
            for var in pvars:
                for level in plevels: 
                  
                    if plot_wind and var == 'wind':
                        u = ggf.get_ua_field(grib,level,coord,shortName='u')
                        v = ggf.get_ua_field(grib,level,coord,shortName='v')
                        winds=[u,v]
                    elif plot_wind and var == '10wind':
                        u = ggf.get_ua_field(grib,level,coord,shortName='10u')
                        v = ggf.get_ua_field(grib,level,coord,shortName='10v')
                        winds=[u,v]
                    else:
                        winds=None

                    # Get the variable from file by passing key,value pairs from the grib file. If you do not want to 
                    # use a particular field to search for the grib record, "None" is also acceptable.
                    if var != 'wind':
                         field=ggf.get_ua_field(grib,level,coord,shortname=var)
                    else:
                         field=u
                    
                    
                    mymap=maps.make_basemap(field,date,var,level,
                                          winds=winds, fhr=fhr)
                    mymap.set_bm()
                    mymap.run()
                    mymap.save_figure(outdir)
        else:
            def_maps=set_def_maps()
            for fig in def_maps:
              # Optional if statement here can control evaluating one type of figure at a time.
              #if fig == '6h_apcp' or fig == '6h_acpcp':
              acc_field=def_maps[fig].get('acc_field',False)
              if (acc_field and fhr != '000') or (not acc_field): 
                level = def_maps[fig].get('lev')
                var   = def_maps[fig].get('pvar')
                hgt   = def_maps[fig].get('hgt')
                plot_wind = def_maps[fig].get('barbs',False)
                plot_line = def_maps[fig].get('pline',False)
                coord     = def_maps[fig].get('vc')
#z                print level,var, hgt,plot_wind, hgt
    
                # Grab winds from file if plotting the barbs
                if plot_wind and var == 'wind': 
                    u = ggf.get_ua_field(grib,level,coord,shortName='u')
                    v = ggf.get_ua_field(grib,level,coord,shortName='v')
                    winds=[u,v]
                elif plot_wind and var == '10wind':
                    u = ggf.get_ua_field(grib,level,coord,shortName='10u')
                    v = ggf.get_ua_field(grib,level,coord,shortName='10v')
                    winds=[u,v]
                else:
                    winds=None
    
                height = None
                if hgt:
                    height = ggf.get_ua_field(grib,level,coord,shortName='gh')

                # Get the variable from file
                if var != 'wind' and var != '10wind':
                     field=ggf.get_ua_field(grib,level,coord,shortName=var)
                     mymap=maps.make_basemap(field,date,var,level,area_flag=outmap, line=plot_line, 
                                  winds=winds, plot_height=height,def_maps=def_maps)
                     mymap.set_bm()
                     mymap.run()
                     mymap.save_figure(outdir)

                else: 
                    # Calculate and plot  wind speed on level
                     spd=gsf.get_spd_field(grib,level,coord,outdir,outmap,fhr,var)

    
main()