diff --git a/modules/aeolus/aeolus_amv.py b/modules/aeolus/aeolus_amv.py index 99a03c0c4c1e751a3312bb16529a1682c182f564..2f02e0f2ce642ed6d9e2d76c4f5adbad5c639d0f 100644 --- a/modules/aeolus/aeolus_amv.py +++ b/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]]),