snapshot...

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