diff --git a/wrf_open/var/src/python/wrf/var/cape.py b/wrf_open/var/src/python/wrf/var/cape.py index e21ed9f..73bb548 100755 --- a/wrf_open/var/src/python/wrf/var/cape.py +++ b/wrf_open/var/src/python/wrf/var/cape.py @@ -4,7 +4,7 @@ import numpy.ma as ma from wrf.var.extension import computetk,computecape from wrf.var.destagger import destagger from wrf.var.constants import Constants, ConversionFactors -from wrf.var.util import extract_vars, hold_dim_fixed +from wrf.var.util import extract_vars __all__ = ["get_2dcape", "get_3dcape"] @@ -27,7 +27,7 @@ def get_2dcape(wrfnc, missing=-999999.0, timeidx=0): tk = computetk(full_p, full_t) geopt = ph + phb - geopt_unstag = destagger(geopt, 0) + geopt_unstag = destagger(geopt, -3) z = geopt_unstag/Constants.G # Convert pressure to hPa @@ -40,21 +40,23 @@ def get_2dcape(wrfnc, missing=-999999.0, timeidx=0): cape_res,cin_res = computecape(p_hpa,tk,qv,z,ter,psfc_hpa, missing,i3dflag,ter_follow) - cape = hold_dim_fixed(cape_res, -3, 0) - cin = hold_dim_fixed(cin_res, -3, 0) - lcl = hold_dim_fixed(cin_res, -3, 1) - lfc = hold_dim_fixed(cin_res, -3, 2) + cape = cape_res[...,0,:,:] + cin = cin_res[...,0,:,:] + lcl = cin_res[...,1,:,:] + lfc = cin_res[...,2,:,:] - resdim = [4] - resdim += cape.shape + left_dims = [x for x in cape_res.shape[0:-3]] + right_dims = [x for x in cape_res.shape[-2:]] + + resdim = left_dims + [4] + right_dims # Make a new output array for the result res = n.zeros((resdim), cape.dtype) - res[0,:] = cape[:] - res[1,:] = cin[:] - res[2,:] = lcl[:] - res[3,:] = lfc[:] + res[...,0,:,:] = cape[:] + res[...,1,:,:] = cin[:] + res[...,2,:,:] = lcl[:] + res[...,3,:,:] = lfc[:] return ma.masked_values(res,missing) @@ -77,7 +79,7 @@ def get_3dcape(wrfnc, missing=-999999.0, timeidx=0): tk = computetk(full_p, full_t) geopt = ph + phb - geopt_unstag = destagger(geopt, 0) + geopt_unstag = destagger(geopt, -3) z = geopt_unstag/Constants.G # Convert pressure to hPa @@ -91,13 +93,17 @@ def get_3dcape(wrfnc, missing=-999999.0, timeidx=0): missing,i3dflag,ter_follow) # Make a new output array for the result - resdim = [2] - resdim += cape.shape - res = n.zeros((resdim), cape.dtype) + left_dims = [x for x in cape.shape[0:-3]] + right_dims = [x for x in cape.shape[-3:]] + + resdim = left_dims + [2] + right_dims + + res = n.zeros(resdim, cape.dtype) - res[0,:] = cape[:] - res[1,:] = cin[:] + res[...,0,:,:,:] = cape[:] + res[...,1,:,:,:] = cin[:] + #return res return ma.masked_values(res, missing) diff --git a/wrf_open/var/src/python/wrf/var/ctt.py b/wrf_open/var/src/python/wrf/var/ctt.py index ad80e32..b364208 100644 --- a/wrf_open/var/src/python/wrf/var/ctt.py +++ b/wrf_open/var/src/python/wrf/var/ctt.py @@ -46,7 +46,7 @@ def get_ctt(wrfnc, units="c", timeidx=0): tk = computetk(full_p, full_t) geopt = ph + phb - geopt_unstag = destagger(geopt, 0) + geopt_unstag = destagger(geopt, -3) ght = geopt_unstag / Constants.G ctt = computectt(p_hpa,tk,qice,qcld,qv,ght,ter,haveqci) diff --git a/wrf_open/var/src/python/wrf/var/dbz.py b/wrf_open/var/src/python/wrf/var/dbz.py index a5cd65b..8313d42 100755 --- a/wrf_open/var/src/python/wrf/var/dbz.py +++ b/wrf_open/var/src/python/wrf/var/dbz.py @@ -60,5 +60,5 @@ def get_dbz(wrfnc, do_varint=False, do_liqskin=False, timeidx=0): def get_max_dbz(wrfnc, do_varint=False, do_liqskin=False, timeidx=0): return n.amax(get_dbz(wrfnc, do_varint, do_liqskin, timeidx), - axis=0) + axis=-3) diff --git a/wrf_open/var/src/python/wrf/var/decorators.py b/wrf_open/var/src/python/wrf/var/decorators.py index 334bac8..9c5f972 100644 --- a/wrf_open/var/src/python/wrf/var/decorators.py +++ b/wrf_open/var/src/python/wrf/var/decorators.py @@ -5,6 +5,7 @@ from itertools import product import numpy as n from wrf.var.units import do_conversion, check_units +from wrf.var.destagger import destagger __all__ = ["convert_units", "handle_left_iter"] @@ -59,7 +60,7 @@ def _left_indexes(dims): def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1, ref_var_name=None, alg_out_fixed_dims=(), - ignore_args=(), ignore_kargs={}): + ignore_args=(), ignore_kargs={}, ref_stag_dim=None): """Decorator to handle iterating over leftmost dimensions when using multiple files and/or multiple times. @@ -78,18 +79,22 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1, directly without any slicing - ignore_kargs - keys of any keyword arguments which should be passed directly without slicing + - ref_stag_dim - in some cases the reference variable dimensions + have a staggered dimension that needs to be corrected when calculating + the output dimensions (e.g. avo) """ def indexing_decorator(func): @wraps(func) def func_wrapper(*args, **kargs): + if ref_var_idx >= 0: ref_var = args[ref_var_idx] else: ref_var = kargs[ref_var_name] ref_var_shape = ref_var.shape - extra_dim_num = len(ref_var_shape) - ref_var_expected_dims + extra_dim_num = ref_var.ndim- ref_var_expected_dims # No special left side iteration, return the function result if (extra_dim_num == 0): @@ -109,10 +114,13 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1, # numerical algorithm has greater dim count than reference, # Note: user must provide this information to decorator right_dims = [dim for dim in alg_out_fixed_dims] + neg_start_idx = -(alg_out_num_dims - len(alg_out_fixed_dims)) right_dims += [ref_var_shape[x] - for x in xrange(-alg_out_num_dims,0,1)] + for x in xrange(neg_start_idx,0,1)] outdims += right_dims - + + if ref_stag_dim is not None: + outdims[ref_stag_dim] -= 1 out_inited = False for left_idxs in _left_indexes(extra_dims): @@ -121,7 +129,6 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1, # the right (e.g. [1,1,:]) left_and_slice_idxs = ([x for x in left_idxs] + [slice(None, None, None)]) - # Slice the args if applicable new_args = [arg[left_and_slice_idxs] if i not in ignore_args else arg @@ -133,7 +140,6 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1, if key not in ignore_kargs else val) for key,val in kargs.iteritems()} - # Call the numerical routine res = func(*new_args, **new_kargs) @@ -147,11 +153,93 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1, else: # This should be a list or a tuple (cape) if not out_inited: - output = [n.zeros(outdims, ref_var.dtype)] + output = [n.zeros(outdims, ref_var.dtype) + for i in xrange(len(res))] out_inited = True - for outarr in output: - outarr[left_and_slice_idxs] = res[:] + for i,outarr in enumerate(res): + (output[i])[left_and_slice_idxs] = outarr[:] + + + return output + + return func_wrapper + + return indexing_decorator + +def uvmet_left_iter(): + """Decorator to handle iterating over leftmost dimensions when using + multiple files and/or multiple times with the uvmet product. + + """ + def indexing_decorator(func): + @wraps(func) + def func_wrapper(*args): + u = args[0] + v = args[1] + lat = args[2] + lon = args[3] + cen_long = args[4] + cone = args[5] + + if u.ndim == lat.ndim: + num_right_dims = 2 + is_3d = False + else: + num_right_dims = 3 + is_3d = True + + is_stag = False + if ((u.shape[-1] != lat.shape[-1]) or + (u.shape[-2] != lat.shape[-2])): + is_stag = True + + if is_3d: + extra_dim_num = u.ndim - 3 + else: + extra_dim_num = u.ndim - 2 + + if is_stag: + u = destagger(u,-1) + v = destagger(v,-2) + + # No special left side iteration, return the function result + if (extra_dim_num == 0): + return func(u,v,lat,lon,cen_long,cone) + + # Start by getting the left-most 'extra' dims + outdims = [u.shape[x] for x in xrange(extra_dim_num)] + extra_dims = list(outdims) # Copy the left-most dims for iteration + + # Append the right-most dimensions + outdims += [2] # For u/v components + + outdims += [u.shape[x] for x in xrange(-num_right_dims,0,1)] + + output = n.zeros(outdims, u.dtype) + + for left_idxs in _left_indexes(extra_dims): + # Make the left indexes plus a single slice object + # The single slice will handle all the dimensions to + # the right (e.g. [1,1,:]) + left_and_slice_idxs = ([x for x in left_idxs] + + [slice(None, None, None)]) + + + new_u = u[left_and_slice_idxs] + new_v = v[left_and_slice_idxs] + new_lat = lat[left_and_slice_idxs] + new_lon = lon[left_and_slice_idxs] + + # Call the numerical routine + res = func(new_u, new_v, new_lat, new_lon, cen_long, cone) + + # Note: The 2D version will return a 3D array with a 1 length + # dimension. Numpy is unable to broadcast this without + # sqeezing first. + res = n.squeeze(res) + + output[left_and_slice_idxs] = res[:] return output @@ -161,3 +249,4 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1, return indexing_decorator + diff --git a/wrf_open/var/src/python/wrf/var/destagger.py b/wrf_open/var/src/python/wrf/var/destagger.py index 43702e1..160f4e4 100755 --- a/wrf_open/var/src/python/wrf/var/destagger.py +++ b/wrf_open/var/src/python/wrf/var/destagger.py @@ -1,6 +1,8 @@ import numpy as n +from wrf.var.util import extract_vars + __all__ = ["destagger", "destagger_windcomp", "destagger_winds"] def destagger(var, stagger_dim): @@ -50,7 +52,9 @@ def destagger_windcomp(wrfnc, comp, timeidx=0): wrfvar = "W" stagdim = -3 - wind_data = wrfnc.variables[wrfvar][timeidx,:,:,:] + #wind_data = wrfnc.variables[wrfvar][timeidx,:,:,:] + ncvars = extract_vars(wrfnc, timeidx, vars=wrfvar) + wind_data = ncvars[wrfvar] return destagger(wind_data, stagdim) def destagger_winds(wrfnc, timeidx=0): diff --git a/wrf_open/var/src/python/wrf/var/extension.py b/wrf_open/var/src/python/wrf/var/extension.py index 16c1768..842cf68 100755 --- a/wrf_open/var/src/python/wrf/var/extension.py +++ b/wrf_open/var/src/python/wrf/var/extension.py @@ -11,7 +11,7 @@ from wrf.var._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d, f_computesrh, f_computeuh, f_computepw, f_computedbz, f_lltoij, f_ijtoll, f_converteta, f_computectt) from wrf.var._wrfcape import f_computecape -from wrf.var.decorators import handle_left_iter +from wrf.var.decorators import handle_left_iter, uvmet_left_iter __all__ = ["FortranException", "computeslp", "computetk", "computetd", "computerh", "computeavo", "computepvo", "computeeth", @@ -86,7 +86,7 @@ def computerh(qv,q,t): res = n.reshape(res, shape, "A") return res.astype("float32") -@handle_left_iter(3,3,0, ignore_args=(6,7)) +@handle_left_iter(3,3,0, ignore_args=(6,7), ref_stag_dim=-1) def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy): res = f_computeabsvort(u.astype("float64").T, v.astype("float64").T, @@ -99,7 +99,7 @@ def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy): return res.astype("float32").T -@handle_left_iter(3,3,0, ignore_args=(8,9)) +@handle_left_iter(3,3,2, ignore_args=(8,9)) def computepvo(u,v,theta,prs,msfu,msfv,msfm,cor,dx,dy): res = f_computepvo(u.astype("float64").T, @@ -124,7 +124,7 @@ def computeeth(qv, tk, p): return res.astype("float32").T -@handle_left_iter(4,2,2, ignore_args=(4,5), alg_out_fixed_dims=(2,)) +@uvmet_left_iter() def computeuvmet(u,v,lat,lon,cen_long,cone): longca = n.zeros((lat.shape[-2], lat.shape[-1]), "float64") longcb = n.zeros((lon.shape[-2], lon.shape[-1]), "float64") @@ -132,7 +132,7 @@ def computeuvmet(u,v,lat,lon,cen_long,cone): # 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])) v = v.reshape((1,v.shape[-2], v.shape[-1])) @@ -154,7 +154,7 @@ def computeuvmet(u,v,lat,lon,cen_long,cone): 0) - return res.astype("float32").T + return n.squeeze(res.astype("float32").T) @handle_left_iter(3,3,0) def computeomega(qv, tk, w, p): @@ -190,7 +190,7 @@ def computewetbulb(p,tk,qv): @handle_left_iter(2,3,0, ignore_args=(4,)) def computesrh(u, v, z, ter, top): - + res = f_computesrh(u.astype("float64").T, v.astype("float64").T, z.astype("float64").T, @@ -276,10 +276,11 @@ def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow): FortranException()) # Don't need to transpose since we only passed a view to fortran - cape = flip_cape.astype("float32") - cin = flip_cin.astype("float32") + cape = flip_cape.astype("float32")[::-1,:,:] + cin = flip_cin.astype("float32")[::-1,:,:] + # Remember to flip cape and cin back to descending p coordinates - return (cape[::-1,:,:],cin[::-1,:,:]) + return (cape, cin) # TODO: This should handle lists of coords def computeij(map_proj,truelat1,truelat2,stdlon, diff --git a/wrf_open/var/src/python/wrf/var/geoht.py b/wrf_open/var/src/python/wrf/var/geoht.py index 612a2ac..edd40cc 100755 --- a/wrf_open/var/src/python/wrf/var/geoht.py +++ b/wrf_open/var/src/python/wrf/var/geoht.py @@ -23,13 +23,13 @@ def _get_geoht(wrfnc, height=True, msl=True, timeidx=0): "NetCDF file") else: geopt_unstag = ght_vars["GHT"] * Constants.G - hgt = destagger(ght_vars["HGT_U"], 1) + hgt = destagger(ght_vars["HGT_U"], -1) else: ph = ph_vars["PH"] phb = ph_vars["PHB"] hgt = ph_vars["HGT"] geopt = ph + phb - geopt_unstag = destagger(geopt, 0) + geopt_unstag = destagger(geopt, -3) if height: if msl: diff --git a/wrf_open/var/src/python/wrf/var/helicity.py b/wrf_open/var/src/python/wrf/var/helicity.py index 190dddb..09220b9 100755 --- a/wrf_open/var/src/python/wrf/var/helicity.py +++ b/wrf_open/var/src/python/wrf/var/helicity.py @@ -40,14 +40,14 @@ def get_srh(wrfnc, top=3000.0, timeidx=0): v = destagger(v_vars["V"], -2) geopt = ph + phb - geopt_unstag = destagger(geopt, 0) + geopt_unstag = destagger(geopt, -3) z = geopt_unstag / Constants.G # Re-ordering from high to low - u1 = u[::-1,:,:] - v1 = v[::-1,:,:] - z1 = z[::-1,:,:] + u1 = u[...,::-1,:,:] + v1 = v[...,::-1,:,:] + z1 = z[...,::-1,:,:] srh = computesrh(u1, v1, z1, ter, top) diff --git a/wrf_open/var/src/python/wrf/var/interp.py b/wrf_open/var/src/python/wrf/var/interp.py index 9468e6e..b261c27 100755 --- a/wrf_open/var/src/python/wrf/var/interp.py +++ b/wrf_open/var/src/python/wrf/var/interp.py @@ -35,7 +35,7 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None, """ # Have a pivot point with an angle to find cross section - if pivot is not None and angle is not None: + if pivot_point is not None and angle is not None: xp = pivot_point[0] yp = pivot_point[1] @@ -132,7 +132,7 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None, # TODO: Add flag to use lat/lon points by doing conversion def get_vertcross(data3d, z, missingval=-99999, - pivot=None,angle=None,start_point=None,end_point=None): + pivot_point=None,angle=None,start_point=None,end_point=None): xdim = z.shape[2] ydim = z.shape[1] diff --git a/wrf_open/var/src/python/wrf/var/latlon.py b/wrf_open/var/src/python/wrf/var/latlon.py index 2564353..c5071cb 100755 --- a/wrf_open/var/src/python/wrf/var/latlon.py +++ b/wrf_open/var/src/python/wrf/var/latlon.py @@ -1,3 +1,5 @@ +from collections import Iterable + from wrf.var.extension import computeij, computell from wrf.var.util import extract_vars, extract_global_attrs diff --git a/wrf_open/var/src/python/wrf/var/omega.py b/wrf_open/var/src/python/wrf/var/omega.py index 2f68190..a78812e 100755 --- a/wrf_open/var/src/python/wrf/var/omega.py +++ b/wrf_open/var/src/python/wrf/var/omega.py @@ -14,7 +14,7 @@ def get_omega(wrfnc, timeidx=0): pb = ncvars["PB"] qv = ncvars["QVAPOR"] - wa = destagger(w, 0) + wa = destagger(w, -3) full_t = t + Constants.T_BASE full_p = p + pb tk = computetk(full_p, full_t) diff --git a/wrf_open/var/src/python/wrf/var/slp.py b/wrf_open/var/src/python/wrf/var/slp.py index ed5626b..cbde168 100755 --- a/wrf_open/var/src/python/wrf/var/slp.py +++ b/wrf_open/var/src/python/wrf/var/slp.py @@ -23,7 +23,7 @@ def get_slp(wrfnc, units="hpa", timeidx=0): qvapor[qvapor < 0] = 0. full_ph = (ph + phb) / Constants.G - destag_ph = destagger(full_ph, 0) + destag_ph = destagger(full_ph, -3) tk = computetk(full_p, full_t) slp = computeslp(destag_ph, tk, full_p, qvapor) diff --git a/wrf_open/var/src/python/wrf/var/times.py b/wrf_open/var/src/python/wrf/var/times.py index 2328d85..aba9b71 100755 --- a/wrf_open/var/src/python/wrf/var/times.py +++ b/wrf_open/var/src/python/wrf/var/times.py @@ -4,11 +4,9 @@ import datetime as dt __all__ = ["get_times"] def _make_time(timearr): - return dt.strptime("".join(timearr[:]), "%Y-%m-%d_%H:%M:%S") + return dt.datetime.strptime("".join(timearr[:]), "%Y-%m-%d_%H:%M:%S") # TODO: handle list of files and multiple times def get_times(wrfnc): times = wrfnc.variables["Times"][:,:] - return [_make_time(times[i,:]) for i in xrange(times.shape[0])] - - \ No newline at end of file + return [_make_time(times[i,:]) for i in xrange(times.shape[0])] diff --git a/wrf_open/var/src/python/wrf/var/util.py b/wrf_open/var/src/python/wrf/var/util.py index c2255f5..892d943 100644 --- a/wrf_open/var/src/python/wrf/var/util.py +++ b/wrf_open/var/src/python/wrf/var/util.py @@ -3,7 +3,7 @@ from collections import Iterable import numpy as n __all__ = ["extract_vars", "extract_global_attrs", "extract_dim", - "hold_dim_fixed", "combine_files"] + "combine_files"] def _is_multi_time(timeidx): if timeidx == -1: @@ -82,28 +82,5 @@ def extract_vars(wrfnc, timeidx, vars): 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 - the array shape is unknown. The values for 'dim' and 'idx' can - be negative. - - For example, on a 4D array with 'dim' set to - -1 and 'idx' set to 3, this is going to return this view: - - var[:,:,:,3] - - """ - - var_shape = var.shape - num_dims = len(var_shape) - full_slice = slice(None, None, None) - - # Make the sequence of slices - dim_slices = [full_slice for x in xrange(num_dims)] - - # Replace the dimension with the fixed index - dim_slices[dim] = idx - - return var[dim_slices] \ No newline at end of file diff --git a/wrf_open/var/src/python/wrf/var/uvmet.py b/wrf_open/var/src/python/wrf/var/uvmet.py index 0f6e66b..6181841 100755 --- a/wrf_open/var/src/python/wrf/var/uvmet.py +++ b/wrf_open/var/src/python/wrf/var/uvmet.py @@ -5,7 +5,7 @@ from wrf.var.destagger import destagger from wrf.var.constants import Constants from wrf.var.wind import _calc_wspd_wdir from wrf.var.decorators import convert_units -from wrf.var.util import extract_vars, extract_global_attrs, hold_dim_fixed +from wrf.var.util import extract_vars, extract_global_attrs __all__=["get_uvmet", "get_uvmet10", "get_uvmet_wspd_wdir", "get_uvmet10_wspd_wdir"] @@ -50,7 +50,7 @@ def get_uvmet(wrfnc, ten_m=False, units ="mps", timeidx=0): else: # For met_em files, this just takes the lowest level winds # (3rd dimension from right is bottom_top) - u = destagger(hold_dim_fixed(uu_vars["UU"], -3, 0), -1) # support met_em files + u = destagger(uu_vars["UU"][...,0,:,:], -1) # support met_em files else: u = u_vars["U10"] @@ -64,7 +64,7 @@ def get_uvmet(wrfnc, ten_m=False, units ="mps", timeidx=0): else: # For met_em files, this just takes the lowest level winds # (3rd dimension from right is bottom_top) - v = destagger(hold_dim_fixed(vv_vars["VV"], -3, 0), -2) # support met_em files + v = destagger(vv_vars["VV"][...,0,:,:], -2) # support met_em files else: v = v_vars["V10"] @@ -141,10 +141,7 @@ def get_uvmet(wrfnc, ten_m=False, units ="mps", timeidx=0): res = computeuvmet(u,v,lat,lon,cen_lon,cone) - if u.ndim == 3: - return res - else: - return res[:,0,:,:] + return res def get_uvmet10(wrfnc, units="mps", timeidx=0): diff --git a/wrf_open/var/test/utests.py b/wrf_open/var/test/utests.py index 47773de..776f025 100644 --- a/wrf_open/var/test/utests.py +++ b/wrf_open/var/test/utests.py @@ -1,6 +1,7 @@ import unittest as ut import numpy.testing as nt import numpy as n +import numpy.ma as ma import os, sys import subprocess @@ -44,12 +45,36 @@ def setUpModule(): # Using helpful information at: # http://eli.thegreenplace.net/2014/04/02/dynamically-generating-python-test-cases -def make_test(varname, wrf_in, referent): +def make_test(varname, wrf_in, referent, multi=False, repeat=3): def test(self): - in_wrfnc = NetCDF(wrf_in) - refnc = NetCDF(referent) - ref_vals = refnc.variables[varname][...] + if not multi: + in_wrfnc = NetCDF(wrf_in) + else: + nc = NetCDF(wrf_in) + in_wrfnc = [nc for i in xrange(repeat)] + + refnc = NetCDF(referent) + + if not multi: + ref_vals = refnc.variables[varname][:] + else: + data = refnc.variables[varname][:] + new_dims = [repeat] + [x for x in data.shape] + masked=False + if (isinstance(data, ma.core.MaskedArray)): + masked=True + + if not masked: + ref_vals = n.zeros(new_dims, data.dtype) + else: + ref_vals = ma.asarray(n.zeros(new_dims, data.dtype)) + + for i in xrange(repeat): + ref_vals[i,:] = data[:] + + if masked: + ref_vals.mask[i,:] = data.mask[:] if (varname == "tc"): my_vals = getvar(in_wrfnc, "temp", units="c") @@ -104,20 +129,16 @@ def make_test(varname, wrf_in, referent): atol = 0 nt.assert_allclose(my_vals, ref_vals, tol, atol) elif (varname == "cape_2d"): - mcape, mcin, lcl, lfc = getvar(in_wrfnc, varname) + cape_2d = getvar(in_wrfnc, varname) tol = 1/100. atol = 0 - nt.assert_allclose(mcape, ref_vals[0,:,:], tol, atol) - nt.assert_allclose(mcin, ref_vals[1,:,:], tol, atol) - nt.assert_allclose(lcl, ref_vals[2,:,:], tol, atol) - nt.assert_allclose(lfc, ref_vals[3,:,:], tol, atol) + nt.assert_allclose(cape_2d, ref_vals, tol, atol) elif (varname == "cape_3d"): - cape, cin = getvar(in_wrfnc, varname) + cape_3d = getvar(in_wrfnc, varname) tol = 1/100. atol = 0 - nt.assert_allclose(cape, ref_vals[0,:,:], tol, atol) - nt.assert_allclose(cin, ref_vals[1,:,:], tol, atol) + nt.assert_allclose(cape_3d, ref_vals, tol, atol) else: my_vals = getvar(in_wrfnc, varname) @@ -144,8 +165,10 @@ if __name__ == "__main__": if var in ignore_vars: continue - test_func = make_test(var, TEST_FILE, OUT_NC_FILE) - setattr(WRFVarsTest, 'test_{0}'.format(var), test_func) + test_func1 = make_test(var, TEST_FILE, OUT_NC_FILE) + test_func2 = make_test(var, TEST_FILE, OUT_NC_FILE, multi=True) + setattr(WRFVarsTest, 'test_{0}'.format(var), test_func1) + setattr(WRFVarsTest, 'test_multi_{0}'.format(var), test_func2) ut.main()