A collection of diagnostic and interpolation routines for use with output from the Weather Research and Forecasting (WRF-ARW) Model.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 
 
 

11 KiB

In [ ]:
from __future__ import print_function

# This jupyter notebook command inserts matplotlib graphics in 
# to the workbook
%matplotlib inline

# Modify these to point to your own files
WRF_DIRECTORY = "/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/moving_nest"

WRF_FILES = ["wrfout_d01_2005-08-28_00:00:00",
             "wrfout_d01_2005-08-28_12:00:00",
             "wrfout_d01_2005-08-29_00:00:00"]


# Do not modify the code below this line
#------------------------------------------------------
# Turn off annoying warnings
import warnings
warnings.filterwarnings('ignore')

# Make sure the environment is good
import numpy
import cartopy
import matplotlib
from netCDF4 import Dataset
from xarray import DataArray
from wrf import (getvar, interplevel, vertcross, 
                 vinterp, ALL_TIMES)
import os

_WRF_FILES = [os.path.abspath(os.path.expanduser(
    os.path.join(WRF_DIRECTORY, f))) for f in WRF_FILES]

# Check that the WRF files exist
try:
    for f in _WRF_FILES:
        if not os.path.exists(f):
            raise ValueError("{} does not exist. "
                "Check for typos or incorrect directory.".format(f))
except ValueError:
    # If the directory ended up in the zip file, then 
    # another 'wrf_tutorial_data' directory might be 
    # used.
    WRF_DIRECTORY = os.path.join(WRF_DIRECTORY, "wrf_tutorial_data")
    _WRF_FILES = [os.path.abspath(os.path.expanduser(
        os.path.join(WRF_DIRECTORY, f))) for f in WRF_FILES]
    for f in _WRF_FILES:
        if not os.path.exists(f):
            raise

            
# Create functions so that the WRF files only need
# to be specified using the WRF_FILES global above
def single_wrf_file():
    global _WRF_FILES
    return _WRF_FILES[0]

def multiple_wrf_files():
    global _WRF_FILES
    return _WRF_FILES

print("All tests passed!")
In [ ]:
def get_kwargs(diagname):
    kwargs = {}
    if diagname == "ctt":
        kwargs = {"fill_nocloud" : True}
    
    return kwargs
In [ ]:
def contour_levels(diagname, diag, numlevels=15):
    levels = numlevels
    extend = "neither"
    if diagname == "ter":
        pass
        #levels = numpy.arange(200.,4000.,250.)
    elif diagname == "avo":
        levels = numpy.arange(10.,75.,5.)
        extend="max"
    elif diagname == "eth":
        levels = numpy.arange(270.,400.,10.)
    elif diagname == "cape_2d":
        levels = numpy.arange(200.,4000.,250.)
        extend = "max"
    elif diagname == "cape_3d":
        levels = numpy.arange(200.,4000.,250.)
        extend = "max"
    elif diagname == "ctt":
        # Note: The MP scheme doesn't produce cloud water so this
        # is just surface temperature
        levels = numpy.arange(-100, 20., 5.0)
        extend = "both"
    elif diagname == "dbz":
        #pass
        levels = numpy.arange(15.,75.,5.)
        extend = "max"
    elif diagname == "mdbz":
        levels = numpy.arange(15.,75.,5.)
        extend = "max"
    elif diagname == "geopt":
        pass
    elif diagname == "helicity":
        levels = numpy.arange(100,500,25)
        extend="max"
    elif diagname == "lat":
        pass
    elif diagname == "lon":
        pass
    elif diagname == "omg":
        pass
    elif diagname == "p":
        levels = numpy.arange(95000.,103000.,250.)
        extend = "min"
    elif diagname == "pressure":
        levels = numpy.arange(950.,1030.,2.5)
        extend = "min"
    elif diagname == "pvo":
        levels = numpy.arange(.5,5.,.25)
    elif diagname == "pw":
        #pass
        levels = numpy.arange(0.1,100,10)
    elif diagname == "rh2":
        levels = numpy.arange(50,101,5)
    elif diagname == "rh":
        levels = numpy.arange(50.,101.,5)
    elif diagname == "slp":
        levels = numpy.arange(950.,1030.,2.5)
        extend = "min"
    elif diagname == "td2":
        levels = numpy.arange(10,40,5)
        extend = "max"
    elif diagname == "td":
        levels = numpy.arange(10,40,5)
        extend = "max"
    elif diagname == "tc":
        levels = numpy.arange(10,40,5)
        extend = "max"
    elif diagname == "theta":
        levels = numpy.arange(270.,350.,2.5)
    elif diagname == "tk":
        levels = numpy.arange(270.,350.,5.)
    elif diagname == "tv":
        levels = numpy.arange(270.,350.,5.)
        #levels = numpy.arange(10,40,5)
    elif diagname == "twb":
        levels = numpy.arange(270.,350.,5.)    
    elif diagname == "updraft_helicity":
        pass
        #levels = numpy.arange(1,100,5)
    elif diagname == "ua":
        levels = numpy.arange(-40,40,5)
        extend = "both"
    elif diagname == "va":
        levels = numpy.arange(-40,40,5)
        extend = "both"
    elif diagname == "wa":
        pass
        #levels = numpy.arange(0.1,40,5)
    elif diagname == "uvmet10":
        levels = numpy.arange(-40,40,5)
        extend = "both"
    elif diagname == "uvmet":
        levels = numpy.arange(-40,40,5)
        extend = "both"
    elif diagname == "z":
        levels = numpy.arange(0,200,10)
        extend="max"
    elif diagname == "cfrac":
        levels = numpy.arange(0.0,1.1,.2)
    elif diagname == "wspd_wdir":
        levels = numpy.arange(-40,40,5)
        extend = "both"
    elif diagname == "wspd_wdir10":
        levels = numpy.arange(-40,40,5)
        extend = "both"
    elif diagname == "uvmet_wspd_wdir":
        levels = numpy.arange(-40,40,5)
        extend = "both"
    elif diagname == "uvmet10_wspd_wdir":
        levels = numpy.arange(-40,40,5)
        extend = "both"
    
    return levels, extend
In [ ]:
import numpy
from matplotlib import pyplot
from matplotlib.cm import get_cmap
from cartopy import crs
from cartopy.feature import NaturalEarthFeature
from netCDF4 import Dataset
from wrf import getvar, to_np, get_cartopy, latlon_coords, cartopy_xlim, cartopy_ylim, Constants

file_path = single_wrf_file()
wrf_file = Dataset(file_path)

for diagname in ("ter","avo", "eth", "cape_2d", "cape_3d", "ctt", "dbz", "mdbz", 
            "geopt", "helicity", "lat", "lon", "omg", "p", "pressure", 
            "pvo", "pw", "rh2", "rh", "slp", "td2", "td", "tc",
            "theta", "tk", "tv", "twb", "updraft_helicity", "ua", "va", 
            "wa", "uvmet10", "uvmet", "z", "wspd_wdir", "wspd_wdir10",
            "uvmet_wspd_wdir", "uvmet10_wspd_wdir", "cfrac"):
    
    # Get the terrain height
    print(diagname)
    kwargs = get_kwargs(diagname)
    diag = getvar(wrf_file, diagname, timeidx=3, **kwargs)

    if diag.ndim == 3:
        diag = diag[0,:]
    elif diag.ndim == 4:
        diag = diag[0,0,:]

    # Get the cartopy object and the lat,lon coords
    cart_proj = get_cartopy(diag)
    lats, lons = latlon_coords(diag)

    # Create a figure and get the GetAxes object
    fig = pyplot.figure(figsize=(10, 7.5))
    ax = pyplot.axes(projection=cart_proj)

    # Download and add the states and coastlines
    # See the cartopy documentation for more on this.
    states = NaturalEarthFeature(category='cultural', 
                                 scale='50m', 
                                 facecolor='none',
                                 name='admin_1_states_provinces_shp')

    ax.add_feature(states, linewidth=.5, edgecolor='black', zorder=3)
    ax.coastlines('50m', linewidth=.8, color='black', zorder=4)

    # Set the contour levels
    levels, extend = contour_levels(diagname, diag, numlevels=15)

    # Make the contour lines and fill them.
    pyplot.contour(to_np(lons), to_np(lats), 
                   to_np(diag), levels=levels, 
                   colors="black",
                   transform=crs.PlateCarree())
    pyplot.contourf(to_np(lons), to_np(lats), 
                    to_np(diag), levels=levels,
                    transform=crs.PlateCarree(),
                    extend=extend,
                    cmap=get_cmap("jet"))

    ax.set_xlim(cartopy_xlim(diag))
    ax.set_ylim(cartopy_ylim(diag))
             
    # Add a color bar. The shrink often needs to be set 
    # by trial and error.
    cb = pyplot.colorbar(ax=ax, shrink=.99)
    
    pyplot.title(diagname)
    pyplot.show()