snapshot...

modules/util/gfs_reader.py

@@ -276,3 +276,51 @@ def get_time_interpolated_point(ds_0, ds_1, time0, time1, fld_name, time, lons,
     intrp_vals = intrp_vals.values
 
     return intrp_vals
+
+
+def get_voxel(xr_dataset, fld_name, lon, lat, press, x_width=5, y_width=5, z_width=3):
+    if lon360:
+        if lon < 0:
+            lon += 360
+    plevs = xr_dataset['pressure levels']
+    fld = xr_dataset[fld_name]
+
+    p_c = value_to_index(plevs.values, press)
+    x_c = value_to_index(lon_coords, lon)
+    y_c = value_to_index(lat_coords, lat)
+    y_h = int(y_width / 2)
+    x_h = int(x_width / 2)
+    p_h = int(z_width / 2)
+
+    sub_fld = fld[y_c-y_h:y_c+y_h+1, x_c-x_h:x_c+x_h+1, p_c-p_h:p_c+p_h+1]
+
+    sub_fld = sub_fld.expand_dims('channel')
+    sub_fld = sub_fld.assign_coords(channel=[fld_name], fakeDim2=lon_coords[x_c-x_h:x_c+x_h+1], fakeDim1=lat_coords[y_c-y_h:y_c+y_h+1], fakeDim0=plevs[p_c-p_h:p_c+p_h+1])
+
+    return sub_fld
+
+
+def get_voxel_s(xr_dataset, fld_name_s, lon, lat, press, x_width=5, y_width=5, z_width=3):
+    if lon360:
+        if lon < 0:
+            lon += 360
+    plevs = xr_dataset['pressure levels']
+
+    p_c = value_to_index(plevs.values, press)
+    x_c = value_to_index(lon_coords, lon)
+    y_c = value_to_index(lat_coords, lat)
+    y_h = int(y_width / 2)
+    x_h = int(x_width / 2)
+    p_h = int(z_width / 2)
+
+    sub_fld_s = []
+    for name in fld_name_s:
+        fld = xr_dataset[name]
+        sub_fld = fld[y_c-y_h:y_c+y_h+1, x_c-x_h:x_c+x_h+1, p_c-p_h:p_c+p_h+1]
+        sub_fld_s.append(sub_fld)
+    sub_fld = xr.concat(sub_fld_s, 'channel')
+
+    sub_fld = sub_fld.assign_coords(channel=fld_name_s, fakeDim2=lon_coords[x_c-x_h:x_c+x_h+1],
+                                    fakeDim1=lat_coords[y_c-y_h:y_c+y_h+1], fakeDim0=plevs[p_c-p_h:p_c+p_h+1])
+
+    return sub_fld
\ No newline at end of file