matter_field.py

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from scipy.fftpack import fftn, ifftn

import camb

h=0.675
pars = camb.set_params(
    H0=100*h, ombh2=0.022, omch2=0.122, mnu=0.06, omk=0, tau=0.06,  
    As=2e-9, ns=0.965, halofit_version='mead'
)
redshifts = [1e4, 1e3, 1e2, 10, 5, 4, 3, 2, 1, 0.75, 0.5, 0.25, 0.2, 0.15, 0.1, 0.05, 0]
pars.set_matter_power(redshifts=redshifts, kmax=2.0)
results = camb.get_results(pars)
kh, z, pk = results.get_matter_power_spectrum(minkh=1e-4, maxkh=1, npoints=200)

# Grid size and resolution
N = 256  # Number of grid points along each dimension
box_size = 100.0  # Physical size of the box in Mpc

# Generate k-space grid
kx = np.fft.fftfreq(N, box_size/N)
ky = np.fft.fftfreq(N, box_size/N)
kz = np.fft.fftfreq(N, box_size/N)
KX, KY, KZ = np.meshgrid(kx, ky, kz, indexing="ij")
k = np.sqrt(KX**2 + KY**2 + KZ**2)  # Wavenumber magnitude

# Generate random complex field in Fourier space
rng = np.random.default_rng(seed=2611)
random_phase = np.exp(1j*rng.uniform(high=2*np.pi, size=(N, N, N)))

# Precompute global min and max for the density field
global_min, global_max = float('inf'), float('-inf')
for frame in range(len(redshifts)):
    P_k = pk[len(redshifts) - frame - 1]
    P_k_3d = np.interp(k.flatten(), kh * h, P_k / h**3).reshape(k.shape)
    delta_k = np.sqrt(P_k_3d) * random_phase
    density_field = np.real(ifftn(delta_k))
    global_min = min(global_min, density_field.min())
    global_max = max(global_max, density_field.max())

# Generate animation
duration = 9
fig, ax = plt.subplots(figsize=(8, 6))
fig.patch.set_facecolor("#e3e0d1")
im = ax.imshow(np.random.rand(N, N), vmin=global_min, vmax=global_max, extent=[0, box_size, 0, box_size], cmap="seismic")
title = ax.set_title("")
cbar = fig.colorbar(im, label=r"$\delta_m(x, y)$", ax=ax)
# From https://matplotlib.org/stable/gallery/specialty_plots/anscombe.html#sphx-glr-gallery-specialty-plots-anscombe-py
bbox = dict(boxstyle='round', fc='blanchedalmond', ec='orange', alpha=0.85)
text = ax.text(2.5, 2.5, "z", bbox=bbox)
plt.xlabel(r"$x$ (Mpc)")
plt.ylabel(r"$y$ (Mpc)")
plt.title(f"Matter density contrast")

def update(frame):
    # Visualizing a slice
    z = redshifts[frame]
    P_k = pk[len(redshifts) - frame - 1]
    P_k_3d = np.interp(k.flatten(), kh * h, P_k / h**3).reshape(k.shape)
    delta_k = np.sqrt(P_k_3d) * random_phase
    density_field = np.real(ifftn(delta_k))
    im.set_data(density_field[:, :, N//2])
    if z > 1: text.set_text(rf"$z = {z:.0f}$")
    else: text.set_text(rf"$z = {z:.2f}$")
    return (im, text)

interval = duration*1000/len(redshifts)
ani = FuncAnimation(fig, update, frames=len(redshifts), interval=interval, blit=False)
# Save the animation as a GIF
ani.save('my_animation.gif', writer='pillow', fps=1000/interval)