Generate mock data from borg, create MB data and prepare likelihood

tests/test_forward.py

@@ -0,0 +1,84 @@
+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')