diff --git a/modules/util/gfs_reader.py b/modules/util/gfs_reader.py index 06b5386ad4e65c968c1a466e6163a6b563c9fb8f..058822285936d109cdf59df878a979c14204103c 100644 --- a/modules/util/gfs_reader.py +++ b/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