forked from 3rdparty/wrf-python
Bill Ladwig
9 years ago
11 changed files with 698 additions and 582 deletions
@ -0,0 +1,227 @@
@@ -0,0 +1,227 @@
|
||||
from __future__ import (absolute_import, division, print_function, |
||||
unicode_literals) |
||||
|
||||
from collections import Iterable |
||||
|
||||
import numpy as np |
||||
|
||||
from .constants import Constants |
||||
from .extension import computeij, computell |
||||
from .util import (extract_vars, extract_global_attrs, |
||||
either, _is_moving_domain, _is_multi_time_req, |
||||
iter_left_indexes, _is_mapping, _is_multi_file, |
||||
viewkeys) |
||||
|
||||
def _lat_varname(wrfnc, stagger): |
||||
if stagger is None or stagger.lower() == "m": |
||||
varname = either("XLAT", "XLAT_M")(wrfnc) |
||||
elif stagger.lower() == "u" or stagger.lower() == "v": |
||||
varname = "XLAT_{}".format(stagger.upper()) |
||||
else: |
||||
raise ValueError("invalid 'stagger' value") |
||||
|
||||
return varname |
||||
|
||||
def _lon_varname(wrfnc, stagger): |
||||
if stagger is None or stagger.lower() == "m": |
||||
varname = either("XLONG", "XLONG_M")(wrfnc) |
||||
elif stagger.lower() == "u" or stagger.lower() == "v": |
||||
varname = "XLONG_{}".format(stagger.upper()) |
||||
else: |
||||
raise ValueError("invalid 'stagger' value") |
||||
|
||||
return varname |
||||
|
||||
def _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache): |
||||
if timeidx < 0: |
||||
raise ValueError("'timeidx' must be greater than 0") |
||||
|
||||
attrs = extract_global_attrs(wrfnc, attrs=("MAP_PROJ", "TRUELAT1", |
||||
"TRUELAT2", "STAND_LON", |
||||
"DX", "DY")) |
||||
map_proj = attrs["MAP_PROJ"] |
||||
truelat1 = attrs["TRUELAT1"] |
||||
truelat2 = attrs["TRUELAT2"] |
||||
stdlon = attrs["STAND_LON"] |
||||
dx = attrs["DX"] |
||||
dy = attrs["DY"] |
||||
|
||||
if map_proj == 6: |
||||
pole_attrs = extract_global_attrs(wrfnc, attrs=("POLE_LAT", |
||||
"POLE_LON")) |
||||
pole_lat = pole_attrs["POLE_LAT"] |
||||
pole_lon = pole_attrs["POLE_LON"] |
||||
latinc = (dy*360.0)/2.0 / Constants.PI/Constants.WRF_EARTH_RADIUS |
||||
loninc = (dx*360.0)/2.0 / Constants.PI/Constants.WRF_EARTH_RADIUS |
||||
else: |
||||
pole_lat = 90.0 |
||||
pole_lon = 0.0 |
||||
latinc = 0.0 |
||||
loninc = 0.0 |
||||
|
||||
latvar = _lat_varname(wrfnc, stagger) |
||||
lonvar = _lon_varname(wrfnc, stagger) |
||||
|
||||
lat_timeidx = timeidx |
||||
|
||||
is_moving = _is_moving_domain(wrfnc, latvar=latvar, lonvar=lonvar) |
||||
|
||||
# Only need one file and one time if the domain is not moving |
||||
if not is_moving: |
||||
if _is_multi_time_req(timeidx): |
||||
lat_timeidx = 0 |
||||
|
||||
if _is_multi_file(wrfnc): |
||||
if not _is_mapping(wrfnc): |
||||
wrfnc = next(iter(wrfnc)) # only need one file |
||||
else: |
||||
wrfnc = wrfnc[next(iter(viewkeys(wrfnc)))] |
||||
return _get_proj_params(wrfnc, timeidx, stagger, |
||||
method, squeeze, cache) |
||||
|
||||
xlat = extract_vars(wrfnc, lat_timeidx, (latvar,), method, squeeze, cache, |
||||
nometa=True)[latvar] |
||||
xlon = extract_vars(wrfnc, lat_timeidx, (lonvar,), method, squeeze, cache, |
||||
nometa=True)[lonvar] |
||||
|
||||
ref_lat = np.ravel(xlat[...,0,0]) |
||||
ref_lon = np.ravel(xlon[...,0,0]) |
||||
|
||||
# Note: fortran index |
||||
known_i = 1.0 |
||||
known_j = 1.0 |
||||
|
||||
return (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, |
||||
pole_lat, pole_lon, known_i, known_j, dx, latinc, loninc) |
||||
|
||||
|
||||
def ll_to_ij(wrfnc, latitude, longitude, timeidx=0, |
||||
stagger=None, method="cat", squeeze=True, cache=None): |
||||
|
||||
(map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon, |
||||
pole_lat,pole_lon,known_i,known_j,dx,latinc, |
||||
loninc) = _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache) |
||||
|
||||
if isinstance(latitude, Iterable): |
||||
lats = np.asarray(latitude) |
||||
lons = np.asarray(longitude) |
||||
|
||||
if lats.ndim > 1: |
||||
lats = lats.ravel() |
||||
|
||||
if lons.ndim > 1: |
||||
lons = lons.ravel() |
||||
|
||||
if (lats.size != lons.size): |
||||
raise ValueError("'latitude' and 'longitude' " |
||||
"must be the same length") |
||||
|
||||
|
||||
if ref_lat.size == 1: |
||||
outdim = [lats.size, 2] |
||||
extra_dims = [outdim[0]] |
||||
else: |
||||
# Moving domain will have moving ref_lats/ref_lons |
||||
outdim = [lats.size, ref_lat.size, 2] |
||||
extra_dims = outdim[0:2] |
||||
|
||||
res = np.empty(outdim, np.float64) |
||||
|
||||
for left_idxs in iter_left_indexes(extra_dims): |
||||
left_and_slice_idxs = left_idxs + (slice(None), ) |
||||
|
||||
if ref_lat.size == 1: |
||||
ref_lat_val = ref_lat[0] |
||||
ref_lon_val = ref_lon[0] |
||||
else: |
||||
ref_lat_val = ref_lat[left_idxs[-1]] |
||||
ref_lon_val = ref_lon[left_idxs[-1]] |
||||
|
||||
lat = lats[left_idxs[0]] |
||||
lon = lons[left_idxs[0]] |
||||
|
||||
ij = computeij(map_proj, truelat1, truelat2, stdlon, |
||||
ref_lat_val, ref_lon_val, pole_lat, pole_lon, |
||||
known_i, known_j, dx, latinc, loninc, |
||||
lat, lon) |
||||
|
||||
res[left_and_slice_idxs] = ij[:] |
||||
|
||||
else: |
||||
|
||||
res = computeij(map_proj, truelat1, truelat2, stdlon, |
||||
ref_lat, ref_lon, pole_lat, pole_lon, |
||||
known_i, known_j, dx, latinc, loninc, |
||||
latitude, longitude) |
||||
|
||||
if squeeze: |
||||
res = res.squeeze() |
||||
|
||||
return res |
||||
|
||||
def ij_to_ll(wrfnc, i, j, timeidx=0, |
||||
stagger=None, method="cat", squeeze=True, cache=None): |
||||
|
||||
(map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon, |
||||
pole_lat,pole_lon,known_i,known_j,dx,latinc, |
||||
loninc) = _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache) |
||||
|
||||
if isinstance(i, Iterable): |
||||
i_arr = np.asarray(i) |
||||
j_arr = np.asarray(j) |
||||
|
||||
if i_arr.ndim > 1: |
||||
i_arr = i_arr.ravel() |
||||
|
||||
if j_arr.ndim > 1: |
||||
j_arr = j_arr.ravel() |
||||
|
||||
if (i_arr.size != j_arr.size): |
||||
raise ValueError("'i' and 'j' " |
||||
"must be the same length") |
||||
|
||||
if ref_lat.size == 1: |
||||
outdim = [i_arr.size, 2] |
||||
extra_dims = [outdim[0]] |
||||
else: |
||||
# Moving domain will have moving ref_lats/ref_lons |
||||
outdim = [i_arr.size, ref_lat.size, 2] |
||||
extra_dims = outdim[0:2] |
||||
|
||||
res = np.empty(outdim, np.float64) |
||||
|
||||
for left_idxs in iter_left_indexes(extra_dims): |
||||
left_and_slice_idxs = left_idxs + (slice(None), ) |
||||
|
||||
if ref_lat.size == 1: |
||||
ref_lat_val = ref_lat[0] |
||||
ref_lon_val = ref_lon[0] |
||||
else: |
||||
ref_lat_val = ref_lat[left_idxs[-1]] |
||||
ref_lon_val = ref_lon[left_idxs[-1]] |
||||
|
||||
i_val = i_arr[left_idxs[0]] |
||||
j_val = j_arr[left_idxs[0]] |
||||
|
||||
ll = computell(map_proj, truelat1, truelat2, |
||||
stdlon, ref_lat_val, ref_lon_val, |
||||
pole_lat, pole_lon, known_i, known_j, |
||||
dx, latinc, loninc, |
||||
i_val, j_val) |
||||
|
||||
res[left_and_slice_idxs] = ll[:] |
||||
|
||||
else: |
||||
i_val = i |
||||
j_val = j |
||||
|
||||
res = computell(map_proj, truelat1, truelat2, |
||||
stdlon, ref_lat, ref_lon, |
||||
pole_lat, pole_lon, known_i, known_j, |
||||
dx, latinc, loninc, |
||||
i_val, j_val) |
||||
|
||||
if squeeze: |
||||
res = res.squeeze() |
||||
|
||||
return res |
||||
@ -0,0 +1,166 @@
@@ -0,0 +1,166 @@
|
||||
|
||||
|
||||
from .util import _unpack_sequence, is_standard_wrf_var, extract_vars |
||||
from .cape import get_2dcape, get_3dcape |
||||
from .ctt import get_ctt |
||||
from .dbz import get_dbz, get_max_dbz |
||||
from .dewpoint import get_dp, get_dp_2m |
||||
from .geoht import get_geopt, get_height |
||||
from .helicity import get_srh, get_uh |
||||
from .latlon import get_lat, get_lon |
||||
from .omega import get_omega |
||||
from .pressure import get_pressure, get_pressure_hpa |
||||
from .pw import get_pw |
||||
from .rh import get_rh, get_rh_2m |
||||
from .slp import get_slp |
||||
from .temp import get_tc, get_eth, get_temp, get_theta, get_tk, get_tv, get_tw |
||||
from .terrain import get_terrain |
||||
from .uvmet import (get_uvmet, get_uvmet10, get_uvmet10_wspd_wdir, |
||||
get_uvmet_wspd_wdir) |
||||
from .vorticity import get_avo, get_pvo |
||||
from .wind import (get_destag_wspd_wdir, get_destag_wspd_wdir10, |
||||
get_u_destag, get_v_destag, get_w_destag) |
||||
from .times import get_times |
||||
|
||||
__all__ = ["getvar"] |
||||
|
||||
# func is the function to call. kargs are required arguments that should |
||||
# not be altered by the user |
||||
_FUNC_MAP = {"cape2d" : get_2dcape, |
||||
"cape3d" : get_3dcape, |
||||
"dbz" : get_dbz, |
||||
"maxdbz" : get_max_dbz, |
||||
"dp" : get_dp, |
||||
"dp2m" : get_dp_2m, |
||||
"height" : get_height, |
||||
"geopt" : get_geopt, |
||||
"srh" : get_srh, |
||||
"uhel" : get_uh, |
||||
"omega" : get_omega, |
||||
"pw" : get_pw, |
||||
"rh" : get_rh, |
||||
"rh2m" : get_rh_2m, |
||||
"slp" : get_slp, |
||||
"theta" : get_theta, |
||||
"temp" : get_temp, |
||||
"tk" : get_tk, |
||||
"tc" : get_tc, |
||||
"theta_e" : get_eth, |
||||
"tv" : get_tv, |
||||
"twb" : get_tw, |
||||
"terrain" : get_terrain, |
||||
"times" : get_times, |
||||
"uvmet" : get_uvmet, |
||||
"uvmet10" : get_uvmet10, |
||||
"avo" : get_avo, |
||||
"pvo" : get_pvo, |
||||
"ua" : get_u_destag, |
||||
"va" : get_v_destag, |
||||
"wa" : get_w_destag, |
||||
"lat" : get_lat, |
||||
"lon" : get_lon, |
||||
"pressure" : get_pressure_hpa, |
||||
"pres" : get_pressure, |
||||
"wspd_wdir" : get_destag_wspd_wdir, |
||||
"wspd_wdir10" : get_destag_wspd_wdir10, |
||||
"wspd_wdir_uvmet" : get_uvmet_wspd_wdir, |
||||
"wspd_wdir_uvmet10" : get_uvmet10_wspd_wdir, |
||||
"ctt" : get_ctt |
||||
} |
||||
|
||||
_VALID_KARGS = {"cape2d" : ["missing"], |
||||
"cape3d" : ["missing"], |
||||
"dbz" : ["do_variant", "do_liqskin"], |
||||
"maxdbz" : ["do_variant", "do_liqskin"], |
||||
"dp" : ["units"], |
||||
"dp2m" : ["units"], |
||||
"height" : ["msl", "units"], |
||||
"geopt" : [], |
||||
"srh" : ["top"], |
||||
"uhel" : ["bottom", "top"], |
||||
"omega" : [], |
||||
"pw" : [], |
||||
"rh" : [], |
||||
"rh2m" : [], |
||||
"slp" : ["units"], |
||||
"temp" : ["units"], |
||||
"tk" : [], |
||||
"tc" : [], |
||||
"theta" : ["units"], |
||||
"theta_e" : ["units"], |
||||
"tv" : ["units"], |
||||
"twb" : ["units"], |
||||
"terrain" : ["units"], |
||||
"times" : [], |
||||
"uvmet" : ["units"], |
||||
"uvmet10" : ["units"], |
||||
"avo" : [], |
||||
"pvo" : [], |
||||
"ua" : ["units"], |
||||
"va" : ["units"], |
||||
"wa" : ["units"], |
||||
"lat" : [], |
||||
"lon" : [], |
||||
"pres" : ["units"], |
||||
"pressure" : ["units"], |
||||
"wspddir" : ["units"], |
||||
"wspddir_uvmet" : ["units"], |
||||
"wspddir_uvmet10" : ["units"], |
||||
"ctt" : [], |
||||
"default" : [] |
||||
} |
||||
|
||||
_ALIASES = {"cape_2d" : "cape2d", |
||||
"cape_3d" : "cape3d", |
||||
"eth" : "theta_e", |
||||
"mdbz" : "maxdbz", |
||||
"geopotential" : "geopt", |
||||
"helicity" : "srh", |
||||
"latitude" : "lat", |
||||
"longitude" : "lon", |
||||
"omg" : "omega", |
||||
"p" : "pres", |
||||
"rh2" : "rh2m", |
||||
"z": "height", |
||||
"ter" : "terrain", |
||||
"updraft_helicity" : "uhel", |
||||
"td" : "dp", |
||||
"td2" : "dp2m" |
||||
} |
||||
|
||||
class ArgumentError(Exception): |
||||
def __init__(self, msg): |
||||
self.msg = msg |
||||
|
||||
def __str__(self): |
||||
return self.msg |
||||
|
||||
def _undo_alias(alias): |
||||
actual = _ALIASES.get(alias, None) |
||||
if actual is None: |
||||
return alias |
||||
else: |
||||
return actual |
||||
|
||||
def _check_kargs(var, kargs): |
||||
for arg in kargs.iterkeys(): |
||||
if arg not in _VALID_KARGS[var]: |
||||
raise ArgumentError("'%s' is an invalid keyword " |
||||
"argument for '%s'" % (arg, var)) |
||||
|
||||
|
||||
def getvar(wrfnc, var, timeidx=0, |
||||
method="cat", squeeze=True, cache=None, |
||||
**kargs): |
||||
|
||||
wrfnc = _unpack_sequence(wrfnc) |
||||
|
||||
if is_standard_wrf_var(wrfnc, var): |
||||
return extract_vars(wrfnc, timeidx, var, method, squeeze, cache)[var] |
||||
|
||||
actual_var = _undo_alias(var) |
||||
if actual_var not in _VALID_KARGS: |
||||
raise ArgumentError("'%s' is not a valid variable name" % (var)) |
||||
|
||||
_check_kargs(actual_var, kargs) |
||||
return _FUNC_MAP[actual_var](wrfnc,timeidx,**kargs) |
||||
Loading…
Reference in new issue