snapshot...

modules/util/gfs_reader.py

@@ -147,12 +147,14 @@ def get_horz_layer_s(xr_dataset, fld_names, press, lon_range=None, lat_range=Non
     return sub_fld
 
 
-def get_time_interpolated_layer(ds_0, ds_1, time_0, time_1, time, fld_name, press, lon_range=None, lat_range=None):
-    lyr_0 = get_horz_layer(ds_0, fld_name, press, lon_range=lon_range, lat_range=lat_range)
-    lyr_1 = get_horz_layer(ds_1, fld_name, press, lon_range=lon_range, lat_range=lat_range)
+def get_time_interpolated_layer(xr_dataset_s, time_s, time, fld_name, press, lon_range=None, lat_range=None, method='linear'):
+    layer_s = []
+    for ds in xr_dataset_s:
+        lyr = get_horz_layer(ds, fld_name, press, lon_range=lon_range, lat_range=lat_range)
+        layer_s.append(lyr)
 
-    lyr = xr.concat([lyr_0, lyr_1], 'time')
-    lyr = lyr.assign_coords(time=[time_0, time_1])
+    lyr = xr.concat(layer_s, 'time')
+    lyr = lyr.assign_coords(time=time_s)
 
     intrp_lyr = lyr.interp(time=time, method='linear')
     intrp_lyr = intrp_lyr.expand_dims('channel')
@@ -161,12 +163,14 @@ def get_time_interpolated_layer(ds_0, ds_1, time_0, time_1, time, fld_name, pres
     return intrp_lyr
 
 
-def get_time_interpolated_layer_s(ds_0, ds_1, time_0, time_1, time, fld_name_s, press, lon_range=None, lat_range=None):
-    lyr_s_0 = get_horz_layer_s(ds_0, fld_name_s, press, lon_range=lon_range, lat_range=lat_range)
-    lyr_s_1 = get_horz_layer_s(ds_1, fld_name_s, press, lon_range=lon_range, lat_range=lat_range)
+def get_time_interpolated_layer_s(xr_dataset_s, time_s, time, fld_name_s, press, lon_range=None, lat_range=None, method='linear'):
+    layer_s = []
+    for ds in xr_dataset_s:
+        lyr = get_horz_layer_s(ds, fld_name_s, press, lon_range=lon_range, lat_range=lat_range, method=method)
+        layer_s.append(lyr)
 
-    lyr = xr.concat([lyr_s_0, lyr_s_1], 'time')
-    lyr = lyr.assign_coords(time=[time_0, time_1])
+    lyr = xr.concat(layer_s, 'time')
+    lyr = lyr.assign_coords(time=time_s)
 
     intrp_lyr = lyr.interp(time=time, method='linear')
 
@@ -230,31 +234,31 @@ def get_point(xr_dataset, fld_name, lons, lats):
     return intrp_fld
 
 
-def get_time_interpolated_vert_profile(ds_0, ds_1, time0, time1, fld_name, time, lons, lats):
+def get_time_interpolated_vert_profile(xr_dataset_s, time_s, fld_name, time, lons, lats, method='linear'):
+    prof_s = []
+    for ds in xr_dataset_s:
+        vp = get_vert_profile(ds, fld_name, lons, lats)
+        prof_s.append(vp)
 
-    prof_0 = get_vert_profile(ds_0, fld_name, lons, lats)
-    prof_1 = get_vert_profile(ds_1, fld_name, lons, lats)
-
-    prof = xr.concat([prof_0, prof_1], 'time')
-    prof = prof.assign_coords(time=[time0, time1])
-
-    intrp_prof = prof.interp(time=time, method='linear')
+    prof = xr.concat(prof_s, 'time')
+    prof = prof.assign_coords(time=time_s)
 
+    intrp_prof = prof.interp(time=time, method=method)
     intrp_prof = intrp_prof.values
 
     return intrp_prof
 
 
-def get_time_interpolated_vert_profile_s(ds_0, ds_1, time0, time1, fld_name_s, time, lons, lats):
-
-    prof_0 = get_vert_profile_s(ds_0, fld_name_s, lons, lats)
-    prof_1 = get_vert_profile_s(ds_1, fld_name_s, lons, lats)
-
-    prof = xr.concat([prof_0, prof_1], 'time')
-    prof = prof.assign_coords(time=[time0, time1])
+def get_time_interpolated_vert_profile_s(xr_dataset_s, time_s, fld_name_s, time, lons, lats, method='linear'):
+    prof_s = []
+    for ds in xr_dataset_s:
+        vp = get_vert_profile(ds, fld_name_s, lons, lats)
+        prof_s.append(vp)
 
-    intrp_prof = prof.interp(time=time, method='linear')
+    prof = xr.concat(prof_s, 'time')
+    prof = prof.assign_coords(time=time_s)
 
+    intrp_prof = prof.interp(time=time, method=method)
     intrp_prof = intrp_prof.values
 
     return intrp_prof