diff --git a/wrf_open/var/src/python/wrf/__init__.pyc b/wrf_open/var/src/python/wrf/__init__.pyc new file mode 100644 index 0000000..3228da8 Binary files /dev/null and b/wrf_open/var/src/python/wrf/__init__.pyc differ diff --git a/wrf_open/var/src/python/wrf/var/__init__.py b/wrf_open/var/src/python/wrf/var/__init__.py index 12a5a0a..42d1417 100755 --- a/wrf_open/var/src/python/wrf/var/__init__.py +++ b/wrf_open/var/src/python/wrf/var/__init__.py @@ -159,7 +159,8 @@ _VALID_ARGS = {"cape2d" : ["missing", "timeidx"], "wspddir" : ["units", "timeidx"], "wspddir_uvmet" : ["units", "timeidx"], "wspddir_uvmet10" : ["units", "timeidx"], - "ctt" : ["timeidx"] + "ctt" : ["timeidx"], + "default" : ["timeidx"] } _ALIASES = {"cape_2d" : "cape2d", @@ -196,13 +197,20 @@ def _undo_alias(alias): return actual def _check_kargs(var, kargs): - for arg, val in kargs.iteritems(): + for arg in kargs.iterkeys(): if arg not in _VALID_ARGS[var]: raise ArgumentError("'%s' is an invalid keyword " "argument for '%s" % (arg, var)) def getvar(wrfnc, var, **kargs): + if is_standard_wrf_var(wrfnc, var): + if "timeidx" in kargs: + timeidx = kargs["timeidx"] + else: + timeidx = 0 + return extract_vars(wrfnc, timeidx, var)[var] + actual_var = _undo_alias(var) if actual_var not in _VALID_ARGS: raise ArgumentError("'%s' is not a valid variable name" % (var)) diff --git a/wrf_open/var/src/python/wrf/var/util.py b/wrf_open/var/src/python/wrf/var/util.py index 892d943..568b472 100644 --- a/wrf_open/var/src/python/wrf/var/util.py +++ b/wrf_open/var/src/python/wrf/var/util.py @@ -3,7 +3,7 @@ from collections import Iterable import numpy as n __all__ = ["extract_vars", "extract_global_attrs", "extract_dim", - "combine_files"] + "combine_files", "is_standard_wrf_var"] def _is_multi_time(timeidx): if timeidx == -1: @@ -80,6 +80,12 @@ def extract_vars(wrfnc, timeidx, vars): return {var:wrfnc.variables[var][:] for var in varlist} else: return {var:combine_files(wrfnc, var, timeidx) for var in varlist} + +def is_standard_wrf_var(wrfnc, var): + multifile = _is_multi_file(wrfnc) + if multifile: + wrfnc = wrfnc[0] + return var in wrfnc.variables diff --git a/wrf_open/var/test/utests.py b/wrf_open/var/test/utests.py index 776f025..cc7cfbc 100644 --- a/wrf_open/var/test/utests.py +++ b/wrf_open/var/test/utests.py @@ -5,13 +5,8 @@ import numpy.ma as ma import os, sys import subprocess -import Ngl, Nio from wrf.var import getvar -# This should work with Nio -from netCDF4 import Dataset as NetCDF -from boto.ec2.instancestatus import Status - NCL_EXE = "/Users/ladwig/nclbuild/6.3.0/bin/ncl" TEST_FILE = "/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00" OUT_NC_FILE = "/tmp/wrftest.nc" @@ -45,14 +40,40 @@ def setUpModule(): # Using helpful information at: # http://eli.thegreenplace.net/2014/04/02/dynamically-generating-python-test-cases -def make_test(varname, wrf_in, referent, multi=False, repeat=3): +def make_test(varname, wrf_in, referent, multi=False, repeat=3, pynio=False): def test(self): + try: + from netCDF4 import Dataset as NetCDF + except: + pass + + try: + from PyNIO import Nio + except: + pass if not multi: - in_wrfnc = NetCDF(wrf_in) + if not pynio: + in_wrfnc = NetCDF(wrf_in) + else: + # Note: Python doesn't like it if you reassign an outer scoped + # variable (wrf_in in this case) + if not wrf_in.endswith(".nc"): + wrf_file = wrf_in + ".nc" + else: + wrf_file = wrf_in + in_wrfnc = Nio.open_file(wrf_file) else: - nc = NetCDF(wrf_in) - in_wrfnc = [nc for i in xrange(repeat)] + if not pynio: + nc = NetCDF(wrf_in) + in_wrfnc = [nc for i in xrange(repeat)] + else: + if not wrf_in.endswith(".nc"): + wrf_file = wrf_in + ".nc" + else: + wrf_file = wrf_in + nc = Nio.open_file(wrf_file) + in_wrfnc = [nc for i in xrange(repeat)] refnc = NetCDF(referent) @@ -161,15 +182,35 @@ if __name__ == "__main__": "theta", "tk", "tv", "twb", "updraft_helicity", "ua", "va", "wa", "uvmet10", "uvmet", "z", "ctt", "cape_2d", "cape_3d"] - for var in wrf_vars: - if var in ignore_vars: - continue + try: + import netCDF4 + except ImportError: + pass + else: + for var in wrf_vars: + if var in ignore_vars: + continue + + test_func1 = make_test(var, TEST_FILE, OUT_NC_FILE) + test_func2 = make_test(var, TEST_FILE, OUT_NC_FILE, multi=True) + setattr(WRFVarsTest, 'test_{0}'.format(var), test_func1) + setattr(WRFVarsTest, 'test_multi_{0}'.format(var), test_func2) + + try: + import PyNIO + except ImportError: + pass + else: + for var in wrf_vars: + if var in ignore_vars: + continue + + test_func1 = make_test(var, TEST_FILE, OUT_NC_FILE,pynio=True) + test_func2 = make_test(var, TEST_FILE, OUT_NC_FILE, multi=True, + pynio=True) + setattr(WRFVarsTest, 'test_pynio_{0}'.format(var), test_func1) + setattr(WRFVarsTest, 'test_pynio_multi_{0}'.format(var), + test_func2) - test_func1 = make_test(var, TEST_FILE, OUT_NC_FILE) - test_func2 = make_test(var, TEST_FILE, OUT_NC_FILE, multi=True) - setattr(WRFVarsTest, 'test_{0}'.format(var), test_func1) - setattr(WRFVarsTest, 'test_multi_{0}'.format(var), test_func2) - - ut.main() \ No newline at end of file