diff --git a/wrf_open/var/src/python/wrf/var/decorators.py b/wrf_open/var/src/python/wrf/var/decorators.py index f0dfcd7..c0ee771 100644 --- a/wrf_open/var/src/python/wrf/var/decorators.py +++ b/wrf_open/var/src/python/wrf/var/decorators.py @@ -34,53 +34,4 @@ def convert_units(unit_type, alg_unit): return convert_decorator -def combine_list_and_times(alg_out_dim): - def combine_decorator(func): - @wraps(func) - def func_wrapper(*args, **kargs): - # Multiple times? - multitime = False - if "timeidx" in kargs: - if kargs["timeidx"] == 0: - multitime = True - - # Multiple files? - multifile = False - if isinstance(args[0], Iterable) and not isinstance(args[0], str): - multifile = True - - # Single file, single time - if not multitime and not multifile: - return func(*args, **kargs) - - # Get the dimensions - if multifile: - wrffile = args[0][0] - else: - wrffile = args[0] - - # TODO: Add PyNIO support - dims = wrffile.dimensions - we_size = len(dims["west_east"]) - sn_size = len(dims["south_north"]) - bt_size = len(dims["bottom_top"]) - time_size = len(dims["Time"]) - - - if alg_out_dim == 2: - pass - elif algout_dim == 3: - pass - elif algout_dim == 4: - pass - else: - raise RuntimeError("invalid algorithm output dimsize") - - - - - return func(*args, **kargs) - - return func_wrapper - - return combine_decorator + diff --git a/wrf_open/var/src/python/wrf/var/precip.py b/wrf_open/var/src/python/wrf/var/precip.py index 294a1b4..399354d 100755 --- a/wrf_open/var/src/python/wrf/var/precip.py +++ b/wrf_open/var/src/python/wrf/var/precip.py @@ -1,5 +1,4 @@ -import numpy as n from wrf.var.util import extract_vars @@ -15,8 +14,8 @@ def get_accum_precip(wrfnc, timeidx=0): return rainsum def get_precip_diff(wrfnc1, wrfnc2, timeidx=0): - vars1 = extract_vars(wrfnc, timeidx, vars=("RAINC", "RAINNC")) - vars2 = extract_vars(wrfnc, timeidx, vars=("RAINC", "RAINNC")) + vars1 = extract_vars(wrfnc1, timeidx, vars=("RAINC", "RAINNC")) + vars2 = extract_vars(wrfnc2, timeidx, vars=("RAINC", "RAINNC")) rainc1 = vars1["RAINC"] rainnc1 = vars1["RAINNC"] diff --git a/wrf_open/var/src/python/wrf/var/util.py b/wrf_open/var/src/python/wrf/var/util.py index d921256..c2255f5 100644 --- a/wrf_open/var/src/python/wrf/var/util.py +++ b/wrf_open/var/src/python/wrf/var/util.py @@ -2,7 +2,8 @@ from collections import Iterable import numpy as n -__all__ = ["extract_vars", "extract_global_attrs", "hold_dim_fixed"] +__all__ = ["extract_vars", "extract_global_attrs", "extract_dim", + "hold_dim_fixed", "combine_files"] def _is_multi_time(timeidx): if timeidx == -1: @@ -35,7 +36,33 @@ def extract_global_attrs(wrfnc, attrs): wrfnc = wrfnc[0] return {attr:_get_attr(wrfnc, attr) for attr in attrlist} + +def extract_dim(wrfnc, dim): + d = wrfnc.dimensions[dim] + if not isinstance(d, int): + return len(d) #netCDF4 + return d # PyNIO +def combine_files(wrflist, var, timeidx): + + multitime = _is_multi_time(timeidx) + numfiles = len(wrflist) + + if not multitime: + time_idx_or_slice = timeidx + else: + time_idx_or_slice = slice(None, None, None) + + first_var = wrflist[0].variables[var][time_idx_or_slice, :] + outdims = [numfiles] + outdims += first_var.shape + outarr = n.zeros(outdims, first_var.dtype) + outarr[0,:] = first_var[:] + + for idx, wrfnc in enumerate(wrflist[1:], start=1): + outarr[idx,:] = wrfnc.variables[var][time_idx_or_slice, :] + + return outarr def extract_vars(wrfnc, timeidx, vars): if isinstance(vars, str): @@ -49,6 +76,11 @@ def extract_vars(wrfnc, timeidx, vars): # Single file, single time if not multitime and not multifile: return {var:wrfnc.variables[var][timeidx,:] for var in varlist} + elif multitime and not multifile: + return {var:wrfnc.variables[var][:] for var in varlist} + else: + return {var:combine_files(wrfnc, var, timeidx) for var in varlist} + def hold_dim_fixed(var, dim, idx): """Generic method to hold a single dimension to a fixed index when @@ -56,7 +88,7 @@ def hold_dim_fixed(var, dim, idx): be negative. For example, on a 4D array with 'dim' set to - -1 and 'idx' set to 3, this is simply going to do this operation: + -1 and 'idx' set to 3, this is going to return this view: var[:,:,:,3]