def plot_excess_temp_map(input):
    DATE = input

    date = DATE.strftime('%Y-%m-%d_%H%M')
    print 'working on', date

    DIR = '/uufs/chpc.utah.edu/common/home/horel-group/archive/%04d%02d%02d/models/' % (
        DATE.year, DATE.month, DATE.day)

    # Open HRRR and get surface temperature and surface pressure
    try:
        grbs = pygrib.open(DIR + 'hrrrX/hrrrX.t%02dz.wrfsfcf00.grib2' %
                           (DATE.hour))
        print "opened", DATE

        temp = grbs.select(name="2 metre temperature")[0].values[
            subset[domain][0]:subset[domain][1],
            subset[domain][2]:subset[domain][3]] - 273.15
        sfcpres = grbs.select(name="Surface pressure")[0].values[
            subset[domain][0]:subset[domain][1],
            subset[domain][2]:subset[domain][3]] / 100
        lat, lon = grbs.select(name="2 metre temperature")[0].latlons()
        lat = lat[subset[domain][0]:subset[domain][1],
                  subset[domain][2]:subset[domain][3]]
        lon = lon[subset[domain][0]:subset[domain][1],
                  subset[domain][2]:subset[domain][3]]

        # Get the 700 hPa temperatures (!! NOT AVAILABLE IN HRRR-X FILE)
        #T700 = grbs.select(name="Temperature")[1].values[subset[domain][0]:subset[domain][1], subset[domain][2]:subset[domain][3]] - 273.15

        # Convert temperatures to potental temperature (theta)
        theta_sfc = TempPress_to_PotTemp(temp, sfcpres)
        #theta_700 = TempPress_to_PotTemp(T700, 700)

        # pluck the value for SLC
        SLC_lat = 40.77
        SLC_lon = -111.95
        pluck_index = pluck_point(SLC_lat, SLC_lon, lat, lon)

        # use that pluck index on the flattened array to ge the value
        SLC_sfc_temp = temp.flatten()[pluck_index]
        SLC_sfc_theta = theta_sfc.flatten()[pluck_index]
        #SLC_700_temp = T700.flatten()[pluck_index]
        #SLC_700_theta = theta_700.flatten()[pluck_index]

        return [SLC_sfc_temp, SLC_sfc_theta]

    except:
        print "could not open", DATE
        return [np.nan, np.nan]
def get_theta_from_grb(DATE):
    """
    Gets the potential temperature at a point from the HRRR surface field
    files for a single time.
    """

    VALID_DATE = DATE + timedelta(hours=fxx)

    DIR = '/uufs/chpc.utah.edu/common/home/horel-group/archive/%04d%02d%02d/models/' % (
        DATE.year, DATE.month, DATE.day)

    # Open HRRR and get surface temperature and surface pressure
    grbs = pygrib.open(DIR + 'hrrr/hrrr.t%02dz.wrfsfcf%02d.grib2' %
                       (DATE.hour, fxx))
    temp = grbs.select(name="2 metre temperature")[0].values[
        subset[domain][0]:subset[domain][1],
        subset[domain][2]:subset[domain][3]] - 273.15
    sfcpres = grbs.select(name="Surface pressure")[0].values[
        subset[domain][0]:subset[domain][1],
        subset[domain][2]:subset[domain][3]] / 100
    lat, lon = grbs.select(name="2 metre temperature")[0].latlons()
    lat = lat[subset[domain][0]:subset[domain][1],
              subset[domain][2]:subset[domain][3]]
    lon = lon[subset[domain][0]:subset[domain][1],
              subset[domain][2]:subset[domain][3]]
    # Get the 700 hPa temperatures
    T700 = grbs.select(name="Temperature")[1].values[
        subset[domain][0]:subset[domain][1],
        subset[domain][2]:subset[domain][3]] - 273.15

    # Convert temperatures to potental temperature (theta)
    theta_sfc = TempPress_to_PotTemp(temp, sfcpres)
    theta_700 = TempPress_to_PotTemp(T700, 700)

    # Pluck value for the station location
    pluck_index = pluck_point(stn_lat, stn_lon, lat, lon)

    # use that pluck index on the flattened array to ge the value
    stn_theta_sfc = theta_sfc.flatten()[pluck_index]
    stn_theta_700 = theta_700.flatten()[pluck_index]

    # Return the values
    print "Got it:", DATE

    # close the grib file
    grbs.close()

    return [VALID_DATE, stn_theta_sfc, stn_theta_700]
        heights = np.column_stack([heights, b['hght'][fxx]])
        press = np.column_stack([press, b['pres'][fxx]])
        temps = np.column_stack([temps, b['temp'][fxx]])
        dwpts = np.column_stack([dwpts, b['dwpt'][fxx]])
    except:
        # It looks like the variable hasn't been created yet, so create it.
        heights = np.array(b['hght'][fxx])
        press = np.array(b['pres'][fxx])
        temps = np.array(b['temp'][fxx])
        dwpts = np.array(b['dwpt'][fxx])

    print 'Got it: f%02d' % (fxx), DATE
    DATE += timedelta(hours=1)

# Derive a few more variables
thetas = TempPress_to_PotTemp(temps, press)
RHs = Tempdwpt_to_RH(temps, dwpts)

# Calculate difference between potental temperture of columns and surface theta
sfcThetas = thetas[0]
diffThetas = thetas - sfcThetas

# Because the contour funciton doesn't like dates, need to convert dates
# to some other happy x axis number. Then the x axis dates can be formatted.
x = matplotlib.dates.date2num(date_list)
x2D = x * np.ones_like(heights)

# === Begin Plots ===

#==============================================================================
#                     Plot: Potential Temperature
Example #4
0
# Get SLC sounding for Salt Lake City
balloon = get_wyoming_sounding(DATE)

# Get WRF from TS list profile data
station = 'KSLC'
wrf_dir = '/uufs/chpc.utah.edu/common/home/horel-group4/model/bblaylock/WRF3.8.1_MYNN/DATA/'
tsfile = station + '.d02.TS'
model_start = datetime(2015, 6, 14, 0)
wrf = get_vert_data(wrf_dir+tsfile,model_start,DATE)

# Get HRRR analysis profile from *modified* BUFR file 
# (modified so all data was on one row for easy reading)
hrrr_file = 'kslc_' + DATE.strftime('%Y%m%d%H') + '.txt'
HRRR = np.genfromtxt(hrrr_file, delimiter=' ', names=True)
HRRR_theta = TempPress_to_PotTemp(HRRR['TMPC'], HRRR['PRES'])
HRRR_u, HRRR_v = wind_calcs.wind_spddir_to_uv(HRRR['SKNT']*0.51444, HRRR['DRCT'])
HRRR_mixr = DwptPress_to_MixRatio(HRRR['DWPC'], HRRR['PRES'])

# Get HRRR 1-hour forecast
f = 'F01'
hrrr_file = 'kslc_'+datetime(2015,6,18,23).strftime('%Y%m%d%H')+'_'+f+'.txt'
HRRR_f01 = np.genfromtxt(hrrr_file, delimiter=' ', names=True)
HRRR_f01_theta = TempPress_to_PotTemp(HRRR_f01['TMPC'], HRRR_f01['PRES'])
HRRR_f01_u, HRRR_f01_v = wind_calcs.wind_spddir_to_uv(HRRR_f01['SKNT']*0.51444, HRRR_f01['DRCT'])
HRRR_f01_mixr = DwptPress_to_MixRatio(HRRR_f01['DWPC'], HRRR_f01['PRES'])

"""
Create 3 subplots for 
    1. Potential Temperature
    2. Mixing Ratio
def plot_excess_temp_map(input):
    m = input[0]
    DATE = input[1]

    date = DATE.strftime('%Y-%m-%d_%H%M')
    print 'working on', date

    DIR = '/uufs/chpc.utah.edu/common/home/horel-group/archive/%04d%02d%02d/models/' % (
        DATE.year, DATE.month, DATE.day)

    # Open HRRR and get surface temperature and surface pressure
    grbs = pygrib.open(DIR + 'hrrr/hrrr.t%02dz.wrfsfcf00.grib2' % (DATE.hour))
    temp = grbs.select(name="2 metre temperature")[0].values[
        subset[domain][0]:subset[domain][1],
        subset[domain][2]:subset[domain][3]] - 273.15
    sfcpres = grbs.select(name="Surface pressure")[0].values[
        subset[domain][0]:subset[domain][1],
        subset[domain][2]:subset[domain][3]] / 100
    lat, lon = grbs.select(name="2 metre temperature")[0].latlons()
    lat = lat[subset[domain][0]:subset[domain][1],
              subset[domain][2]:subset[domain][3]]
    lon = lon[subset[domain][0]:subset[domain][1],
              subset[domain][2]:subset[domain][3]]

    # Get the 700 hPa temperatures
    T700 = grbs.select(name="Temperature")[1].values[
        subset[domain][0]:subset[domain][1],
        subset[domain][2]:subset[domain][3]] - 273.15

    # Convert temperatures to potental temperature (theta)
    theta_sfc = TempPress_to_PotTemp(temp, sfcpres)
    theta_700 = TempPress_to_PotTemp(T700, 700)

    # Draw map and plot temperatures
    fig = plt.figure(figsize=[13, 5])
    # first pannel
    fig.add_subplot(131)

    m.drawstates()
    m.drawcoastlines()
    m.drawcounties(zorder=100)
    m.pcolormesh(lon, lat, theta_sfc, cmap='Spectral_r', vmax=305, vmin=270)
    cb = plt.colorbar(orientation='horizontal',
                      shrink=.9,
                      pad=.03,
                      ticks=range(270, 306, 5),
                      extend="both")
    cb.set_label('Potential Temperature')
    plt.title('Surface theta')

    # second pannel
    fig.add_subplot(132)

    m.drawstates()
    m.drawcoastlines()
    m.drawcounties(zorder=100)
    m.pcolormesh(lon, lat, theta_700, cmap='Spectral_r', vmax=305, vmin=270)
    cb = plt.colorbar(orientation='horizontal',
                      shrink=.9,
                      pad=.03,
                      ticks=range(270, 306, 5),
                      extend="both")
    cb.set_label('Potential Temperature')
    plt.title('700 hPa theta')

    # third pannel
    ax = fig.add_subplot(133)

    m.drawstates()
    m.drawcoastlines()
    m.drawcounties(zorder=100)
    m.pcolormesh(lon, lat, theta_700 - theta_sfc, cmap='Reds', vmax=30, vmin=0)
    cb = plt.colorbar(orientation='horizontal',
                      shrink=.9,
                      pad=.03,
                      ticks=range(0, 31, 5),
                      extend='max')
    cb.set_label('Temperature Excess')
    plt.title('theta 700 - theta sfc')
    plt.suptitle('HRRR anlys: ' + str(DATE))

    # savefig
    plt.savefig('./figs/excess_temp/' + date)