310 lines
No EOL
11 KiB
Python
310 lines
No EOL
11 KiB
Python
import aquila_borg as borg
|
|
import configparser
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
import h5py
|
|
import os
|
|
import jax.numpy as jnp
|
|
|
|
import borg_velocity.likelihood as likelihood
|
|
import borg_velocity.forwards as forwards
|
|
import borg_velocity.utils as utils
|
|
import borg_velocity.projection as projection
|
|
|
|
ini_file = '../conf/supranta_ini.ini'
|
|
mock_dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic/'
|
|
# mock_dirname = None
|
|
test_scaling = False
|
|
test_sigma8 = False
|
|
test_omegam = False
|
|
test_alpha = False
|
|
test_lam = False
|
|
test_muA = False
|
|
test_bulk = False
|
|
test_sample = True
|
|
|
|
# Input box
|
|
box_in = borg.forward.BoxModel()
|
|
config = configparser.ConfigParser()
|
|
config.read(ini_file)
|
|
box_in.L = (float(config['system']['L0']), float(config['system']['L1']), float(config['system']['L2']))
|
|
box_in.N = (int(config['system']['N0']), int(config['system']['N1']), int(config['system']['N2']))
|
|
box_in.xmin = (float(config['system']['corner0']), float(config['system']['corner1']), float(config['system']['corner2']))
|
|
|
|
# Setup BORG forward model and likelihood
|
|
model = likelihood.build_gravity_model(None, box_in, ini_file=ini_file)
|
|
cosmo = utils.get_cosmopar(ini_file)
|
|
model.setCosmoParams(cosmo)
|
|
fwd_param = forwards.NullForward(box_in)
|
|
fwd_vel = likelihood.fwd_vel
|
|
mylike = likelihood.VelocityBORGLikelihood(model, fwd_param, fwd_vel, ini_file)
|
|
|
|
# Create mock data
|
|
state = borg.likelihood.MarkovState()
|
|
mylike.initializeLikelihood(state)
|
|
mylike.updateCosmology(cosmo)
|
|
if mock_dirname is None:
|
|
print('Creating mock data')
|
|
np.random.seed(2)
|
|
s_hat = np.fft.rfftn(np.random.randn(*box_in.N)) / box_in.Ntot ** (0.5)
|
|
mylike.generateMockData(s_hat, state)
|
|
else:
|
|
print('Loading mock data from', mock_dirname)
|
|
cwd = os.getcwd()
|
|
os.chdir(mock_dirname)
|
|
with h5py.File('mock_data.h5') as f:
|
|
s_hat = f['scalars/s_hat_field'][:]
|
|
mylike.loadMockData(state)
|
|
os.chdir(cwd)
|
|
|
|
mylike.logLikelihoodComplex(s_hat, None)
|
|
mylike.gradientLikelihoodComplex(s_hat)
|
|
|
|
quit()
|
|
|
|
if test_scaling:
|
|
all_scale = np.linspace(0.2, 1.8, 100)
|
|
all_lkl = np.empty(all_scale.shape)
|
|
for i, scale in enumerate(all_scale):
|
|
all_lkl[i] = mylike.logLikelihoodComplex(scale * s_hat, None)
|
|
fid_lkl = mylike.logLikelihoodComplex(s_hat, None)
|
|
all_lkl -= fid_lkl
|
|
all_lkl = np.exp(-all_lkl)
|
|
|
|
fig, ax = plt.subplots(1, 1, figsize=(5,5))
|
|
ax.plot(all_scale, all_lkl)
|
|
ax.axhline(y=0, color='k')
|
|
ax.axvline(x=1, color='k')
|
|
ax.set_xlabel(r'$\hat{s}$ scaling')
|
|
ax.set_ylabel(r'$\mathcal{L}$')
|
|
ax.set_xlim(all_scale.min(), all_scale.max())
|
|
fig.tight_layout()
|
|
fig.savefig('../figs/scaling_test.png')
|
|
fig.clf()
|
|
plt.close(fig)
|
|
|
|
# Test sigma8
|
|
if test_sigma8:
|
|
all_sigma8 = np.linspace(0.5, 1.2, 100)
|
|
all_lkl = np.empty(all_sigma8.shape)
|
|
cosmo_true = mylike.fwd.getCosmoParams()
|
|
cosmo = mylike.fwd.getCosmoParams()
|
|
for i, sigma8 in enumerate(all_sigma8):
|
|
cosmo.sigma8 = sigma8
|
|
mylike.updateCosmology(cosmo)
|
|
all_lkl[i] = mylike.logLikelihoodComplex(s_hat, None)
|
|
mylike.updateCosmology(cosmo_true)
|
|
fid_lkl = mylike.logLikelihoodComplex(s_hat, None)
|
|
all_lkl -= fid_lkl
|
|
all_lkl = np.exp(-all_lkl)
|
|
|
|
fig, ax = plt.subplots(1, 1, figsize=(5,5))
|
|
ax.plot(all_sigma8, all_lkl)
|
|
ax.axhline(y=0, color='k')
|
|
ax.axvline(x=cosmo_true.sigma8, color='k')
|
|
ax.set_xlabel(r'$\sigma_8$')
|
|
ax.set_ylabel(r'$\mathcal{L}$')
|
|
fig.tight_layout()
|
|
fig.savefig('../figs/sigma8_test.png')
|
|
fig.clf()
|
|
plt.close(fig)
|
|
|
|
# Test omegam
|
|
if test_omegam:
|
|
all_omegam = np.linspace(0.1, 0.6, 100)
|
|
all_lkl = np.empty(all_omegam.shape)
|
|
cosmo_true = mylike.fwd.getCosmoParams()
|
|
cosmo = mylike.fwd.getCosmoParams()
|
|
for i, omegam in enumerate(all_omegam):
|
|
cosmo.omega_m = omegam
|
|
mylike.updateCosmology(cosmo)
|
|
all_lkl[i] = mylike.logLikelihoodComplex(s_hat, None)
|
|
mylike.updateCosmology(cosmo_true)
|
|
fid_lkl = mylike.logLikelihoodComplex(s_hat, None)
|
|
all_lkl -= fid_lkl
|
|
all_lkl = np.exp(-all_lkl)
|
|
|
|
fig, ax = plt.subplots(1, 1, figsize=(5,5))
|
|
ax.plot(all_omegam, all_lkl)
|
|
ax.axhline(y=0, color='k')
|
|
ax.axvline(x=cosmo_true.omega_m, color='k')
|
|
ax.set_xlabel(r'$\Omega_{\rm m}$')
|
|
ax.set_ylabel(r'$\mathcal{L}$')
|
|
fig.tight_layout()
|
|
fig.savefig('../figs/omegam_test.png')
|
|
fig.clf()
|
|
plt.close(fig)
|
|
|
|
# Test bias model
|
|
if test_alpha:
|
|
all_alpha = np.linspace(-2.0, 5.0, 100)
|
|
all_lkl = np.empty(all_alpha.shape)
|
|
for i, alpha in enumerate(all_alpha):
|
|
mylike.fwd_param.setModelParams({'alpha0':alpha})
|
|
all_lkl[i] = mylike.logLikelihoodComplex(s_hat, None)
|
|
mylike.fwd_param.setModelParams({'alpha0':mylike.alpha[0]})
|
|
fid_lkl = mylike.logLikelihoodComplex(s_hat, None)
|
|
all_lkl -= fid_lkl
|
|
all_lkl = np.exp(-all_lkl)
|
|
|
|
fig, ax = plt.subplots(1, 1, figsize=(5,5))
|
|
ax.plot(all_alpha, all_lkl)
|
|
ax.axhline(y=0, color='k')
|
|
ax.axvline(x=mylike.alpha[0], color='k')
|
|
ax.set_xlabel(r'$\alpha_0$')
|
|
ax.set_ylabel(r'$\mathcal{L}$')
|
|
fig.tight_layout()
|
|
fig.savefig('../figs/alpha_test.png')
|
|
fig.clf()
|
|
plt.close(fig)
|
|
|
|
# Test radial sampling
|
|
if test_lam:
|
|
all_lam = np.linspace(0.0, 10.0, 100)
|
|
all_lkl = np.empty(all_lam.shape)
|
|
for i, lam in enumerate(all_lam):
|
|
mylike.fwd_param.setModelParams({'lam0':lam})
|
|
all_lkl[i] = mylike.logLikelihoodComplex(s_hat, None)
|
|
mylike.fwd_param.setModelParams({'lam0':mylike.lam[0]})
|
|
fid_lkl = mylike.logLikelihoodComplex(s_hat, None)
|
|
all_lkl -= fid_lkl
|
|
all_lkl = np.exp(-all_lkl)
|
|
|
|
fig, ax = plt.subplots(1, 1, figsize=(5,5))
|
|
ax.plot(all_lam, all_lkl)
|
|
ax.axhline(y=0, color='k')
|
|
ax.axvline(x=mylike.lam[0], color='k')
|
|
ax.set_xlabel(r'$\lambda_0$')
|
|
ax.set_ylabel(r'$\mathcal{L}$')
|
|
fig.tight_layout()
|
|
fig.savefig('../figs/lam_test.png')
|
|
fig.clf()
|
|
plt.close(fig)
|
|
|
|
|
|
# Test bias model
|
|
if test_muA:
|
|
all_muA = np.linspace(0.95, 1.05, 100)
|
|
all_lkl = np.empty(all_muA.shape)
|
|
for i, muA in enumerate(all_muA):
|
|
mylike.fwd_param.setModelParams({'mua0':muA})
|
|
all_lkl[i] = mylike.logLikelihoodComplex(s_hat, None)
|
|
mylike.fwd_param.setModelParams({'mua0':mylike.muA[0]})
|
|
fid_lkl = mylike.logLikelihoodComplex(s_hat, None)
|
|
all_lkl -= fid_lkl
|
|
all_lkl = np.exp(-all_lkl)
|
|
|
|
fig, ax = plt.subplots(1, 1, figsize=(5,5))
|
|
ax.plot(all_muA, all_lkl)
|
|
ax.axhline(y=0, color='k')
|
|
ax.axvline(x=mylike.muA[0], color='k')
|
|
ax.set_xlabel(r'$\mu_0$')
|
|
ax.set_ylabel(r'$\mathcal{L}$')
|
|
fig.tight_layout()
|
|
fig.savefig('../figs/muA_test.png')
|
|
fig.clf()
|
|
plt.close(fig)
|
|
|
|
|
|
# Test bulk flow
|
|
if test_bulk:
|
|
all_vx = np.linspace(-100, 100, 50)
|
|
all_lkl = np.empty(all_vx.shape)
|
|
|
|
for i, vx in enumerate(all_vx):
|
|
mylike.fwd_param.setModelParams({'bulk_flow_x':vx})
|
|
all_lkl[i] = mylike.logLikelihoodComplex(s_hat, None)
|
|
mylike.fwd_param.setModelParams({'bulk_flow_x':mylike.bulk_flow[0]})
|
|
# fid_lkl = mylike.logLikelihoodComplex(s_hat, None)
|
|
fid_lkl = np.amin(all_lkl)
|
|
all_lkl -= fid_lkl
|
|
all_lkl = np.exp(-all_lkl)
|
|
|
|
sigma_bulk = utils.get_sigma_bulk(mylike.L[0], mylike.fwd.getCosmoParams())
|
|
ref_lkl = np.exp(-all_vx ** 2 / 2 / sigma_bulk ** 2)
|
|
ref_lkl *= np.amax(all_lkl) / np.amax(ref_lkl)
|
|
|
|
fig, ax = plt.subplots(1, 1, figsize=(5,5))
|
|
ax.plot(all_vx, all_lkl, label='Posterior')
|
|
ax.plot(all_vx, ref_lkl, ls='--', label=r'$\Lambda$CDM prior')
|
|
ax.axhline(y=0, color='k')
|
|
ax.axvline(x=mylike.muA[0], color='k')
|
|
ax.legend()
|
|
ax.set_xlabel(r'$v_{{\rm bulk}, x}$')
|
|
ax.set_ylabel(r'$\mathcal{L}$')
|
|
fig.tight_layout()
|
|
fig.savefig('../figs/bulk_test.png')
|
|
fig.clf()
|
|
plt.close(fig)
|
|
|
|
|
|
if test_sample:
|
|
with h5py.File(f'{mock_dirname}/mcmc_96.h5') as f:
|
|
sample_s_hat = f['scalars/s_hat_field'][:]
|
|
mylike.logLikelihoodComplex(s_hat, None)
|
|
np.savez(f'{mock_dirname}/mock_lkl_ind.npz', *mylike.lkl_ind)
|
|
mylike.logLikelihoodComplex(sample_s_hat, None)
|
|
np.savez(f'{mock_dirname}/mcmc_lkl_ind.npz', *mylike.lkl_ind)
|
|
|
|
print(s_hat.shape)
|
|
N = s_hat.shape[0]
|
|
|
|
# Run BORG density field - s_hat
|
|
print('\nTrue')
|
|
output_density = np.zeros((N,N,N))
|
|
mylike.fwd.forwardModel_v2(s_hat)
|
|
mylike.fwd.getDensityFinal(output_density)
|
|
output_velocity = mylike.fwd_vel.getVelocityField()
|
|
np.save(f'{mock_dirname}/mock_vel.npy', output_velocity)
|
|
tracer_vr = [None] * mylike.nsamp
|
|
tracer_vel = [None] * mylike.nsamp
|
|
print('3D', output_velocity.min(), output_velocity.max())
|
|
for i in range(mylike.nsamp):
|
|
tracer_vel[i] = projection.interp_field(
|
|
output_velocity,
|
|
mylike.MB_pos[i],
|
|
mylike.L[0],
|
|
jnp.array(mylike.fwd.getOutputBoxModel().xmin),
|
|
1,
|
|
use_jitted=True,
|
|
)
|
|
tracer_vr[i] = projection.project_radial(
|
|
tracer_vel[i],
|
|
mylike.MB_pos[i],
|
|
jnp.zeros(3,)
|
|
)
|
|
print(i, tracer_vr[i].min(), tracer_vr[i].max(), tracer_vel[i].min(), tracer_vel[i].max())
|
|
np.savez(f'{mock_dirname}/mock_tracer_vr.npz', *tracer_vr)
|
|
np.savez(f'{mock_dirname}/mock_tracer_vel.npz', *tracer_vel)
|
|
|
|
# Run BORG density field - sample_s_hat
|
|
print('\nMCMC Sample')
|
|
output_density = np.zeros((N,N,N))
|
|
mylike.fwd.forwardModel_v2(sample_s_hat)
|
|
mylike.fwd.getDensityFinal(output_density)
|
|
output_velocity = mylike.fwd_vel.getVelocityField()
|
|
np.save(f'{mock_dirname}/mcmc_vel.npy', output_velocity)
|
|
tracer_vr = [None] * mylike.nsamp
|
|
tracer_vel = [None] * mylike.nsamp
|
|
print('3D', output_velocity.min(), output_velocity.max())
|
|
for i in range(mylike.nsamp):
|
|
tracer_vel[i] = projection.interp_field(
|
|
output_velocity,
|
|
mylike.MB_pos[i],
|
|
mylike.L[0],
|
|
jnp.array(mylike.fwd.getOutputBoxModel().xmin),
|
|
1,
|
|
use_jitted=True,
|
|
)
|
|
tracer_vr[i] = projection.project_radial(
|
|
tracer_vel[i],
|
|
mylike.MB_pos[i],
|
|
jnp.zeros(3,)
|
|
)
|
|
print(i, tracer_vr[i].min(), tracer_vr[i].max(), tracer_vel[i].min(), tracer_vel[i].max())
|
|
np.savez(f'{mock_dirname}/mcmc_tracer_vr.npz', *tracer_vr)
|
|
np.savez(f'{mock_dirname}/mcmc_tracer_vel.npz', *tracer_vel)
|
|
|
|
# Save the MB points
|
|
np.savez(f'{mock_dirname}/MB_pos.npz', *mylike.MB_pos)
|
|
|