diff --git a/abacusnbody/hod/abacus_hod.py b/abacusnbody/hod/abacus_hod.py index 71b07eca..f29c26c1 100644 --- a/abacusnbody/hod/abacus_hod.py +++ b/abacusnbody/hod/abacus_hod.py @@ -258,18 +258,10 @@ def staging(self): simname = Path(self.sim_name) sim_dir = Path(self.sim_dir) mock_dir = output_dir / simname / ('z%4.3f' % self.z_mock) + # create mock_dir if not created + mock_dir.mkdir(parents=True, exist_ok=True) subsample_dir = Path(self.subsample_dir) / simname / ('z%4.3f' % self.z_mock) - # Check if the simulation directory exists - if not (sim_dir / simname).exists(): - raise FileNotFoundError( - f'Simulation directory {sim_dir / simname} not found.' - ) - - # Check if the subsample directory exists - if not subsample_dir.exists(): - raise FileNotFoundError(f'Subsample directory {subsample_dir} not found.') - # load header to read parameters if self.halo_lc: halo_info_fns = [ @@ -1398,7 +1390,8 @@ def compute_power( clustering['mu_binc'] = power['mu_mid'][0] return clustering - def apply_zcv(self, mock_dict, config, load_presaved=False): + # AB: adding 'cross' flag + def apply_zcv(self, mock_dict, config, load_presaved=False, cross=False): """ Apply control variates reduction of the variance to a power spectrum observable. """ @@ -1410,9 +1403,9 @@ def apply_zcv(self, mock_dict, config, load_presaved=False): # compute real space and redshift space # assert config['HOD_params']['want_rsd'], "Currently want_rsd=False not implemented" - assert len(mock_dict.keys()) == 1, ( - 'Currently implemented only a single tracer' - ) # should make a dict of dicts, but need cross + # assert len(mock_dict.keys()) == 1, ( + # 'Currently implemented only a single tracer' + # ) # should make a dict of dicts, but need cross assert len(config['power_params']['poles']) <= 3, ( 'Currently implemented only multipoles 0, 2, 4; need to change ZeNBu' ) @@ -1527,25 +1520,42 @@ def apply_zcv(self, mock_dict, config, load_presaved=False): # run version with rsd or without rsd for tr in mock_dict.keys(): # obtain the positions - tracer_pos = ( - np.vstack( - (mock_dict[tr]['x'], mock_dict[tr]['y'], mock_dict[tr]['z']) - ).T - ).astype(np.float32) - del mock_dict + if tr == 'matter': # AB + matter_pos = ( + np.vstack( + (mock_dict[tr]['x'], mock_dict[tr]['y'], mock_dict[tr]['z']) + ).T + ).astype(np.float32) + else: + tracer_pos = ( + np.vstack( + (mock_dict[tr]['x'], mock_dict[tr]['y'], mock_dict[tr]['z']) + ).T + ).astype(np.float32) + # del mock_dict gc.collect() # get power spectra for this tracer + if cross: # AB + pk_rsd_tr_dict, pk_rsd_m_dict = get_tracer_power( + tracer_pos, + config['HOD_params']['want_rsd'], + config, + cross=True, + matter_pos=matter_pos, + ) + else: pk_rsd_tr_dict = get_tracer_power( tracer_pos, config['HOD_params']['want_rsd'], config ) - pk_rsd_ij_dict = asdf.open(power_rsd_ij_fn)['data'] - assert np.allclose(k_binc, pk_rsd_ij_dict['k_binc']), ( - f'Mismatching file: {str(power_rsd_ij_fn)}' - ) - assert np.allclose(mu_binc, pk_rsd_ij_dict['mu_binc']), ( - f'Mismatching file: {str(power_rsd_ij_fn)}' - ) + + pk_rsd_ij_dict = asdf.open(power_rsd_ij_fn)['data'] + assert np.allclose(k_binc, pk_rsd_ij_dict['k_binc']), ( + f'Mismatching file: {str(power_rsd_ij_fn)}' + ) + assert np.allclose(mu_binc, pk_rsd_ij_dict['mu_binc']), ( + f'Mismatching file: {str(power_rsd_ij_fn)}' + ) # run version without rsd if rsd was requested if config['HOD_params']['want_rsd']: mock_dict = self.run_hod( @@ -1559,31 +1569,63 @@ def apply_zcv(self, mock_dict, config, load_presaved=False): ) for tr in mock_dict.keys(): # obtain the positions - tracer_pos = ( - np.vstack( - (mock_dict[tr]['x'], mock_dict[tr]['y'], mock_dict[tr]['z']) - ).T - ).astype(np.float32) - del mock_dict + if tr == 'matter': # AB + matter_pos = ( + np.vstack( + ( + mock_dict[tr]['x'], + mock_dict[tr]['y'], + mock_dict[tr]['z'], + ) + ).T + ).astype(np.float32) + else: + tracer_pos = ( + np.vstack( + ( + mock_dict[tr]['x'], + mock_dict[tr]['y'], + mock_dict[tr]['z'], + ) + ).T + ).astype(np.float32) + # del mock_dict gc.collect() - # get power spectra for this tracer + # get power spectra for this tracer + if cross: # AB + pk_tr_dict, pk_m_dict = get_tracer_power( + tracer_pos, + want_rsd=False, + config=config, + cross=True, + matter_pos=matter_pos, + ) + else: pk_tr_dict = get_tracer_power( tracer_pos, want_rsd=False, config=config ) - pk_ij_dict = asdf.open(power_ij_fn)['data'] - assert np.allclose(k_binc, pk_ij_dict['k_binc']), ( - f'Mismatching file: {str(power_ij_fn)}' - ) - assert np.allclose(mu_binc, pk_ij_dict['mu_binc']), ( - f'Mismatching file: {str(power_ij_fn)}' - ) + + pk_ij_dict = asdf.open(power_ij_fn)['data'] + assert np.allclose(k_binc, pk_ij_dict['k_binc']), ( + f'Mismatching file: {str(power_ij_fn)}' + ) + assert np.allclose(mu_binc, pk_ij_dict['mu_binc']), ( + f'Mismatching file: {str(power_ij_fn)}' + ) else: - pk_tr_dict, pk_ij_dict = None, None + pk_tr_dict, pk_m_dict, pk_ij_dict = None, None, None # run the final part and save zcv_dict = run_zcv( - pk_rsd_tr_dict, pk_rsd_ij_dict, pk_tr_dict, pk_ij_dict, config + pk_rsd_tr_dict, + pk_rsd_m_dict, # AB + pk_rsd_ij_dict, + pk_tr_dict, + pk_m_dict, + pk_ij_dict, + config, + cross, ) return zcv_dict diff --git a/abacusnbody/hod/zcv/tools_cv.py b/abacusnbody/hod/zcv/tools_cv.py index c904fd05..3262d81e 100644 --- a/abacusnbody/hod/zcv/tools_cv.py +++ b/abacusnbody/hod/zcv/tools_cv.py @@ -443,26 +443,45 @@ def loss(bias): return out -def read_power_dict(power_tr_dict, power_ij_dict, want_rsd, keynames, poles): +# AB: Adding arguments "power_m_dict" and "cross" +def read_power_dict( + power_tr_dict, power_m_dict, power_ij_dict, want_rsd, keynames, poles, cross +): """ ZCV: Function for reading the power spectra and saving them in the same format as Zenbu. """ + if not cross: # AB + pk_tm, pk_mm, pk_ij_zm = None, None, None + k = power_tr_dict['k_binc'].flatten() mu = np.zeros((len(k), 1)) if want_rsd: pk_tt = np.zeros((1, len(poles), len(k))) pk_ij_zz = np.zeros((15, len(poles), len(k))) pk_ij_zt = np.zeros((5, len(poles), len(k))) + if cross: # AB + pk_tm = np.zeros((1, len(poles), len(k))) + pk_mm = np.zeros((1, len(poles), len(k))) + pk_ij_zm = np.zeros((5, len(poles), len(k))) else: pk_tt = np.zeros((1, len(k), 1)) pk_ij_zz = np.zeros((15, len(k), 1)) pk_ij_zt = np.zeros((5, len(k), 1)) - + if cross: # AB + pk_tm = np.zeros((1, len(k), 1)) + pk_mm = np.zeros((1, len(k), 1)) + pk_ij_zm = np.zeros((5, len(k), 1)) if want_rsd: pk_tt[0, :, :] = power_tr_dict['P_ell_tr_tr'].reshape(len(poles), len(k)) + if cross: # AB + pk_mm[0, :, :] = power_m_dict['P_ell_m_m'].reshape(len(poles), len(k)) + pk_tm[0, :, :] = power_tr_dict['P_ell_tr_m'].reshape(len(poles), len(k)) nmodes = power_tr_dict['N_ell_tr_tr'].flatten() else: pk_tt[0, :, :] = power_tr_dict['P_kmu_tr_tr'].reshape(len(k), 1) + if cross: # AB + pk_mm[0, :, :] = power_m_dict['P_kmu_m_m'].reshape(len(k), 1) + pk_tm[0, :, :] = power_tr_dict['P_kmu_tr_m'].reshape(len(k), 1) nmodes = power_tr_dict['N_kmu_tr_tr'].flatten() count = 0 @@ -470,11 +489,19 @@ def read_power_dict(power_tr_dict, power_ij_dict, want_rsd, keynames, poles): if want_rsd: pk_ij_zt[i, :, :] = power_tr_dict[f'P_ell_{keynames[i]}_tr'].reshape( len(poles), len(k) - ) + ) # AB + if cross: + pk_ij_zm[i, :, :] = power_tr_dict[f'P_ell_{keynames[i]}_m'].reshape( + len(poles), len(k) + ) else: pk_ij_zt[i, :, :] = power_tr_dict[f'P_kmu_{keynames[i]}_tr'].reshape( len(k), 1 ) + if cross: + pk_ij_zm[i, :, :] = power_m_dict[f'P_kmu_{keynames[i]}_m'].reshape( + len(k), 1 + ) for j in range(len(keynames)): if i < j: continue @@ -494,7 +521,11 @@ def read_power_dict(power_tr_dict, power_ij_dict, want_rsd, keynames, poles): np.sum(pk_ij_zz == 0.0), np.sum(pk_ij_zt == 0.0), ) - return k, mu, pk_tt, pk_ij_zz, pk_ij_zt, nmodes + + if cross: # AB + print(np.sum(pk_tm == 0.0), np.sum(pk_mm == 0.0), np.sum(pk_ij_zm == 0.0)) + + return k, mu, pk_tt, pk_tm, pk_mm, pk_ij_zz, pk_ij_zt, pk_ij_zm, nmodes # AB def get_cfg(sim_name, z_this, nmesh): @@ -531,7 +562,17 @@ def get_cfg(sim_name, z_this, nmesh): return cfg -def run_zcv(power_rsd_tr_dict, power_rsd_ij_dict, power_tr_dict, power_ij_dict, config): +# AB: Adding arguments "power_rsd_m_dict", "power_m_dict", and "cross", which will need to be passed to apply_zcv in abacus_hod.py +def run_zcv( + power_rsd_tr_dict, + power_rsd_m_dict, + power_rsd_ij_dict, + power_tr_dict, + power_m_dict, + power_ij_dict, + config, + cross, +): """ Apply Zel'dovich control variates (ZCV) reduction to some measured power spectrum. """ @@ -604,17 +645,49 @@ def run_zcv(power_rsd_tr_dict, power_rsd_ij_dict, power_tr_dict, power_ij_dict, power_tr_dict, power_ij_dict = power_rsd_tr_dict, power_rsd_ij_dict # load real-space version (used for bias fitting) - k, mu, pk_tt_real, pk_ij_zz_real, pk_ij_zt_real, nmodes = read_power_dict( - power_tr_dict, power_ij_dict, want_rsd=False, keynames=keynames, poles=poles - ) + # AB: changes made below entailed modifying read_power_dict function (above) + # pk_tm_real, pk_mm_real, pk_ij_zt_real, and pk_ij_zm_real not currently used, but the new read_power_dict + # returns 9 quantities + ( + k, + mu, + pk_tt_real, + pk_tm_real, + pk_mm_real, + pk_ij_zz_real, + pk_ij_zt_real, + pk_ij_zm_real, + nmodes, + ) = read_power_dict( + power_tr_dict, + power_m_dict, + power_ij_dict, + want_rsd=False, + keynames=keynames, + poles=poles, + # cross=cross + cross=False, + ) # power_tr_dict also passed as input to this function (run_zcv) # load either real or redshift - k, mu, pk_tt_poles, pk_ij_zz_poles, pk_ij_zt_poles, nmodes = read_power_dict( + ( + k, + mu, + pk_tt_poles, + pk_tm_poles, + pk_mm_poles, + pk_ij_zz_poles, + pk_ij_zt_poles, + pk_ij_zm_poles, + nmodes, + ) = read_power_dict( power_rsd_tr_dict, + power_rsd_m_dict, power_rsd_ij_dict, want_rsd=want_rsd, keynames=keynames, poles=poles, + cross=cross, ) assert np.isclose(k, k_binc).all() @@ -630,12 +703,18 @@ def run_zcv(power_rsd_tr_dict, power_rsd_ij_dict, power_tr_dict, power_ij_dict, # decide what to input depending on whether rsd requested or not if want_rsd: pk_tt_input = pk_tt_poles[0, ...] + pk_tm_input = pk_tm_poles[0, ...] # AB + pk_mm_input = pk_mm_poles[0, ...] # AB pk_ij_zz_input = pk_ij_zz_poles pk_ij_zt_input = pk_ij_zt_poles + pk_ij_zm_input = pk_ij_zm_poles # AB else: pk_tt_input = pk_tt_poles[0, :, 0] + pk_tm_input = pk_tm_poles[0, :, 0] # AB + pk_mm_input = pk_mm_poles[0, :, 0] # AB pk_ij_zz_input = pk_ij_zz_poles[:, :, 0] pk_ij_zt_input = pk_ij_zt_poles[:, :, 0] + pk_ij_zm_input = pk_ij_zm_poles[:, :, 0] # AB # load the presaved window function data = np.load(window_fn) @@ -708,6 +787,46 @@ def run_zcv(power_rsd_tr_dict, power_rsd_ij_dict, power_tr_dict, power_ij_dict, # beta needs to be smooth for best results pk_nn_betasmooth = pk_tt_input - beta_smooth * (pk_zz - pk_zenbu) + # AB: doing everything done above but for cross + if cross: + idx_cross = [0, 1, 3, 6, 10] + + pk_zz_cross = combine_cross_spectra( + k_binc, pk_ij_zz_input[idx_cross, :], bias_vec[1:], rsd=want_rsd + ) + pk_zenbu_cross = combine_cross_spectra( + k_binc, pk_ij_zenbu[idx_cross, :], bias_vec[1:], rsd=want_rsd + ) + pk_nz_gm = pk_ij_zt_input[0, :] + pk_zn_gm = combine_cross_spectra( + k_binc, pk_ij_zm_input, bias_vec[1:], rsd=want_rsd + ) + pk_nz_mm = pk_ij_zm_input[0, :] + + cov_zn_cross = pk_zn * pk_nz_mm + pk_nz_gm * pk_zn_gm + var_z_cross = pk_zz * pk_ij_zz_input[0] + pk_zz_cross**2 + var_n_cross = pk_tt_input * pk_mm_input + pk_tm_input**2 + + beta_cross = cov_zn_cross / np.sqrt(var_z_cross * var_n_cross) + beta_damp_cross = 0.5 * (1 - np.tanh((k_binc - k0) / dk_cv)) * beta_cross + beta_damp_cross = np.atleast_2d(beta_damp_cross) + beta_damp_cross[beta_damp_cross != beta_damp_cross] = 0 + beta_damp_cross[:, : k_binc.searchsorted(beta1_k)] = 1 + beta_smooth_cross = np.zeros_like(beta_damp_cross) + for i in range(beta_smooth.shape[0]): + try: + beta_smooth_cross[i, :] = savgol_filter( + beta_damp_cross.T[:, i], sg_window, 3 + ) + except ValueError: + warnings.warn( + 'This message should only appear when doing a smoke test.' + ) + + pk_nn_betasmooth_cross = pk_tm_input - beta_smooth_cross * ( + pk_zz_cross - pk_zenbu_cross + ) + # save results to a dictionary zcv_dict = {} zcv_dict['k_binc'] = k_binc @@ -721,6 +840,13 @@ def run_zcv(power_rsd_tr_dict, power_rsd_ij_dict, power_tr_dict, power_ij_dict, zcv_dict['Pk_tr_tr_ell_zcv'] = pk_nn_betasmooth zcv_dict['Pk_ZD_ZD_ell_ZeNBu'] = pk_zenbu zcv_dict['bias'] = bias_vec[1:] + if cross: # AB + zcv_dict['Pk_ZD_ZD_m'] = pk_ij_zz_input[0] + zcv_dict['Pk_m_m_ell'] = pk_mm_input + zcv_dict['Pk_tr_m_ell'] = pk_tm_input + zcv_dict['Pk_tr_m_ell_zcv'] = pk_nn_betasmooth_cross + zcv_dict['Pk_ZD_ZD_cross_ell'] = pk_zz_cross + zcv_dict['Pk_ZD_ZD_cross_ell_ZeNBu'] = pk_zenbu_cross return zcv_dict diff --git a/abacusnbody/hod/zcv/tracer_power.py b/abacusnbody/hod/zcv/tracer_power.py index 6e5f15f7..1f288dda 100644 --- a/abacusnbody/hod/zcv/tracer_power.py +++ b/abacusnbody/hod/zcv/tracer_power.py @@ -26,7 +26,17 @@ ) from e -def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power=False): +# AB: The name of this function is now a bit misleading because it is also computing matter power and cross-power, +# but I wanted to keep everything within this function for ease of implementing +def get_tracer_power( + tracer_pos, + want_rsd, + config, + want_save=True, + save_3D_power=False, + cross=False, + matter_pos=None, +): """ Compute the auto- and cross-correlation between a galaxy catalog (`tracer_pos`) and the advected Zel'dovich fields. @@ -42,12 +52,16 @@ def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power save_3D_power : bool, optional save the 3D power spectra in individual ASDF files. Default is False. + cross: bool + flag indicating whether to compute cross-spectra + matter_pos: array_like + matter positions with shape (N, 3) Returns ------- - pk_tr_dict : dict - dictionary containing the auto- and cross-power spectra of - the tracer with the 5 fields. + pk_tr_dict, pk_m_dict : dicts + dictionaries containing the auto- and cross-power spectra of + the tracer and matter, respectively, with the 5 fields. """ # read zcv parameters advected_dir = config['zcv_params']['zcv_dir'] # input of advected fields @@ -74,6 +88,7 @@ def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power # get a few parameters for the simulation meta = get_meta(sim_name, redshift=z_this) Lbox = meta['BoxSize'] + print(Lbox) z_ic = meta['InitialRedshift'] # k_Ny = np.pi*nmesh/Lbox @@ -89,6 +104,11 @@ def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power pk_tr_dict['k_binc'] = k_binc pk_tr_dict['mu_binc'] = mu_binc + if cross: # AB + pk_m_dict = {} + pk_m_dict['k_binc'] = k_binc + pk_m_dict['mu_binc'] = mu_binc + # create save directory save_dir = Path(tracer_dir) / sim_name save_z_dir = save_dir / f'z{z_this:.3f}' @@ -160,7 +180,18 @@ def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power tr_field_fft = get_field_fft( tracer_pos, Lbox, nmesh, paste, w, W, compensated, interlaced ) - del tracer_pos + # del tracer_pos + + if cross: # AB + matter_pos += Lbox / 2.0 # I think necessary for cross correlations + matter_pos %= Lbox + print( + 'min/max matter pos', matter_pos.min(), matter_pos.max(), matter_pos.shape + ) + m_field_fft = get_field_fft( + matter_pos, Lbox, nmesh, paste, w, W, compensated, interlaced + ) + gc.collect() if want_save: @@ -218,6 +249,31 @@ def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power pk_tr_dict['P_ell_tr_tr'] = P['binned_poles'] pk_tr_dict['N_ell_tr_tr'] = P['N_mode_poles'] + if cross: # AB + # AB: try using calc_power instead + P_matter = calc_pk_from_deltak( + m_field_fft, + Lbox, + k_bin_edges, + mu_bin_edges, + field2_fft=None, + poles=np.asarray(poles), + ) + + P_cross = calc_pk_from_deltak( + tr_field_fft, + Lbox, + k_bin_edges, + mu_bin_edges, + field2_fft=m_field_fft, + poles=np.asarray(poles), + ) + + pk_tr_dict['P_kmu_tr_m'] = P_cross['power'] + # pk_tr_dict['P_ell_tr_m'] = P_cross['poles'] + pk_m_dict['P_kmu_m_m'] = P_matter['power'] + # pk_m_dict['P_ell_m_m'] = P_matter['poles'] + # loop over fields for i in range(len(keynames)): print('Computing cross-correlation of tracer and ', keynames[i]) @@ -236,6 +292,12 @@ def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power pk3d *= field_D[i] # pk3d *= Lbox**3 + if cross: + pk3d_m = np.array( + (field_fft_i * np.conj(m_field_fft)).real, dtype=np.float32 + ) + pk3d_m *= field_D[i] + # record pk_tr_dict = {} pk_tr_dict[f'P_k3D_{keynames[i]}_tr'] = pk3d @@ -250,6 +312,18 @@ def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power ) power_tr_fns.append(power_tr_fn) compress_asdf(str(power_tr_fn), pk_tr_dict, header) + + if cross: # AB + pk_m_dict = {} + pk_m_dict[f'P_k3D_{keynames[i]}_m'] = pk3d_m + + power_m_fn = ( + Path(save_z_dir) + / f'power{rsd_str}_{keynames[i]}_m_nmesh{nmesh:d}.asdf' + ) + + compress_asdf(str(power_m_fn), pk_m_dict, header) + del field_fft_i gc.collect() @@ -264,16 +338,38 @@ def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power field2_fft=tr_field_fft, poles=np.asarray(poles), ) + + if cross: # AB + P_matter = calc_pk_from_deltak( + field_fft_i[f'{keynames[i]}_Re'] + + 1j * field_fft_i[f'{keynames[i]}_Im'], + Lbox, + k_bin_edges, + mu_bin_edges, + field2_fft=m_field_fft, + poles=np.asarray(poles), + ) + P['power'] *= field_D[i] P['binned_poles'] *= field_D[i] pk_tr_dict[f'P_kmu_{keynames[i]}_tr'] = P['power'] pk_tr_dict[f'N_kmu_{keynames[i]}_tr'] = P['N_mode'] pk_tr_dict[f'P_ell_{keynames[i]}_tr'] = P['binned_poles'] pk_tr_dict[f'N_ell_{keynames[i]}_tr'] = P['N_mode_poles'] + + if cross: # AB + P_matter['power'] *= field_D[i] + P_matter['binned_poles'] *= field_D[i] + + pk_m_dict[f'P_kmu_{keynames[i]}_m'] = P_matter['power'] + pk_m_dict[f'N_kmu_{keynames[i]}_m'] = P_matter['N_mode'] + pk_m_dict[f'P_ell_{keynames[i]}_m'] = P_matter['binned_poles'] + pk_m_dict[f'N_ell_{keynames[i]}_m'] = P_matter['N_mode_poles'] + del field_fft_i gc.collect() - if save_3D_power: + if save_3D_power: # AB: The below line should probably be updated return power_tr_fns header = {} @@ -283,7 +379,12 @@ def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power header['kcut'] = kcut if want_save: compress_asdf(str(power_tr_fn), pk_tr_dict, header) - return pk_tr_dict + # if cross: #AB + # compress_asdf(str(power_m_fn), pk_m_dict, header) + if cross: + return pk_tr_dict, pk_m_dict + else: + return pk_tr_dict def get_recon_power(