Browse Source

Added thread local coordinate and is_moving caching so that this is only done once. Modified the latlon/xy routines to put the product on the left like NCL. Updated unit testing to add the latlon routines. Updated ipynb scripts.

main
Bill Ladwig 10 years ago
parent
commit
7e503a145b
  1. 60
      ncl_reference/WRFUserARW.ncl
  2. 5541
      ncl_reference/WRFUsersherrie.ncl
  3. 11
      src/wrf/cache.py
  4. 12
      src/wrf/cape.py
  5. 5
      src/wrf/cloudfrac.py
  6. 95
      src/wrf/config.py
  7. 10
      src/wrf/ctt.py
  8. 14
      src/wrf/dbz.py
  9. 9
      src/wrf/dewpoint.py
  10. 19
      src/wrf/geoht.py
  11. 17
      src/wrf/helicity.py
  12. 26
      src/wrf/interp.py
  13. 29
      src/wrf/latlon.py
  14. 77
      src/wrf/latlonutils.py
  15. 47
      src/wrf/metadecorators.py
  16. 4
      src/wrf/omega.py
  17. 13
      src/wrf/pressure.py
  18. 4587
      src/wrf/psadlookup.py
  19. 4
      src/wrf/pw.py
  20. 8
      src/wrf/rh.py
  21. 15
      src/wrf/routines.py
  22. 4
      src/wrf/slp.py
  23. 30
      src/wrf/temp.py
  24. 5
      src/wrf/terrain.py
  25. 13
      src/wrf/times.py
  26. 651
      src/wrf/util.py
  27. 32
      src/wrf/uvmet.py
  28. 2
      src/wrf/version.py
  29. 8
      src/wrf/vorticity.py
  30. 24
      src/wrf/wind.py
  31. 2
      test/ipynb/WRF_Workshop_Demo.ipynb
  32. 25
      test/ipynb/WRF_python_demo.ipynb
  33. 26
      test/ncl_get_var.ncl
  34. 58
      test/test_multi_cache.py
  35. 120
      test/utests.py

60
ncl_reference/WRFUserARW.ncl

@ -1378,64 +1378,18 @@ begin @@ -1378,64 +1378,18 @@ begin
end if
if( any( variable .eq. "cfrac" )) then
print("calculating wrf cloud fraction")
if(typeof(file_handle).eq."file") then
ISFILE = True
nc_file = file_handle
else if(typeof(file_handle).eq."list") then
ISFILE = False
nc_file = file_handle[0]
else
print("wrf_user_getvar: error: the first argument must be a file or a list of files opened with addfile or addfiles")
return
end if
end if
time_in = time
if ( time .eq. -1 ) then
if(ISFILE) then
P = nc_file->P
PB = nc_file->PB
rh = wrf_user_getvar(nc_file,"rh",-1)
else
P = file_handle[:]->P
PB = file_handle[:]->PB
rh = wrf_user_getvar(file_handle,"rh",-1)
end if
else
if(ISFILE) then
P = nc_file->P(time_in,:,:,:)
PB = nc_file->PB(time_in,:,:,:)
rh = wrf_user_getvar(file_handle,"rh",time_in)
else
P = file_handle[:]->P(time_in,:,:,:)
PB = file_handle[:]->PB(time_in,:,:,:)
rh = wrf_user_getvar(file_handle,"rh",time_in)
end if
end if
P = _get_wrf_var(file_handle,"P",time)
PB = _get_wrf_var(file_handle,"PB",time)
rh = wrf_user_getvar(file_handle,"rh",time)
pres = (P + PB)
thedims = dimsizes(pres)
if(time .eq. -1) then
nt = thedims(0)
nz = thedims(1)
ns = thedims(2)
ew = thedims(3)
else
nt = 1
nz = thedims(0)
ns = thedims(1)
ew = thedims(2)
end if
out_array = wrf_cloud_fraction(pres,rh)
copy_VarAtts_except(rh,out_array,(/"MemoryOrder","description","units"/))
out_array@units = "%"
out_array@description = "Low, Mid, High Clouds"
out_array@MemoryOrder = "XY"
return(out_array)
end if ;variable is cfrac

5541
ncl_reference/WRFUsersherrie.ncl

File diff suppressed because it is too large Load Diff

11
src/wrf/cache.py

@ -11,6 +11,9 @@ _local_storage = local() @@ -11,6 +11,9 @@ _local_storage = local()
def cache_item(key, product, value):
global _local_storage
if key is None:
return
try:
cache = _local_storage.cache
except AttributeError:
@ -30,6 +33,9 @@ def cache_item(key, product, value): @@ -30,6 +33,9 @@ def cache_item(key, product, value):
def get_cached_item(key, product):
if key is None:
return None
cache = getattr(_local_storage, "cache", None)
if cache is None:
@ -40,8 +46,9 @@ def get_cached_item(key, product): @@ -40,8 +46,9 @@ def get_cached_item(key, product):
if prod_dict is None:
return None
return prod_dict.get(product, None)
result = prod_dict.get(product, None)
return result
def _get_cache():
return getattr(_local_storage, "cache", None)

12
src/wrf/cape.py

@ -21,15 +21,13 @@ from .metadecorators import set_cape_metadata @@ -21,15 +21,13 @@ from .metadecorators import set_cape_metadata
# units="J/kg ; J/kg ; m ; m",
# MemoryOrder="XY")
@set_cape_metadata(is2d=True)
def get_2dcape(wrfnc, timeidx=0, method="cat",
squeeze=True, cache=None, meta=True,
missing=Constants.DEFAULT_FILL):
def get_2dcape(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True, _key=None, missing=Constants.DEFAULT_FILL):
"""Return the 2d fields of cape, cin, lcl, and lfc"""
varnames = ("T", "P", "PB", "QVAPOR", "PH","PHB", "HGT", "PSFC")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
t = ncvars["T"]
p = ncvars["P"]
@ -86,11 +84,11 @@ def get_2dcape(wrfnc, timeidx=0, method="cat", @@ -86,11 +84,11 @@ def get_2dcape(wrfnc, timeidx=0, method="cat",
@set_cape_metadata(is2d=False)
def get_3dcape(wrfnc, timeidx=0, method="cat",
squeeze=True, cache=None, meta=True,
missing=Constants.DEFAULT_FILL):
_key=None, missing=Constants.DEFAULT_FILL):
"""Return the 3d fields of cape and cin"""
varnames = ("T", "P", "PB", "QVAPOR", "PH", "PHB", "HGT", "PSFC")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
t = ncvars["T"]
p = ncvars["P"]
pb = ncvars["PB"]

5
src/wrf/cloudfrac.py

@ -8,10 +8,11 @@ from .util import extract_vars @@ -8,10 +8,11 @@ from .util import extract_vars
@set_cloudfrac_metadata()
def get_cloudfrac(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True):
cache=None, meta=True, _key=None):
vars = extract_vars(wrfnc, timeidx, ("P", "PB", "QVAPOR", "T"),
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False,
_key=_key)
p = vars["P"]
pb = vars["PB"]

95
src/wrf/config.py

@ -1,90 +1,103 @@ @@ -1,90 +1,103 @@
from __future__ import (absolute_import, division, print_function,
unicode_literals)
from threading import local
_local_config = local()
_local_config.xarray_enabled = True
_local_config.cartopy_enabled = True
_local_config.basemap_enabled = True
_local_config.pyngl_enabled = True
_local_config.cache_size = 20
try:
from xarray import DataArray
except ImportError:
_xarray_enabled = False
else:
_xarray_enabled = True
_local_config.xarray_enabled = False
try:
from cartopy import crs
except ImportError:
_cartopy_enabled = False
else:
_cartopy_enabled = True
_local_config.cartopy_enabled = False
try:
from mpl_toolkits.basemap import Basemap
except ImportError:
_basemap_enabled = False
else:
_basemap_enabled = True
_local_config.basemap_enabled = False
try:
from Ngl import Resources
except ImportError:
_pyngl_enabled = False
else:
_pyngl_enabled = True
_cache_size = 5
_local_config.pyngl_enabled = False
def xarray_enabled():
global _xarray_enabled
return _xarray_enabled
global _local_config
return _local_config.xarray_enabled
def disable_xarray():
global _xarray_enabled
_xarray_enabled = False
global _local_config
_local_config.xarray_enabled = False
def enable_xarray():
global _xarray_enabled
_xarray_enabled = True
global _local_config
_local_config.xarray_enabled = True
def cartopy_enabled():
global _cartopy_enabled
return _cartopy_enabled
global _local_config
return _local_config.cartopy_enabled
def enable_cartopy():
global _cartopy_enabled
_cartopy_enabled = True
global _local_config
_local_config.cartopy_enabled = True
def disable_cartopy():
global _cartopy_enabled
_cartopy_enabled = True
global _local_config
_local_config.cartopy_enabled = True
def basemap_enabled():
global _basemap_enabled
return _basemap_enabled
global _local_config
return _local_config.basemap_enabled
def enable_basemap():
global _basemap_enabled
_basemap_enabled = True
global _local_config
_local_config.basemap_enabled = True
def disable_basemap():
global _basemap_enabled
_basemap_enabled = True
global _local_config
_local_config.basemap_enabled = True
def pyngl_enabled():
global _pyngl_enabled
return _pyngl_enabled
global _local_config
return _local_config.pyngl_enabled
def enable_pyngl():
global _pyngl_enabled
_pyngl_enabled = True
global _local_config
_local_config.pyngl_enabled = True
def disable_pyngl():
global _pyngl_enabled
_pyngl_enabled = True
global _local_config
_local_config.pyngl_enabled = True
def set_cache_size(size):
global _cache_size
_cache_size = size
global _local_config
_local_config.cache_size = size
def get_cache_size():
return int(_cache_size)
global _local_config
return int(_local_config.cache_size)

10
src/wrf/ctt.py

@ -18,14 +18,14 @@ from .util import extract_vars @@ -18,14 +18,14 @@ from .util import extract_vars
MemoryOrder="XY")
@convert_units("temp", "c")
def get_ctt(wrfnc, timeidx=0, method="cat",
squeeze=True, cache=None, meta=True,
squeeze=True, cache=None, meta=True, _key=None,
units="c"):
"""Return the cloud top temperature.
"""
varnames = ("T", "P", "PB", "PH", "PHB", "HGT", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
t = ncvars["T"]
p = ncvars["P"]
pb = ncvars["PB"]
@ -37,7 +37,8 @@ def get_ctt(wrfnc, timeidx=0, method="cat", @@ -37,7 +37,8 @@ def get_ctt(wrfnc, timeidx=0, method="cat",
haveqci = 1
try:
icevars = extract_vars(wrfnc, timeidx, "QICE",
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False,
_key=_key)
except KeyError:
qice = n.zeros(qv.shape, qv.dtype)
haveqci = 0
@ -46,7 +47,8 @@ def get_ctt(wrfnc, timeidx=0, method="cat", @@ -46,7 +47,8 @@ def get_ctt(wrfnc, timeidx=0, method="cat",
try:
cldvars = extract_vars(wrfnc, timeidx, "QCLOUD",
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False,
_key=_key)
except KeyError:
raise RuntimeError("'QCLOUD' not found in NetCDF file")
else:

14
src/wrf/dbz.py

@ -14,7 +14,7 @@ from .metadecorators import copy_and_set_metadata @@ -14,7 +14,7 @@ from .metadecorators import copy_and_set_metadata
description="radar reflectivity",
units="dBz")
def get_dbz(wrfnc, timeidx=0, method="cat",
squeeze=True, cache=None, meta=True,
squeeze=True, cache=None, meta=True, _key=None,
use_varint=False, use_liqskin=False):
""" Return the dbz
@ -28,7 +28,7 @@ def get_dbz(wrfnc, timeidx=0, method="cat", @@ -28,7 +28,7 @@ def get_dbz(wrfnc, timeidx=0, method="cat",
"""
varnames = ("T", "P", "PB", "QVAPOR", "QRAIN")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
t = ncvars["T"]
p = ncvars["P"]
pb = ncvars["PB"]
@ -37,7 +37,8 @@ def get_dbz(wrfnc, timeidx=0, method="cat", @@ -37,7 +37,8 @@ def get_dbz(wrfnc, timeidx=0, method="cat",
try:
snowvars = extract_vars(wrfnc, timeidx, "QSNOW",
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False,
_key=_key)
except KeyError:
qs = np.zeros(qv.shape, qv.dtype)
else:
@ -45,7 +46,8 @@ def get_dbz(wrfnc, timeidx=0, method="cat", @@ -45,7 +46,8 @@ def get_dbz(wrfnc, timeidx=0, method="cat",
try:
graupvars = extract_vars(wrfnc, timeidx, "QGRAUP",
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False,
_key=_key)
except KeyError:
qg = np.zeros(qv.shape, qv.dtype)
else:
@ -69,9 +71,9 @@ def get_dbz(wrfnc, timeidx=0, method="cat", @@ -69,9 +71,9 @@ def get_dbz(wrfnc, timeidx=0, method="cat",
units="dBz",
MemoryOrder="XY")
def get_max_dbz(wrfnc, timeidx=0, method="cat",
squeeze=True, cache=None, meta=True,
squeeze=True, cache=None, meta=True, _key=None,
do_varint=False, do_liqskin=False):
return np.amax(get_dbz(wrfnc, timeidx, method, squeeze, cache, meta,
do_varint, do_liqskin),
_key, do_varint, do_liqskin),
axis=-3)

9
src/wrf/dewpoint.py

@ -12,11 +12,11 @@ from .util import extract_vars @@ -12,11 +12,11 @@ from .util import extract_vars
description="dew point temperature")
@convert_units("temp", "c")
def get_dp(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True, units="c"):
cache=None, meta=True, _key=None, units="c"):
varnames=("P", "PB", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
p = ncvars["P"]
pb = ncvars["PB"]
@ -33,11 +33,10 @@ def get_dp(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -33,11 +33,10 @@ def get_dp(wrfnc, timeidx=0, method="cat", squeeze=True,
description="2m dew point temperature")
@convert_units("temp", "c")
def get_dp_2m(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
units="c"):
cache=None, meta=True, _key=None, units="c"):
varnames=("PSFC", "Q2")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
# Algorithm requires hPa
psfc = .01*(ncvars["PSFC"])

19
src/wrf/geoht.py

@ -8,7 +8,7 @@ from .metadecorators import set_height_metadata @@ -8,7 +8,7 @@ from .metadecorators import set_height_metadata
from .util import extract_vars, either
def _get_geoht(wrfnc, timeidx, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
height=True, msl=True):
"""Return the geopotential in units of m2 s-2 if height is False,
otherwise return the geopotential height in meters. If height is True,
@ -20,7 +20,8 @@ def _get_geoht(wrfnc, timeidx, method="cat", squeeze=True, @@ -20,7 +20,8 @@ def _get_geoht(wrfnc, timeidx, method="cat", squeeze=True,
varname = either("PH", "GHT")(wrfnc)
if varname == "PH":
ph_vars = extract_vars(wrfnc, timeidx, ("PH", "PHB", "HGT"),
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False,
_key=_key)
ph = ph_vars["PH"]
phb = ph_vars["PHB"]
hgt = ph_vars["HGT"]
@ -28,7 +29,8 @@ def _get_geoht(wrfnc, timeidx, method="cat", squeeze=True, @@ -28,7 +29,8 @@ def _get_geoht(wrfnc, timeidx, method="cat", squeeze=True,
geopt_unstag = destagger(geopt, -3)
else:
ght_vars = extract_vars(wrfnc, timeidx, ("GHT", "HGT_M"),
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False,
_key=_key)
geopt_unstag = ght_vars["GHT"] * Constants.G
hgt = ght_vars["HGT_M"]
@ -50,16 +52,17 @@ def _get_geoht(wrfnc, timeidx, method="cat", squeeze=True, @@ -50,16 +52,17 @@ def _get_geoht(wrfnc, timeidx, method="cat", squeeze=True,
@set_height_metadata(geopt=True)
def get_geopt(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True):
return _get_geoht(wrfnc, timeidx, method, squeeze, cache, meta,
False, True,)
meta=True, _key=None):
return _get_geoht(wrfnc, timeidx, method, squeeze, cache, meta, _key,
False, True)
@set_height_metadata(geopt=False)
@convert_units("height", "m")
def get_height(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
msl=True, units="m"):
return _get_geoht(wrfnc, timeidx, method, squeeze, cache, meta, True, msl)
return _get_geoht(wrfnc, timeidx, method, squeeze, cache, meta, _key,
True, msl)

17
src/wrf/helicity.py

@ -13,11 +13,12 @@ from .metadecorators import copy_and_set_metadata @@ -13,11 +13,12 @@ from .metadecorators import copy_and_set_metadata
description="storm relative helicity",
units="m-2/s-2")
def get_srh(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True, top=3000.0):
cache=None, meta=True, _key=None, top=3000.0):
# Top can either be 3000 or 1000 (for 0-1 srh or 0-3 srh)
ncvars = extract_vars(wrfnc, timeidx, ("HGT", "PH", "PHB"),
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False,
_key=_key)
ter = ncvars["HGT"]
ph = ncvars["PH"]
@ -26,12 +27,12 @@ def get_srh(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -26,12 +27,12 @@ def get_srh(wrfnc, timeidx=0, method="cat", squeeze=True,
# As coded in NCL, but not sure this is possible
varname = either("U", "UU")(wrfnc)
u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
u = destagger(u_vars[varname], -1)
varname = either("V", "VV")(wrfnc)
v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
v = destagger(v_vars[varname], -2)
geopt = ph + phb
@ -52,11 +53,11 @@ def get_srh(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -52,11 +53,11 @@ def get_srh(wrfnc, timeidx=0, method="cat", squeeze=True,
description="updraft helicity",
units="m-2/s-2")
def get_uh(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
bottom=2000.0, top=5000.0):
ncvars = extract_vars(wrfnc, timeidx, ("W", "PH", "PHB", "MAPFAC_M"),
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False, _key=None)
wstag = ncvars["W"]
ph = ncvars["PH"]
@ -70,12 +71,12 @@ def get_uh(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -70,12 +71,12 @@ def get_uh(wrfnc, timeidx=0, method="cat", squeeze=True,
# As coded in NCL, but not sure this is possible
varname = either("U", "UU")(wrfnc)
u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
u = destagger(u_vars[varname], -1)
varname = either("V", "VV")(wrfnc)
v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
v = destagger(v_vars[varname], -2)
zp = ph + phb

26
src/wrf/interp.py

@ -12,7 +12,7 @@ from .extension import (_interpz3d, _interp2dxy, _interp1d, _vertcross, @@ -12,7 +12,7 @@ from .extension import (_interpz3d, _interp2dxy, _interp1d, _vertcross,
_interpline, _smooth2d, _monotonic, _vintrp)
from .metadecorators import set_interp_metadata
from .util import extract_vars, is_staggered
from .util import extract_vars, is_staggered, get_id
from .interputils import get_xy, get_xy_z_params
from .constants import Constants, ConversionFactors
from .terrain import get_terrain
@ -284,6 +284,7 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, @@ -284,6 +284,7 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False,
`numpy.ndarray` object with no metadata.
"""
_key = get_id(wrfnc)
# Remove case sensitivity
field_type = field_type.lower() if field_type is not None else "none"
@ -344,24 +345,30 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, @@ -344,24 +345,30 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False,
# Extract vriables
#timeidx = -1 # Should this be an argument?
ncvars = extract_vars(wrfnc, timeidx, ("PSFC", "QVAPOR", "F"),
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False, _key=_key)
sfp = ncvars["PSFC"] * ConversionFactors.PA_TO_HPA
qv = ncvars["QVAPOR"]
coriolis = ncvars["F"]
terht = get_terrain(wrfnc, timeidx, units="m",
method=method, squeeze=squeeze, cache=cache)
method=method, squeeze=squeeze, cache=cache,
meta=False, _key=_key)
t = get_theta(wrfnc, timeidx, units="k",
method=method, squeeze=squeeze, cache=cache)
method=method, squeeze=squeeze, cache=cache,
meta=False, _key=_key)
tk = get_temp(wrfnc, timeidx, units="k",
method=method, squeeze=squeeze, cache=cache)
method=method, squeeze=squeeze, cache=cache,
meta=False, _key=_key)
p = get_pressure(wrfnc, timeidx, units="pa",
method=method, squeeze=squeeze, cache=cache)
method=method, squeeze=squeeze, cache=cache,
meta=False, _key=_key)
ght = get_height(wrfnc, timeidx, msl=True, units="m",
method=method, squeeze=squeeze, cache=cache)
method=method, squeeze=squeeze, cache=cache,
meta=False, _key=_key)
ht_agl = get_height(wrfnc, timeidx, msl=False, units="m",
method=method, squeeze=squeeze, cache=cache)
method=method, squeeze=squeeze, cache=cache,
meta=False, _key=_key)
smsfp = _smooth2d(sfp, 3)
@ -400,7 +407,8 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, @@ -400,7 +407,8 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False,
idir = 1
delta = 0.01
eth = get_eth(wrfnc, timeidx)
eth = get_eth(wrfnc, timeidx, method=method, squeeze=squeeze,
cache=cache, meta=False, _key=_key)
p_hpa = p * ConversionFactors.PA_TO_HPA

29
src/wrf/latlon.py

@ -1,58 +1,63 @@ @@ -1,58 +1,63 @@
from __future__ import (absolute_import, division, print_function,
unicode_literals)
from .util import extract_vars
from .util import extract_vars, get_id
from .latlonutils import (_lat_varname, _lon_varname, _ll_to_xy, _xy_to_ll)
from .metadecorators import set_latlon_metadata
def get_lat(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
stagger=None):
varname = _lat_varname(wrfnc, stagger)
lat_var = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,
meta)
meta, _key)
return lat_var[varname]
def get_lon(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
stagger=None):
varname = _lon_varname(wrfnc, stagger)
lon_var = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,
meta)
meta, _key)
return lon_var[varname]
# TODO: Do we need the user to know about method, squeeze, cache for this?
# Can either use wrfnc as a single file or sequence, or provide
# projection parameters (which don't allow for moving domains)
@set_latlon_metadata(xy=True)
def ll_to_xy(wrfnc, latitude, longitude, timeidx=0, stagger=None, method="cat",
squeeze=True, cache=None, meta=True):
squeeze=True, cache=None, meta=True, as_int=True):
_key = get_id(wrfnc)
return _ll_to_xy(latitude, longitude, wrfnc, timeidx, stagger, method,
squeeze, cache, **{})
squeeze, cache, _key, as_int, **{})
@set_latlon_metadata(xy=True)
def ll_to_xy_proj(latitude, longitude, meta=True, squeeze=True, **projparams):
return _ll_to_xy(latitude, longitude, None, 0, squeeze, "cat", True, None,
**projparams)
def ll_to_xy_proj(latitude, longitude, meta=True, squeeze=True, as_int=True,
**projparams):
return _ll_to_xy(latitude, longitude, None, 0, squeeze, "cat", True, None,
None, as_int, **projparams)
@set_latlon_metadata(xy=False)
def xy_to_ll(wrfnc, x, y, timeidx=0, stagger=None, method="cat", squeeze=True,
cache=None, meta=True):
_key = get_id(wrfnc)
return _xy_to_ll(x, y, wrfnc, timeidx, stagger, method, squeeze, cache,
**{})
_key, **{})
@set_latlon_metadata(xy=False)
def xy_to_ll_proj(x, y, meta=True, squeeze=True, **projparams):
return _xy_to_ll(x, y, None, 0, None, "cat", squeeze, None, **projparams)
return _xy_to_ll(x, y, None, 0, None, "cat", squeeze, None, None,
**projparams)

77
src/wrf/latlonutils.py

@ -32,7 +32,7 @@ def _lon_varname(wrfnc, stagger): @@ -32,7 +32,7 @@ def _lon_varname(wrfnc, stagger):
return varname
def _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache):
def _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache, _key):
if timeidx < 0:
raise ValueError("'timeidx' must be greater than 0")
@ -75,14 +75,16 @@ def _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache): @@ -75,14 +75,16 @@ def _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache):
if not is_mapping(wrfnc):
wrfnc = next(iter(wrfnc)) # only need one file
else:
wrfnc = wrfnc[next(iter(viewkeys(wrfnc)))]
first_entry = next(iter(viewkeys(wrfnc)))
wrfnc = wrfnc[first_entry]
key = _key[first_entry]
return _get_proj_params(wrfnc, timeidx, stagger,
method, squeeze, cache)
method, squeeze, cache, key)
xlat = extract_vars(wrfnc, lat_timeidx, (latvar,), method, squeeze, cache,
meta=False)[latvar]
meta=False, _key=_key)[latvar]
xlon = extract_vars(wrfnc, lat_timeidx, (lonvar,), method, squeeze, cache,
meta=False)[lonvar]
meta=False, _key=_key)[lonvar]
ref_lat = np.ravel(xlat[..., 0, 0])
ref_lon = np.ravel(xlon[..., 0, 0])
@ -164,13 +166,13 @@ def _kwarg_proj_params(projparams): @@ -164,13 +166,13 @@ def _kwarg_proj_params(projparams):
# Will return 0-based indexes
def _ll_to_xy(latitude, longitude, wrfnc=None, timeidx=0,
stagger=None, method="cat", squeeze=True, cache=None,
**projparms):
_key=None, as_int=True, **projparms):
if wrfnc is not None:
(map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
pole_lat, pole_lon, known_x, known_y, dx, dy, latinc,
loninc) = _get_proj_params(wrfnc, timeidx, stagger, method, squeeze,
cache)
cache, _key)
else:
(map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
pole_lat, pole_lon, known_x, known_y, dx, dy, latinc,
@ -191,18 +193,22 @@ def _ll_to_xy(latitude, longitude, wrfnc=None, timeidx=0, @@ -191,18 +193,22 @@ def _ll_to_xy(latitude, longitude, wrfnc=None, timeidx=0,
"must be the same length")
if ref_lat.size == 1:
outdim = [lats.size, 2]
extra_dims = [outdim[0]]
outdim = [2, lats.size]
#outdim = [lats.size, 2]
extra_dims = [outdim[1]]
else:
# Moving domain will have moving ref_lats/ref_lons
outdim = [lats.size, ref_lat.size, 2]
extra_dims = outdim[0:2]
outdim = [2, ref_lat.size, lats.size]
#outdim = [lats.size, ref_lat.size, 2]
extra_dims = outdim[1:]
result = np.empty(outdim, np.float64)
for left_idxs in iter_left_indexes(extra_dims):
left_and_slice_idxs = left_idxs + (slice(None), )
#left_and_slice_idxs = left_idxs + (slice(None), )
# Left indexes is a misnomer, since these will be on the right
x_idxs = (0,) + left_idxs
y_idxs = (1,) + left_idxs
if ref_lat.size == 1:
ref_lat_val = ref_lat[0]
ref_lon_val = ref_lon[0]
@ -218,29 +224,41 @@ def _ll_to_xy(latitude, longitude, wrfnc=None, timeidx=0, @@ -218,29 +224,41 @@ def _ll_to_xy(latitude, longitude, wrfnc=None, timeidx=0,
known_x, known_y, dx, dy, latinc, loninc,
lat, lon)
result[left_and_slice_idxs] = xy[:]
# Note: comes back from fortran as y,x
result[x_idxs] = xy[1]
result[y_idxs] = xy[0]
else:
result = np.empty((2,), np.float64)
result = _lltoxy(map_proj, truelat1, truelat2, stdlon,
fort_out = _lltoxy(map_proj, truelat1, truelat2, stdlon,
ref_lat, ref_lon, pole_lat, pole_lon,
known_x, known_y, dx, dy, latinc, loninc,
latitude, longitude)
# Note, comes back from fortran as y,x. So, need to swap them.
result[0] = fort_out[1]
result[1] = fort_out[0]
# Make indexes 0-based
result = result - 1
if as_int:
result = np.rint(result).astype(int)
return result
# X and Y should be 0-based
def _xy_to_ll(x, y, wrfnc=None, timeidx=0, stagger=None,
method="cat", squeeze=True, cache=None, **projparams):
method="cat", squeeze=True, cache=None, _key=None,
**projparams):
if wrfnc is not None:
(map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
pole_lat, pole_lon, known_x, known_y, dx, dy, latinc,
loninc) = _get_proj_params(wrfnc, timeidx, stagger, method, squeeze,
cache)
cache, _key)
else:
(map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
pole_lat, pole_lon, known_x, known_y, dx, dy, latinc,
@ -262,21 +280,26 @@ def _xy_to_ll(x, y, wrfnc=None, timeidx=0, stagger=None, @@ -262,21 +280,26 @@ def _xy_to_ll(x, y, wrfnc=None, timeidx=0, stagger=None,
y_arr = y_arr.ravel()
if (x_arr.size != y_arr.size):
raise ValueError("'x' and 'y' "
"must be the same length")
raise ValueError("'x' and 'y' must be the same length")
if ref_lat.size == 1:
outdim = [x_arr.size, 2]
extra_dims = [outdim[0]]
#outdim = [x_arr.size, 2]
#extra_dims = [outdim[0]]
outdim = [2, x_arr.size]
extra_dims = [outdim[1]]
else:
# Moving domain will have moving ref_lats/ref_lons
outdim = [x_arr.size, ref_lat.size, 2]
extra_dims = outdim[0:2]
#outdim = [x_arr.size, ref_lat.size, 2]
#extra_dims = outdim[0:2]
outdim = [2, ref_lat.size, x_arr.size]
extra_dims = outdim[1:]
result = np.empty(outdim, np.float64)
for left_idxs in iter_left_indexes(extra_dims):
left_and_slice_idxs = left_idxs + (slice(None), )
#left_and_slice_idxs = left_idxs + (slice(None), )
lat_idxs = (0,) + left_idxs
lon_idxs = (1,) + left_idxs
if ref_lat.size == 1:
ref_lat_val = ref_lat[0]
@ -292,7 +315,9 @@ def _xy_to_ll(x, y, wrfnc=None, timeidx=0, stagger=None, @@ -292,7 +315,9 @@ def _xy_to_ll(x, y, wrfnc=None, timeidx=0, stagger=None,
ref_lon_val, pole_lat, pole_lon, known_x, known_y,
dx, dy, latinc, loninc, x_val, y_val)
result[left_and_slice_idxs] = ll[:]
#result[left_and_slice_idxs] = ll[:]
result[lat_idxs] = ll[0]
result[lon_idxs] = ll[1]
else:
# Convert 0-based to 1-based for Fortran

47
src/wrf/metadecorators.py

@ -39,7 +39,8 @@ def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, @@ -39,7 +39,8 @@ def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None,
return wrapped(*args, **kwargs)
argvars = from_args(wrapped, ("wrfnc", "timeidx", "method",
"squeeze", "cache", "units", "meta"),
"squeeze", "cache", "units", "meta",
"_key"),
*args, **kwargs)
wrfnc = argvars["wrfnc"]
@ -48,6 +49,7 @@ def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, @@ -48,6 +49,7 @@ def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None,
method = argvars["method"]
squeeze = argvars["squeeze"]
cache = argvars["cache"]
_key = argvars["_key"]
if cache is None:
cache = {}
@ -65,7 +67,7 @@ def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, @@ -65,7 +67,7 @@ def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None,
if var_to_copy is None:
var_to_copy = extract_vars(wrfnc, timeidx, (_copy_varname,),
method, squeeze, cache,
meta=True)[_copy_varname]
meta=True, _key=_key)[_copy_varname]
# Make a copy so we don't modify a user supplied cache
new_cache = dict(cache)
@ -148,7 +150,8 @@ def set_wind_metadata(copy_varname, name, description, @@ -148,7 +150,8 @@ def set_wind_metadata(copy_varname, name, description,
return wrapped(*args, **kwargs)
argvars = from_args(wrapped, ("wrfnc", "timeidx", "units",
"method", "squeeze", "ten_m", "cache"),
"method", "squeeze", "ten_m", "cache",
"_key"),
*args, **kwargs)
wrfnc = argvars["wrfnc"]
timeidx = argvars["timeidx"]
@ -157,6 +160,7 @@ def set_wind_metadata(copy_varname, name, description, @@ -157,6 +160,7 @@ def set_wind_metadata(copy_varname, name, description,
squeeze = argvars["squeeze"]
ten_m = argvars["ten_m"]
cache = argvars["cache"]
_key = argvars["_key"]
if cache is None:
cache = {}
@ -167,7 +171,7 @@ def set_wind_metadata(copy_varname, name, description, @@ -167,7 +171,7 @@ def set_wind_metadata(copy_varname, name, description,
copy_var = extract_vars(wrfnc, timeidx, _copy_varname,
method, squeeze, cache,
meta=True)[_copy_varname]
meta=True, _key=_key)[_copy_varname]
# Make a copy so we don't modify a user supplied cache
new_cache = dict(cache)
@ -242,7 +246,7 @@ def set_cape_metadata(is2d): @@ -242,7 +246,7 @@ def set_cape_metadata(is2d):
return wrapped(*args, **kwargs)
argvars = from_args(wrapped, ("wrfnc", "timeidx", "method", "squeeze",
"cache", "missing"),
"cache", "_key", "missing"),
*args, **kwargs)
wrfnc = argvars["wrfnc"]
timeidx = argvars["timeidx"]
@ -250,12 +254,13 @@ def set_cape_metadata(is2d): @@ -250,12 +254,13 @@ def set_cape_metadata(is2d):
squeeze = argvars["squeeze"]
cache = argvars["cache"]
missing = argvars["missing"]
_key = argvars["_key"]
if cache is None:
cache = {}
_copy_varname = "P"
copy_var = extract_vars(wrfnc, timeidx, _copy_varname, method, squeeze,
cache, meta=True)[_copy_varname]
cache, meta=True, _key=_key)[_copy_varname]
# Make a copy so we don't modify a user supplied cache
new_cache = dict(cache)
@ -335,19 +340,20 @@ def set_cloudfrac_metadata(): @@ -335,19 +340,20 @@ def set_cloudfrac_metadata():
return wrapped(*args, **kwargs)
argvars = from_args(wrapped, ("wrfnc", "timeidx", "method", "squeeze",
"cache"),
"cache", "_key"),
*args, **kwargs)
wrfnc = argvars["wrfnc"]
timeidx = argvars["timeidx"]
method = argvars["method"]
squeeze = argvars["squeeze"]
cache = argvars["cache"]
_key = argvars["_key"]
if cache is None:
cache = {}
_copy_varname = "P"
copy_var = extract_vars(wrfnc, timeidx, _copy_varname, method, squeeze,
cache, meta=True)[_copy_varname]
cache, meta=True, _key=_key)[_copy_varname]
# Make a copy so we don't modify a user supplied cache
new_cache = dict(cache)
@ -416,18 +422,17 @@ def set_latlon_metadata(xy=False): @@ -416,18 +422,17 @@ def set_latlon_metadata(xy=False):
# Want to preserve the input coordinate pair in metadata
if result.ndim == 1:
result = result[np.newaxis, :]
result = result[:, np.newaxis]
argnames = ["x", "y"] if not xy else ["latitude", "longitude"]
argnames.append("squeeze")
outname = "latlon" if not xy else "xy"
if result.ndim == 2:
dimnames = (["idx", "lat_lon"] if not xy
else ["idx", "x_y"])
dimnames = (["lat_lon", "idx"] if not xy else ["x_y", "idx"])
else:
dimnames = (["idx", "domain", "lat_lon"] if not xy
else ["idx", "domain", "x_y"])
dimnames = (["lat_lon", "domain_idx", "idx"] if not xy
else ["x_y", "domain_idx", "idx"])
argvars = from_args(wrapped, argnames, *args, **kwargs)
@ -440,13 +445,14 @@ def set_latlon_metadata(xy=False): @@ -440,13 +445,14 @@ def set_latlon_metadata(xy=False):
coords = {}
if not xy:
coords["xy_coord"] = (dimnames[0], [CoordPair(x=x[0], y=x[1])
coords["xy_coord"] = (dimnames[-1], [CoordPair(x=x[0], y=x[1])
for x in zip(arr1, arr2)])
coords[dimnames[-1]] = ["lat", "lon"]
coords[dimnames[0]] = ["lat", "lon"]
else:
coords["latlon_coord"] = (dimnames[0], [CoordPair(lat=x[0], lon=x[1])
coords["latlon_coord"] = (dimnames[-1], [CoordPair(lat=x[0],
lon=x[1])
for x in zip(arr1, arr2)])
coords[dimnames[-1]] = ["x", "y"]
coords[dimnames[0]] = ["x", "y"]
da = DataArray(result, name=outname, dims=dimnames, coords=coords)
@ -469,7 +475,8 @@ def set_height_metadata(geopt=False): @@ -469,7 +475,8 @@ def set_height_metadata(geopt=False):
return wrapped(*args, **kwargs)
argvars = from_args(wrapped, ("wrfnc", "timeidx", "method",
"squeeze", "units", "msl", "cache"),
"squeeze", "units", "msl", "cache",
"_key"),
*args, **kwargs)
wrfnc = argvars["wrfnc"]
timeidx = argvars["timeidx"]
@ -478,6 +485,7 @@ def set_height_metadata(geopt=False): @@ -478,6 +485,7 @@ def set_height_metadata(geopt=False):
squeeze = argvars["squeeze"]
msl = argvars["msl"]
cache = argvars["cache"]
_key = argvars["_key"]
if cache is None:
cache = {}
@ -486,7 +494,8 @@ def set_height_metadata(geopt=False): @@ -486,7 +494,8 @@ def set_height_metadata(geopt=False):
# 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, meta=True)
method, squeeze, cache, meta=True,
_key=_key)
ht_metadata_var = ht_var[ht_metadata_varname]
# Make a copy so we don't modify a user supplied cache

4
src/wrf/omega.py

@ -13,10 +13,10 @@ from .metadecorators import copy_and_set_metadata @@ -13,10 +13,10 @@ from .metadecorators import copy_and_set_metadata
description="omega",
units="Pa/s")
def get_omega(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True):
meta=True, _key=None):
varnames=("T", "P", "W", "PB", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
t = ncvars["T"]
p = ncvars["P"]
w = ncvars["W"]

13
src/wrf/pressure.py

@ -10,26 +10,29 @@ from .util import extract_vars, either @@ -10,26 +10,29 @@ from .util import extract_vars, either
description="pressure")
@convert_units("pressure", "pa")
def get_pressure(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="pa"):
varname = either("P", "PRES")(wrfnc)
if varname == "P":
p_vars = extract_vars(wrfnc, timeidx, ("P", "PB"),
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False,
_key=_key)
p = p_vars["P"]
pb = p_vars["PB"]
pres = p + pb
else:
pres = extract_vars(wrfnc, timeidx, "PRES",
method, squeeze, cache, meta=False)["PRES"]
method, squeeze, cache, meta=False,
_key=_key)["PRES"]
return pres
def get_pressure_hpa(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="hpa"):
return get_pressure(wrfnc, timeidx, method, squeeze, cache, meta, units)
return get_pressure(wrfnc, timeidx, method, squeeze, cache, meta, _key,
units)

4587
src/wrf/psadlookup.py

File diff suppressed because it is too large Load Diff

4
src/wrf/pw.py

@ -14,10 +14,10 @@ from .metadecorators import copy_and_set_metadata @@ -14,10 +14,10 @@ from .metadecorators import copy_and_set_metadata
MemoryOrder="XY",
units="kg m-2")
def get_pw(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True):
meta=True, _key=None):
varnames=("T", "P", "PB", "PH", "PHB", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
t = ncvars["T"]
p = ncvars["P"]

8
src/wrf/rh.py

@ -12,10 +12,10 @@ from .metadecorators import copy_and_set_metadata @@ -12,10 +12,10 @@ from .metadecorators import copy_and_set_metadata
description="relative humidity",
delete_attrs=("units",))
def get_rh(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True):
meta=True, _key=None):
varnames=("T", "P", "PB", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
t = ncvars["T"]
p = ncvars["P"]
pb = ncvars["PB"]
@ -33,10 +33,10 @@ def get_rh(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, @@ -33,10 +33,10 @@ def get_rh(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
description="2m relative humidity",
delete_attrs=("units",))
def get_rh_2m(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True):
meta=True, _key=None):
varnames=("T2", "PSFC", "Q2")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
t2 = ncvars["T2"]
psfc = ncvars["PSFC"]
q2 = ncvars["Q2"]

15
src/wrf/routines.py

@ -1,7 +1,8 @@ @@ -1,7 +1,8 @@
from __future__ import (absolute_import, division, print_function,
unicode_literals)
from .util import get_iterable, is_standard_wrf_var, extract_vars, viewkeys
from .util import (get_iterable, is_standard_wrf_var, extract_vars, viewkeys,
get_id)
from .cape import get_2dcape, get_3dcape
from .ctt import get_ctt
from .dbz import get_dbz, get_max_dbz
@ -284,17 +285,21 @@ def getvar(wrfnc, varname, timeidx=0, @@ -284,17 +285,21 @@ def getvar(wrfnc, varname, timeidx=0,
"""
_key = get_id(wrfnc)
wrfnc = get_iterable(wrfnc)
if is_standard_wrf_var(wrfnc, varname):
if is_standard_wrf_var(wrfnc, varname) and varname != "Times":
return extract_vars(wrfnc, timeidx, varname,
method, squeeze, cache, meta)[varname]
method, squeeze, cache, meta, _key)[varname]
elif varname == "Times":
varname = "times" # Diverting to the get_times routine
actual_var = _undo_alias(varname)
if actual_var not in _VALID_KARGS:
raise ArgumentError("'%s' is not a valid variable name" % (varname))
_check_kargs(actual_var, kargs)
return _FUNC_MAP[actual_var](wrfnc, timeidx,
method, squeeze, cache, meta, **kargs)
return _FUNC_MAP[actual_var](wrfnc, timeidx, method, squeeze, cache,
meta, _key, **kargs)

4
src/wrf/slp.py

@ -16,11 +16,11 @@ from .util import extract_vars @@ -16,11 +16,11 @@ from .util import extract_vars
MemoryOrder="XY")
@convert_units("pressure", "hpa")
def get_slp(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="hpa"):
varnames=("T", "P", "PB", "QVAPOR", "PH", "PHB")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
t = ncvars["T"]
p = ncvars["P"]

30
src/wrf/temp.py

@ -13,12 +13,12 @@ from .util import extract_vars @@ -13,12 +13,12 @@ from .util import extract_vars
description="potential temperature")
@convert_units("temp", "k")
def get_theta(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="k"):
varnames = ("T",)
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
t = ncvars["T"]
full_t = t + Constants.T_BASE
@ -28,13 +28,13 @@ def get_theta(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -28,13 +28,13 @@ def get_theta(wrfnc, timeidx=0, method="cat", squeeze=True,
description="temperature")
@convert_units("temp", "k")
def get_temp(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="k"):
"""Return the temperature in Kelvin or Celsius"""
varnames=("T", "P", "PB")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
t = ncvars["T"]
p = ncvars["P"]
pb = ncvars["PB"]
@ -49,13 +49,13 @@ def get_temp(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -49,13 +49,13 @@ def get_temp(wrfnc, timeidx=0, method="cat", squeeze=True,
description="equivalent potential temperature")
@convert_units("temp", "k")
def get_eth(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="k"):
"Return equivalent potential temperature (Theta-e) in Kelvin"
varnames=("T", "P", "PB", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
t = ncvars["T"]
p = ncvars["P"]
pb = ncvars["PB"]
@ -73,13 +73,13 @@ def get_eth(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -73,13 +73,13 @@ def get_eth(wrfnc, timeidx=0, method="cat", squeeze=True,
description="virtual temperature")
@convert_units("temp", "k")
def get_tv(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="k"):
"Return the virtual temperature (tv) in Kelvin or Celsius"
varnames=("T", "P", "PB", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
t = ncvars["T"]
p = ncvars["P"]
@ -98,13 +98,13 @@ def get_tv(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -98,13 +98,13 @@ def get_tv(wrfnc, timeidx=0, method="cat", squeeze=True,
description="wetbulb temperature")
@convert_units("temp", "k")
def get_tw(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="k"):
"Return the wetbulb temperature (tw)"
varnames=("T", "P", "PB", "QVAPOR")
ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
t = ncvars["T"]
p = ncvars["P"]
pb = ncvars["PB"]
@ -119,12 +119,14 @@ def get_tw(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -119,12 +119,14 @@ def get_tw(wrfnc, timeidx=0, method="cat", squeeze=True,
return tw
def get_tk(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True):
return get_temp(wrfnc, timeidx, method, squeeze, cache, meta, units="k")
meta=True, _key=None):
return get_temp(wrfnc, timeidx, method, squeeze, cache, meta, _key,
units="k")
def get_tc(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True):
return get_temp(wrfnc, timeidx, method, squeeze, cache, meta, units="c")
meta=True, _key=None):
return get_temp(wrfnc, timeidx, method, squeeze, cache, meta, _key,
units="c")

5
src/wrf/terrain.py

@ -11,11 +11,12 @@ from .util import extract_vars, either @@ -11,11 +11,12 @@ from .util import extract_vars, either
description="terrain height")
@convert_units("height", "m")
def get_terrain(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=False, units="m"):
cache=None, meta=False, _key=None, units="m"):
varname = either("HGT", "HGT_M")(wrfnc)
return extract_vars(wrfnc, timeidx, varname,
method, squeeze, cache, meta=False)[varname]
method, squeeze, cache, meta=False,
_key=_key)[varname]

13
src/wrf/times.py

@ -4,5 +4,14 @@ from __future__ import (absolute_import, division, print_function, @@ -4,5 +4,14 @@ from __future__ import (absolute_import, division, print_function,
from .util import extract_times
def get_times(wrfnc, timeidx=0):
return extract_times(wrfnc, timeidx)
def get_times(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True, _key=None):
return extract_times(wrfnc, timeidx, method, squeeze, cache,
meta=meta, do_xtime=False)
def get_xtimes(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True, _key=None):
return extract_times(wrfnc, timeidx, method, squeeze, cache,
meta=meta, do_xtime=True)

651
src/wrf/util.py

@ -29,6 +29,7 @@ from .projection import getproj, NullProjection @@ -29,6 +29,7 @@ from .projection import getproj, NullProjection
from .constants import Constants, ALL_TIMES
from .py3compat import (viewitems, viewkeys, viewvalues, isstr, py2round,
py3range, ucode)
from .cache import cache_item, get_cached_item
if xarray_enabled():
from xarray import DataArray
@ -302,7 +303,7 @@ def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar): @@ -302,7 +303,7 @@ def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar):
def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"),
lonvar=either("XLONG", "XLONG_M")):
lonvar=either("XLONG", "XLONG_M"), _key=None):
if isinstance(latvar, either):
latvar = latvar(wrfseq)
@ -320,25 +321,14 @@ def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"), @@ -320,25 +321,14 @@ def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"),
wrf_iter = iter(wrfseq)
first_wrfnc = next(wrf_iter)
else:
entry = wrfseq[next(iter(viewkeys(wrfseq)))]
return is_moving_domain(entry, varname, latvar, lonvar)
# Quick check on pressure coordinates, bypassing the need to search the
# domain corner points
# try:
# coord_names = getattr(first_wrfnc.variables["P"],
# "coordinates").split()
# except KeyError:
# pass
# else:
# if "XTIME" in coord_names:
# return True
# else:
# return False
# The long way of checking all lat/lon corner points
# There doesn't appear to be a shortcut, so this should probably
# be stored in a cache somewhere
# Currently only checking the first dict entry.
dict_key = next(iter(viewkeys(wrfseq)))
entry = wrfseq[dict_key]
key = _key[dict_key] if _key is not None else None
return is_moving_domain(entry, varname, latvar, lonvar, key)
# The long way of checking all lat/lon corner points. Doesn't appear
# to be a shortcut in the netcdf files.
if varname is not None:
try:
coord_names = getattr(first_wrfnc.variables[varname],
@ -348,29 +338,20 @@ def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"), @@ -348,29 +338,20 @@ def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"),
# arguments
lon_coord = lonvar
lat_coord = latvar
else:
# If the XTIME variable is found to be a coordinate variable,
# then it's a moving domain file
# try:
# xtime_coord = coord_names[2]
# except IndexError:
# # XTIME is not a coordinate variable, if the variable is in the
# # file, then this is not a moving domain file
# if "XTIME" in first_wrfnc.variables:
# return False
#
# else:
# # XTIME is a coordinate, so this is a moving domain file
# if xtime_coord == "XTIME":
# return True
else:
lon_coord = coord_names[0]
lat_coord = coord_names[1]
else:
lon_coord = lonvar
lat_coord = latvar
# Probably a met_em file, need to search all the files
# See if there is a cached value
product = "is_moving_{}_{}".format(lat_coord, lon_coord)
moving = get_cached_item(_key, product)
if moving is not None:
return moving
# Need to search all the files
lats = first_wrfnc.variables[lat_coord]
lons = first_wrfnc.variables[lon_coord]
@ -397,9 +378,13 @@ def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"), @@ -397,9 +378,13 @@ def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"),
else:
if _corners_moved(wrfnc, ll_corner, ur_corner,
lat_coord, lon_coord):
cache_item(_key, product, True)
return True
cache_item(_key, product, False)
return False
def _get_global_attr(wrfnc, attr):
@ -441,7 +426,7 @@ def extract_dim(wrfnc, dim): @@ -441,7 +426,7 @@ def extract_dim(wrfnc, dim):
return len(d) #netCDF4
return d # PyNIO
def _combine_dict(wrfdict, varname, timeidx, method, meta):
def _combine_dict(wrfdict, varname, timeidx, method, meta, _key):
"""Dictionary combination creates a new left index for each key, then
does a cat or join for the list of files for that key"""
keynames = []
@ -451,11 +436,14 @@ def _combine_dict(wrfdict, varname, timeidx, method, meta): @@ -451,11 +436,14 @@ def _combine_dict(wrfdict, varname, timeidx, method, meta):
first_key = next(key_iter)
keynames.append(first_key)
is_moving = is_moving_domain(wrfdict, varname)
is_moving = is_moving_domain(wrfdict, varname, _key=_key)
# Not quite sure how to handle coord caching with dictionaries, so
# disabling it for now by setting _key to None.
first_array = _extract_var(wrfdict[first_key], varname,
timeidx, is_moving=is_moving, method=method,
squeeze=False, cache=None, meta=meta)
squeeze=False, cache=None, meta=meta,
_key=_key[first_key])
# Create the output data numpy array based on the first array
@ -474,7 +462,8 @@ def _combine_dict(wrfdict, varname, timeidx, method, meta): @@ -474,7 +462,8 @@ def _combine_dict(wrfdict, varname, timeidx, method, meta):
keynames.append(key)
vardata = _extract_var(wrfdict[key], varname, timeidx,
is_moving=is_moving, method=method,
squeeze=False, cache=None, meta=meta)
squeeze=False, cache=None, meta=meta,
_key=_key[key])
if outdata.shape[1:] != vardata.shape:
raise ValueError("data sequences must have the "
@ -486,8 +475,41 @@ def _combine_dict(wrfdict, varname, timeidx, method, meta): @@ -486,8 +475,41 @@ def _combine_dict(wrfdict, varname, timeidx, method, meta):
outname = str(first_array.name)
# Note: assumes that all entries in dict have same coords
outcoords = OrderedDict(first_array.coords)
outdims = ["key"] + list(first_array.dims)
outcoords["key"] = keynames
# First find and store all the existing key coord names/values
# This is applicable only if there are nested dictionaries.
key_coordnames = []
coord_vals = []
existing_cnt = 0
while True:
key_coord_name = "key_{}".format(existing_cnt)
if key_coord_name not in first_array.dims:
break
key_coordnames.append(key_coord_name)
coord_vals.append(npvalues(first_array.coords[key_coord_name]))
existing_cnt += 1
# Now add the key coord name and values for THIS dictionary.
# Put the new key_n name at the bottom, but the new values will
# be at the top to be associated with key_0 (left most). This
# effectively shifts the existing 'key_n' coordinate values to the
# right one dimension so *this* dicionary's key coordinate values
# are at 'key_0'.
key_coordnames.append(key_coord_name)
coord_vals.insert(0, keynames)
# make it so that key_0 is leftmost
outdims = key_coordnames + list(first_array.dims[existing_cnt:])
# Create the new 'key_n', value pairs
for coordname, coordval in zip(key_coordnames, coord_vals):
outcoords[coordname] = coordval
outattrs = OrderedDict(first_array.attrs)
outarr = DataArray(outdata, name=outname, coords=outcoords,
@ -497,6 +519,7 @@ def _combine_dict(wrfdict, varname, timeidx, method, meta): @@ -497,6 +519,7 @@ def _combine_dict(wrfdict, varname, timeidx, method, meta):
return outarr
def _find_coord_names(coords):
try:
lat_coord = [name for name in _COORD_VARS[0::2] if name in coords][0]
@ -532,11 +555,19 @@ def _find_max_time_size(wrfseq): @@ -532,11 +555,19 @@ def _find_max_time_size(wrfseq):
return max_times
def _build_data_array(wrfnc, varname, timeidx, is_moving_domain):
def _build_data_array(wrfnc, varname, timeidx, is_moving_domain, is_multifile,
_key):
# Note: wrfnc is always a single netcdf file object
# is_moving_domain and is_multifile are arguments indicating if the
# single file came from a sequence, and if that sequence is has a moving
# domain. Both arguments are used mainly for coordinate extraction and
# caching.
multitime = is_multi_time_req(timeidx)
time_idx_or_slice = timeidx if not multitime else slice(None)
var = wrfnc.variables[varname]
data = var[time_idx_or_slice, :]
time_coord = None
# Want to preserve the time dimension
if not multitime:
@ -551,14 +582,19 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain): @@ -551,14 +582,19 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain):
# WRF files
coord_attr = getattr(var, "coordinates")
except AttributeError:
if is_coordvar(varname):
# Coordinate variable (most likely XLAT or XLONG)
lat_coord, lon_coord = get_coord_pairs(varname)
time_coord = None
if is_moving_domain and has_time_coord(wrfnc):
if has_time_coord(wrfnc):
time_coord = "XTIME"
elif is_time_coord_var(varname):
lon_coord = None
lat_coord = None
time_coord = None
else:
try:
# met_em files
@ -584,18 +620,60 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain): @@ -584,18 +620,60 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain):
# Handle lat/lon coordinates and projection information if available
if lon_coord is not None and lat_coord is not None:
lon_coord_var = wrfnc.variables[lon_coord]
lat_coord_var = wrfnc.variables[lat_coord]
time_coord_var = (wrfnc.variables[time_coord]
if time_coord is not None
else None)
# Using a cache for coordinate variables so the extraction only happens
# once.
lon_coord_dimkey = lon_coord + "_dim"
lon_coord_valkey = lon_coord + "_val"
lat_coord_dimkey = lat_coord + "_dim"
lat_coord_valkey = lat_coord + "_val"
lon_coord_dims = get_cached_item(_key, lon_coord_dimkey)
lon_coord_vals = get_cached_item(_key, lon_coord_valkey)
if lon_coord_dims is None or lon_coord_vals is None:
lon_var = wrfnc.variables[lon_coord]
lon_coord_dims = lon_var.dimensions
lon_coord_vals = lon_var[:]
# Only cache here if the domain is not moving, otherwise
# caching is handled by cat/join
if not is_moving_domain:
cache_item(_key, lon_coord_dimkey, lon_coord_dims)
cache_item(_key, lon_coord_valkey, lon_coord_vals)
lat_coord_dims = get_cached_item(_key, lat_coord_dimkey)
lat_coord_vals = get_cached_item(_key, lat_coord_valkey)
if lat_coord_dims is None or lat_coord_vals is None:
lat_var = wrfnc.variables[lat_coord]
lat_coord_dims = lat_var.dimensions
lat_coord_vals = lat_var[:]
# Only cache here if the domain is not moving, otherwise
# caching is done in cat/join
if not is_moving_domain:
cache_item(_key, lat_coord_dimkey, lat_coord_dims)
cache_item(_key, lat_coord_valkey, lat_coord_vals)
time_coord_vals = None
if time_coord is not None:
# If not from a multifile sequence, then cache the time
# coordinate. Otherwise, handled in cat/join/
if not is_multifile:
time_coord_vals = get_cached_item(_key, time_coord)
if time_coord_vals is None:
time_coord_vals = wrfnc.variables[time_coord][:]
if not is_multifile:
cache_item(_key, time_coord, time_coord_vals)
else:
time_coord_vals = wrfnc.variables[time_coord][:]
if multitime:
if is_moving_domain:
# Special case with a moving domain in a multi-time file,
# otherwise the projection parameters don't change
coords[lon_coord] = lon_coord_var.dimensions, lon_coord_var[:]
coords[lat_coord] = lat_coord_var.dimensions, lat_coord_var[:]
coords[lon_coord] = lon_coord_dims, lon_coord_vals
coords[lat_coord] = lat_coord_dims, lat_coord_vals
# Returned lats/lons arrays will have a time dimension, so proj
# will need to be a list due to moving corner points
@ -603,26 +681,32 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain): @@ -603,26 +681,32 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain):
timeidx,
varname)
proj = [getproj(lats=lats[i,:],
lons=lons[i,:],
**proj_params) for i in py3range(lats.shape[0])]
lons=lons[i,:],
**proj_params) for i in py3range(lats.shape[0])]
if time_coord is not None:
coords[time_coord] = (lon_coord_var.dimensions[0],
time_coord_var[:])
else:
coords[lon_coord] = (lon_coord_var.dimensions[1:],
lon_coord_var[0,:])
coords[lat_coord] = (lat_coord_var.dimensions[1:],
lat_coord_var[0,:])
coords[lon_coord] = (lon_coord_dims[1:],
lon_coord_vals[0,:])
coords[lat_coord] = (lat_coord_dims[1:],
lat_coord_vals[0,:])
# Domain not moving, so just get the first time
lats, lons, proj_params = get_proj_params(wrfnc, 0, varname)
proj = getproj(lats=lats, lons=lons, **proj_params)
if time_coord is not None:
coords[time_coord] = (lon_coord_dims[0], time_coord_vals)
else:
coords[lon_coord] = (lon_coord_var.dimensions[1:],
lon_coord_var[timeidx,:])
coords[lat_coord] = (lat_coord_var.dimensions[1:],
lat_coord_var[timeidx,:])
coords[lon_coord] = (lon_coord_dims[1:],
lon_coord_vals[timeidx,:])
coords[lat_coord] = (lat_coord_dims[1:],
lat_coord_vals[timeidx,:])
if time_coord is not None:
coords[time_coord] = (lon_coord_dims[0],
[time_coord_vals[timeidx]])
lats, lons, proj_params = get_proj_params(wrfnc, 0, varname)
proj = getproj(lats=lats, lons=lons, **proj_params)
@ -630,7 +714,10 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain): @@ -630,7 +714,10 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain):
if dimnames[0] == "Time":
coords[dimnames[0]] = extract_times(wrfnc, timeidx)
t = extract_times(wrfnc, timeidx, meta=False, do_xtime=False)
if not multitime:
t = [t]
coords[dimnames[0]] = t
data_array = DataArray(data, name=varname, dims=dimnames, coords=coords,
attrs=attrs)
@ -639,7 +726,7 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain): @@ -639,7 +726,7 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain):
return data_array
def _find_forward(wrfseq, varname, timeidx, is_moving, meta):
def _find_forward(wrfseq, varname, timeidx, is_moving, meta, _key):
wrf_iter = iter(wrfseq)
comboidx = 0
@ -657,18 +744,18 @@ def _find_forward(wrfseq, varname, timeidx, is_moving, meta): @@ -657,18 +744,18 @@ def _find_forward(wrfseq, varname, timeidx, is_moving, meta):
if meta:
return _build_data_array(wrfnc, varname, filetimeidx,
is_moving)
is_moving, True, _key)
else:
result = wrfnc.variables[varname][filetimeidx, :]
return result[np.newaxis, :] # So that nosqueee works
return result[np.newaxis, :] # So that nosqueeze works
else:
comboidx += numtimes
raise IndexError("invalid time index")
raise IndexError("timeidx {} is out of bounds".format(timeidx))
def _find_reverse(wrfseq, varname, timeidx, is_moving, meta):
def _find_reverse(wrfseq, varname, timeidx, is_moving, meta, _key):
try:
revwrfseq = reversed(wrfseq)
except TypeError:
@ -695,35 +782,36 @@ def _find_reverse(wrfseq, varname, timeidx, is_moving, meta): @@ -695,35 +782,36 @@ def _find_reverse(wrfseq, varname, timeidx, is_moving, meta):
if meta:
return _build_data_array(wrfnc, varname, filetimeidx,
is_moving)
is_moving, True, _key)
else:
result = wrfnc.variables[varname][filetimeidx, :]
return result[np.newaxis, :] # So that nosqueeze works
else:
comboidx += numtimes
raise IndexError("invalid time index")
raise IndexError("timeidx {} is out of bounds".format(timeidx))
def _find_arr_for_time(wrfseq, varname, timeidx, is_moving, meta):
def _find_arr_for_time(wrfseq, varname, timeidx, is_moving, meta, _key):
if timeidx >= 0:
return _find_forward(wrfseq, varname, timeidx, is_moving, meta)
return _find_forward(wrfseq, varname, timeidx, is_moving, meta, _key)
else:
return _find_reverse(wrfseq, varname, timeidx, is_moving, meta)
return _find_reverse(wrfseq, varname, timeidx, is_moving, meta, _key)
# TODO: implement in C
def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta):
def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta, _key):
if is_moving is None:
is_moving = is_moving_domain(wrfseq, varname)
is_moving = is_moving_domain(wrfseq, varname, _key=_key)
file_times = extract_times(wrfseq, ALL_TIMES)
file_times = extract_times(wrfseq, ALL_TIMES, meta=False, do_xtime=False)
multitime = is_multi_time_req(timeidx)
# For single times, just need to find the ncfile and appropriate
# time index, and return that array
if not multitime:
return _find_arr_for_time(wrfseq, varname, timeidx, is_moving, meta)
return _find_arr_for_time(wrfseq, varname, timeidx, is_moving, meta,
_key)
#time_idx_or_slice = timeidx if not multitime else slice(None)
@ -733,7 +821,7 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta): @@ -733,7 +821,7 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta):
if xarray_enabled() and meta:
first_var = _build_data_array(next(wrf_iter), varname,
ALL_TIMES, is_moving)
ALL_TIMES, is_moving, True, _key)
else:
first_var = (next(wrf_iter)).variables[varname][:]
@ -750,27 +838,62 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta): @@ -750,27 +838,62 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta):
outdata[startidx:endidx, :] = first_var[:]
if xarray_enabled() and meta and is_moving:
if xarray_enabled() and meta:
latname, lonname, timename = _find_coord_names(first_var.coords)
outcoorddims = outdims[0:1] + outdims[-2:]
if latname is not None:
outlats = np.empty(outcoorddims, first_var.dtype)
outlats[startidx:endidx, :] = first_var.coords[latname][:]
timecached = False
latcached = False
loncached = False
projcached = False
outxtimes = None
outlats = None
outlons = None
outprojs = None
timekey = timename+"_cat" if timename is not None else None
latkey = latname + "_cat" if latname is not None else None
lonkey = lonname + "_cat" if lonname is not None else None
projkey = "projection_cat" if is_moving else None
if timename is not None:
outxtimes = get_cached_item(_key, timekey)
if outxtimes is None:
outxtimes = np.empty(outdims[0])
outxtimes[startidx:endidx] = npvalues(first_var.coords[timename][:])
else:
timecached = True
if is_moving:
outcoorddims = outdims[0:1] + outdims[-2:]
if lonname is not None:
outlons = np.empty(outcoorddims, first_var.dtype)
outlons[startidx:endidx, :] = first_var.coords[lonname][:]
if latname is not None:
# Try to pull from the coord cache
outlats = get_cached_item(_key, latkey)
if outlats is None:
outlats = np.empty(outcoorddims, first_var.dtype)
outlats[startidx:endidx, :] = npvalues(first_var.coords[latname][:])
else:
latcached = True
if lonname is not None:
outlons = get_cached_item(_key, lonkey)
if outlons is None:
outlons = np.empty(outcoorddims, first_var.dtype)
outlons[startidx:endidx, :] = npvalues(first_var.coords[lonname][:])
else:
loncached = True
# Projections also need to be aggregated
if timename is not None:
outxtimes = np.empty(outdims[0])
outxtimes[startidx:endidx] = first_var.coords[timename][:]
outprojs = get_cached_item(_key, projkey)
if outprojs is None:
outprojs = np.empty(outdims[0], np.object)
outprojs[startidx:endidx] = np.asarray(
first_var.attrs["projection"], np.object)[:]
else:
projcached = True
# Projections also need to be aggregated
outprojs = np.empty(outdims[0], np.object)
outprojs[startidx:endidx] = np.asarray(first_var.attrs["projection"],
np.object)[:]
startidx = endidx
while True:
@ -787,31 +910,46 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta): @@ -787,31 +910,46 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta):
outdata[startidx:endidx, :] = vardata[:]
if xarray_enabled() and meta and is_moving:
if latname is not None:
latdata = wrfnc.variables[latname][:]
outlats[startidx:endidx, :] = latdata[:]
if lonname is not None:
londata = wrfnc.variables[lonname][:]
outlons[startidx:endidx, :] = londata[:]
if timename is not None:
if xarray_enabled() and meta:
# XTIME new in 3.7
if timename is not None and not timecached:
xtimedata = wrfnc.variables[timename][:]
outxtimes[startidx:endidx] = xtimedata[:]
lats, lons, proj_params = get_proj_params(wrfnc,
ALL_TIMES,
varname)
projs = [getproj(lats=lats[i,:],
lons=lons[i,:],
**proj_params) for i in py3range(lats.shape[0])]
outprojs[startidx:endidx] = np.asarray(projs, np.object)[:]
if is_moving:
if latname is not None and not latcached:
latdata = wrfnc.variables[latname][:]
outlats[startidx:endidx, :] = latdata[:]
if lonname is not None and not loncached:
londata = wrfnc.variables[lonname][:]
outlons[startidx:endidx, :] = londata[:]
if not projcached:
lats, lons, proj_params = get_proj_params(wrfnc,
ALL_TIMES,
varname)
projs = [getproj(lats=lats[i,:],
lons=lons[i,:],
**proj_params) for i in py3range(lats.shape[0])]
outprojs[startidx:endidx] = np.asarray(projs,
np.object)[:]
startidx = endidx
if xarray_enabled() and meta:
# Cache the coords if applicable
if not latcached and outlats is not None:
cache_item(_key, latkey, outlats)
if not loncached and outlons is not None:
cache_item(_key, lonkey, outlons)
if not projcached and outprojs is not None:
cache_item(_key, projkey, outprojs)
if not timecached and outxtimes is not None:
cache_item(_key, timekey, outxtimes)
outname = ucode(first_var.name)
outattrs = OrderedDict(first_var.attrs)
outcoords = OrderedDict(first_var.coords)
@ -827,7 +965,11 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta): @@ -827,7 +965,11 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta):
outcoords["datetime"] = outdimnames[0], file_times
# If the domain is moving, need to create the lat/lon/xtime coords
if timename is not None:
outxtimes = outxtimes[:]
outcoords[timename] = outdimnames[0], outxtimes
# If the domain is moving, need to create the lat/lon coords
# since they can't be copied
if is_moving:
outlatdims = [outdimnames[0]] + outdimnames[-2:]
@ -838,10 +980,7 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta): @@ -838,10 +980,7 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta):
if lonname is not None:
outlons = outlons[:]
outcoords[lonname] = outlatdims, outlons
if timename is not None:
outxtimes = outxtimes[:]
outcoords[timename] = outdimnames[0], outxtimes
outattrs["projection"] = outprojs[:]
outdata = outdata[:]
@ -855,6 +994,7 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta): @@ -855,6 +994,7 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta):
return outarr
def _get_numfiles(wrfseq):
try:
return len(wrfseq)
@ -862,10 +1002,11 @@ def _get_numfiles(wrfseq): @@ -862,10 +1002,11 @@ def _get_numfiles(wrfseq):
wrf_iter = iter(wrfseq)
return sum(1 for _ in wrf_iter)
# TODO: implement in C
def _join_files(wrfseq, varname, timeidx, is_moving, meta):
def _join_files(wrfseq, varname, timeidx, is_moving, meta, _key):
if is_moving is None:
is_moving = is_moving_domain(wrfseq, varname)
is_moving = is_moving_domain(wrfseq, varname, _key=_key)
multitime = is_multi_time_req(timeidx)
numfiles = _get_numfiles(wrfseq)
maxtimes = _find_max_time_size(wrfseq)
@ -880,7 +1021,8 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta): @@ -880,7 +1021,8 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta):
numtimes = extract_dim(wrfnc, "Time")
if xarray_enabled() and meta:
first_var = _build_data_array(wrfnc, varname, ALL_TIMES, is_moving)
first_var = _build_data_array(wrfnc, varname, ALL_TIMES, is_moving,
True, _key)
time_coord = np.full((numfiles, maxtimes), int(NaT), "datetime64[ns]")
time_coord[file_idx, 0:numtimes] = first_var.coords["Time"][:]
else:
@ -900,31 +1042,67 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta): @@ -900,31 +1042,67 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta):
outdata[file_idx, 0:numtimes, :] = first_var[:]
# Create the secondary coordinate arrays
if xarray_enabled() and meta and is_moving:
if xarray_enabled() and meta:
latname, lonname, timename = _find_coord_names(first_var.coords)
outcoorddims = outdims[0:2] + outdims[-2:]
if latname is not None:
outlats = np.full(outcoorddims, Constants.DEFAULT_FILL,
first_var.dtype)
outlats[file_idx, 0:numtimes, :] = first_var.coords[latname][:]
if lonname is not None:
outlons = np.full(outcoorddims, Constants.DEFAULT_FILL,
first_var.dtype)
outlons[file_idx, 0:numtimes, :] = first_var.coords[lonname][:]
outcoorddims = outdims[0:2] + outdims[-2:]
timecached = False
latcached = False
loncached = False
projcached = False
outxtimes = None
outlats = None
outlons = None
outprojs = None
timekey = timename+"_join" if timename is not None else None
latkey = latname + "_join" if latname is not None else None
lonkey = lonname + "_join" if lonname is not None else None
projkey = "projection_join" if is_moving else None
if timename is not None:
outxtimes = np.full(outdims[0:2], Constants.DEFAULT_FILL,
outxtimes = get_cached_item(_key, timekey)
if outxtimes is None:
outxtimes = np.full(outdims[0:2], Constants.DEFAULT_FILL,
first_var.dtype)
outxtimes[file_idx, 0:numtimes] = first_var.coords[timename][:]
# Projections also need two dimensions
outprojs = np.full(outdims[0:2], NullProjection(), np.object)
outxtimes[file_idx, 0:numtimes] = first_var.coords[timename][:]
else:
timecached = True
outprojs[file_idx, 0:numtimes] = np.asarray(
first_var.attrs["projection"],
np.object)[:]
if is_moving:
if latname is not None:
outlats = get_cached_item(_key, latkey)
if outlats is None:
outlats = np.full(outcoorddims, Constants.DEFAULT_FILL,
first_var.dtype)
outlats[file_idx, 0:numtimes, :] = (
first_var.coords[latname][:])
else:
latcached = True
if lonname is not None:
outlons = get_cached_item(_key, lonkey)
if outlons is None:
outlons = np.full(outcoorddims, Constants.DEFAULT_FILL,
first_var.dtype)
outlons[file_idx, 0:numtimes, :] = (
first_var.coords[lonname][:])
else:
loncached = True
# Projections also need two dimensions
outprojs = get_cached_item(_key, projkey)
if outprojs is None:
outprojs = np.full(outdims[0:2], NullProjection(), np.object)
outprojs[file_idx, 0:numtimes] = np.asarray(
first_var.attrs["projection"],
np.object)[:]
else:
projcached = True
file_idx=1
while True:
@ -944,43 +1122,56 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta): @@ -944,43 +1122,56 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta):
outdata[file_idx, 0:numtimes, :] = outvar[:]
if xarray_enabled() and meta:
file_times = extract_times(wrfnc, ALL_TIMES)
# For join, the times are a function of fileidx
file_times = extract_times(wrfnc, ALL_TIMES, meta=False,
do_xtime=False)
time_coord[file_idx, 0:numtimes] = np.asarray(file_times,
"datetime64[ns]")[:]
if xarray_enabled() and meta and is_moving:
if latname is not None:
latdata = wrfnc.variables[latname][:]
outlats[file_idx, 0:numtimes, :] = latdata[:]
if lonname is not None:
londata = wrfnc.variables[lonname][:]
outlons[file_idx, 0:numtimes, :] = londata[:]
if timename is not None:
if timename is not None and not timecached:
xtimedata = wrfnc.variables[timename][:]
outxtimes[file_idx, 0:numtimes] = xtimedata[:]
lats, lons, proj_params = get_proj_params(wrfnc,
ALL_TIMES,
varname)
projs = [getproj(lats=lats[i,:],
lons=lons[i,:],
**proj_params) for i in py3range(lats.shape[0])]
outprojs[file_idx, 0:numtimes] = (
np.asarray(projs, np.object)[:])
if is_moving:
if latname is not None and not latcached:
latdata = wrfnc.variables[latname][:]
outlats[file_idx, 0:numtimes, :] = latdata[:]
if lonname is not None and not loncached:
londata = wrfnc.variables[lonname][:]
outlons[file_idx, 0:numtimes, :] = londata[:]
if not projcached:
lats, lons, proj_params = get_proj_params(wrfnc,
ALL_TIMES,
varname)
projs = [getproj(lats=lats[i,:],
lons=lons[i,:],
**proj_params) for i in py3range(lats.shape[0])]
outprojs[file_idx, 0:numtimes] = (
np.asarray(projs, np.object)[:])
# Need to update coords here
file_idx += 1
file_idx += 1
# If any of the output files contain less than the max number of times,
# then a mask array is needed to flag all the missing arrays with
# missing values
if file_times_less_than_max:
outdata = np.ma.masked_values(outdata, Constants.DEFAULT_FILL)
if xarray_enabled() and meta:
# Cache the coords if applicable
if not latcached and outlats is not None:
cache_item(_key, latkey, outlats)
if not loncached and outlons is not None:
cache_item(_key, lonkey, outlons)
if not projcached and outprojs is not None:
cache_item(_key, projkey, outprojs)
if not timecached and outxtimes is not None:
cache_item(_key, timekey, outxtimes)
outname = ucode(first_var.name)
outcoords = OrderedDict(first_var.coords)
outattrs = OrderedDict(first_var.attrs)
@ -1000,7 +1191,13 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta): @@ -1000,7 +1191,13 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta):
outattrs["_FillValue"] = Constants.DEFAULT_FILL
outattrs["missing_value"] = Constants.DEFAULT_FILL
# If the domain is moving, need to create the lat/lon/xtime coords
if timename is not None:
outxtimes = outxtimes[:, time_idx_or_slice]
if not multitime:
outxtimes = outxtimes[:, np.newaxis]
outcoords[timename] = outdimnames[0:2], outxtimes[:]
# If the domain is moving, need to create the lat/lon coords
# since they can't be copied
if is_moving:
outlatdims = outdimnames[0:2] + outdimnames[-2:]
@ -1015,11 +1212,6 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta): @@ -1015,11 +1212,6 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta):
if not multitime:
outlons = outlons[:, np.newaxis, :]
outcoords[lonname] = outlatdims, outlons
if timename is not None:
outxtimes = outxtimes[:, time_idx_or_slice]
if not multitime:
outxtimes = outxtimes[:, np.newaxis]
outcoords[timename] = outdimnames[0:2], outxtimes[:]
if not multitime:
outattrs["projection"] = outprojs[:, timeidx]
@ -1043,19 +1235,20 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta): @@ -1043,19 +1235,20 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta):
return outarr
def combine_files(wrfseq, varname, timeidx, is_moving=None,
method="cat", squeeze=True, meta=True):
method="cat", squeeze=True, meta=True,
_key=None):
# Handles generators, single files, lists, tuples, custom classes
wrfseq = get_iterable(wrfseq)
# Dictionary is unique
if is_mapping(wrfseq):
outarr = _combine_dict(wrfseq, varname, timeidx, method, meta)
outarr = _combine_dict(wrfseq, varname, timeidx, method, meta, _key)
elif method.lower() == "cat":
outarr = _cat_files(wrfseq, varname, timeidx, is_moving,
squeeze, meta)
squeeze, meta, _key)
elif method.lower() == "join":
outarr = _join_files(wrfseq, varname, timeidx, is_moving, meta)
outarr = _join_files(wrfseq, varname, timeidx, is_moving, meta, _key)
else:
raise ValueError("method must be 'cat' or 'join'")
@ -1064,7 +1257,7 @@ def combine_files(wrfseq, varname, timeidx, is_moving=None, @@ -1064,7 +1257,7 @@ def combine_files(wrfseq, varname, timeidx, is_moving=None,
# Cache is a dictionary of already extracted variables
def _extract_var(wrfnc, varname, timeidx, is_moving,
method, squeeze, cache, meta):
method, squeeze, cache, meta, _key):
# Mainly used internally so variables don't get extracted multiple times,
# particularly to copy metadata. This can be slow.
if cache is not None:
@ -1081,11 +1274,16 @@ def _extract_var(wrfnc, varname, timeidx, is_moving, @@ -1081,11 +1274,16 @@ def _extract_var(wrfnc, varname, timeidx, is_moving,
multitime = is_multi_time_req(timeidx)
multifile = is_multi_file(wrfnc)
if is_time_coord_var(varname):
return extract_times(wrfnc, timeidx, method, squeeze, cache,
meta, do_xtime=True)
if not multifile:
if xarray_enabled() and meta:
if is_moving is None:
is_moving = is_moving_domain(wrfnc, varname)
result = _build_data_array(wrfnc, varname, timeidx, is_moving)
is_moving = is_moving_domain(wrfnc, varname, _key=_key)
result = _build_data_array(wrfnc, varname, timeidx, is_moving,
multifile, _key)
else:
if not multitime:
result = wrfnc.variables[varname][timeidx,:]
@ -1095,20 +1293,20 @@ def _extract_var(wrfnc, varname, timeidx, is_moving, @@ -1095,20 +1293,20 @@ def _extract_var(wrfnc, varname, timeidx, is_moving,
else:
# Squeeze handled in this routine, so just return it
return combine_files(wrfnc, varname, timeidx, is_moving,
method, squeeze, meta)
method, squeeze, meta, _key)
return result.squeeze() if squeeze else result
def extract_vars(wrfnc, timeidx, varnames, method="cat", squeeze=True,
cache=None, meta=True):
cache=None, meta=True, _key=None):
if isstr(varnames):
varlist = [varnames]
else:
varlist = varnames
return {var:_extract_var(wrfnc, var, timeidx, None,
method, squeeze, cache, meta)
method, squeeze, cache, meta, _key)
for var in varlist}
# Python 3 compatability
@ -1120,28 +1318,80 @@ def _make_time(timearr): @@ -1120,28 +1318,80 @@ def _make_time(timearr):
return dt.datetime.strptime("".join(npbytes_to_str(timearr)),
"%Y-%m-%d_%H:%M:%S")
def _file_times(wrfnc, timeidx):
multitime = is_multi_time_req(timeidx)
if multitime:
def _file_times(wrfnc, do_xtime):
if not do_xtime:
times = wrfnc.variables["Times"][:,:]
for i in py3range(times.shape[0]):
yield _make_time(times[i,:])
else:
times = wrfnc.variables["Times"][timeidx,:]
yield _make_time(times)
xtimes = wrfnc.variables["XTIME"][:]
for i in py3range(xtimes.shape[0]):
yield xtimes[i]
def _extract_time_map(wrfnc, timeidx, do_xtime, meta=False):
return {key : extract_times(wrfseq, timeidx, do_xtime, meta)
for key, wrfseq in viewitems(wrfnc)}
def extract_times(wrfnc, timeidx):
def extract_times(wrfnc, timeidx, method="cat", squeeze=True, cache=None,
meta=False, do_xtime=False):
if is_mapping(wrfnc):
return _extract_time_map(wrfnc, timeidx, do_xtime)
multitime = is_multi_time_req(timeidx)
multi_file = is_multi_file(wrfnc)
if not multi_file:
wrf_list = [wrfnc]
else:
wrf_list = wrfnc
return [file_time
for wrf_file in wrf_list
for file_time in _file_times(wrf_file, timeidx)]
try:
if method.lower() == "cat":
time_list = [file_time
for wrf_file in wrf_list
for file_time in _file_times(wrf_file, do_xtime)]
elif method.lower() == "join":
time_list = [[file_time
for file_time in _file_times(wrf_file, do_xtime)]
for wrf_file in wrf_list]
else:
raise ValueError("invalid method argument '{}'".format(method))
except KeyError:
return None # Thrown for pre-3.7 XTIME not existing
if xarray_enabled() and meta:
outattrs = OrderedDict()
outcoords = None
if method.lower() == "cat":
outdimnames = ["Time"]
else:
outdimnames = ["fileidx", "Time"]
if not do_xtime:
outname = "times"
outattrs["description"] = "model times [np.datetime64]"
else:
ncfile = next(iter(wrf_list))
var = ncfile.variables["XTIME"]
outattrs.update(var.__dict__)
outname = "XTIME"
outarr = DataArray(time_list, name=outname, coords=outcoords,
dims=outdimnames, attrs=outattrs)
else:
outarr = np.asarray(time_list)
if not multitime:
return outarr[timeidx]
return outarr
def is_standard_wrf_var(wrfnc, var):
multifile = is_multi_file(wrfnc)
@ -1385,8 +1635,21 @@ def arg_location(func, argname, args, kwargs): @@ -1385,8 +1635,21 @@ def arg_location(func, argname, args, kwargs):
return _arg_location(func, argname, args, kwargs)
def psafilepath():
return os.path.join(os.path.dirname(__file__), "data", "psadilookup.dat")
def get_id(seq):
if not is_mapping(seq):
return id(seq)
# For each key in the mapping, recurisvely call get_id until
# until a non-mapping is found
return {key : get_id(val) for key,val in viewitems(seq)}

32
src/wrf/uvmet.py

@ -17,31 +17,31 @@ from .util import extract_vars, extract_global_attrs, either @@ -17,31 +17,31 @@ from .util import extract_vars, extract_global_attrs, either
@convert_units("wind", "mps")
def _get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
ten_m=False, units ="mps"):
""" Return a tuple of u,v with the winds rotated in to earth space"""
if not ten_m:
varname = either("U", "UU")(wrfnc)
u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
u = destagger(u_vars[varname], -1)
varname = either("V", "VV")(wrfnc)
v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
v = destagger(v_vars[varname], -2)
else:
varname = either("U10", "UU")(wrfnc)
u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
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,
meta=False)
meta=False, _key=_key)
v = (v_vars[varname] if varname == "V10" else
destagger(v_vars[varname][...,0,:,:], -2))
@ -97,12 +97,14 @@ def _get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -97,12 +97,14 @@ def _get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True,
varname = either("XLAT_M", "XLAT")(wrfnc)
xlat_var = extract_vars(wrfnc, timeidx, varname,
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False,
_key=_key)
lat = xlat_var[varname]
varname = either("XLONG_M", "XLONG")(wrfnc)
xlon_var = extract_vars(wrfnc, timeidx, varname,
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False,
_key=_key)
lon = xlon_var[varname]
if map_proj == 1:
@ -132,10 +134,10 @@ def _get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -132,10 +134,10 @@ def _get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True,
two_d=False,
wspd_wdir=False)
def get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="mps"):
return _get_uvmet(wrfnc, timeidx, method, squeeze, cache, meta,
return _get_uvmet(wrfnc, timeidx, method, squeeze, cache, meta, _key,
False, units)
@ -145,10 +147,10 @@ def get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -145,10 +147,10 @@ def get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True,
two_d=True,
wspd_wdir=False)
def get_uvmet10(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="mps"):
return _get_uvmet(wrfnc, timeidx, method, squeeze, cache, meta,
return _get_uvmet(wrfnc, timeidx, method, squeeze, cache, meta, _key,
True, units)
@ -158,11 +160,11 @@ def get_uvmet10(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -158,11 +160,11 @@ def get_uvmet10(wrfnc, timeidx=0, method="cat", squeeze=True,
two_d=False,
wspd_wdir=True)
def get_uvmet_wspd_wdir(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="mps"):
uvmet = _get_uvmet(wrfnc, timeidx, method, squeeze,
cache, meta, False, units)
cache, meta, _key, False, units)
return _calc_wspd_wdir(uvmet[0,...,:,:,:], uvmet[1,...,:,:,:],
False, units)
@ -174,10 +176,10 @@ def get_uvmet_wspd_wdir(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -174,10 +176,10 @@ def get_uvmet_wspd_wdir(wrfnc, timeidx=0, method="cat", squeeze=True,
two_d=True,
wspd_wdir=True)
def get_uvmet10_wspd_wdir(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="mps"):
uvmet10 = _get_uvmet(wrfnc, timeidx, method, squeeze, cache, meta,
uvmet10 = _get_uvmet(wrfnc, timeidx, method, squeeze, cache, meta, _key,
True, units)
return _calc_wspd_wdir(uvmet10[0,...,:,:], uvmet10[1,...,:,:], True, units)

2
src/wrf/version.py

@ -1,2 +1,2 @@ @@ -1,2 +1,2 @@
__version__ = "1.0a1"
__version__ = "1.0a2"

8
src/wrf/vorticity.py

@ -10,11 +10,11 @@ from .metadecorators import copy_and_set_metadata @@ -10,11 +10,11 @@ from .metadecorators import copy_and_set_metadata
description="absolute vorticity",
units="10-5 s-1")
def get_avo(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True):
meta=True, _key=None):
ncvars = extract_vars(wrfnc, timeidx, ("U", "V", "MAPFAC_U",
"MAPFAC_V", "MAPFAC_M",
"F"),
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False, _key=_key)
attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
u = ncvars["U"]
@ -34,12 +34,12 @@ def get_avo(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, @@ -34,12 +34,12 @@ def get_avo(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
description="potential vorticity",
units="PVU")
def get_pvo(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True):
meta=True, _key=None):
ncvars = extract_vars(wrfnc, timeidx, ("U", "V", "T", "P",
"PB", "MAPFAC_U",
"MAPFAC_V", "MAPFAC_M",
"F"),
method, squeeze, cache, meta=False)
method, squeeze, cache, meta=False, _key=_key)
attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
u = ncvars["U"]

24
src/wrf/wind.py

@ -52,11 +52,11 @@ def _calc_wspd_wdir(u, v, two_d, units): @@ -52,11 +52,11 @@ def _calc_wspd_wdir(u, v, two_d, units):
wspd_wdir=False)
@convert_units("wind", "mps")
def get_u_destag(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="mps"):
varname = either("U", "UU")(wrfnc)
u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
u = destagger(u_vars[varname], -1)
return u
@ -70,11 +70,11 @@ def get_u_destag(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -70,11 +70,11 @@ def get_u_destag(wrfnc, timeidx=0, method="cat", squeeze=True,
wspd_wdir=False)
@convert_units("wind", "mps")
def get_v_destag(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="mps"):
varname = either("V", "VV")(wrfnc)
v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
v = destagger(v_vars[varname], -2)
return v
@ -88,10 +88,10 @@ def get_v_destag(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -88,10 +88,10 @@ def get_v_destag(wrfnc, timeidx=0, method="cat", squeeze=True,
wspd_wdir=False)
@convert_units("wind", "mps")
def get_w_destag(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True,
cache=None, meta=True, _key=None,
units="mps"):
w_vars = extract_vars(wrfnc, timeidx, "W", method, squeeze, cache,
meta=False)
meta=False, _key=_key)
w = destagger(w_vars["W"], -3)
return w
@ -102,16 +102,16 @@ def get_w_destag(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -102,16 +102,16 @@ def get_w_destag(wrfnc, timeidx=0, method="cat", squeeze=True,
two_d=False,
wspd_wdir=True)
def get_destag_wspd_wdir(wrfnc, timeidx=0, method="cat",
squeeze=True, cache=None, meta=True,
squeeze=True, cache=None, meta=True, _key=None,
units="mps"):
varname = either("U", "UU")(wrfnc)
u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
u = destagger(u_vars[varname], -1)
varname = either("V", "VV")(wrfnc)
v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
v = destagger(v_vars[varname], -2)
return _calc_wspd_wdir(u, v, False, units)
@ -123,18 +123,18 @@ def get_destag_wspd_wdir(wrfnc, timeidx=0, method="cat", @@ -123,18 +123,18 @@ def get_destag_wspd_wdir(wrfnc, timeidx=0, method="cat",
two_d=False,
wspd_wdir=True)
def get_destag_wspd_wdir10(wrfnc, timeidx=0, method="cat",
squeeze=True, cache=None, meta=True,
squeeze=True, cache=None, meta=True, _key=None,
units="mps"):
varname = either("U10", "UU")(wrfnc)
u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,
meta=False)
meta=False, _key=_key)
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,
meta=False)
meta=False, _key=_key)
v = (v_vars[varname] if varname == "V10" else
destagger(v_vars[varname][...,0,:,:], -2))

2
test/ipynb/WRF_Workshop_Demo.ipynb

@ -293,7 +293,7 @@ @@ -293,7 +293,7 @@
"angle = 90.0\n",
"\n",
"# Compute the vertical cross-section interpolation. Also, include the lat/lon points along the cross-section.\n",
"p_vertx = vertcross(p, z, pivot_point=pivot_point, angle=angle, include_latlon=True)\n",
"p_vertx = vertcross(p, z, pivot_point=pivot_point, angle=angle, latlon=True)\n",
"\n",
"# Extract the numpy array\n",
"p_vert_array = npvalues(p_vertx)\n",

25
test/ipynb/WRF_python_demo.ipynb

@ -565,8 +565,9 @@ @@ -565,8 +565,9 @@
"source": [
"from wrf.latlon import xy_to_ll, ll_to_xy \n",
"\n",
"a = xy_to_ll(ncfile, 400, 200)\n",
"a = xy_to_ll(ncfile, [400,105], [200,205])\n",
"b = ll_to_xy(ncfile, 45.5, -110.8)\n",
"b = ll_to_xy(ncfile, 45.5, -110.8, as_int=True)\n",
"\n",
"# Note: Lists/Dictionaries of files will add a new dimension ('domain') only if the domain is moving\n",
"c = xy_to_ll([ncfile, ncfile, ncfile], [400,105], [200,205])\n",
@ -991,11 +992,29 @@ @@ -991,11 +992,29 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true
"collapsed": false
},
"outputs": [],
"source": [
"\n"
"from wrf.latlon import xy_to_ll, ll_to_xy \n",
"\n",
"a = xy_to_ll(ncfiles, 400, 200)\n",
"a = xy_to_ll(ncfiles, [400,105], [200,205])\n",
"b = ll_to_xy(ncfiles, 45.5, -110.8, as_int=True)\n",
"\n",
"# Note: Lists/Dictionaries of files will add a new dimension ('domain') only if the domain is moving\n",
"c = xy_to_ll(ncfiles, [400,105], [200,205])\n",
"d = xy_to_ll({\"label1\" : ncfiles,\n",
" \"label2\" : ncfiles}, \n",
" [400,105], [200,205])\n",
"\n",
"print(a)\n",
"print(\"\\n\")\n",
"print(b)\n",
"print(\"\\n\")\n",
"print(c)\n",
"print(\"\\n\")\n",
"print(d)\n"
]
},
{

26
test/ncl_get_var.ncl

@ -326,6 +326,32 @@ @@ -326,6 +326,32 @@
filevarattdef(fout,"fld_thetae_pres", fld8_intrp)
fout->fld_thetae_pres = (/fld8_intrp/)
; lat/lon to x/y and x/y to lat/lon routines
lats = (/-55, -60, -65 /)
lons = (/25, 30, 35 /)
i_s = (/10, 100, 150 /)
j_s = (/10, 100, 150 /)
ij = wrf_user_ll_to_ij(input_file, lons, lats, True)
ll = wrf_user_ij_to_ll(input_file, i_s, j_s, True)
ij_dims = dimsizes(ij)
ll_dims = dimsizes(ll)
filedimdef(fout, (/"i_j", "ij_idx"/), \
(/ij_dims(0), ij_dims(1)/), \
(/False,False/))
filedimdef(fout, (/"lat_lon", "ll_idx"/), \
(/ll_dims(0), ll_dims(1)/), \
(/False,False/))
filevardef(fout, "ij", typeof(ij), (/"i_j", "ij_idx"/))
filevardef(fout, "ll", typeof(ll), (/"lat_lon", "ll_idx"/))
filevarattdef(fout,"ij", ij)
filevarattdef(fout,"ll", ll)
fout->ij = (/ij/)
fout->ll = (/ll/)
delete(fout)

58
test/test_multi_cache.py

@ -0,0 +1,58 @@ @@ -0,0 +1,58 @@
import time
from wrf.cache import _get_cache
from wrf import getvar
from netCDF4 import Dataset as nc
#a = nc("/Users/ladwig/Documents/wrf_files/wrf_vortex_single/wrfout_d02_2005-08-28_00:00:00")
#b = nc("/Users/ladwig/Documents/wrf_files/wrf_vortex_single/wrfout_d02_2005-08-28_03:00:00")
a = nc("/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d02_2005-08-28_00:00:00")
b = nc("/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d02_2005-08-28_12:00:00")
q = {"outoutoutout" : {"outoutout" : {"outout" : {"out1" : {"blah" : [a,b], "blah2" : [a,b]}, "out2" : {"blah" : [a,b], "blah2" : [a,b]} } } } }
t1 = time.time()
c = getvar(q, "rh", method="cat", timeidx=None, squeeze=True)
t2 = time.time()
print (c)
print ("time taken: {}".format((t2-t1)*1000.))
t1 = time.time()
c = getvar(q, "rh", method="cat", timeidx=None, squeeze=False)
t2 = time.time()
print (c)
print ("time taken: {}".format((t2-t1)*1000.))
t1 = time.time()
c = getvar(q, "rh", method="cat", timeidx=1, squeeze=True)
t2 = time.time()
print (c)
print ("time taken: {}".format((t2-t1)*1000.))
t1 = time.time()
c = getvar(q, "rh", method="cat", timeidx=1, squeeze=False)
t2 = time.time()
print(c)
print ("time taken: {}".format((t2-t1)*1000.))
t1 = time.time()
c = getvar(q, "rh", method="join", timeidx=None, squeeze=True)
t2 = time.time()
print (c)
print ("time taken: {}".format((t2-t1)*1000.))
t1 = time.time()
c = getvar(q, "rh", method="join", timeidx=None, squeeze=False)
t2 = time.time()
print(c)
print ("time taken: {}".format((t2-t1)*1000.))
t1 = time.time()
c = getvar(q, "rh", method="join", timeidx=1, squeeze=True)
t2 = time.time()
print (c)
print ("time taken: {}".format((t2-t1)*1000.))
t1 = time.time()
c = getvar(q, "rh", method="join", timeidx=1, squeeze=False)
t2 = time.time()
print (c)
print ("time taken: {}".format((t2-t1)*1000.))

120
test/utests.py

@ -6,7 +6,8 @@ import os, sys @@ -6,7 +6,8 @@ import os, sys
import subprocess
from wrf import (getvar, interplevel, interpline, vertcross, vinterp,
disable_xarray, xarray_enabled, npvalues)
disable_xarray, xarray_enabled, npvalues,
xy_to_ll, ll_to_xy )
NCL_EXE = "/Users/ladwig/nclbuild/6.3.0/bin/ncl"
TEST_FILE = "/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00"
@ -457,11 +458,100 @@ def make_interp_test(varname, wrf_in, referent, multi=False, @@ -457,11 +458,100 @@ def make_interp_test(varname, wrf_in, referent, multi=False,
return test
def make_latlon_test(testid, wrf_in, referent, single, multi=False, repeat=3,
pynio=False):
def test(self):
try:
from netCDF4 import Dataset as NetCDF
except:
pass
try:
from PyNIO import Nio
except:
pass
if not multi:
timeidx = 0
if not pynio:
in_wrfnc = NetCDF(wrf_in)
else:
# Note: Python doesn't like it if you reassign an outer scoped
# variable (wrf_in in this case)
if not wrf_in.endswith(".nc"):
wrf_file = wrf_in + ".nc"
else:
wrf_file = wrf_in
in_wrfnc = Nio.open_file(wrf_file)
else:
timeidx = None
if not pynio:
nc = NetCDF(wrf_in)
in_wrfnc = [nc for i in xrange(repeat)]
else:
if not wrf_in.endswith(".nc"):
wrf_file = wrf_in + ".nc"
else:
wrf_file = wrf_in
nc = Nio.open_file(wrf_file)
in_wrfnc = [nc for i in xrange(repeat)]
refnc = NetCDF(referent)
if testid == "xy":
# Since this domain is not moving, the reference values are the
# same whether there are multiple or single files
ref_vals = refnc.variables["ij"][:]
# Lats/Lons taken from NCL script, just hard-coding for now
lats = [-55, -60, -65]
lons = [25, 30, 35]
# Just call with a single lat/lon
if single:
xy = ll_to_xy(in_wrfnc, lats[0], lons[0])
xy = xy + 1 # NCL uses fortran indexing
ref = ref_vals[:,0]
nt.assert_allclose(npvalues(xy), ref)
else:
xy = ll_to_xy(in_wrfnc, lats, lons)
xy = xy + 1 # NCL uses fortran indexing
ref = ref_vals[:]
nt.assert_allclose(npvalues(xy), ref)
else:
# Since this domain is not moving, the reference values are the
# same whether there are multiple or single files
ref_vals = refnc.variables["ll"][:]
# i_s, j_s taken from NCL script, just hard-coding for now
# NCL uses 1-based indexing for this, so need to subtract 1
i_s = np.asarray([10, 100, 150], int) - 1
j_s = np.asarray([10, 100, 150], int) - 1
if single:
ll = xy_to_ll(in_wrfnc, i_s[0], j_s[0])
ref = ref_vals[::-1,0]
nt.assert_allclose(npvalues(ll), ref)
else:
ll = xy_to_ll(in_wrfnc, i_s, j_s)
ref = ref_vals[::-1,:]
nt.assert_allclose(npvalues(ll), ref)
return test
class WRFVarsTest(ut.TestCase):
longMessage = True
class WRFInterpTest(ut.TestCase):
longMessage = True
class WRFLatLonTest(ut.TestCase):
longMessage = True
if __name__ == "__main__":
@ -472,6 +562,7 @@ if __name__ == "__main__": @@ -472,6 +562,7 @@ if __name__ == "__main__":
"theta", "tk", "tv", "twb", "updraft_helicity", "ua", "va",
"wa", "uvmet10", "uvmet", "z", "cfrac"]
interp_methods = ["interplevel", "vertcross", "interpline", "vinterp"]
latlon_tests = ["xy", "ll"]
try:
import netCDF4
@ -497,6 +588,19 @@ if __name__ == "__main__": @@ -497,6 +588,19 @@ if __name__ == "__main__":
setattr(WRFInterpTest, 'test_multi_{0}'.format(method),
test_interp_func2)
for testid in latlon_tests:
for single in (True, False):
for multi in (True, False):
test_ll_func = make_latlon_test(testid, TEST_FILE,
OUT_NC_FILE,
single=single, multi=multi,
repeat=3, pynio=False)
multistr = "" if not multi else "_multi"
singlestr = "_nosingle" if not single else "_single"
test_name = "test_{}{}{}".format(testid, singlestr,
multistr)
setattr(WRFLatLonTest, test_name, test_ll_func)
try:
import PyNIO
except ImportError:
@ -522,6 +626,20 @@ if __name__ == "__main__": @@ -522,6 +626,20 @@ if __name__ == "__main__":
test_interp_func1)
setattr(WRFInterpTest, 'test_pynio_multi_{0}'.format(method),
test_interp_func2)
for testid in latlon_tests:
for single in (True, False):
for multi in (True, False):
test_ll_func = make_latlon_test(testid, TEST_FILE,
OUT_NC_FILE,
single=single, multi=multi,
repeat=3, pynio=False)
multistr = "" if not multi else "_multi"
singlestr = "_nosingle" if not single else "_single"
test_name = "test_pynio_{}{}{}".format(testid,
singlestr,
multistr)
setattr(WRFLatLonTest, test_name, test_ll_func)
ut.main()
Loading…
Cancel
Save