84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
import numpy as np
|
|
import aquila_borg as borg
|
|
import configparser
|
|
import matplotlib.pyplot as plt
|
|
|
|
import borg_velocity.likelihood as likelihood
|
|
import borg_velocity.utils as utils
|
|
import borg_velocity.forwards as forwards
|
|
|
|
ini_file = '../conf/basic_ini.ini'
|
|
|
|
# Setup config
|
|
config = configparser.ConfigParser()
|
|
config.read(ini_file)
|
|
|
|
# Cosmology
|
|
cosmo = utils.get_cosmopar(ini_file)
|
|
|
|
# Input box
|
|
box_in = borg.forward.BoxModel()
|
|
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']))
|
|
|
|
# Load forward model
|
|
fwd_model = likelihood.build_gravity_model(None, box_in, ini_file=ini_file)
|
|
fwd_model.setCosmoParams(cosmo)
|
|
|
|
# Make some initial conditions
|
|
s_hat = np.fft.rfftn(np.random.randn(*box_in.N)) / box_in.Ntot ** (0.5)
|
|
|
|
# Get the real version of this
|
|
s_real = np.fft.irfftn(s_hat, norm="ortho")
|
|
|
|
# Run BORG density field
|
|
output_density = np.zeros(box_in.N)
|
|
fwd_model.forwardModel_v2(s_hat)
|
|
fwd_model.getDensityFinal(output_density)
|
|
|
|
# Get growth rate
|
|
cosmology = borg.cosmo.Cosmology(cosmo)
|
|
af = float(config['model']['af'])
|
|
f = cosmology.gplus(af) # dD / da
|
|
f *= af / cosmology.d_plus(af) # f = dlnD / dlna
|
|
|
|
# Get velocity
|
|
smooth_R = float(config['model']['smooth_R'])
|
|
output_vel = forwards.dens2vel_linear(output_density, f, box_in.L[0], smooth_R)
|
|
print(output_vel.shape)
|
|
|
|
|
|
# Plot the IC and dens fields
|
|
fig, axs = plt.subplots(2, 3, figsize=(15,10))
|
|
for i, field in enumerate([s_real, np.log10(2 + output_density)]):
|
|
vmin, vmax = field.min(), field.max()
|
|
pc = axs[i,0].pcolor(field[box_in.N[0]//2], vmin=vmin, vmax=vmax)
|
|
fig.colorbar(pc, ax=axs[i,0])
|
|
pc = axs[i,1].pcolor(field[:,box_in.N[1]//2,:], vmin=vmin, vmax=vmax)
|
|
fig.colorbar(pc, ax=axs[i,1])
|
|
pc = axs[i,2].pcolor(field[:,:,box_in.N[2]//2], vmin=vmin, vmax=vmax)
|
|
fig.colorbar(pc, ax=axs[i,2])
|
|
axs[0,1].set_title('Initial Conditions')
|
|
axs[1,1].set_title(r'$\log_{10} (2 + \delta)$')
|
|
for ax in axs.flatten():
|
|
ax.set_aspect('equal')
|
|
fig.savefig('../figs/test_dens_ic.png')
|
|
|
|
|
|
# Plot the velocity fields
|
|
fig, axs = plt.subplots(3, 3, figsize=(15,15))
|
|
vmin, vmax = output_vel.min(), output_vel.max()
|
|
for i in range(3):
|
|
pc = axs[i,0].pcolor(output_vel[i,box_in.N[0]//2], vmin=vmin, vmax=vmax)
|
|
fig.colorbar(pc, ax=axs[i,0])
|
|
pc = axs[i,1].pcolor(output_vel[i,:,box_in.N[1]//2,:], vmin=vmin, vmax=vmax)
|
|
fig.colorbar(pc, ax=axs[i,1])
|
|
pc = axs[i,2].pcolor(output_vel[i,:,:,box_in.N[2]//2], vmin=vmin, vmax=vmax)
|
|
fig.colorbar(pc, ax=axs[i,2])
|
|
axs[0,1].set_title(r'$v_x$')
|
|
axs[1,1].set_title(r'$v_y$')
|
|
axs[2,1].set_title(r'$v_z$')
|
|
for ax in axs.flatten():
|
|
ax.set_aspect('equal')
|
|
fig.savefig('../figs/test_vel.png')
|