Browse Source

Added projection support. Added configuration to allow cartopy, basemap, xarray, etc to be optional. Began adding metadata support.

main
Bill Ladwig 9 years ago
parent
commit
a887fa1a44
  1. 60
      wrf_open/var/src/python/wrf/var/__init__.py
  2. 4
      wrf_open/var/src/python/wrf/var/cape.py
  3. 81
      wrf_open/var/src/python/wrf/var/config.py
  4. 4
      wrf_open/var/src/python/wrf/var/constants.py
  5. 4
      wrf_open/var/src/python/wrf/var/ctt.py
  6. 2
      wrf_open/var/src/python/wrf/var/dbz.py
  7. 140
      wrf_open/var/src/python/wrf/var/decorators.py
  8. 4
      wrf_open/var/src/python/wrf/var/dewpoint.py
  9. 324
      wrf_open/var/src/python/wrf/var/extension.py
  10. 4
      wrf_open/var/src/python/wrf/var/geoht.py
  11. 20
      wrf_open/var/src/python/wrf/var/helicity.py
  12. 2
      wrf_open/var/src/python/wrf/var/interp.py
  13. 16
      wrf_open/var/src/python/wrf/var/latlon.py
  14. 3
      wrf_open/var/src/python/wrf/var/omega.py
  15. 6
      wrf_open/var/src/python/wrf/var/precip.py
  16. 4
      wrf_open/var/src/python/wrf/var/pressure.py
  17. 748
      wrf_open/var/src/python/wrf/var/projection.py
  18. BIN
      wrf_open/var/src/python/wrf/var/psadlookup.pyc
  19. 4
      wrf_open/var/src/python/wrf/var/pw.py
  20. 4
      wrf_open/var/src/python/wrf/var/rh.py
  21. 2
      wrf_open/var/src/python/wrf/var/slp.py
  22. 13
      wrf_open/var/src/python/wrf/var/temp.py
  23. 4
      wrf_open/var/src/python/wrf/var/terrain.py
  24. 317
      wrf_open/var/src/python/wrf/var/util.py
  25. 30
      wrf_open/var/src/python/wrf/var/uvmet.py
  26. 4
      wrf_open/var/src/python/wrf/var/vorticity.py
  27. 214
      wrf_open/var/test/projtest.py
  28. 2
      wrf_open/var/test/utests.py

60
wrf_open/var/src/python/wrf/var/__init__.py

@ -1,5 +1,7 @@
import warnings import warnings
from config import *
import config
from extension import * from extension import *
import extension import extension
from util import * from util import *
@ -54,35 +56,39 @@ from times import *
import times import times
from units import * from units import *
import units import units
from projection import *
import projection
__all__ = ["getvar"] __all__ = ["getvar"]
__all__ += extension.__all__ __all__.extend(config.__all__)
__all__ += util.__all__ __all__.extend( extension.__all__)
__all__ += cape.__all__ __all__.extend(util.__all__)
__all__ += constants.__all__ __all__.extend(cape.__all__)
__all__ += ctt.__all__ __all__.extend(constants.__all__)
__all__ += dbz.__all__ __all__.extend(ctt.__all__)
__all__ += destag.__all__ __all__.extend(dbz.__all__)
__all__ += dewpoint.__all__ __all__.extend(destag.__all__)
__all__ += etaconv.__all__ __all__.extend(dewpoint.__all__)
__all__ += geoht.__all__ __all__.extend(etaconv.__all__)
__all__ += helicity.__all__ __all__.extend(geoht.__all__)
__all__ += interp.__all__ __all__.extend(helicity.__all__)
__all__ += latlon.__all__ __all__.extend(interp.__all__)
__all__ += omega.__all__ __all__.extend(latlon.__all__)
__all__ += precip.__all__ __all__.extend(omega.__all__)
__all__ += psadlookup.__all__ __all__.extend(precip.__all__)
__all__ += pw.__all__ __all__.extend(psadlookup.__all__)
__all__ += rh.__all__ __all__.extend(pw.__all__)
__all__ += slp.__all__ __all__.extend(rh.__all__)
__all__ += temp.__all__ __all__.extend(slp.__all__)
__all__ += terrain.__all__ __all__.extend(temp.__all__)
__all__ += uvmet.__all__ __all__.extend(terrain.__all__)
__all__ += vorticity.__all__ __all__.extend(uvmet.__all__)
__all__ += wind.__all__ __all__.extend(vorticity.__all__)
__all__ += times.__all__ __all__.extend(wind.__all__)
__all__ += pressure.__all__ __all__.extend(times.__all__)
__all__ += units.__all__ __all__.extend(pressure.__all__)
__all__.extend(units.__all__)
__all__.extend(projection.__all__)
# func is the function to call. kargs are required arguments that should # func is the function to call. kargs are required arguments that should
# not be altered by the user # not be altered by the user

4
wrf_open/var/src/python/wrf/var/cape.py

@ -10,7 +10,7 @@ __all__ = ["get_2dcape", "get_3dcape"]
def get_2dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL): def get_2dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL):
"""Return the 2d fields of cape, cin, lcl, and lfc""" """Return the 2d fields of cape, cin, lcl, and lfc"""
ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", "PH", ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR", "PH",
"PHB", "HGT", "PSFC")) "PHB", "HGT", "PSFC"))
t = ncvars["T"] t = ncvars["T"]
@ -63,7 +63,7 @@ def get_2dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL):
def get_3dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL): def get_3dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL):
"""Return the 3d fields of cape and cin""" """Return the 3d fields of cape and cin"""
ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR",
"PH", "PHB", "HGT", "PSFC")) "PH", "PHB", "HGT", "PSFC"))
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]

81
wrf_open/var/src/python/wrf/var/config.py

@ -0,0 +1,81 @@
try:
from xarray import DataArray
except ImportError:
_XARRAY_ENABLED = False
else:
_XARRAY_ENABLED = True
try:
from cartopy import crs
except ImportError:
_CARTOPY_ENABLED = False
else:
_CARTOPY_ENABLED = True
try:
from mpl_toolkits.basemap import Basemap
except ImportError:
_BASEMAP_ENABLED = False
else:
_BASEMAP_ENABLED = True
try:
from Ngl import Resources
except ImportError:
_PYNGL_ENABLED = False
else:
_PYNGL_ENABLED = True
__all__ = ["xarray_enabled", "disable_xarray", "enable_xarray",
"cartopy_enabled", "enable_cartopy", "enable_cartopy",
"basemap_enabled", "disable_basemap", "enable_basemap",
"pyngl_enabled", "enable_pyngl", "disable_pyngl"]
def xarray_enabled():
global _XARRAY_ENABLED
return _XARRAY_ENABLED
def disable_xarray():
global _XARRAY_ENABLED
_XARRAY_ENABLED = False
def enable_xarray():
global _XARRAY_ENABLED
_XARRAY_ENABLED = True
def cartopy_enabled():
global _CARTOPY_ENABLED
return _CARTOPY_ENABLED
def enable_cartopy():
global _CARTOPY_ENABLED
_CARTOPY_ENABLED = True
def disable_cartopy():
global _CARTOPY_ENABLED
_CARTOPY_ENABLED = True
def basemap_enabled():
global _BASEMAP_ENABLED
return _BASEMAP_ENABLED
def enable_basemap():
global _BASEMAP_ENABLED
_BASEMAP_ENABLED = True
def disable_basemap():
global _BASEMAP_ENABLED
_BASEMAP_ENABLED = True
def pyngl_enabled():
global _PYNGL_ENABLED
return _PYNGL_ENABLED
def enable_pyngl():
global _PYNGL_ENABLED
_PYNGL_ENABLED = True
def disable_pyngl():
global _PYNGL_ENABLED
_PYNGL_ENABLED = True

4
wrf_open/var/src/python/wrf/var/constants.py

@ -7,9 +7,9 @@ class Constants(object):
G = 9.81 G = 9.81
TCK0 = 273.15 TCK0 = 273.15
T_BASE = 300.0 # In WRF the base temperature is always 300 (not var T00) T_BASE = 300.0 # In WRF the base temperature is always 300 (not var T00)
PI = 3.14159265 PI = 3.141592653589793
DEFAULT_FILL = 9.9692099683868690e+36 # Value is from netcdf.h DEFAULT_FILL = 9.9692099683868690e+36 # Value is from netcdf.h
WRF_EARTH_RADIUS = 6370000.
class ConversionFactors(object): class ConversionFactors(object):
PA_TO_HPA = .01 PA_TO_HPA = .01

4
wrf_open/var/src/python/wrf/var/ctt.py

@ -14,8 +14,8 @@ def get_ctt(wrfnc, timeidx=0, units="c"):
"""Return the cloud top temperature. """Return the cloud top temperature.
""" """
ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "PH" ,"PHB", ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "PH" ,
"HGT", "QVAPOR")) "PHB", "HGT", "QVAPOR"))
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
pb = ncvars["PB"] pb = ncvars["PB"]

2
wrf_open/var/src/python/wrf/var/dbz.py

@ -17,7 +17,7 @@ def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False):
as liquid) as liquid)
""" """
ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR",
"QRAIN")) "QRAIN"))
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]

140
wrf_open/var/src/python/wrf/var/decorators.py

@ -7,10 +7,51 @@ import numpy.ma as ma
from wrf.var.units import do_conversion, check_units from wrf.var.units import do_conversion, check_units
from wrf.var.destag import destagger from wrf.var.destag import destagger
from wrf.var.util import iter_left_indexes from wrf.var.util import iter_left_indexes
from wrf.var.config import xarray_enabled
if xarray_enabled():
from xarray import DataArray
__all__ = ["convert_units", "handle_left_iter", "uvmet_left_iter", __all__ = ["convert_units", "handle_left_iter", "uvmet_left_iter",
"handle_casting"] "handle_casting" "set_metadata"]
# TODO: In python 3.x, getargspec is deprecated
class from_args(object):
def __init__(self, argname):
self.argname
def __call__(self, func, *args, **kargs):
"""Parses the function args and kargs looking for the desired argument
value. Otherwise, the value is taken from the default keyword argument
using the arg spec.
"""
# If units are provided to the method call, use them.
# Otherwise, need to parse the argspec to find what the default
# value is since wraps does not preserve this.
argspec = getargspec(func)
if self.argname not in argspec.args and self.argname not in kargs:
return None
try:
result_idx = argspec.args.index(self.argname)
except ValueError:
result_idx = None
if (self.argname in kargs):
result = kargs[self.argname]
elif (len(args) > result_idx and result_idx is not None):
result = args[result_idx]
else:
arg_idx_from_right = (len(argspec.args) -
argspec.args.index(self.argname))
result = argspec.defaults[-arg_idx_from_right]
return result
def convert_units(unit_type, alg_unit): def convert_units(unit_type, alg_unit):
"""A decorator that applies unit conversion to a function's output array. """A decorator that applies unit conversion to a function's output array.
@ -23,25 +64,8 @@ def convert_units(unit_type, alg_unit):
def convert_decorator(func): def convert_decorator(func):
@wraps(func) @wraps(func)
def func_wrapper(*args, **kargs): def func_wrapper(*args, **kargs):
# If units are provided to the method call, use them.
# Otherwise, need to parse the argspec to find what the default desired_units = from_args("units")(func, *args, **kargs)
# value is since wraps does not preserve this.
argspec = getargspec(func)
try:
unit_idx = argspec.args.index("units")
except ValueError:
unit_idx = None
if ("units" in kargs):
desired_units = kargs["units"]
elif (len(args) > unit_idx and unit_idx is not None):
desired_units = args[unit_idx]
else:
arg_idx_from_right = (len(argspec.args)
- argspec.args.index("units"))
desired_units = argspec.defaults[-arg_idx_from_right]
check_units(desired_units, unit_type) check_units(desired_units, unit_type)
# Unit conversion done here # Unit conversion done here
@ -59,7 +83,7 @@ def _calc_out_dims(outvar, left_dims):
def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1,
ref_var_name=None, ref_var_name=None,
ignore_args=(), ignore_kargs=()): ignore_args=None, ignore_kargs=None):
"""Decorator to handle iterating over leftmost dimensions when using """Decorator to handle iterating over leftmost dimensions when using
multiple files and/or multiple times. multiple files and/or multiple times.
@ -84,6 +108,8 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1,
def indexing_decorator(func): def indexing_decorator(func):
@wraps(func) @wraps(func)
def func_wrapper(*args, **kargs): def func_wrapper(*args, **kargs):
ignore_args = ignore_args if ignore_args is not None else ()
ignore_kargs = ignore_kargs if ignore_kargs is not None else ()
if ref_var_idx >= 0: if ref_var_idx >= 0:
ref_var = args[ref_var_idx] ref_var = args[ref_var_idx]
@ -254,7 +280,7 @@ def uvmet_left_iter():
return indexing_decorator return indexing_decorator
def handle_casting(ref_idx=0, arg_idxs=(), karg_names=(),dtype=n.float64): def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=n.float64):
"""Decorator to handle casting to/from required dtype used in """Decorator to handle casting to/from required dtype used in
numerical routine. numerical routine.
@ -262,6 +288,9 @@ def handle_casting(ref_idx=0, arg_idxs=(), karg_names=(),dtype=n.float64):
def cast_decorator(func): def cast_decorator(func):
@wraps(func) @wraps(func)
def func_wrapper(*args, **kargs): def func_wrapper(*args, **kargs):
arg_idxs = arg_idxs if arg_idxs is not None else ()
karg_names = karg_names if karg_names is not None else ()
orig_type = args[ref_idx].dtype orig_type = args[ref_idx].dtype
new_args = [arg.astype(dtype) new_args = [arg.astype(dtype)
@ -271,7 +300,6 @@ def handle_casting(ref_idx=0, arg_idxs=(), karg_names=(),dtype=n.float64):
new_kargs = {key:(val.astype(dtype) new_kargs = {key:(val.astype(dtype)
if key in karg_names else val) if key in karg_names else val)
for key,val in kargs.iteritems()} for key,val in kargs.iteritems()}
res = func(*new_args, **new_kargs) res = func(*new_args, **new_kargs)
@ -288,4 +316,70 @@ def handle_casting(ref_idx=0, arg_idxs=(), karg_names=(),dtype=n.float64):
return cast_decorator return cast_decorator
def _extract_and_transpose(arg):
if xarray_enabled():
if isinstance(arg, DataArray):
arg = arg.values
if isinstance(arg, n.ndarray):
if not arg.flags.F_CONTIGUOUS:
return arg.T
return arg
def handle_extract_transpose():
"""Decorator to extract the data array from a DataArray and also
transposes if the data is not fortran contiguous.
"""
def extract_transpose_decorator(func):
@wraps(func)
def func_wrapper(*args, **kargs):
new_args = [_extract_and_transpose(arg) for arg in args]
new_kargs = {key:_extract_and_transpose(val)
for key,val in kargs.iteritems()}
res = func(*new_args, **new_kargs)
if isinstance(res, n.ndarray):
if res.flags.F_CONTIGUOUS:
return res.T
else:
return tuple(x.T if x.flags.F_CONTIGUOUS else x for x in res)
return res
return func_wrapper
return extract_transpose_decorator
def set_metadata(copy_from=None, ignore=None, **fixed_kargs):
"""Decorator to set the attributes for a WRF method.
"""
def attr_decorator(func):
@wraps(func)
def func_wrapper(*args, **kargs):
if not xarray_enabled():
return func(*args, **kargs)
result = func(*args, **kargs)
units = from_args("units")(func, *args, **kargs)
if units is not None:
result.attrs["units"] = units
for argname, val in fixed_kargs.iteritems():
result[argname] = val
return result
return func_wrapper
return attr_decorator

4
wrf_open/var/src/python/wrf/var/dewpoint.py

@ -7,7 +7,7 @@ __all__ = ["get_dp", "get_dp_2m"]
@convert_units("temp", "c") @convert_units("temp", "c")
def get_dp(wrfnc, timeidx=0, units="c"): def get_dp(wrfnc, timeidx=0, units="c"):
ncvars = extract_vars(wrfnc, timeidx, vars=("P", "PB", "QVAPOR")) ncvars = extract_vars(wrfnc, timeidx, varnames=("P", "PB", "QVAPOR"))
p = ncvars["P"] p = ncvars["P"]
pb = ncvars["PB"] pb = ncvars["PB"]
@ -22,7 +22,7 @@ def get_dp(wrfnc, timeidx=0, units="c"):
@convert_units("temp", "c") @convert_units("temp", "c")
def get_dp_2m(wrfnc, timeidx=0, units="c"): def get_dp_2m(wrfnc, timeidx=0, units="c"):
ncvars = extract_vars(wrfnc, timeidx, vars=("PSFC", "Q2")) ncvars = extract_vars(wrfnc, timeidx, varnames=("PSFC", "Q2"))
# Algorithm requires hPa # Algorithm requires hPa
psfc = .01*(ncvars["PSFC"]) psfc = .01*(ncvars["PSFC"])

324
wrf_open/var/src/python/wrf/var/extension.py

@ -12,7 +12,7 @@ from wrf.var._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d,
f_monotonic, f_filter2d, f_vintrp) f_monotonic, f_filter2d, f_vintrp)
from wrf.var._wrfcape import f_computecape from wrf.var._wrfcape import f_computecape
from wrf.var.decorators import (handle_left_iter, uvmet_left_iter, from wrf.var.decorators import (handle_left_iter, uvmet_left_iter,
handle_casting) handle_casting, handle_extract_transpose)
__all__ = ["FortranException", "computeslp", "computetk", "computetd", __all__ = ["FortranException", "computeslp", "computetk", "computetd",
"computerh", "computeavo", "computepvo", "computeeth", "computerh", "computeavo", "computepvo", "computeeth",
@ -26,21 +26,24 @@ class FortranException(Exception):
@handle_left_iter(3,0, ignore_args=(2,3)) @handle_left_iter(3,0, ignore_args=(2,3))
@handle_casting(arg_idxs=(0,1)) @handle_casting(arg_idxs=(0,1))
@handle_extract_transpose()
def interpz3d(data3d,zdata,desiredloc,missingval): def interpz3d(data3d,zdata,desiredloc,missingval):
res = f_interpz3d(data3d.T, res = f_interpz3d(data3d,
zdata.T, zdata,
desiredloc, desiredloc,
missingval) missingval)
return res.T return res
@handle_left_iter(3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1)) @handle_casting(arg_idxs=(0,1))
@handle_extract_transpose()
def interp2dxy(data3d,xy): def interp2dxy(data3d,xy):
res = f_interp2dxy(data3d.T, res = f_interp2dxy(data3d,
xy.T) xy)
return res.T return res
@handle_casting(arg_idxs=(0,1,2)) @handle_casting(arg_idxs=(0,1,2))
@handle_extract_transpose()
def interp1d(v_in,z_in,z_out,missingval): def interp1d(v_in,z_in,z_out,missingval):
res = f_interp1d(v_in, res = f_interp1d(v_in,
z_in, z_in,
@ -51,113 +54,121 @@ def interp1d(v_in,z_in,z_out,missingval):
@handle_left_iter(3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1,2,3)) @handle_casting(arg_idxs=(0,1,2,3))
@handle_extract_transpose()
def computeslp(z,t,p,q): def computeslp(z,t,p,q):
t_surf = n.zeros((z.shape[-2], z.shape[-1]), "float64") t_surf = n.zeros((z.shape[-2], z.shape[-1]), "float64", order="F")
t_sea_level = n.zeros((z.shape[-2], z.shape[-1]), "float64") t_sea_level = n.zeros((z.shape[-2], z.shape[-1]), "float64", order="F")
level = n.zeros((z.shape[-2], z.shape[-1]), "int32") level = n.zeros((z.shape[-2], z.shape[-1]), "int32", order="F")
res = f_computeslp(z.T, res = f_computeslp(z,
t.T, t,
p.T, p,
q.T, q,
t_sea_level.T, t_sea_level, # Should come in with fortran ordering
t_surf.T, t_surf,
level.T, level,
FortranException()) FortranException())
return res.T return res
@handle_left_iter(3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1)) @handle_casting(arg_idxs=(0,1))
@handle_extract_transpose()
def computetk(pres, theta): def computetk(pres, theta):
# No need to transpose here since operations on 1D array # No need to transpose here since operations on 1D array
shape = pres.shape shape = pres.shape
res = f_computetk(pres.flatten("A"), res = f_computetk(pres.ravel(order="A"),
theta.flatten("A")) theta.ravel(order="A"))
res = n.reshape(res, shape, "A") res = n.reshape(res, shape)
return res return res
# Note: No left iteration decorator needed with 1D arrays # Note: No left iteration decorator needed with 1D arrays
@handle_casting(arg_idxs=(0,1)) @handle_casting(arg_idxs=(0,1))
@handle_extract_transpose()
def computetd(pressure,qv_in): def computetd(pressure,qv_in):
shape = pressure.shape shape = pressure.shape
res = f_computetd(pressure.flatten("A"), res = f_computetd(pressure.ravel(order="A"),
qv_in.flatten("A")) qv_in.ravel(order="A"))
res = n.reshape(res, shape, "A") res = n.reshape(res, shape)
return res return res
# Note: No decorator needed with 1D arrays # Note: No decorator needed with 1D arrays
@handle_casting(arg_idxs=(0,1,2)) @handle_casting(arg_idxs=(0,1,2))
@handle_extract_transpose()
def computerh(qv,q,t): def computerh(qv,q,t):
shape = qv.shape shape = qv.shape
res = f_computerh(qv.flatten("A"), res = f_computerh(qv.ravel(order="A"),
q.flatten("A"), q.ravel(order="A"),
t.flatten("A")) t.ravel(order="A"))
res = n.reshape(res, shape, "A") res = n.reshape(res, shape)
return res return res
@handle_left_iter(3,0, ignore_args=(6,7)) @handle_left_iter(3,0, ignore_args=(6,7))
@handle_casting(arg_idxs=(0,1,2,3,4,5)) @handle_casting(arg_idxs=(0,1,2,3,4,5))
@handle_extract_transpose()
def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy): def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy):
res = f_computeabsvort(u.T, res = f_computeabsvort(u,
v.T, v,
msfu.T, msfu,
msfv.T, msfv,
msfm.T, msfm,
cor.T, cor,
dx, dx,
dy) dy)
return res.T return res
@handle_left_iter(3,2, ignore_args=(8,9)) @handle_left_iter(3,2, ignore_args=(8,9))
@handle_casting(arg_idxs=(0,1,2,3,4,5,6,7)) @handle_casting(arg_idxs=(0,1,2,3,4,5,6,7))
@handle_extract_transpose()
def computepvo(u,v,theta,prs,msfu,msfv,msfm,cor,dx,dy): def computepvo(u,v,theta,prs,msfu,msfv,msfm,cor,dx,dy):
res = f_computepvo(u.T, res = f_computepvo(u,
v.T, v,
theta.T, theta,
prs.T, prs,
msfu.T, msfu,
msfv.T, msfv,
msfm.T, msfm,
cor.T, cor,
dx, dx,
dy) dy)
return res.T return res
@handle_left_iter(3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1,2)) @handle_casting(arg_idxs=(0,1,2))
@handle_extract_transpose()
def computeeth(qv, tk, p): def computeeth(qv, tk, p):
res = f_computeeth(qv.T, res = f_computeeth(qv,
tk.T, tk,
p.T) p)
return res.T return res
@uvmet_left_iter() @uvmet_left_iter()
@handle_casting(arg_idxs=(0,1,2,3)) @handle_casting(arg_idxs=(0,1,2,3))
@handle_extract_transpose()
def computeuvmet(u,v,lat,lon,cen_long,cone): def computeuvmet(u,v,lat,lon,cen_long,cone):
longca = n.zeros((lat.shape[-2], lat.shape[-1]), "float64") longca = n.zeros((lat.shape[-2], lat.shape[-1]), "float64", order="F")
longcb = n.zeros((lon.shape[-2], lon.shape[-1]), "float64") longcb = n.zeros((lon.shape[-2], lon.shape[-1]), "float64", order="F")
rpd = Constants.PI/180. rpd = Constants.PI/180.
# Make the 2D array a 3D array with 1 dimension # Make the 2D array a 3D array with 1 dimension
if u.ndim < 3: if u.ndim < 3:
u = u.reshape((1,u.shape[-2], u.shape[-1])) u = u.reshape((u.shape[-2], u.shape[-1], 1), order="F")
v = v.reshape((1,v.shape[-2], v.shape[-1])) v = v.reshape((v.shape[-2], v.shape[-1], 1), order="F")
# istag will always be false since winds are destaggered already # istag will always be false since winds are destaggered already
# Missing values don't appear to be used, so setting to 0 # Missing values don't appear to be used, so setting to 0
res = f_computeuvmet(u.T, res = f_computeuvmet(u,
v.T, v,
longca.T, longca,
longcb.T, longcb,
lon.T, lon,
lat.T, lat,
cen_long, cen_long,
cone, cone,
rpd, rpd,
@ -168,113 +179,123 @@ def computeuvmet(u,v,lat,lon,cen_long,cone):
0) 0)
return n.squeeze(res.T) return n.squeeze(res)
@handle_left_iter(3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1,2,3)) @handle_casting(arg_idxs=(0,1,2,3))
@handle_extract_transpose()
def computeomega(qv, tk, w, p): def computeomega(qv, tk, w, p):
res = f_computeomega(qv.T, res = f_computeomega(qv,
tk.T, tk,
w.T, w,
p.T) p)
#return res.T #return res.T
return res.T return res
@handle_left_iter(3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1)) @handle_casting(arg_idxs=(0,1))
@handle_extract_transpose()
def computetv(tk,qv): def computetv(tk,qv):
res = f_computetv(tk.T, res = f_computetv(tk,
qv.T) qv)
return res.T return res
@handle_left_iter(3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1,2)) @handle_casting(arg_idxs=(0,1,2))
@handle_extract_transpose()
def computewetbulb(p,tk,qv): def computewetbulb(p,tk,qv):
PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables() PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables()
PSADITMK = PSADITMK.T
res = f_computewetbulb(p.T, res = f_computewetbulb(p,
tk.T, tk,
qv.T, qv,
PSADITHTE, PSADITHTE,
PSADIPRS, PSADIPRS,
PSADITMK.T, PSADITMK,
FortranException()) FortranException())
return res.T return res
@handle_left_iter(3,0, ignore_args=(4,)) @handle_left_iter(3,0, ignore_args=(4,))
@handle_casting(arg_idxs=(0,1,2,3)) @handle_casting(arg_idxs=(0,1,2,3))
@handle_extract_transpose()
def computesrh(u, v, z, ter, top): def computesrh(u, v, z, ter, top):
res = f_computesrh(u.T, res = f_computesrh(u,
v.T, v,
z.T, z,
ter.T, ter,
top) top)
return res.T return res
@handle_left_iter(3,2, ignore_args=(5,6,7,8)) @handle_left_iter(3,2, ignore_args=(5,6,7,8))
@handle_casting(arg_idxs=(0,1,2,3,4)) @handle_casting(arg_idxs=(0,1,2,3,4))
@handle_extract_transpose()
def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top): def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top):
tem1 = n.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), "float64") tem1 = n.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), "float64", order="F")
tem2 = n.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), "float64") tem2 = n.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), "float64", order="F")
res = f_computeuh(zp.T, res = f_computeuh(zp,
mapfct.T, mapfct,
dx, dx,
dy, dy,
bottom, bottom,
top, top,
u.T, u,
v.T, v,
wstag.T, wstag,
tem1.T, tem1,
tem2.T) tem2)
return res.T return res
@handle_left_iter(3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1,2,3)) @handle_casting(arg_idxs=(0,1,2,3))
@handle_extract_transpose()
def computepw(p,tv,qv,ht): def computepw(p,tv,qv,ht):
# Note, dim -3 is height, we only want y and x # Note, dim -3 is height, we only want y and x
zdiff = n.zeros((p.shape[-2], p.shape[-1]), "float64") zdiff = n.zeros((p.shape[-2], p.shape[-1]), "float64", order="F")
res = f_computepw(p.T, res = f_computepw(p,
tv.T, tv,
qv.T, qv,
ht.T, ht,
zdiff.T) zdiff)
return res.T return res
@handle_left_iter(3,0, ignore_args=(6,7,8)) @handle_left_iter(3,0, ignore_args=(6,7,8))
@handle_casting(arg_idxs=(0,1,2,3,4,5)) @handle_casting(arg_idxs=(0,1,2,3,4,5))
@handle_extract_transpose()
def computedbz(p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin): def computedbz(p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin):
res = f_computedbz(p.T, res = f_computedbz(p,
tk.T, tk,
qv.T, qv,
qr.T, qr,
qs.T, qs,
qg.T, qg,
sn0, sn0,
ivarint, ivarint,
iliqskin) iliqskin)
return res.T return res
@handle_left_iter(3,0,ignore_args=(6,7,8)) @handle_left_iter(3,0,ignore_args=(6,7,8))
@handle_casting(arg_idxs=(0,1,2,3,4,5)) @handle_casting(arg_idxs=(0,1,2,3,4,5))
@handle_extract_transpose()
def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow): def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow):
flip_cape = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), flip_cape = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]),
"float64") "float64", order="F")
flip_cin = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), flip_cin = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]),
"float64") "float64", order="F")
PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables() PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables()
PSADITMK = PSADITMK.T
# The fortran routine needs pressure to be ascending in z-direction, # The fortran routine needs pressure to be ascending in z-direction,
# along with tk,qv,and ht. # along with tk,qv,and ht.
@ -283,17 +304,17 @@ def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow):
flip_qv = qv[::-1,:,:] flip_qv = qv[::-1,:,:]
flip_ht = ht[::-1,:,:] flip_ht = ht[::-1,:,:]
f_computecape(flip_p.T, f_computecape(flip_p,
flip_tk.T, flip_tk,
flip_qv.T, flip_qv,
flip_ht.T, flip_ht,
ter.T, ter,
sfp.T, sfp,
flip_cape.T, flip_cape,
flip_cin.T, flip_cin,
PSADITHTE, PSADITHTE,
PSADIPRS, PSADIPRS,
PSADITMK.T, PSADITMK,
missing, missing,
i3dflag, i3dflag,
ter_follow, ter_follow,
@ -304,7 +325,6 @@ def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow):
cape = flip_cape[::-1,:,:] cape = flip_cape[::-1,:,:]
cin = flip_cin[::-1,:,:] cin = flip_cin[::-1,:,:]
return (cape, cin) return (cape, cin)
# TODO: This should handle lists of coords # TODO: This should handle lists of coords
@ -333,37 +353,40 @@ def computell(map_proj,truelat1,truelat2,stdlon,lat1,lon1,
@handle_left_iter(3,0, ignore_args=(3,)) @handle_left_iter(3,0, ignore_args=(3,))
@handle_casting(arg_idxs=(0,1,2)) @handle_casting(arg_idxs=(0,1,2))
@handle_extract_transpose()
def computeeta(full_t, znu, psfc, ptop): def computeeta(full_t, znu, psfc, ptop):
pcalc = n.zeros(full_t.shape, "float64") pcalc = n.zeros(full_t.shape, "float64", order="F")
mean_t = n.zeros(full_t.shape, "float64") mean_t = n.zeros(full_t.shape, "float64", order="F")
temp_t = n.zeros(full_t.shape, "float64") temp_t = n.zeros(full_t.shape, "float64", order="F")
res = f_converteta(full_t.T, res = f_converteta(full_t,
znu.T, znu,
psfc.T, psfc,
ptop, ptop,
pcalc.T, pcalc,
mean_t.T, mean_t,
temp_t.T) temp_t)
return res.T return res
@handle_left_iter(3,0,ignore_args=(7,)) @handle_left_iter(3,0,ignore_args=(7,))
@handle_casting(arg_idxs=(0,1,2,3,4,5,6)) @handle_casting(arg_idxs=(0,1,2,3,4,5,6))
@handle_extract_transpose()
def computectt(p_hpa,tk,qice,qcld,qv,ght,ter,haveqci): def computectt(p_hpa,tk,qice,qcld,qv,ght,ter,haveqci):
res = f_computectt(p_hpa.T, res = f_computectt(p_hpa,
tk.T, tk,
qice.T, qice,
qcld.T, qcld,
qv.T, qv,
ght.T, ght,
ter.T, ter,
haveqci) haveqci)
return res.T return res
@handle_left_iter(2,0,ignore_args=(1,)) @handle_left_iter(2,0,ignore_args=(1,))
@handle_casting(arg_idxs=(0,)) @handle_casting(arg_idxs=(0,))
@handle_extract_transpose()
def smooth2d(field, passes): def smooth2d(field, passes):
# Unlike NCL, this routine will not modify the values in place, but # Unlike NCL, this routine will not modify the values in place, but
# copies the original data before modifying it. This allows the decorator # copies the original data before modifying it. This allows the decorator
@ -375,10 +398,10 @@ def smooth2d(field, passes):
missing = Constants.DEFAULT_FILL missing = Constants.DEFAULT_FILL
field_copy = field.copy() field_copy = field.copy()
field_tmp = n.zeros(field_copy.shape, field_copy.dtype) field_tmp = n.zeros(field_copy.shape, field_copy.dtype, order="F")
f_filter2d(field_copy.T, f_filter2d(field_copy,
field_tmp.T, field_tmp,
missing, missing,
passes) passes)
@ -388,32 +411,33 @@ def smooth2d(field, passes):
@handle_left_iter(3,0,ignore_args=(3,4,5)) @handle_left_iter(3,0,ignore_args=(3,4,5))
@handle_casting(arg_idxs=(0,1,2)) @handle_casting(arg_idxs=(0,1,2))
@handle_extract_transpose()
def monotonic(var,lvprs,coriolis,idir,delta,icorsw): def monotonic(var,lvprs,coriolis,idir,delta,icorsw):
res = f_monotonic(var.T, res = f_monotonic(var,
lvprs.T, lvprs,
coriolis.T, coriolis,
idir, idir,
delta, delta,
icorsw) icorsw)
return res.T return res
# TODO: Check those decorators, they might be wrong
@handle_left_iter(3,0,ignore_args=(9,10,11,12,13,14)) @handle_left_iter(3,0,ignore_args=(9,10,11,12,13,14))
@handle_casting(arg_idxs=(0,1,2,3,4,5,6,7,8,9)) @handle_casting(arg_idxs=(0,1,2,3,4,5,6,7,8,9))
@handle_extract_transpose()
def vintrp(field,pres,tk,qvp,ght,terrain,sfp,smsfp, def vintrp(field,pres,tk,qvp,ght,terrain,sfp,smsfp,
vcarray,interp_levels,icase,extrap,vcor,logp, vcarray,interp_levels,icase,extrap,vcor,logp,
missing): missing):
res = f_vintrp(field.T, res = f_vintrp(field,
pres.T, pres,
tk.T, tk,
qvp.T, qvp,
ght.T, ght,
terrain.T, terrain,
sfp.T, sfp,
smsfp.T, smsfp,
vcarray.T, vcarray,
interp_levels, interp_levels,
icase, icase,
extrap, extrap,
@ -421,7 +445,7 @@ def vintrp(field,pres,tk,qvp,ght,terrain,sfp,smsfp,
logp, logp,
missing) missing)
return res.T return res

4
wrf_open/var/src/python/wrf/var/geoht.py

@ -14,10 +14,10 @@ def _get_geoht(wrfnc, timeidx, height=True, msl=True):
""" """
try: try:
ph_vars = extract_vars(wrfnc, timeidx, ("PH", "PHB", "HGT")) ph_vars = extract_vars(wrfnc, timeidx, varnames=("PH", "PHB", "HGT"))
except KeyError: except KeyError:
try: try:
ght_vars = extract_vars(wrfnc, timeidx, ("GHT", "HGT_U")) ght_vars = extract_vars(wrfnc, timeidx, varnames=("GHT", "HGT_U"))
except KeyError: except KeyError:
raise RuntimeError("Cannot calculate height with variables in " raise RuntimeError("Cannot calculate height with variables in "
"NetCDF file") "NetCDF file")

20
wrf_open/var/src/python/wrf/var/helicity.py

@ -9,17 +9,17 @@ __all__ = ["get_srh", "get_uh"]
def get_srh(wrfnc, timeidx=0, top=3000.0): def get_srh(wrfnc, timeidx=0, top=3000.0):
# Top can either be 3000 or 1000 (for 0-1 srh or 0-3 srh) # Top can either be 3000 or 1000 (for 0-1 srh or 0-3 srh)
ncvars = extract_vars(wrfnc, timeidx, vars=("HGT", "PH", "PHB")) ncvars = extract_vars(wrfnc, timeidx, varnames=("HGT", "PH", "PHB"))
ter = ncvars["HGT"] ter = ncvars["HGT"]
ph = ncvars["PH"] ph = ncvars["PH"]
phb = ncvars["PHB"] phb = ncvars["PHB"]
try: try:
u_vars = extract_vars(wrfnc, timeidx, vars="U") u_vars = extract_vars(wrfnc, timeidx, varnames="U")
except KeyError: except KeyError:
try: try:
uu_vars = extract_vars(wrfnc, timeidx, vars="UU") uu_vars = extract_vars(wrfnc, timeidx, varnames="UU")
except KeyError: except KeyError:
raise RuntimeError("No valid wind data found in NetCDF file") raise RuntimeError("No valid wind data found in NetCDF file")
else: else:
@ -28,10 +28,10 @@ def get_srh(wrfnc, timeidx=0, top=3000.0):
u = destagger(u_vars["U"], -1) u = destagger(u_vars["U"], -1)
try: try:
v_vars = extract_vars(wrfnc, timeidx, vars="V") v_vars = extract_vars(wrfnc, timeidx, varnames="V")
except KeyError: except KeyError:
try: try:
vv_vars = extract_vars(wrfnc, timeidx, vars="VV") vv_vars = extract_vars(wrfnc, timeidx, varnames="VV")
except KeyError: except KeyError:
raise RuntimeError("No valid wind data found in NetCDF file") raise RuntimeError("No valid wind data found in NetCDF file")
else: else:
@ -55,7 +55,7 @@ def get_srh(wrfnc, timeidx=0, top=3000.0):
def get_uh(wrfnc, timeidx=0, bottom=2000.0, top=5000.0): def get_uh(wrfnc, timeidx=0, bottom=2000.0, top=5000.0):
ncvars = extract_vars(wrfnc, timeidx, vars=("W", "PH", "PHB", "MAPFAC_M")) ncvars = extract_vars(wrfnc, timeidx, varnames=("W", "PH", "PHB", "MAPFAC_M"))
wstag = ncvars["W"] wstag = ncvars["W"]
ph = ncvars["PH"] ph = ncvars["PH"]
@ -67,10 +67,10 @@ def get_uh(wrfnc, timeidx=0, bottom=2000.0, top=5000.0):
dy = attrs["DY"] dy = attrs["DY"]
try: try:
u_vars = extract_vars(wrfnc, timeidx, vars="U") u_vars = extract_vars(wrfnc, timeidx, varnames="U")
except KeyError: except KeyError:
try: try:
uu_vars = extract_vars(wrfnc, timeidx, vars="UU") uu_vars = extract_vars(wrfnc, timeidx, varnames="UU")
except KeyError: except KeyError:
raise RuntimeError("No valid wind data found in NetCDF file") raise RuntimeError("No valid wind data found in NetCDF file")
else: else:
@ -79,10 +79,10 @@ def get_uh(wrfnc, timeidx=0, bottom=2000.0, top=5000.0):
u = destagger(u_vars["U"], -1) u = destagger(u_vars["U"], -1)
try: try:
v_vars = extract_vars(wrfnc, timeidx, vars="V") v_vars = extract_vars(wrfnc, timeidx, varnames="V")
except KeyError: except KeyError:
try: try:
vv_vars = extract_vars(wrfnc, timeidx, vars="VV") vv_vars = extract_vars(wrfnc, timeidx, varnames="VV")
except KeyError: except KeyError:
raise RuntimeError("No valid wind data found in NetCDF file") raise RuntimeError("No valid wind data found in NetCDF file")
else: else:

2
wrf_open/var/src/python/wrf/var/interp.py

@ -324,7 +324,7 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False,
# Extract vriables # Extract vriables
timeidx = -1 # Should this be an argument? timeidx = -1 # Should this be an argument?
ncvars = extract_vars(wrfnc, timeidx, vars=("PSFC", "QVAPOR", "F")) ncvars = extract_vars(wrfnc, timeidx, varnames=("PSFC", "QVAPOR", "F"))
sfp = ncvars["PSFC"] * ConversionFactors.PA_TO_HPA sfp = ncvars["PSFC"] * ConversionFactors.PA_TO_HPA
qv = ncvars["QVAPOR"] qv = ncvars["QVAPOR"]

16
wrf_open/var/src/python/wrf/var/latlon.py

@ -1,5 +1,6 @@
from collections import Iterable from collections import Iterable
from wrf.var.constants import Constants
from wrf.var.extension import computeij, computell from wrf.var.extension import computeij, computell
from wrf.var.util import extract_vars, extract_global_attrs from wrf.var.util import extract_vars, extract_global_attrs
@ -8,10 +9,10 @@ __all__ = ["get_lat", "get_lon", "get_ij", "get_ll"]
def get_lat(wrfnc, timeidx=0): def get_lat(wrfnc, timeidx=0):
try: try:
lat_vars = extract_vars(wrfnc, timeidx, vars="XLAT") lat_vars = extract_vars(wrfnc, timeidx, varnames="XLAT")
except KeyError: except KeyError:
try: try:
latm_vars = extract_vars(wrfnc, timeidx, vars="XLAT_M") latm_vars = extract_vars(wrfnc, timeidx, varnames="XLAT_M")
except: except:
raise RuntimeError("Latitude variable not found in NetCDF file") raise RuntimeError("Latitude variable not found in NetCDF file")
else: else:
@ -23,10 +24,10 @@ def get_lat(wrfnc, timeidx=0):
def get_lon(wrfnc, timeidx=0): def get_lon(wrfnc, timeidx=0):
try: try:
lon_vars = extract_vars(wrfnc, timeidx, vars="XLONG") lon_vars = extract_vars(wrfnc, timeidx, varnames="XLONG")
except KeyError: except KeyError:
try: try:
lonm_vars = extract_vars(wrfnc, timeidx, vars="XLONG_M") lonm_vars = extract_vars(wrfnc, timeidx, varnames="XLONG_M")
except: except:
raise RuntimeError("Latitude variable not found in NetCDF file") raise RuntimeError("Latitude variable not found in NetCDF file")
else: else:
@ -42,22 +43,21 @@ def _get_proj_params(wrfnc, timeidx):
attrs = extract_global_attrs(wrfnc, attrs=("MAP_PROJ", "TRUELAT1", attrs = extract_global_attrs(wrfnc, attrs=("MAP_PROJ", "TRUELAT1",
"TRUELAT2", "STAND_LON", "TRUELAT2", "STAND_LON",
"DX", "DY", "STAND_LON")) "DX", "DY"))
map_proj = attrs["MAP_PROJ"] map_proj = attrs["MAP_PROJ"]
truelat1 = attrs["TRUELAT1"] truelat1 = attrs["TRUELAT1"]
truelat2 = attrs["TRUELAT2"] truelat2 = attrs["TRUELAT2"]
stdlon = attrs["STAND_LON"] stdlon = attrs["STAND_LON"]
dx = attrs["DX"] dx = attrs["DX"]
dy = attrs["DY"] dy = attrs["DY"]
stdlon = attrs["STAND_LON"]
if map_proj == 6: if map_proj == 6:
pole_attrs = extract_global_attrs(wrfnc, attrs=("POLE_LAT", pole_attrs = extract_global_attrs(wrfnc, attrs=("POLE_LAT",
"POLE_LON")) "POLE_LON"))
pole_lat = pole_attrs["POLE_LAT"] pole_lat = pole_attrs["POLE_LAT"]
pole_lon = pole_attrs["POLE_LON"] pole_lon = pole_attrs["POLE_LON"]
latinc = (dy*360.0)/2.0/3.141592653589793/6370000. latinc = (dy*360.0)/2.0 / Constants.PI/Constants.WRF_EARTH_RADIUS
loninc = (dx*360.0)/2.0/3.141592653589793/6370000. loninc = (dx*360.0)/2.0 / Constants.PI/Constants.WRF_EARTH_RADIUS
else: else:
pole_lat = 90.0 pole_lat = 90.0
pole_lon = 0.0 pole_lon = 0.0

3
wrf_open/var/src/python/wrf/var/omega.py

@ -7,7 +7,8 @@ from wrf.var.util import extract_vars
__all__ = ["get_omega"] __all__ = ["get_omega"]
def get_omega(wrfnc, timeidx=0): def get_omega(wrfnc, timeidx=0):
ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "W", "PB", "QVAPOR")) ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "W",
"PB", "QVAPOR"))
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
w = ncvars["W"] w = ncvars["W"]

6
wrf_open/var/src/python/wrf/var/precip.py

@ -5,7 +5,7 @@ from wrf.var.util import extract_vars
__all__ = ["get_accum_precip", "get_precip_diff"] __all__ = ["get_accum_precip", "get_precip_diff"]
def get_accum_precip(wrfnc, timeidx=0): def get_accum_precip(wrfnc, timeidx=0):
ncvars = extract_vars(wrfnc, timeidx, vars=("RAINC", "RAINNC")) ncvars = extract_vars(wrfnc, timeidx, varnames=("RAINC", "RAINNC"))
rainc = ncvars["RAINC"] rainc = ncvars["RAINC"]
rainnc = ncvars["RAINNC"] rainnc = ncvars["RAINNC"]
@ -14,8 +14,8 @@ def get_accum_precip(wrfnc, timeidx=0):
return rainsum return rainsum
def get_precip_diff(wrfnc1, wrfnc2, timeidx=0): def get_precip_diff(wrfnc1, wrfnc2, timeidx=0):
vars1 = extract_vars(wrfnc1, timeidx, vars=("RAINC", "RAINNC")) vars1 = extract_vars(wrfnc1, timeidx, varnames=("RAINC", "RAINNC"))
vars2 = extract_vars(wrfnc2, timeidx, vars=("RAINC", "RAINNC")) vars2 = extract_vars(wrfnc2, timeidx, varnames=("RAINC", "RAINNC"))
rainc1 = vars1["RAINC"] rainc1 = vars1["RAINC"]
rainnc1 = vars1["RAINNC"] rainnc1 = vars1["RAINNC"]

4
wrf_open/var/src/python/wrf/var/pressure.py

@ -8,10 +8,10 @@ __all__ = ["get_pressure", "get_pressure_hpa"]
def get_pressure(wrfnc, timeidx=0, units="pa"): def get_pressure(wrfnc, timeidx=0, units="pa"):
try: try:
p_vars = extract_vars(wrfnc, timeidx, vars=("P", "PB")) p_vars = extract_vars(wrfnc, timeidx, varnames=("P", "PB"))
except KeyError: except KeyError:
try: try:
pres_vars = extract_vars(wrfnc, timeidx, vars="PRES") pres_vars = extract_vars(wrfnc, timeidx, varnames="PRES")
except: except:
raise RuntimeError("pressure variable not found in NetCDF file") raise RuntimeError("pressure variable not found in NetCDF file")
else: else:

748
wrf_open/var/src/python/wrf/var/projection.py

@ -0,0 +1,748 @@
import numpy as np
import math
from wrf.var.config import basemap_enabled, cartopy_enabled, pyngl_enabled
from wrf.var.constants import Constants
if cartopy_enabled():
from cartopy import crs
if basemap_enabled():
from mpl_toolkits.basemap import Basemap
if pyngl_enabled():
from Ngl import Resources
__all__ = ["WrfProj", "LambertConformalProj", "MercatorProj",
"PolarStereographicProj", "LatLonProj", "RotLatLonProj",
"getproj"]
if cartopy_enabled():
class MercatorWithLatTsProj(crs.Mercator):
def __init__(self, central_longitude=0.0,
latitude_true_scale=0.0,
min_latitude=-80.0,
max_latitude=84.0,
globe=None):
proj4_params = [("proj", "merc"),
("lon_0", central_longitude),
("lat_ts", latitude_true_scale),
("k", 1),
("units", "m")]
super(crs.Mercator, self).__init__(proj4_params, globe=globe)
# Calculate limits.
limits = self.transform_points(crs.Geodetic(),
np.array([-180, 180]) + central_longitude,
np.array([min_latitude, max_latitude]))
# When using a latitude of true scale, the min/max x-limits get set
# to the same value, so make sure the left one is negative
xlimits = limits[..., 0]
if xlimits[0] == xlimits[1]:
if xlimits[0] < 0:
xlimits[1] = -xlimits[1]
else:
xlimits[0] = -xlimits[0]
self._xlimits = tuple(xlimits)
self._ylimits = tuple(limits[..., 1])
self._threshold = np.diff(self.x_limits)[0] / 720
def _ismissing(val):
return val is None or val > 90. or val < -90.
class WrfProj(object):
def __init__(self, bottom_left=None, top_right=None,
lats=None, lons=None, **proj_params):
if bottom_left is not None and top_right is not None:
self.ll_lat = bottom_left[0]
self.ll_lon = bottom_left[1]
self.ur_lat = top_right[0]
self.ur_lon = top_right[1]
self.bottom_left = bottom_left
self.top_right = top_right
elif lats is not None and lons is not None:
self.ll_lat = lats[0,0]
self.ur_lat = lats[-1,-1]
self.ll_lon = lons[0,0]
self.ur_lon = lons[-1,-1]
self.bottom_left = [self.ll_lat, self.ll_lon]
self.top_right = [self.ur_lat, self.ur_lon]
else:
raise ValueError("invalid corner point arguments")
# These indicate the center of the nest/domain, not necessarily the
# center of the projection
self._cen_lat = proj_params.get("CEN_LAT", None)
self._cen_lon = proj_params.get("CEN_LON", None)
self.truelat1 = proj_params.get("TRUELAT1", None)
self.truelat2 = (proj_params.get("TRUELAT2", None)
if not _ismissing(proj_params.get("TRUELAT2", None))
else None)
self.moad_cen_lat = proj_params.get("MOAD_CEN_LAT", None)
self.stand_lon = proj_params.get("STAND_LON", None)
self.pole_lat = proj_params.get("POLE_LAT", None)
self.pole_lon = proj_params.get("POLE_LON", None)
# Just in case...
if self.moad_cen_lat is None:
self.moad_cen_lat = self._cen_lat
if self.stand_lon is None:
self.stand_lon = self._cen_lon
@property
def _basemap(self):
return None
@property
def _cf_params(self):
return None
@property
def _cartopy(self):
return None
@property
def _cart_extents(self):
return ([self.ll_lon, self.ur_lon], [self.ll_lat, self.ur_lat])
@property
def _pyngl(self):
return None
@property
def _proj4(self):
return None
@property
def _globe(self):
return (None if not cartopy_enabled()
else crs.Globe(ellipse=None,
semimajor_axis=Constants.WRF_EARTH_RADIUS,
semiminor_axis=Constants.WRF_EARTH_RADIUS))
def cartopy_xlim(self):
"""Return the x extents in projected coordinates (for cartopy)"""
return self._cart_extents[0]
def cartopy_ylim(self):
"""Return the y extents in projected coordinates (for cartopy)"""
return self._cart_extents[1]
def __repr__(self):
args = ("bottom_left={}, top_right={}, "
"stand_lon={}, moad_cen_lat={}, "
"pole_lat={}, pole_lon={}".format((self.ll_lat, self.ll_lon),
(self.ur_lat, self.ur_lon),
self.stand_lon,
self.moad_cen_lat,
self.pole_lat,
self.pole_lon))
return "{}({})".format(self.__class__.__name__, args)
def basemap(self):
"""Return a mpl_toolkits.basemap.Basemap instance for the
projection"""
if not basemap_enabled():
raise RuntimeError("'mpl_toolkits.basemap' is not "
"installed or is disabled")
return self._basemap
def cartopy(self):
"""Return a cartopy.crs.Projection subclass for the
projection"""
if not cartopy_enabled():
raise RuntimeError("'cartopy' is not "
"installed or is disabled")
return self._cartopy
def pyngl(self):
"""Return the PyNGL resources for the projection"""
if not pyngl_enabled():
raise RuntimeError("'pyngl' is not "
"installed or is disabled")
return self._pyngl
def proj4(self):
"""Return the proj4 string for the map projection"""
return self._proj4
def cf(self):
"""Return a dictionary with the NetCDF CF parameters for the
projection"""
return self._cf_params
class LambertConformalProj(WrfProj):
def __init__(self, bottom_left=None, top_right=None,
lats=None, lons=None, **proj_params):
super(LambertConformalProj, self).__init__(bottom_left,
top_right, lats, lons, **proj_params)
self._std_parallels = [self.truelat1]
if self.truelat2 is not None:
self._std_parallels.append(self.truelat2)
@property
def _cf_params(self):
_cf_params = {}
_cf_params["grid_mapping_name"] = "lambert_conformal_conic";
_cf_params["standard_parallel"] = self._std_parallels
_cf_params["longitude_of_central_meridian"] = self.stand_lon
_cf_params["latitude_of_projection_origin"] = self.moad_cen_lat
_cf_params["semi_major_axis"] = Constants.WRF_EARTH_RADIUS
return _cf_params
@property
def _pyngl(self):
if not pyngl_enabled():
return None
truelat2 = (self.truelat1
if _ismissing(self.truelat2)
else self.truelat2)
_pyngl = Resources()
_pyngl.mpProjection = "LambertConformal"
_pyngl.mpDataBaseVersion = "MediumRes"
_pyngl.mpLimitMode = "Corners"
_pyngl.mpLeftCornerLonF = self.ll_lon
_pyngl.mpLeftCornerLatF = self.ll_lat
_pyngl.mpRightCornerLonF = self.ur_lon
_pyngl.mpRightCornerLatF = self.ur_lat
_pyngl.mpLambertMeridianF = self.stand_lon
_pyngl.mpLambertParallel1F = self.truelat1
_pyngl.mpLambertParallel2F = truelat2
return _pyngl
@property
def _basemap(self):
if not basemap_enabled():
return None
_basemap = Basemap(projection = "lcc",
lon_0 = self.stand_lon,
lat_0 = self.moad_cen_lat,
lat_1 = self.truelat1,
lat_2 = self.truelat2,
llcrnrlat = self.ll_lat,
urcrnrlat = self.ur_lat,
llcrnrlon = self.ll_lon,
urcrnrlon = self.ur_lon,
rsphere = Constants.WRF_EARTH_RADIUS,
resolution = 'l')
return _basemap
@property
def _cartopy(self):
if not cartopy_enabled():
return None
_cartopy = crs.LambertConformal(
central_longitude = self.stand_lon,
central_latitude = self.moad_cen_lat,
standard_parallels = self._std_parallels,
globe = self._globe)
return _cartopy
@property
def _cart_extents(self):
# Need to modify the extents for the new projection
pc = crs.PlateCarree()
xs, ys, zs = self._cartopy.transform_points(pc,
np.array([self.ll_lon, self.ur_lon]),
np.array([self.ll_lat, self.ur_lat])).T
_xlimits = xs.tolist()
_ylimits = ys.tolist()
return (_xlimits, _ylimits)
@property
def _proj4(self):
truelat2 = (self.truelat1
if _ismissing(self.truelat2)
else self.truelat2)
_proj4 = ("+proj=lcc +units=meters +a={} +b={} +lat_1={} "
"+lat_2={} +lat_0={} +lon_0={}".format(
Constants.WRF_EARTH_RADIUS,
Constants.WRF_EARTH_RADIUS,
self.truelat1,
truelat2,
self.moad_cen_lat,
self.stand_lon))
return _proj4
class MercatorProj(WrfProj):
def __init__(self, bottom_left=None, top_right=None,
lats=None, lons=None, **proj_params):
super(MercatorProj, self).__init__(bottom_left, top_right,
lats, lons, **proj_params)
self._lat_ts = (None
if self.truelat1 == 0. or _ismissing(self.truelat1)
else self.truelat1)
@property
def _cf_params(self):
_cf_params = {}
_cf_params["grid_mapping_name"] = "mercator"
_cf_params["longitude_of_projection_origin"] = self.stand_lon
_cf_params["standard_parallel"] = self.truelat1
return _cf_params
@property
def _pyngl(self):
if not pyngl_enabled():
return None
_pyngl = Resources()
_pyngl.mpProjection = "Mercator"
_pyngl.mpDataBaseVersion = "MediumRes"
_pyngl.mpLimitMode = "Corners"
_pyngl.mpLeftCornerLonF = self.ll_lon
_pyngl.mpLeftCornerLatF = self.ll_lat
_pyngl.mpRightCornerLonF = self.ur_lon
_pyngl.mpRightCornerLatF = self.ur_lat
_pyngl.mpCenterLatF = 0.0
_pyngl.mpCenterLonF = self.stand_lon
return _pyngl
@property
def _basemap(self):
if not basemap_enabled():
return None
_basemap = Basemap(projection = "merc",
lon_0 = self.stand_lon,
lat_0 = self.moad_cen_lat,
lat_ts = self._lat_ts,
llcrnrlat = self.ll_lat,
urcrnrlat = self.ur_lat,
llcrnrlon = self.ll_lon,
urcrnrlon = self.ur_lon,
rsphere = Constants.WRF_EARTH_RADIUS,
resolution = 'l')
return _basemap
@property
def _cartopy(self):
if not cartopy_enabled():
return None
if self._lat_ts == 0.0:
_cartopy = crs.Mercator(
central_longitude = self.stand_lon,
globe = self._globe)
else:
_cartopy = MercatorWithLatTsProj(
central_longitude = self.stand_lon,
latitude_true_scale = self._lat_ts,
globe = self._globe)
return _cartopy
@property
def _cart_extents(self):
# Need to modify the extents for the new projection
pc = crs.PlateCarree()
xs, ys, zs = self._cartopy.transform_points(pc,
np.array([self.ll_lon, self.ur_lon]),
np.array([self.ll_lat, self.ur_lat])).T
_xlimits = xs.tolist()
_ylimits = ys.tolist()
return (_xlimits, _ylimits)
@property
def _proj4(self):
_proj4 = ("+proj=merc +units=meters +a={} +b={} "
"+lon_0={} +lat_ts={}".format(
Constants.WRF_EARTH_RADIUS,
Constants.WRF_EARTH_RADIUS,
self.stand_lon,
self._lat_ts))
return _proj4
class PolarStereographicProj(WrfProj):
def __init__(self, bottom_left=None, top_right=None,
lats=None, lons=None, **proj_params):
super(PolarStereographicProj, self).__init__(bottom_left,
top_right, lats, lons, **proj_params)
self._hemi = -90. if self.truelat1 < 0 else 90.
self._lat_ts = (None
if _ismissing(self.truelat1)
else self.truelat1)
@property
def _cf_params(self):
_cf_params = {}
_cf_params["grid_mapping_name"] = "polar_stereographic"
_cf_params["straight_vertical_longitude_from_pole"] = (
self.stand_lon)
_cf_params["standard_parallel"] = self.truelat1
_cf_params["latitude_of_projection_origin"] = self._hemi
return _cf_params
@property
def _pyngl(self):
if not pyngl_enabled():
return None
_pyngl = Resources()
_pyngl.mpProjection = "Stereographic"
_pyngl.mpDataBaseVersion = "MediumRes"
_pyngl.mpLimitMode = "Corners"
_pyngl.mpLeftCornerLonF = self.ll_lon
_pyngl.mpLeftCornerLatF = self.ll_lat
_pyngl.mpRightCornerLonF = self.ur_lon
_pyngl.mpRightCornerLatF = self.ur_lat
_pyngl.mpCenterLonF = self.stand_lon
if self._hemi > 0:
_pyngl.mpCenterLatF = 90.0
else:
_pyngl.mpCenterLatF = -90.0
return _pyngl
@property
def _basemap(self):
if not basemap_enabled():
return None
_basemap = Basemap(projection = "stere",
lon_0 = self.stand_lon,
lat_0 = self._hemi,
lat_ts = self._lat_ts,
llcrnrlat = self.ll_lat,
urcrnrlat = self.ur_lat,
llcrnrlon = self.ll_lon,
urcrnrlon = self.ur_lon,
rsphere = Constants.WRF_EARTH_RADIUS,
resolution = 'l')
return _basemap
@property
def _cartopy(self):
if not cartopy_enabled():
return None
_cartopy = crs.Stereographic(central_latitude=self._hemi,
central_longitude=self.stand_lon,
true_scale_latitude=self._lat_ts,
globe=self._globe)
return _cartopy
@property
def _cart_extents(self):
# Need to modify the extents for the new projection
pc = crs.PlateCarree()
xs, ys, zs = self._cartopy.transform_points(pc,
np.array([self.ll_lon, self.ur_lon]),
np.array([self.ll_lat, self.ur_lat])).T
_xlimits = xs.tolist()
_ylimits = ys.tolist()
return (_xlimits, _ylimits)
@property
def _proj4(self):
_proj4 = ("+proj=stere +units=meters +a={} +b={} "
"+lat0={} +lon_0={} +lat_ts={}".format(
Constants.WRF_EARTH_RADIUS,
Constants.WRF_EARTH_RADIUS,
self._hemi,
self.stand_lon,
self._lat_ts))
return _proj4
class LatLonProj(WrfProj):
def __init__(self, bottom_left=None, top_right=None,
lats=None, lons=None, **proj_params):
super(LatLonProj, self).__init__(bottom_left, top_right,
lats, lons, **proj_params)
@property
def _cf_params(self):
_cf_params = {}
_cf_params["grid_mapping_name"] = "latitude_longitude"
return _cf_params
@property
def _pyngl(self):
if not pyngl_enabled():
return None
_pyngl = Resources()
_pyngl.mpProjection = "CylindricalEquidistant"
_pyngl.mpDataBaseVersion = "MediumRes"
_pyngl.mpLimitMode = "Corners"
_pyngl.mpLeftCornerLonF = self.ll_lon
_pyngl.mpLeftCornerLatF = self.ll_lat
_pyngl.mpRightCornerLonF = self.ur_lon
_pyngl.mpRightCornerLatF = self.ur_lat
_pyngl.mpCenterLonF = self.stand_lon
_pyngl.mpCenterLatF = self.moad_cen_lat
return _pyngl
@property
def _basemap(self):
if not basemap_enabled():
return None
_basemap = Basemap(projection = "cyl",
lon_0 = self.stand_lon,
lat_0 = self.moad_cen_lat,
llcrnrlat = self.ll_lat,
urcrnrlat = self.ur_lat,
llcrnrlon = self.ll_lon,
urcrnrlon = self.ur_lon,
rsphere = Constants.WRF_EARTH_RADIUS,
resolution = 'l')
return _basemap
@property
def _cartopy(self):
if not cartopy_enabled():
return None
_cartopy = crs.PlateCarree(central_longitude=self.stand_lon,
globe=self._globe)
return _cartopy
@property
def _cart_extents(self):
return ([self.ll_lon, self.ur_lon], [self.ll_lat, self.ur_lat])
@property
def _proj4(self):
_proj4 = ("+proj=eqc +units=meters +a={} +b={} "
"+lon_0={}".format(Constants.WRF_EARTH_RADIUS,
Constants.WRF_EARTH_RADIUS,
self.stand_lon))
return _proj4
# Notes (may not be correct since this projection confuses me):
# Each projection system handles this differently.
# 1) In WRF, if following the WPS instructions, POLE_LON is mainly used to
# determine north or south hemisphere. In other words, it determines if
# the globe is tipped toward or away from you. If a non-0 or non-180
# value is chosen, PyNGL cannot plot it.
# 2) In WRF, POLE_LAT is always positive, but should be negative in the
# proj4 based systems when using the southern hemisphere projections.
# 3) In cartopy, pole_longitude is used to describe the dateline, which
# is 180 degrees away from the normal central (standard) longitude
# (e.g. center of the projection), according to the cartopy developer.
# 4) In basemap, lon_0 should be set to the central (standard) longitude.
# 5) In either cartopy, basemap or pyngl, I'm not sure that projections with
# a pole_lon not equal to 0 or 180 can be plotted. Hopefully people
# follow the WPS instructions, otherwise I need to see a sample file and
# a lot of rum.
# 6) For items in 3 - 4, the "longitude" (lon_0 or pole_longitude) is
# determined by WRF's
# STAND_LON values, with the following calculations based on hemisphere:
# BASEMAP: NH: -STAND_LON; SH: 180.0 - STAND_LON
# CARTOPY: NH: -STAND_LON - 180.; SH: -STAND_LON
# 9) For PYNGL/NCL, you only need to set the center lat and center lon,
# Center lat is the offset of the pole from +/- 90 degrees. Center
# lon is -STAND_LON in NH and 180.0 - STAND_LON in SH.
# 10) It also appears that NetCDF CF has no clear documentation on what
# each parameter means. Going to assume it is the same as basemap, since
# basemap appears to mirror the WMO way of doing things (tilt earth, then
# spin globe).
# 11) Basemap and cartopy produce projections that differ in their extent
# calculations by either using negative values or 0-360 (basemap). For
# this reason, the proj4 string for this class will use cartopy's values
# to keep things in the -180 to 180, -90 to 90 range.
# 12) This projection makes me sad.
class RotLatLonProj(WrfProj):
def __init__(self, bottom_left=None, top_right=None,
lats=None, lons=None, **proj_params):
super(RotLatLonProj, self).__init__(bottom_left, top_right,
lats, lons, **proj_params)
# Need to determine hemisphere, typically pole_lon is 0 for southern
# hemisphere, 180 for northern hemisphere. If not, going to have
# to guess based on other parameters, but hopefully people follow
# the WPS instructions and this never happens.
self._north = True
if self.pole_lon is not None:
if self.pole_lon == 0.:
self._north = False
elif self.pole_lon != 180.:
if self.moad_cen_lat is not None and self.moad_cen_lat < 0.0:
# Only probably true
self._north = False
else:
if self.moad_cen_lat is not None and self.moad_cen_lat < 0.0:
# Only probably true
self._north = False
if self.pole_lat is not None and self.stand_lon is not None:
self._pyngl_cen_lat = (90. - self.pole_lat if self._north
else self.pole_lat - 90.0)
self._pyngl_cen_lon = (-self.stand_lon if self._north
else 180.0 - self.stand_lon)
self._bm_lon_0 = (-self.stand_lon if self._north
else 180.0 - self.stand_lon)
self._bm_cart_pole_lat = (self.pole_lat if self._north
else -self.pole_lat )
# The important point is that pole longitude is the position
# of the dateline of the new projection, not its central
# longitude (per the creator of cartopy). This is based on
# how it's handled by agencies like WMO, but not proj4.
self._cart_pole_lon = (-self.stand_lon - 180.0 if self._north
else -self.stand_lon)
else:
self._pyngl_cen_lat = self.moad_cen_lat
self._pyngl_cen_lon = self.stand_lon
self._bm_cart_pole_lat = (90.0 - self.moad_cen_lat if self._north
else -90.0 - self.moad_cen_lat)
self._bm_lon_0 = (-self.stand_lon if self._north
else 180.0 - self.stand_lon)
self._cart_pole_lon = (-self.stand_lon - 180.0 if self._north
else -self.stand_lon)
@property
def _cf_params(self):
_cf_params = {}
# Assuming this follows the same guidelines as cartopy
_cf_params["grid_mapping_name"] = "rotated_latitude_longitude"
_cf_params["grid_north_pole_latitude"] = self._bm_cart_pole_lat
_cf_params["grid_north_pole_longitude"] = self.pole_lon
_cf_params["north_pole_grid_longitude"] = self._bm_lon_0
return _cf_params
@property
def _pyngl(self):
if not pyngl_enabled():
return None
_pyngl = Resources()
_pyngl.mpProjection = "CylindricalEquidistant"
_pyngl.mpDataBaseVersion = "MediumRes"
_pyngl.mpLimitMode = "Corners"
_pyngl.mpLeftCornerLonF = self.ll_lon
_pyngl.mpLeftCornerLatF = self.ll_lat
_pyngl.mpRightCornerLonF = self.ur_lon
_pyngl.mpRightCornerLatF = self.ur_lat
_pyngl.mpCenterLatF = self._pyngl_cen_lat
_pyngl.mpCenterLonF = self._pyngl_cen_lon
return _pyngl
@property
def _basemap(self):
if not basemap_enabled():
return None
_basemap = Basemap(projection = "rotpole",
o_lat_p = self._bm_cart_pole_lat,
o_lon_p = self.pole_lon,
llcrnrlat = self.ll_lat,
urcrnrlat = self.ur_lat,
llcrnrlon = self.ll_lon,
urcrnrlon = self.ur_lon,
lon_0 = self._bm_lon_0,
rsphere = Constants.WRF_EARTH_RADIUS,
resolution = 'l')
return _basemap
@property
def _cartopy(self):
if not cartopy_enabled():
return None
_cartopy = crs.RotatedPole(
pole_longitude=self._cart_pole_lon,
pole_latitude=self._bm_cart_pole_lat,
central_rotated_longitude=(
180.0 - self.pole_lon), # Probably
globe = self._globe)
return _cartopy
@property
def _cart_extents(self):
# Need to modify the extents for the new projection
pc = crs.PlateCarree()
xs, ys, zs = self._cartopy.transform_points(pc,
np.array([self.ll_lon, self.ur_lon]),
np.array([self.ll_lat, self.ur_lat])).T
_xlimits = xs.tolist()
_ylimits = ys.tolist()
return (_xlimits, _ylimits)
@property
def _proj4(self):
_proj4 = ("+proj=ob_tran +o_proj=latlon "
"+a={} +b={} +to_meter={} +o_lon_p={} +o_lat_p={} "
"+lon_0={}".format(Constants.WRF_EARTH_RADIUS,
Constants.WRF_EARTH_RADIUS,
math.radians(1),
180.0 - self.pole_lon,
self._bm_cart_pole_lat,
180.0 + self._cart_pole_lon))
return _proj4
def getproj(bottom_left=None, top_right=None,
lats=None, lons=None, **proj_params):
proj_type = proj_params.get("MAP_PROJ", 0)
if proj_type == 1:
return LambertConformalProj(bottom_left, top_right,
lats, lons, **proj_params)
elif proj_type == 2:
return PolarStereographicProj(bottom_left, top_right,
lats, lons, **proj_params)
elif proj_type == 3:
return MercatorProj(bottom_left, top_right,
lats, lons, **proj_params)
elif proj_type == 0 or proj_type == 6:
if (proj_params.get("POLE_LAT", None) == 90.
and proj_params.get("POLE_LON", None) == 0.):
return LatLonProj(bottom_left, top_right,
lats, lons, **proj_params)
else:
return RotLatLonProj(bottom_left, top_right,
lats, lons, **proj_params)
else:
# Unknown projection
return WrfProj(bottom_left, top_right,
lats, lons, **proj_params)

BIN
wrf_open/var/src/python/wrf/var/psadlookup.pyc

Binary file not shown.

4
wrf_open/var/src/python/wrf/var/pw.py

@ -6,8 +6,8 @@ from wrf.var.util import extract_vars
__all__ = ["get_pw"] __all__ = ["get_pw"]
def get_pw(wrfnc, timeidx=0): def get_pw(wrfnc, timeidx=0):
ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "PH", "PHB", ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "PH",
"QVAPOR")) "PHB", "QVAPOR"))
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]

4
wrf_open/var/src/python/wrf/var/rh.py

@ -6,7 +6,7 @@ from wrf.var.util import extract_vars
__all__ = ["get_rh", "get_rh_2m"] __all__ = ["get_rh", "get_rh_2m"]
def get_rh(wrfnc, timeidx=0): def get_rh(wrfnc, timeidx=0):
ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR")) ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR"))
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
pb = ncvars["PB"] pb = ncvars["PB"]
@ -21,7 +21,7 @@ def get_rh(wrfnc, timeidx=0):
return rh return rh
def get_rh_2m(wrfnc, timeidx=0): def get_rh_2m(wrfnc, timeidx=0):
ncvars = extract_vars(wrfnc, timeidx, vars=("T2", "PSFC", "Q2")) ncvars = extract_vars(wrfnc, timeidx, varnames=("T2", "PSFC", "Q2"))
t2 = ncvars["T2"] t2 = ncvars["T2"]
psfc = ncvars["PSFC"] psfc = ncvars["PSFC"]
q2 = ncvars["Q2"] q2 = ncvars["Q2"]

2
wrf_open/var/src/python/wrf/var/slp.py

@ -8,7 +8,7 @@ __all__ = ["get_slp"]
@convert_units("pressure", "hpa") @convert_units("pressure", "hpa")
def get_slp(wrfnc, timeidx=0, units="hpa"): def get_slp(wrfnc, timeidx=0, units="hpa"):
ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR",
"PH", "PHB")) "PH", "PHB"))
t = ncvars["T"] t = ncvars["T"]

13
wrf_open/var/src/python/wrf/var/temp.py

@ -1,15 +1,16 @@
from wrf.var.constants import Constants from wrf.var.constants import Constants
from wrf.var.extension import computetk, computeeth, computetv, computewetbulb from wrf.var.extension import computetk, computeeth, computetv, computewetbulb
from wrf.var.decorators import convert_units from wrf.var.decorators import convert_units, set_metadata
from wrf.var.util import extract_vars from wrf.var.util import extract_vars
__all__ = ["get_theta", "get_temp", "get_eth", "get_tv", "get_tw", __all__ = ["get_theta", "get_temp", "get_eth", "get_tv", "get_tw",
"get_tk", "get_tc"] "get_tk", "get_tc"]
@set_metadata()
@convert_units("temp", "k") @convert_units("temp", "k")
def get_theta(wrfnc, timeidx=0, units="k"): def get_theta(wrfnc, timeidx=0, units="k"):
ncvars = extract_vars(wrfnc, timeidx, vars="T") ncvars = extract_vars(wrfnc, timeidx, varnames="T")
t = ncvars["T"] t = ncvars["T"]
full_t = t + Constants.T_BASE full_t = t + Constants.T_BASE
@ -19,7 +20,7 @@ def get_theta(wrfnc, timeidx=0, units="k"):
def get_temp(wrfnc, timeidx=0, units="k"): def get_temp(wrfnc, timeidx=0, units="k"):
"""Return the temperature in Kelvin or Celsius""" """Return the temperature in Kelvin or Celsius"""
ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB")) ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB"))
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
pb = ncvars["PB"] pb = ncvars["PB"]
@ -34,7 +35,7 @@ def get_temp(wrfnc, timeidx=0, units="k"):
def get_eth(wrfnc, timeidx=0, units="k"): def get_eth(wrfnc, timeidx=0, units="k"):
"Return equivalent potential temperature (Theta-e) in Kelvin" "Return equivalent potential temperature (Theta-e) in Kelvin"
ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR")) ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR"))
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
pb = ncvars["PB"] pb = ncvars["PB"]
@ -52,7 +53,7 @@ def get_eth(wrfnc, timeidx=0, units="k"):
def get_tv(wrfnc, timeidx=0, units="k"): def get_tv(wrfnc, timeidx=0, units="k"):
"Return the virtual temperature (tv) in Kelvin or Celsius" "Return the virtual temperature (tv) in Kelvin or Celsius"
ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR")) ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR"))
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
@ -72,7 +73,7 @@ def get_tv(wrfnc, timeidx=0, units="k"):
def get_tw(wrfnc, timeidx=0, units="k"): def get_tw(wrfnc, timeidx=0, units="k"):
"Return the wetbulb temperature (tw)" "Return the wetbulb temperature (tw)"
ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR")) ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR"))
t = ncvars["T"] t = ncvars["T"]
p = ncvars["P"] p = ncvars["P"]
pb = ncvars["PB"] pb = ncvars["PB"]

4
wrf_open/var/src/python/wrf/var/terrain.py

@ -8,10 +8,10 @@ __all__ = ["get_terrain"]
def get_terrain(wrfnc, timeidx=0, units="m"): def get_terrain(wrfnc, timeidx=0, units="m"):
try: try:
hgt_vars = extract_vars(wrfnc, timeidx, vars="HGT") hgt_vars = extract_vars(wrfnc, timeidx, varnames="HGT")
except KeyError: except KeyError:
try: try:
hgt_m_vars = extract_vars(wrfnc, timeidx, vars="HGT_M") hgt_m_vars = extract_vars(wrfnc, timeidx, varnames="HGT_M")
except KeyError: except KeyError:
raise RuntimeError("height variable not found in NetCDF file") raise RuntimeError("height variable not found in NetCDF file")
else: else:

317
wrf_open/var/src/python/wrf/var/util.py

@ -1,13 +1,20 @@
from collections import Iterable from collections import Iterable, Mapping, OrderedDict
from itertools import product from itertools import product
import datetime as dt import datetime as dt
import warnings
import numpy as n import numpy as n
from wrf.var.config import xarray_enabled
from wrf.var.projection import getproj
if xarray_enabled():
from xarray import DataArray
__all__ = ["extract_vars", "extract_global_attrs", "extract_dim", __all__ = ["extract_vars", "extract_global_attrs", "extract_dim",
"combine_files", "is_standard_wrf_var", "extract_times", "combine_files", "is_standard_wrf_var", "extract_times",
"iter_left_indexes", "get_left_indexes", "get_right_slices", "iter_left_indexes", "get_left_indexes", "get_right_slices",
"is_staggered"] "is_staggered", "get_proj_params"]
def _is_multi_time(timeidx): def _is_multi_time(timeidx):
if timeidx == -1: if timeidx == -1:
@ -19,8 +26,25 @@ def _is_multi_file(wrfnc):
return True return True
return False return False
def _is_mapping(wrfnc):
if isinstance(wrfnc, Mapping):
return True
return False
def _is_moving_domain(wrfnc, latvar="XLAT", lonvar="XLONG"):
lat1 = wrfnc.variables[latvar][0,:]
lat2 = wrfnc.variables[latvar][-1,:]
lon1 = wrfnc.variables[lonvar][0,:]
lon2 = wrfnc.variables[lonvar][-1,:]
if (lat1[0,0] != lat2[0,0] or lat1[-1,-1] != lat2[-1,-1] or
lon1[0,0] != lon2[0,0] or lon1[-1,-1] != lon2[-1,-1]):
return True
return False
def _get_attr(wrfnc, attr): def _get_attr(wrfnc, attr):
val = getattr(wrfnc, attr) val = getattr(wrfnc, attr, None)
# PyNIO puts single values in to an array # PyNIO puts single values in to an array
if isinstance(val,n.ndarray): if isinstance(val,n.ndarray):
@ -37,61 +61,274 @@ def extract_global_attrs(wrfnc, attrs):
multifile = _is_multi_file(wrfnc) multifile = _is_multi_file(wrfnc)
if multifile: if multifile:
wrfnc = wrfnc[0] if not _is_mapping(wrfnc):
wrfnc = wrfnc[0]
else:
wrfnc = wrfnc[next(wrfnc.iterkeys())]
return {attr:_get_attr(wrfnc, attr) for attr in attrlist} return {attr:_get_attr(wrfnc, attr) for attr in attrlist}
def extract_dim(wrfnc, dim): def extract_dim(wrfnc, dim):
if _is_multi_file(wrfnc): if _is_multi_file(wrfnc):
wrfnc = wrfnc[0] if not _is_mapping(wrfnc):
wrfnc = wrfnc[0]
else:
wrfnc = wrfnc[next(wrfnc.iterkeys())]
d = wrfnc.dimensions[dim] d = wrfnc.dimensions[dim]
if not isinstance(d, int): if not isinstance(d, int):
return len(d) #netCDF4 return len(d) #netCDF4
return d # PyNIO return d # PyNIO
def combine_files(wrflist, var, timeidx):
# TODO
def _combine_dict(wrfseq, var, timeidx, method):
"""Dictionary combination creates a new left index for each key, then
does a cat or join for the list of files for that key"""
multitime = _is_multi_time(timeidx) multitime = _is_multi_time(timeidx)
numfiles = len(wrflist) numfiles = len(wrfseq)
if not multitime: if not multitime:
time_idx_or_slice = timeidx time_idx_or_slice = timeidx
else: else:
time_idx_or_slice = slice(None, None, None) time_idx_or_slice = slice(None, None, None)
first_var = n.squeeze(wrflist[0].variables[var][time_idx_or_slice, :]) keys = (list(x for x in xrange(numfiles)) if not _is_mapping(wrfseq) else
list(key for key in wrfseq.iterkeys()))
# Check if
if not xarray_enabled():
first_var = wrfseq[keys[0]].variables[var][time_idx_or_slice, :]
else:
first_var = _build_data_array(wrfseq[keys[0]], var, timeidx)
# Create the output data numpy array based on the first array
outdims = [numfiles] outdims = [numfiles]
outdims += first_var.shape outdims += first_var.shape
outarr = n.zeros(outdims, first_var.dtype) outdata = n.zeros(outdims, first_var.dtype)
outarr[0,:] = first_var[:] outdata[0,:] = first_var[:]
for idx, key in enumerate(keys[1:], start=1):
outdata[idx,:] = wrfseq[key].variables[var][time_idx_or_slice, :]
for idx, wrfnc in enumerate(wrflist[1:], start=1): if not xarray_enabled():
outarr[idx,:] = n.squeeze(wrfnc.variables[var][time_idx_or_slice, :]) outarr = outdata
else:
outname = str(first_var.name)
outcoords = dict(first_var.coords)
outdims = ["sequence"] + list(first_var.dims)
outcoords["sequence"] = keys
outattrs = dict(first_var.attrs)
outarr = DataArray(outdata, name=outname, coords=outcoords,
dims=outdims, attrs=outattrs)
return outarr
def _cat_files(wrfseq, var, timeidx):
file_times = extract_times(wrfseq, timeidx)
multitime = _is_multi_time(timeidx)
time_idx_or_slice = timeidx if not multitime else slice(None, None, None)
first_var = (_build_data_array(wrfseq[0], var, timeidx)
if xarray_enabled() else
wrfseq[0].variables[var][time_idx_or_slice, :])
outdims = [len(file_times)]
# Making a new time dim, so ignore this one
outdims += first_var.shape[1:]
outdata = n.zeros(outdims, first_var.dtype)
numtimes = first_var.shape[0]
startidx = 0
endidx = numtimes
outdata[startidx:endidx, :] = first_var[:]
startidx = endidx
for wrfnc in wrfseq[1:]:
vardata = wrfnc.variables[var][time_idx_or_slice, :]
if multitime:
numtimes = vardata.shape[0]
else:
numtimes = 1
endidx = startidx + numtimes
outdata[startidx:endidx, :] = vardata[:]
startidx = endidx
if xarray_enabled():
# FIXME: If it's a moving nest, then the coord arrays need to have same
# time indexes as the whole data set
outname = str(first_var.name)
outattrs = dict(first_var.attrs)
outcoords = dict(first_var.coords)
outdimnames = list(first_var.dims)
outcoords[outdimnames[0]] = file_times # New time dimension values
outarr = DataArray(outdata, name=outname, coords=outcoords,
dims=outdimnames, attrs=outattrs)
else:
outarr = outdata
return outarr
def _join_files(wrfseq, var, timeidx):
multitime = _is_multi_time(timeidx)
numfiles = len(wrfseq)
if not multitime:
time_idx_or_slice = timeidx
else:
time_idx_or_slice = slice(None, None, None)
if xarray_enabled():
first_var = _build_data_array(wrfseq[0], var, timeidx)
else:
first_var = wrfseq[0].variables[var][time_idx_or_slice, :]
# Create the output data numpy array based on the first array
outdims = [numfiles]
outdims += first_var.shape
outdata = n.zeros(outdims, first_var.dtype)
outdata[0,:] = first_var[:]
for idx, wrfnc in enumerate(wrfseq[1:], 1):
print idx
outdata[idx,:] = wrfnc.variables[var][time_idx_or_slice, :]
if xarray_enabled():
outname = str(first_var.name)
outcoords = dict(first_var.coords)
outattrs = dict(first_var.attrs)
# New dimensions
outdimnames = ["file_idx"] + list(first_var.dims)
outcoords["file_idx"] = [i for i in xrange(numfiles)]
outarr = DataArray(outdata, name=outname, coords=outcoords,
dims=outdimnames, attrs=outattrs)
else:
outarr = outdata
return outarr return outarr
def extract_vars(wrfnc, timeidx, vars): def combine_files(wrfseq, var, timeidx, method="cat", squeeze=True):
if isinstance(vars, str): # Dictionary is unique
varlist = [vars] if _is_mapping(wrfseq):
outarr = _combine_dict(wrfseq, var, timeidx, method)
if method.lower() == "cat":
outarr = _cat_files(wrfseq, var, timeidx)
elif method.lower() == "join":
outarr = _join_files(wrfseq, var, timeidx)
else: else:
varlist = vars raise ValueError("method must be 'cat' or 'join'")
if squeeze:
return outarr.squeeze()
return outarr
# Note, always returns the full data set with the time dimension included
def _build_data_array(wrfnc, varname, timeidx):
multitime = _is_multi_time(timeidx)
var = wrfnc.variables[varname]
data = var[:]
attrs = OrderedDict(var.__dict__)
dimnames = var.dimensions
# Add the coordinate variables here.
coord_names = getattr(var, "coordinates").split()
lon_coord = coord_names[0]
lat_coord = coord_names[1]
coords = {}
lon_coord_var = wrfnc.variables[lon_coord]
lat_coord_var = wrfnc.variables[lat_coord]
if multitime:
if _is_moving_domain(wrfnc, lat_coord, lon_coord):
# Special case with a moving domain in a multi-time file,
# otherwise the projection parameters don't change
coords[lon_coord] = lon_coord_var.dimensions, lon_coord_var[:]
coords[lat_coord] = lat_coord_var.dimensions, lat_coord_var[:]
# Returned lats/lons arrays will have a time dimension, so proj
# will need to be a list due to moving corner points
lats, lons, proj_params = get_proj_params(wrfnc,
timeidx,
varname)
proj = [getproj(lats=lats[i,:],
lons=lons[i,:],
**proj_params) for i in xrange(lats.shape[0])]
else:
coords[lon_coord] = (lon_coord_var.dimensions[1:],
lon_coord_var[0,:])
coords[lat_coord] = (lat_coord_var.dimensions[1:],
lat_coord_var[0,:])
# Domain not moving, so just get the first time
lats, lons, proj_params = get_proj_params(wrfnc, 0, varname)
proj = getproj(lats=lats, lons=lons, **proj_params)
else:
coords[lon_coord] = (lon_coord_var.dimensions[1:],
lon_coord_var[timeidx,:])
coords[lat_coord] = (lat_coord_var.dimensions[1:],
lat_coord_var[timeidx,:])
lats, lons, proj_params = get_proj_params(wrfnc, 0, varname)
proj = getproj(lats=lats, lons=lons, **proj_params)
coords[dimnames[0]] = extract_times(wrfnc, timeidx)
attrs["projection"] = proj
data_array = DataArray(data, name=varname, dims=dimnames, coords=coords,
attrs=attrs)
return data_array
def _extract_var(wrfnc, varname, timeidx, method, squeeze):
multitime = _is_multi_time(timeidx) multitime = _is_multi_time(timeidx)
multifile = _is_multi_file(wrfnc) multifile = _is_multi_file(wrfnc)
# Single file, single time if not multifile:
if not multitime and not multifile: if xarray_enabled():
return {var:wrfnc.variables[var][timeidx,:] for var in varlist} return _build_data_array(wrfnc, varname, timeidx)
elif multitime and not multifile: else:
return {var:n.squeeze(wrfnc.variables[var][:]) for var in varlist} if not multitime:
return wrfnc.variables[varname][timeidx,:]
else:
return wrfnc.variables[varname][:]
else:
return combine_files(wrfnc, varname, timeidx, method)
def extract_vars(wrfnc, timeidx, varnames, method="cat", squeeze=True):
if isinstance(varnames, str):
varlist = [varnames]
else: else:
return {var:combine_files(wrfnc, var, timeidx) for var in varlist} varlist = varnames
return {var:_extract_var(wrfnc, var, timeidx, method, squeeze)
for var in varlist}
def _make_time(timearr): def _make_time(timearr):
return dt.datetime.strptime("".join(timearr[:]), "%Y-%m-%d_%H:%M:%S") return dt.datetime.strptime("".join(timearr[:]), "%Y-%m-%d_%H:%M:%S")
def _file_times(wrfnc,timeidx): def _file_times(wrfnc, timeidx):
multitime = _is_multi_time(timeidx) multitime = _is_multi_time(timeidx)
if multitime: if multitime:
times = wrfnc.variables["Times"][:,:] times = wrfnc.variables["Times"][:,:]
@ -102,7 +339,7 @@ def _file_times(wrfnc,timeidx):
yield _make_time(times) yield _make_time(times)
def extract_times(wrfnc,timeidx): def extract_times(wrfnc, timeidx):
multi_file = _is_multi_file(wrfnc) multi_file = _is_multi_file(wrfnc)
if not multi_file: if not multi_file:
wrf_list = [wrfnc] wrf_list = [wrfnc]
@ -168,6 +405,34 @@ def get_right_slices(var, right_ndims, fixed_val=0):
return [slice(None,None,None)] * right_ndims return [slice(None,None,None)] * right_ndims
return [fixed_val]*extra_dim_num + [slice(None,None,None)]*right_ndims return [fixed_val]*extra_dim_num + [slice(None,None,None)]*right_ndims
def get_proj_params(wrfnc, timeidx=0, varname=None):
proj_params = extract_global_attrs(wrfnc, attrs=("MAP_PROJ",
"CEN_LAT", "CEN_LON",
"TRUELAT1", "TRUELAT2",
"MOAD_CEN_LAT", "STAND_LON",
"POLE_LAT", "POLE_LON"))
multitime = _is_multi_time(timeidx)
if not multitime:
time_idx_or_slice = timeidx
else:
time_idx_or_slice = slice(None, None, None)
if varname is not None:
coord_names = getattr(wrfnc.variables[varname], "coordinates").split()
lon_coord = coord_names[0]
lat_coord = coord_names[1]
else:
lat_coord = "XLAT"
lon_coord = "XLONG"
return (wrfnc.variables[lat_coord][time_idx_or_slice,:],
wrfnc.variables[lon_coord][time_idx_or_slice,:],
proj_params)

30
wrf_open/var/src/python/wrf/var/uvmet.py

@ -16,10 +16,10 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"):
if not ten_m: if not ten_m:
try: try:
u_vars = extract_vars(wrfnc, timeidx, vars="U") u_vars = extract_vars(wrfnc, timeidx, varnames="U")
except KeyError: except KeyError:
try: try:
uu_vars = extract_vars(wrfnc, timeidx, vars="UU") uu_vars = extract_vars(wrfnc, timeidx, varnames="UU")
except KeyError: except KeyError:
raise RuntimeError("No valid wind data found in NetCDF file") raise RuntimeError("No valid wind data found in NetCDF file")
else: else:
@ -28,10 +28,10 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"):
u = destagger(u_vars["U"], -1) u = destagger(u_vars["U"], -1)
try: try:
v_vars = extract_vars(wrfnc, timeidx, vars="V") v_vars = extract_vars(wrfnc, timeidx, varnames="V")
except KeyError: except KeyError:
try: try:
vv_vars = extract_vars(wrfnc, timeidx, vars="VV") vv_vars = extract_vars(wrfnc, timeidx, varnames="VV")
except KeyError: except KeyError:
raise RuntimeError("No valid wind data found in NetCDF file") raise RuntimeError("No valid wind data found in NetCDF file")
else: else:
@ -41,10 +41,10 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"):
else: else:
try: try:
u_vars = extract_vars(wrfnc, timeidx, vars="U10") u_vars = extract_vars(wrfnc, timeidx, varnames="U10")
except KeyError: except KeyError:
try: try:
uu_vars = extract_vars(wrfnc, timeidx, vars="UU") uu_vars = extract_vars(wrfnc, timeidx, varnames="UU")
except KeyError: except KeyError:
raise RuntimeError("No valid wind data found in NetCDF file") raise RuntimeError("No valid wind data found in NetCDF file")
else: else:
@ -55,10 +55,10 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"):
u = u_vars["U10"] u = u_vars["U10"]
try: try:
v_vars = extract_vars(wrfnc, timeidx, vars="V10") v_vars = extract_vars(wrfnc, timeidx, varnames="V10")
except KeyError: except KeyError:
try: try:
vv_vars = extract_vars(wrfnc, timeidx, vars="VV") vv_vars = extract_vars(wrfnc, timeidx, varnames="VV")
except KeyError: except KeyError:
raise RuntimeError("No valid wind data found in NetCDF file") raise RuntimeError("No valid wind data found in NetCDF file")
else: else:
@ -82,12 +82,10 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"):
# No rotation needed for Mercator and Lat/Lon # No rotation needed for Mercator and Lat/Lon
return u,v return u,v
elif map_proj in (1,2): elif map_proj in (1,2):
lat_attrs = extract_global_attrs(wrfnc, attrs=("CEN_LAT", lat_attrs = extract_global_attrs(wrfnc, attrs=("TRUELAT1",
"TRUELAT1", "TRUELAT2"))
"TRUELAT2"))
radians_per_degree = Constants.PI/180.0 radians_per_degree = Constants.PI/180.0
# Rotation needed for Lambert and Polar Stereographic # Rotation needed for Lambert and Polar Stereographic
cen_lat = lat_attrs["CEN_LAT"]
true_lat1 = lat_attrs["TRUELAT1"] true_lat1 = lat_attrs["TRUELAT1"]
true_lat2 = lat_attrs["TRUELAT2"] true_lat2 = lat_attrs["TRUELAT2"]
@ -104,10 +102,10 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"):
cen_lon = lon_attrs["STAND_LON"] cen_lon = lon_attrs["STAND_LON"]
try: try:
xlat_m_vars = extract_vars(wrfnc, timeidx, vars="XLAT_M") xlat_m_vars = extract_vars(wrfnc, timeidx, varnames="XLAT_M")
except KeyError: except KeyError:
try: try:
xlat_vars = extract_vars(wrfnc, timeidx, vars="XLAT") xlat_vars = extract_vars(wrfnc, timeidx, varnames="XLAT")
except KeyError: except KeyError:
raise RuntimeError("xlat not found in NetCDF file") raise RuntimeError("xlat not found in NetCDF file")
else: else:
@ -116,10 +114,10 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"):
lat = xlat_m_vars["XLAT_M"] lat = xlat_m_vars["XLAT_M"]
try: try:
xlon_m_vars = extract_vars(wrfnc, timeidx, vars="XLONG_M") xlon_m_vars = extract_vars(wrfnc, timeidx, varnames="XLONG_M")
except KeyError: except KeyError:
try: try:
xlon_vars = extract_vars(wrfnc, timeidx, vars="XLONG") xlon_vars = extract_vars(wrfnc, timeidx, varnames="XLONG")
except KeyError: except KeyError:
raise RuntimeError("xlong not found in NetCDF file") raise RuntimeError("xlong not found in NetCDF file")
else: else:

4
wrf_open/var/src/python/wrf/var/vorticity.py

@ -4,7 +4,7 @@ from wrf.var.util import extract_vars, extract_global_attrs
__all__ = ["get_avo", "get_pvo"] __all__ = ["get_avo", "get_pvo"]
def get_avo(wrfnc, timeidx=0): def get_avo(wrfnc, timeidx=0):
ncvars = extract_vars(wrfnc, timeidx, vars=("U", "V", "MAPFAC_U", ncvars = extract_vars(wrfnc, timeidx, varnames=("U", "V", "MAPFAC_U",
"MAPFAC_V", "MAPFAC_M", "MAPFAC_V", "MAPFAC_M",
"F")) "F"))
@ -23,7 +23,7 @@ def get_avo(wrfnc, timeidx=0):
def get_pvo(wrfnc, timeidx=0): def get_pvo(wrfnc, timeidx=0):
ncvars = extract_vars(wrfnc, timeidx, vars=("U", "V", "T", "P", ncvars = extract_vars(wrfnc, timeidx, varnames=("U", "V", "T", "P",
"PB", "MAPFAC_U", "PB", "MAPFAC_U",
"MAPFAC_V", "MAPFAC_M", "MAPFAC_V", "MAPFAC_M",
"F")) "F"))

214
wrf_open/var/test/projtest.py

@ -0,0 +1,214 @@
import unittest as ut
from os.path import join, basename
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from netCDF4 import Dataset as NetCDF
PYNGL = True
try:
import Ngl
except ImportError:
PYNGL = False
BASEMAP = True
try:
import mpl_toolkits.basemap
except ImportError:
BASEMAP = False
CARTOPY = True
try:
from cartopy import crs, feature
except ImportError:
CARTOPY = False
from wrf.var import (get_proj_params, getproj, RotLatLonProj,
PolarStereographicProj)
FILE_DIR = "/Users/ladwig/Documents/wrf_files/"
WRF_FILES = [
join(FILE_DIR,"wrfout_d01_2016-02-25_18_00_00"),
join(FILE_DIR, "wrfout_d01_2008-09-29_23-30-00"),
join(FILE_DIR, "wrfout_d01_2010-06-13_21:00:00")]
def nz_proj():
lats = np.array([[-47.824014, -47.824014],
[-32.669853, -32.669853]])
lons = np.array([[163.839595, -179.693502],
[163.839595, -179.693502]])
params = {"MAP_PROJ" : 6,
"CEN_LAT" : -41.814869,
"CEN_LON" : 179.693502,
"TRUELAT1" : 0,
"TRUELAT2": 0,
"MOAD_CEN_LAT" : -41.814869,
"STAND_LON" : 180.0 - 179.693502,
"POLE_LAT" : 48.185131,
"POLE_LON" : 0.0}
return lats, lons, RotLatLonProj(lats=lats, lons=lons, **params)
def argentina_proj():
lats = np.array([[-57.144064, -57.144064],
[-21.154470, -21.154470]])
lons = np.array([[-86.893797, -37.089724],
[-86.893797, -37.089724]])
params = {"MAP_PROJ" : 6,
"CEN_LAT" : -39.222954,
"CEN_LON" : -65.980109,
"TRUELAT1" : 0,
"TRUELAT2": 0,
"MOAD_CEN_LAT" : -39.222954,
"STAND_LON" : 180.0 - -65.980109,
"POLE_LAT" : 90 + -39.222954,
"POLE_LON" : 0.0}
return lats, lons, RotLatLonProj(lats=lats, lons=lons, **params)
def south_polar_proj():
lats = np.array([[-30.0,-30.0],
[-30.0,-30.0]])
lons = np.array([[-120, 60],
[-120, 60]])
params = {"MAP_PROJ" : 2,
"CEN_LAT" : -90.0,
"CEN_LON" : 0,
"TRUELAT1" : -10.0,
"MOAD_CEN_LAT" : -90.0,
"STAND_LON" : 0}
return lats, lons, PolarStereographicProj(lats=lats, lons=lons, **params)
def north_polar_proj():
lats = np.array([[30.0,30.0],
[30.0,30.0]])
lons = np.array([[-45, 140],
[-45, 140]])
params = {"MAP_PROJ" : 2,
"CEN_LAT" : 90.0,
"CEN_LON" : 10,
"TRUELAT1" : 10.0,
"MOAD_CEN_LAT" : 90.0,
"STAND_LON" : 10}
return lats, lons, PolarStereographicProj(lats=lats, lons=lons, **params)
def dateline_rot_proj():
lats = np.array([[60.627974, 60.627974],
[71.717521, 71.717521]])
lons = np.array([[170.332771, -153.456292],
[170.332771, -153.456292]])
params = {"MAP_PROJ" : 6,
"CEN_LAT" : 66.335764,
"CEN_LON" : -173.143792,
"TRUELAT1" : 0,
"TRUELAT2": 0,
"MOAD_CEN_LAT" : 66.335764,
"STAND_LON" : 173.143792,
"POLE_LAT" : 90.0 - 66.335764,
"POLE_LON" : 180.0}
return lats, lons, RotLatLonProj(lats=lats, lons=lons, **params)
class WRFProjTest(ut.TestCase):
longMessage = True
def make_test(wrf_file=None, fixed_case=None):
if wrf_file is not None:
ncfile = NetCDF(wrf_file)
lats, lons, proj_params = get_proj_params(ncfile)
proj = getproj(lats=lats, lons=lons, **proj_params)
name_suffix = basename(wrf_file)
elif fixed_case is not None:
name_suffix = fixed_case
if fixed_case == "south_rot":
lats, lons, proj = nz_proj()
elif fixed_case == "arg_rot":
lats, lons, proj = argentina_proj()
elif fixed_case == "south_polar":
lats, lons, proj = south_polar_proj()
elif fixed_case == "north_polar":
lats, lons, proj = north_polar_proj()
elif fixed_case == "dateline_rot":
lats,lons,proj = dateline_rot_proj()
print "wrf proj4: {}".format(proj.proj4())
if PYNGL:
# PyNGL plotting
wks_type = "png"
wks = Ngl.open_wks(wks_type,"pyngl_{}".format(name_suffix))
mpres = proj.pyngl()
map = Ngl.map(wks,mpres)
Ngl.delete_wks(wks)
if BASEMAP:
# Basemap plotting
fig = plt.figure(figsize=(10,10))
ax = fig.add_axes([0.1,0.1,0.8,0.8])
# Define and plot the meridians and parallels
min_lat = np.amin(lats)
max_lat = np.amax(lats)
min_lon = np.amin(lons)
max_lon = np.amax(lons)
parallels = np.arange(np.floor(min_lat), np.ceil(max_lat),
(max_lat - min_lat)/5.0)
meridians = np.arange(np.floor(min_lon), np.ceil(max_lon),
(max_lon - min_lon)/5.0)
bm = proj.basemap()
bm.drawcoastlines(linewidth=.5)
#bm.drawparallels(parallels,labels=[1,1,1,1],fontsize=10)
#bm.drawmeridians(meridians,labels=[1,1,1,1],fontsize=10)
print "basemap proj4: {}".format(bm.proj4string)
plt.savefig("basemap_{}.png".format(name_suffix))
plt.close(fig)
if CARTOPY:
# Cartopy plotting
fig = plt.figure(figsize=(10,10))
ax = plt.axes([0.1,0.1,0.8,0.8], projection=proj.cartopy())
print "cartopy proj4: {}".format(proj.cartopy().proj4_params)
ax.coastlines('50m', linewidth=0.8)
#print proj.x_extents()
#print proj.y_extents()
ax.set_xlim(proj.cartopy_xlim())
ax.set_ylim(proj.cartopy_ylim())
ax.gridlines()
plt.savefig("cartopy_{}.png".format(name_suffix))
plt.close(fig)
if __name__ == "__main__":
for wrf_file in WRF_FILES:
test_func = make_test(wrf_file=wrf_file)
setattr(WRFProjTest, "test_proj", test_func)
test_func2 = make_test(fixed_case="south_rot")
setattr(WRFProjTest, "test_south_rot", test_func2)
test_func3 = make_test(fixed_case="arg_rot")
setattr(WRFProjTest, "test_arg_rot", test_func3)
test_func4 = make_test(fixed_case="south_polar")
setattr(WRFProjTest, "test_south_polar", test_func4)
test_func5 = make_test(fixed_case="north_polar")
setattr(WRFProjTest, "test_north_polar", test_func5)
test_func6 = make_test(fixed_case="dateline_rot")
setattr(WRFProjTest, "test_dateline_rot", test_func6)
ut.main()

2
wrf_open/var/test/utests.py

@ -440,7 +440,7 @@ if __name__ == "__main__":
if var in ignore_vars: if var in ignore_vars:
continue continue
test_func1 = make_test(var, TEST_FILE, OUT_NC_FILE,pynio=True) test_func1 = make_test(var, TEST_FILE, OUT_NC_FILE, pynio=True)
test_func2 = make_test(var, TEST_FILE, OUT_NC_FILE, multi=True, test_func2 = make_test(var, TEST_FILE, OUT_NC_FILE, multi=True,
pynio=True) pynio=True)
setattr(WRFVarsTest, 'test_pynio_{0}'.format(var), test_func1) setattr(WRFVarsTest, 'test_pynio_{0}'.format(var), test_func1)

Loading…
Cancel
Save