def interpz(self, zcoord, levels, dim_name='', fill_value=np.NaN, use_multiprocessing=True): """Interpolates the array along the vertical dimension Parameters ---------- zcoord: DataArray the z coordinates of the variable. Must be of same dimensions levels: 1dArray the levels at which to interpolate dim_name: str the name of the new dimension fill_value : np.NaN or 'extrapolate', optional how to handle levels below the topography. Default is to mark them as invalid, but you might want the have them extrapolated. use_multiprocessing: bool set to false if, for some reason, you don't want to use mp Returns ------- a new DataArray with the interpolated data """ if self.z_dim is None: raise RuntimeError('zdimension not recognized') data = wrftools.interp3d(self._obj.values, zcoord.values, np.atleast_1d(levels), fill_value=fill_value, use_multiprocessing=use_multiprocessing) dims = list(self._obj.dims) zd = np.nonzero([self.z_dim == d for d in dims])[0][0] dims[zd] = dim_name coords = dict(self._obj.coords) coords.pop(self.z_dim, None) coords[dim_name] = np.atleast_1d(levels) out = xr.DataArray(data, name=self._obj.name, dims=dims, coords=coords) out.attrs['pyproj_srs'] = self.grid.proj.srs if not np.asarray(levels).shape: out = out.isel(**{dim_name: 0}) return out
def interpz(self, zcoord, levels, dim_name='', use_multiprocessing=True): """Interpolates the array along the vertical dimension Parameters ---------- zcoord: DataArray the z coordinates of the variable. Must be of same dimensions levels: 1dArray the levels at which to interpolate dim_name: str the name of the new dimension use_multiprocessing: bool set to false if, for some reason, you don't want to use mp Returns ------- a new DataArray with the interpolated data """ if self.z_dim is None: raise RuntimeError('zdimension not recognized') data = wrftools.interp3d(self._obj.values, zcoord.values, np.atleast_1d(levels), use_multiprocessing= use_multiprocessing) dims = list(self._obj.dims) zd = np.nonzero([self.z_dim == d for d in dims])[0][0] dims[zd] = dim_name coords = dict(self._obj.coords) coords.pop(self.z_dim, None) coords[dim_name] = np.atleast_1d(levels) out = xr.DataArray(data, name=self._obj.name, dims=dims, coords=coords) out.attrs['pyproj_srs'] = self.grid.proj.srs if not np.asarray(levels).shape: out = out.isel(**{dim_name:0}) return out