def calc_analysis_gmt(analysis_var_obj, trange=(1880, 2000), center_trange=(1900, 2000), detrend=False): gm_analysis = utils2.global_mean2(analysis_var_obj.data, analysis_var_obj.lat) gm_analysis = center_to_time_range(analysis_var_obj.time, gm_analysis, 0, trange=center_trange) analysis_tidx = ((analysis_var_obj.time >= trange[0]) & (analysis_var_obj.time <= trange[1])) gm_analysis = gm_analysis[analysis_tidx] if detrend: linfit_line, coef = detrend_data( analysis_var_obj.time[analysis_tidx][:, None], gm_analysis[:, None], ret_coef=True) gm_analysis -= linfit_line.squeeze() return gm_analysis
def test_global_mean2(ncf_data): dat = ncf_data['air'][0:4] lat = ncf_data['lat'][:] lon = ncf_data['lon'][:] longrid, latgrid = np.meshgrid(lon, lat) gm_time, _, _ = Utils.global_hemispheric_means(dat, lat) gm0, _, _ = Utils.global_hemispheric_means(dat[0], lat) # with time gm_test = Utils.global_mean2(dat, lat) np.testing.assert_allclose(gm_test, gm_time) # flattened lat w/ time flat_dat = dat.reshape(4, 94*192) gm_test = Utils.global_mean2(flat_dat, latgrid.flatten()) np.testing.assert_allclose(gm_test, gm_time) # no time gm_test = Utils.global_mean2(dat[0], lat) np.testing.assert_allclose(gm_test, gm0) # no time flattened spatial gm_test = Utils.global_mean2(dat[0].flatten(), latgrid.flatten()) np.testing.assert_allclose(gm_test, gm0) # NaN values dat[:, 0, :] = np.nan gm_nan_time, _, _ = Utils.global_hemispheric_means(dat, lat) gm_nan_test = Utils.global_mean2(dat, lat) np.testing.assert_allclose(gm_nan_test, gm_nan_time) # Test hemispheric gm_time, nhm_time, shm_time = Utils.global_hemispheric_means(dat, lat) gm_test, nhm_test, shm_test = Utils.global_mean2(dat, lat, output_hemispheric=True) np.testing.assert_allclose(gm_test, gm_time) np.testing.assert_allclose(nhm_test, nhm_time) np.testing.assert_allclose(shm_test, shm_time)
def compile_ens_var(parent_dir, out_dir, out_fname, a_d_vals=None, r_iters=None, ignore_npz=False): if exists(join(out_dir, out_fname)) and not ignore_npz: print('Loading pre-compiled ensemble variance metrics.') return np.load(join(out_dir, out_fname)) # Get reconstruction iteration directory if r_iters is not None: print('Joining on r iters ', r_iters) parent_iters = [join(parent_dir, 'r{:d}'.format(r)) for r in r_iters] else: parent_iters = glob.glob(join(parent_dir, 'r*')) ens_var = None gm_ens_var = None pri_ens_var = None pri_gm_ens_var = None lats = None lons = None times = None for i, parent in enumerate(parent_iters): print('Compiling iteration {:d}/{:d}'.format(i + 1, len(parent_iters))) # Directories for each parameter value if a_d_vals is not None: ad_dir = 'a{:1.2g}_d{:1.2f}' ad_dir2 = 'a{:1.1f}_d{:1.2f}' param_iters = [] for a, d in a_d_vals: curr_ad = ad_dir.format(a, d) if a == 1.0 or a == 0.0: if not exists(join(parent, curr_ad)): curr_ad = ad_dir2.format(a, d) param_iters.append(join(parent, curr_ad)) else: param_iters = [parent] for j, f in enumerate(param_iters): try: # Load analysis ensemble variance analy_var = np.load( join(f, 'ensemble_variance_tas_sfc_Amon.npz')) if times is None: times = analy_var['years'] lats = analy_var['lat'] lons = analy_var['lon'] var_shape = [len(parent_iters), len(param_iters)] + \ list(analy_var['xav'].shape) ens_var = np.zeros(var_shape) * np.nan gm_ens_var = np.zeros(var_shape[:3]) * np.nan pri_ens_var = np.zeros_like(ens_var) * np.nan pri_gm_ens_var = np.zeros_like(gm_ens_var) * np.nan ens_var[i, j] = analy_var['xav'] gm_ens_var[i, j] = utils2.global_mean2(ens_var[i, j], lats) prior_var = np.load(join(f, 'prior_ensvar_tas_sfc_Amon.npz')) pri_ens_var[i, j] = prior_var['xbv'] pri_gm_ens_var[i, j] = utils2.global_mean2( pri_ens_var[i, j], lats) except IOError as e: print(e) ens_var = ens_var.mean(axis=0).astype(np.float32) gm_ens_var = gm_ens_var.mean(axis=0) pri_ens_var = pri_ens_var.mean(axis=0).astype(np.float32) pri_gm_ens_var = pri_gm_ens_var.mean(axis=0) res_dict = { 'times': times, 'lats': lats, 'lons': lons, 'ens_var': ens_var.squeeze(), 'gm_ens_var': gm_ens_var.squeeze(), 'pri_ens_var': pri_ens_var.squeeze(), 'pri_gm_ens_var': pri_gm_ens_var.squeeze() } if not exists(out_dir): os.makedirs(out_dir) np.savez(join(out_dir, out_fname), **res_dict) return res_dict