dbartlett/borg_velocity
0
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Ensure tracers are inside box

Browse source
This commit is contained in:
Deaglan Bartlett 2024-11-04 15:44:20 +01:00
parent ea7019b2db
commit cbcd1fce0f
17 changed files with 735 additions and 52 deletions
Download patch file Download diff file Expand all files Collapse all files

BIN
tests/s_hat.npy Normal file
View file

Binary file not shown.

115
tests/test_likelihood.py
View file

@ -2,19 +2,26 @@ 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'
test_scaling = True
test_sigma8 = True
test_omegam = True
test_alpha = True
test_lam = True
test_muA = True
test_bulk = True
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()
@ -36,12 +43,27 @@ mylike = likelihood.VelocityBORGLikelihood(model, fwd_param, fwd_vel, ini_file)
state = borg.likelihood.MarkovState()
mylike.initializeLikelihood(state)
mylike.updateCosmology(cosmo)
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)
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.5, 1.5, 100)
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)
@ -214,4 +236,75 @@ if test_bulk:
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)