snapshot...

modules/aeolus/aeolus_amv.py

@@ -579,12 +579,25 @@ def subset_by_time(match_file, dt_str_0, dt_str_1, amv_var_names):
         a_nc_vars.append(rootgrp[vname])
     nvars = len(a_nc_vars)
 
+    mf_vars = list(rootgrp.variables.keys())
+
     p_lon_v = rootgrp['prof_longitude']
     p_lat_v = rootgrp['prof_latitude']
-    p_azm_v = rootgrp['prof_azm']
-    p_spd_v = rootgrp['prof_spd']
-    p_hhb_v = rootgrp['prof_hhb']
-    p_hht_v = rootgrp['prof_hht']
+    p_vars = []
+    p_var_names = []
+    if 'prof_hhb' in mf_vars:
+        p_vars.append(rootgrp['prof_hhb'])
+        p_var_names.append('layer_bot')
+    if 'prof_hht' in mf_vars:
+        p_vars.append(rootgrp['prof_hht'])
+        p_var_names.append('layer_top')
+    if 'prof_spd' in mf_vars:
+        p_vars.append(rootgrp['prof_spd'])
+        p_var_names.append('speed')
+    if 'prof_azm' in mf_vars:
+        p_vars.append(rootgrp['prof_azm'])
+        p_var_names.append('azimuth')
+    npvars = len(p_vars)
 
     dto = datetime.datetime.strptime(dt_str_0, '%Y-%m-%d_%H:%M').replace(tzinfo=timezone.utc)
     dto.replace(tzinfo=timezone.utc)
@@ -614,7 +627,7 @@ def subset_by_time(match_file, dt_str_0, dt_str_1, amv_var_names):
     mx_nlevs = np.max(n_levs_per_prof[time_idxs[0]:time_idxs[0]+n_times])
 
     amvs = np.zeros((n_times, mx_namvs, nvars))
-    profs = np.zeros((n_times, mx_nlevs, 4))
+    profs = np.zeros((n_times, mx_nlevs, npvars))
     amvs.fill(np.nan)
     profs.fill(np.nan)
 
@@ -633,15 +646,13 @@ def subset_by_time(match_file, dt_str_0, dt_str_1, amv_var_names):
         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]
-        profs[idx, 0:n_levs, 2] = p_hhb_v[c:d]
-        profs[idx, 0:n_levs, 3] = p_hht_v[c:d]
+        for k in range(nvars):
+            profs[idx, 0:n_levs, k] = p_vars[k][c:d]
 
         accum_amv += n_amvs
         accum_prf += n_levs
 
-    coords = {'num_profs': times[time_idxs], 'num_params': ['speed', 'azimuth', 'layer_bot', 'layer_top']}
+    coords = {'num_profs': times[time_idxs], 'num_params': p_var_names}
     prof_da = xr.DataArray(profs, coords=coords, dims=['num_profs', 'max_num_levels', 'num_params'])
 
     coords = {'num_profs': times[time_idxs], 'num_params': amv_var_names}