snapshot...

modules/util/gfs_reader.py

@@ -444,9 +444,7 @@ def get_voxel_s(xr_dataset, fld_name_s, lon, lat, press, x_width=3, y_width=3, z
     return sub_fld
 
 
-def get_volume(xr_dataset, fld_name, press_range=None, lon_range=None, lat_range=None):
-    # p_idx = value_to_index(plevs, press)
-    fld = xr_dataset[fld_name]
+def get_volume(xr_dataset, fld_name_s, press_range=None, lon_range=None, lat_range=None):
 
     x_lo = 0
     x_hi = NX
@@ -479,9 +477,13 @@ def get_volume(xr_dataset, fld_name, press_range=None, lon_range=None, lat_range
         z_lo = value_to_index(plevs, press_range[0])
         z_hi = value_to_index(plevs, press_range[1])
 
-    sub_fld = fld[y_lo:y_hi, x_lo:x_hi, z_lo:z_hi]
-    sub_fld = sub_fld.expand_dims('channel')
+    sub_fld_s = []
+    for name in fld_name_s:
+        fld = xr_dataset[name]
+        sub_fld = fld[y_lo:y_hi, x_lo:x_hi, z_lo:z_hi]
+        sub_fld_s.append(sub_fld)
+    sub_fld = xr.concat(sub_fld_s, 'channel')
 
-    sub_fld = sub_fld.assign_coords(channel=[fld_name], fakeDim2=lon_coords[x_lo:x_hi], fakeDim1=lat_coords[y_lo:y_hi], fakeDim0=plevs[z_lo:z_hi])
+    sub_fld = sub_fld.assign_coords(channel=fld_name_s, fakeDim2=lon_coords[x_lo:x_hi], fakeDim1=lat_coords[y_lo:y_hi], fakeDim0=plevs[z_lo:z_hi])
 
     return sub_fld