diff --git a/src/wrf/interp.py b/src/wrf/interp.py index 771d353..b7d939e 100755 --- a/src/wrf/interp.py +++ b/src/wrf/interp.py @@ -8,7 +8,8 @@ from .extension import (_interpz3d, _vertcross, _interpline, _smooth2d, _monotonic, _vintrp) from .metadecorators import set_interp_metadata -from .util import extract_vars, is_staggered, get_id +from .util import extract_vars, is_staggered, get_id, npvalues +from .py3compat import py3range from .interputils import get_xy, get_xy_z_params from .constants import Constants, ConversionFactors from .terrain import get_terrain @@ -54,10 +55,22 @@ def interplevel(field3d, z, desiredlev, missing=Constants.DEFAULT_FILL, `numpy.ndarray` object with no metadata. """ - r1 = _interpz3d(field3d, z, desiredlev, missing) - masked_r1 = ma.masked_values (r1, missing) + # Some fields like uvmet have an extra left dimension for the product + # type, we'll handle that iteration here. + multi = True if field3d.ndim - z.ndim == 1 else False - return masked_r1 + if not multi: + result = _interpz3d(field3d, z, desiredlev, missing) + else: + outshape = field3d.shape[0:-3] + field3d.shape[-2:] + result = np.empty(outshape, dtype=field3d.dtype) + + for i in py3range(field3d.shape[0]): + result[i,:] = ( + _interpz3d(field3d[i,:], z, desiredlev, missing)[:]) + + return ma.masked_values (result, missing) + @set_interp_metadata("cross") def vertcross(field3d, z, missing=Constants.DEFAULT_FILL, @@ -125,18 +138,29 @@ def vertcross(field3d, z, missing=Constants.DEFAULT_FILL, `numpy.ndarray` object with no metadata. """ + # Some fields like uvmet have an extra left dimension for the product + # type, we'll handle that iteration here. + multi = True if field3d.ndim - z.ndim == 1 else False try: xy = cache["xy"] var2dz = cache["var2dz"] z_var2d = cache["z_var2d"] except (KeyError, TypeError): - xy, var2dz, z_var2d = get_xy_z_params(z, pivot_point, angle, + xy, var2dz, z_var2d = get_xy_z_params(npvalues(z), pivot_point, angle, start_point, end_point) - - res = _vertcross(field3d, xy, var2dz, z_var2d, missing) - return ma.masked_values(res, missing) + if not multi: + result = _vertcross(field3d, xy, var2dz, z_var2d, missing) + else: + outshape = field3d.shape[0:-3] + (z_var2d.shape[0], xy.shape[0]) + result = np.empty(outshape, dtype=field3d.dtype) + + for i in py3range(field3d.shape[0]): + result[i,:] = _vertcross(field3d[i,:], xy, var2dz, z_var2d, + missing)[:] + + return ma.masked_values(result, missing) @set_interp_metadata("line") @@ -219,8 +243,8 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, Input WRF ARW NetCDF data as a `netCDF4.Dataset`, `Nio.NioFile` or an iterable sequence of the aforementioned types. - field2d : `xarray.DataArray` or `numpy.ndarray` - A two-dimensional field. + field : `xarray.DataArray` or `numpy.ndarray` + A three-dimensional field. vert_coord : {'pressure', 'pres', 'p', 'ght_msl', 'ght_agl', 'theta', 'th', 'theta-e', 'thetae', 'eth'} A string indicating the vertical coordinate type to interpolate to. diff --git a/src/wrf/interputils.py b/src/wrf/interputils.py index 1af701f..645e4ef 100644 --- a/src/wrf/interputils.py +++ b/src/wrf/interputils.py @@ -35,6 +35,11 @@ def calc_xy(xdim, ydim, pivot_point=None, angle=None, xp = pivot_point[-2] yp = pivot_point[-1] + if xp >= xdim or yp >= ydim: + raise ValueError("pivot point {} is outside of domain " + "with shape {}".format(pivot_point, + (xdim, ydim))) + if (angle > 315.0 or angle < 45.0 or ((angle > 135.0) and (angle < 225.0))): @@ -101,6 +106,15 @@ def calc_xy(xdim, ydim, pivot_point=None, angle=None, y0 = start_point[-1] x1 = end_point[-2] y1 = end_point[-1] + + if x0 >= xdim or y0 >= ydim: + raise ValueError("start_point {} is outside of domain " + "with shape {}".format(start_point, (xdim, ydim))) + + if x1 >= xdim or y1 >= ydim: + raise ValueError("end_point {} is outside of domain " + "with shape {}".format(end_point, (xdim, ydim))) + if ( x1 > xdim-1 ): x1 = xdim if ( y1 > ydim-1): diff --git a/src/wrf/util.py b/src/wrf/util.py index 3ab74e2..742b16d 100644 --- a/src/wrf/util.py +++ b/src/wrf/util.py @@ -290,8 +290,8 @@ class from_var(object): def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar): - lats = wrfnc.variables[latvar] - lons = wrfnc.variables[lonvar] + lats = wrfnc.variables[latvar][:] + lons = wrfnc.variables[lonvar][:] # Need to check all times for i in py3range(lats.shape[-3]): @@ -362,8 +362,8 @@ def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"), return moving # Need to search all the files - lats = first_wrfnc.variables[lat_coord] - lons = first_wrfnc.variables[lon_coord] + lats = first_wrfnc.variables[lat_coord][:] + lons = first_wrfnc.variables[lon_coord][:] zero_idxs = [0]*len(lats.shape) # PyNIO doesn't have ndim last_idxs = list(zero_idxs) @@ -380,6 +380,13 @@ def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"), ll_corner = (lat0, lon0) ur_corner = (lat1, lon1) + # Need to check if the first file is moving, might be a single + # file with multiple times + if _corners_moved(first_wrfnc, ll_corner, ur_corner, lat_coord, lon_coord): + cache_item(_key, product, True) + return True + + # Now check any additional files while True: try: wrfnc = next(wrf_iter) @@ -391,7 +398,6 @@ def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"), cache_item(_key, product, True) return True - cache_item(_key, product, False)