snapshot..

modules/aeolus/aeolus_amv.py

@@ -592,10 +592,11 @@ def create_amv_to_aeolus_match_file(aeolus_files_dir, amv_files_dir, outfile=Non
 # match_file: pathname for the product file
 # dt_str_0: start time (YYYY-MM-DD_HH:MM)
 # dt_str_1: end time (YYYY-MM-DD_HH:MM)
+# amv_var_names: list of amv parameters (see match_file) to subset
 # returns: Xarray.DataArray
 # amvs[nprofs, max_num_amvs_per_prof, num_of_params], profs[nprofs, max_num_levs_per_prof, num_of_params],
 # prof_locs[nprofs, (lon, lat)
-def subset_by_time(match_file, dt_str_0, dt_str_1):
+def subset_by_time(match_file, dt_str_0, dt_str_1, amv_var_names):
     rootgrp = Dataset(match_file, 'r', format='NETCDF4')
     all_dims = rootgrp.dimensions
     t_var = rootgrp['time']
@@ -604,10 +605,10 @@ def subset_by_time(match_file, dt_str_0, dt_str_1):
     n_amvs_per_prof = rootgrp['num_amvs_per_prof'][:]
     n_levs_per_prof = rootgrp['num_levs_per_prof'][:]
 
-    a_dir_v = rootgrp['amv_dir']
-    a_spd_v = rootgrp['amv_spd']
-    a_dst_v = rootgrp['amv_dist']
-    a_prs_v = rootgrp['amv_pres']
+    a_nc_vars = []
+    for vname in amv_var_names:
+        a_nc_vars.append(rootgrp[vname])
+    nvars = len(a_nc_vars)
 
     p_lon_v = rootgrp['prof_longitude']
     p_lat_v = rootgrp['prof_latitude']
@@ -643,8 +644,8 @@ def subset_by_time(match_file, dt_str_0, dt_str_1):
     mx_namvs = np.max(n_amvs_per_prof[time_idxs[0]:time_idxs[0]+n_times])
     mx_nlevs = np.max(n_levs_per_prof[time_idxs[0]:time_idxs[0]+n_times])
 
-    amvs = np.zeros((n_times, mx_namvs, 4))
-    profs = np.zeros((n_times, mx_nlevs, 4))
+    amvs = np.zeros((n_times, mx_namvs, nvars))
+    profs = np.zeros((n_times, mx_nlevs, nvars))
     amvs.fill(np.nan)
     profs.fill(np.nan)
 
@@ -660,10 +661,8 @@ def subset_by_time(match_file, dt_str_0, dt_str_1):
         c = accum_prf
         d = accum_prf + n_levs
 
-        amvs[idx, 0:n_amvs, 0] = a_spd_v[a:b]
-        amvs[idx, 0:n_amvs, 1] = a_dir_v[a:b]
-        amvs[idx, 0:n_amvs, 2] = a_prs_v[a:b]
-        amvs[idx, 0:n_amvs, 3] = a_dst_v[a:b]
+        for k in range(nvars):
+            amvs[idx, 0:n_amvs, k] = a_nc_vars[k][a:b]
 
         profs[idx, 0:n_levs, 0] = p_spd_v[c:d]
         profs[idx, 0:n_levs, 1] = p_azm_v[c:d]
@@ -676,7 +675,7 @@ def subset_by_time(match_file, dt_str_0, dt_str_1):
     coords = {'num_profs': times[time_idxs], 'num_params': ['speed', 'azimuth', 'layer_bot', 'layer_top']}
     prof_da = xr.DataArray(profs, coords=coords, dims=['num_profs', 'max_num_levels', 'num_params'])
 
-    coords = {'num_profs': times[time_idxs], 'num_params': ['speed', 'azimuth', 'pressure', 'distance']}
+    coords = {'num_profs': times[time_idxs], 'num_params': amv_var_names}
     amvs_da = xr.DataArray(amvs, coords=coords, dims=['num_profs', 'max_num_amvs', 'num_params'])
 
     prof_locs_da = xr.DataArray(np.column_stack([lons[time_idxs], lats[time_idxs]]),