nilyin
/
wrf-python
forked from 3rdparty/wrf-python
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
Browse Source

Added metadata decorators (untested and buggy)

main
Bill Ladwig 9 years ago
parent
a887fa1a44
commit
f862a793e2
27 changed files with 1585 additions and 921 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 128
      wrf_open/var/src/python/wrf/var/__init__.py
  2. 44
      wrf_open/var/src/python/wrf/var/cape.py
  3. 23
      wrf_open/var/src/python/wrf/var/ctt.py
  4. 33
      wrf_open/var/src/python/wrf/var/dbz.py
  5. 435
      wrf_open/var/src/python/wrf/var/decorators.py
  6. 43
      wrf_open/var/src/python/wrf/var/destag.py
  7. 22
      wrf_open/var/src/python/wrf/var/dewpoint.py
  8. 109
      wrf_open/var/src/python/wrf/var/extension.py
  9. 41
      wrf_open/var/src/python/wrf/var/geoht.py
  10. 89
      wrf_open/var/src/python/wrf/var/helicity.py
  11. 44
      wrf_open/var/src/python/wrf/var/interp.py
  12. 216
      wrf_open/var/src/python/wrf/var/latlon.py
  13. 18
      wrf_open/var/src/python/wrf/var/omega.py
  14. 2
      wrf_open/var/src/python/wrf/var/precip.py
  15. 33
      wrf_open/var/src/python/wrf/var/pressure.py
  16. 4
      wrf_open/var/src/python/wrf/var/projection.py
  17. 20
      wrf_open/var/src/python/wrf/var/pw.py
  18. 23
      wrf_open/var/src/python/wrf/var/rh.py
  19. 20
      wrf_open/var/src/python/wrf/var/slp.py
  20. 61
      wrf_open/var/src/python/wrf/var/temp.py
  21. 26
      wrf_open/var/src/python/wrf/var/terrain.py
  22. 2
      wrf_open/var/src/python/wrf/var/times.py
  23. 2
      wrf_open/var/src/python/wrf/var/units.py
  24. 355
      wrf_open/var/src/python/wrf/var/util.py
  25. 140
      wrf_open/var/src/python/wrf/var/uvmet.py
  26. 21
      wrf_open/var/src/python/wrf/var/vorticity.py
  27. 86
      wrf_open/var/src/python/wrf/var/wind.py

128
wrf_open/var/src/python/wrf/var/__init__.py
Unescape View File

@ -1,63 +1,61 @@ @@ -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__) @@ -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, @@ -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): @@ -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:

44
wrf_open/var/src/python/wrf/var/cape.py
Unescape View File

@ -1,17 +1,26 @@ @@ -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): @@ -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"]

23
wrf_open/var/src/python/wrf/var/ctt.py
Unescape View File

@ -1,21 +1,24 @@ @@ -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"): @@ -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"): @@ -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:

33
wrf_open/var/src/python/wrf/var/dbz.py
Unescape View File

@ -1,12 +1,17 @@ @@ -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): @@ -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): @@ -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): @@ -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)

435
wrf_open/var/src/python/wrf/var/decorators.py
Unescape View File

@ -1,13 +1,16 @@ @@ -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,40 +18,60 @@ if xarray_enabled(): @@ -15,40 +18,60 @@ 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.
if isinstance(argnames, str):
arglist = [argnames]
else:
arglist = argnames
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)
if self.argname not in argspec.args and self.argname not in kargs:
print argspec
if argname not in argspec.args and argname not in kwargs:
return None
try:
result_idx = argspec.args.index(self.argname)
result_idx = argspec.args.index(argname)
except ValueError:
result_idx = None
if (self.argname in kargs):
result = kargs[self.argname]
result = None
if (argname in kwargs):
result = kwargs, argname
elif (len(args) > result_idx and result_idx is not None):
result = args[result_idx]
result = args, result_idx
else:
arg_idx_from_right = (len(argspec.args) -
argspec.args.index(self.argname))
result = argspec.defaults[-arg_idx_from_right]
argspec.args.index(argname))
result = list(argspec.defaults), -arg_idx_from_right
return result
@ -61,19 +84,19 @@ def convert_units(unit_type, alg_unit): @@ -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):
# def convert_decorator(func):
@wrapt.decorator
def func_wrapper(wrapped, instance, args, kwargs):
desired_units = from_args("units")(func, *args, **kargs)
desired_units = from_args(wrapped, "units", *args, **kwargs)["units"]
check_units(desired_units, unit_type)
# Unit conversion done here
return do_conversion(func(*args, **kargs), unit_type,
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,23 +128,23 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, @@ -105,23 +128,23 @@ 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):
# 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 = kargs[ref_var_name]
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 func(*args, **kargs)
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)]
@ -141,22 +164,22 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, @@ -141,22 +164,22 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1,
# 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()}
for key,val in viewitems(kwargs)}
# Call the numerical routine
res = func(*new_args, **new_kargs)
res = wrapped(*new_args, **new_kargs)
if isinstance(res, n.ndarray):
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 = n.zeros(outdims, ref_var.dtype)
output = np.empty(outdims, ref_var.dtype)
masked = False
else:
output = ma.MaskedArray(
n.zeros(outdims, ref_var.dtype),
mask=n.zeros(outdims, n.bool_),
np.zeros(outdims, ref_var.dtype),
mask=np.zeros(outdims, np.bool_),
fill_value=res.fill_value)
masked = True
@ -172,13 +195,13 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, @@ -172,13 +195,13 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1,
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)
output = [np.empty(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_),
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
@ -197,16 +220,16 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, @@ -197,16 +220,16 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1,
return func_wrapper
return indexing_decorator
# 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):
#def indexing_decorator(func):
@wrapt.decorator
def func_wrapper(wrapped, instance, args, kwargs):
u = args[0]
v = args[1]
lat = args[2]
@ -237,7 +260,7 @@ def uvmet_left_iter(): @@ -237,7 +260,7 @@ def uvmet_left_iter():
# No special left side iteration, return the function result
if (extra_dim_num == 0):
return func(u,v,lat,lon,cen_long,cone)
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)]
@ -248,7 +271,7 @@ def uvmet_left_iter(): @@ -248,7 +271,7 @@ def uvmet_left_iter():
outdims += [u.shape[x] for x in xrange(-num_right_dims,0,1)]
output = n.zeros(outdims, u.dtype)
output = np.empty(outdims, u.dtype)
for left_idxs in iter_left_indexes(extra_dims):
# Make the left indexes plus a single slice object
@ -264,12 +287,12 @@ def uvmet_left_iter(): @@ -264,12 +287,12 @@ def uvmet_left_iter():
new_lon = lon[left_and_slice_idxs]
# Call the numerical routine
res = func(new_u, new_v, new_lat, new_lon, cen_long, cone)
res = wrapped(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)
res = np.squeeze(res)
output[left_and_slice_idxs] = res[:]
@ -278,16 +301,16 @@ def uvmet_left_iter(): @@ -278,16 +301,16 @@ def uvmet_left_iter():
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):
# 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 ()
@ -299,11 +322,11 @@ def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=n.float64): @@ -299,11 +322,11 @@ def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=n.float64):
new_kargs = {key:(val.astype(dtype)
if key in karg_names else val)
for key,val in kargs.iteritems()}
for key,val in viewitems()}
res = func(*new_args, **new_kargs)
res = wrapped(*new_args, **new_kargs)
if isinstance(res, n.ndarray):
if isinstance(res, np.ndarray):
if res.dtype == orig_type:
return res
return res.astype(orig_type)
@ -314,14 +337,14 @@ def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=n.float64): @@ -314,14 +337,14 @@ def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=n.float64):
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,18 +355,18 @@ def handle_extract_transpose(): @@ -332,18 +355,18 @@ def handle_extract_transpose():
transposes if the data is not fortran contiguous.
"""
def extract_transpose_decorator(func):
@wraps(func)
def func_wrapper(*args, **kargs):
# 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 kargs.iteritems()}
for key,val in viewitems(kwargs)}
res = func(*new_args, **new_kargs)
res = wrapped(*new_args, **new_kargs)
if isinstance(res, n.ndarray):
if isinstance(res, np.ndarray):
if res.flags.F_CONTIGUOUS:
return res.T
else:
@ -353,32 +376,284 @@ def handle_extract_transpose(): @@ -353,32 +376,284 @@ def handle_extract_transpose():
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):
@wrapt.decorator
def func_wrapper(wrapped, instance, args, kwargs):
if not xarray_enabled():
return func(*args, **kargs)
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 = {}
result = func(*args, **kargs)
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)
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)
outattrs.update(var_to_copy.attrs)
if name is not None:
outname = name
units = from_args("units")(func, *args, **kargs)
if units is not None:
result.attrs["units"] = units
outattrs["units"] = units
for argname, val in fixed_kargs.iteritems():
result[argname] = val
for argname, val in viewitems(fixed_attrs):
outattrs[argname] = val
return result
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"])
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)

43
wrf_open/var/src/python/wrf/var/destag.py
Unescape View File

@ -1,6 +1,6 @@ @@ -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): @@ -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))

22
wrf_open/var/src/python/wrf/var/dewpoint.py
Unescape View File

@ -1,13 +1,17 @@ @@ -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"]
@ -20,9 +24,13 @@ def get_dp(wrfnc, timeidx=0, units="c"): @@ -20,9 +24,13 @@ 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"])

109
wrf_open/var/src/python/wrf/var/extension.py
Unescape View File

@ -1,8 +1,8 @@ @@ -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, @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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,

41
wrf_open/var/src/python/wrf/var/geoht.py
Unescape View File

@ -1,11 +1,12 @@ @@ -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): @@ -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): @@ -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)

89
wrf_open/var/src/python/wrf/var/helicity.py
Unescape View File

@ -1,43 +1,34 @@ @@ -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): @@ -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): @@ -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

44
wrf_open/var/src/python/wrf/var/interp.py
Unescape View File

@ -3,20 +3,20 @@ from math import floor, ceil @@ -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): @@ -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, @@ -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, @@ -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, @@ -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, @@ -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, @@ -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)

216
wrf_open/var/src/python/wrf/var/latlon.py
Unescape View File

@ -1,43 +1,57 @@ @@ -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
__all__ = ["get_lat", "get_lon", "get_ij", "get_ll"]
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
def get_lat(wrfnc, timeidx=0):
if xarray_enabled():
from xarray import DataArray
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"]
else:
xlat = lat_vars["XLAT"]
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")
__all__ = ["get_lat", "get_lon", "get_ij", "get_ll"]
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:
xlon = lonm_vars["XLONG_M"]
raise ValueError("invalid 'stagger' value")
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]
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 xlon
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): @@ -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): @@ -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): @@ -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

18
wrf_open/var/src/python/wrf/var/omega.py
Unescape View File

@ -1,14 +1,18 @@ @@ -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"]

2
wrf_open/var/src/python/wrf/var/precip.py
Unescape View File

@ -1,6 +1,6 @@ @@ -1,6 +1,6 @@
from wrf.var.util import extract_vars
from .util import extract_vars
__all__ = ["get_accum_precip", "get_precip_diff"]

33
wrf_open/var/src/python/wrf/var/pressure.py
Unescape View File

@ -1,30 +1,31 @@ @@ -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)

4
wrf_open/var/src/python/wrf/var/projection.py
Unescape View File

@ -1,8 +1,8 @@ @@ -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

20
wrf_open/var/src/python/wrf/var/pw.py
Unescape View File

@ -1,13 +1,17 @@ @@ -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): @@ -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)

23
wrf_open/var/src/python/wrf/var/rh.py
Unescape View File

@ -1,12 +1,17 @@ @@ -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): @@ -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"]

20
wrf_open/var/src/python/wrf/var/slp.py
Unescape View File

@ -1,15 +1,19 @@ @@ -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"]

61
wrf_open/var/src/python/wrf/var/temp.py
Unescape View File

@ -1,26 +1,34 @@ @@ -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"): @@ -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"]
@ -49,11 +61,15 @@ def get_eth(wrfnc, timeidx=0, units="k"): @@ -49,11 +61,15 @@ def get_eth(wrfnc, timeidx=0, units="k"):
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"): @@ -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"): @@ -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)

26
wrf_open/var/src/python/wrf/var/terrain.py
Unescape View File

@ -1,24 +1,18 @@ @@ -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"):
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]
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

2
wrf_open/var/src/python/wrf/var/times.py
Unescape View File

@ -1,5 +1,5 @@ @@ -1,5 +1,5 @@
from wrf.var.util import extract_times
from .util import extract_times
__all__ = ["get_times"]

2
wrf_open/var/src/python/wrf/var/units.py
Unescape View File

@ -1,5 +1,5 @@ @@ -1,5 +1,5 @@
from wrf.var.constants import Constants, ConversionFactors
from .constants import Constants, ConversionFactors
__all__ = ["check_units", "do_conversion"]

355
wrf_open/var/src/python/wrf/var/util.py
Unescape View File

@ -1,12 +1,11 @@ @@ -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(): @@ -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): @@ -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
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
lat0 = lats[zero_idxs]
lat1 = lats[last_idxs]
lon0 = lons[zero_idxs]
lon1 = lons[last_idxs]
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]):
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): @@ -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): @@ -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)
first_array = _extract_var(wrfdict[first_key], varname,
timeidx, method, squeeze=False)
keys = (list(x for x in xrange(numfiles)) if not _is_mapping(wrfseq) else
list(key for key in wrfseq.iterkeys()))
# 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[:]
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)
for idx, key in enumerate(keys[1:], start=1):
outdata[idx,:] = wrfseq[key].variables[var][time_idx_or_slice, :]
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
# TODO: implement in C
def _cat_files(wrfseq, varname, timeidx):
is_moving = _is_moving_domain(wrfseq, varname)
def _cat_files(wrfseq, var, timeidx):
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,8 +217,14 @@ def _cat_files(wrfseq, var, timeidx): @@ -153,8 +217,14 @@ 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, :]
while True:
try:
wrfnc = next(wrf_iter)
except StopIteration:
break
else:
vardata = wrfnc.variables[varname][time_idx_or_slice, :]
if multitime:
numtimes = vardata.shape[0]
else:
@ -170,8 +240,8 @@ def _cat_files(wrfseq, var, timeidx): @@ -170,8 +240,8 @@ def _cat_files(wrfseq, var, timeidx):
# 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): @@ -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): @@ -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,42 +306,57 @@ def _join_files(wrfseq, var, timeidx): @@ -225,42 +306,57 @@ 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()
# 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 = {}
coords = OrderedDict()
# 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):
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[:]
@ -291,38 +387,53 @@ def _build_data_array(wrfnc, varname, timeidx): @@ -291,38 +387,53 @@ def _build_data_array(wrfnc, varname, 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)
def extract_vars(wrfnc, timeidx, varnames, method="cat", squeeze=True):
return result.squeeze() if squeeze else result
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):
@ -430,6 +541,78 @@ def get_proj_params(wrfnc, timeidx=0, varname=None): @@ -430,6 +541,78 @@ def get_proj_params(wrfnc, timeidx=0, varname=None):
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

140
wrf_open/var/src/python/wrf/var/uvmet.py
Unescape View File

@ -1,72 +1,39 @@ @@ -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)
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("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)
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"]
@ -101,37 +68,25 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"): @@ -101,37 +68,25 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"):
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)))
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:
@ -142,16 +97,27 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"): @@ -142,16 +97,27 @@ 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)
@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)
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)

21
wrf_open/var/src/python/wrf/var/vorticity.py
Unescape View File

@ -1,12 +1,17 @@ @@ -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): @@ -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"]

86
wrf_open/var/src/python/wrf/var/wind.py
Unescape View File

@ -1,42 +1,88 @@ @@ -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)

Write Preview
Loading…
Cancel
Save