diff --git a/src/wrf/decorators.py b/src/wrf/decorators.py index a93ec1f..443c09f 100644 --- a/src/wrf/decorators.py +++ b/src/wrf/decorators.py @@ -8,7 +8,7 @@ import numpy as np import numpy.ma as ma from .units import do_conversion, check_units -from .util import (iter_left_indexes, viewitems, from_args, npvalues, range2) +from .util import (iter_left_indexes, viewitems, from_args, npvalues, py3range) from .config import xarray_enabled if xarray_enabled(): @@ -86,7 +86,7 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, return wrapped(*args, **kwargs) # Start by getting the left-most 'extra' dims - extra_dims = [ref_var_shape[x] for x in range2(extra_dim_num)] + extra_dims = [ref_var_shape[x] for x in py3range(extra_dim_num)] out_inited = False for left_idxs in iter_left_indexes(extra_dims): @@ -135,14 +135,14 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, outdims = _calc_out_dims(res[0], extra_dims) if not isinstance(res[0], ma.MaskedArray): output = [np.empty(outdims, ref_var.dtype) - for i in range2(len(res))] + for i in py3range(len(res))] masked = False else: output = [ma.MaskedArray( np.zeros(outdims, ref_var.dtype), mask=np.zeros(outdims, np.bool_), fill_value=res[0].fill_value) - for i in range2(len(res))] + for i in py3range(len(res))] masked = True out_inited = True diff --git a/src/wrf/extension.py b/src/wrf/extension.py index 7f1164d..20ff98e 100755 --- a/src/wrf/extension.py +++ b/src/wrf/extension.py @@ -16,7 +16,7 @@ from ._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d, from ._wrfcape import f_computecape from .decorators import (handle_left_iter, handle_casting, handle_extract_transpose) -from .util import range2 +from .util import py3range from .uvdecorator import uvmet_left_iter __all__ = ["FortranException", "computeslp", "computetk", "computetd", @@ -68,7 +68,7 @@ def computevertcross(field3d, xy, var2dz, z_var2d, missingval): var2d = np.empty((z_var2d.size, xy.shape[0]), dtype=var2dz.dtype) var2dtmp = interp2dxy(field3d, xy) - for i in range2(xy.shape[0]): + for i in py3range(xy.shape[0]): var2d[:,i] = interp1d(var2dtmp[:,i], var2dz[:,i], z_var2d, missingval) return var2d diff --git a/src/wrf/interputils.py b/src/wrf/interputils.py index 57e0c27..50ab620 100644 --- a/src/wrf/interputils.py +++ b/src/wrf/interputils.py @@ -6,7 +6,7 @@ from math import floor, ceil import numpy as np from .extension import interp2dxy -from .util import range2 +from .util import py3range __all__ = ["to_positive_idxs", "calc_xy", "get_xy_z_params", "get_xy"] @@ -119,7 +119,7 @@ def calc_xy(xdim, ydim, pivot_point=None, angle=None, dx = dx/npts dy = dy/npts - for i in range2(npts): + for i in py3range(npts): xy[i,0] = x0 + i*dx xy[i,1] = y0 + i*dy @@ -154,7 +154,7 @@ def get_xy_z_params(z, pivot_point=None, angle=None, z_var2d = np.zeros((nlevels), dtype=z.dtype) z_var2d[0] = z_min - for i in range2(1,nlevels): + for i in py3range(1,nlevels): z_var2d[i] = z_var2d[0] + i*dz return xy, var2dz, z_var2d diff --git a/src/wrf/metadecorators.py b/src/wrf/metadecorators.py index 8f80a07..785f13f 100644 --- a/src/wrf/metadecorators.py +++ b/src/wrf/metadecorators.py @@ -9,7 +9,7 @@ import numpy.ma as ma from .util import (viewkeys, viewitems, extract_vars, combine_with, either, from_args, arg_location, - _is_coord_var, CoordPair, npvalues, range2, ucode) + _is_coord_var, CoordPair, npvalues, py3range, ucode) from .interputils import get_xy_z_params, get_xy from .latlonutils import ij_to_ll, ll_to_ij from .config import xarray_enabled @@ -458,7 +458,7 @@ def _set_cross_meta(wrapped, instance, args, kwargs): outattrs = OrderedDict() outdimnames = list(field3d.dims) outcoords.update(field3d.coords) - for i in range2(-3,0,1): + for i in py3range(-3,0,1): outdimnames.remove(field3d.dims[i]) del outcoords[field3d.dims[i]] @@ -481,7 +481,7 @@ def _set_cross_meta(wrapped, instance, args, kwargs): pass outcoords["xy_loc"] = ("xy", [CoordPair(xy[i,0], xy[i,1]) - for i in range2(xy.shape[-2])]) + for i in py3range(xy.shape[-2])]) outcoords["vertical"] = z_var2d[:] @@ -549,7 +549,7 @@ def _set_line_meta(wrapped, instance, args, kwargs): outattrs = OrderedDict() outdimnames = list(field2d.dims) outcoords.update(field2d.coords) - for i in range2(-2,0,1): + for i in py3range(-2,0,1): outdimnames.remove(field2d.dims[i]) del outcoords[field2d.dims[i]] @@ -570,7 +570,7 @@ def _set_line_meta(wrapped, instance, args, kwargs): pass outcoords["xy_loc"] = ("xy", [CoordPair(xy[i,0], xy[i,1]) - for i in range2(xy.shape[-2])]) + for i in py3range(xy.shape[-2])]) else: outname = "field2d_line" @@ -652,7 +652,7 @@ def _set_2dxy_meta(wrapped, instance, args, kwargs): outattrs = OrderedDict() outdimnames = list(field3d.dims) outcoords.update(field3d.coords) - for i in range2(-2,0,1): + for i in py3range(-2,0,1): outdimnames.remove(field3d.dims[i]) del outcoords[field3d.dims[i]] @@ -662,7 +662,7 @@ def _set_2dxy_meta(wrapped, instance, args, kwargs): outname = "{0}_xy".format(field3d.name) outcoords["xy_loc"] = ("xy", [CoordPair(xy[i,0], xy[i,1]) - for i in range2(xy.shape[-2])]) + for i in py3range(xy.shape[-2])]) for key in ("MemoryOrder",): try: diff --git a/src/wrf/util.py b/src/wrf/util.py index c15b512..6a644ae 100644 --- a/src/wrf/util.py +++ b/src/wrf/util.py @@ -36,7 +36,7 @@ __all__ = ["extract_vars", "extract_global_attrs", "extract_dim", "combine_files", "is_standard_wrf_var", "extract_times", "iter_left_indexes", "get_left_indexes", "get_right_slices", "is_staggered", "get_proj_params", "viewitems", "viewkeys", - "viewvalues", "py2round", "combine_with", "either", + "viewvalues", "py2round", "py3range", "combine_with", "either", "from_args", "arg_location", "args_to_list", "npvalues", "CoordPair"] @@ -198,7 +198,7 @@ def py2round(x, d=0): return round(x, d) -def range2(*args): +def py3range(*args): if version_info >= (3,): return range(*args) @@ -285,7 +285,7 @@ def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar): lons = wrfnc.variables[lonvar] # Need to check all times - for i in range2(lats.shape[-3]): + for i in py3range(lats.shape[-3]): start_idxs = [0]*len(lats.shape) # PyNIO does not support ndim start_idxs[-3] = i start_idxs = tuple(start_idxs) @@ -603,7 +603,7 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain): varname) proj = [getproj(lats=lats[i,:], lons=lons[i,:], - **proj_params) for i in range2(lats.shape[0])] + **proj_params) for i in py3range(lats.shape[0])] if time_coord is not None: coords[time_coord] = (lon_coord_var.dimensions[0], @@ -801,7 +801,7 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta): varname) projs = [getproj(lats=lats[i,:], lons=lons[i,:], - **proj_params) for i in range2(lats.shape[0])] + **proj_params) for i in py3range(lats.shape[0])] outprojs[startidx:endidx] = np.asarray(projs, np.object)[:] @@ -962,7 +962,7 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta): varname) projs = [getproj(lats=lats[i,:], lons=lons[i,:], - **proj_params) for i in range2(lats.shape[0])] + **proj_params) for i in py3range(lats.shape[0])] outprojs[file_idx, 0:numtimes] = ( np.asarray(projs, np.object)[:]) @@ -982,7 +982,7 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta): outattrs = OrderedDict(first_var.attrs) # New dimensions outdimnames = ["file"] + list(first_var.dims) - outcoords["file"] = [i for i in range2(numfiles)] + outcoords["file"] = [i for i in py3range(numfiles)] # Time needs to be multi dimensional, so use the default dimension del outcoords["Time"] @@ -1121,7 +1121,7 @@ def _file_times(wrfnc, timeidx): multitime = _is_multi_time_req(timeidx) if multitime: times = wrfnc.variables["Times"][:,:] - for i in range2(times.shape[0]): + for i in py3range(times.shape[0]): yield _make_time(times[i,:]) else: times = wrfnc.variables["Times"][timeidx,:] @@ -1177,7 +1177,7 @@ def get_left_indexes(ref_var, expected_dims): if (extra_dim_num == 0): return [] - return tuple([ref_var.shape[x] for x in range2(extra_dim_num)]) + return tuple([ref_var.shape[x] for x in py3range(extra_dim_num)]) def iter_left_indexes(dims): """A generator which yields the iteration tuples for a sequence of @@ -1192,7 +1192,7 @@ def iter_left_indexes(dims): - dims - a sequence of dimensions sizes (e.g. ndarry.shape) """ - arg = [range2(dim) for dim in dims] + arg = [py3range(dim) for dim in dims] for idxs in product(*arg): yield idxs diff --git a/src/wrf/uvdecorator.py b/src/wrf/uvdecorator.py index cc4ccc6..9b1b664 100644 --- a/src/wrf/uvdecorator.py +++ b/src/wrf/uvdecorator.py @@ -6,7 +6,7 @@ import numpy as np import wrapt from .destag import destagger -from .util import iter_left_indexes, range2 +from .util import iter_left_indexes, py3range __all__ = ["uvmet_left_iter"] @@ -53,13 +53,13 @@ def uvmet_left_iter(): return wrapped(u, v, lat, lon, cen_long, cone) # Start by getting the left-most 'extra' dims - outdims = [u.shape[x] for x in range2(extra_dim_num)] + outdims = [u.shape[x] for x in py3range(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 range2(-num_right_dims,0,1)] + outdims += [u.shape[x] for x in py3range(-num_right_dims,0,1)] output = np.empty(outdims, u.dtype) diff --git a/test/utests.py b/test/utests.py index 1042084..d0dbb2e 100644 --- a/test/utests.py +++ b/test/utests.py @@ -12,6 +12,10 @@ NCL_EXE = "/Users/ladwig/nclbuild/6.3.0/bin/ncl" TEST_FILE = "/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00" OUT_NC_FILE = "/tmp/wrftest.nc" +# Python 3 +if sys.version_info > (3,): + xrange = range + def setUpModule(): ncarg_root = os.environ.get("NCARG_ROOT", None) if ncarg_root is None: