diff --git a/modules/aeolus/aeolus_amv.py b/modules/aeolus/aeolus_amv.py index fa008ae6116b0de108a0e649e586a6f631d1e4e1..1e75a0f5fde989e7016d5e8be93e8d49f995babf 100644 --- a/modules/aeolus/aeolus_amv.py +++ b/modules/aeolus/aeolus_amv.py @@ -484,11 +484,12 @@ def run_best_fit_all(): 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) - analyze2(amvs_list, bf_list, raob_match_list, bf_gfs_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) def get_product_at_locs(raob_to_amv_dct, ts, files, filepath=None): @@ -566,6 +567,73 @@ def get_product_at_locs1x1(raob_to_amv_dct, ts, files, filepath=None): return m_dct +def run_best_fit2(raob_to_amv_dct, raob_dct, gfs_filename=None): + bfs_dct ={} + keys = list(raob_to_amv_dct.keys()) + + do_gfs_best_fit = False + gfs_press = None + gfs_spd = None + gfs_dir = None + + if gfs_filename is not None: + locs = np.array(keys) + do_gfs_best_fit = True + xr_dataset = xr.open_dataset(gfs_filename) + gfs_press = xr_dataset['pressure levels'] + gfs_press = gfs_press.values + gfs_press = gfs_press[::-1] + + uv_wind = get_vert_profile_s(xr_dataset, ['u-wind', 'v-wind'], locs[:, 1], locs[:, 0], method='nearest') + uv_wind = uv_wind.values + wspd, wdir = spd_dir_from_uv(uv_wind[0, :, :], uv_wind[1, :, :]) + wspd = wspd.magnitude + wdir = wdir.magnitude + gfs_spd = wspd[:, ::-1] + gfs_dir = wdir[:, ::-1] + + for key_idx, key in enumerate(keys): + bf_list = [] + raob_match_list = [] + bf_gfs_list = [] + + raob = raob_dct.get(key) + raob_prs = raob[:, 0] + raob_spd = raob[:, 2] + raob_dir = raob[:, 3] + amvs = raob_to_amv_dct.get(key) + num_amvs = amvs.shape[1] + for i in range(num_amvs): + amv_lon = amvs[0, i] + amv_lat = amvs[1, i] + amv_prs = amvs[4, i] + amv_spd = amvs[5, i] + amv_dir = amvs[6, i] + + bf = best_fit(amv_spd, amv_dir, amv_prs, amv_lat, amv_lon, raob_spd, raob_dir, raob_prs) + bf_list.append(bf) + + pdiff = amv_prs - raob_prs + lev_idx = np.argmin(np.abs(pdiff)) + if np.abs(pdiff[lev_idx]) > 100.0: + tup = (raob_spd[lev_idx], raob_dir[lev_idx], raob_prs[lev_idx], -9) + else: + tup = (raob_spd[lev_idx], raob_dir[lev_idx], raob_prs[lev_idx], 0) + raob_match_list.append(tup) + + if do_gfs_best_fit: + bf = best_fit(amv_spd, amv_dir, amv_prs, amv_lat, amv_lon, gfs_spd[key_idx], gfs_dir[key_idx], gfs_press) + bf_gfs_list.append(bf) + + bf_nd = np.array(bf_list) + raob_match_nd = np.array(raob_match_list) + bf_gfs_nd = np.array(bf_gfs_list) + + bfs_dct[key] = (bf_nd, raob_match_nd, bf_gfs_nd) + + return bfs_dct + + def run_best_fit(raob_to_amv_dct, raob_dct, gfs_filename=None, amvs_list=[], bf_list=[], raob_match_list=[], bf_gfs_list=[]): keys = list(raob_to_amv_dct.keys())