diff --git a/wrf_open/var/src/python/wrf/var/__init__.py b/wrf_open/var/src/python/wrf/var/__init__.py index 0539df8..557771e 100755 --- a/wrf_open/var/src/python/wrf/var/__init__.py +++ b/wrf_open/var/src/python/wrf/var/__init__.py @@ -1,63 +1,61 @@ import warnings -from config import * -import config -from extension import * -import extension -from util import * -import util -from cape import * -import cape -from constants import * -import constants -from ctt import * -import ctt -from dbz import * -import dbz -from destag import * -import destag -from dewpoint import * -import dewpoint -from etaconv import * -import etaconv -from geoht import * -import geoht -from helicity import * -import helicity -from interp import * -import interp -from latlon import * -import latlon -from omega import * -import omega -from precip import * -import precip -from pressure import * -import pressure -from psadlookup import * -import psadlookup -from pw import * -import pw -from rh import * -import rh -from slp import * -import slp -from temp import * -import temp -from terrain import * -import terrain -from uvmet import * -import uvmet -from vorticity import * -import vorticity -from wind import * -import wind -from times import * -import times -from units import * -import units -from projection import * -import projection +from . import config +from .config import * +from . import extension +from .extension import * +from . import util +from .util import * +from . import cape +from .cape import * +from . import constants +from .constants import * +from . import ctt +from .ctt import * +from . import dbz +from .dbz import * +from . import destag +from .destag import * +from . import dewpoint +from .dewpoint import * +from . import geoht +from .geoht import * +from . import helicity +from .helicity import * +from . import interp +from .interp import * +from . import latlon +from .latlon import * +from . import omega +from .omega import * +from . import precip +from .precip import * +from . import pressure +from .pressure import * +from . import psadlookup +from .psadlookup import * +from . import pw +from .pw import * +from . import rh +from .rh import * +from . import slp +from .slp import * +from . import temp +from .temp import * +from . import terrain +from .terrain import * +from . import uvmet +from .uvmet import * +from . import vorticity +from .vorticity import * +from . import wind +from .wind import * +from . import times +from .times import * +from . import units +from .units import * +from . import projection +from .projection import * __all__ = ["getvar"] __all__.extend(config.__all__) @@ -69,7 +67,6 @@ __all__.extend(ctt.__all__) __all__.extend(dbz.__all__) __all__.extend(destag.__all__) __all__.extend(dewpoint.__all__) -__all__.extend(etaconv.__all__) __all__.extend(geoht.__all__) __all__.extend(helicity.__all__) __all__.extend(interp.__all__) @@ -127,9 +124,10 @@ _FUNC_MAP = {"cape2d" : get_2dcape, "lon" : get_lon, "pressure" : get_pressure_hpa, "pres" : get_pressure, - "wspddir" : get_destag_wspd_wdir, - "wspddir_uvmet" : get_uvmet_wspd_wdir, - "wspddir_uvmet10" : get_uvmet10_wspd_wdir, + "wspd_wdir" : get_destag_wspd_wdir, + "wspd_wdir10" : get_destag_wspd_wdir10, + "wspd_wdir_uvmet" : get_uvmet_wspd_wdir, + "wspd_wdir_uvmet10" : get_uvmet10_wspd_wdir, "ctt" : get_ctt } @@ -214,9 +212,11 @@ def _check_kargs(var, kargs): "argument for '%s'" % (arg, var)) -def getvar(wrfnc, var, timeidx=0, **kargs): +def getvar(wrfnc, var, timeidx=0, + method="cat", squeeze=True, cache=None, + **kargs): if is_standard_wrf_var(wrfnc, var): - return extract_vars(wrfnc, timeidx, var)[var] + return extract_vars(wrfnc, timeidx, var, method, squeeze, cache)[var] actual_var = _undo_alias(var) if actual_var not in _VALID_KARGS: diff --git a/wrf_open/var/src/python/wrf/var/cape.py b/wrf_open/var/src/python/wrf/var/cape.py index 320227b..2cc134c 100755 --- a/wrf_open/var/src/python/wrf/var/cape.py +++ b/wrf_open/var/src/python/wrf/var/cape.py @@ -1,17 +1,26 @@ import numpy as n import numpy.ma as ma -from wrf.var.extension import computetk,computecape -from wrf.var.destag import destagger -from wrf.var.constants import Constants, ConversionFactors -from wrf.var.util import extract_vars +from .extension import computetk,computecape +from .destag import destagger +from .constants import Constants, ConversionFactors +from .util import extract_vars, combine_with +from .decorators import copy_and_set_metadata __all__ = ["get_2dcape", "get_3dcape"] - -def get_2dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL): +@copy_and_set_metadata(copy_varname="T", + name="cape_2d", + dimnames=combine_with("T", remove_dims=("bottom_top",), + insert_before="south_north", + new_dimnames=("mcape_mcin_lcl_lfc",)), + description="mcape ; mcin ; lcl ; lfc", + units="J/kg ; J/kg ; m ; m", + MemoryOrder = "XY") +def get_2dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL, + method="cat", squeeze=True, cache=None): """Return the 2d fields of cape, cin, lcl, and lfc""" - ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR", "PH", - "PHB", "HGT", "PSFC")) + varnames = ("T", "P", "PB", "QVAPOR", "PH","PHB", "HGT", "PSFC") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) t = ncvars["T"] p = ncvars["P"] @@ -51,20 +60,27 @@ def get_2dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL): resdim = left_dims + [4] + right_dims # Make a new output array for the result - res = n.zeros((resdim), cape.dtype) + res = n.zeros(resdim, cape.dtype) res[...,0,:,:] = cape[:] res[...,1,:,:] = cin[:] res[...,2,:,:] = lcl[:] res[...,3,:,:] = lfc[:] - return ma.masked_values(res,missing) + return ma.masked_values(res, missing) -def get_3dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL): +@copy_and_set_metadata(copy_varname="T", name="cape_3d", + dimnames=combine_with("T", + insert_before="bottom_top", + new_dimnames=("cape_cin",)), + description="cape ; cin", + units="J kg-1 ; J kg-1", + MemoryOrder = "XY") +def get_3dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL, + method="cat", squeeze=True, cache=None): """Return the 3d fields of cape and cin""" - - ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR", - "PH", "PHB", "HGT", "PSFC")) + varnames = ("T", "P", "PB", "QVAPOR", "PH", "PHB", "HGT", "PSFC") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) t = ncvars["T"] p = ncvars["P"] pb = ncvars["PB"] diff --git a/wrf_open/var/src/python/wrf/var/ctt.py b/wrf_open/var/src/python/wrf/var/ctt.py index 8cf0a2f..0f06c82 100644 --- a/wrf_open/var/src/python/wrf/var/ctt.py +++ b/wrf_open/var/src/python/wrf/var/ctt.py @@ -1,21 +1,24 @@ import numpy as n -from wrf.var.extension import computectt, computetk -from wrf.var.constants import Constants, ConversionFactors -from wrf.var.destag import destagger -from wrf.var.decorators import convert_units -from wrf.var.util import extract_vars +from .extension import computectt, computetk +from .constants import Constants, ConversionFactors +from .destag import destagger +from .decorators import convert_units, copy_and_set_metadata +from .util import extract_vars __all__ = ["get_ctt"] +@copy_and_set_metadata(copy_varname="T", name="ctt", + description="cloud top temperature") @convert_units("temp", "c") -def get_ctt(wrfnc, timeidx=0, units="c"): +def get_ctt(wrfnc, timeidx=0, units="c", method="cat", squeeze=True, + cache=None): """Return the cloud top temperature. """ - ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "PH" , - "PHB", "HGT", "QVAPOR")) + varnames = ("T", "P", "PB", "PH", "PHB", "HGT", "QVAPOR") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) t = ncvars["T"] p = ncvars["P"] pb = ncvars["PB"] @@ -26,7 +29,7 @@ def get_ctt(wrfnc, timeidx=0, units="c"): haveqci = 1 try: - icevars = extract_vars(wrfnc, timeidx, vars="QICE") + icevars = extract_vars(wrfnc, timeidx, "QICE", method, squeeze, cache) except KeyError: qice = n.zeros(qv.shape, qv.dtype) haveqci = 0 @@ -34,7 +37,7 @@ def get_ctt(wrfnc, timeidx=0, units="c"): qice = icevars["QICE"] * 1000.0 #g/kg try: - cldvars = extract_vars(wrfnc, timeidx, vars="QCLOUD") + cldvars = extract_vars(wrfnc, timeidx, "QCLOUD", method, squeeze, cache) except KeyError: raise RuntimeError("'QCLOUD' not found in NetCDF file") else: diff --git a/wrf_open/var/src/python/wrf/var/dbz.py b/wrf_open/var/src/python/wrf/var/dbz.py index d68a10f..ff0c67f 100755 --- a/wrf_open/var/src/python/wrf/var/dbz.py +++ b/wrf_open/var/src/python/wrf/var/dbz.py @@ -1,12 +1,17 @@ import numpy as n -from wrf.var.extension import computedbz,computetk -from wrf.var.constants import Constants -from wrf.var.util import extract_vars +from .extension import computedbz,computetk +from .constants import Constants +from .util import extract_vars, combine_with +from .decorators import copy_and_set_metadata __all__ = ["get_dbz", "get_max_dbz"] -def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False): +@copy_and_set_metadata(copy_varname="T", name="dbz", + description="radar reflectivity", + units="dBz") +def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False, + method="cat", squeeze=True, cache=None): """ Return the dbz do_varint - do variable intercept (if False, constants are used. Otherwise, @@ -17,8 +22,8 @@ def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False): as liquid) """ - ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR", - "QRAIN")) + varnames = ("T", "P", "PB", "QVAPOR", "QRAIN") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) t = ncvars["T"] p = ncvars["P"] pb = ncvars["PB"] @@ -26,14 +31,16 @@ def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False): qr = ncvars["QRAIN"] try: - snowvars = extract_vars(wrfnc, timeidx, vars="QSNOW") + snowvars = extract_vars(wrfnc, timeidx, "QSNOW", + method, squeeze, cache) except KeyError: qs = n.zeros(qv.shape, "float") else: qs = snowvars["QSNOW"] try: - graupvars = extract_vars(wrfnc, timeidx, vars="QGRAUP") + graupvars = extract_vars(wrfnc, timeidx, "QGRAUP", + method, squeeze, cache) except KeyError: qg = n.zeros(qv.shape, "float") else: @@ -58,7 +65,13 @@ def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False): return computedbz(full_p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin) -def get_max_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False): - return n.amax(get_dbz(wrfnc, do_varint, do_liqskin, timeidx), +@copy_and_set_metadata(copy_varname="T", name="dbz", + dimnames=combine_with("T", remove_dims=("bottom_top",)), + description="maximum radar reflectivity", + units="dBz") +def get_max_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False, + method="cat", squeeze=True, cache=None): + return n.amax(get_dbz(wrfnc, do_varint, do_liqskin, timeidx, method, + method, squeeze, cache), axis=-3) diff --git a/wrf_open/var/src/python/wrf/var/decorators.py b/wrf_open/var/src/python/wrf/var/decorators.py index d5de0a8..8153ee3 100644 --- a/wrf_open/var/src/python/wrf/var/decorators.py +++ b/wrf_open/var/src/python/wrf/var/decorators.py @@ -1,13 +1,16 @@ -from functools import wraps +#from functools import wraps +import wrapt from inspect import getargspec +from collections import OrderedDict -import numpy as n +import numpy as np import numpy.ma as ma -from wrf.var.units import do_conversion, check_units -from wrf.var.destag import destagger -from wrf.var.util import iter_left_indexes -from wrf.var.config import xarray_enabled +from .units import do_conversion, check_units +from .destag import destagger +from .util import (iter_left_indexes, viewitems, extract_vars, + combine_with, either) +from .config import xarray_enabled if xarray_enabled(): from xarray import DataArray @@ -15,43 +18,63 @@ if xarray_enabled(): __all__ = ["convert_units", "handle_left_iter", "uvmet_left_iter", - "handle_casting" "set_metadata"] + "handle_casting" "copy_and_set_metadata", "set_wind_metadata", + "set_latlon_metadata", "set_height_metadata"] # TODO: In python 3.x, getargspec is deprecated -class from_args(object): - def __init__(self, argname): - self.argname - - def __call__(self, func, *args, **kargs): +def from_args(func, argnames, *args, **kwargs): """Parses the function args and kargs looking for the desired argument value. Otherwise, the value is taken from the default keyword argument using the arg spec. """ - # If units are provided to the method call, use them. - # Otherwise, need to parse the argspec to find what the default - # value is since wraps does not preserve this. - argspec = getargspec(func) + if isinstance(argnames, str): + arglist = [argnames] + else: + arglist = argnames - if self.argname not in argspec.args and self.argname not in kargs: - return None - - try: - result_idx = argspec.args.index(self.argname) - except ValueError: - result_idx = None - - if (self.argname in kargs): - result = kargs[self.argname] - elif (len(args) > result_idx and result_idx is not None): - result = args[result_idx] - else: - arg_idx_from_right = (len(argspec.args) - - argspec.args.index(self.argname)) - result = argspec.defaults[-arg_idx_from_right] + result = {} + for argname in arglist: + arg_loc = arg_location(func, argname, args, kwargs) + if arg_loc is not None: + result[argname] = arg_loc[0][arg_loc[1]] + else: + result[argname] = None + return result + +def arg_location(func, argname, args, kwargs): + """Parses the function args, kargs and signature looking for the desired + argument location (either in args, kargs, or argspec.defaults), + and returns a tuple of argument_sequence, location. + + """ + + argspec = getargspec(func) + print argspec + + if argname not in argspec.args and argname not in kwargs: + return None + + try: + result_idx = argspec.args.index(argname) + except ValueError: + result_idx = None + + result = None + if (argname in kwargs): + result = kwargs, argname + elif (len(args) > result_idx and result_idx is not None): + result = args, result_idx + else: + arg_idx_from_right = (len(argspec.args) - + argspec.args.index(argname)) + result = list(argspec.defaults), -arg_idx_from_right + + return result + def convert_units(unit_type, alg_unit): """A decorator that applies unit conversion to a function's output array. @@ -61,19 +84,19 @@ def convert_units(unit_type, alg_unit): - alg_unit - the units that the function returns by default """ - def convert_decorator(func): - @wraps(func) - def func_wrapper(*args, **kargs): - - desired_units = from_args("units")(func, *args, **kargs) - check_units(desired_units, unit_type) - - # Unit conversion done here - return do_conversion(func(*args, **kargs), unit_type, - alg_unit, desired_units) - return func_wrapper +# def convert_decorator(func): + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + + desired_units = from_args(wrapped, "units", *args, **kwargs)["units"] + check_units(desired_units, unit_type) + + # Unit conversion done here + return do_conversion(wrapped(*args, **kwargs), unit_type, + alg_unit, desired_units) + return func_wrapper - return convert_decorator +# return convert_decorator def _calc_out_dims(outvar, left_dims): @@ -105,223 +128,223 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, the output dimensions (e.g. avo) """ - def indexing_decorator(func): - @wraps(func) - def func_wrapper(*args, **kargs): - ignore_args = ignore_args if ignore_args is not None else () - ignore_kargs = ignore_kargs if ignore_kargs is not None else () - - if ref_var_idx >= 0: - ref_var = args[ref_var_idx] - else: - ref_var = kargs[ref_var_name] - - ref_var_shape = ref_var.shape - extra_dim_num = ref_var.ndim - ref_var_expected_dims - - # No special left side iteration, return the function result - if (extra_dim_num == 0): - return func(*args, **kargs) - - # Start by getting the left-most 'extra' dims - extra_dims = [ref_var_shape[x] for x in xrange(extra_dim_num)] - - out_inited = False - for left_idxs in iter_left_indexes(extra_dims): - # Make the left indexes plus a single slice object - # The single slice will handle all the dimensions to - # the right (e.g. [1,1,:]) - left_and_slice_idxs = ([x for x in left_idxs] + - [slice(None, None, None)]) - # Slice the args if applicable - new_args = [arg[left_and_slice_idxs] - if i not in ignore_args else arg - for i,arg in enumerate(args)] +# def indexing_decorator(func): + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + ignore_args = ignore_args if ignore_args is not None else () + ignore_kargs = ignore_kargs if ignore_kargs is not None else () + + if ref_var_idx >= 0: + ref_var = args[ref_var_idx] + else: + ref_var = kwargs[ref_var_name] + + ref_var_shape = ref_var.shape + extra_dim_num = ref_var.ndim - ref_var_expected_dims + + # No special left side iteration, return the function result + if (extra_dim_num == 0): + return wrapped(*args, **kwargs) + + # Start by getting the left-most 'extra' dims + extra_dims = [ref_var_shape[x] for x in xrange(extra_dim_num)] + + out_inited = False + for left_idxs in iter_left_indexes(extra_dims): + # Make the left indexes plus a single slice object + # The single slice will handle all the dimensions to + # the right (e.g. [1,1,:]) + left_and_slice_idxs = ([x for x in left_idxs] + + [slice(None, None, None)]) + # Slice the args if applicable + new_args = [arg[left_and_slice_idxs] + if i not in ignore_args else arg + for i,arg in enumerate(args)] + + # Slice the kargs if applicable + new_kargs = {key:(val[left_and_slice_idxs] + if key not in ignore_kargs else val) + for key,val in viewitems(kwargs)} + + # Call the numerical routine + res = wrapped(*new_args, **new_kargs) + + if isinstance(res, np.ndarray): + # Output array + if not out_inited: + outdims = _calc_out_dims(res, extra_dims) + if not isinstance(res, ma.MaskedArray): + output = np.empty(outdims, ref_var.dtype) + masked = False + else: + output = ma.MaskedArray( + np.zeros(outdims, ref_var.dtype), + mask=np.zeros(outdims, np.bool_), + fill_value=res.fill_value) + masked = True - # Slice the kargs if applicable - new_kargs = {key:(val[left_and_slice_idxs] - if key not in ignore_kargs else val) - for key,val in kargs.iteritems()} + out_inited = True - # Call the numerical routine - res = func(*new_args, **new_kargs) + if not masked: + output[left_and_slice_idxs] = res[:] + else: + output.data[left_and_slice_idxs] = res.data[:] + output.mask[left_and_slice_idxs] = res.mask[:] - if isinstance(res, n.ndarray): - # Output array - if not out_inited: - outdims = _calc_out_dims(res, extra_dims) - if not isinstance(res, ma.MaskedArray): - output = n.zeros(outdims, ref_var.dtype) - masked = False - else: - output = ma.MaskedArray( - n.zeros(outdims, ref_var.dtype), - mask=n.zeros(outdims, n.bool_), - fill_value=res.fill_value) - masked = True - - out_inited = True - - if not masked: - output[left_and_slice_idxs] = res[:] + else: # This should be a list or a tuple (cape) + if not out_inited: + outdims = _calc_out_dims(res[0], extra_dims) + if not isinstance(res[0], ma.MaskedArray): + output = [np.empty(outdims, ref_var.dtype) + for i in xrange(len(res))] + masked = False else: - output.data[left_and_slice_idxs] = res.data[:] - output.mask[left_and_slice_idxs] = res.mask[:] - - else: # This should be a list or a tuple (cape) - if not out_inited: - outdims = _calc_out_dims(res[0], extra_dims) - if not isinstance(res[0], ma.MaskedArray): - output = [n.zeros(outdims, ref_var.dtype) - for i in xrange(len(res))] - masked = False - else: - output = [ma.MaskedArray( - n.zeros(outdims, ref_var.dtype), - mask=n.zeros(outdims, n.bool_), - fill_value=res[0].fill_value) - for i in xrange(len(res))] - masked = True - - out_inited = True + output = [ma.MaskedArray( + np.zeros(outdims, ref_var.dtype), + mask=np.zeros(outdims, np.bool_), + fill_value=res[0].fill_value) + for i in xrange(len(res))] + masked = True - for i,outarr in enumerate(res): - if not masked: - output[i][left_and_slice_idxs] = outarr[:] - else: - output[i].data[left_and_slice_idxs] = outarr.data[:] - output[i].mask[left_and_slice_idxs] = outarr.mask[:] + out_inited = True + for i,outarr in enumerate(res): + if not masked: + output[i][left_and_slice_idxs] = outarr[:] + else: + output[i].data[left_and_slice_idxs] = outarr.data[:] + output[i].mask[left_and_slice_idxs] = outarr.mask[:] - return output - return func_wrapper + return output - return indexing_decorator + return func_wrapper + +# return indexing_decorator def uvmet_left_iter(): """Decorator to handle iterating over leftmost dimensions when using multiple files and/or multiple times with the uvmet product. """ - def indexing_decorator(func): - @wraps(func) - def func_wrapper(*args): - u = args[0] - v = args[1] - lat = args[2] - lon = args[3] - cen_long = args[4] - cone = args[5] - - if u.ndim == lat.ndim: - num_right_dims = 2 - is_3d = False - else: - num_right_dims = 3 - is_3d = True - - is_stag = False - if ((u.shape[-1] != lat.shape[-1]) or - (u.shape[-2] != lat.shape[-2])): - is_stag = True - - if is_3d: - extra_dim_num = u.ndim - 3 - else: - extra_dim_num = u.ndim - 2 - - if is_stag: - u = destagger(u,-1) - v = destagger(v,-2) - - # No special left side iteration, return the function result - if (extra_dim_num == 0): - return func(u,v,lat,lon,cen_long,cone) - - # Start by getting the left-most 'extra' dims - outdims = [u.shape[x] for x in xrange(extra_dim_num)] - extra_dims = list(outdims) # Copy the left-most dims for iteration + #def indexing_decorator(func): + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + u = args[0] + v = args[1] + lat = args[2] + lon = args[3] + cen_long = args[4] + cone = args[5] + + if u.ndim == lat.ndim: + num_right_dims = 2 + is_3d = False + else: + num_right_dims = 3 + is_3d = True + + is_stag = False + if ((u.shape[-1] != lat.shape[-1]) or + (u.shape[-2] != lat.shape[-2])): + is_stag = True + + if is_3d: + extra_dim_num = u.ndim - 3 + else: + extra_dim_num = u.ndim - 2 + + if is_stag: + u = destagger(u,-1) + v = destagger(v,-2) + + # No special left side iteration, return the function result + if (extra_dim_num == 0): + return wrapped(u,v,lat,lon,cen_long,cone) + + # Start by getting the left-most 'extra' dims + outdims = [u.shape[x] for x in xrange(extra_dim_num)] + extra_dims = list(outdims) # Copy the left-most dims for iteration + + # Append the right-most dimensions + outdims += [2] # For u/v components + + outdims += [u.shape[x] for x in xrange(-num_right_dims,0,1)] + + output = np.empty(outdims, u.dtype) + + for left_idxs in iter_left_indexes(extra_dims): + # Make the left indexes plus a single slice object + # The single slice will handle all the dimensions to + # the right (e.g. [1,1,:]) + left_and_slice_idxs = ([x for x in left_idxs] + + [slice(None, None, None)]) + - # Append the right-most dimensions - outdims += [2] # For u/v components + new_u = u[left_and_slice_idxs] + new_v = v[left_and_slice_idxs] + new_lat = lat[left_and_slice_idxs] + new_lon = lon[left_and_slice_idxs] - outdims += [u.shape[x] for x in xrange(-num_right_dims,0,1)] + # Call the numerical routine + res = wrapped(new_u, new_v, new_lat, new_lon, cen_long, cone) - output = n.zeros(outdims, u.dtype) + # Note: The 2D version will return a 3D array with a 1 length + # dimension. Numpy is unable to broadcast this without + # sqeezing first. + res = np.squeeze(res) - for left_idxs in iter_left_indexes(extra_dims): - # Make the left indexes plus a single slice object - # The single slice will handle all the dimensions to - # the right (e.g. [1,1,:]) - left_and_slice_idxs = ([x for x in left_idxs] + - [slice(None, None, None)]) - - - new_u = u[left_and_slice_idxs] - new_v = v[left_and_slice_idxs] - new_lat = lat[left_and_slice_idxs] - new_lon = lon[left_and_slice_idxs] - - # Call the numerical routine - res = func(new_u, new_v, new_lat, new_lon, cen_long, cone) - - # Note: The 2D version will return a 3D array with a 1 length - # dimension. Numpy is unable to broadcast this without - # sqeezing first. - res = n.squeeze(res) - - output[left_and_slice_idxs] = res[:] - + output[left_and_slice_idxs] = res[:] - return output - return func_wrapper + return output + + return func_wrapper - return indexing_decorator +# return indexing_decorator -def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=n.float64): +def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=np.float64): """Decorator to handle casting to/from required dtype used in numerical routine. """ - def cast_decorator(func): - @wraps(func) - def func_wrapper(*args, **kargs): - arg_idxs = arg_idxs if arg_idxs is not None else () - karg_names = karg_names if karg_names is not None else () - - orig_type = args[ref_idx].dtype - - new_args = [arg.astype(dtype) - if i in arg_idxs else arg - for i,arg in enumerate(args)] - - new_kargs = {key:(val.astype(dtype) - if key in karg_names else val) - for key,val in kargs.iteritems()} - - res = func(*new_args, **new_kargs) - - if isinstance(res, n.ndarray): - if res.dtype == orig_type: - return res - return res.astype(orig_type) - else: # got back a sequence of arrays - return tuple(arr.astype(orig_type) - if arr.dtype != orig_type else arr - for arr in res) - - return func_wrapper +# def cast_decorator(func): + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + arg_idxs = arg_idxs if arg_idxs is not None else () + karg_names = karg_names if karg_names is not None else () + + orig_type = args[ref_idx].dtype + + new_args = [arg.astype(dtype) + if i in arg_idxs else arg + for i,arg in enumerate(args)] + + new_kargs = {key:(val.astype(dtype) + if key in karg_names else val) + for key,val in viewitems()} + + res = wrapped(*new_args, **new_kargs) + + if isinstance(res, np.ndarray): + if res.dtype == orig_type: + return res + return res.astype(orig_type) + else: # got back a sequence of arrays + return tuple(arr.astype(orig_type) + if arr.dtype != orig_type else arr + for arr in res) + + return func_wrapper - return cast_decorator +# return cast_decorator def _extract_and_transpose(arg): if xarray_enabled(): if isinstance(arg, DataArray): arg = arg.values - if isinstance(arg, n.ndarray): + if isinstance(arg, np.ndarray): if not arg.flags.F_CONTIGUOUS: return arg.T @@ -332,53 +355,305 @@ def handle_extract_transpose(): transposes if the data is not fortran contiguous. """ - def extract_transpose_decorator(func): - @wraps(func) - def func_wrapper(*args, **kargs): - - new_args = [_extract_and_transpose(arg) for arg in args] - - new_kargs = {key:_extract_and_transpose(val) - for key,val in kargs.iteritems()} - - res = func(*new_args, **new_kargs) - - if isinstance(res, n.ndarray): - if res.flags.F_CONTIGUOUS: - return res.T - else: - return tuple(x.T if x.flags.F_CONTIGUOUS else x for x in res) +# def extract_transpose_decorator(func): + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + + new_args = [_extract_and_transpose(arg) for arg in args] + + new_kargs = {key:_extract_and_transpose(val) + for key,val in viewitems(kwargs)} - return res + res = wrapped(*new_args, **new_kargs) - return func_wrapper + if isinstance(res, np.ndarray): + if res.flags.F_CONTIGUOUS: + return res.T + else: + return tuple(x.T if x.flags.F_CONTIGUOUS else x for x in res) + + return res + + return func_wrapper - return extract_transpose_decorator +# return extract_transpose_decorator + -def set_metadata(copy_from=None, ignore=None, **fixed_kargs): - """Decorator to set the attributes for a WRF method. +def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, + dimnames=None, coords=None, **fixed_attrs): + """Decorator to set the metadata for a WRF method. + + A cache is inserted/updated to include the extracted variable that will + have its metadata copied. This prevents the variable being extracted more + than once. This extraction can be slow with sequences of large files. """ - def attr_decorator(func): - @wraps(func) - def func_wrapper(*args, **kargs): - if not xarray_enabled(): - return func(*args, **kargs) + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + if not xarray_enabled(): + return wrapped(*args, **kwargs) + + argvars = from_args(wrapped, ("wrfnc", "timeidx", "method", + "squeeze", "cache", "units"), + *args, **kwargs) + wrfnc = argvars["wrfnc"] + timeidx = argvars["timeidx"] + units = argvars["units"] + method = argvars["method"] + squeeze = argvars["squeeze"] + cache = argvars["cache"] + if cache is None: + cache = {} + + + if (callable(copy_varname)): + copy_varname = copy_varname(wrfnc) + + # Extract the copy_from argument + var_to_copy = None if cache is None else cache.get(copy_varname, + None) + - result = func(*args, **kargs) + if var_to_copy is None: + var_to_copy = extract_vars(wrfnc, timeidx, (copy_varname,), + method, squeeze, cache)[copy_varname] + + # Make a copy so we don't modify a user supplied cache + new_cache = dict(cache) + new_cache[copy_varname] = var_to_copy + + # Don't modify the original args/kargs. The args need to be a list + # so it can be modified. + new_args = list(args) + new_kargs = dict(kwargs) + cache_argseq, cache_argloc = arg_location(wrapped, "cache", + new_args, new_kargs) + + cache_argseq[cache_argloc] = new_cache + + result = wrapped(*new_args, **new_kargs) + + outname = "" + outdimnames = list() + outcoords = OrderedDict() + outattrs = OrderedDict() + + if copy_varname is not None: + outname = str(var_to_copy.name) + + if dimnames is not None: + if isinstance(combine_with): + outdimnames, outcoords = dimnames(copy_varname) + else: + outdimnames = dimnames + outcoords = coords + else: + outdimnames += var_to_copy.dims + outcoords.update(var_to_copy.coords) - units = from_args("units")(func, *args, **kargs) - if units is not None: - result.attrs["units"] = units + outattrs.update(var_to_copy.attrs) + + if name is not None: + outname = name + + if units is not None: + outattrs["units"] = units - for argname, val in fixed_kargs.iteritems(): - result[argname] = val + for argname, val in viewitems(fixed_attrs): + outattrs[argname] = val + + if delete_attrs is not None: + for attr in delete_attrs: + del outattrs[attr] + + if isinstance(ma.MaskedArray): + outattrs["_FillValue"] = result.fill_value + outattrs["missing_value"] = result.fill_value + + return DataArray(result, name=outname, coords=outcoords, + dims=outdimnames, attrs=outattrs) + + return func_wrapper + + +def set_wind_metadata(wspd_wdir=False): + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + if not xarray_enabled(): + return wrapped(*args, **kwargs) + + argvars = from_args(wrapped, ("wrfnc", "timeidx", "units", + "method", "squeeze", "ten_m", "cache"), + *args, **kwargs) + wrfnc = argvars["wrfnc"] + timeidx = argvars["timeidx"] + units = argvars["units"] + method = argvars["method"] + squeeze = argvars["squeeze"] + ten_m = argvars["ten_m"] + cache = argvars["cache"] + if cache is None: + cache = {} + + lat_var = either("XLAT", "XLAT_M")(wrfnc) + xlat_var = extract_vars(wrfnc, timeidx, lat_var, + method, squeeze, cache) + lat = xlat_var[lat_var] + + lon_var = either("XLONG", "XLONG_M") + xlon_var = extract_vars(wrfnc, timeidx, lon_var, + method, squeeze, cache) + lon = xlon_var[lon_var] + + pres_var = either("P", "PRES") + p_var = extract_vars(wrfnc, timeidx, lon_var, method, squeeze, cache) + pres = p_var[pres_var] + + + # Make a copy so we don't modify a user supplied cache + new_cache = dict(cache) + new_cache[lat_var] = lat + new_cache[lon_var] = lon + new_cache[pres_var] = pres + + # Don't modify the original args/kargs. The args need to be a list + # so it can be modified. + new_args = list(args) + new_kargs = dict(kwargs) + cache_argseq, cache_argloc = arg_location(wrapped, "cache", + new_args, new_kargs) + + cache_argseq[cache_argloc] = new_cache + + result = wrapped(*new_args, **new_kargs) + + outcoords = OrderedDict() + outattrs = OrderedDict() + outname = "uvmet" if not ten_m else "uvmet10" + + outdimnames = list(pres.dims) + outcoords.update(pres.coords) + outattrs.update(pres.attrs) + + if not wspd_wdir: + outattrs["description"] = ("earth rotated u,v" if not ten_m else + "10m earth rotated u,v") + if not ten_m: + outdimnames.insert(-3, "u_v") + else: + outdimnames.insert(-2, "u_v") + else: + outattrs["description"] = ("earth rotated wspd,wdir" if not ten_m + else "10m earth rotated wspd,wdir") + if not ten_m: + outdimnames.insert(-3, "wspd_wdir") + else: + outdimnames.insert(-2, "wspd_wdir") + + if units is not None: + outattrs["units"] = units + + return DataArray(result, name=outname, coords=outcoords, + dims=outdimnames, attrs=outattrs) + + return func_wrapper + +def set_latlon_metadata(ij=False): + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + if not xarray_enabled(): + return wrapped(*args, **kwargs) + + res = wrapped(*args, **kwargs) + + argnames = ("latitude", "longitude") if not ij else ("i", "j") + outname = "latlon" if not ij else "ij" + + if res.ndim <= 2: + dimnames = (["lat_lon", "i_j"] if not ij + else ["i_j", "lat_lon"]) + else: + dimnames = (["lat_lon", "domain", "i_j"] if not ij + else ["i_j", "domain", "lat_lon"]) - return result - return func_wrapper + argvars = from_args(wrapped, argnames, *args, **kwargs) + + var1 = argvars[argnames[0]] + var2 = argvars[argnames[1]] + + arr1 = np.asarray(var1).ravel() + arr2 = np.asarray(var2).ravel() + + coords = {} + coords[dimnames[0]] = [x for x in zip(arr1, arr2)] + + return DataArray(res, name=outname, dims=dimnames, coords=coords) + + return func_wrapper - return attr_decorator +def set_height_metadata(geopt=False): + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + if not xarray_enabled(): + return wrapped(*args, **kwargs) + + argvars = from_args(wrapped, ("wrfnc", "timeidx", "method", + "squeeze", "units", "msl", "cache"), + *args, **kwargs) + wrfnc = argvars["wrfnc"] + timeidx = argvars["timeidx"] + units = argvars["units"] + method = argvars["method"] + squeeze = argvars["squeeze"] + msl = argvars["msl"] + cache = argvars["cache"] + if cache is None: + cache = {} + + # For height, either copy the met_em GHT variable or copy and modify + # pressure (which has the same dims as destaggered height) + ht_metadata_varname = either("P", "GHT")(wrfnc) + ht_var = extract_vars(wrfnc, timeidx, ht_metadata_varname, + method, squeeze, cache) + ht_metadata_var = ht_var[ht_metadata_varname] + + # Make a copy so we don't modify a user supplied cache + new_cache = dict(cache) + new_cache[ht_metadata_var] = ht_metadata_var + + # Don't modify the original args/kargs. The args need to be a list + # so it can be modified. + new_args = list(args) + new_kargs = dict(kwargs) + cache_argseq, cache_argloc = arg_location(wrapped, "cache", + new_args, new_kargs) + + cache_argseq[cache_argloc] = new_cache + + result = wrapped(*new_args, **new_kargs) + + outcoords = OrderedDict() + outattrs = OrderedDict() + outdimnames = list(ht_metadata_var.dims) + outcoords.update(ht_metadata_var.coords) + outattrs.update(ht_metadata_var.attrs) + + if geopt: + outname = "geopt" + outattrs["units"] = "m2 s-2" + outattrs["description"] = "full model geopotential" + else: + outname = "height" if msl else "height_agl" + outattrs["units"] = units + height_type = "MSL" if msl else "AGL" + outattrs["description"] = "model height ({})".format(height_type) + + + return DataArray(result, name=outname, + dims=outdimnames, coords=outcoords) + + diff --git a/wrf_open/var/src/python/wrf/var/destag.py b/wrf_open/var/src/python/wrf/var/destag.py index 7652ce1..2ab2a6f 100755 --- a/wrf_open/var/src/python/wrf/var/destag.py +++ b/wrf_open/var/src/python/wrf/var/destag.py @@ -1,6 +1,6 @@ -from wrf.var.util import extract_vars +from .util import extract_vars __all__ = ["destagger", "destagger_windcomp", "destagger_winds"] @@ -40,24 +40,25 @@ def destagger(var, stagger_dim): return result -def destagger_windcomp(wrfnc, comp, timeidx=0): - if comp.lower() == "u": - wrfvar = "U" - stagdim = -1 - elif comp.lower() == "v": - wrfvar = "V" - stagdim = -2 - elif comp.lower() == "w": - wrfvar = "W" - stagdim = -3 - - #wind_data = wrfnc.variables[wrfvar][timeidx,:,:,:] - ncvars = extract_vars(wrfnc, timeidx, vars=wrfvar) - wind_data = ncvars[wrfvar] - return destagger(wind_data, stagdim) - -def destagger_winds(wrfnc, timeidx=0): - return (destagger_windcomp(wrfnc, "u", timeidx), - destagger_windcomp(wrfnc, "v", timeidx), - destagger_windcomp(wrfnc, "w", timeidx)) +# def destagger_windcomp(wrfnc, comp, timeidx=0, +# method="cat", squeeze=True, cache=None): +# if comp.lower() == "u": +# wrfvar = "U" +# stagdim = -1 +# elif comp.lower() == "v": +# wrfvar = "V" +# stagdim = -2 +# elif comp.lower() == "w": +# wrfvar = "W" +# stagdim = -3 +# +# ncvars = extract_vars(wrfnc, timeidx, wrfvar, method, squeeze, cache) +# wind_data = ncvars[wrfvar] +# +# return destagger(wind_data, stagdim) +# +# def destagger_winds(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): +# return (destagger_windcomp(wrfnc, "u", timeidx, method, squeeze, cache), +# destagger_windcomp(wrfnc, "v", timeidx, method, squeeze, cache), +# destagger_windcomp(wrfnc, "w", timeidx, method, squeeze, cache)) \ No newline at end of file diff --git a/wrf_open/var/src/python/wrf/var/dewpoint.py b/wrf_open/var/src/python/wrf/var/dewpoint.py index 14f09ff..4b55047 100755 --- a/wrf_open/var/src/python/wrf/var/dewpoint.py +++ b/wrf_open/var/src/python/wrf/var/dewpoint.py @@ -1,13 +1,17 @@ -from wrf.var.extension import computetd -from wrf.var.decorators import convert_units -from wrf.var.util import extract_vars +from .extension import computetd +from .decorators import convert_units, copy_and_set_metadata +from .util import extract_vars __all__ = ["get_dp", "get_dp_2m"] +@copy_and_set_metadata(copy_varname="QVAPOR", name="td", + description="dew point temperature") @convert_units("temp", "c") -def get_dp(wrfnc, timeidx=0, units="c"): +def get_dp(wrfnc, timeidx=0, units="c", + method="cat", squeeze=True, cache=None): - ncvars = extract_vars(wrfnc, timeidx, varnames=("P", "PB", "QVAPOR")) + varnames=("P", "PB", "QVAPOR") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) p = ncvars["P"] pb = ncvars["PB"] @@ -19,10 +23,14 @@ def get_dp(wrfnc, timeidx=0, units="c"): td = computetd(full_p, qvapor) return td - + +@copy_and_set_metadata(copy_varname="Q2", name="td2", + description="2m dew point temperature") @convert_units("temp", "c") -def get_dp_2m(wrfnc, timeidx=0, units="c"): - ncvars = extract_vars(wrfnc, timeidx, varnames=("PSFC", "Q2")) +def get_dp_2m(wrfnc, timeidx=0, units="c", + method="cat", squeeze=True, cache=None): + varnames=("PSFC", "Q2") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) # Algorithm requires hPa psfc = .01*(ncvars["PSFC"]) diff --git a/wrf_open/var/src/python/wrf/var/extension.py b/wrf_open/var/src/python/wrf/var/extension.py index 1981b06..47af45f 100755 --- a/wrf_open/var/src/python/wrf/var/extension.py +++ b/wrf_open/var/src/python/wrf/var/extension.py @@ -1,8 +1,8 @@ -import numpy as n +import numpy as np -from wrf.var.constants import Constants -from wrf.var.psadlookup import get_lookup_tables -from wrf.var._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d, +from .constants import Constants +from .psadlookup import get_lookup_tables +from ._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d, f_computeslp, f_computetk, f_computetd, f_computerh, f_computeabsvort,f_computepvo, f_computeeth, f_computeuvmet, @@ -10,8 +10,8 @@ from wrf.var._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d, f_computesrh, f_computeuh, f_computepw, f_computedbz, f_lltoij, f_ijtoll, f_converteta, f_computectt, f_monotonic, f_filter2d, f_vintrp) -from wrf.var._wrfcape import f_computecape -from wrf.var.decorators import (handle_left_iter, uvmet_left_iter, +from ._wrfcape import f_computecape +from .decorators import (handle_left_iter, uvmet_left_iter, handle_casting, handle_extract_transpose) __all__ = ["FortranException", "computeslp", "computetk", "computetd", @@ -27,7 +27,7 @@ class FortranException(Exception): @handle_left_iter(3,0, ignore_args=(2,3)) @handle_casting(arg_idxs=(0,1)) @handle_extract_transpose() -def interpz3d(data3d,zdata,desiredloc,missingval): +def interpz3d(data3d, zdata, desiredloc, missingval): res = f_interpz3d(data3d, zdata, desiredloc, @@ -44,7 +44,7 @@ def interp2dxy(data3d,xy): @handle_casting(arg_idxs=(0,1,2)) @handle_extract_transpose() -def interp1d(v_in,z_in,z_out,missingval): +def interp1d(v_in, z_in, z_out, missingval): res = f_interp1d(v_in, z_in, z_out, @@ -55,10 +55,10 @@ def interp1d(v_in,z_in,z_out,missingval): @handle_left_iter(3,0) @handle_casting(arg_idxs=(0,1,2,3)) @handle_extract_transpose() -def computeslp(z,t,p,q): - t_surf = n.zeros((z.shape[-2], z.shape[-1]), "float64", order="F") - t_sea_level = n.zeros((z.shape[-2], z.shape[-1]), "float64", order="F") - level = n.zeros((z.shape[-2], z.shape[-1]), "int32", order="F") +def computeslp(z, t, p, q): + t_surf = np.zeros((z.shape[-2], z.shape[-1]), np.float64, order="F") + t_sea_level = np.zeros((z.shape[-2], z.shape[-1]), np.float64, order="F") + level = np.zeros((z.shape[-2], z.shape[-1]), np.int32, order="F") res = f_computeslp(z, t, @@ -79,34 +79,34 @@ def computetk(pres, theta): shape = pres.shape res = f_computetk(pres.ravel(order="A"), theta.ravel(order="A")) - res = n.reshape(res, shape) + res = np.reshape(res, shape) return res # Note: No left iteration decorator needed with 1D arrays @handle_casting(arg_idxs=(0,1)) @handle_extract_transpose() -def computetd(pressure,qv_in): +def computetd(pressure, qv_in): shape = pressure.shape res = f_computetd(pressure.ravel(order="A"), qv_in.ravel(order="A")) - res = n.reshape(res, shape) + res = np.reshape(res, shape) return res # Note: No decorator needed with 1D arrays @handle_casting(arg_idxs=(0,1,2)) @handle_extract_transpose() -def computerh(qv,q,t): +def computerh(qv, q, t): shape = qv.shape res = f_computerh(qv.ravel(order="A"), q.ravel(order="A"), t.ravel(order="A")) - res = n.reshape(res, shape) + res = np.reshape(res, shape) return res @handle_left_iter(3,0, ignore_args=(6,7)) @handle_casting(arg_idxs=(0,1,2,3,4,5)) @handle_extract_transpose() -def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy): +def computeavo(u, v, msfu, msfv, msfm, cor, dx, dy): res = f_computeabsvort(u, v, msfu, @@ -121,7 +121,7 @@ def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy): @handle_left_iter(3,2, ignore_args=(8,9)) @handle_casting(arg_idxs=(0,1,2,3,4,5,6,7)) @handle_extract_transpose() -def computepvo(u,v,theta,prs,msfu,msfv,msfm,cor,dx,dy): +def computepvo(u, v, theta, prs, msfu, msfv, msfm, cor, dx, dy): res = f_computepvo(u, v, @@ -150,9 +150,9 @@ def computeeth(qv, tk, p): @uvmet_left_iter() @handle_casting(arg_idxs=(0,1,2,3)) @handle_extract_transpose() -def computeuvmet(u,v,lat,lon,cen_long,cone): - longca = n.zeros((lat.shape[-2], lat.shape[-1]), "float64", order="F") - longcb = n.zeros((lon.shape[-2], lon.shape[-1]), "float64", order="F") +def computeuvmet(u, v, lat, lon, cen_long, cone): + longca = np.zeros((lat.shape[-2], lat.shape[-1]), np.float64, order="F") + longcb = np.zeros((lon.shape[-2], lon.shape[-1]), np.float64, order="F") rpd = Constants.PI/180. @@ -179,7 +179,7 @@ def computeuvmet(u,v,lat,lon,cen_long,cone): 0) - return n.squeeze(res) + return np.squeeze(res) @handle_left_iter(3,0) @handle_casting(arg_idxs=(0,1,2,3)) @@ -197,7 +197,7 @@ def computeomega(qv, tk, w, p): @handle_left_iter(3,0) @handle_casting(arg_idxs=(0,1)) @handle_extract_transpose() -def computetv(tk,qv): +def computetv(tk, qv): res = f_computetv(tk, qv) @@ -206,7 +206,7 @@ def computetv(tk,qv): @handle_left_iter(3,0) @handle_casting(arg_idxs=(0,1,2)) @handle_extract_transpose() -def computewetbulb(p,tk,qv): +def computewetbulb(p, tk, qv): PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables() PSADITMK = PSADITMK.T @@ -238,8 +238,10 @@ def computesrh(u, v, z, ter, top): @handle_extract_transpose() def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top): - tem1 = n.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), "float64", order="F") - tem2 = n.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), "float64", order="F") + tem1 = np.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), np.float64, + order="F") + tem2 = np.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), np.float64, + order="F") res = f_computeuh(zp, mapfct, @@ -258,9 +260,9 @@ def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top): @handle_left_iter(3,0) @handle_casting(arg_idxs=(0,1,2,3)) @handle_extract_transpose() -def computepw(p,tv,qv,ht): +def computepw(p, tv, qv, ht): # Note, dim -3 is height, we only want y and x - zdiff = n.zeros((p.shape[-2], p.shape[-1]), "float64", order="F") + zdiff = np.zeros((p.shape[-2], p.shape[-1]), np.float64, order="F") res = f_computepw(p, tv, qv, @@ -272,7 +274,7 @@ def computepw(p,tv,qv,ht): @handle_left_iter(3,0, ignore_args=(6,7,8)) @handle_casting(arg_idxs=(0,1,2,3,4,5)) @handle_extract_transpose() -def computedbz(p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin): +def computedbz(p, tk, qv, qr, qs, qg, sn0, ivarint, iliqskin): res = f_computedbz(p, tk, @@ -289,11 +291,11 @@ def computedbz(p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin): @handle_left_iter(3,0,ignore_args=(6,7,8)) @handle_casting(arg_idxs=(0,1,2,3,4,5)) @handle_extract_transpose() -def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow): - flip_cape = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), - "float64", order="F") - flip_cin = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), - "float64", order="F") +def computecape(p_hpa, tk, qv, ht, ter, sfp, missing, i3dflag, ter_follow): + flip_cape = np.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), + np.float64, order="F") + flip_cin = np.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), + np.float64, order="F") PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables() PSADITMK = PSADITMK.T @@ -327,37 +329,34 @@ def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow): return (cape, cin) -# TODO: This should handle lists of coords -def computeij(map_proj,truelat1,truelat2,stdlon, - lat1,lon1,pole_lat,pole_lon, - knowni,knownj,dx,latinc,loninc,lat,lon): +def computeij(map_proj, truelat1, truelat2, stdlon, + lat1, lon1, pole_lat, pole_lon, + knowni, knownj, dx, latinc, loninc, lat, lon): res = f_lltoij(map_proj,truelat1,truelat2,stdlon, lat1,lon1,pole_lat,pole_lon, knowni,knownj,dx,latinc,loninc,lat,lon, FortranException()) - return res[0],res[1] + return res -# TODO: This should handle lists of coords -def computell(map_proj,truelat1,truelat2,stdlon,lat1,lon1, - pole_lat,pole_lon,knowni,knownj,dx,latinc, - loninc,i,j): +def computell(map_proj, truelat1, truelat2, stdlon, lat1, lon1, + pole_lat, pole_lon, knowni, knownj, dx, latinc, + loninc, i, j): res = f_ijtoll(map_proj,truelat1,truelat2,stdlon,lat1,lon1, pole_lat,pole_lon,knowni,knownj,dx,latinc, loninc,i,j,FortranException()) - # Want lon,lat - return res[1],res[0] + return res @handle_left_iter(3,0, ignore_args=(3,)) @handle_casting(arg_idxs=(0,1,2)) @handle_extract_transpose() def computeeta(full_t, znu, psfc, ptop): - pcalc = n.zeros(full_t.shape, "float64", order="F") - mean_t = n.zeros(full_t.shape, "float64", order="F") - temp_t = n.zeros(full_t.shape, "float64", order="F") + pcalc = np.zeros(full_t.shape, np.float64, order="F") + mean_t = np.zeros(full_t.shape, np.float64, order="F") + temp_t = np.zeros(full_t.shape, np.float64, order="F") res = f_converteta(full_t, znu, @@ -372,7 +371,7 @@ def computeeta(full_t, znu, psfc, ptop): @handle_left_iter(3,0,ignore_args=(7,)) @handle_casting(arg_idxs=(0,1,2,3,4,5,6)) @handle_extract_transpose() -def computectt(p_hpa,tk,qice,qcld,qv,ght,ter,haveqci): +def computectt(p_hpa, tk, qice, qcld, qv, ght, ter, haveqci): res = f_computectt(p_hpa, tk, qice, @@ -392,13 +391,13 @@ def smooth2d(field, passes): # copies the original data before modifying it. This allows the decorator # to work properly and also behaves like the other methods. - if isinstance(field, n.ma.MaskedArray): + if isinstance(field, np.ma.MaskedArray): missing = field.fill_value else: missing = Constants.DEFAULT_FILL field_copy = field.copy() - field_tmp = n.zeros(field_copy.shape, field_copy.dtype, order="F") + field_tmp = np.zeros(field_copy.shape, field_copy.dtype, order="F") f_filter2d(field_copy, field_tmp, @@ -412,7 +411,7 @@ def smooth2d(field, passes): @handle_left_iter(3,0,ignore_args=(3,4,5)) @handle_casting(arg_idxs=(0,1,2)) @handle_extract_transpose() -def monotonic(var,lvprs,coriolis,idir,delta,icorsw): +def monotonic(var, lvprs, coriolis, idir, delta, icorsw): res = f_monotonic(var, lvprs, coriolis, @@ -425,8 +424,8 @@ def monotonic(var,lvprs,coriolis,idir,delta,icorsw): @handle_left_iter(3,0,ignore_args=(9,10,11,12,13,14)) @handle_casting(arg_idxs=(0,1,2,3,4,5,6,7,8,9)) @handle_extract_transpose() -def vintrp(field,pres,tk,qvp,ght,terrain,sfp,smsfp, - vcarray,interp_levels,icase,extrap,vcor,logp, +def vintrp(field, pres, tk, qvp, ght, terrain, sfp, smsfp, + vcarray, interp_levels, icase, extrap, vcor, logp, missing): res = f_vintrp(field, diff --git a/wrf_open/var/src/python/wrf/var/geoht.py b/wrf_open/var/src/python/wrf/var/geoht.py index 6dce908..dc40b47 100755 --- a/wrf_open/var/src/python/wrf/var/geoht.py +++ b/wrf_open/var/src/python/wrf/var/geoht.py @@ -1,11 +1,12 @@ -from wrf.var.constants import Constants -from wrf.var.destag import destagger -from wrf.var.decorators import convert_units -from wrf.var.util import extract_vars +from .constants import Constants +from .destag import destagger +from .decorators import convert_units, set_height_metadata +from .util import extract_vars, either __all__ = ["get_geopt", "get_height"] -def _get_geoht(wrfnc, timeidx, height=True, msl=True): +def _get_geoht(wrfnc, timeidx, height=True, msl=True, + method="cat", squeeze=True, cache=None): """Return the geopotential in units of m2 s-2 if height is False, otherwise return the geopotential height in meters. If height is True, then if msl is True the result will be in MSL, otherwise AGL (the terrain @@ -13,23 +14,18 @@ def _get_geoht(wrfnc, timeidx, height=True, msl=True): """ - try: + varname = either("PH", "GHT")(wrfnc) + if varname == "PH": ph_vars = extract_vars(wrfnc, timeidx, varnames=("PH", "PHB", "HGT")) - except KeyError: - try: - ght_vars = extract_vars(wrfnc, timeidx, varnames=("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, -3) + else: + ght_vars = extract_vars(wrfnc, timeidx, varnames=("GHT", "HGT_M")) + geopt_unstag = ght_vars["GHT"] * Constants.G + hgt = ght_vars["HGT_M"] if height: if msl: @@ -46,11 +42,14 @@ def _get_geoht(wrfnc, timeidx, height=True, msl=True): else: return geopt_unstag -def get_geopt(wrfnc, timeidx=0): - return _get_geoht(wrfnc, timeidx, False) +@set_height_metadata(geopt=True) +def get_geopt(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): + return _get_geoht(wrfnc, timeidx, False, True, method, squeeze, cache) +@set_height_metadata(geopt=False) @convert_units("height", "m") -def get_height(wrfnc, timeidx=0, msl=True, units="m"): - z = _get_geoht(wrfnc, timeidx, True, msl) - return z +def get_height(wrfnc, timeidx=0, msl=True, units="m", + method="cat", squeeze=True, cache=None): + + return _get_geoht(wrfnc, timeidx, True, msl, method, squeeze, cache) diff --git a/wrf_open/var/src/python/wrf/var/helicity.py b/wrf_open/var/src/python/wrf/var/helicity.py index 1483204..47b86cc 100755 --- a/wrf_open/var/src/python/wrf/var/helicity.py +++ b/wrf_open/var/src/python/wrf/var/helicity.py @@ -1,43 +1,34 @@ -from wrf.var.constants import Constants +from .constants import Constants -from wrf.var.extension import computesrh, computeuh -from wrf.var.destag import destagger -from wrf.var.util import extract_vars, extract_global_attrs +from .extension import computesrh, computeuh +from .destag import destagger +from .util import extract_vars, extract_global_attrs, either +from .decorators import copy_and_set_metadata __all__ = ["get_srh", "get_uh"] -def get_srh(wrfnc, timeidx=0, top=3000.0): +@copy_and_set_metadata(copy_varname="HGT", name="srh", + description="storm relative helicity", + units="m-2/s-2") +def get_srh(wrfnc, timeidx=0, top=3000.0, + method="cat", squeeze=True, cache=None): # Top can either be 3000 or 1000 (for 0-1 srh or 0-3 srh) - ncvars = extract_vars(wrfnc, timeidx, varnames=("HGT", "PH", "PHB")) + ncvars = extract_vars(wrfnc, timeidx, ("HGT", "PH", "PHB"), + method, squeeze, cache) ter = ncvars["HGT"] ph = ncvars["PH"] phb = ncvars["PHB"] - try: - u_vars = extract_vars(wrfnc, timeidx, varnames="U") - except KeyError: - try: - uu_vars = extract_vars(wrfnc, timeidx, varnames="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, varnames="V") - except KeyError: - try: - vv_vars = extract_vars(wrfnc, timeidx, varnames="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) + # As coded in NCL, but not sure this is possible + varname = either("U", "UU")(wrfnc) + u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + u = destagger(u_vars[varname], -1) + + varname = either("V", "VV")(wrfnc) + v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + v = destagger(v_vars[varname], -2) geopt = ph + phb geopt_unstag = destagger(geopt, -3) @@ -53,9 +44,14 @@ def get_srh(wrfnc, timeidx=0, top=3000.0): return srh -def get_uh(wrfnc, timeidx=0, bottom=2000.0, top=5000.0): +@copy_and_set_metadata(copy_varname="MAPFAC_M", name="updraft_helicity", + description="updraft helicity", + units="m-2/s-2") +def get_uh(wrfnc, timeidx=0, bottom=2000.0, top=5000.0, + method="cat", squeeze=True, cache=None): - ncvars = extract_vars(wrfnc, timeidx, varnames=("W", "PH", "PHB", "MAPFAC_M")) + ncvars = extract_vars(wrfnc, timeidx, ("W", "PH", "PHB", "MAPFAC_M"), + method, squeeze, cache) wstag = ncvars["W"] ph = ncvars["PH"] @@ -66,29 +62,14 @@ def get_uh(wrfnc, timeidx=0, bottom=2000.0, top=5000.0): dx = attrs["DX"] dy = attrs["DY"] - try: - u_vars = extract_vars(wrfnc, timeidx, varnames="U") - except KeyError: - try: - uu_vars = extract_vars(wrfnc, timeidx, varnames="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, varnames="V") - except KeyError: - try: - vv_vars = extract_vars(wrfnc, timeidx, varnames="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) + # As coded in NCL, but not sure this is possible + varname = either("U", "UU")(wrfnc) + u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + u = destagger(u_vars[varname], -1) + + varname = either("V", "VV")(wrfnc) + v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + v = destagger(v_vars[varname], -2) zp = ph + phb diff --git a/wrf_open/var/src/python/wrf/var/interp.py b/wrf_open/var/src/python/wrf/var/interp.py index b7b7bb3..802c91e 100755 --- a/wrf_open/var/src/python/wrf/var/interp.py +++ b/wrf_open/var/src/python/wrf/var/interp.py @@ -3,20 +3,20 @@ from math import floor, ceil import numpy as n import numpy.ma as ma -from wrf.var.extension import (interpz3d,interp2dxy,interp1d, +from .extension import (interpz3d,interp2dxy,interp1d, smooth2d,monotonic,vintrp) -from wrf.var.decorators import handle_left_iter, handle_casting -from wrf.var.util import extract_vars, is_staggered -from wrf.var.constants import Constants, ConversionFactors -from wrf.var.terrain import get_terrain -from wrf.var.geoht import get_height -from wrf.var.temp import get_theta, get_temp, get_eth -from wrf.var.pressure import get_pressure +from .decorators import handle_left_iter, handle_casting +from .util import extract_vars, is_staggered +from .constants import Constants, ConversionFactors +from .terrain import get_terrain +from .geoht import get_height +from .temp import get_theta, get_temp, get_eth +from .pressure import get_pressure __all__ = ["interplevel", "vertcross", "interpline", "vinterp"] # Note: Extension decorator is good enough to handle left dims -def interplevel(data3d,zdata,desiredloc,missingval=Constants.DEFAULT_FILL): +def interplevel(data3d, zdata, desiredloc, missingval=Constants.DEFAULT_FILL): """Return the horizontally interpolated data at the provided level data3d - the 3D field to interpolate @@ -28,11 +28,13 @@ def interplevel(data3d,zdata,desiredloc,missingval=Constants.DEFAULT_FILL): """ r1 = interpz3d(data3d, zdata, desiredloc, missingval) masked_r1 = ma.masked_values (r1, missingval) + return masked_r1 def _to_positive_idxs(shape, coord): - if (coord[-2] >= 0 and coord[-1] >=0): + if (coord[-2] >= 0 and coord[-1] >= 0): return coord + return [x if (x >= 0) else shape[i]+x for (i,x) in enumerate(coord) ] def _get_xy(xdim, ydim, pivot_point=None, angle=None, @@ -41,7 +43,8 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None, xdim - maximum x-dimension ydim - maximum y-dimension - pivot_point - a pivot point of (south_north, west_east) (must be used with angle) + pivot_point - a pivot point of (south_north, west_east) + (must be used with angle) angle - the angle through the pivot point in degrees start_point - a start_point sequence of [south_north1, west_east1] end_point - an end point sequence of [south_north2, west_east2] @@ -148,15 +151,16 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None, "start_point", "end_point")) @handle_casting(arg_idxs=(0,1)) def vertcross(data3d, z, missingval=Constants.DEFAULT_FILL, - pivot_point=None,angle=None, - start_point=None,end_point=None): + pivot_point=None, angle=None, + start_point=None, end_point=None): """Return the vertical cross section for a 3D field, interpolated to a verical plane defined by a horizontal line. Arguments: data3d - a 3D data field z - 3D height field - pivot_point - a pivot point of (south_north,west_east) (must be used with angle) + pivot_point - a pivot point of (south_north,west_east) + (must be used with angle) angle - the angle through the pivot point in degrees start_point - a start_point tuple of (south_north1,west_east1) end_point - an end point tuple of (south_north2,west_east2) @@ -362,7 +366,7 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, p_hpa = p * ConversionFactors.PA_TO_HPA - vcord_array = monotonic(t,p_hpa,coriolis,idir,delta,icorsw) + vcord_array = monotonic(t, p_hpa, coriolis, idir, delta, icorsw) # We only extrapolate temperature fields below ground # if we are interpolating to pressure or height vertical surfaces. @@ -379,7 +383,7 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, p_hpa = p * ConversionFactors.PA_TO_HPA - vcord_array = monotonic(eth,p_hpa,coriolis,idir,delta,icorsw) + vcord_array = monotonic(eth, p_hpa, coriolis, idir, delta, icorsw) # We only extrapolate temperature fields below ground if we are # interpolating to pressure or height vertical surfaces icase = 0 @@ -390,11 +394,11 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, else: missing = Constants.DEFAULT_FILL - res = vintrp(field,p,tk,qv,ght,terht,sfp,smsfp, - vcord_array,interp_levels, - icase,extrap,vcor,log_p_int,missing) + res = vintrp(field, p, tk, qv, ght, terht, sfp, smsfp, + vcord_array, interp_levels, + icase, extrap, vcor, log_p_int, missing) - return ma.masked_values(res,missing) + return ma.masked_values(res, missing) diff --git a/wrf_open/var/src/python/wrf/var/latlon.py b/wrf_open/var/src/python/wrf/var/latlon.py index c2932f4..912965e 100755 --- a/wrf_open/var/src/python/wrf/var/latlon.py +++ b/wrf_open/var/src/python/wrf/var/latlon.py @@ -1,43 +1,57 @@ from collections import Iterable -from wrf.var.constants import Constants -from wrf.var.extension import computeij, computell -from wrf.var.util import extract_vars, extract_global_attrs +import numpy as np + +from .config import xarray_enabled +from .constants import Constants +from .extension import computeij, computell +from .util import (extract_vars, extract_global_attrs, + either, _is_moving_domain, _is_multi_time, + iter_left_indexes) +from .decorators import set_latlon_metadata + +if xarray_enabled(): + from xarray import DataArray __all__ = ["get_lat", "get_lon", "get_ij", "get_ll"] -def get_lat(wrfnc, timeidx=0): - - try: - lat_vars = extract_vars(wrfnc, timeidx, varnames="XLAT") - except KeyError: - try: - latm_vars = extract_vars(wrfnc, timeidx, varnames="XLAT_M") - except: - raise RuntimeError("Latitude variable not found in NetCDF file") - else: - xlat = latm_vars["XLAT_M"] +def _lat_varname(stagger): + if stagger is None or stagger.lower() == "m": + varname = either("XLAT", "XLAT_M") + elif stagger.lower() == "u" or stagger.lower() == "v": + varname = "XLAT_{}".format(stagger.upper()) else: - xlat = lat_vars["XLAT"] + raise ValueError("invalid 'stagger' value") - return xlat - -def get_lon(wrfnc, timeidx=0): - try: - lon_vars = extract_vars(wrfnc, timeidx, varnames="XLONG") - except KeyError: - try: - lonm_vars = extract_vars(wrfnc, timeidx, varnames="XLONG_M") - except: - raise RuntimeError("Latitude variable not found in NetCDF file") - else: - xlon = lonm_vars["XLONG_M"] + return varname + +def _lon_varname(stagger): + if stagger is None or stagger.lower() == "m": + varname = either("XLONG", "XLONG_M") + elif stagger.lower() == "u" or stagger.lower() == "v": + varname = "XLONG_{}".format(stagger.upper()) else: - xlon = lon_vars["XLONG"] + raise ValueError("invalid 'stagger' value") + + return varname + +def get_lat(wrfnc, timeidx=0, stagger=None, + method="cat", squeeze=True, cache=None): + + varname = _lat_varname(stagger) + lat_var = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + + return lat_var[varname] - return xlon +def get_lon(wrfnc, timeidx=0, stagger=None, + method="cat", squeeze=True, cache=None): + + varname = _lon_varname(stagger) + lon_var = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + + return lon_var[varname] -def _get_proj_params(wrfnc, timeidx): +def _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache): if timeidx < 0: raise ValueError("'timeidx' must be greater than 0") @@ -64,12 +78,22 @@ def _get_proj_params(wrfnc, timeidx): latinc = 0.0 loninc = 0.0 - xlat = get_lat(wrfnc, timeidx) - xlon = get_lon(wrfnc, timeidx) + latvar = _lat_varname(stagger) + lonvar = _lon_varname(stagger) - ref_lat = xlat[0,0] - ref_lon = xlon[0,0] + lat_timeidx = timeidx + # Only need all the lats/lons if it's a moving domain file/files + if _is_multi_time(timeidx): + if not _is_moving_domain(wrfnc, latvar=latvar, lonvar=lonvar): + lat_timeidx = 0 + + xlat = get_lat(wrfnc, lat_timeidx, stagger, method, squeeze, cache) + xlon = get_lon(wrfnc, lat_timeidx, stagger, method, squeeze, cache) + ref_lat = np.ravel(xlat[...,0,0]) + ref_lon = np.ravel(xlon[...,0,0]) + + # Note: fortran index known_i = 1.0 known_j = 1.0 @@ -78,20 +102,70 @@ def _get_proj_params(wrfnc, timeidx): 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") - +@set_latlon_metadata(ij=True) +def get_ij(wrfnc, latitude, longitude, timeidx=0, + stagger=None, method="cat", squeeze=True, cache=None): + (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) + if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str): + lats = np.asarray(latitude) + lons = np.asarray(longitude) + + if lats.ndim > 1: + lats = lats.ravel() + + if lons.ndim > 1: + lons = lons.ravel() + + if (lats.size != lons.size): + raise ValueError("'latitude' and 'longitude' " + "must be the same length") + + + if ref_lat.size == 1: + outdim = [lats.size, 2] + extra_dims = outdim[0] + else: + # Moving domain will have moving ref_lats/ref_lons + outdim = [lats.size, ref_lat.size, 2] + extra_dims = outdim[0:2] + + res = np.empty(outdim, np.float64) + + for left_idxs in iter_left_indexes(extra_dims): + left_and_slice_idxs = left_idxs + [slice(None, None, None)] + + if ref_lat.size == 1: + ref_lat_val = ref_lat[0] + ref_lon_val = ref_lon[0] + else: + ref_lat_val = ref_lat[left_idxs[-1]] + ref_lon_val = ref_lon[left_idxs[-1]] + + lat = lats[left_idxs[0]] + lon = lons[left_idxs[0]] + + ij = computeij(map_proj, truelat1, truelat2, stdlon, + ref_lat_val, ref_lon_val, pole_lat, pole_lon, + known_i, known_j, dx, latinc, loninc, + lat, lon) + + res[left_and_slice_idxs] = ij[:] + + else: + + res = 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 + return res + + +@set_latlon_metadata(ij=False) 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") @@ -100,9 +174,59 @@ def get_ll(wrfnc, i, j, timeidx=0): 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, - loninc,i,j) + if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str): + i_arr = np.asarray(i) + j_arr = np.asarray(j) + + if i_arr.ndim > 1: + i_arr = i_arr.ravel() + + if j_arr.ndim > 1: + j_arr = j_arr.ravel() + + if (i_arr.size != j_arr.size): + raise ValueError("'i' and 'j' " + "must be the same length") + + if ref_lat.size == 1: + outdim = [i_arr.size, 2] + extra_dims = outdim[0] + else: + # Moving domain will have moving ref_lats/ref_lons + outdim = [i_arr.size, ref_lat.size, 2] + extra_dims = outdim[0:2] + + res = np.empty(outdim, np.float64) + + for left_idxs in iter_left_indexes(extra_dims): + left_and_slice_idxs = left_idxs + [slice(None, None, None)] + + if ref_lat.size == 1: + ref_lat_val = ref_lat[0] + ref_lon_val = ref_lon[0] + else: + ref_lat_val = ref_lat[left_idxs[-1]] + ref_lon_val = ref_lon[left_idxs[-1]] + + i_val = i_arr[left_idxs[0]] + j_val = j_arr[left_idxs[0]] + + ll = computell(map_proj, truelat1, truelat2, + stdlon, ref_lat_val, ref_lon_val, + pole_lat, pole_lon, known_i, known_j, + dx, latinc, loninc, + i_val, j_val) + + res[left_and_slice_idxs] = ll[:] + + else: + res = computell(map_proj, truelat1, truelat2, + stdlon, ref_lat, ref_lon, + pole_lat, pole_lon, known_i, known_j, + dx, latinc, loninc, + i_val, j_val) + + return res \ No newline at end of file diff --git a/wrf_open/var/src/python/wrf/var/omega.py b/wrf_open/var/src/python/wrf/var/omega.py index a88c6d8..6bf3235 100755 --- a/wrf_open/var/src/python/wrf/var/omega.py +++ b/wrf_open/var/src/python/wrf/var/omega.py @@ -1,14 +1,18 @@ -from wrf.var.constants import Constants -from wrf.var.destag import destagger -from wrf.var.extension import computeomega,computetk -from wrf.var.util import extract_vars +from .constants import Constants +from .destag import destagger +from .extension import computeomega,computetk +from .util import extract_vars +from .decorators import copy_and_set_metadata __all__ = ["get_omega"] -def get_omega(wrfnc, timeidx=0): - ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "W", - "PB", "QVAPOR")) +@copy_and_set_metadata(copy_varname="T", name="omega", + description="omega", + units="Pa/s") +def get_omega(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): + varnames=("T", "P", "W", "PB", "QVAPOR") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) t = ncvars["T"] p = ncvars["P"] w = ncvars["W"] diff --git a/wrf_open/var/src/python/wrf/var/precip.py b/wrf_open/var/src/python/wrf/var/precip.py index 86756de..332ea2d 100755 --- a/wrf_open/var/src/python/wrf/var/precip.py +++ b/wrf_open/var/src/python/wrf/var/precip.py @@ -1,6 +1,6 @@ -from wrf.var.util import extract_vars +from .util import extract_vars __all__ = ["get_accum_precip", "get_precip_diff"] diff --git a/wrf_open/var/src/python/wrf/var/pressure.py b/wrf_open/var/src/python/wrf/var/pressure.py index 5482fa6..b528f77 100755 --- a/wrf_open/var/src/python/wrf/var/pressure.py +++ b/wrf_open/var/src/python/wrf/var/pressure.py @@ -1,30 +1,31 @@ -from wrf.var.decorators import convert_units -from wrf.var.util import extract_vars +from .decorators import convert_units, copy_and_set_metadata +from .util import extract_vars, either __all__ = ["get_pressure", "get_pressure_hpa"] +@copy_and_set_metadata(copy_varname=either("P", "PRES"), name="pressure", + description="pressure") @convert_units("pressure", "pa") -def get_pressure(wrfnc, timeidx=0, units="pa"): - - try: - p_vars = extract_vars(wrfnc, timeidx, varnames=("P", "PB")) - except KeyError: - try: - pres_vars = extract_vars(wrfnc, timeidx, varnames="PRES") - except: - raise RuntimeError("pressure variable not found in NetCDF file") - else: - pres = pres_vars["PRES"] - else: +def get_pressure(wrfnc, timeidx=0, units="pa", + method="cat", squeeze=True, cache=None): + + varname = either("P", "PRES")(wrfnc) + if varname == "P": + p_vars = extract_vars(wrfnc, timeidx, ("P", "PB"), + method, squeeze, cache) p = p_vars["P"] pb = p_vars["PB"] pres = p + pb + else: + pres = extract_vars(wrfnc, timeidx, "PRES", + method, squeeze, cache)["PRES"] return pres -def get_pressure_hpa(wrfnc, timeidx=0, units="hpa"): - return get_pressure(wrfnc, timeidx, units=units) +def get_pressure_hpa(wrfnc, timeidx=0, units="hpa", + method="cat", squeeze=True, cache=None): + return get_pressure(wrfnc, timeidx, units, method, squeeze, cache) diff --git a/wrf_open/var/src/python/wrf/var/projection.py b/wrf_open/var/src/python/wrf/var/projection.py index 2db9fcd..a16306d 100644 --- a/wrf_open/var/src/python/wrf/var/projection.py +++ b/wrf_open/var/src/python/wrf/var/projection.py @@ -1,8 +1,8 @@ import numpy as np import math -from wrf.var.config import basemap_enabled, cartopy_enabled, pyngl_enabled -from wrf.var.constants import Constants +from .config import basemap_enabled, cartopy_enabled, pyngl_enabled +from .constants import Constants if cartopy_enabled(): from cartopy import crs diff --git a/wrf_open/var/src/python/wrf/var/pw.py b/wrf_open/var/src/python/wrf/var/pw.py index 98f8585..667bd53 100755 --- a/wrf_open/var/src/python/wrf/var/pw.py +++ b/wrf_open/var/src/python/wrf/var/pw.py @@ -1,13 +1,17 @@ -from wrf.var.extension import computepw,computetv,computetk -from wrf.var.constants import Constants -from wrf.var.util import extract_vars +from .extension import computepw,computetv,computetk +from .constants import Constants +from .util import extract_vars +from .decorators import copy_and_set_metadata __all__ = ["get_pw"] -def get_pw(wrfnc, timeidx=0): - ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "PH", - "PHB", "QVAPOR")) +@copy_and_set_metadata(copy_varname="T", name="pw", + description="precipitable water", + units="kg m-2") +def get_pw(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): + varnames=("T", "P", "PB", "PH", "PHB", "QVAPOR") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) t = ncvars["T"] p = ncvars["P"] @@ -22,9 +26,9 @@ def get_pw(wrfnc, timeidx=0): full_t = t + Constants.T_BASE tk = computetk(full_p, full_t) - tv = computetv(tk,qv) + tv = computetv(tk, qv) - return computepw(full_p,tv,qv,ht) + return computepw(full_p, tv, qv, ht) diff --git a/wrf_open/var/src/python/wrf/var/rh.py b/wrf_open/var/src/python/wrf/var/rh.py index 222aaf9..65e245a 100755 --- a/wrf_open/var/src/python/wrf/var/rh.py +++ b/wrf_open/var/src/python/wrf/var/rh.py @@ -1,12 +1,17 @@ -from wrf.var.constants import Constants -from wrf.var.extension import computerh, computetk -from wrf.var.util import extract_vars +from .constants import Constants +from .extension import computerh, computetk +from .util import extract_vars +from .decorators import copy_and_set_metadata __all__ = ["get_rh", "get_rh_2m"] -def get_rh(wrfnc, timeidx=0): - ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR")) +@copy_and_set_metadata(copy_varname="T", name="rh", + description="relative humidity", + delete_attrs=("units",)) +def get_rh(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): + varnames=("T", "P", "PB", "QVAPOR") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) t = ncvars["T"] p = ncvars["P"] pb = ncvars["PB"] @@ -20,8 +25,12 @@ def get_rh(wrfnc, timeidx=0): return rh -def get_rh_2m(wrfnc, timeidx=0): - ncvars = extract_vars(wrfnc, timeidx, varnames=("T2", "PSFC", "Q2")) +@copy_and_set_metadata(copy_varname="T2", name="rh2", + description="2m relative humidity", + delete_attrs=("units",)) +def get_rh_2m(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): + varnames=("T2", "PSFC", "Q2") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) t2 = ncvars["T2"] psfc = ncvars["PSFC"] q2 = ncvars["Q2"] diff --git a/wrf_open/var/src/python/wrf/var/slp.py b/wrf_open/var/src/python/wrf/var/slp.py index b51595c..1171f73 100755 --- a/wrf_open/var/src/python/wrf/var/slp.py +++ b/wrf_open/var/src/python/wrf/var/slp.py @@ -1,15 +1,19 @@ -from wrf.var.extension import computeslp, computetk -from wrf.var.constants import Constants -from wrf.var.destag import destagger -from wrf.var.decorators import convert_units -from wrf.var.util import extract_vars +from .extension import computeslp, computetk +from .constants import Constants +from .destag import destagger +from .decorators import convert_units, copy_and_set_metadata +from .util import extract_vars __all__ = ["get_slp"] +@copy_and_set_metadata(copy_varname="T", name="slp", + remove_dims=("bottom_top",), + description="sea level pressure") @convert_units("pressure", "hpa") -def get_slp(wrfnc, timeidx=0, units="hpa"): - ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR", - "PH", "PHB")) +def get_slp(wrfnc, timeidx=0, units="hpa", + method="cat", squeeze=True, cache=None): + varnames=("T", "P", "PB", "QVAPOR", "PH", "PHB") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) t = ncvars["T"] p = ncvars["P"] diff --git a/wrf_open/var/src/python/wrf/var/temp.py b/wrf_open/var/src/python/wrf/var/temp.py index 9a53a7e..390a3cf 100755 --- a/wrf_open/var/src/python/wrf/var/temp.py +++ b/wrf_open/var/src/python/wrf/var/temp.py @@ -1,26 +1,34 @@ -from wrf.var.constants import Constants -from wrf.var.extension import computetk, computeeth, computetv, computewetbulb -from wrf.var.decorators import convert_units, set_metadata -from wrf.var.util import extract_vars +from .constants import Constants +from .extension import computetk, computeeth, computetv, computewetbulb +from .decorators import convert_units, copy_and_set_metadata +from .util import extract_vars __all__ = ["get_theta", "get_temp", "get_eth", "get_tv", "get_tw", "get_tk", "get_tc"] -@set_metadata() +@copy_and_set_metadata(copy_varname="T", name="theta", + description="potential temperature") @convert_units("temp", "k") -def get_theta(wrfnc, timeidx=0, units="k"): - ncvars = extract_vars(wrfnc, timeidx, varnames="T") +def get_theta(wrfnc, timeidx=0, units="k", + method="cat", squeeze=True, cache=None): + varnames = ("T",) + + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) t = ncvars["T"] full_t = t + Constants.T_BASE return full_t +@copy_and_set_metadata(copy_varname="T", name="temp", + description="temperature") @convert_units("temp", "k") -def get_temp(wrfnc, timeidx=0, units="k"): +def get_temp(wrfnc, timeidx=0, units="k", method="cat", squeeze=True, + cache=None): """Return the temperature in Kelvin or Celsius""" - ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB")) + varnames=("T", "P", "PB") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) t = ncvars["T"] p = ncvars["P"] pb = ncvars["PB"] @@ -31,11 +39,15 @@ def get_temp(wrfnc, timeidx=0, units="k"): return tk +@copy_and_set_metadata(copy_varname="T", name="theta_e", + description="equivalent potential temperature") @convert_units("temp", "k") -def get_eth(wrfnc, timeidx=0, units="k"): +def get_eth(wrfnc, timeidx=0, units="k", method="cat", squeeze=True, + cache=None): "Return equivalent potential temperature (Theta-e) in Kelvin" - ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR")) + varnames=("T", "P", "PB", "QVAPOR") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) t = ncvars["T"] p = ncvars["P"] pb = ncvars["PB"] @@ -48,12 +60,16 @@ def get_eth(wrfnc, timeidx=0, units="k"): eth = computeeth(qv, tk, full_p) return eth - + +@copy_and_set_metadata(copy_varname="T", name="tv", + description="virtual temperature") @convert_units("temp", "k") -def get_tv(wrfnc, timeidx=0, units="k"): +def get_tv(wrfnc, timeidx=0, units="k", method="cat", squeeze=True, + cache=None): "Return the virtual temperature (tv) in Kelvin or Celsius" - ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR")) + varnames=("T", "P", "PB", "QVAPOR") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) t = ncvars["T"] p = ncvars["P"] @@ -68,12 +84,15 @@ def get_tv(wrfnc, timeidx=0, units="k"): return tv - +@copy_and_set_metadata(copy_varname="T", name="twb", + description="wetbulb temperature") @convert_units("temp", "k") -def get_tw(wrfnc, timeidx=0, units="k"): +def get_tw(wrfnc, timeidx=0, units="k", method="cat", squeeze=True, + cache=None): "Return the wetbulb temperature (tw)" - ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR")) + varnames=("T", "P", "PB", "QVAPOR") + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) t = ncvars["T"] p = ncvars["P"] pb = ncvars["PB"] @@ -87,11 +106,13 @@ def get_tw(wrfnc, timeidx=0, units="k"): return tw -def get_tk(wrfnc, timeidx=0): - return get_temp(wrfnc, timeidx, units="k") +def get_tk(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): + return get_temp(wrfnc, timeidx, units="k", + method=method, squeeze=squeeze, cache=cache) -def get_tc(wrfnc, timeidx=0): - return get_temp(wrfnc, timeidx, units="c") +def get_tc(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): + return get_temp(wrfnc, timeidx, units="c", + method=method, squeeze=squeeze, cache=cache) diff --git a/wrf_open/var/src/python/wrf/var/terrain.py b/wrf_open/var/src/python/wrf/var/terrain.py index 3ced26d..5942d34 100755 --- a/wrf_open/var/src/python/wrf/var/terrain.py +++ b/wrf_open/var/src/python/wrf/var/terrain.py @@ -1,24 +1,18 @@ -from wrf.var.decorators import convert_units -from wrf.var.util import extract_vars +from .decorators import convert_units, copy_and_set_metadata +from .util import extract_vars, either __all__ = ["get_terrain"] +# Need to handle either +@copy_and_set_metadata(copy_varname=either("HGT", "HGT_M"), name="terrain", + description="terrain height") @convert_units("height", "m") -def get_terrain(wrfnc, timeidx=0, units="m"): - - try: - hgt_vars = extract_vars(wrfnc, timeidx, varnames="HGT") - except KeyError: - try: - hgt_m_vars = extract_vars(wrfnc, timeidx, varnames="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 +def get_terrain(wrfnc, timeidx=0, units="m", method="cat", squeeze=True, + cache=None): + varname = either("HGT", "HGT_M")(wrfnc) + return extract_vars(wrfnc, timeidx, varname, + method, squeeze, cache)[varname] + \ No newline at end of file diff --git a/wrf_open/var/src/python/wrf/var/times.py b/wrf_open/var/src/python/wrf/var/times.py index e8b3174..7393845 100755 --- a/wrf_open/var/src/python/wrf/var/times.py +++ b/wrf_open/var/src/python/wrf/var/times.py @@ -1,5 +1,5 @@ -from wrf.var.util import extract_times +from .util import extract_times __all__ = ["get_times"] diff --git a/wrf_open/var/src/python/wrf/var/units.py b/wrf_open/var/src/python/wrf/var/units.py index 4399723..e8fe358 100755 --- a/wrf_open/var/src/python/wrf/var/units.py +++ b/wrf_open/var/src/python/wrf/var/units.py @@ -1,5 +1,5 @@ -from wrf.var.constants import Constants, ConversionFactors +from .constants import Constants, ConversionFactors __all__ = ["check_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 index 330f791..ef26359 100644 --- a/wrf_open/var/src/python/wrf/var/util.py +++ b/wrf_open/var/src/python/wrf/var/util.py @@ -1,12 +1,11 @@ from collections import Iterable, Mapping, OrderedDict from itertools import product import datetime as dt -import warnings -import numpy as n +import numpy as np -from wrf.var.config import xarray_enabled -from wrf.var.projection import getproj +from .config import xarray_enabled +from .projection import getproj if xarray_enabled(): from xarray import DataArray @@ -14,7 +13,8 @@ if xarray_enabled(): __all__ = ["extract_vars", "extract_global_attrs", "extract_dim", "combine_files", "is_standard_wrf_var", "extract_times", "iter_left_indexes", "get_left_indexes", "get_right_slices", - "is_staggered", "get_proj_params"] + "is_staggered", "get_proj_params", "viewitems", "viewkeys", + "viewvalues"] def _is_multi_time(timeidx): if timeidx == -1: @@ -31,23 +31,77 @@ def _is_mapping(wrfnc): return True return False -def _is_moving_domain(wrfnc, latvar="XLAT", lonvar="XLONG"): - lat1 = wrfnc.variables[latvar][0,:] - lat2 = wrfnc.variables[latvar][-1,:] - lon1 = wrfnc.variables[lonvar][0,:] - lon2 = wrfnc.variables[lonvar][-1,:] +def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar): + lats = wrfnc.variables[latvar] + lons = wrfnc.variables[lonvar] + # Need to check all times + for i in xrange(lats.shape[-3]): + start_idxs = [0]*lats.ndim + start_idxs[-3] = i + + end_idxs = [-1]*lats.ndim + end_idxs[-3] = i + + if (first_ll_corner[0] != lats[start_idxs] or + first_ll_corner[1] != lons[start_idxs] or + first_ur_corner[0] != lats[end_idxs] or + first_ur_corner[1] != lons[end_idxs]): + return True - if (lat1[0,0] != lat2[0,0] or lat1[-1,-1] != lat2[-1,-1] or - lon1[0,0] != lon2[0,0] or lon1[-1,-1] != lon2[-1,-1]): - return True + return False + +def _is_moving_domain(wrfseq, varname=None, latvar="XLAT", lonvar="XLONG"): + # In case it's just a single file + if not _is_multi_file(wrfseq): + wrfseq = [wrfseq] + + # Slow, but safe. Compare the corner points to the first item and see + # any move. User iterator protocol in case wrfseq is not a list/tuple. + wrf_iter = iter(wrfseq) + + first_wrfnc = next(wrf_iter) + + if varname is not None: + coord_names = getattr(first_wrfnc.variables[varname], + "coordinates").split() + lon_coord = coord_names[0] + lat_coord = coord_names[1] + else: + lon_coord = lonvar + lat_coord = latvar + + lats = first_wrfnc.variables[lat_coord] + lons = first_wrfnc.variables[lon_coord] + + zero_idxs = [0] * first_wrfnc.variables[lat_coord].ndim + last_idxs = list(zero_idxs) + last_idxs[-2:] = [-1]*2 - return False + lat0 = lats[zero_idxs] + lat1 = lats[last_idxs] + lon0 = lons[zero_idxs] + lon1 = lons[last_idxs] + + ll_corner = (lat0, lon0) + ur_corner = (lat1, lon1) + + while True: + try: + wrfnc = next(wrf_iter) + except StopIteration: + break + else: + if _corners_moved(wrfnc, ll_corner, ur_corner, + lat_coord, lon_coord): + return True + + return False -def _get_attr(wrfnc, attr): +def _get_global_attr(wrfnc, attr): val = getattr(wrfnc, attr, None) # PyNIO puts single values in to an array - if isinstance(val,n.ndarray): + if isinstance(val,np.ndarray): if len(val) == 1: return val[0] return val @@ -66,7 +120,7 @@ def extract_global_attrs(wrfnc, attrs): else: wrfnc = wrfnc[next(wrfnc.iterkeys())] - return {attr:_get_attr(wrfnc, attr) for attr in attrlist} + return {attr:_get_global_attr(wrfnc, attr) for attr in attrlist} def extract_dim(wrfnc, dim): if _is_multi_file(wrfnc): @@ -80,71 +134,81 @@ def extract_dim(wrfnc, dim): return len(d) #netCDF4 return d # PyNIO - -# TODO -def _combine_dict(wrfseq, var, timeidx, method): +def _combine_dict(wrfdict, varname, timeidx, method): """Dictionary combination creates a new left index for each key, then does a cat or join for the list of files for that key""" + keynames = [] + numkeys = len(wrfdict) - multitime = _is_multi_time(timeidx) - numfiles = len(wrfseq) + key_iter = iter(viewkeys(wrfdict)) + first_key = next(key_iter) + keynames.append(first_key) - if not multitime: - time_idx_or_slice = timeidx - else: - time_idx_or_slice = slice(None, None, None) - - keys = (list(x for x in xrange(numfiles)) if not _is_mapping(wrfseq) else - list(key for key in wrfseq.iterkeys())) + first_array = _extract_var(wrfdict[first_key], varname, + timeidx, method, squeeze=False) - # Check if - - if not xarray_enabled(): - first_var = wrfseq[keys[0]].variables[var][time_idx_or_slice, :] - else: - first_var = _build_data_array(wrfseq[keys[0]], var, timeidx) # Create the output data numpy array based on the first array - outdims = [numfiles] - outdims += first_var.shape - outdata = n.zeros(outdims, first_var.dtype) - outdata[0,:] = first_var[:] - - for idx, key in enumerate(keys[1:], start=1): - outdata[idx,:] = wrfseq[key].variables[var][time_idx_or_slice, :] + outdims = [numkeys] + outdims += first_array.shape + outdata = np.empty(outdims, first_array.dtype) + outdata[0,:] = first_array[:] + + idx = 1 + while True: + try: + key = next(key_iter) + except StopIteration: + break + else: + keynames.append(key) + vardata = _extract_var(wrfdict[key], varname, timeidx, + method, squeeze=False) + + if outdata.shape[1:] != vardata.shape: + raise ValueError("data sequences must have the " + "same size for all dictionary keys") + outdata[idx,:] = vardata.values[:] + idx += 1 if not xarray_enabled(): outarr = outdata else: - outname = str(first_var.name) - outcoords = dict(first_var.coords) - outdims = ["sequence"] + list(first_var.dims) - outcoords["sequence"] = keys - outattrs = dict(first_var.attrs) + outname = str(first_array.name) + # Note: assumes that all entries in dict have same coords + outcoords = OrderedDict(first_array.coords) + outdims = ["key"] + list(first_array.dims) + outcoords["key"] = keynames + outattrs = OrderedDict(first_array.attrs) outarr = DataArray(outdata, name=outname, coords=outcoords, dims=outdims, attrs=outattrs) return outarr - -def _cat_files(wrfseq, var, timeidx): +# TODO: implement in C +def _cat_files(wrfseq, varname, timeidx): + is_moving = _is_moving_domain(wrfseq, varname) + file_times = extract_times(wrfseq, timeidx) multitime = _is_multi_time(timeidx) time_idx_or_slice = timeidx if not multitime else slice(None, None, None) - first_var = (_build_data_array(wrfseq[0], var, timeidx) + # wrfseq might be a generator + wrf_iter = iter(wrfseq) + + first_var = (_build_data_array(next(wrf_iter), varname, timeidx, is_moving) if xarray_enabled() else - wrfseq[0].variables[var][time_idx_or_slice, :]) + wrfseq[0].variables[varname][time_idx_or_slice, :]) outdims = [len(file_times)] # Making a new time dim, so ignore this one outdims += first_var.shape[1:] - outdata = n.zeros(outdims, first_var.dtype) + outdata = np.empty(outdims, first_var.dtype) numtimes = first_var.shape[0] startidx = 0 @@ -153,25 +217,31 @@ def _cat_files(wrfseq, var, timeidx): outdata[startidx:endidx, :] = first_var[:] startidx = endidx - for wrfnc in wrfseq[1:]: - vardata = wrfnc.variables[var][time_idx_or_slice, :] - if multitime: - numtimes = vardata.shape[0] + while True: + try: + wrfnc = next(wrf_iter) + except StopIteration: + break else: - numtimes = 1 + vardata = wrfnc.variables[varname][time_idx_or_slice, :] - endidx = startidx + numtimes - - outdata[startidx:endidx, :] = vardata[:] - - startidx = endidx - + if multitime: + numtimes = vardata.shape[0] + else: + numtimes = 1 + + endidx = startidx + numtimes + + outdata[startidx:endidx, :] = vardata[:] + + startidx = endidx + if xarray_enabled(): # FIXME: If it's a moving nest, then the coord arrays need to have same # time indexes as the whole data set outname = str(first_var.name) - outattrs = dict(first_var.attrs) - outcoords = dict(first_var.coords) + outattrs = OrderedDict(first_var.attrs) + outcoords = OrderedDict(first_var.coords) outdimnames = list(first_var.dims) outcoords[outdimnames[0]] = file_times # New time dimension values @@ -183,8 +253,9 @@ def _cat_files(wrfseq, var, timeidx): return outarr -def _join_files(wrfseq, var, timeidx): - +# TODO: implement in C +def _join_files(wrfseq, varname, timeidx): + is_moving = _is_moving_domain(wrfseq, varname) multitime = _is_multi_time(timeidx) numfiles = len(wrfseq) @@ -193,26 +264,36 @@ def _join_files(wrfseq, var, timeidx): else: time_idx_or_slice = slice(None, None, None) + # wrfseq might be a generator + wrf_iter = iter(wrfseq) + if xarray_enabled(): - first_var = _build_data_array(wrfseq[0], var, timeidx) + first_var = _build_data_array(next(wrf_iter), varname, + timeidx, is_moving) else: - first_var = wrfseq[0].variables[var][time_idx_or_slice, :] + first_var = (next(wrf_iter)).variables[varname][time_idx_or_slice, :] # Create the output data numpy array based on the first array outdims = [numfiles] outdims += first_var.shape - outdata = n.zeros(outdims, first_var.dtype) + outdata = np.empty(outdims, first_var.dtype) outdata[0,:] = first_var[:] - - for idx, wrfnc in enumerate(wrfseq[1:], 1): - print idx - outdata[idx,:] = wrfnc.variables[var][time_idx_or_slice, :] + + idx=1 + while True: + try: + wrfnc = next(wrf_iter) + except StopIteration: + break + else: + outdata[idx,:] = wrfnc.variables[varname][time_idx_or_slice, :] + idx += 1 if xarray_enabled(): outname = str(first_var.name) - outcoords = dict(first_var.coords) - outattrs = dict(first_var.attrs) + outcoords = OrderedDict(first_var.coords) + outattrs = OrderedDict(first_var.attrs) # New dimensions outdimnames = ["file_idx"] + list(first_var.dims) outcoords["file_idx"] = [i for i in xrange(numfiles)] @@ -225,104 +306,134 @@ def _join_files(wrfseq, var, timeidx): return outarr -def combine_files(wrfseq, var, timeidx, method="cat", squeeze=True): +def combine_files(wrfseq, varname, timeidx, method="cat", squeeze=True): # Dictionary is unique if _is_mapping(wrfseq): - outarr = _combine_dict(wrfseq, var, timeidx, method) - - if method.lower() == "cat": - outarr = _cat_files(wrfseq, var, timeidx) + outarr = _combine_dict(wrfseq, varname, timeidx, method) + elif method.lower() == "cat": + outarr = _cat_files(wrfseq, varname, timeidx) elif method.lower() == "join": - outarr = _join_files(wrfseq, var, timeidx) + outarr = _join_files(wrfseq, varname, timeidx) else: raise ValueError("method must be 'cat' or 'join'") - if squeeze: - return outarr.squeeze() - - return outarr + return outarr.squeeze() if squeeze else outarr # Note, always returns the full data set with the time dimension included -def _build_data_array(wrfnc, varname, timeidx): +def _build_data_array(wrfnc, varname, timeidx, is_moving_domain): multitime = _is_multi_time(timeidx) var = wrfnc.variables[varname] data = var[:] attrs = OrderedDict(var.__dict__) dimnames = var.dimensions - # Add the coordinate variables here. - coord_names = getattr(var, "coordinates").split() - lon_coord = coord_names[0] - lat_coord = coord_names[1] + # WRF variables will have a coordinates attribute. MET_EM files have + # a stagger attribute which indicates the coordinate variable. + try: + # WRF files + coord_attr = getattr(var, "coordinates") + except KeyError: + try: + # met_em files + stag_attr = getattr(var, "stagger") + except KeyError: + lon_coord = None + lat_coord = None + else: + # For met_em files, use the stagger name to get the lat/lon var + lat_coord = "XLAT_{}".format(stag_attr) + lon_coord = "XLONG_{}".format(stag_attr) + else: + coord_names = coord_attr.split() + lon_coord = coord_names[0] + lat_coord = coord_names[1] + + coords = OrderedDict() - coords = {} - lon_coord_var = wrfnc.variables[lon_coord] - lat_coord_var = wrfnc.variables[lat_coord] + # Handle lat/lon coordinates and projection information if available + if lon_coord is not None and lat_coord is not None: + lon_coord_var = wrfnc.variables[lon_coord] + lat_coord_var = wrfnc.variables[lat_coord] - if multitime: - if _is_moving_domain(wrfnc, lat_coord, lon_coord): - # Special case with a moving domain in a multi-time file, - # otherwise the projection parameters don't change - coords[lon_coord] = lon_coord_var.dimensions, lon_coord_var[:] - coords[lat_coord] = lat_coord_var.dimensions, lat_coord_var[:] - - # Returned lats/lons arrays will have a time dimension, so proj - # will need to be a list due to moving corner points - lats, lons, proj_params = get_proj_params(wrfnc, - timeidx, - varname) - proj = [getproj(lats=lats[i,:], - lons=lons[i,:], - **proj_params) for i in xrange(lats.shape[0])] + if multitime: + if is_moving_domain: + # Special case with a moving domain in a multi-time file, + # otherwise the projection parameters don't change + coords[lon_coord] = lon_coord_var.dimensions, lon_coord_var[:] + coords[lat_coord] = lat_coord_var.dimensions, lat_coord_var[:] + + # Returned lats/lons arrays will have a time dimension, so proj + # will need to be a list due to moving corner points + lats, lons, proj_params = get_proj_params(wrfnc, + timeidx, + varname) + proj = [getproj(lats=lats[i,:], + lons=lons[i,:], + **proj_params) for i in xrange(lats.shape[0])] + else: + coords[lon_coord] = (lon_coord_var.dimensions[1:], + lon_coord_var[0,:]) + coords[lat_coord] = (lat_coord_var.dimensions[1:], + lat_coord_var[0,:]) + + # Domain not moving, so just get the first time + lats, lons, proj_params = get_proj_params(wrfnc, 0, varname) + proj = getproj(lats=lats, lons=lons, **proj_params) else: coords[lon_coord] = (lon_coord_var.dimensions[1:], - lon_coord_var[0,:]) + lon_coord_var[timeidx,:]) coords[lat_coord] = (lat_coord_var.dimensions[1:], - lat_coord_var[0,:]) - - # Domain not moving, so just get the first time + lat_coord_var[timeidx,:]) lats, lons, proj_params = get_proj_params(wrfnc, 0, varname) proj = getproj(lats=lats, lons=lons, **proj_params) - else: - coords[lon_coord] = (lon_coord_var.dimensions[1:], - lon_coord_var[timeidx,:]) - coords[lat_coord] = (lat_coord_var.dimensions[1:], - lat_coord_var[timeidx,:]) - lats, lons, proj_params = get_proj_params(wrfnc, 0, varname) - proj = getproj(lats=lats, lons=lons, **proj_params) + attrs["projection"] = proj - coords[dimnames[0]] = extract_times(wrfnc, timeidx) - attrs["projection"] = proj + if dimnames[0] == "Time": + coords[dimnames[0]] = extract_times(wrfnc, timeidx) data_array = DataArray(data, name=varname, dims=dimnames, coords=coords, attrs=attrs) return data_array -def _extract_var(wrfnc, varname, timeidx, method, squeeze): +# Cache is a dictionary of already extracted variables +def _extract_var(wrfnc, varname, timeidx, method, squeeze, cache): + # Mainly used internally so variables don't get extracted multiple times, + # particularly to copy metadata. This can be slow. + if cache is not None: + try: + return cache[varname] + except KeyError: + pass + + is_moving = _is_moving_domain(wrfnc, varname) multitime = _is_multi_time(timeidx) multifile = _is_multi_file(wrfnc) if not multifile: if xarray_enabled(): - return _build_data_array(wrfnc, varname, timeidx) + result = _build_data_array(wrfnc, varname, timeidx, is_moving) else: if not multitime: - return wrfnc.variables[varname][timeidx,:] + result = wrfnc.variables[varname][timeidx,:] else: - return wrfnc.variables[varname][:] + result = wrfnc.variables[varname][:] else: - return combine_files(wrfnc, varname, timeidx, method) + # Squeeze handled in this routine, so just return it + return combine_files(wrfnc, varname, timeidx, method, squeeze) + + return result.squeeze() if squeeze else result -def extract_vars(wrfnc, timeidx, varnames, method="cat", squeeze=True): +def extract_vars(wrfnc, timeidx, varnames, method="cat", squeeze=True, + cache=None): if isinstance(varnames, str): varlist = [varnames] else: varlist = varnames - return {var:_extract_var(wrfnc, var, timeidx, method, squeeze) + return {var:_extract_var(wrfnc, var, timeidx, method, squeeze, cache) for var in varlist} def _make_time(timearr): @@ -429,6 +540,78 @@ def get_proj_params(wrfnc, timeidx=0, varname=None): return (wrfnc.variables[lat_coord][time_idx_or_slice,:], wrfnc.variables[lon_coord][time_idx_or_slice,:], proj_params) + +# Dictionary python 2-3 compatibility stuff +def viewitems(d): + func = getattr(d, "viewitems", None) + if func is None: + func = d.items + return func() + + +def viewkeys(d): + func = getattr(d, "viewkeys", None) + if func is None: + func = d.keys + return func() + + +def viewvalues(d): + func = getattr(d, "viewvalues", None) + if func is None: + func = d.values + return func() + +# Helper utilities for metadata +class either(object): + def __init__(self, *varnames): + self.varnames = varnames + + def __call__(self, wrfnc): + if _is_multi_file(wrfnc): + wrfnc = next(iter(wrfnc)) + + for varname in self.varnames: + if varname in wrfnc: + return varname + + raise ValueError("{} are not valid variable names".format( + self.varnames)) + +class combine_with: + # Remove remove_idx first, then insert_idx is applied to removed set + def __init__(self, varname, remove_dims=None, insert_before=None, + new_dimnames=None, new_coords=None): + self.varname = varname + self.remove_dims = remove_dims + self.insert_before = insert_before + self.new_dimnames = new_dimnames + self.new_coords = new_coords + + def __call__(self, var): + new_dims = list(var.dims) + new_coords = OrderedDict(var.coords) + + if self.remove_dims is not None: + for dim in self.remove_dims: + new_dims.remove(dim) + del new_coords[dim] + + if self.insert_before is not None: + insert_idx = new_dims.index(self.insert_before) + new_dims = (new_dims[0:insert_idx] + self.new_dimnames + + new_dims[insert_idx:]) + elif self.new_dimnames is not None: + new_dims = self.new_dimnames + + if self.new_coords is not None: + new_coords.update(self.new_coords) + + return new_dims, new_coords + + + + diff --git a/wrf_open/var/src/python/wrf/var/uvmet.py b/wrf_open/var/src/python/wrf/var/uvmet.py index 2be79e7..c650bbe 100755 --- a/wrf_open/var/src/python/wrf/var/uvmet.py +++ b/wrf_open/var/src/python/wrf/var/uvmet.py @@ -1,72 +1,39 @@ from math import fabs, log, tan, sin, cos -from wrf.var.extension import computeuvmet -from wrf.var.destag 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 +from .extension import computeuvmet +from .destag import destagger +from .constants import Constants +from .wind import _calc_wspd_wdir +from .decorators import convert_units, set_wind_metadata +from .util import extract_vars, extract_global_attrs, either __all__=["get_uvmet", "get_uvmet10", "get_uvmet_wspd_wdir", "get_uvmet10_wspd_wdir"] +@set_wind_metadata(wspd_wdir=False) @convert_units("wind", "mps") -def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"): +def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps", + method="cat", squeeze=True, cache=None): """ Return a tuple of u,v with the winds rotated in to earth space""" if not ten_m: - try: - u_vars = extract_vars(wrfnc, timeidx, varnames="U") - except KeyError: - try: - uu_vars = extract_vars(wrfnc, timeidx, varnames="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) + varname = either("U", "UU")(wrfnc) + u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + u = destagger(u_vars[varname], -1) - try: - v_vars = extract_vars(wrfnc, timeidx, varnames="V") - except KeyError: - try: - vv_vars = extract_vars(wrfnc, timeidx, varnames="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) - + varname = either("V", "VV")(wrfnc) + v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + v = destagger(v_vars[varname], -2) else: - try: - u_vars = extract_vars(wrfnc, timeidx, varnames="U10") - except KeyError: - try: - uu_vars = extract_vars(wrfnc, timeidx, varnames="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(uu_vars["UU"][...,0,:,:], -1) # support met_em files - else: - u = u_vars["U10"] + varname = either("U10", "UU") + u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + u = (u_vars[varname] if varname == "U10" else + destagger(u_vars[varname][...,0,:,:], -1)) - try: - v_vars = extract_vars(wrfnc, timeidx, varnames="V10") - except KeyError: - try: - vv_vars = extract_vars(wrfnc, timeidx, varnames="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(vv_vars["VV"][...,0,:,:], -2) # support met_em files - else: - v = v_vars["V10"] + varname = either("V10", "VV") + v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + v = (v_vars[varname] if varname == "V10" else + destagger(v_vars[varname][...,0,:,:], -2)) map_proj_attrs = extract_global_attrs(wrfnc, attrs="MAP_PROJ") map_proj = map_proj_attrs["MAP_PROJ"] @@ -100,38 +67,26 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"): cen_lon = cen_lon_attrs["CEN_LON"] else: cen_lon = lon_attrs["STAND_LON"] - - try: - xlat_m_vars = extract_vars(wrfnc, timeidx, varnames="XLAT_M") - except KeyError: - try: - xlat_vars = extract_vars(wrfnc, timeidx, varnames="XLAT") - except KeyError: - raise RuntimeError("xlat not found in NetCDF file") - else: - lat = xlat_vars["XLAT"] - else: - lat = xlat_m_vars["XLAT_M"] - - try: - xlon_m_vars = extract_vars(wrfnc, timeidx, varnames="XLONG_M") - except KeyError: - try: - xlon_vars = extract_vars(wrfnc, timeidx, varnames="XLONG") - except KeyError: - raise RuntimeError("xlong not found in NetCDF file") - else: - lon = xlon_vars["XLONG"] - else: - lon = xlon_m_vars["XLONG_M"] - + + + varname = either("XLAT_M", "XLAT")(wrfnc) + xlat_var = extract_vars(wrfnc, timeidx, varname, + method, squeeze, cache) + lat = xlat_var[varname] + + varname = either("XLONG_M", "XLONG") + xlon_var = extract_vars(wrfnc, timeidx, varname, + method, squeeze, cache) + lon = xlon_var[varname] + if map_proj == 1: if((fabs(true_lat1 - true_lat2) > 0.1) and (fabs(true_lat2 - 90.) > 0.1)): cone = (log(cos(true_lat1*radians_per_degree)) - - log(cos(true_lat2*radians_per_degree))) - cone = cone / (log(tan((45.-fabs(true_lat1/2.))*radians_per_degree)) - - log(tan((45.-fabs(true_lat2/2.))*radians_per_degree))) + - log(cos(true_lat2*radians_per_degree))) + cone = (cone / + (log(tan((45.-fabs(true_lat1/2.))*radians_per_degree)) + - log(tan((45.-fabs(true_lat2/2.))*radians_per_degree)))) else: cone = sin(fabs(true_lat1)*radians_per_degree) else: @@ -141,17 +96,28 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"): return res - -def get_uvmet10(wrfnc, timeidx=0, units="mps"): + +def get_uvmet10(wrfnc, timeidx=0, units="mps", + method="cat", squeeze=True, cache=None): return get_uvmet(wrfnc, timeidx, True, units) -def get_uvmet_wspd_wdir(wrfnc, timeidx=0, units="mps"): - u,v = get_uvmet(wrfnc, timeidx, False, units) - return _calc_wspd_wdir(u, v, units) +@set_wind_metadata(wspd_wdir=True) +def get_uvmet_wspd_wdir(wrfnc, timeidx=0, units="mps", + method="cat", squeeze=True, cache=None): + + uvmet = get_uvmet(wrfnc, timeidx, False, units, method, squeeze, cache) + + return _calc_wspd_wdir(uvmet[0,:], uvmet[1,:], False, units) + -def get_uvmet10_wspd_wdir(wrfnc, timeidx=0, units="mps"): - u,v = get_uvmet10(wrfnc, timeidx, units="mps") - return _calc_wspd_wdir(u, v, units) +@set_wind_metadata(wspd_wdir=True) +def get_uvmet10_wspd_wdir(wrfnc, timeidx=0, units="mps", + method="cat", squeeze=True, cache=None): + + uvmet10 = get_uvmet10(wrfnc, timeidx, units="mps", method, squeeze, cache) + + return _calc_wspd_wdir(uvmet10[0,:], uvmet10[1,:], True, units) + diff --git a/wrf_open/var/src/python/wrf/var/vorticity.py b/wrf_open/var/src/python/wrf/var/vorticity.py index 7a52fb1..e19845d 100755 --- a/wrf_open/var/src/python/wrf/var/vorticity.py +++ b/wrf_open/var/src/python/wrf/var/vorticity.py @@ -1,12 +1,17 @@ -from wrf.var.extension import computeavo, computepvo -from wrf.var.util import extract_vars, extract_global_attrs +from .extension import computeavo, computepvo +from .util import extract_vars, extract_global_attrs +from .decorators import copy_and_set_metadata __all__ = ["get_avo", "get_pvo"] -def get_avo(wrfnc, timeidx=0): +@copy_and_set_metadata(copy_varname="F", name="avo", + description="absolute vorticity", + units="10-5 s-1") +def get_avo(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): ncvars = extract_vars(wrfnc, timeidx, varnames=("U", "V", "MAPFAC_U", "MAPFAC_V", "MAPFAC_M", - "F")) + "F"), + method, squeeze, cache) attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY")) u = ncvars["U"] @@ -22,11 +27,15 @@ def get_avo(wrfnc, timeidx=0): return computeavo(u,v,msfu,msfv,msfm,cor,dx,dy) -def get_pvo(wrfnc, timeidx=0): +@copy_and_set_metadata(copy_varname="T", name="pvo", + description="potential vorticity", + units="PVU") +def get_pvo(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): ncvars = extract_vars(wrfnc, timeidx, varnames=("U", "V", "T", "P", "PB", "MAPFAC_U", "MAPFAC_V", "MAPFAC_M", - "F")) + "F"), + method, squeeze, cache) attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY")) u = ncvars["U"] diff --git a/wrf_open/var/src/python/wrf/var/wind.py b/wrf_open/var/src/python/wrf/var/wind.py index 08a8f46..384a68e 100755 --- a/wrf_open/var/src/python/wrf/var/wind.py +++ b/wrf_open/var/src/python/wrf/var/wind.py @@ -1,42 +1,88 @@ -import numpy as n +import numpy as np -from wrf.var.constants import Constants -from wrf.var.destag import destagger_windcomp -from wrf.var.decorators import convert_units +from .constants import Constants +from .destag import destagger +from .util import extract_vars, either +from .decorators import convert_units, set_wind_metadata __all__ = ["get_u_destag", "get_v_destag", "get_w_destag", - "get_destag_wspd_wdir"] + "get_destag_wspd_wdir", "get_destag_wspd_wdir10"] @convert_units("wind", "mps") def _calc_wspd(u, v, units="mps"): - return n.sqrt(u**2 + v**2) + return np.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) + wdir = 270.0 - np.arctan2(v,u) * (180.0/Constants.PI) + return np.remainder(wdir, 360.0) -def _calc_wspd_wdir(u, v, units): - return (_calc_wspd(u,v, units), _calc_wdir(u,v)) +def _calc_wspd_wdir(u, v, two_d, units): + wspd = _calc_wspd(u, v, units) + wdir = _calc_wdir(u, v) + + idx_end = -2 if two_d else -3 + + outdims = list(wspd.shape[0:idx_end]) + [2] + list(wspd.shape[idx_end:]) + res = np.zeros(outdims, wspd.dtype) + res[...,0,:] = wspd[:] + res[...,1,:] = wdir[:] + + return res +@set_wind_metadata(wspd_wdir=False) @convert_units("wind", "mps") -def get_u_destag(wrfnc, timeidx=0, units="mps"): - u = destagger_windcomp(wrfnc,"u", timeidx) +def get_u_destag(wrfnc, timeidx=0, units="mps", + method="cat", squeeze=True, cache=None): + varname = either("U", "UU")(wrfnc) + u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + u = destagger(u_vars[varname], -1) + return u +@set_wind_metadata(wspd_wdir=False) @convert_units("wind", "mps") -def get_v_destag(wrfnc, timeidx=0, units="mps"): - v = destagger_windcomp(wrfnc,"v", timeidx) +def get_v_destag(wrfnc, timeidx=0, units="mps", + method="cat", squeeze=True, cache=None): + varname = either("V", "VV")(wrfnc) + v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + v = destagger(v_vars[varname], -2) return v +@set_wind_metadata(wspd_wdir=False) @convert_units("wind", "mps") -def get_w_destag(wrfnc, timeidx=0, units="mps"): - w = destagger_windcomp(wrfnc,"w", timeidx) +def get_w_destag(wrfnc, timeidx=0, units="mps", + method="cat", squeeze=True, cache=None): + w_vars = extract_vars(wrfnc, timeidx, "W", method, squeeze, cache) + w = destagger(w_vars["W"], -3) return w -def get_destag_wspd_wdir(wrfnc, timeidx=0, units="mps"): - u = destagger_windcomp(wrfnc,"u", timeidx) - v = destagger_windcomp(wrfnc,"v", timeidx) +@set_wind_metadata(wspd_wdir=True) +def get_destag_wspd_wdir(wrfnc, timeidx=0, units="mps", + method="cat", squeeze=True, cache=None): + varname = either("U", "UU")(wrfnc) + u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + u = destagger(u_vars[varname], -1) + + varname = either("V", "VV")(wrfnc) + v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + v = destagger(v_vars[varname], -2) + + return _calc_wspd_wdir(u, v, False, units) + +@set_wind_metadata(wspd_wdir=True) +def get_destag_wspd_wdir10(wrfnc, timeidx=0, units="mps", + method="cat", squeeze=True, cache=None): + + varname = either("U10", "UU")(wrfnc) + u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + u = (u_vars[varname] if varname == "U10" else + destagger(u_vars[varname][...,0,:,:], -1)) + + varname = either("V10", "VV")(wrfnc) + v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + v = (v_vars[varname] if varname == "V10" else + destagger(v_vars[varname][...,0,:,:], -2)) - return _calc_wspd_wdir(u,v,units) + return _calc_wspd_wdir(u,v,True,units)