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
Unified 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 @@
import warnings import warnings
from config import * from . import config
import config from .config import *
from extension import * from . import extension
import extension from .extension import *
from util import * from . import util
import util from .util import *
from cape import * from . import cape
import cape from .cape import *
from constants import * from . import constants
import constants from .constants import *
from ctt import * from . import ctt
import ctt from .ctt import *
from dbz import * from . import dbz
import dbz from .dbz import *
from destag import * from . import destag
import destag from .destag import *
from dewpoint import * from . import dewpoint
import dewpoint from .dewpoint import *
from etaconv import * from . import geoht
import etaconv from .geoht import *
from geoht import * from . import helicity
import geoht from .helicity import *
from helicity import * from . import interp
import helicity from .interp import *
from interp import * from . import latlon
import interp from .latlon import *
from latlon import * from . import omega
import latlon from .omega import *
from omega import * from . import precip
import omega from .precip import *
from precip import * from . import pressure
import precip from .pressure import *
from pressure import * from . import psadlookup
import pressure from .psadlookup import *
from psadlookup import * from . import pw
import psadlookup from .pw import *
from pw import * from . import rh
import pw from .rh import *
from rh import * from . import slp
import rh from .slp import *
from slp import * from . import temp
import slp from .temp import *
from temp import * from . import terrain
import temp from .terrain import *
from terrain import * from . import uvmet
import terrain from .uvmet import *
from uvmet import * from . import vorticity
import uvmet from .vorticity import *
from vorticity import * from . import wind
import vorticity from .wind import *
from wind import * from . import times
import wind from .times import *
from times import * from . import units
import times from .units import *
from units import * from . import projection
import units from .projection import *
from projection import *
import projection
__all__ = ["getvar"] __all__ = ["getvar"]
__all__.extend(config.__all__) __all__.extend(config.__all__)
@ -69,7 +67,6 @@ __all__.extend(ctt.__all__)
__all__.extend(dbz.__all__) __all__.extend(dbz.__all__)
__all__.extend(destag.__all__) __all__.extend(destag.__all__)
__all__.extend(dewpoint.__all__) __all__.extend(dewpoint.__all__)
__all__.extend(etaconv.__all__)
__all__.extend(geoht.__all__) __all__.extend(geoht.__all__)
__all__.extend(helicity.__all__) __all__.extend(helicity.__all__)
__all__.extend(interp.__all__) __all__.extend(interp.__all__)
@ -127,9 +124,10 @@ _FUNC_MAP = {"cape2d" : get_2dcape,
"lon" : get_lon, "lon" : get_lon,
"pressure" : get_pressure_hpa, "pressure" : get_pressure_hpa,
"pres" : get_pressure, "pres" : get_pressure,
"wspddir" : get_destag_wspd_wdir, "wspd_wdir" : get_destag_wspd_wdir,
"wspddir_uvmet" : get_uvmet_wspd_wdir, "wspd_wdir10" : get_destag_wspd_wdir10,
"wspddir_uvmet10" : get_uvmet10_wspd_wdir, "wspd_wdir_uvmet" : get_uvmet_wspd_wdir,
"wspd_wdir_uvmet10" : get_uvmet10_wspd_wdir,
"ctt" : get_ctt "ctt" : get_ctt
} }
@ -214,9 +212,11 @@ def _check_kargs(var, kargs):
"argument for '%s'" % (arg, var)) "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): 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) actual_var = _undo_alias(var)
if actual_var not in _VALID_KARGS: if actual_var not in _VALID_KARGS:

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

@ -1,17 +1,26 @@
import numpy as n import numpy as n
import numpy.ma as ma import numpy.ma as ma
from wrf.var.extension import computetk,computecape from .extension import computetk,computecape
from wrf.var.destag import destagger from .destag import destagger
from wrf.var.constants import Constants, ConversionFactors from .constants import Constants, ConversionFactors
from wrf.var.util import extract_vars from .util import extract_vars, combine_with
from .decorators import copy_and_set_metadata
__all__ = ["get_2dcape", "get_3dcape"] __all__ = ["get_2dcape", "get_3dcape"]
@copy_and_set_metadata(copy_varname="T",
def get_2dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL): 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""" """Return the 2d fields of cape, cin, lcl, and lfc"""
ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR", "PH", varnames = ("T", "P", "PB", "QVAPOR", "PH","PHB", "HGT", "PSFC")
"PHB", "HGT", "PSFC")) ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
@ -51,20 +60,27 @@ def get_2dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL):
resdim = left_dims + [4] + right_dims resdim = left_dims + [4] + right_dims
# Make a new output array for the result # Make a new output array for the result
res = n.zeros((resdim), cape.dtype) res = n.zeros(resdim, cape.dtype)
res[...,0,:,:] = cape[:] res[...,0,:,:] = cape[:]
res[...,1,:,:] = cin[:] res[...,1,:,:] = cin[:]
res[...,2,:,:] = lcl[:] res[...,2,:,:] = lcl[:]
res[...,3,:,:] = lfc[:] 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""" """Return the 3d fields of cape and cin"""
varnames = ("T", "P", "PB", "QVAPOR", "PH", "PHB", "HGT", "PSFC")
ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR", ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
"PH", "PHB", "HGT", "PSFC"))
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
pb = ncvars["PB"] pb = ncvars["PB"]

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

@ -1,21 +1,24 @@
import numpy as n import numpy as n
from wrf.var.extension import computectt, computetk from .extension import computectt, computetk
from wrf.var.constants import Constants, ConversionFactors from .constants import Constants, ConversionFactors
from wrf.var.destag import destagger from .destag import destagger
from wrf.var.decorators import convert_units from .decorators import convert_units, copy_and_set_metadata
from wrf.var.util import extract_vars from .util import extract_vars
__all__ = ["get_ctt"] __all__ = ["get_ctt"]
@copy_and_set_metadata(copy_varname="T", name="ctt",
description="cloud top temperature")
@convert_units("temp", "c") @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. """Return the cloud top temperature.
""" """
ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "PH" , varnames = ("T", "P", "PB", "PH", "PHB", "HGT", "QVAPOR")
"PHB", "HGT", "QVAPOR")) ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
pb = ncvars["PB"] pb = ncvars["PB"]
@ -26,7 +29,7 @@ def get_ctt(wrfnc, timeidx=0, units="c"):
haveqci = 1 haveqci = 1
try: try:
icevars = extract_vars(wrfnc, timeidx, vars="QICE") icevars = extract_vars(wrfnc, timeidx, "QICE", method, squeeze, cache)
except KeyError: except KeyError:
qice = n.zeros(qv.shape, qv.dtype) qice = n.zeros(qv.shape, qv.dtype)
haveqci = 0 haveqci = 0
@ -34,7 +37,7 @@ def get_ctt(wrfnc, timeidx=0, units="c"):
qice = icevars["QICE"] * 1000.0 #g/kg qice = icevars["QICE"] * 1000.0 #g/kg
try: try:
cldvars = extract_vars(wrfnc, timeidx, vars="QCLOUD") cldvars = extract_vars(wrfnc, timeidx, "QCLOUD", method, squeeze, cache)
except KeyError: except KeyError:
raise RuntimeError("'QCLOUD' not found in NetCDF file") raise RuntimeError("'QCLOUD' not found in NetCDF file")
else: else:

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

@ -1,12 +1,17 @@
import numpy as n import numpy as n
from wrf.var.extension import computedbz,computetk from .extension import computedbz,computetk
from wrf.var.constants import Constants from .constants import Constants
from wrf.var.util import extract_vars from .util import extract_vars, combine_with
from .decorators import copy_and_set_metadata
__all__ = ["get_dbz", "get_max_dbz"] __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 """ Return the dbz
do_varint - do variable intercept (if False, constants are used. Otherwise, do_varint - do variable intercept (if False, constants are used. Otherwise,
@ -17,8 +22,8 @@ def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False):
as liquid) as liquid)
""" """
ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR", varnames = ("T", "P", "PB", "QVAPOR", "QRAIN")
"QRAIN")) ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
pb = ncvars["PB"] pb = ncvars["PB"]
@ -26,14 +31,16 @@ def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False):
qr = ncvars["QRAIN"] qr = ncvars["QRAIN"]
try: try:
snowvars = extract_vars(wrfnc, timeidx, vars="QSNOW") snowvars = extract_vars(wrfnc, timeidx, "QSNOW",
method, squeeze, cache)
except KeyError: except KeyError:
qs = n.zeros(qv.shape, "float") qs = n.zeros(qv.shape, "float")
else: else:
qs = snowvars["QSNOW"] qs = snowvars["QSNOW"]
try: try:
graupvars = extract_vars(wrfnc, timeidx, vars="QGRAUP") graupvars = extract_vars(wrfnc, timeidx, "QGRAUP",
method, squeeze, cache)
except KeyError: except KeyError:
qg = n.zeros(qv.shape, "float") qg = n.zeros(qv.shape, "float")
else: else:
@ -58,7 +65,13 @@ def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False):
return computedbz(full_p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin) 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): @copy_and_set_metadata(copy_varname="T", name="dbz",
return n.amax(get_dbz(wrfnc, do_varint, do_liqskin, timeidx), 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) axis=-3)

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

@ -1,13 +1,16 @@
from functools import wraps #from functools import wraps
import wrapt
from inspect import getargspec from inspect import getargspec
from collections import OrderedDict
import numpy as n import numpy as np
import numpy.ma as ma import numpy.ma as ma
from wrf.var.units import do_conversion, check_units from .units import do_conversion, check_units
from wrf.var.destag import destagger from .destag import destagger
from wrf.var.util import iter_left_indexes from .util import (iter_left_indexes, viewitems, extract_vars,
from wrf.var.config import xarray_enabled combine_with, either)
from .config import xarray_enabled
if xarray_enabled(): if xarray_enabled():
from xarray import DataArray from xarray import DataArray
@ -15,40 +18,60 @@ if xarray_enabled():
__all__ = ["convert_units", "handle_left_iter", "uvmet_left_iter", __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 # TODO: In python 3.x, getargspec is deprecated
class from_args(object): def from_args(func, argnames, *args, **kwargs):
def __init__(self, argname):
self.argname
def __call__(self, func, *args, **kargs):
"""Parses the function args and kargs looking for the desired argument """Parses the function args and kargs looking for the desired argument
value. Otherwise, the value is taken from the default keyword argument value. Otherwise, the value is taken from the default keyword argument
using the arg spec. using the arg spec.
""" """
# If units are provided to the method call, use them. if isinstance(argnames, str):
# Otherwise, need to parse the argspec to find what the default arglist = [argnames]
# value is since wraps does not preserve this. 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) 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 return None
try: try:
result_idx = argspec.args.index(self.argname) result_idx = argspec.args.index(argname)
except ValueError: except ValueError:
result_idx = None result_idx = None
if (self.argname in kargs): result = None
result = kargs[self.argname] if (argname in kwargs):
result = kwargs, argname
elif (len(args) > result_idx and result_idx is not None): elif (len(args) > result_idx and result_idx is not None):
result = args[result_idx] result = args, result_idx
else: else:
arg_idx_from_right = (len(argspec.args) - arg_idx_from_right = (len(argspec.args) -
argspec.args.index(self.argname)) argspec.args.index(argname))
result = argspec.defaults[-arg_idx_from_right] result = list(argspec.defaults), -arg_idx_from_right
return result return result
@ -61,19 +84,19 @@ def convert_units(unit_type, alg_unit):
- alg_unit - the units that the function returns by default - alg_unit - the units that the function returns by default
""" """
def convert_decorator(func): # def convert_decorator(func):
@wraps(func) @wrapt.decorator
def func_wrapper(*args, **kargs): 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) check_units(desired_units, unit_type)
# Unit conversion done here # Unit conversion done here
return do_conversion(func(*args, **kargs), unit_type, return do_conversion(wrapped(*args, **kwargs), unit_type,
alg_unit, desired_units) alg_unit, desired_units)
return func_wrapper return func_wrapper
return convert_decorator # return convert_decorator
def _calc_out_dims(outvar, left_dims): def _calc_out_dims(outvar, left_dims):
@ -105,23 +128,23 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1,
the output dimensions (e.g. avo) the output dimensions (e.g. avo)
""" """
def indexing_decorator(func): # def indexing_decorator(func):
@wraps(func) @wrapt.decorator
def func_wrapper(*args, **kargs): def func_wrapper(wrapped, instance, args, kwargs):
ignore_args = ignore_args if ignore_args is not None else () ignore_args = ignore_args if ignore_args is not None else ()
ignore_kargs = ignore_kargs if ignore_kargs is not None else () ignore_kargs = ignore_kargs if ignore_kargs is not None else ()
if ref_var_idx >= 0: if ref_var_idx >= 0:
ref_var = args[ref_var_idx] ref_var = args[ref_var_idx]
else: else:
ref_var = kargs[ref_var_name] ref_var = kwargs[ref_var_name]
ref_var_shape = ref_var.shape ref_var_shape = ref_var.shape
extra_dim_num = ref_var.ndim - ref_var_expected_dims extra_dim_num = ref_var.ndim - ref_var_expected_dims
# No special left side iteration, return the function result # No special left side iteration, return the function result
if (extra_dim_num == 0): if (extra_dim_num == 0):
return func(*args, **kargs) return wrapped(*args, **kwargs)
# Start by getting the left-most 'extra' dims # Start by getting the left-most 'extra' dims
extra_dims = [ref_var_shape[x] for x in xrange(extra_dim_num)] 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,
# Slice the kargs if applicable # Slice the kargs if applicable
new_kargs = {key:(val[left_and_slice_idxs] new_kargs = {key:(val[left_and_slice_idxs]
if key not in ignore_kargs else val) if key not in ignore_kargs else val)
for key,val in kargs.iteritems()} for key,val in viewitems(kwargs)}
# Call the numerical routine # 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 # Output array
if not out_inited: if not out_inited:
outdims = _calc_out_dims(res, extra_dims) outdims = _calc_out_dims(res, extra_dims)
if not isinstance(res, ma.MaskedArray): if not isinstance(res, ma.MaskedArray):
output = n.zeros(outdims, ref_var.dtype) output = np.empty(outdims, ref_var.dtype)
masked = False masked = False
else: else:
output = ma.MaskedArray( output = ma.MaskedArray(
n.zeros(outdims, ref_var.dtype), np.zeros(outdims, ref_var.dtype),
mask=n.zeros(outdims, n.bool_), mask=np.zeros(outdims, np.bool_),
fill_value=res.fill_value) fill_value=res.fill_value)
masked = True masked = True
@ -172,13 +195,13 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1,
if not out_inited: if not out_inited:
outdims = _calc_out_dims(res[0], extra_dims) outdims = _calc_out_dims(res[0], extra_dims)
if not isinstance(res[0], ma.MaskedArray): 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))] for i in xrange(len(res))]
masked = False masked = False
else: else:
output = [ma.MaskedArray( output = [ma.MaskedArray(
n.zeros(outdims, ref_var.dtype), np.zeros(outdims, ref_var.dtype),
mask=n.zeros(outdims, n.bool_), mask=np.zeros(outdims, np.bool_),
fill_value=res[0].fill_value) fill_value=res[0].fill_value)
for i in xrange(len(res))] for i in xrange(len(res))]
masked = True masked = True
@ -197,16 +220,16 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1,
return func_wrapper return func_wrapper
return indexing_decorator # return indexing_decorator
def uvmet_left_iter(): def uvmet_left_iter():
"""Decorator to handle iterating over leftmost dimensions when using """Decorator to handle iterating over leftmost dimensions when using
multiple files and/or multiple times with the uvmet product. multiple files and/or multiple times with the uvmet product.
""" """
def indexing_decorator(func): #def indexing_decorator(func):
@wraps(func) @wrapt.decorator
def func_wrapper(*args): def func_wrapper(wrapped, instance, args, kwargs):
u = args[0] u = args[0]
v = args[1] v = args[1]
lat = args[2] lat = args[2]
@ -237,7 +260,7 @@ def uvmet_left_iter():
# No special left side iteration, return the function result # No special left side iteration, return the function result
if (extra_dim_num == 0): 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 # Start by getting the left-most 'extra' dims
outdims = [u.shape[x] for x in xrange(extra_dim_num)] outdims = [u.shape[x] for x in xrange(extra_dim_num)]
@ -248,7 +271,7 @@ def uvmet_left_iter():
outdims += [u.shape[x] for x in xrange(-num_right_dims,0,1)] 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): for left_idxs in iter_left_indexes(extra_dims):
# Make the left indexes plus a single slice object # Make the left indexes plus a single slice object
@ -264,12 +287,12 @@ def uvmet_left_iter():
new_lon = lon[left_and_slice_idxs] new_lon = lon[left_and_slice_idxs]
# Call the numerical routine # 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 # Note: The 2D version will return a 3D array with a 1 length
# dimension. Numpy is unable to broadcast this without # dimension. Numpy is unable to broadcast this without
# sqeezing first. # sqeezing first.
res = n.squeeze(res) res = np.squeeze(res)
output[left_and_slice_idxs] = res[:] output[left_and_slice_idxs] = res[:]
@ -278,16 +301,16 @@ def uvmet_left_iter():
return func_wrapper 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 """Decorator to handle casting to/from required dtype used in
numerical routine. numerical routine.
""" """
def cast_decorator(func): # def cast_decorator(func):
@wraps(func) @wrapt.decorator
def func_wrapper(*args, **kargs): def func_wrapper(wrapped, instance, args, kwargs):
arg_idxs = arg_idxs if arg_idxs is not None else () arg_idxs = arg_idxs if arg_idxs is not None else ()
karg_names = karg_names if karg_names 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):
new_kargs = {key:(val.astype(dtype) new_kargs = {key:(val.astype(dtype)
if key in karg_names else val) 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: if res.dtype == orig_type:
return res return res
return res.astype(orig_type) return res.astype(orig_type)
@ -314,14 +337,14 @@ def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=n.float64):
return func_wrapper return func_wrapper
return cast_decorator # return cast_decorator
def _extract_and_transpose(arg): def _extract_and_transpose(arg):
if xarray_enabled(): if xarray_enabled():
if isinstance(arg, DataArray): if isinstance(arg, DataArray):
arg = arg.values arg = arg.values
if isinstance(arg, n.ndarray): if isinstance(arg, np.ndarray):
if not arg.flags.F_CONTIGUOUS: if not arg.flags.F_CONTIGUOUS:
return arg.T return arg.T
@ -332,18 +355,18 @@ def handle_extract_transpose():
transposes if the data is not fortran contiguous. transposes if the data is not fortran contiguous.
""" """
def extract_transpose_decorator(func): # def extract_transpose_decorator(func):
@wraps(func) @wrapt.decorator
def func_wrapper(*args, **kargs): def func_wrapper(wrapped, instance, args, kwargs):
new_args = [_extract_and_transpose(arg) for arg in args] new_args = [_extract_and_transpose(arg) for arg in args]
new_kargs = {key:_extract_and_transpose(val) 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: if res.flags.F_CONTIGUOUS:
return res.T return res.T
else: else:
@ -353,32 +376,284 @@ def handle_extract_transpose():
return func_wrapper 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): @wrapt.decorator
@wraps(func) def func_wrapper(wrapped, instance, args, kwargs):
def func_wrapper(*args, **kargs):
if not xarray_enabled(): 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: if units is not None:
result.attrs["units"] = units outattrs["units"] = units
for argname, val in fixed_kargs.iteritems(): for argname, val in viewitems(fixed_attrs):
result[argname] = val 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 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 @@
from wrf.var.util import extract_vars from .util import extract_vars
__all__ = ["destagger", "destagger_windcomp", "destagger_winds"] __all__ = ["destagger", "destagger_windcomp", "destagger_winds"]
@ -40,24 +40,25 @@ def destagger(var, stagger_dim):
return result return result
def destagger_windcomp(wrfnc, comp, timeidx=0): # def destagger_windcomp(wrfnc, comp, timeidx=0,
if comp.lower() == "u": # method="cat", squeeze=True, cache=None):
wrfvar = "U" # if comp.lower() == "u":
stagdim = -1 # wrfvar = "U"
elif comp.lower() == "v": # stagdim = -1
wrfvar = "V" # elif comp.lower() == "v":
stagdim = -2 # wrfvar = "V"
elif comp.lower() == "w": # stagdim = -2
wrfvar = "W" # elif comp.lower() == "w":
stagdim = -3 # wrfvar = "W"
# stagdim = -3
#wind_data = wrfnc.variables[wrfvar][timeidx,:,:,:] #
ncvars = extract_vars(wrfnc, timeidx, vars=wrfvar) # ncvars = extract_vars(wrfnc, timeidx, wrfvar, method, squeeze, cache)
wind_data = ncvars[wrfvar] # wind_data = ncvars[wrfvar]
return destagger(wind_data, stagdim) #
# return destagger(wind_data, stagdim)
def destagger_winds(wrfnc, timeidx=0): #
return (destagger_windcomp(wrfnc, "u", timeidx), # def destagger_winds(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None):
destagger_windcomp(wrfnc, "v", timeidx), # return (destagger_windcomp(wrfnc, "u", timeidx, method, squeeze, cache),
destagger_windcomp(wrfnc, "w", timeidx)) # 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 @@
from wrf.var.extension import computetd from .extension import computetd
from wrf.var.decorators import convert_units from .decorators import convert_units, copy_and_set_metadata
from wrf.var.util import extract_vars from .util import extract_vars
__all__ = ["get_dp", "get_dp_2m"] __all__ = ["get_dp", "get_dp_2m"]
@copy_and_set_metadata(copy_varname="QVAPOR", name="td",
description="dew point temperature")
@convert_units("temp", "c") @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"] p = ncvars["P"]
pb = ncvars["PB"] pb = ncvars["PB"]
@ -20,9 +24,13 @@ def get_dp(wrfnc, timeidx=0, units="c"):
td = computetd(full_p, qvapor) td = computetd(full_p, qvapor)
return td return td
@copy_and_set_metadata(copy_varname="Q2", name="td2",
description="2m dew point temperature")
@convert_units("temp", "c") @convert_units("temp", "c")
def get_dp_2m(wrfnc, timeidx=0, units="c"): def get_dp_2m(wrfnc, timeidx=0, units="c",
ncvars = extract_vars(wrfnc, timeidx, varnames=("PSFC", "Q2")) method="cat", squeeze=True, cache=None):
varnames=("PSFC", "Q2")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
# Algorithm requires hPa # Algorithm requires hPa
psfc = .01*(ncvars["PSFC"]) psfc = .01*(ncvars["PSFC"])

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

@ -1,8 +1,8 @@
import numpy as n import numpy as np
from wrf.var.constants import Constants from .constants import Constants
from wrf.var.psadlookup import get_lookup_tables from .psadlookup import get_lookup_tables
from wrf.var._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d, from ._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d,
f_computeslp, f_computetk, f_computetd, f_computerh, f_computeslp, f_computetk, f_computetd, f_computerh,
f_computeabsvort,f_computepvo, f_computeeth, f_computeabsvort,f_computepvo, f_computeeth,
f_computeuvmet, f_computeuvmet,
@ -10,8 +10,8 @@ from wrf.var._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d,
f_computesrh, f_computeuh, f_computepw, f_computedbz, f_computesrh, f_computeuh, f_computepw, f_computedbz,
f_lltoij, f_ijtoll, f_converteta, f_computectt, f_lltoij, f_ijtoll, f_converteta, f_computectt,
f_monotonic, f_filter2d, f_vintrp) f_monotonic, f_filter2d, f_vintrp)
from wrf.var._wrfcape import f_computecape from ._wrfcape import f_computecape
from wrf.var.decorators import (handle_left_iter, uvmet_left_iter, from .decorators import (handle_left_iter, uvmet_left_iter,
handle_casting, handle_extract_transpose) handle_casting, handle_extract_transpose)
__all__ = ["FortranException", "computeslp", "computetk", "computetd", __all__ = ["FortranException", "computeslp", "computetk", "computetd",
@ -27,7 +27,7 @@ class FortranException(Exception):
@handle_left_iter(3,0, ignore_args=(2,3)) @handle_left_iter(3,0, ignore_args=(2,3))
@handle_casting(arg_idxs=(0,1)) @handle_casting(arg_idxs=(0,1))
@handle_extract_transpose() @handle_extract_transpose()
def interpz3d(data3d,zdata,desiredloc,missingval): def interpz3d(data3d, zdata, desiredloc, missingval):
res = f_interpz3d(data3d, res = f_interpz3d(data3d,
zdata, zdata,
desiredloc, desiredloc,
@ -44,7 +44,7 @@ def interp2dxy(data3d,xy):
@handle_casting(arg_idxs=(0,1,2)) @handle_casting(arg_idxs=(0,1,2))
@handle_extract_transpose() @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, res = f_interp1d(v_in,
z_in, z_in,
z_out, z_out,
@ -55,10 +55,10 @@ def interp1d(v_in,z_in,z_out,missingval):
@handle_left_iter(3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1,2,3)) @handle_casting(arg_idxs=(0,1,2,3))
@handle_extract_transpose() @handle_extract_transpose()
def computeslp(z,t,p,q): def computeslp(z, t, p, q):
t_surf = n.zeros((z.shape[-2], z.shape[-1]), "float64", order="F") t_surf = np.zeros((z.shape[-2], z.shape[-1]), np.float64, order="F")
t_sea_level = n.zeros((z.shape[-2], z.shape[-1]), "float64", order="F") t_sea_level = np.zeros((z.shape[-2], z.shape[-1]), np.float64, order="F")
level = n.zeros((z.shape[-2], z.shape[-1]), "int32", order="F") level = np.zeros((z.shape[-2], z.shape[-1]), np.int32, order="F")
res = f_computeslp(z, res = f_computeslp(z,
t, t,
@ -79,34 +79,34 @@ def computetk(pres, theta):
shape = pres.shape shape = pres.shape
res = f_computetk(pres.ravel(order="A"), res = f_computetk(pres.ravel(order="A"),
theta.ravel(order="A")) theta.ravel(order="A"))
res = n.reshape(res, shape) res = np.reshape(res, shape)
return res return res
# Note: No left iteration decorator needed with 1D arrays # Note: No left iteration decorator needed with 1D arrays
@handle_casting(arg_idxs=(0,1)) @handle_casting(arg_idxs=(0,1))
@handle_extract_transpose() @handle_extract_transpose()
def computetd(pressure,qv_in): def computetd(pressure, qv_in):
shape = pressure.shape shape = pressure.shape
res = f_computetd(pressure.ravel(order="A"), res = f_computetd(pressure.ravel(order="A"),
qv_in.ravel(order="A")) qv_in.ravel(order="A"))
res = n.reshape(res, shape) res = np.reshape(res, shape)
return res return res
# Note: No decorator needed with 1D arrays # Note: No decorator needed with 1D arrays
@handle_casting(arg_idxs=(0,1,2)) @handle_casting(arg_idxs=(0,1,2))
@handle_extract_transpose() @handle_extract_transpose()
def computerh(qv,q,t): def computerh(qv, q, t):
shape = qv.shape shape = qv.shape
res = f_computerh(qv.ravel(order="A"), res = f_computerh(qv.ravel(order="A"),
q.ravel(order="A"), q.ravel(order="A"),
t.ravel(order="A")) t.ravel(order="A"))
res = n.reshape(res, shape) res = np.reshape(res, shape)
return res return res
@handle_left_iter(3,0, ignore_args=(6,7)) @handle_left_iter(3,0, ignore_args=(6,7))
@handle_casting(arg_idxs=(0,1,2,3,4,5)) @handle_casting(arg_idxs=(0,1,2,3,4,5))
@handle_extract_transpose() @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, res = f_computeabsvort(u,
v, v,
msfu, msfu,
@ -121,7 +121,7 @@ def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy):
@handle_left_iter(3,2, ignore_args=(8,9)) @handle_left_iter(3,2, ignore_args=(8,9))
@handle_casting(arg_idxs=(0,1,2,3,4,5,6,7)) @handle_casting(arg_idxs=(0,1,2,3,4,5,6,7))
@handle_extract_transpose() @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, res = f_computepvo(u,
v, v,
@ -150,9 +150,9 @@ def computeeth(qv, tk, p):
@uvmet_left_iter() @uvmet_left_iter()
@handle_casting(arg_idxs=(0,1,2,3)) @handle_casting(arg_idxs=(0,1,2,3))
@handle_extract_transpose() @handle_extract_transpose()
def computeuvmet(u,v,lat,lon,cen_long,cone): def computeuvmet(u, v, lat, lon, cen_long, cone):
longca = n.zeros((lat.shape[-2], lat.shape[-1]), "float64", order="F") longca = np.zeros((lat.shape[-2], lat.shape[-1]), np.float64, order="F")
longcb = n.zeros((lon.shape[-2], lon.shape[-1]), "float64", order="F") longcb = np.zeros((lon.shape[-2], lon.shape[-1]), np.float64, order="F")
rpd = Constants.PI/180. rpd = Constants.PI/180.
@ -179,7 +179,7 @@ def computeuvmet(u,v,lat,lon,cen_long,cone):
0) 0)
return n.squeeze(res) return np.squeeze(res)
@handle_left_iter(3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1,2,3)) @handle_casting(arg_idxs=(0,1,2,3))
@ -197,7 +197,7 @@ def computeomega(qv, tk, w, p):
@handle_left_iter(3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1)) @handle_casting(arg_idxs=(0,1))
@handle_extract_transpose() @handle_extract_transpose()
def computetv(tk,qv): def computetv(tk, qv):
res = f_computetv(tk, res = f_computetv(tk,
qv) qv)
@ -206,7 +206,7 @@ def computetv(tk,qv):
@handle_left_iter(3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1,2)) @handle_casting(arg_idxs=(0,1,2))
@handle_extract_transpose() @handle_extract_transpose()
def computewetbulb(p,tk,qv): def computewetbulb(p, tk, qv):
PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables() PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables()
PSADITMK = PSADITMK.T PSADITMK = PSADITMK.T
@ -238,8 +238,10 @@ def computesrh(u, v, z, ter, top):
@handle_extract_transpose() @handle_extract_transpose()
def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top): 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") tem1 = np.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), np.float64,
tem2 = n.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), "float64", order="F") order="F")
tem2 = np.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), np.float64,
order="F")
res = f_computeuh(zp, res = f_computeuh(zp,
mapfct, mapfct,
@ -258,9 +260,9 @@ def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top):
@handle_left_iter(3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1,2,3)) @handle_casting(arg_idxs=(0,1,2,3))
@handle_extract_transpose() @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 # 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, res = f_computepw(p,
tv, tv,
qv, qv,
@ -272,7 +274,7 @@ def computepw(p,tv,qv,ht):
@handle_left_iter(3,0, ignore_args=(6,7,8)) @handle_left_iter(3,0, ignore_args=(6,7,8))
@handle_casting(arg_idxs=(0,1,2,3,4,5)) @handle_casting(arg_idxs=(0,1,2,3,4,5))
@handle_extract_transpose() @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, res = f_computedbz(p,
tk, tk,
@ -289,11 +291,11 @@ def computedbz(p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin):
@handle_left_iter(3,0,ignore_args=(6,7,8)) @handle_left_iter(3,0,ignore_args=(6,7,8))
@handle_casting(arg_idxs=(0,1,2,3,4,5)) @handle_casting(arg_idxs=(0,1,2,3,4,5))
@handle_extract_transpose() @handle_extract_transpose()
def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow): 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]), flip_cape = np.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]),
"float64", order="F") np.float64, order="F")
flip_cin = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), flip_cin = np.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]),
"float64", order="F") np.float64, order="F")
PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables() PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables()
PSADITMK = PSADITMK.T PSADITMK = PSADITMK.T
@ -327,37 +329,34 @@ def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow):
return (cape, cin) return (cape, cin)
# TODO: This should handle lists of coords def computeij(map_proj, truelat1, truelat2, stdlon,
def computeij(map_proj,truelat1,truelat2,stdlon, lat1, lon1, pole_lat, pole_lon,
lat1,lon1,pole_lat,pole_lon, knowni, knownj, dx, latinc, loninc, lat, lon):
knowni,knownj,dx,latinc,loninc,lat,lon):
res = f_lltoij(map_proj,truelat1,truelat2,stdlon, res = f_lltoij(map_proj,truelat1,truelat2,stdlon,
lat1,lon1,pole_lat,pole_lon, lat1,lon1,pole_lat,pole_lon,
knowni,knownj,dx,latinc,loninc,lat,lon, knowni,knownj,dx,latinc,loninc,lat,lon,
FortranException()) FortranException())
return res[0],res[1] return res
# TODO: This should handle lists of coords def computell(map_proj, truelat1, truelat2, stdlon, lat1, lon1,
def computell(map_proj,truelat1,truelat2,stdlon,lat1,lon1, pole_lat, pole_lon, knowni, knownj, dx, latinc,
pole_lat,pole_lon,knowni,knownj,dx,latinc, loninc, i, j):
loninc,i,j):
res = f_ijtoll(map_proj,truelat1,truelat2,stdlon,lat1,lon1, res = f_ijtoll(map_proj,truelat1,truelat2,stdlon,lat1,lon1,
pole_lat,pole_lon,knowni,knownj,dx,latinc, pole_lat,pole_lon,knowni,knownj,dx,latinc,
loninc,i,j,FortranException()) loninc,i,j,FortranException())
# Want lon,lat return res
return res[1],res[0]
@handle_left_iter(3,0, ignore_args=(3,)) @handle_left_iter(3,0, ignore_args=(3,))
@handle_casting(arg_idxs=(0,1,2)) @handle_casting(arg_idxs=(0,1,2))
@handle_extract_transpose() @handle_extract_transpose()
def computeeta(full_t, znu, psfc, ptop): def computeeta(full_t, znu, psfc, ptop):
pcalc = n.zeros(full_t.shape, "float64", order="F") pcalc = np.zeros(full_t.shape, np.float64, order="F")
mean_t = n.zeros(full_t.shape, "float64", order="F") mean_t = np.zeros(full_t.shape, np.float64, order="F")
temp_t = n.zeros(full_t.shape, "float64", order="F") temp_t = np.zeros(full_t.shape, np.float64, order="F")
res = f_converteta(full_t, res = f_converteta(full_t,
znu, znu,
@ -372,7 +371,7 @@ def computeeta(full_t, znu, psfc, ptop):
@handle_left_iter(3,0,ignore_args=(7,)) @handle_left_iter(3,0,ignore_args=(7,))
@handle_casting(arg_idxs=(0,1,2,3,4,5,6)) @handle_casting(arg_idxs=(0,1,2,3,4,5,6))
@handle_extract_transpose() @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, res = f_computectt(p_hpa,
tk, tk,
qice, qice,
@ -392,13 +391,13 @@ def smooth2d(field, passes):
# copies the original data before modifying it. This allows the decorator # copies the original data before modifying it. This allows the decorator
# to work properly and also behaves like the other methods. # 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 missing = field.fill_value
else: else:
missing = Constants.DEFAULT_FILL missing = Constants.DEFAULT_FILL
field_copy = field.copy() 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, f_filter2d(field_copy,
field_tmp, field_tmp,
@ -412,7 +411,7 @@ def smooth2d(field, passes):
@handle_left_iter(3,0,ignore_args=(3,4,5)) @handle_left_iter(3,0,ignore_args=(3,4,5))
@handle_casting(arg_idxs=(0,1,2)) @handle_casting(arg_idxs=(0,1,2))
@handle_extract_transpose() @handle_extract_transpose()
def monotonic(var,lvprs,coriolis,idir,delta,icorsw): def monotonic(var, lvprs, coriolis, idir, delta, icorsw):
res = f_monotonic(var, res = f_monotonic(var,
lvprs, lvprs,
coriolis, coriolis,
@ -425,8 +424,8 @@ def monotonic(var,lvprs,coriolis,idir,delta,icorsw):
@handle_left_iter(3,0,ignore_args=(9,10,11,12,13,14)) @handle_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_casting(arg_idxs=(0,1,2,3,4,5,6,7,8,9))
@handle_extract_transpose() @handle_extract_transpose()
def vintrp(field,pres,tk,qvp,ght,terrain,sfp,smsfp, def vintrp(field, pres, tk, qvp, ght, terrain, sfp, smsfp,
vcarray,interp_levels,icase,extrap,vcor,logp, vcarray, interp_levels, icase, extrap, vcor, logp,
missing): missing):
res = f_vintrp(field, res = f_vintrp(field,

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

@ -1,11 +1,12 @@
from wrf.var.constants import Constants from .constants import Constants
from wrf.var.destag import destagger from .destag import destagger
from wrf.var.decorators import convert_units from .decorators import convert_units, set_height_metadata
from wrf.var.util import extract_vars from .util import extract_vars, either
__all__ = ["get_geopt", "get_height"] __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, """Return the geopotential in units of m2 s-2 if height is False,
otherwise return the geopotential height in meters. If height is True, 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 then if msl is True the result will be in MSL, otherwise AGL (the terrain
@ -13,23 +14,18 @@ def _get_geoht(wrfnc, timeidx, height=True, msl=True):
""" """
try: varname = either("PH", "GHT")(wrfnc)
if varname == "PH":
ph_vars = extract_vars(wrfnc, timeidx, varnames=("PH", "PHB", "HGT")) 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"] ph = ph_vars["PH"]
phb = ph_vars["PHB"] phb = ph_vars["PHB"]
hgt = ph_vars["HGT"] hgt = ph_vars["HGT"]
geopt = ph + phb geopt = ph + phb
geopt_unstag = destagger(geopt, -3) 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 height:
if msl: if msl:
@ -46,11 +42,14 @@ def _get_geoht(wrfnc, timeidx, height=True, msl=True):
else: else:
return geopt_unstag return geopt_unstag
def get_geopt(wrfnc, timeidx=0): @set_height_metadata(geopt=True)
return _get_geoht(wrfnc, timeidx, False) 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") @convert_units("height", "m")
def get_height(wrfnc, timeidx=0, msl=True, units="m"): def get_height(wrfnc, timeidx=0, msl=True, units="m",
z = _get_geoht(wrfnc, timeidx, True, msl) method="cat", squeeze=True, cache=None):
return z
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 @@
from wrf.var.constants import Constants from .constants import Constants
from wrf.var.extension import computesrh, computeuh from .extension import computesrh, computeuh
from wrf.var.destag import destagger from .destag import destagger
from wrf.var.util import extract_vars, extract_global_attrs from .util import extract_vars, extract_global_attrs, either
from .decorators import copy_and_set_metadata
__all__ = ["get_srh", "get_uh"] __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) # 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"] ter = ncvars["HGT"]
ph = ncvars["PH"] ph = ncvars["PH"]
phb = ncvars["PHB"] phb = ncvars["PHB"]
try: # As coded in NCL, but not sure this is possible
u_vars = extract_vars(wrfnc, timeidx, varnames="U") varname = either("U", "UU")(wrfnc)
except KeyError: u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
try: u = destagger(u_vars[varname], -1)
uu_vars = extract_vars(wrfnc, timeidx, varnames="UU")
except KeyError: varname = either("V", "VV")(wrfnc)
raise RuntimeError("No valid wind data found in NetCDF file") v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
else: v = destagger(v_vars[varname], -2)
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)
geopt = ph + phb geopt = ph + phb
geopt_unstag = destagger(geopt, -3) geopt_unstag = destagger(geopt, -3)
@ -53,9 +44,14 @@ def get_srh(wrfnc, timeidx=0, top=3000.0):
return srh 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"] wstag = ncvars["W"]
ph = ncvars["PH"] ph = ncvars["PH"]
@ -66,29 +62,14 @@ def get_uh(wrfnc, timeidx=0, bottom=2000.0, top=5000.0):
dx = attrs["DX"] dx = attrs["DX"]
dy = attrs["DY"] dy = attrs["DY"]
try: # As coded in NCL, but not sure this is possible
u_vars = extract_vars(wrfnc, timeidx, varnames="U") varname = either("U", "UU")(wrfnc)
except KeyError: u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
try: u = destagger(u_vars[varname], -1)
uu_vars = extract_vars(wrfnc, timeidx, varnames="UU")
except KeyError: varname = either("V", "VV")(wrfnc)
raise RuntimeError("No valid wind data found in NetCDF file") v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
else: v = destagger(v_vars[varname], -2)
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)
zp = ph + phb 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
import numpy as n import numpy as n
import numpy.ma as ma import numpy.ma as ma
from wrf.var.extension import (interpz3d,interp2dxy,interp1d, from .extension import (interpz3d,interp2dxy,interp1d,
smooth2d,monotonic,vintrp) smooth2d,monotonic,vintrp)
from wrf.var.decorators import handle_left_iter, handle_casting from .decorators import handle_left_iter, handle_casting
from wrf.var.util import extract_vars, is_staggered from .util import extract_vars, is_staggered
from wrf.var.constants import Constants, ConversionFactors from .constants import Constants, ConversionFactors
from wrf.var.terrain import get_terrain from .terrain import get_terrain
from wrf.var.geoht import get_height from .geoht import get_height
from wrf.var.temp import get_theta, get_temp, get_eth from .temp import get_theta, get_temp, get_eth
from wrf.var.pressure import get_pressure from .pressure import get_pressure
__all__ = ["interplevel", "vertcross", "interpline", "vinterp"] __all__ = ["interplevel", "vertcross", "interpline", "vinterp"]
# Note: Extension decorator is good enough to handle left dims # 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 """Return the horizontally interpolated data at the provided level
data3d - the 3D field to interpolate data3d - the 3D field to interpolate
@ -28,11 +28,13 @@ def interplevel(data3d,zdata,desiredloc,missingval=Constants.DEFAULT_FILL):
""" """
r1 = interpz3d(data3d, zdata, desiredloc, missingval) r1 = interpz3d(data3d, zdata, desiredloc, missingval)
masked_r1 = ma.masked_values (r1, missingval) masked_r1 = ma.masked_values (r1, missingval)
return masked_r1 return masked_r1
def _to_positive_idxs(shape, coord): 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 coord
return [x if (x >= 0) else shape[i]+x for (i,x) in enumerate(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, 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 xdim - maximum x-dimension
ydim - maximum y-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 angle - the angle through the pivot point in degrees
start_point - a start_point sequence of [south_north1, west_east1] start_point - a start_point sequence of [south_north1, west_east1]
end_point - an end point sequence of [south_north2, west_east2] end_point - an end point sequence of [south_north2, west_east2]
@ -148,15 +151,16 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None,
"start_point", "end_point")) "start_point", "end_point"))
@handle_casting(arg_idxs=(0,1)) @handle_casting(arg_idxs=(0,1))
def vertcross(data3d, z, missingval=Constants.DEFAULT_FILL, def vertcross(data3d, z, missingval=Constants.DEFAULT_FILL,
pivot_point=None,angle=None, pivot_point=None, angle=None,
start_point=None,end_point=None): start_point=None, end_point=None):
"""Return the vertical cross section for a 3D field, interpolated """Return the vertical cross section for a 3D field, interpolated
to a verical plane defined by a horizontal line. to a verical plane defined by a horizontal line.
Arguments: Arguments:
data3d - a 3D data field data3d - a 3D data field
z - 3D height 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 angle - the angle through the pivot point in degrees
start_point - a start_point tuple of (south_north1,west_east1) start_point - a start_point tuple of (south_north1,west_east1)
end_point - an end point tuple of (south_north2,west_east2) end_point - an end point tuple of (south_north2,west_east2)
@ -362,7 +366,7 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False,
p_hpa = p * ConversionFactors.PA_TO_HPA 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 # We only extrapolate temperature fields below ground
# if we are interpolating to pressure or height vertical surfaces. # if we are interpolating to pressure or height vertical surfaces.
@ -379,7 +383,7 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False,
p_hpa = p * ConversionFactors.PA_TO_HPA 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 # We only extrapolate temperature fields below ground if we are
# interpolating to pressure or height vertical surfaces # interpolating to pressure or height vertical surfaces
icase = 0 icase = 0
@ -390,11 +394,11 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False,
else: else:
missing = Constants.DEFAULT_FILL missing = Constants.DEFAULT_FILL
res = vintrp(field,p,tk,qv,ght,terht,sfp,smsfp, res = vintrp(field, p, tk, qv, ght, terht, sfp, smsfp,
vcord_array,interp_levels, vcord_array, interp_levels,
icase,extrap,vcor,log_p_int,missing) 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 @@
from collections import Iterable from collections import Iterable
from wrf.var.constants import Constants import numpy as np
from wrf.var.extension import computeij, computell
from wrf.var.util import extract_vars, extract_global_attrs
__all__ = ["get_lat", "get_lon", "get_ij", "get_ll"] 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: __all__ = ["get_lat", "get_lon", "get_ij", "get_ll"]
lat_vars = extract_vars(wrfnc, timeidx, varnames="XLAT")
except KeyError: def _lat_varname(stagger):
try: if stagger is None or stagger.lower() == "m":
latm_vars = extract_vars(wrfnc, timeidx, varnames="XLAT_M") varname = either("XLAT", "XLAT_M")
except: elif stagger.lower() == "u" or stagger.lower() == "v":
raise RuntimeError("Latitude variable not found in NetCDF file") varname = "XLAT_{}".format(stagger.upper())
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")
else: 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: 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: if timeidx < 0:
raise ValueError("'timeidx' must be greater than 0") raise ValueError("'timeidx' must be greater than 0")
@ -64,12 +78,22 @@ def _get_proj_params(wrfnc, timeidx):
latinc = 0.0 latinc = 0.0
loninc = 0.0 loninc = 0.0
xlat = get_lat(wrfnc, timeidx) latvar = _lat_varname(stagger)
xlon = get_lon(wrfnc, timeidx) lonvar = _lon_varname(stagger)
ref_lat = xlat[0,0] lat_timeidx = timeidx
ref_lon = xlon[0,0] # 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_i = 1.0
known_j = 1.0 known_j = 1.0
@ -78,20 +102,70 @@ def _get_proj_params(wrfnc, timeidx):
loninc) loninc)
# TODO: longitude and latitude can also be lists @set_latlon_metadata(ij=True)
def get_ij(wrfnc, longitude, latitude, timeidx=0): def get_ij(wrfnc, latitude, longitude, timeidx=0,
if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str): stagger=None, method="cat", squeeze=True, cache=None):
raise TypeError("'get_ij' is only applicabe for single files")
(map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon, (map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon,
pole_lat,pole_lon,known_i,known_j,dx,latinc, pole_lat,pole_lon,known_i,known_j,dx,latinc,
loninc) = _get_proj_params(wrfnc, timeidx) loninc) = _get_proj_params(wrfnc, timeidx)
return computeij(map_proj,truelat1,truelat2,stdlon, if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str):
ref_lat,ref_lon,pole_lat,pole_lon, lats = np.asarray(latitude)
known_i,known_j,dx,latinc,loninc,latitude,longitude) 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): def get_ll(wrfnc, i, j, timeidx=0):
if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str): if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str):
raise TypeError("'get_ll' is only applicabe for single files") raise TypeError("'get_ll' is only applicabe for single files")
@ -100,9 +174,59 @@ def get_ll(wrfnc, i, j, timeidx=0):
pole_lat,pole_lon,known_i,known_j,dx,latinc, pole_lat,pole_lon,known_i,known_j,dx,latinc,
loninc) = _get_proj_params(wrfnc, timeidx) loninc) = _get_proj_params(wrfnc, timeidx)
return computell(map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon, if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str):
pole_lat,pole_lon,known_i,known_j,dx,latinc, i_arr = np.asarray(i)
loninc,i,j) 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 @@
from wrf.var.constants import Constants from .constants import Constants
from wrf.var.destag import destagger from .destag import destagger
from wrf.var.extension import computeomega,computetk from .extension import computeomega,computetk
from wrf.var.util import extract_vars from .util import extract_vars
from .decorators import copy_and_set_metadata
__all__ = ["get_omega"] __all__ = ["get_omega"]
def get_omega(wrfnc, timeidx=0): @copy_and_set_metadata(copy_varname="T", name="omega",
ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "W", description="omega",
"PB", "QVAPOR")) 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"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
w = ncvars["W"] w = ncvars["W"]

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

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

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

@ -1,30 +1,31 @@
from wrf.var.decorators import convert_units from .decorators import convert_units, copy_and_set_metadata
from wrf.var.util import extract_vars from .util import extract_vars, either
__all__ = ["get_pressure", "get_pressure_hpa"] __all__ = ["get_pressure", "get_pressure_hpa"]
@copy_and_set_metadata(copy_varname=either("P", "PRES"), name="pressure",
description="pressure")
@convert_units("pressure", "pa") @convert_units("pressure", "pa")
def get_pressure(wrfnc, timeidx=0, units="pa"): def get_pressure(wrfnc, timeidx=0, units="pa",
method="cat", squeeze=True, cache=None):
try:
p_vars = extract_vars(wrfnc, timeidx, varnames=("P", "PB")) varname = either("P", "PRES")(wrfnc)
except KeyError: if varname == "P":
try: p_vars = extract_vars(wrfnc, timeidx, ("P", "PB"),
pres_vars = extract_vars(wrfnc, timeidx, varnames="PRES") method, squeeze, cache)
except:
raise RuntimeError("pressure variable not found in NetCDF file")
else:
pres = pres_vars["PRES"]
else:
p = p_vars["P"] p = p_vars["P"]
pb = p_vars["PB"] pb = p_vars["PB"]
pres = p + pb pres = p + pb
else:
pres = extract_vars(wrfnc, timeidx, "PRES",
method, squeeze, cache)["PRES"]
return pres return pres
def get_pressure_hpa(wrfnc, timeidx=0, units="hpa"): def get_pressure_hpa(wrfnc, timeidx=0, units="hpa",
return get_pressure(wrfnc, timeidx, units=units) 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 @@
import numpy as np import numpy as np
import math import math
from wrf.var.config import basemap_enabled, cartopy_enabled, pyngl_enabled from .config import basemap_enabled, cartopy_enabled, pyngl_enabled
from wrf.var.constants import Constants from .constants import Constants
if cartopy_enabled(): if cartopy_enabled():
from cartopy import crs from cartopy import crs

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

@ -1,13 +1,17 @@
from wrf.var.extension import computepw,computetv,computetk from .extension import computepw,computetv,computetk
from wrf.var.constants import Constants from .constants import Constants
from wrf.var.util import extract_vars from .util import extract_vars
from .decorators import copy_and_set_metadata
__all__ = ["get_pw"] __all__ = ["get_pw"]
def get_pw(wrfnc, timeidx=0): @copy_and_set_metadata(copy_varname="T", name="pw",
ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "PH", description="precipitable water",
"PHB", "QVAPOR")) 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"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
@ -22,9 +26,9 @@ def get_pw(wrfnc, timeidx=0):
full_t = t + Constants.T_BASE full_t = t + Constants.T_BASE
tk = computetk(full_p, full_t) 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 @@
from wrf.var.constants import Constants from .constants import Constants
from wrf.var.extension import computerh, computetk from .extension import computerh, computetk
from wrf.var.util import extract_vars from .util import extract_vars
from .decorators import copy_and_set_metadata
__all__ = ["get_rh", "get_rh_2m"] __all__ = ["get_rh", "get_rh_2m"]
def get_rh(wrfnc, timeidx=0): @copy_and_set_metadata(copy_varname="T", name="rh",
ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR")) 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"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
pb = ncvars["PB"] pb = ncvars["PB"]
@ -20,8 +25,12 @@ def get_rh(wrfnc, timeidx=0):
return rh return rh
def get_rh_2m(wrfnc, timeidx=0): @copy_and_set_metadata(copy_varname="T2", name="rh2",
ncvars = extract_vars(wrfnc, timeidx, varnames=("T2", "PSFC", "Q2")) 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"] t2 = ncvars["T2"]
psfc = ncvars["PSFC"] psfc = ncvars["PSFC"]
q2 = ncvars["Q2"] q2 = ncvars["Q2"]

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

@ -1,15 +1,19 @@
from wrf.var.extension import computeslp, computetk from .extension import computeslp, computetk
from wrf.var.constants import Constants from .constants import Constants
from wrf.var.destag import destagger from .destag import destagger
from wrf.var.decorators import convert_units from .decorators import convert_units, copy_and_set_metadata
from wrf.var.util import extract_vars from .util import extract_vars
__all__ = ["get_slp"] __all__ = ["get_slp"]
@copy_and_set_metadata(copy_varname="T", name="slp",
remove_dims=("bottom_top",),
description="sea level pressure")
@convert_units("pressure", "hpa") @convert_units("pressure", "hpa")
def get_slp(wrfnc, timeidx=0, units="hpa"): def get_slp(wrfnc, timeidx=0, units="hpa",
ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR", method="cat", squeeze=True, cache=None):
"PH", "PHB")) varnames=("T", "P", "PB", "QVAPOR", "PH", "PHB")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]

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

@ -1,26 +1,34 @@
from wrf.var.constants import Constants from .constants import Constants
from wrf.var.extension import computetk, computeeth, computetv, computewetbulb from .extension import computetk, computeeth, computetv, computewetbulb
from wrf.var.decorators import convert_units, set_metadata from .decorators import convert_units, copy_and_set_metadata
from wrf.var.util import extract_vars from .util import extract_vars
__all__ = ["get_theta", "get_temp", "get_eth", "get_tv", "get_tw", __all__ = ["get_theta", "get_temp", "get_eth", "get_tv", "get_tw",
"get_tk", "get_tc"] "get_tk", "get_tc"]
@set_metadata() @copy_and_set_metadata(copy_varname="T", name="theta",
description="potential temperature")
@convert_units("temp", "k") @convert_units("temp", "k")
def get_theta(wrfnc, timeidx=0, units="k"): def get_theta(wrfnc, timeidx=0, units="k",
ncvars = extract_vars(wrfnc, timeidx, varnames="T") method="cat", squeeze=True, cache=None):
varnames = ("T",)
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
t = ncvars["T"] t = ncvars["T"]
full_t = t + Constants.T_BASE full_t = t + Constants.T_BASE
return full_t return full_t
@copy_and_set_metadata(copy_varname="T", name="temp",
description="temperature")
@convert_units("temp", "k") @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""" """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"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
pb = ncvars["PB"] pb = ncvars["PB"]
@ -31,11 +39,15 @@ def get_temp(wrfnc, timeidx=0, units="k"):
return tk return tk
@copy_and_set_metadata(copy_varname="T", name="theta_e",
description="equivalent potential temperature")
@convert_units("temp", "k") @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" "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"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
pb = ncvars["PB"] pb = ncvars["PB"]
@ -49,11 +61,15 @@ def get_eth(wrfnc, timeidx=0, units="k"):
return eth return eth
@copy_and_set_metadata(copy_varname="T", name="tv",
description="virtual temperature")
@convert_units("temp", "k") @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" "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"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
@ -68,12 +84,15 @@ def get_tv(wrfnc, timeidx=0, units="k"):
return tv return tv
@copy_and_set_metadata(copy_varname="T", name="twb",
description="wetbulb temperature")
@convert_units("temp", "k") @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)" "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"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
pb = ncvars["PB"] pb = ncvars["PB"]
@ -87,11 +106,13 @@ def get_tw(wrfnc, timeidx=0, units="k"):
return tw return tw
def get_tk(wrfnc, timeidx=0): def get_tk(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None):
return get_temp(wrfnc, timeidx, units="k") return get_temp(wrfnc, timeidx, units="k",
method=method, squeeze=squeeze, cache=cache)
def get_tc(wrfnc, timeidx=0): def get_tc(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None):
return get_temp(wrfnc, timeidx, units="c") 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 @@
from wrf.var.decorators import convert_units from .decorators import convert_units, copy_and_set_metadata
from wrf.var.util import extract_vars from .util import extract_vars, either
__all__ = ["get_terrain"] __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") @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 @@
from wrf.var.util import extract_times from .util import extract_times
__all__ = ["get_times"] __all__ = ["get_times"]

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

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

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

@ -1,12 +1,11 @@
from collections import Iterable, Mapping, OrderedDict from collections import Iterable, Mapping, OrderedDict
from itertools import product from itertools import product
import datetime as dt import datetime as dt
import warnings
import numpy as n import numpy as np
from wrf.var.config import xarray_enabled from .config import xarray_enabled
from wrf.var.projection import getproj from .projection import getproj
if xarray_enabled(): if xarray_enabled():
from xarray import DataArray from xarray import DataArray
@ -14,7 +13,8 @@ if xarray_enabled():
__all__ = ["extract_vars", "extract_global_attrs", "extract_dim", __all__ = ["extract_vars", "extract_global_attrs", "extract_dim",
"combine_files", "is_standard_wrf_var", "extract_times", "combine_files", "is_standard_wrf_var", "extract_times",
"iter_left_indexes", "get_left_indexes", "get_right_slices", "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): def _is_multi_time(timeidx):
if timeidx == -1: if timeidx == -1:
@ -31,23 +31,77 @@ def _is_mapping(wrfnc):
return True return True
return False return False
def _is_moving_domain(wrfnc, latvar="XLAT", lonvar="XLONG"): def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar):
lat1 = wrfnc.variables[latvar][0,:] lats = wrfnc.variables[latvar]
lat2 = wrfnc.variables[latvar][-1,:] lons = wrfnc.variables[lonvar]
lon1 = wrfnc.variables[lonvar][0,:] # Need to check all times
lon2 = wrfnc.variables[lonvar][-1,:] 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 ll_corner = (lat0, lon0)
lon1[0,0] != lon2[0,0] or lon1[-1,-1] != lon2[-1,-1]): 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 True
return False return False
def _get_attr(wrfnc, attr): def _get_global_attr(wrfnc, attr):
val = getattr(wrfnc, attr, None) val = getattr(wrfnc, attr, None)
# PyNIO puts single values in to an array # PyNIO puts single values in to an array
if isinstance(val,n.ndarray): if isinstance(val,np.ndarray):
if len(val) == 1: if len(val) == 1:
return val[0] return val[0]
return val return val
@ -66,7 +120,7 @@ def extract_global_attrs(wrfnc, attrs):
else: else:
wrfnc = wrfnc[next(wrfnc.iterkeys())] 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): def extract_dim(wrfnc, dim):
if _is_multi_file(wrfnc): if _is_multi_file(wrfnc):
@ -80,71 +134,81 @@ def extract_dim(wrfnc, dim):
return len(d) #netCDF4 return len(d) #netCDF4
return d # PyNIO return d # PyNIO
def _combine_dict(wrfdict, varname, timeidx, method):
# TODO
def _combine_dict(wrfseq, var, timeidx, method):
"""Dictionary combination creates a new left index for each key, then """Dictionary combination creates a new left index for each key, then
does a cat or join for the list of files for that key""" does a cat or join for the list of files for that key"""
keynames = []
numkeys = len(wrfdict)
multitime = _is_multi_time(timeidx) key_iter = iter(viewkeys(wrfdict))
numfiles = len(wrfseq) first_key = next(key_iter)
keynames.append(first_key)
if not multitime: first_array = _extract_var(wrfdict[first_key], varname,
time_idx_or_slice = timeidx timeidx, method, squeeze=False)
else:
time_idx_or_slice = slice(None, None, None)
keys = (list(x for x in xrange(numfiles)) if not _is_mapping(wrfseq) else
list(key for key in wrfseq.iterkeys()))
# 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 # Create the output data numpy array based on the first array
outdims = [numfiles] outdims = [numkeys]
outdims += first_var.shape outdims += first_array.shape
outdata = n.zeros(outdims, first_var.dtype) outdata = np.empty(outdims, first_array.dtype)
outdata[0,:] = first_var[:] 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): if outdata.shape[1:] != vardata.shape:
outdata[idx,:] = wrfseq[key].variables[var][time_idx_or_slice, :] raise ValueError("data sequences must have the "
"same size for all dictionary keys")
outdata[idx,:] = vardata.values[:]
idx += 1
if not xarray_enabled(): if not xarray_enabled():
outarr = outdata outarr = outdata
else: else:
outname = str(first_var.name) outname = str(first_array.name)
outcoords = dict(first_var.coords) # Note: assumes that all entries in dict have same coords
outdims = ["sequence"] + list(first_var.dims) outcoords = OrderedDict(first_array.coords)
outcoords["sequence"] = keys outdims = ["key"] + list(first_array.dims)
outattrs = dict(first_var.attrs) outcoords["key"] = keynames
outattrs = OrderedDict(first_array.attrs)
outarr = DataArray(outdata, name=outname, coords=outcoords, outarr = DataArray(outdata, name=outname, coords=outcoords,
dims=outdims, attrs=outattrs) dims=outdims, attrs=outattrs)
return outarr 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) file_times = extract_times(wrfseq, timeidx)
multitime = _is_multi_time(timeidx) multitime = _is_multi_time(timeidx)
time_idx_or_slice = timeidx if not multitime else slice(None, None, None) 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 if xarray_enabled() else
wrfseq[0].variables[var][time_idx_or_slice, :]) wrfseq[0].variables[varname][time_idx_or_slice, :])
outdims = [len(file_times)] outdims = [len(file_times)]
# Making a new time dim, so ignore this one # Making a new time dim, so ignore this one
outdims += first_var.shape[1:] outdims += first_var.shape[1:]
outdata = n.zeros(outdims, first_var.dtype) outdata = np.empty(outdims, first_var.dtype)
numtimes = first_var.shape[0] numtimes = first_var.shape[0]
startidx = 0 startidx = 0
@ -153,8 +217,14 @@ def _cat_files(wrfseq, var, timeidx):
outdata[startidx:endidx, :] = first_var[:] outdata[startidx:endidx, :] = first_var[:]
startidx = endidx startidx = endidx
for wrfnc in wrfseq[1:]: while True:
vardata = wrfnc.variables[var][time_idx_or_slice, :] try:
wrfnc = next(wrf_iter)
except StopIteration:
break
else:
vardata = wrfnc.variables[varname][time_idx_or_slice, :]
if multitime: if multitime:
numtimes = vardata.shape[0] numtimes = vardata.shape[0]
else: else:
@ -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 # FIXME: If it's a moving nest, then the coord arrays need to have same
# time indexes as the whole data set # time indexes as the whole data set
outname = str(first_var.name) outname = str(first_var.name)
outattrs = dict(first_var.attrs) outattrs = OrderedDict(first_var.attrs)
outcoords = dict(first_var.coords) outcoords = OrderedDict(first_var.coords)
outdimnames = list(first_var.dims) outdimnames = list(first_var.dims)
outcoords[outdimnames[0]] = file_times # New time dimension values outcoords[outdimnames[0]] = file_times # New time dimension values
@ -183,8 +253,9 @@ def _cat_files(wrfseq, var, timeidx):
return outarr 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) multitime = _is_multi_time(timeidx)
numfiles = len(wrfseq) numfiles = len(wrfseq)
@ -193,26 +264,36 @@ def _join_files(wrfseq, var, timeidx):
else: else:
time_idx_or_slice = slice(None, None, None) time_idx_or_slice = slice(None, None, None)
# wrfseq might be a generator
wrf_iter = iter(wrfseq)
if xarray_enabled(): 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: 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 # Create the output data numpy array based on the first array
outdims = [numfiles] outdims = [numfiles]
outdims += first_var.shape outdims += first_var.shape
outdata = n.zeros(outdims, first_var.dtype) outdata = np.empty(outdims, first_var.dtype)
outdata[0,:] = first_var[:] outdata[0,:] = first_var[:]
for idx, wrfnc in enumerate(wrfseq[1:], 1): idx=1
print idx while True:
outdata[idx,:] = wrfnc.variables[var][time_idx_or_slice, :] try:
wrfnc = next(wrf_iter)
except StopIteration:
break
else:
outdata[idx,:] = wrfnc.variables[varname][time_idx_or_slice, :]
idx += 1
if xarray_enabled(): if xarray_enabled():
outname = str(first_var.name) outname = str(first_var.name)
outcoords = dict(first_var.coords) outcoords = OrderedDict(first_var.coords)
outattrs = dict(first_var.attrs) outattrs = OrderedDict(first_var.attrs)
# New dimensions # New dimensions
outdimnames = ["file_idx"] + list(first_var.dims) outdimnames = ["file_idx"] + list(first_var.dims)
outcoords["file_idx"] = [i for i in xrange(numfiles)] outcoords["file_idx"] = [i for i in xrange(numfiles)]
@ -225,42 +306,57 @@ def _join_files(wrfseq, var, timeidx):
return outarr 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 # Dictionary is unique
if _is_mapping(wrfseq): if _is_mapping(wrfseq):
outarr = _combine_dict(wrfseq, var, timeidx, method) outarr = _combine_dict(wrfseq, varname, timeidx, method)
elif method.lower() == "cat":
if method.lower() == "cat": outarr = _cat_files(wrfseq, varname, timeidx)
outarr = _cat_files(wrfseq, var, timeidx)
elif method.lower() == "join": elif method.lower() == "join":
outarr = _join_files(wrfseq, var, timeidx) outarr = _join_files(wrfseq, varname, timeidx)
else: else:
raise ValueError("method must be 'cat' or 'join'") raise ValueError("method must be 'cat' or 'join'")
if squeeze: return outarr.squeeze() if squeeze else outarr
return outarr.squeeze()
return outarr
# Note, always returns the full data set with the time dimension included # 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) multitime = _is_multi_time(timeidx)
var = wrfnc.variables[varname] var = wrfnc.variables[varname]
data = var[:] data = var[:]
attrs = OrderedDict(var.__dict__) attrs = OrderedDict(var.__dict__)
dimnames = var.dimensions dimnames = var.dimensions
# Add the coordinate variables here. # WRF variables will have a coordinates attribute. MET_EM files have
coord_names = getattr(var, "coordinates").split() # 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] lon_coord = coord_names[0]
lat_coord = coord_names[1] 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] lon_coord_var = wrfnc.variables[lon_coord]
lat_coord_var = wrfnc.variables[lat_coord] lat_coord_var = wrfnc.variables[lat_coord]
if multitime: 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, # Special case with a moving domain in a multi-time file,
# otherwise the projection parameters don't change # otherwise the projection parameters don't change
coords[lon_coord] = lon_coord_var.dimensions, lon_coord_var[:] coords[lon_coord] = lon_coord_var.dimensions, lon_coord_var[:]
@ -291,38 +387,53 @@ def _build_data_array(wrfnc, varname, timeidx):
lats, lons, proj_params = get_proj_params(wrfnc, 0, varname) lats, lons, proj_params = get_proj_params(wrfnc, 0, varname)
proj = getproj(lats=lats, lons=lons, **proj_params) 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, data_array = DataArray(data, name=varname, dims=dimnames, coords=coords,
attrs=attrs) attrs=attrs)
return data_array 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) multitime = _is_multi_time(timeidx)
multifile = _is_multi_file(wrfnc) multifile = _is_multi_file(wrfnc)
if not multifile: if not multifile:
if xarray_enabled(): if xarray_enabled():
return _build_data_array(wrfnc, varname, timeidx) result = _build_data_array(wrfnc, varname, timeidx, is_moving)
else: else:
if not multitime: if not multitime:
return wrfnc.variables[varname][timeidx,:] result = wrfnc.variables[varname][timeidx,:]
else: else:
return wrfnc.variables[varname][:] result = wrfnc.variables[varname][:]
else: 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): if isinstance(varnames, str):
varlist = [varnames] varlist = [varnames]
else: else:
varlist = varnames 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} for var in varlist}
def _make_time(timearr): def _make_time(timearr):
@ -430,6 +541,78 @@ def get_proj_params(wrfnc, timeidx=0, varname=None):
wrfnc.variables[lon_coord][time_idx_or_slice,:], wrfnc.variables[lon_coord][time_idx_or_slice,:],
proj_params) 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 @@
from math import fabs, log, tan, sin, cos from math import fabs, log, tan, sin, cos
from wrf.var.extension import computeuvmet from .extension import computeuvmet
from wrf.var.destag import destagger from .destag import destagger
from wrf.var.constants import Constants from .constants import Constants
from wrf.var.wind import _calc_wspd_wdir from .wind import _calc_wspd_wdir
from wrf.var.decorators import convert_units from .decorators import convert_units, set_wind_metadata
from wrf.var.util import extract_vars, extract_global_attrs from .util import extract_vars, extract_global_attrs, either
__all__=["get_uvmet", "get_uvmet10", "get_uvmet_wspd_wdir", __all__=["get_uvmet", "get_uvmet10", "get_uvmet_wspd_wdir",
"get_uvmet10_wspd_wdir"] "get_uvmet10_wspd_wdir"]
@set_wind_metadata(wspd_wdir=False)
@convert_units("wind", "mps") @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""" """ Return a tuple of u,v with the winds rotated in to earth space"""
if not ten_m: if not ten_m:
try: varname = either("U", "UU")(wrfnc)
u_vars = extract_vars(wrfnc, timeidx, varnames="U") u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
except KeyError: u = destagger(u_vars[varname], -1)
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("V", "VV")(wrfnc)
v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
v = destagger(v_vars[varname], -2)
else: else:
try: varname = either("U10", "UU")
u_vars = extract_vars(wrfnc, timeidx, varnames="U10") u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
except KeyError: u = (u_vars[varname] if varname == "U10" else
try: destagger(u_vars[varname][...,0,:,:], -1))
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"]
try: varname = either("V10", "VV")
v_vars = extract_vars(wrfnc, timeidx, varnames="V10") v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
except KeyError: v = (v_vars[varname] if varname == "V10" else
try: destagger(v_vars[varname][...,0,:,:], -2))
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"]
map_proj_attrs = extract_global_attrs(wrfnc, attrs="MAP_PROJ") map_proj_attrs = extract_global_attrs(wrfnc, attrs="MAP_PROJ")
map_proj = map_proj_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"):
else: else:
cen_lon = lon_attrs["STAND_LON"] 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: varname = either("XLAT_M", "XLAT")(wrfnc)
xlon_m_vars = extract_vars(wrfnc, timeidx, varnames="XLONG_M") xlat_var = extract_vars(wrfnc, timeidx, varname,
except KeyError: method, squeeze, cache)
try: lat = xlat_var[varname]
xlon_vars = extract_vars(wrfnc, timeidx, varnames="XLONG")
except KeyError: varname = either("XLONG_M", "XLONG")
raise RuntimeError("xlong not found in NetCDF file") xlon_var = extract_vars(wrfnc, timeidx, varname,
else: method, squeeze, cache)
lon = xlon_vars["XLONG"] lon = xlon_var[varname]
else:
lon = xlon_m_vars["XLONG_M"]
if map_proj == 1: if map_proj == 1:
if((fabs(true_lat1 - true_lat2) > 0.1) and if((fabs(true_lat1 - true_lat2) > 0.1) and
(fabs(true_lat2 - 90.) > 0.1)): (fabs(true_lat2 - 90.) > 0.1)):
cone = (log(cos(true_lat1*radians_per_degree)) cone = (log(cos(true_lat1*radians_per_degree))
- log(cos(true_lat2*radians_per_degree))) - log(cos(true_lat2*radians_per_degree)))
cone = cone / (log(tan((45.-fabs(true_lat1/2.))*radians_per_degree)) cone = (cone /
- log(tan((45.-fabs(true_lat2/2.))*radians_per_degree))) (log(tan((45.-fabs(true_lat1/2.))*radians_per_degree))
- log(tan((45.-fabs(true_lat2/2.))*radians_per_degree))))
else: else:
cone = sin(fabs(true_lat1)*radians_per_degree) cone = sin(fabs(true_lat1)*radians_per_degree)
else: else:
@ -142,16 +97,27 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"):
return res 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) return get_uvmet(wrfnc, timeidx, True, units)
def get_uvmet_wspd_wdir(wrfnc, timeidx=0, units="mps"): @set_wind_metadata(wspd_wdir=True)
u,v = get_uvmet(wrfnc, timeidx, False, units) def get_uvmet_wspd_wdir(wrfnc, timeidx=0, units="mps",
return _calc_wspd_wdir(u, v, units) 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 @@
from wrf.var.extension import computeavo, computepvo from .extension import computeavo, computepvo
from wrf.var.util import extract_vars, extract_global_attrs from .util import extract_vars, extract_global_attrs
from .decorators import copy_and_set_metadata
__all__ = ["get_avo", "get_pvo"] __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", ncvars = extract_vars(wrfnc, timeidx, varnames=("U", "V", "MAPFAC_U",
"MAPFAC_V", "MAPFAC_M", "MAPFAC_V", "MAPFAC_M",
"F")) "F"),
method, squeeze, cache)
attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY")) attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
u = ncvars["U"] u = ncvars["U"]
@ -22,11 +27,15 @@ def get_avo(wrfnc, timeidx=0):
return computeavo(u,v,msfu,msfv,msfm,cor,dx,dy) 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", ncvars = extract_vars(wrfnc, timeidx, varnames=("U", "V", "T", "P",
"PB", "MAPFAC_U", "PB", "MAPFAC_U",
"MAPFAC_V", "MAPFAC_M", "MAPFAC_V", "MAPFAC_M",
"F")) "F"),
method, squeeze, cache)
attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY")) attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
u = ncvars["U"] u = ncvars["U"]

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

@ -1,42 +1,88 @@
import numpy as n import numpy as np
from wrf.var.constants import Constants from .constants import Constants
from wrf.var.destag import destagger_windcomp from .destag import destagger
from wrf.var.decorators import convert_units from .util import extract_vars, either
from .decorators import convert_units, set_wind_metadata
__all__ = ["get_u_destag", "get_v_destag", "get_w_destag", __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") @convert_units("wind", "mps")
def _calc_wspd(u, v, units="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): def _calc_wdir(u, v):
wdir = 270.0 - n.arctan2(v,u) * (180.0/Constants.PI) wdir = 270.0 - np.arctan2(v,u) * (180.0/Constants.PI)
return n.remainder(wdir, 360.0) return np.remainder(wdir, 360.0)
def _calc_wspd_wdir(u, v, units): def _calc_wspd_wdir(u, v, two_d, units):
return (_calc_wspd(u,v, units), _calc_wdir(u,v)) 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") @convert_units("wind", "mps")
def get_u_destag(wrfnc, timeidx=0, units="mps"): def get_u_destag(wrfnc, timeidx=0, units="mps",
u = destagger_windcomp(wrfnc,"u", timeidx) 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 return u
@set_wind_metadata(wspd_wdir=False)
@convert_units("wind", "mps") @convert_units("wind", "mps")
def get_v_destag(wrfnc, timeidx=0, units="mps"): def get_v_destag(wrfnc, timeidx=0, units="mps",
v = destagger_windcomp(wrfnc,"v", timeidx) 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 return v
@set_wind_metadata(wspd_wdir=False)
@convert_units("wind", "mps") @convert_units("wind", "mps")
def get_w_destag(wrfnc, timeidx=0, units="mps"): def get_w_destag(wrfnc, timeidx=0, units="mps",
w = destagger_windcomp(wrfnc,"w", timeidx) method="cat", squeeze=True, cache=None):
w_vars = extract_vars(wrfnc, timeidx, "W", method, squeeze, cache)
w = destagger(w_vars["W"], -3)
return w return w
def get_destag_wspd_wdir(wrfnc, timeidx=0, units="mps"): @set_wind_metadata(wspd_wdir=True)
u = destagger_windcomp(wrfnc,"u", timeidx) def get_destag_wspd_wdir(wrfnc, timeidx=0, units="mps",
v = destagger_windcomp(wrfnc,"v", timeidx) 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