diff --git a/modules/aeolus/aeolus_amv.py b/modules/aeolus/aeolus_amv.py index 1e75a0f5fde989e7016d5e8be93e8d49f995babf..6c759f85798decf2901d301c5118e7a2a3257dc6 100644 --- a/modules/aeolus/aeolus_amv.py +++ b/modules/aeolus/aeolus_amv.py @@ -481,18 +481,19 @@ def run_best_fit_all(): bf_list = [] raob_match_list = [] bf_gfs_list = [] + amv_prod_list = [] for k, file in enumerate(raob_files): raob_dct, ts = get_raob_dict_cdf(raob_dir+file, raob_dtsts[k]) m_d = match_amvs_to_raobs(raob_dct, ts, amv_files=amv_files) - # amvs_list, bf_list, raob_match_list, bf_gfs_list = run_best_fit(m_d, raob_dct, gfs_dir+gfs_files[k], - # amvs_list=amvs_list, bf_list=bf_list, raob_match_list=raob_match_list, bf_gfs_list=bf_gfs_list) + amvs_list, bf_list, raob_match_list, bf_gfs_list = run_best_fit(m_d, raob_dct, gfs_dir+gfs_files[k], + amvs_list=amvs_list, bf_list=bf_list, raob_match_list=raob_match_list, bf_gfs_list=bf_gfs_list) # analyze2(amvs_list, bf_list, raob_match_list, bf_gfs_list) - prd_dct = get_product_at_locs(m_d, ts, prd_files) - bfs_dct = run_best_fit2(m_d, raob_dct, gfs_dir+gfs_files[k]) - # analyze2(amvs_list, bf_list, raob_match_list, bf_gfs_list) + prd_dct = get_product_at_locs(m_d, ts, prd_files, amv_prod_list=amv_prod_list) + # bfs_dct = run_best_fit2(m_d, raob_dct, gfs_dir+gfs_files[k]) + analyze2(amvs_list, bf_list, raob_match_list, bf_gfs_list, amv_prod_list) -def get_product_at_locs(raob_to_amv_dct, ts, files, filepath=None): +def get_product_at_locs(raob_to_amv_dct, ts, files, amv_prod_list=[], filepath=None): keys = list(raob_to_amv_dct.keys()) m_dct = {} @@ -524,6 +525,7 @@ def get_product_at_locs(raob_to_amv_dct, ts, files, filepath=None): aaa[:, k] = param_nd[:, ll[k], cc[k]] m_dct[key] = aaa + amv_prod_list.append(aaa) ds.close() @@ -641,12 +643,11 @@ def run_best_fit(raob_to_amv_dct, raob_dct, gfs_filename=None, amvs_list=[], bf_ gfs_press = None gfs_spd = None gfs_dir = None - #bf_gfs_list = None + if gfs_filename is not None: locs = np.array(keys) do_gfs_best_fit = True xr_dataset = xr.open_dataset(gfs_filename) - #bf_gfs_list = [] gfs_press = xr_dataset['pressure levels'] gfs_press = gfs_press.values gfs_press = gfs_press[::-1] @@ -659,12 +660,8 @@ def run_best_fit(raob_to_amv_dct, raob_dct, gfs_filename=None, amvs_list=[], bf_ gfs_spd = wspd[:, ::-1] gfs_dir = wdir[:, ::-1] - #amvs_list = [] - #bf_list = [] - #raob_match_list = [] for key_idx, key in enumerate(keys): raob = raob_dct.get(key) - nlevs = raob.shape[0] raob_prs = raob[:, 0] raob_spd = raob[:, 2] raob_dir = raob[:, 3] @@ -696,14 +693,21 @@ def run_best_fit(raob_to_amv_dct, raob_dct, gfs_filename=None, amvs_list=[], bf_ return amvs_list, bf_list, raob_match_list, bf_gfs_list -def analyze2(amvs_list, bf_list, raob_match_list, bf_gfs_list): +def analyze2(amvs_list, bf_list, raob_match_list, bf_gfs_list, amv_prod_list): amvs = np.concatenate(amvs_list, axis=1) amvs = np.transpose(amvs, axes=[1, 0]) + print(amvs.shape) bfs = np.stack(bf_list, axis=0) + print(bfs.shape) raob_match = np.stack(raob_match_list, axis=0) + print(raob_match.shape) if len(bf_gfs_list) > 0: bfs_gfs = np.stack(bf_gfs_list, axis=0) + if len(amv_prod_list) > 0: + amv_prd = np.concatenate(amv_prod_list, axis=1) + amv_prd = np.transpose(amv_prd, axes=[1, 0]) + print(amv_prd.shape) good_amvs = amvs num_good = good_amvs.shape[0]