From 417a049ea0c3e94b35fe845c41e7dfc70678266a Mon Sep 17 00:00:00 2001 From: Alexa Bartlett Date: Fri, 3 Oct 2025 13:16:58 -0700 Subject: [PATCH 1/2] uploading modified files to apply ZCV to cross-spectrum --- abacusnbody/hod/abacus_hod.py | 110 +++++++++++++++++----------- abacusnbody/hod/zcv/tools_cv.py | 109 ++++++++++++++++++++++++--- abacusnbody/hod/zcv/tracer_power.py | 107 +++++++++++++++++++++++++-- 3 files changed, 265 insertions(+), 61 deletions(-) diff --git a/abacusnbody/hod/abacus_hod.py b/abacusnbody/hod/abacus_hod.py index 71b07eca..2c1a9475 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,38 @@ 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 +1565,51 @@ def apply_zcv(self, mock_dict, config, load_presaved=False): ) for tr in mock_dict.keys(): # obtain the positions - tracer_pos = ( - np.vstack( + if tr=='matter': #AB + matter_pos = ( + np.vstack( (mock_dict[tr]['x'], mock_dict[tr]['y'], mock_dict[tr]['z']) - ).T - ).astype(np.float32) - del mock_dict + ).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..d55770a1 100644 --- a/abacusnbody/hod/zcv/tools_cv.py +++ b/abacusnbody/hod/zcv/tools_cv.py @@ -442,27 +442,43 @@ def loss(bias): out = minimize(loss, 1.0) 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 +486,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 +518,15 @@ 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): @@ -530,8 +562,8 @@ def get_cfg(sim_name, z_this, nmesh): cfg['z_ic'] = z_ic 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 +636,29 @@ 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 +674,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 +758,36 @@ 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 +801,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..fc513bdf 100644 --- a/abacusnbody/hod/zcv/tracer_power.py +++ b/abacusnbody/hod/zcv/tracer_power.py @@ -12,6 +12,7 @@ get_k_mu_edges, get_delta_mu2, get_W_compensated, + calc_power ) from abacusnbody.metadata import get_meta @@ -25,8 +26,9 @@ '"pip install abacusutils[all]" to install zcv dependencies.' ) 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 +44,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 +80,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 +96,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 +172,16 @@ 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 +239,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 +282,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 +302,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 +328,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 +369,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( From d9f670f1de7882197ef0fade937aabf574be728b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 3 Oct 2025 20:36:00 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- abacusnbody/hod/abacus_hod.py | 44 +++++++--- abacusnbody/hod/zcv/tools_cv.py | 131 ++++++++++++++++++---------- abacusnbody/hod/zcv/tracer_power.py | 78 +++++++++-------- 3 files changed, 159 insertions(+), 94 deletions(-) diff --git a/abacusnbody/hod/abacus_hod.py b/abacusnbody/hod/abacus_hod.py index 2c1a9475..f29c26c1 100644 --- a/abacusnbody/hod/abacus_hod.py +++ b/abacusnbody/hod/abacus_hod.py @@ -1390,7 +1390,7 @@ def compute_power( clustering['mu_binc'] = power['mu_mid'][0] return clustering -# AB: adding 'cross' flag + # 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. @@ -1520,7 +1520,7 @@ def apply_zcv(self, mock_dict, config, load_presaved=False, cross=False): # run version with rsd or without rsd for tr in mock_dict.keys(): # obtain the positions - if tr=='matter': #AB + if tr == 'matter': # AB matter_pos = ( np.vstack( (mock_dict[tr]['x'], mock_dict[tr]['y'], mock_dict[tr]['z']) @@ -1536,11 +1536,15 @@ def apply_zcv(self, mock_dict, config, load_presaved=False, cross=False): gc.collect() # get power spectra for this tracer - if cross: #AB + 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 + tracer_pos, + config['HOD_params']['want_rsd'], + config, + cross=True, + matter_pos=matter_pos, ) - else: + else: pk_rsd_tr_dict = get_tracer_power( tracer_pos, config['HOD_params']['want_rsd'], config ) @@ -1565,31 +1569,43 @@ def apply_zcv(self, mock_dict, config, load_presaved=False, cross=False): ) for tr in mock_dict.keys(): # obtain the positions - if tr=='matter': #AB + if tr == 'matter': # AB matter_pos = ( np.vstack( - (mock_dict[tr]['x'], mock_dict[tr]['y'], mock_dict[tr]['z']) + ( + 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']) + ( + 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 + 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 + 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)}' @@ -1602,14 +1618,14 @@ def apply_zcv(self, mock_dict, config, load_presaved=False, cross=False): # run the final part and save zcv_dict = run_zcv( - pk_rsd_tr_dict, - pk_rsd_m_dict, #AB + pk_rsd_tr_dict, + pk_rsd_m_dict, # AB pk_rsd_ij_dict, pk_tr_dict, pk_m_dict, pk_ij_dict, config, - cross + cross, ) return zcv_dict diff --git a/abacusnbody/hod/zcv/tools_cv.py b/abacusnbody/hod/zcv/tools_cv.py index d55770a1..3262d81e 100644 --- a/abacusnbody/hod/zcv/tools_cv.py +++ b/abacusnbody/hod/zcv/tools_cv.py @@ -442,12 +442,15 @@ def loss(bias): out = minimize(loss, 1.0) return out -#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): + +# 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 + if not cross: # AB pk_tm, pk_mm, pk_ij_zm = None, None, None k = power_tr_dict['k_binc'].flatten() @@ -456,27 +459,27 @@ def read_power_dict(power_tr_dict, power_m_dict, power_ij_dict, want_rsd, keynam 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))) + 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 + 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 + 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 + 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() @@ -486,7 +489,7 @@ def read_power_dict(power_tr_dict, power_m_dict, power_ij_dict, want_rsd, keynam if want_rsd: pk_ij_zt[i, :, :] = power_tr_dict[f'P_ell_{keynames[i]}_tr'].reshape( len(poles), len(k) - ) #AB + ) # AB if cross: pk_ij_zm[i, :, :] = power_tr_dict[f'P_ell_{keynames[i]}_m'].reshape( len(poles), len(k) @@ -519,14 +522,10 @@ def read_power_dict(power_tr_dict, power_m_dict, power_ij_dict, want_rsd, keynam np.sum(pk_ij_zt == 0.0), ) - if cross: #AB - print( - np.sum(pk_tm == 0.0), - np.sum(pk_mm == 0.0), - np.sum(pk_ij_zm == 0.0) - ) + 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 + 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): @@ -562,8 +561,18 @@ def get_cfg(sim_name, z_this, nmesh): cfg['z_ic'] = z_ic return cfg -#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): + +# 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. """ @@ -636,10 +645,20 @@ def run_zcv(power_rsd_tr_dict, power_rsd_m_dict, power_rsd_ij_dict, power_tr_dic power_tr_dict, power_ij_dict = power_rsd_tr_dict, power_rsd_ij_dict # load real-space version (used for bias fitting) - #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 + # 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( + ( + 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, @@ -647,18 +666,28 @@ def run_zcv(power_rsd_tr_dict, power_rsd_m_dict, power_rsd_ij_dict, power_tr_dic keynames=keynames, poles=poles, # cross=cross - cross=False - ) #power_tr_dict also passed as input to this function (run_zcv) + 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_tm_poles, pk_mm_poles, pk_ij_zz_poles, pk_ij_zt_poles, pk_ij_zm_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 + cross=cross, ) assert np.isclose(k, k_binc).all() @@ -674,18 +703,18 @@ def run_zcv(power_rsd_tr_dict, power_rsd_m_dict, power_rsd_ij_dict, power_tr_dic # 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_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 + 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_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 + pk_ij_zm_input = pk_ij_zm_poles[:, :, 0] # AB # load the presaved window function data = np.load(window_fn) @@ -758,22 +787,27 @@ def run_zcv(power_rsd_tr_dict, power_rsd_m_dict, power_rsd_ij_dict, power_tr_dic # 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 + # 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_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 + 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_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 @@ -781,12 +815,17 @@ def run_zcv(power_rsd_tr_dict, power_rsd_m_dict, power_rsd_ij_dict, power_tr_dic 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) + 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) + 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 = {} @@ -801,7 +840,7 @@ def run_zcv(power_rsd_tr_dict, power_rsd_m_dict, power_rsd_ij_dict, power_tr_dic 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 + 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 diff --git a/abacusnbody/hod/zcv/tracer_power.py b/abacusnbody/hod/zcv/tracer_power.py index fc513bdf..1f288dda 100644 --- a/abacusnbody/hod/zcv/tracer_power.py +++ b/abacusnbody/hod/zcv/tracer_power.py @@ -12,7 +12,6 @@ get_k_mu_edges, get_delta_mu2, get_W_compensated, - calc_power ) from abacusnbody.metadata import get_meta @@ -26,9 +25,18 @@ '"pip install abacusutils[all]" to install zcv dependencies.' ) from e -#AB: The name of this function is now a bit misleading because it is also computing matter power and cross-power, + +# 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): +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. @@ -96,7 +104,7 @@ 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 + if cross: # AB pk_m_dict = {} pk_m_dict['k_binc'] = k_binc pk_m_dict['mu_binc'] = mu_binc @@ -174,10 +182,12 @@ def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power ) # del tracer_pos - if cross: #AB + 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) + 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 ) @@ -239,24 +249,24 @@ 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 + 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), + 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), + 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'] @@ -284,7 +294,7 @@ def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power if cross: pk3d_m = np.array( - (field_fft_i * np.conj(m_field_fft)).real, dtype=np.float32 + (field_fft_i * np.conj(m_field_fft)).real, dtype=np.float32 ) pk3d_m *= field_D[i] @@ -303,13 +313,13 @@ 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 + 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' + 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) @@ -329,15 +339,15 @@ def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power poles=np.asarray(poles), ) - if cross: #AB + 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), + 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] @@ -347,7 +357,7 @@ def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power 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 + if cross: # AB P_matter['power'] *= field_D[i] P_matter['binned_poles'] *= field_D[i] @@ -359,7 +369,7 @@ def get_tracer_power(tracer_pos, want_rsd, config, want_save=True, save_3D_power del field_fft_i gc.collect() - if save_3D_power: #AB: The below line should probably be updated + if save_3D_power: # AB: The below line should probably be updated return power_tr_fns header = {}