snapshot...

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