diff --git a/modules/aeolus/aeolus_amv.py b/modules/aeolus/aeolus_amv.py index 335c2efad233da188e997d99ff3d1ee86d95abc0..2f56668ea35bbab99b147d5170ff49412e1e6a9c 100644 --- a/modules/aeolus/aeolus_amv.py +++ b/modules/aeolus/aeolus_amv.py @@ -437,3 +437,68 @@ def create_amv_to_aeolus_match_file(aeolus_files_dir, amv_files_dir, outfile=Non if outfile is not None: create_file(outfile, m_d, a_d) + +def get_subset(pathname, t_0, t_1): + rootgrp = Dataset(pathname, 'r', format='NETCDF4') + all_dims = rootgrp.dimensions + t_var = rootgrp['time'] + + n_profs = len(all_dims['num_aeolus_profs']) + 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'] + + p_azm_v = rootgrp['prof_azm'] + p_spd_v = rootgrp['prof_spd'] + p_hhb_v = rootgrp['prof_hhb'] + p_hht_v = rootgrp['prof_hht'] + + times = t_var[:] + time_idxs = np.arange(n_profs) + valid = np.logical_and(times >= t_0, times < t_1) + time_idxs = time_idxs[valid] + n_times = time_idxs.shape[0] + + prf_idx_start = np.sum(n_levs_per_prof[0:time_idxs[0]]) + amv_idx_start = np.sum(n_amvs_per_prof[0:time_idxs[0]]) + + mx_namvs = np.max(n_amvs_per_prof[amv_idx_start:amv_idx_start+n_times]) + mx_nlevs = np.max(n_levs_per_prof[prf_idx_start:prf_idx_start+n_times]) + + amvs = np.zeros((n_times, mx_namvs, 4)) + profs = np.zeros((n_times, mx_nlevs, 4)) + amvs.fill(np.nan) + profs.fill(np.nan) + + accum_prf = 0 + accum_amv = 0 + for idx, t_i in enumerate(time_idxs): + n_amvs = n_amvs_per_prof[t_i] + n_levs = n_levs_per_prof[t_i] + + a = accum_amv + b = accum_amv + n_amvs + + 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] + + 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] + + accum_amv += n_amvs + accum_prf += n_levs + + return amvs, profs, times[time_idxs] + +