import numpy as np import glob import configparser import h5py as h5 import ast import Pk_library as PKL import os from tqdm import tqdm import sys from contextlib import contextmanager @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, iter_min=0): """ Obtain evenly-spaced sample of MCMC steps to make movie from """ all_mcmc = glob.glob(dirname + '/mcmc_*.h5') x = [m[len(dirname + '/mcmc_'):-3] for m in all_mcmc] all_mcmc = np.sort([int(m[len(dirname + '/mcmc_'):-3]) for m in all_mcmc]) if iter_max >= 0: all_mcmc = all_mcmc[all_mcmc <= iter_max] all_mcmc = all_mcmc[all_mcmc >= iter_min] 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 load_param_samples(ini_name, dirname, nframe, iter_max, iter_min): config = configparser.ConfigParser() config.read(ini_name) to_sample = [] for k,v in config['block_loop'].items(): if v.strip() == 'false': i = k.index('_sampler') if k[:i] not in ['hades', 'bias', 'nmean']: to_sample.append(k[:i]) print("TO SAMPLE", to_sample) nsamp = int(config['run']['nsamp']) new_to_sample = [] for s in to_sample: if s in ['omega_m', 'sigma8', 'sig_v']: new_to_sample.append(s) elif s == 'bulk_flow': for d in ['_x', '_y', '_z']: new_to_sample.append(f'{s}{d}') else: for i in range(nsamp): new_to_sample.append(f'{s}{i}') # This is desired list to sample to_sample = new_to_sample # Which steps to use all_mcmc = get_mcmc_steps(dirname, nframe, iter_max, iter_min=iter_min) sampler = config['sampling']['algorithm'].lower() samples = np.empty((len(to_sample),len(all_mcmc))) print('MY SAMPLER IS', sampler) if sampler == 'slice': for i in tqdm(range(len(all_mcmc))): with h5.File(f'{dirname}/mcmc_{all_mcmc[i]}.h5', 'r') as f: for j, s in enumerate(to_sample): if 'model_params_' + s in f['scalars'].keys(): samples[j,i] = f['scalars/model_params_' + s][:][0] elif 'model_params_cosmology.' + s in f['scalars'].keys(): samples[j,i] = f['scalars/model_params_cosmology.' + s][:][0] elif s == 'sig_v': samples[j,i] = float(config['model'][s]) elif s.startswith('bulk_flow'): if s[-1] == 'x': samples[j,i] = np.array(ast.literal_eval(config['model']['bulk_flow']))[0] elif s[-1] == 'y': samples[j,i] = np.array(ast.literal_eval(config['model']['bulk_flow']))[1] elif s[-1] == 'z': samples[j,i] = np.array(ast.literal_eval(config['model']['bulk_flow']))[2] else: raise NotImplementedError else: if s in config[f'cosmology'].keys(): samples[j,i] = float(config['cosmology'][s]) else: print("NOT THERE") samples[j,i] = float(config[f'sample_{s[-1]}'][s[:-1]]) elif sampler in ['hmc', 'mvslice', 'transformedblackjax', 'blackjax']: if sampler in ['hmc', 'transformedblackjax']: key_name = 'attributes' key_name = 'model_params' elif sampler in ['mvslice', 'blackjax']: key_name = 'model_paramsattributes' # Get order in which model parameters are stored if os.path.isfile(f'{dirname}/model_params.txt'): with open(f'{dirname}/model_params.txt', 'r') as file: model_params = [line.strip() for line in file] else: model_params = [] print(model_params) for i in tqdm(range(len(all_mcmc))): with h5.File(f'{dirname}/mcmc_{all_mcmc[i]}.h5', 'r') as f: if key_name in f['scalars'].keys(): data = f[f'scalars/{key_name}'][:] else: data = None for j, s in enumerate(to_sample): if s in model_params: samples[j,i] = data[model_params.index(s)] elif 'model_params_cosmology.' + s in f['scalars'].keys(): samples[j,i] = f['scalars/model_params_cosmology.' + s][:][0] elif s == 'sig_v': samples[j,i] = float(config['model'][s]) elif s in config[f'cosmology'].keys(): samples[j,i] = float(config['cosmology'][s]) elif s.startswith('bulk_flow'): idx = {'x':0, 'y':1, 'z':2} idx = idx[s[-1]] samples[j,i] = np.array(ast.literal_eval(config['model']['bulk_flow']))[idx] else: samples[j,i] = float(config[f'sample_{s[-1]}'][s[:-1]]) else: raise NotImplementedError return to_sample, all_mcmc, samples def get_truths(ini_name, to_sample): config = configparser.ConfigParser() config.read(ini_name) truths = [None] * len(to_sample) for i, s in enumerate(to_sample): if s in config[f'cosmology'].keys(): truths[i] = float(config['cosmology'][s]) elif s == 'sig_v': truths[i] = float(config['model'][s]) elif s.startswith('bulk_flow'): idx = {'x':0, 'y':1, 'z':2} idx = idx[s[-1]] truths[i] = np.array(ast.literal_eval(config['model']['bulk_flow']))[idx] else: truths[i] = float(config[f'sample_{s[-1]}'][s[:-1]]) return truths def crop_field(ini_name, field): config = configparser.ConfigParser() config.read(ini_name) Rmax = float(config['mock']['R_max']) xmin = float(config['system']['corner0']) L = float(config['system']['L0']) N = int(config['system']['N0']) x = np.linspace(xmin, xmin+L, N) m = np.abs(x) < Rmax L = x[m].max() - x[m].min() return field[m][:, m][:, :, m], L def compute_ensemble_mean_field(ini_name, dirname, nframe, iter_max, iter_min, cut_field=True): """ Compute the mean and std deviation of the inferred density field """ print('Computing ensemble mean field') # Which steps to use all_mcmc = get_mcmc_steps(dirname, nframe, iter_max, iter_min=iter_min) #COMPUTE THE MEAN-DENSITY FIELD for i in tqdm(range(len(all_mcmc))): idx = all_mcmc[i] with h5.File(dirname + "/mcmc_%d.h5" % idx,'r') as mcmc_file: temp_field = np.array(mcmc_file['scalars/BORG_final_density'][...],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(all_mcmc)) std_field = std_field/np.float64(len(all_mcmc)) # < delta^2 > std_field = np.sqrt(std_field - mean_field **2) # < delta^2 > - < delta >^2 # Cut the density field if needed if cut_field: mean_field, _ = crop_field(ini_name, mean_field) std_field, _ = crop_field(ini_name, std_field) return mean_field, std_field def get_mock_field(ini_name, dirname, which_field='delta', cut_field=True): with h5.File(f'{dirname}/mock_data.h5', 'r') as f: if which_field == 'delta': dens = f['scalars/BORG_final_density'][:] elif which_field == 'ics': dens = f['scalars/s_field'][:] if cut_field: dens, L = crop_field(ini_name, dens) else: config = configparser.ConfigParser() config.read(ini_name) L = float(config['system']['L0']) return dens, L def get_velmass_field(ini_file, field_type): """ Load the true field from the Velmass directory. Uses the ini file given in dirname field_type can be one of 'delta', 'vx', 'vy', 'vz', 'ics' """ config = configparser.ConfigParser() config.read(ini_file) dirname = config['mock']['velmass_dirname'] if field_type == 'ics': f = np.load('/data101/bartlett/fsigma8/VELMASS/velmass_density.npy') else: if field_type == 'delta': key = 'd' else: key = field_type f = np.load(dirname + '/../density.npz')[key] if field_type == 'delta': f = f / np.mean(f) - 1 return f def get_velmass_Lbox(ini_file): """ Load the true box length from the Velmass directory. Uses the ini file given in dirname """ config = configparser.ConfigParser() config.read(ini_file) velmass_dirname = config['mock']['velmass_dirname'] # Get box size with open(velmass_dirname + '/auto-rockstar.cfg') as f: data = [r for r in f] Lbox = [r for r in data if r.startswith('BOX_SIZE')][0].strip() Lbox = float(Lbox.split('=')[1]) return Lbox def get_borg_Lbox(ini_file): """ Load the box length used by BORG Uses the ini file given in dirname """ config = configparser.ConfigParser() config.read(ini_file) Lbox = float(config['system']['L0']) return Lbox def get_borg_N(ini_file): """ Load the box length used by BORG Uses the ini file given in dirname """ config = configparser.ConfigParser() config.read(ini_file) N = int(config['system']['N0']) return N def get_borg_Rmax(ini_file): """ Load the Rmax used by BORG Uses the ini file given in dirname """ config = configparser.ConfigParser() config.read(ini_file) R_max = float(config['mock']['R_max']) return R_max def get_borg_corner(ini_file): """ Load the corner used by BORG Uses the ini file given in dirname """ config = configparser.ConfigParser() config.read(ini_file) corn = float(config['system']['corner0']) return corn def crop_velmass_to_borg(ini_file, field_type): """ Load the true field from the Velmass directory. Then crop the field to only contain the region used by BORG Uses the ini file given in dirname field_type can be one of 'delta', 'vx', 'vy', 'vz', or 'ics' """ true_field = get_velmass_field(ini_file, field_type) Lbox_true = get_velmass_Lbox(ini_file) Lbox_borg = get_borg_Lbox(ini_file) N = true_field.shape[0] imin = int((1 - Lbox_borg / Lbox_true) * N / 2) imax = int((1 + Lbox_borg / Lbox_true) * N / 2) true_field = true_field[imin:imax,imin:imax,imin:imax] return true_field def get_spectra(ini_name, dirname, nframe, iter_max, iter_min, which_field='delta', cut_field=True, mock_type='borg', MAS=''): # Which steps to use all_mcmc = get_mcmc_steps(dirname, nframe, iter_max, iter_min=iter_min) if MAS == '': if which_field == 'delta': MAS = "CIC" elif which_field == 'ics': MAS = None else: raise NotImplementedError # Compute original power spectrum if mock_type == 'borg': delta1, boxsize = get_mock_field(ini_name, dirname, which_field=which_field, cut_field=cut_field) elif mock_type == 'velmass': delta1 = crop_velmass_to_borg(ini_name, which_field) if cut_field: delta1, boxsize = crop_field(ini_name, delta1) else: boxsize = get_borg_Lbox(ini_name) else: raise NotImplementedError # Turn to overdensity if which_field == 'delta': 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(all_mcmc), len(k))) all_r = np.zeros((len(all_mcmc), len(k))) for i in tqdm(range(len(all_mcmc))): idx = all_mcmc[i] with h5.File(dirname + "/mcmc_%d.h5" % idx,'r') as mcmc_file: if which_field == 'delta': delta2= np.array(mcmc_file['scalars/BORG_final_density'][...],dtype=np.float64) elif which_field == 'ics': delta2 = np.array(mcmc_file['scalars/s_field'][...],dtype=np.float64) else: raise NotImplementedError if cut_field: delta2, _ = crop_field(ini_name, delta2) 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 def get_both_fields(ini_name, dirname, step, which_field='delta', cut_field=True, mock_type='borg'): # Mock if mock_type == 'borg': delta1, boxsize = get_mock_field(ini_name, dirname, which_field=which_field, cut_field=cut_field) elif mock_type == 'velmass': delta1 = crop_velmass_to_borg(ini_name, which_field) if cut_field: delta1, boxsize = crop_field(ini_name, delta1) else: boxsize = get_borg_Lbox(ini_name) else: raise NotImplementedError # Step with h5.File(dirname + "/mcmc_%d.h5" % step,'r') as mcmc_file: if which_field == 'delta': delta2= np.array(mcmc_file['scalars/BORG_final_density'][...],dtype=np.float64) elif which_field == 'ics': delta2 = np.array(mcmc_file['scalars/s_field'][...],dtype=np.float64) else: raise NotImplementedError if cut_field: delta2, _ = crop_field(ini_name, delta2) return delta1, delta2 def get_likelihood_values(dirname, nframe, iter_max, iter_min): all_mcmc = get_mcmc_steps(dirname, nframe, iter_max, iter_min=iter_min) all_logL = np.zeros(len(all_mcmc)) all_logprior = np.zeros(len(all_mcmc)) for i in tqdm(range(len(all_mcmc))): with h5.File(f'{dirname}/mcmc_{all_mcmc[i]}.h5', 'r') as f: s_hat = f['scalars/s_hat_field'][:] all_logL[i] = f['scalars/hmc_Elh'][:] all_logprior[i] = f['scalars/hmc_Eprior'][:] return all_logL, all_logprior