diff --git a/modules/util/gfs_reader.py b/modules/util/gfs_reader.py index 058822285936d109cdf59df878a979c14204103c..2df215ece7f408ae48126cbce92e9da14281dbf8 100644 --- a/modules/util/gfs_reader.py +++ b/modules/util/gfs_reader.py @@ -5,6 +5,7 @@ import os import numpy as np import xarray as xr # from util.util import value_to_index, homedir +from metpy.units import units # gfs_directory = '/ships22/cloud/Ancil_Data/clavrx_ancil_data/dynamic/gfs/' homedir = os.path.expanduser('~') + '/' @@ -102,12 +103,9 @@ def get_bounding_gfs_files(timestamp): def get_horz_layer(xr_dataset, fld_name, press, lon_range=None, lat_range=None): p_idx = value_to_index(plevs, press) - fld = xr_dataset[fld_name] - x_lo = 0 - x_hi = NX - y_lo = 0 - y_hi = NY + x_lo, x_hi = 0, NX + y_lo, y_hi = 0, NY if lon_range is not None: lon_lo = lon_range[0] @@ -125,12 +123,13 @@ def get_horz_layer(xr_dataset, fld_name, press, lon_range=None, lat_range=None): y_lo = value_to_index(lat_coords, lat_lo) y_hi = value_to_index(lat_coords, lat_hi) - sub_fld = fld[y_lo:y_hi, x_lo:x_hi, p_idx] - sub_fld = sub_fld.expand_dims('channel') - - sub_fld = sub_fld.assign_coords(channel=[fld_name], fakeDim2=lon_coords[x_lo:x_hi], fakeDim1=lat_coords[y_lo:y_hi]) + nda = xr_dataset[fld_name].values + sub_nda = nda[y_lo:y_hi, x_lo:x_hi, p_idx] + xra = xr.DataArray(sub_nda, dims=['latitude', 'longitude'], + coords={"latitude": lat_coords[y_lo:y_hi], "longitude": lon_coords[x_lo:x_hi]}, + attrs={"description": fld_name, "units": 'm/s'}) - return sub_fld + return xra def get_horz_layer_s(xr_dataset, fld_names, press, lon_range=None, lat_range=None): @@ -444,14 +443,11 @@ 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_s, press_range=None, lon_range=None, lat_range=None): +def get_volume(xr_dataset, fld_name, press_range=None, lon_range=None, lat_range=None): - x_lo = 0 - x_hi = NX - y_lo = 0 - y_hi = NY - z_lo = 0 - z_hi = NZ + x_lo, x_hi = 0, NX + y_lo, y_hi = 0, NY + z_lo, z_hi = 0, NZ if lon_range is not None: lon_lo = lon_range[0] @@ -477,13 +473,10 @@ def get_volume(xr_dataset, fld_name_s, press_range=None, lon_range=None, lat_ran z_lo = value_to_index(plevs, press_range[0]) z_hi = value_to_index(plevs, press_range[1]) - 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') + nda = xr_dataset[fld_name].values + sub_nda = nda[y_lo:y_hi, x_lo:x_hi, z_lo:z_hi] + xra = xr.DataArray(sub_nda, dims=['Latitude', 'Longitude', 'Pressure'], + coords={"Latitude": lat_coords[y_lo:y_hi], "Longitude": lon_coords[x_lo:x_hi], "Pressure": plevs[z_lo:z_hi]}, + attrs={"description": fld_name, "units": 'm s-1'}) - 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 + return xra