161 lines
5.1 KiB
Python
161 lines
5.1 KiB
Python
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
|
||
import pylab as pl
|
||
pl.ion()
|
||
|
||
@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
|
||
elif which_field.startwith('BORG_final_velocity'):
|
||
MAS = "CIC"
|
||
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[f'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
|
||
|
||
def plot_dens_fields(mean_field, std_field, true_field, slice):
|
||
"""
|
||
"""
|
||
sides = mean_field.shape[0]
|
||
x = np.arange(sides)
|
||
y = x
|
||
X, Y = np.meshgrid(x, y)
|
||
fig, axes = pl.subplots(figsize=(16,7), ncols=3)
|
||
_, _, _, im0 = axes[0].hist2d(X.ravel(), Y.ravel(), bins=(sides, sides), weights=true_field[slice].ravel())
|
||
axes[0].set_xlabel('x')
|
||
axes[0].set_ylabel('y')
|
||
axes[0].set_title('Generated mock density')
|
||
pl.colorbar(im0)
|
||
|
||
_, _, _, im1 = axes[1].hist2d(X.ravel(), Y.ravel(), bins=(sides, sides), weights=mean_field[slice].ravel())
|
||
axes[1].set_xlabel('x')
|
||
axes[1].set_ylabel('y')
|
||
axes[1].set_title('Mean reconstructed density')
|
||
pl.colorbar(im1)
|
||
|
||
_, _, _, im2 = axes[2].hist2d(X.ravel(), Y.ravel(), bins=(sides, sides), weights=std_field[slice].ravel())
|
||
axes[2].set_xlabel('x')
|
||
axes[2].set_ylabel('y')
|
||
axes[2].set_title('Std of reconstructed density')
|
||
pl.colorbar(im2)
|
||
|