diff --git a/analysis.py b/analysis.py new file mode 100644 index 0000000..2022004 --- /dev/null +++ b/analysis.py @@ -0,0 +1,130 @@ +from tqdm import tqdm +import h5py as h5 +import numpy as np +import configparser +import sys +from contextlib import contextmanager +import os +import Pk_library as PKL +import glob + +@contextmanager +def suppress_stdout(): + with open(os.devnull, 'w') as devnull: + old_stdout = sys.stdout + sys.stdout = devnull + try: + yield + finally: + sys.stdout = old_stdout + + +def get_mcmc_steps(dirname, nframe, iter_max=-1, iter_min=-1): + """ + Obtain evenly-spaced sample of MCMC steps + + Args: + dirname (str): Directory containing MCMC files + nframe (int): Number of frames to sample + iter_max (int): Maximum MCMC step to sample + iter_min (int): Minimum MCMC step to sample + + Returns: + all_mcmc (list): List of MCMC steps to sample + + """ + + # Get all MCMC steps + all_mcmc = glob.glob(dirname + '/mcmc_*.h5') + all_mcmc = np.sort([int(os.path.basename(m)[len('mcmc_'):-3]) for m in all_mcmc]) + + # Get in range + if iter_max >= 0: + all_mcmc = all_mcmc[all_mcmc <= iter_max] + all_mcmc = all_mcmc[all_mcmc >= iter_min] + + # Subsample + if nframe > 0: + max_out = max(all_mcmc) + min_out = min(all_mcmc) + step = max(int((max_out - min_out+1) / nframe), 1) + all_mcmc = all_mcmc[::step] + if max_out not in all_mcmc: + all_mcmc = np.concatenate([all_mcmc, [max_out]]) + + return all_mcmc + + +def compute_ensemble_mean_field(ini_name, dirname, mcmc_steps, which_field='BORG_final_density'): + """ + Compute the mean and std deviation of the inferred density field + """ + + print('Computing ensemble mean field') + + #COMPUTE THE MEAN-DENSITY FIELD + for i in tqdm(range(len(mcmc_steps))): + idx = mcmc_steps[i] + with h5.File(dirname + "/mcmc_%d.h5" % idx,'r') as mcmc_file: + temp_field = np.array(mcmc_file[f'scalars/{which_field}'][...],dtype=np.float64) + if i == 0: + mean_field = np.array(np.full(temp_field.shape,0),dtype=np.float64) + std_field = np.array(np.full(temp_field.shape,0),dtype=np.float64) + mean_field += temp_field + std_field += temp_field*temp_field + mean_field = mean_field/np.float64(len(mcmc_steps)) + std_field = std_field/np.float64(len(mcmc_steps)) # < delta^2 > + std_field = np.sqrt(std_field - mean_field **2) # < delta^2 > - < delta >^2 + + return mean_field, std_field + + +def get_mock_field(ini_name, dirname, which_field='BORG_final_density'): + + with h5.File(f'{dirname}/mock_data.h5', 'r') as f: + dens = f[f'scalars/{which_field}'][:] + + config = configparser.ConfigParser() + config.read(ini_name) + L = float(config['system']['L0']) + + return dens, L + + +def get_spectra(ini_name, dirname, mcmc_steps, which_field='BORG_final_density', MAS=''): + + if MAS == '': + if which_field == 'BORG_final_density': + MAS = "CIC" + elif which_field == 'ics': + MAS = None + else: + raise NotImplementedError + + # Compute original power spectrum + delta1, boxsize = get_mock_field(ini_name, dirname, which_field=which_field) + + # Turn to overdensity + if which_field == 'BORG_final_density': + delta1 = (1. + delta1) / (1. + delta1).mean() - 1. + + print("BOXSIZE", boxsize) + Pk = PKL.Pk(delta1.astype(np.float32), boxsize, axis=0, MAS=MAS, threads=1, verbose=True) + k = Pk.k3D + Pk_true = Pk.Pk[:,0] + + # Get other spectra + all_pk = np.zeros((len(mcmc_steps), len(k))) + all_r = np.zeros((len(mcmc_steps), len(k))) + for i in tqdm(range(len(mcmc_steps))): + idx = mcmc_steps[i] + + with h5.File(dirname + "/mcmc_%d.h5" % idx,'r') as mcmc_file: + delta2 = np.array(mcmc_file['scalars/{which_field}'][...],dtype=np.float64) + + with suppress_stdout(): + Pk = PKL.XPk([delta1.astype(np.float32),delta2.astype(np.float32)], boxsize, axis=0, MAS=[MAS, MAS], threads=1) + all_pk[i,:] = Pk.Pk[:,0,1] #monopole of field 2 + all_r[i,:] = Pk.XPk[:,0,0] / np.sqrt(Pk.Pk[:,0,1] * Pk.Pk[:,0,0]) + + return k, Pk_true, all_pk, all_r