Added full variable extraction with lists and multiple times

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
+

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"]
     

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]