diff --git a/modules/util/gfs_reader.py b/modules/util/gfs_reader.py index f6f6f8c8a9b713fc1cd834a05535664da6a0a929..343c5afcfd396069d251a0385c777701d092e1f2 100644 --- a/modules/util/gfs_reader.py +++ b/modules/util/gfs_reader.py @@ -232,7 +232,7 @@ def get_vert_profile_s(xr_dataset, fld_name_s, lons, lats, method='linear'): return intrp_fld -def get_point(xr_dataset, fld_name, lons, lats, method='nearest'): +def get_point(xr_dataset, fld_name, lons, lats, pres_s=None, method='nearest'): if lon360: # convert -180/+180 to 0,360 lons = np.where(lons < 0, lons + 360, lons) # convert -180,180 to 0,360 @@ -240,17 +240,24 @@ def get_point(xr_dataset, fld_name, lons, lats, method='nearest'): lon_coords = np.linspace(0, 359.5, xr_dataset.fakeDim2.size) fld = xr_dataset[fld_name] - fld = fld.assign_coords(fakeDim2=lon_coords, fakeDim1=lat_coords) + + if pres_s is not None: + fld = fld.assign_coords(fakeDim2=lon_coords, fakeDim1=lat_coords, fakeDim0=plevs) + else: + fld = fld.assign_coords(fakeDim2=lon_coords, fakeDim1=lat_coords) dim1 = xr.DataArray(lats, dims='k') dim2 = xr.DataArray(lons, dims='k') - - intrp_fld = fld.interp(fakeDim1=dim1, fakeDim2=dim2, method=method) + if pres_s is not None: + dim0 = xr.DataArray(pres_s, dims='k') + intrp_fld = fld.interp(fakeDim0=dim0, fakeDim1=dim1, fakeDim2=dim2, method=method) + else: + intrp_fld = fld.interp(fakeDim1=dim1, fakeDim2=dim2, method=method) return intrp_fld -def get_point_s(xr_dataset, fld_name_s, lons, lats, method='nearest'): +def get_point_s(xr_dataset, fld_name_s, lons, lats, pres_s=None, method='nearest'): if lon360: # convert -180/+180 to 0,360 lons = np.where(lons < 0, lons + 360, lons) # convert -180,180 to 0,360 @@ -260,14 +267,20 @@ def get_point_s(xr_dataset, fld_name_s, lons, lats, method='nearest'): fld_s = [] for fld_name in fld_name_s: fld = xr_dataset[fld_name] - fld = fld.assign_coords(fakeDim2=lon_coords, fakeDim1=lat_coords) + if pres_s is not None: + fld = fld.assign_coords(fakeDim2=lon_coords, fakeDim1=lat_coords, fakeDim0=plevs) + else: + fld = fld.assign_coords(fakeDim2=lon_coords, fakeDim1=lat_coords) fld_s.append(fld) fld = xr.concat(fld_s, 'fld_dim') dim1 = xr.DataArray(lats, dims='k') dim2 = xr.DataArray(lons, dims='k') - - intrp_fld = fld.interp(fakeDim1=dim1, fakeDim2=dim2, method=method) + if pres_s is not None: + dim0 = xr.DataArray(pres_s, dims='k') + intrp_fld = fld.interp(fakeDim0=dim0, fakeDim1=dim1, fakeDim2=dim2, method=method) + else: + intrp_fld = fld.interp(fakeDim1=dim1, fakeDim2=dim2, method=method) return intrp_fld