From 230a532061547b59b96eb1802f876e7880b21911 Mon Sep 17 00:00:00 2001 From: Bill Ladwig Date: Fri, 4 Dec 2015 12:34:18 -0700 Subject: [PATCH] now uses extract methods to pull variables and attributes out of netcdf files --- wrf_open/var/src/python/wrf/var/__init__.py | 3 + wrf_open/var/src/python/wrf/var/cape.py | 39 ++--- wrf_open/var/src/python/wrf/var/ctt.py | 35 +++-- wrf_open/var/src/python/wrf/var/dbz.py | 32 +++-- wrf_open/var/src/python/wrf/var/decorators.py | 44 +++++- wrf_open/var/src/python/wrf/var/destagger.py | 9 +- wrf_open/var/src/python/wrf/var/dewpoint.py | 14 +- wrf_open/var/src/python/wrf/var/etaconv.py | 1 + wrf_open/var/src/python/wrf/var/geoht.py | 25 ++-- wrf_open/var/src/python/wrf/var/helicity.py | 92 ++++++++---- wrf_open/var/src/python/wrf/var/latlon.py | 126 +++++++++-------- wrf_open/var/src/python/wrf/var/omega.py | 12 +- wrf_open/var/src/python/wrf/var/precip.py | 17 ++- wrf_open/var/src/python/wrf/var/pressure.py | 18 ++- wrf_open/var/src/python/wrf/var/pw.py | 15 +- wrf_open/var/src/python/wrf/var/rh.py | 18 +-- wrf_open/var/src/python/wrf/var/slp.py | 15 +- wrf_open/var/src/python/wrf/var/temp.py | 41 +++--- wrf_open/var/src/python/wrf/var/terrain.py | 18 ++- wrf_open/var/src/python/wrf/var/times.py | 1 + wrf_open/var/src/python/wrf/var/units.py | 3 +- wrf_open/var/src/python/wrf/var/util.py | 77 ++++++++++ wrf_open/var/src/python/wrf/var/uvmet.py | 133 +++++++++++++----- wrf_open/var/src/python/wrf/var/vorticity.py | 52 ++++--- wrf_open/var/src/python/wrf/var/wind.py | 7 +- 25 files changed, 581 insertions(+), 266 deletions(-) create mode 100644 wrf_open/var/src/python/wrf/var/util.py diff --git a/wrf_open/var/src/python/wrf/var/__init__.py b/wrf_open/var/src/python/wrf/var/__init__.py index be75209..12a5a0a 100755 --- a/wrf_open/var/src/python/wrf/var/__init__.py +++ b/wrf_open/var/src/python/wrf/var/__init__.py @@ -2,6 +2,8 @@ import warnings from extension import * import extension +from util import * +import util from cape import * import cape from constants import * @@ -55,6 +57,7 @@ import units __all__ = ["getvar"] __all__ += extension.__all__ +__all__ += util.__all__ __all__ += cape.__all__ __all__ += constants.__all__ __all__ += ctt.__all__ diff --git a/wrf_open/var/src/python/wrf/var/cape.py b/wrf_open/var/src/python/wrf/var/cape.py index 7294ea8..cbec7f1 100755 --- a/wrf_open/var/src/python/wrf/var/cape.py +++ b/wrf_open/var/src/python/wrf/var/cape.py @@ -3,19 +3,23 @@ import numpy.ma as ma from wrf.var.extension import computetk,computecape from wrf.var.destagger import destagger from wrf.var.constants import Constants, ConversionFactors +from wrf.var.util import extract_vars __all__ = ["get_2dcape", "get_3dcape"] def get_2dcape(wrfnc, missing=-999999.0, timeidx=0): """Return the 2d fields of cape, cin, lcl, and lfc""" - t = wrfnc.variables["T"][timeidx,:,:,:] - p = wrfnc.variables["P"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] - qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:] - ph = wrfnc.variables["PH"][timeidx,:,:,:] - phb = wrfnc.variables["PHB"][timeidx,:,:,:] - ter = wrfnc.variables["HGT"][timeidx,:,:] - psfc = wrfnc.variables["PSFC"][timeidx,:,:] + vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", "PH", + "PHB", "HGT", "PSFC")) + + t = vars["T"] + p = vars["P"] + pb = vars["PB"] + qv = vars["QVAPOR"] + ph = vars["PH"] + phb = vars["PHB"] + ter = vars["HGT"] + psfc = vars["PSFC"] full_t = t + Constants.T_BASE full_p = p + pb @@ -49,14 +53,17 @@ def get_2dcape(wrfnc, missing=-999999.0, timeidx=0): def get_3dcape(wrfnc, missing=-999999.0, timeidx=0): """Return the 3d fields of cape and cin""" - t = wrfnc.variables["T"][timeidx,:,:,:] - p = wrfnc.variables["P"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] - qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:] - ph = wrfnc.variables["PH"][timeidx,:,:,:] - phb = wrfnc.variables["PHB"][timeidx,:,:,:] - ter = wrfnc.variables["HGT"][timeidx,:,:] - psfc = wrfnc.variables["PSFC"][timeidx,:,:] + + vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", "PH", + "PHB", "HGT", "PSFC")) + t = vars["T"] + p = vars["P"] + pb = vars["PB"] + qv = vars["QVAPOR"] + ph = vars["PH"] + phb = vars["PHB"] + ter = vars["HGT"] + psfc = vars["PSFC"] full_t = t + Constants.T_BASE full_p = p + pb diff --git a/wrf_open/var/src/python/wrf/var/ctt.py b/wrf_open/var/src/python/wrf/var/ctt.py index ac49d7b..37aec68 100644 --- a/wrf_open/var/src/python/wrf/var/ctt.py +++ b/wrf_open/var/src/python/wrf/var/ctt.py @@ -5,6 +5,7 @@ from wrf.var.extension import computectt, computetk from wrf.var.constants import Constants, ConversionFactors from wrf.var.destagger import destagger from wrf.var.decorators import convert_units +from wrf.var.util import extract_vars __all__ = ["get_ctt"] @@ -13,25 +14,31 @@ def get_ctt(wrfnc, units="c", timeidx=0): """Return the cloud top temperature. """ - t = wrfnc.variables["T"][timeidx,:,:,:] - p = wrfnc.variables["P"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] - ph = wrfnc.variables["PH"][timeidx,:,:,:] - phb = wrfnc.variables["PHB"][timeidx,:,:,:] - ter = wrfnc.variables["HGT"][timeidx,:,:] - qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:] * 1000.0 # g/kg + vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "PH" ,"PHB", + "HGT", "QVAPOR")) + t = vars["T"] + p = vars["P"] + pb = vars["PB"] + ph = vars["PH"] + phb = vars["PHB"] + ter = vars["HGT"] + qv = vars["QVAPOR"] * 1000.0 # g/kg haveqci = 1 - if "QICE" in wrfnc.variables: - qice = wrfnc.variables["QICE"][timeidx,:,:,:] * 1000.0 #g/kg - else: + try: + icevars = extract_vars(wrfnc, timeidx, vars="QICE") + except KeyError: qice = n.zeros(qv.shape, qv.dtype) haveqci = 0 - - if "QCLOUD" in wrfnc.variables: - qcld = wrfnc.variables["QCLOUD"][timeidx,:,:,:] * 1000.0 #g/kg else: + qice = icevars["QICE"] * 1000.0 #g/kg + + try: + cldvars = extract_vars(wrfnc, timeidx, vars="QCLOUD") + except KeyError: raise RuntimeError("'QCLOUD' not found in NetCDF file") + else: + qcld = cldvars["QCLOUD"] * 1000.0 #g/kg full_p = p + pb p_hpa = full_p * ConversionFactors.PA_TO_HPA @@ -44,4 +51,4 @@ def get_ctt(wrfnc, units="c", timeidx=0): ctt = computectt(p_hpa,tk,qice,qcld,qv,ght,ter,haveqci) - return ctt \ No newline at end of file + return ctt diff --git a/wrf_open/var/src/python/wrf/var/dbz.py b/wrf_open/var/src/python/wrf/var/dbz.py index c464845..8ee4454 100755 --- a/wrf_open/var/src/python/wrf/var/dbz.py +++ b/wrf_open/var/src/python/wrf/var/dbz.py @@ -2,6 +2,7 @@ import numpy as n from wrf.var.extension import computedbz,computetk from wrf.var.constants import Constants +from wrf.var.util import extract_vars __all__ = ["get_dbz", "get_max_dbz"] @@ -16,22 +17,27 @@ def get_dbz(wrfnc, do_varint=False, do_liqskin=False, timeidx=0): as liquid) """ - t = wrfnc.variables["T"][timeidx,:,:,:] - p = wrfnc.variables["P"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] + vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", + "QRAIN")) + t = vars["T"] + p = vars["P"] + pb = vars["PB"] + qv = vars["QVAPOR"] + qr = vars["QRAIN"] - qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:] - qr = wrfnc.variables["QRAIN"][timeidx,:,:,:] - - if "QSNOW" in wrfnc.variables: - qs = wrfnc.variables["QSNOW"][timeidx,:,:,:] + try: + snowvars = extract_vars(wrfnc, timeidx, vars="QSNOW") + except KeyError: + qs = n.zeros(qv.shape, "float") else: - qs = n.zeros((qv.shape[0], qv.shape[1], qv.shape[2]), "float") - - if "QGRAUP" in wrfnc.variables: - qg = wrfnc.variables["QGRAUP"][timeidx,:,:,:] + qs = snowvars["QSNOW"] + + try: + graupvars = extract_vars(wrfnc, timeidx, vars="QGRAUP") + except KeyError: + qg = n.zeros(qv.shape, "float") else: - qg = n.zeros((qv.shape[0], qv.shape[1], qv.shape[2]), "float") + qg = graupvars["QGRAUP"] # If qsnow is all 0, set sn0 to 1 sn0 = 0 diff --git a/wrf_open/var/src/python/wrf/var/decorators.py b/wrf_open/var/src/python/wrf/var/decorators.py index e096efa..f0dfcd7 100644 --- a/wrf_open/var/src/python/wrf/var/decorators.py +++ b/wrf_open/var/src/python/wrf/var/decorators.py @@ -1,5 +1,6 @@ from functools import wraps from inspect import getargspec +from collections import Iterable from wrf.var.units import do_conversion, check_units @@ -37,7 +38,48 @@ def combine_list_and_times(alg_out_dim): def combine_decorator(func): @wraps(func) def func_wrapper(*args, **kargs): - argspec = getargspec(func) + # Multiple times? + multitime = False + if "timeidx" in kargs: + if kargs["timeidx"] == 0: + multitime = True + + # Multiple files? + multifile = False + if isinstance(args[0], Iterable) and not isinstance(args[0], str): + multifile = True + + # Single file, single time + if not multitime and not multifile: + return func(*args, **kargs) + + # Get the dimensions + if multifile: + wrffile = args[0][0] + else: + wrffile = args[0] + + # TODO: Add PyNIO support + dims = wrffile.dimensions + we_size = len(dims["west_east"]) + sn_size = len(dims["south_north"]) + bt_size = len(dims["bottom_top"]) + time_size = len(dims["Time"]) + + + if alg_out_dim == 2: + pass + elif algout_dim == 3: + pass + elif algout_dim == 4: + pass + else: + raise RuntimeError("invalid algorithm output dimsize") + + + + + return func(*args, **kargs) return func_wrapper diff --git a/wrf_open/var/src/python/wrf/var/destagger.py b/wrf_open/var/src/python/wrf/var/destagger.py index 9b7b57e..43702e1 100755 --- a/wrf_open/var/src/python/wrf/var/destagger.py +++ b/wrf_open/var/src/python/wrf/var/destagger.py @@ -9,7 +9,8 @@ def destagger(var, stagger_dim): Arguments: - var is a numpy array for the variable - stagger_dim is the dimension of the numpy array to de-stagger - (e.g. 0, 1, 2) + (e.g. 0, 1, 2). Note: negative values are acceptable to choose + a dimensions from the right hand side (e.g. -1, -2, -3) """ var_shape = var.shape @@ -41,13 +42,13 @@ def destagger(var, stagger_dim): def destagger_windcomp(wrfnc, comp, timeidx=0): if comp.lower() == "u": wrfvar = "U" - stagdim = 2 + stagdim = -1 elif comp.lower() == "v": wrfvar = "V" - stagdim = 1 + stagdim = -2 elif comp.lower() == "w": wrfvar = "W" - stagdim = 0 + stagdim = -3 wind_data = wrfnc.variables[wrfvar][timeidx,:,:,:] return destagger(wind_data, stagdim) diff --git a/wrf_open/var/src/python/wrf/var/dewpoint.py b/wrf_open/var/src/python/wrf/var/dewpoint.py index 258d333..7fc92a0 100755 --- a/wrf_open/var/src/python/wrf/var/dewpoint.py +++ b/wrf_open/var/src/python/wrf/var/dewpoint.py @@ -1,14 +1,17 @@ from wrf.var.extension import computetd from wrf.var.decorators import convert_units +from wrf.var.util import extract_vars __all__ = ["get_dp", "get_dp_2m"] @convert_units("temp", "c") def get_dp(wrfnc, units="c", timeidx=0): - p = wrfnc.variables["P"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] - qvapor = wrfnc.variables["QVAPOR"][timeidx,:,:,:] + vars = extract_vars(wrfnc, timeidx, vars=("P", "PB", "QVAPOR")) + + p = vars["P"] + pb = vars["PB"] + qvapor = vars["QVAPOR"] # Algorithm requires hPa full_p = .01*(p + pb) @@ -19,10 +22,11 @@ def get_dp(wrfnc, units="c", timeidx=0): @convert_units("temp", "c") def get_dp_2m(wrfnc, units="c", timeidx=0): + vars = extract_vars(wrfnc, timeidx, vars=("PSFC", "Q2")) # Algorithm requires hPa - psfc = .01*(wrfnc.variables["PSFC"][timeidx,:,:]) - q2 = wrfnc.variables["Q2"][timeidx,:,:] + psfc = .01*(vars["PSFC"]) + q2 = vars["Q2"] q2[q2 < 0] = 0 td = computetd(psfc, q2) diff --git a/wrf_open/var/src/python/wrf/var/etaconv.py b/wrf_open/var/src/python/wrf/var/etaconv.py index 829ee8b..4b2885a 100755 --- a/wrf_open/var/src/python/wrf/var/etaconv.py +++ b/wrf_open/var/src/python/wrf/var/etaconv.py @@ -3,6 +3,7 @@ import numpy as n from wrf.var.extension import computeeta from wrf.var.constants import Constants from wrf.var.decorators import convert_units +from wrf.var.util import extract_vars #__all__ = ["convert_eta"] __all__ = [] diff --git a/wrf_open/var/src/python/wrf/var/geoht.py b/wrf_open/var/src/python/wrf/var/geoht.py index e45b6fa..612a2ac 100755 --- a/wrf_open/var/src/python/wrf/var/geoht.py +++ b/wrf_open/var/src/python/wrf/var/geoht.py @@ -1,6 +1,7 @@ from wrf.var.constants import Constants from wrf.var.destagger import destagger from wrf.var.decorators import convert_units +from wrf.var.util import extract_vars __all__ = ["get_geopt", "get_height"] @@ -11,16 +12,24 @@ def _get_geoht(wrfnc, height=True, msl=True, timeidx=0): height is subtracted). """ - - if "PH" in wrfnc.variables: - ph = wrfnc.variables["PH"][timeidx,:,:,:] - phb = wrfnc.variables["PHB"][timeidx,:,:,:] - hgt = wrfnc.variables["HGT"][timeidx,:,:] + + try: + ph_vars = extract_vars(wrfnc, timeidx, ("PH", "PHB", "HGT")) + except KeyError: + try: + ght_vars = extract_vars(wrfnc, timeidx, ("GHT", "HGT_U")) + except KeyError: + raise RuntimeError("Cannot calculate height with variables in " + "NetCDF file") + else: + geopt_unstag = ght_vars["GHT"] * Constants.G + hgt = destagger(ght_vars["HGT_U"], 1) + else: + ph = ph_vars["PH"] + phb = ph_vars["PHB"] + hgt = ph_vars["HGT"] geopt = ph + phb geopt_unstag = destagger(geopt, 0) - elif "GHT" in wrfnc.variables: # met_em files - geopt_unstag = wrfnc.variables["GHT"][timeidx,:,:,:] * Constants.G - hgt = destagger(wrfnc.variables["HGT_U"][timidx,:,:], 1) if height: if msl: diff --git a/wrf_open/var/src/python/wrf/var/helicity.py b/wrf_open/var/src/python/wrf/var/helicity.py index 96b9fb4..25fb277 100755 --- a/wrf_open/var/src/python/wrf/var/helicity.py +++ b/wrf_open/var/src/python/wrf/var/helicity.py @@ -2,26 +2,43 @@ from wrf.var.constants import Constants from wrf.var.extension import computesrh, computeuh from wrf.var.destagger import destagger +from wrf.var.util import extract_vars, extract_global_attrs __all__ = ["get_srh", "get_uh"] def get_srh(wrfnc, top=3000.0, timeidx=0): # Top can either be 3000 or 1000 (for 0-1 srh or 0-3 srh) - if "U" in wrfnc.variables: - u = destagger(wrfnc.variables["U"][timeidx,:,:,:], 2) - elif "UU" in wrfnc.variables: - u = destagger(wrfnc.variables["UU"][timeidx,:,:,:], 2) # support met_em files - - if "V" in wrfnc.variables: - v = destagger(wrfnc.variables["V"][timeidx,:,:,:], 1) - elif "VV" in wrfnc.variables: - v = destagger(wrfnc.variables["VV"][timeidx,:,:,:], 1) + vars = extract_vars(wrfnc, timeidx, vars=("HGT", "PH", "PHB")) - ter = wrfnc.variables["HGT"][timeidx,:,:] - ph = wrfnc.variables["PH"][timeidx,:,:,:] - phb = wrfnc.variables["PHB"][timeidx,:,:,:] + ter = vars["HGT"] + ph = vars["PH"] + phb = vars["PHB"] + try: + u_vars = extract_vars(wrfnc, timeidx, vars="U") + except KeyError: + try: + uu_vars = extract_vars(wrfnc, timeidx, vars="UU") + except KeyError: + raise RuntimeError("No valid wind data found in NetCDF file") + else: + u = destagger(uu_vars["UU"], -1) # support met_em files + else: + u = destagger(u_vars["U"], -1) + + try: + v_vars = extract_vars(wrfnc, timeidx, vars="V") + except KeyError: + try: + vv_vars = extract_vars(wrfnc, timeidx, vars="VV") + except KeyError: + raise RuntimeError("No valid wind data found in NetCDF file") + else: + v = destagger(vv_vars["VV"], -2) # support met_em files + else: + v = destagger(v_vars["V"], -2) + geopt = ph + phb geopt_unstag = destagger(geopt, 0) @@ -38,26 +55,43 @@ def get_srh(wrfnc, top=3000.0, timeidx=0): def get_uh(wrfnc, bottom=2000.0, top=5000.0, timeidx=0): - if "U" in wrfnc.variables: - u = destagger(wrfnc.variables["U"][timeidx,:,:,:], 2) - elif "UU" in wrfnc.variables: - u = destagger(wrfnc.variables["UU"][timeidx,:,:,:], 2) # support met_em files + vars = extract_vars(wrfnc, timeidx, vars=("W", "PH", "PHB", "MAPFAC_M")) + + wstag = vars["W"] + ph = vars["PH"] + phb = vars["PHB"] + mapfct = vars["MAPFAC_M"] + + attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY")) + dx = attrs["DX"] + dy = attrs["DY"] + + try: + u_vars = extract_vars(wrfnc, timeidx, vars="U") + except KeyError: + try: + uu_vars = extract_vars(wrfnc, timeidx, vars="UU") + except KeyError: + raise RuntimeError("No valid wind data found in NetCDF file") + else: + u = destagger(uu_vars["UU"], -1) # support met_em files + else: + u = destagger(u_vars["U"], -1) - if "V" in wrfnc.variables: - v = destagger(wrfnc.variables["V"][timeidx,:,:,:], 1) - elif "VV" in wrfnc.variables: - v = destagger(wrfnc.variables["VV"][timeidx,:,:,:], 1) - - wstag = wrfnc.variables["W"][timeidx,:,:,:] - ph = wrfnc.variables["PH"][timeidx,:,:,:] - phb = wrfnc.variables["PHB"][timeidx,:,:,:] + try: + v_vars = extract_vars(wrfnc, timeidx, vars="V") + except KeyError: + try: + vv_vars = extract_vars(wrfnc, timeidx, vars="VV") + except KeyError: + raise RuntimeError("No valid wind data found in NetCDF file") + else: + v = destagger(vv_vars["VV"], -2) # support met_em files + else: + v = destagger(v_vars["V"], -2) + zp = ph + phb - mapfct = wrfnc.variables["MAPFAC_M"][timeidx,:,:] - dx = wrfnc.getncattr("DX") - dy = wrfnc.getncattr("DY") - - uh = computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top) return uh diff --git a/wrf_open/var/src/python/wrf/var/latlon.py b/wrf_open/var/src/python/wrf/var/latlon.py index 1360f01..2564353 100755 --- a/wrf_open/var/src/python/wrf/var/latlon.py +++ b/wrf_open/var/src/python/wrf/var/latlon.py @@ -1,35 +1,59 @@ from wrf.var.extension import computeij, computell +from wrf.var.util import extract_vars, extract_global_attrs __all__ = ["get_lat", "get_lon", "get_ij", "get_ll"] def get_lat(wrfnc, timeidx=0): - if "XLAT" in wrfnc.variables: - xlat = wrfnc.variables["XLAT"][timeidx,:,:] - elif "XLAT_M" in wrfnc.variables: - xlat = wrfnc.variables["XLAT_M"][timeidx,:,:] + + try: + lat_vars = extract_vars(wrfnc, timeidx, vars="XLAT") + except KeyError: + try: + latm_vars = extract_vars(wrfnc, timeidx, vars="XLAT_M") + except: + raise RuntimeError("Latitude variable not found in NetCDF file") + else: + xlat = latm_vars["XLAT_M"] + else: + xlat = lat_vars["XLAT"] return xlat def get_lon(wrfnc, timeidx=0): - if "XLONG" in wrfnc.variables: - xlon = wrfnc.variables["XLONG"][timeidx,:,:] - elif "XLONG_M" in wrfnc.variables: - xlon = wrfnc.variables["XLONG_M"][timeidx,:,:] + try: + lon_vars = extract_vars(wrfnc, timeidx, vars="XLONG") + except KeyError: + try: + lonm_vars = extract_vars(wrfnc, timeidx, vars="XLONG_M") + except: + raise RuntimeError("Latitude variable not found in NetCDF file") + else: + xlon = lonm_vars["XLONG_M"] + else: + xlon = lon_vars["XLONG"] return xlon -def get_ij(wrfnc, longitude, latitude, timeidx=0): - map_proj = wrfnc.getncattr("MAP_PROJ") - truelat1 = wrfnc.getncattr("TRUELAT1") - truelat2 = wrfnc.getncattr("TRUELAT2") - stdlon = wrfnc.getncattr("STAND_LON") - dx = wrfnc.getncattr("DX") - dy = wrfnc.getncattr("DY") - stdlon = wrfnc.getncattr("STAND_LON") +def _get_proj_params(wrfnc, timeidx): + if timeidx < 0: + raise ValueError("'timeidx' must be greater than 0") + + attrs = extract_global_attrs(wrfnc, attrs=("MAP_PROJ", "TRUELAT1", + "TRUELAT2", "STAND_LON", + "DX", "DY", "STAND_LON")) + map_proj = attrs["MAP_PROJ"] + truelat1 = attrs["TRUELAT1"] + truelat2 = attrs["TRUELAT2"] + stdlon = attrs["STAND_LON"] + dx = attrs["DX"] + dy = attrs["DY"] + stdlon = attrs["STAND_LON"] if map_proj == 6: - pole_lat = wrfnc.getncattr("POLE_LAT") - pole_lon = wrfnc.getncattr("POLE_LON") + pole_attrs = extract_global_attrs(wrfnc, attrs=("POLE_LAT", + "POLE_LON")) + pole_lat = pole_attrs["POLE_LAT"] + pole_lon = pole_attrs["POLE_LON"] latinc = (dy*360.0)/2.0/3.141592653589793/6370000. loninc = (dx*360.0)/2.0/3.141592653589793/6370000. else: @@ -38,15 +62,8 @@ def get_ij(wrfnc, longitude, latitude, timeidx=0): latinc = 0.0 loninc = 0.0 - if "XLAT" in wrfnc.variables: - xlat = wrfnc.variables["XLAT"][timeidx,:,:] - elif "XLAT_M" in wrfnc.variables: - xlat = wrfnc.variables["XLAT_M"][timeidx,:,:] - - if "XLONG" in wrfnc.variables: - xlon = wrfnc.variables["XLONG"][timeidx,:,:] - elif "XLONG_M" in wrfnc.variables: - xlon = wrfnc.variables["XLONG_M"][timeidx,:,:] + xlat = get_lat(wrfnc, timeidx) + xlon = get_lon(wrfnc, timeidx) ref_lat = xlat[0,0] ref_lon = xlon[0,0] @@ -54,47 +71,32 @@ def get_ij(wrfnc, longitude, latitude, timeidx=0): known_i = 1.0 known_j = 1.0 + return (map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon, + pole_lat,pole_lon,known_i,known_j,dx,latinc, + loninc) + + +# TODO: longitude and latitude can also be lists +def get_ij(wrfnc, longitude, latitude, timeidx=0): + if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str): + raise TypeError("'get_ij' is only applicabe for single files") + + (map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon, + pole_lat,pole_lon,known_i,known_j,dx,latinc, + loninc) = _get_proj_params(wrfnc, timeidx) + return computeij(map_proj,truelat1,truelat2,stdlon, ref_lat,ref_lon,pole_lat,pole_lon, known_i,known_j,dx,latinc,loninc,latitude,longitude) - +# TODO: i and j can also be lists def get_ll(wrfnc, i, j, timeidx=0): + if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str): + raise TypeError("'get_ll' is only applicabe for single files") - map_proj = wrfnc.getncattr("MAP_PROJ") - truelat1 = wrfnc.getncattr("TRUELAT1") - truelat2 = wrfnc.getncattr("TRUELAT2") - stdlon = wrfnc.getncattr("STAND_LON") - dx = wrfnc.getncattr("DX") - dy = wrfnc.getncattr("DY") - stdlon = wrfnc.getncattr("STAND_LON") - - if map_proj == 6: - pole_lat = wrfnc.getncattr("POLE_LAT") - pole_lon = wrfnc.getncattr("POLE_LON") - latinc = (dy*360.0)/2.0/3.141592653589793/6370000. - loninc = (dx*360.0)/2.0/3.141592653589793/6370000. - else: - pole_lat = 90.0 - pole_lon = 0.0 - latinc = 0.0 - loninc = 0.0 - - if "XLAT" in wrfnc.variables: - xlat = wrfnc.variables["XLAT"][timeidx,:,:] - elif "XLAT_M" in wrfnc.variables: - xlat = wrfnc.variables["XLAT_M"][timeidx,:,:] - - if "XLONG" in wrfnc.variables: - xlon = wrfnc.variables["XLONG"][timeidx,:,:] - elif "XLONG_M" in wrfnc.variables: - xlon = wrfnc.variables["XLONG_M"][timeidx,:,:] - - ref_lat = xlat[0,0] - ref_lon = xlon[0,0] - - known_i = 1.0 - known_j = 1.0 + (map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon, + pole_lat,pole_lon,known_i,known_j,dx,latinc, + loninc) = _get_proj_params(wrfnc, timeidx) return computell(map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon, pole_lat,pole_lon,known_i,known_j,dx,latinc, diff --git a/wrf_open/var/src/python/wrf/var/omega.py b/wrf_open/var/src/python/wrf/var/omega.py index f5bcd56..b2eed83 100755 --- a/wrf_open/var/src/python/wrf/var/omega.py +++ b/wrf_open/var/src/python/wrf/var/omega.py @@ -2,15 +2,17 @@ from wrf.var.constants import Constants from wrf.var.destagger import destagger from wrf.var.extension import computeomega,computetk +from wrf.var.util import extract_vars __all__ = ["get_omega"] def get_omega(wrfnc, timeidx=0): - t = wrfnc.variables["T"][timeidx,:,:,:] - p = wrfnc.variables["P"][timeidx,:,:,:] - w = wrfnc.variables["W"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] - qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:] + vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "W", "PB", "QVAPOR")) + t = vars["T"] + p = vars["P"] + w = vars["W"] + pb = vars["PB"] + qv = vars["QVAPOR"] wa = destagger(w, 0) full_t = t + Constants.T_BASE diff --git a/wrf_open/var/src/python/wrf/var/precip.py b/wrf_open/var/src/python/wrf/var/precip.py index 2ce40cd..294a1b4 100755 --- a/wrf_open/var/src/python/wrf/var/precip.py +++ b/wrf_open/var/src/python/wrf/var/precip.py @@ -1,22 +1,27 @@ import numpy as n +from wrf.var.util import extract_vars + __all__ = ["get_accum_precip", "get_precip_diff"] def get_accum_precip(wrfnc, timeidx=0): - rainc = wrfnc.variables["RAINC"][timeidx,:,:] - rainnc = wrfnc.variables["RAINNC"][timeidx,:,:] + vars = extract_vars(wrfnc, timeidx, vars=("RAINC", "RAINNC")) + rainc = vars["RAINC"] + rainnc = vars["RAINNC"] rainsum = rainc + rainnc return rainsum def get_precip_diff(wrfnc1, wrfnc2, timeidx=0): - rainc1 = wrfnc1.variables["RAINC"][timeidx,:,:] - rainnc1 = wrfnc1.variables["RAINNC"][timeidx,:,:] + vars1 = extract_vars(wrfnc, timeidx, vars=("RAINC", "RAINNC")) + vars2 = extract_vars(wrfnc, timeidx, vars=("RAINC", "RAINNC")) + rainc1 = vars1["RAINC"] + rainnc1 = vars1["RAINNC"] - rainc2 = wrfnc2.variables["RAINC"][timeidx,:,:] - rainnc2 = wrfnc2.variables["RAINNC"][timeidx,:,:] + rainc2 = vars2["RAINC"] + rainnc2 = vars2["RAINNC"] rainsum1 = rainc1 + rainnc1 rainsum2 = rainc2 + rainnc2 diff --git a/wrf_open/var/src/python/wrf/var/pressure.py b/wrf_open/var/src/python/wrf/var/pressure.py index da1f0b6..2ba4235 100755 --- a/wrf_open/var/src/python/wrf/var/pressure.py +++ b/wrf_open/var/src/python/wrf/var/pressure.py @@ -1,18 +1,26 @@ from wrf.var.constants import Constants from wrf.var.decorators import convert_units +from wrf.var.util import extract_vars __all__ = ["get_pressure"] @convert_units("pressure", "pa") def get_pressure(wrfnc, units="hpa", timeidx=0): - if "P" in wrfnc.variables: - p = wrfnc.variables["P"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] + try: + p_vars = extract_vars(wrfnc, timeidx, vars=("P", "PB")) + except KeyError: + try: + pres_vars = extract_vars(wrfnc, timeidx, vars="PRES") + except: + raise RuntimeError("pressure variable not found in NetCDF file") + else: + pres = pres_vars["PRES"] + else: + p = p_vars["P"] + pb = p_vars["PB"] pres = p + pb - elif "PRES" in wrfnc.variables: - pres = wrfnc.variables["PRES"][timeidx,:,:,:] return pres diff --git a/wrf_open/var/src/python/wrf/var/pw.py b/wrf_open/var/src/python/wrf/var/pw.py index 2eb4bc6..6ecc8c9 100755 --- a/wrf_open/var/src/python/wrf/var/pw.py +++ b/wrf_open/var/src/python/wrf/var/pw.py @@ -1,17 +1,20 @@ from wrf.var.extension import computepw,computetv,computetk from wrf.var.constants import Constants +from wrf.var.util import extract_vars __all__ = ["get_pw"] def get_pw(wrfnc, timeidx=0): + vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "PH", "PHB", + "QVAPOR")) - t = wrfnc.variables["T"][timeidx,:,:,:] - p = wrfnc.variables["P"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] - ph = wrfnc.variables["PH"][timeidx,:,:,:] - phb = wrfnc.variables["PHB"][timeidx,:,:,:] - qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:] + t = vars["T"] + p = vars["P"] + pb = vars["PB"] + ph = vars["PH"] + phb = vars["PHB"] + qv = vars["QVAPOR"] # Change this to use real virtual temperature! full_p = p + pb diff --git a/wrf_open/var/src/python/wrf/var/rh.py b/wrf_open/var/src/python/wrf/var/rh.py index 44c94a5..90575a5 100755 --- a/wrf_open/var/src/python/wrf/var/rh.py +++ b/wrf_open/var/src/python/wrf/var/rh.py @@ -1,15 +1,16 @@ from wrf.var.constants import Constants from wrf.var.extension import computerh, computetk +from wrf.var.util import extract_vars __all__ = ["get_rh", "get_rh_2m"] def get_rh(wrfnc, timeidx=0): - t = wrfnc.variables["T"][timeidx,:,:,:] - #t00 = wrfnc.variables["T00"][timeidx] - p = wrfnc.variables["P"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] - qvapor = wrfnc.variables["QVAPOR"][timeidx,:,:,:] + vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR")) + t = vars["T"] + p = vars["P"] + pb = vars["PB"] + qvapor = vars["QVAPOR"] full_t = t + Constants.T_BASE full_p = p + pb @@ -20,9 +21,10 @@ def get_rh(wrfnc, timeidx=0): return rh def get_rh_2m(wrfnc, timeidx=0): - t2 = wrfnc.variables["T2"][timeidx,:,:] - psfc = wrfnc.variables["PSFC"][timeidx,:,:] - q2 = wrfnc.variables["Q2"][timeidx,:,:] + vars = extract_vars(wrfnc, timeidx, vars=("T2", "PSFC", "Q2")) + t2 = vars["T2"] + psfc = vars["PSFC"] + q2 = vars["Q2"] q2[q2 < 0] = 0 rh = computerh(q2, psfc, t2) diff --git a/wrf_open/var/src/python/wrf/var/slp.py b/wrf_open/var/src/python/wrf/var/slp.py index be92438..8c4320d 100755 --- a/wrf_open/var/src/python/wrf/var/slp.py +++ b/wrf_open/var/src/python/wrf/var/slp.py @@ -2,18 +2,21 @@ from wrf.var.extension import computeslp, computetk from wrf.var.constants import Constants from wrf.var.destagger import destagger from wrf.var.decorators import convert_units +from wrf.var.util import extract_vars __all__ = ["get_slp"] @convert_units("pressure", "hpa") def get_slp(wrfnc, units="hpa", timeidx=0): + vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", + "PH", "PHB")) - t = wrfnc.variables["T"][timeidx,:,:,:] - p = wrfnc.variables["P"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] - qvapor = wrfnc.variables["QVAPOR"][timeidx,:,:,:] - ph = wrfnc.variables["PH"][timeidx,:,:,:] - phb = wrfnc.variables["PHB"][timeidx,:,:,:] + t = vars["T"] + p = vars["P"] + pb = vars["PB"] + qvapor = vars["QVAPOR"] + ph = vars["PH"] + phb = vars["PHB"] full_t = t + Constants.T_BASE full_p = p + pb diff --git a/wrf_open/var/src/python/wrf/var/temp.py b/wrf_open/var/src/python/wrf/var/temp.py index a223cae..10686d8 100755 --- a/wrf_open/var/src/python/wrf/var/temp.py +++ b/wrf_open/var/src/python/wrf/var/temp.py @@ -2,12 +2,14 @@ from wrf.var.constants import Constants from wrf.var.extension import computetk, computeeth, computetv, computewetbulb from wrf.var.decorators import convert_units +from wrf.var.util import extract_vars __all__ = ["get_theta", "get_temp", "get_eth", "get_tv", "get_tw"] @convert_units("temp", "k") def get_theta(wrfnc, units="k", timeidx=0): - t = wrfnc.variables["T"][timeidx,:,:,:] + vars = extract_vars(wrfnc, timeidx, vars="T") + t = vars["T"] full_t = t + Constants.T_BASE return full_t @@ -16,9 +18,10 @@ def get_theta(wrfnc, units="k", timeidx=0): def get_temp(wrfnc, units="k", timeidx=0): """Return the temperature in Kelvin or Celsius""" - t = wrfnc.variables["T"][timeidx,:,:,:] - p = wrfnc.variables["P"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] + vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB")) + t = vars["T"] + p = vars["P"] + pb = vars["PB"] full_t = t + Constants.T_BASE full_p = p + pb @@ -30,10 +33,11 @@ def get_temp(wrfnc, units="k", timeidx=0): def get_eth(wrfnc, units="k", timeidx=0): "Return equivalent potential temperature (Theta-e) in Kelvin" - t = wrfnc.variables["T"][timeidx,:,:,:] - p = wrfnc.variables["P"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] - qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:] + vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR")) + t = vars["T"] + p = vars["P"] + pb = vars["PB"] + qv = vars["QVAPOR"] full_t = t + Constants.T_BASE full_p = p + pb @@ -47,10 +51,12 @@ def get_eth(wrfnc, units="k", timeidx=0): def get_tv(wrfnc, units="k", timeidx=0): "Return the virtual temperature (tv) in Kelvin or Celsius" - t = wrfnc.variables["T"][timeidx,:,:,:] - p = wrfnc.variables["P"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] - qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:] + vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR")) + + t = vars["T"] + p = vars["P"] + pb = vars["PB"] + qv = vars["QVAPOR"] full_t = t + Constants.T_BASE full_p = p + pb @@ -64,11 +70,12 @@ def get_tv(wrfnc, units="k", timeidx=0): @convert_units("temp", "k") def get_tw(wrfnc, units="k", timeidx=0): "Return the wetbulb temperature (tw)" - - t = wrfnc.variables["T"][timeidx,:,:,:] - p = wrfnc.variables["P"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] - qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:] + + vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR")) + t = vars["T"] + p = vars["P"] + pb = vars["PB"] + qv = vars["QVAPOR"] full_t = t + Constants.T_BASE full_p = p + pb diff --git a/wrf_open/var/src/python/wrf/var/terrain.py b/wrf_open/var/src/python/wrf/var/terrain.py index 078343f..a3fc6a5 100755 --- a/wrf_open/var/src/python/wrf/var/terrain.py +++ b/wrf_open/var/src/python/wrf/var/terrain.py @@ -1,15 +1,23 @@ from wrf.var.decorators import convert_units +from wrf.var.util import extract_vars __all__ = ["get_terrain"] @convert_units("height", "m") def get_terrain(wrfnc, units="m", timeidx=0): - - if "HGT" in wrfnc.variables: - hgt = wrfnc.variables["HGT"][timeidx,:,:] - elif "HGT_M": - hgt = wrfnc.variables["HGT_M"][timeidx,:,:] + + try: + hgt_vars = extract_vars(wrfnc, timeidx, vars="HGT") + except KeyError: + try: + hgt_m_vars = extract_vars(wrfnc, timeidx, vars="HGT_M") + except KeyError: + raise RuntimeError("height variable not found in NetCDF file") + else: + hgt = hgt_m_vars["HGT_M"] + else: + hgt = hgt_vars["HGT"] return hgt diff --git a/wrf_open/var/src/python/wrf/var/times.py b/wrf_open/var/src/python/wrf/var/times.py index 770fbab..2328d85 100755 --- a/wrf_open/var/src/python/wrf/var/times.py +++ b/wrf_open/var/src/python/wrf/var/times.py @@ -6,6 +6,7 @@ __all__ = ["get_times"] def _make_time(timearr): return dt.strptime("".join(timearr[:]), "%Y-%m-%d_%H:%M:%S") +# TODO: handle list of files and multiple times def get_times(wrfnc): times = wrfnc.variables["Times"][:,:] return [_make_time(times[i,:]) for i in xrange(times.shape[0])] diff --git a/wrf_open/var/src/python/wrf/var/units.py b/wrf_open/var/src/python/wrf/var/units.py index 483e588..97055ae 100755 --- a/wrf_open/var/src/python/wrf/var/units.py +++ b/wrf_open/var/src/python/wrf/var/units.py @@ -1,7 +1,7 @@ from wrf.var.constants import Constants, ConversionFactors -__all__ = ["check_units", "do_conversion", "convert_units"] +__all__ = ["check_units", "do_conversion"] # Handles unit conversions that only differ by multiplication factors def _apply_conv_fact(var, vartype, var_unit, dest_unit): @@ -119,7 +119,6 @@ def do_conversion(var, vartype, var_unit, dest_unit): else: return _apply_temp_conv(var, var_unit, dest_unit) -convert_units = do_conversion diff --git a/wrf_open/var/src/python/wrf/var/util.py b/wrf_open/var/src/python/wrf/var/util.py new file mode 100644 index 0000000..d921256 --- /dev/null +++ b/wrf_open/var/src/python/wrf/var/util.py @@ -0,0 +1,77 @@ +from collections import Iterable + +import numpy as n + +__all__ = ["extract_vars", "extract_global_attrs", "hold_dim_fixed"] + +def _is_multi_time(timeidx): + if timeidx == -1: + return True + return False + +def _is_multi_file(wrfnc): + if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str): + return True + return False + +def _get_attr(wrfnc, attr): + val = getattr(wrfnc, attr) + + # PyNIO puts single values in to an array + if isinstance(val,n.ndarray): + if len(val) == 1: + return val[0] + return val + +def extract_global_attrs(wrfnc, attrs): + if isinstance(attrs, str): + attrlist = [attrs] + else: + attrlist = attrs + + multifile = _is_multi_file(wrfnc) + + if multifile: + wrfnc = wrfnc[0] + + return {attr:_get_attr(wrfnc, attr) for attr in attrlist} + + +def extract_vars(wrfnc, timeidx, vars): + if isinstance(vars, str): + varlist = [vars] + else: + varlist = vars + + multitime = _is_multi_time(timeidx) + multifile = _is_multi_file(wrfnc) + + # Single file, single time + if not multitime and not multifile: + return {var:wrfnc.variables[var][timeidx,:] for var in varlist} + +def hold_dim_fixed(var, dim, idx): + """Generic method to hold a single dimension to a fixed index when + the array shape is unknown. The values for 'dim' and 'idx' can + be negative. + + For example, on a 4D array with 'dim' set to + -1 and 'idx' set to 3, this is simply going to do this operation: + + var[:,:,:,3] + + """ + + var_shape = var.shape + num_dims = len(var_shape) + full_slice = slice(None, None, None) + + # Make the sequence of slices + dim_slices = [full_slice for x in xrange(num_dims)] + + # Replace the dimension with the fixed index + dim_slices[dim] = idx + + return var[dim_slices] + + \ No newline at end of file diff --git a/wrf_open/var/src/python/wrf/var/uvmet.py b/wrf_open/var/src/python/wrf/var/uvmet.py index 0a8fec0..fb037bc 100755 --- a/wrf_open/var/src/python/wrf/var/uvmet.py +++ b/wrf_open/var/src/python/wrf/var/uvmet.py @@ -5,6 +5,7 @@ from wrf.var.destagger import destagger from wrf.var.constants import Constants from wrf.var.wind import _calc_wspd_wdir from wrf.var.decorators import convert_units +from wrf.var.util import extract_vars, extract_global_attrs, hold_dim_fixed __all__=["get_uvmet", "get_uvmet10", "get_uvmet_wspd_wdir", "get_uvmet10_wspd_wdir"] @@ -14,27 +15,71 @@ def get_uvmet(wrfnc, ten_m=False, units ="mps", timeidx=0): """ Return a tuple of u,v with the winds rotated in to earth space""" if not ten_m: - if "U" in wrfnc.variables: - u = destagger(wrfnc.variables["U"][timeidx,:,:,:], 2) - elif "UU" in wrfnc.variables: - u = destagger(wrfnc.variables["UU"][timeidx,:,:,:], 2) # support met_em files + try: + u_vars = extract_vars(wrfnc, timeidx, vars="U") + except KeyError: + try: + uu_vars = extract_vars(wrfnc, timeidx, vars="UU") + except KeyError: + raise RuntimeError("No valid wind data found in NetCDF file") + else: + u = destagger(uu_vars["UU"], -1) # support met_em files + else: + u = destagger(u_vars["U"], -1) + + try: + v_vars = extract_vars(wrfnc, timeidx, vars="V") + except KeyError: + try: + vv_vars = extract_vars(wrfnc, timeidx, vars="VV") + except KeyError: + raise RuntimeError("No valid wind data found in NetCDF file") + else: + v = destagger(vv_vars["VV"], -2) # support met_em files + else: + v = destagger(v_vars["V"], -2) - if "V" in wrfnc.variables: - v = destagger(wrfnc.variables["V"][timeidx,:,:,:], 1) - elif "VV" in wrfnc.variables: - v = destagger(wrfnc.variables["VV"][timeidx,:,:,:], 1) else: - if "U10" in wrfnc.variables: - u = wrfnc.variables["U10"][timeidx,:,:] - elif "UU" in wrfnc.variables: - u = destagger(wrfnc.variables["UU"][timeidx,0,:,:], 1) # support met_em files - - if "V10" in wrfnc.variables: - v = wrfnc.variables["V10"][timeidx,:,:] - elif "VV" in wrfnc.variables: - v = destagger(wrfnc.variables["VV"][timeidx,0,:,:], 0) # support met_em files +# # Original code, keeping for reference +# if "U10" in wrfnc.variables: +# u = wrfnc.variables["U10"][timeidx,:,:] +# elif "UU" in wrfnc.variables: +# u = destagger(wrfnc.variables["UU"][timeidx,0,:,:], 1) # support met_em files +# +# if "V10" in wrfnc.variables: +# v = wrfnc.variables["V10"][timeidx,:,:] +# elif "VV" in wrfnc.variables: +# v = destagger(wrfnc.variables["VV"][timeidx,0,:,:], 0) # support met_em files + try: + u_vars = extract_vars(wrfnc, timeidx, vars="U10") + except KeyError: + try: + uu_vars = extract_vars(wrfnc, timeidx, vars="UU") + except KeyError: + raise RuntimeError("No valid wind data found in NetCDF file") + else: + # For met_em files, this just takes the lowest level winds + # (3rd dimension from right is bottom_top) + u = destagger(hold_dim_fixed(uu_vars["UU"], -3, 0), -1) # support met_em files + else: + u = u_vars["U10"] - map_proj = wrfnc.getncattr("MAP_PROJ") + try: + v_vars = extract_vars(wrfnc, timeidx, vars="V10") + except KeyError: + try: + vv_vars = extract_vars(wrfnc, timeidx, vars="VV") + except KeyError: + raise RuntimeError("No valid wind data found in NetCDF file") + else: + # For met_em files, this just takes the lowest level winds + # (3rd dimension from right is bottom_top) + v = destagger(hold_dim_fixed(vv_vars["VV"], -3, 0), -2) # support met_em files + else: + v = v_vars["V10"] + + map_proj_attrs = extract_global_attrs(wrfnc, attrs="MAP_PROJ") + map_proj = map_proj_attrs["MAP_PROJ"] # 1 - Lambert # 2 - Polar Stereographic @@ -47,26 +92,50 @@ def get_uvmet(wrfnc, ten_m=False, units ="mps", timeidx=0): # No rotation needed for Mercator and Lat/Lon return u,v elif map_proj in (1,2): + lat_attrs = extract_global_attrs(wrfnc, attrs=("CEN_LAT", + "TRUELAT1", + "TRUELAT2")) radians_per_degree = Constants.PI/180.0 # Rotation needed for Lambert and Polar Stereographic - cen_lat = wrfnc.getncattr("CEN_LAT") - if "STAND_LON" in wrfnc.ncattrs(): - cen_lon = wrfnc.getncattr("STAND_LON") - else: - cen_lon = wrfnc.getncattr("CEN_LON") - - true_lat1 = wrfnc.getncattr("TRUELAT1") - true_lat2 = wrfnc.getncattr("TRUELAT2") + cen_lat = lat_attrs["CEN_LAT"] + true_lat1 = lat_attrs["TRUELAT1"] + true_lat2 = lat_attrs["TRUELAT2"] - if "XLAT_M" in wrfnc.variables: - lat = wrfnc.variables["XLAT_M"][timeidx,:,:] + try: + lon_attrs = extract_global_attrs(wrfnc, attrs="STAND_LON") + except AttributeError: + try: + cen_lon_attrs = extract_global_attrs(wrfnc, attrs="CEN_LON") + except AttributeError: + raise RuntimeError("longitude attributes not found in NetCDF") + else: + cen_lon = cen_lon_attrs["CEN_LON"] + else: + cen_lon = lon_attrs["STAND_LON"] + + try: + xlat_m_vars = extract_vars(wrfnc, timeidx, vars="XLAT_M") + except KeyError: + try: + xlat_vars = extract_vars(wrfnc, timeidx, vars="XLAT") + except KeyError: + raise RuntimeError("xlat not found in NetCDF file") + else: + lat = xlat_vars["XLAT"] else: - lat = wrfnc.variables["XLAT"][timeidx,:,:] + lat = xlat_m_vars["XLAT_M"] - if "XLONG_M" in wrfnc.variables: - lon = wrfnc.variables["XLONG_M"][timeidx,:,:] + try: + xlon_m_vars = extract_vars(wrfnc, timeidx, vars="XLONG_M") + except KeyError: + try: + xlon_vars = extract_vars(wrfnc, timeidx, vars="XLONG") + except KeyError: + raise RuntimeError("xlong not found in NetCDF file") + else: + lon = xlon_vars["XLONG"] else: - lon = wrfnc.variables["XLONG"][timeidx,:,:] + lon = xlon_m_vars["XLONG_M"] if map_proj == 1: if((fabs(true_lat1 - true_lat2) > 0.1) and diff --git a/wrf_open/var/src/python/wrf/var/vorticity.py b/wrf_open/var/src/python/wrf/var/vorticity.py index 1072031..69ce159 100755 --- a/wrf_open/var/src/python/wrf/var/vorticity.py +++ b/wrf_open/var/src/python/wrf/var/vorticity.py @@ -1,32 +1,46 @@ from wrf.var.extension import computeavo, computepvo +from wrf.var.util import extract_vars, extract_global_attrs __all__ = ["get_avo", "get_pvo"] def get_avo(wrfnc, timeidx=0): - u = wrfnc.variables["U"][timeidx,:,:,:] - v = wrfnc.variables["V"][timeidx,:,:,:] - msfu = wrfnc.variables["MAPFAC_U"][timeidx,:,:] - msfv = wrfnc.variables["MAPFAC_V"][timeidx,:,:] - msfm = wrfnc.variables["MAPFAC_M"][timeidx,:,:] - cor = wrfnc.variables["F"][timeidx,:,:] - dx = wrfnc.getncattr("DX") - dy = wrfnc.getncattr("DY") + vars = extract_vars(wrfnc, timeidx, vars=("U", "V", "MAPFAC_U", + "MAPFAC_V", "MAPFAC_M", + "F")) + + attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY")) + u = vars["U"] + v = vars["V"] + msfu = vars["MAPFAC_U"] + msfv = vars["MAPFAC_V"] + msfm = vars["MAPFAC_M"] + cor = vars["F"] + + dx = attrs["DX"] + dy = attrs["DY"] return computeavo(u,v,msfu,msfv,msfm,cor,dx,dy) def get_pvo(wrfnc, timeidx=0): - u = wrfnc.variables["U"][timeidx,:,:,:] - v = wrfnc.variables["V"][timeidx,:,:,:] - t = wrfnc.variables["T"][timeidx,:,:,:] - p = wrfnc.variables["P"][timeidx,:,:,:] - pb = wrfnc.variables["PB"][timeidx,:,:,:] - msfu = wrfnc.variables["MAPFAC_U"][timeidx,:,:] - msfv = wrfnc.variables["MAPFAC_V"][timeidx,:,:] - msfm = wrfnc.variables["MAPFAC_M"][timeidx,:,:] - cor = wrfnc.variables["F"][timeidx,:,:] - dx = wrfnc.getncattr("DX") - dy = wrfnc.getncattr("DY") + vars = extract_vars(wrfnc, timeidx, vars=("U", "V", "T", "P", + "PB", "MAPFAC_U", + "MAPFAC_V", "MAPFAC_M", + "F")) + attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY")) + + u = vars["U"] + v = vars["V"] + t = vars["T"] + p = vars["P"] + pb = vars["PB"] + msfu = vars["MAPFAC_U"] + msfv = vars["MAPFAC_V"] + msfm = vars["MAPFAC_M"] + cor = vars["F"] + + dx = attrs["DX"] + dy = attrs["DY"] full_t = t + 300 full_p = p + pb diff --git a/wrf_open/var/src/python/wrf/var/wind.py b/wrf_open/var/src/python/wrf/var/wind.py index adcb078..4631b0d 100755 --- a/wrf_open/var/src/python/wrf/var/wind.py +++ b/wrf_open/var/src/python/wrf/var/wind.py @@ -4,21 +4,22 @@ import numpy as n from wrf.var.constants import Constants from wrf.var.destagger import destagger_windcomp from wrf.var.decorators import convert_units +from wrf.var.util import extract_vars __all__ = ["get_u_destag", "get_v_destag", "get_w_destag", "get_destag_wspd_wdir"] -def _calc_wspd(u, v): +@convert_units("wind", "mps") +def _calc_wspd(u, v, units="mps"): return n.sqrt(u**2 + v**2) def _calc_wdir(u, v): wdir = 270.0 - n.arctan2(v,u) * (180.0/Constants.PI) return n.remainder(wdir, 360.0) -@convert_units("wind", "mps") def _calc_wspd_wdir(u, v, units="mps"): check_units(units, "wind") - return (_calc_wspd(u,v), _calc_wdir(u,v)) + return (_calc_wspd(u,v, units), _calc_wdir(u,v)) @convert_units("wind", "mps") def get_u_destag(wrfnc, units="mps", timeidx=0):