833 KiB
833 KiB
Tristan Hoellinger
Institut d'Astrophysique de Paris
tristan.hoellinger@iap.fr
Non-regression tests towards implementing P3M gravity¶
Set up the environment and parameters¶
In [1]:
from wip3m.params import *
In [2]:
workdir = ROOT_PATH
output_path = OUTPUT_PATH
# L = 64 # Box size in Mpc/h
# L = 32 # Box size in Mpc/h
L = 16 # Box size in Mpc/h
# L = 8 # Box size in Mpc/h
N = 32 # Density grid size
Np = 32 # Number of dark matter particles per spatial dimension
Npm = 64 # PM grid size
n_Tiles = 8
go_beyond_Nyquist_ss = True # for the summary statistics
force = True
force_hard = True
run_id = "run13"
# Good set of parameters for the force diagnostic
# nPairsForceDiagnostic_pm = nPairsForceDiagnostic_spm = nPairsForceDiagnostic_p3m = 5
# nBinsForceDiagnostic = 30
# maxTrialsForceDiagnostic = int(2e9)
# Faster force diagnostic
nPairsForceDiagnostic_pm = nPairsForceDiagnostic_spm = nPairsForceDiagnostic_p3m = 3
nBinsForceDiagnostic = 20
maxTrialsForceDiagnostic = int(1e8)
# Simulation parameters
# nsteps_cola = 20
# nsteps_pm = 200
# nsteps_spm = 200
# nsteps_p3m = 200
nsteps_cola = 20
nsteps_pm = 50
nsteps_spm = 50
nsteps_p3m = 50
# nsteps_cola = 10
# nsteps_pm = 20
# nsteps_spm = 20
# nsteps_p3m = 20
In principle nothing needs to be changed below this cell.
In [3]:
# Automatic reloading of modules
%load_ext autoreload
%autoreload 2
from os.path import isfile
from pathlib import Path
import numpy as np
from pysbmy.power import PowerSpectrum
from pysbmy.fft import FourierGrid, read_FourierGrid
from pysbmy.field import read_field
from pysbmy.correlations import get_autocorrelation
from wip3m.tools import get_k_max, generate_sim_params, generate_white_noise_Field
from wip3m.params import params_planck_kmax_missing, cosmo_small_to_full_dict, z2a, BASELINE_SEEDPHASE
from wip3m.plot_utils import *
In [4]:
corner = -L / 2.0
RedshiftLPT = 19.0
RedshiftFCs = 0.0
ai = z2a(RedshiftLPT)
af = z2a(RedshiftFCs)
k_max = get_k_max(L, N) # k_max in h/Mpc
print(f"k_max = {k_max}")
cosmo = params_planck_kmax_missing.copy()
cosmo["k_max"] = k_max
wd = workdir + run_id + "/"
simdir = output_path + run_id + "/"
logdir = simdir + "logs/"
if force_hard:
import shutil
if Path(simdir).exists():
shutil.rmtree(simdir)
if Path(wd).exists():
shutil.rmtree(wd)
Path(wd).mkdir(parents=True, exist_ok=True)
Path(logdir).mkdir(parents=True, exist_ok=True)
input_white_noise_file = simdir + "input_white_noise.h5"
input_seed_phase_file = simdir + "seed"
ICs_path = simdir + "initial_density.h5"
simpath = simdir
# Path to the input matter power spectrum (generated later)
input_power_file = simdir + "input_power.h5"
# Paths the the force diagnostic CSVs
OutputForceDiagnostic_pm = simdir + "force_diagnostic_pm.txt"
OutputForceDiagnostic_cola = simdir + "force_diagnostic_cola.txt"
OutputForceDiagnostic_spm = simdir + "force_diagnostic_spm.txt"
OutputForceDiagnostic_p3m = simdir + "force_diagnostic_p3m.txt"
Generate the parameter files¶
The first preparatory step is to generate all the parameter files required for all the simulations.
To this end we use the generate_sim_params function defined in params.py.
In [5]:
common_params = {
"Np": Np,
"N": N,
"L": L,
"corner0": corner,
"corner1": corner,
"corner2": corner,
"h": cosmo["h"],
"Omega_m": cosmo["Omega_m"],
"Omega_b": cosmo["Omega_b"],
"n_s": cosmo["n_s"],
"sigma8": cosmo["sigma8"],
}
lpt_params = common_params.copy()
lpt_params["method"] = "lpt"
lpt_params["InputPowerSpectrum"] = input_power_file
lpt_params["ICsMode"] = 1
# 0 : the codes generates white noise, then initial conditions
# 1 : external white noise specified, the code multiplies by the power spectrum
# 2 : external initial conditions specified
lpt_params["InputWhiteNoise"] = input_white_noise_file
pm_params = common_params.copy()
pm_params["method"] = "pm"
pm_params["TimeStepDistribution"] = 0
pm_params["ai"] = ai
pm_params["af"] = af
pm_params["RedshiftLPT"] = RedshiftLPT
pm_params["RedshiftFCs"] = RedshiftFCs
pm_params["Npm"] = Npm
pm_params["nsteps"] = nsteps_pm
pm_params["RunForceDiagnostic"] = True
pm_params["nPairsForceDiagnostic"] = nPairsForceDiagnostic_pm
pm_params["nBinsForceDiagnostic"] = nBinsForceDiagnostic
pm_params["OutputForceDiagnostic"] = OutputForceDiagnostic_pm
pm_params["maxTrialsForceDiagnostic"] = maxTrialsForceDiagnostic
cola_params = common_params.copy()
cola_params["method"] = "cola"
cola_params["TimeStepDistribution"] = 0
cola_params["ai"] = ai
cola_params["af"] = af
cola_params["RedshiftLPT"] = RedshiftLPT
cola_params["RedshiftFCs"] = RedshiftFCs
cola_params["Npm"] = Npm
cola_params["nsteps"] = nsteps_cola
cola_params["RunForceDiagnostic"] = False
cola_params["RunForceDiagnostic"] = True
cola_params["nPairsForceDiagnostic"] = nPairsForceDiagnostic_pm
cola_params["nBinsForceDiagnostic"] = nBinsForceDiagnostic
cola_params["OutputForceDiagnostic"] = OutputForceDiagnostic_cola
cola_params["maxTrialsForceDiagnostic"] = maxTrialsForceDiagnostic
spm_params = common_params.copy()
spm_params["method"] = "spm"
spm_params["EvolutionMode"] = 5
spm_params["TimeStepDistribution"] = 0
spm_params["ai"] = ai
spm_params["af"] = af
spm_params["RedshiftLPT"] = RedshiftLPT
spm_params["RedshiftFCs"] = RedshiftFCs
spm_params["Npm"] = Npm
spm_params["nsteps"] = nsteps_spm
spm_params["n_Tiles"] = n_Tiles
spm_params["RunForceDiagnostic"] = True
spm_params["nPairsForceDiagnostic"] = nPairsForceDiagnostic_spm
spm_params["nBinsForceDiagnostic"] = nBinsForceDiagnostic
spm_params["OutputForceDiagnostic"] = OutputForceDiagnostic_spm
spm_params["maxTrialsForceDiagnostic"] = maxTrialsForceDiagnostic
p3m_params = common_params.copy()
p3m_params["method"] = "p3m"
p3m_params["EvolutionMode"] = 4
p3m_params["TimeStepDistribution"] = 0
p3m_params["ai"] = ai
p3m_params["af"] = af
p3m_params["RedshiftLPT"] = RedshiftLPT
p3m_params["RedshiftFCs"] = RedshiftFCs
p3m_params["Npm"] = Npm
p3m_params["nsteps"] = nsteps_p3m
p3m_params["n_Tiles"] = n_Tiles
p3m_params["RunForceDiagnostic"] = True
p3m_params["nPairsForceDiagnostic"] = nPairsForceDiagnostic_p3m
p3m_params["nBinsForceDiagnostic"] = nBinsForceDiagnostic
p3m_params["OutputForceDiagnostic"] = OutputForceDiagnostic_p3m
p3m_params["maxTrialsForceDiagnostic"] = maxTrialsForceDiagnostic
In [6]:
reset_plotting() # Default style for Simbelmynë
generate_sim_params(lpt_params, ICs_path, wd, simdir, None, force)
print(f"PM nsteps = {nsteps_pm}:")
file_ext = f"nsteps{nsteps_pm}" # "pm" is already in the filename
generate_sim_params(pm_params, ICs_path, wd, simdir, file_ext, force)
print(f"COLA nsteps = {nsteps_cola}:")
file_ext = f"nsteps{nsteps_cola}" # "cola" is already in the filename
generate_sim_params(cola_params, ICs_path, wd, simdir, file_ext, force)
print(f"SPM nsteps = {nsteps_spm}:")
file_ext = f"nsteps{nsteps_spm}" # "spm" is already in the filename
generate_sim_params(spm_params, ICs_path, wd, simdir, file_ext, force)
print(f"P3M nsteps = {nsteps_p3m}:")
file_ext = f"nsteps{nsteps_p3m}" # "p3m" is already in the filename
generate_sim_params(p3m_params, ICs_path, wd, simdir, file_ext, force)
setup_plotting() # Reset plotting style for this project
Generate the initial phase¶
In [7]:
generate_white_noise_Field(
L=L,
size=N,
corner=corner,
seedphase=BASELINE_SEEDPHASE,
fname_whitenoise=input_white_noise_file,
seedname_whitenoise=input_seed_phase_file,
force_phase=force,
)
Generating the input power spectrum¶
The second preparatory step is to compute the initial power spectrum to be used in the simulations, given the cosmological parameters and prescription specified in params.py. The power spectrum is saved in input_power_file.
In [8]:
# If cosmo["WhichSpectrum"] == "class", then the module classy is required.
if not isfile(input_power_file) or force:
Pk = PowerSpectrum(L, L, L, N, N, N, cosmo_small_to_full_dict(cosmo))
Pk.write(input_power_file)
In [9]:
# k grid used to compute the final overdensity power spectrum
Pinit = 100
trim_threshold = 100 # Merge bins until this minimum number of modes per bin is reached
log_kmin = np.log10(2 * np.pi / (np.sqrt(3) * L)) # Minimum non-zero k in h/Mpc
if go_beyond_Nyquist_ss:
k_max_ss = get_k_max(L, N)
else:
k_max_ss = get_k_max(L, N) / np.sqrt(3) # 1D Nyquist frequency
Pbins_left_bnds = np.logspace(log_kmin, np.log10(k_max_ss), Pinit + 1, dtype=np.float32)
Pbins_left_bnds = Pbins_left_bnds[:-1]
input_ss_file = simdir + "input_ss_k_grid.h5"
Gk = FourierGrid(
L,
L,
L,
N,
N,
N,
k_modes=Pbins_left_bnds,
kmax=k_max_ss,
trim_bins=True,
trim_threshold=trim_threshold,
)
Gk.write(input_ss_file)
Running the simulations¶
We are now ready to run the actual simulations using the Simbelmynë executable. Warning: the following may take some time, even in relatively low dimension, and should not be run on a login node.
In [10]:
%%capture
if not isfile(ICs_path) or not isfile(simdir + "lpt_density.h5") or not isfile(simdir + "lpt_particles.gadget3") or force:
!simbelmyne {wd}example_lpt.sbmy {logdir}lpt.txt
file_ext = f"nsteps{nsteps_pm}" # "pm" is already in the filename
if not isfile(simdir + f"{file_ext}_final_density_pm.h5") or force:
!simbelmyne {wd}{file_ext}_example_pm.sbmy {logdir}{file_ext}_pm.txt
file_ext = f"nsteps{nsteps_cola}" # "cola" is already in the filename
if not isfile(simdir + f"{file_ext}_final_density_cola.h5") or force:
!simbelmyne {wd}{file_ext}_example_cola.sbmy {logdir}{file_ext}_cola.txt
file_ext = f"nsteps{nsteps_spm}" # "spm" is already in the filename
if not isfile(simdir + f"{file_ext}_final_density_spm.h5") or force:
!simbelmyne {wd}{file_ext}_example_spm.sbmy {logdir}{file_ext}_spm.txt
file_ext = f"nsteps{nsteps_p3m}" # "p3m" is already in the filename
if not isfile(simdir + f"{file_ext}_final_density_p3m.h5") or force:
!simbelmyne {wd}{file_ext}_example_p3m.sbmy {logdir}{file_ext}_p3m.txt
The logs can be monitored in the corresponding files in the logdir directory.
Plot results¶
Plot the evolved dark matter density fields¶
In [11]:
slice_ijk = (N // 2, slice(None), slice(None))
DELTA_LPT = read_field(simdir + "lpt_density.h5").data[slice_ijk]
DELTA_COLA = read_field(simdir + f"nsteps{nsteps_cola}_final_density_cola.h5").data[slice_ijk]
DELTA_PM = read_field(simdir + f"nsteps{nsteps_pm}_final_density_pm.h5").data[slice_ijk]
DELTA_SPM = read_field(simdir + f"nsteps{nsteps_spm}_final_density_spm.h5").data[slice_ijk]
DELTA_P3M = read_field(simdir + f"nsteps{nsteps_p3m}_final_density_p3m.h5").data[slice_ijk]
diff_p3m_pm = DELTA_P3M - DELTA_PM
diff_p3m_spm = DELTA_P3M - DELTA_SPM
In [12]:
print(f"max(DELTA_PM) = {np.max(DELTA_PM)}, min(DELTA_PM) = {np.min(DELTA_PM)}")
print(f"max(DELTA_P3M) = {np.max(DELTA_P3M)}, min(DELTA_P3M) = {np.min(DELTA_P3M)}")
print(f"max(diff) = {np.max(diff_p3m_pm)}, min(diff) = {np.min(diff_p3m_pm)}")
In [13]:
# fields = ["pm", "spm", "p3m", "diff_p3m_pm"] # fields to plot
fields = ["lpt", "cola", "pm", "spm", "p3m", "diff_p3m_pm", "diff_p3m_spm"] # fields to plot
figname = "_".join(fields)
slices_dict = {
"lpt": DELTA_LPT,
"cola": DELTA_COLA,
"pm": DELTA_PM,
"spm": DELTA_SPM,
"p3m": DELTA_P3M,
"diff_p3m_pm": diff_p3m_pm,
"diff_p3m_spm": diff_p3m_spm,
}
titles_dict = {
"lpt": "LPT",
"cola": f"COLA $n_\\mathrm{{steps}}={nsteps_cola}$",
"pm": f"PM $n_\\mathrm{{steps}}={nsteps_pm}$",
"spm": f"sPM $n_\\mathrm{{steps}}={nsteps_spm}$",
"p3m": f"P3M $n_\\mathrm{{steps}}={nsteps_p3m}$",
"diff_p3m_pm": r"$\delta_{\rm P3M}-\delta_{\rm PM}$",
"diff_p3m_spm": r"$\delta_{\rm P3M}-\delta_{\rm sPM}$",
}
npanels = len(fields)
fig, axs = plt.subplots(1, npanels, figsize=(3 * npanels, 4), sharey=True)
ims = []
for i, key in enumerate(fields):
ax = axs[i]
data = slices_dict[key]
title = titles_dict[key]
if key.startswith("diff"):
norm = TwoSlopeNorm(vmin=-np.log(1 + np.abs(np.min(data))), vcenter=0, vmax=np.log10(1 + np.abs(np.max(data))))
im = ax.imshow(
np.sign(data) * np.log(1 + np.abs(data)), cmap="RdBu_r", norm=norm
)
else:
im = ax.imshow(np.log10(2 + data), cmap=cmap)
ims.append((im, key))
ax.set_title(title, fontsize=fs_titles)
for spine in ax.spines.values():
spine.set_visible(False)
axs[0].set_yticks([0, N // 2, N])
axs[0].set_yticklabels([f"{-L/2:.0f}", "0", f"{L/2:.0f}"], fontsize=fs)
axs[0].set_ylabel(r"Mpc/$h$", size=GLOBAL_FS_SMALL)
for i, ax in enumerate(axs):
ax.set_xticks([0, N // 2, N])
ax.set_xticklabels([f"{-L/2:.0f}", "0", f"{L/2:.0f}"], fontsize=fs)
ax.set_xlabel(r"Mpc/$h$", size=GLOBAL_FS_SMALL)
for ax, (im, key) in zip(axs, ims):
divider = make_axes_locatable(ax)
cax = divider.append_axes("bottom", size="5%", pad=0.6)
cb = fig.colorbar(im, cax=cax, orientation="horizontal")
if key.startswith("diff"):
cb.set_label(r"$\textrm{sgn}\left(\Delta\delta\right)\log_{10}(1 + |\Delta\delta|)$", fontsize=fs)
else:
cb.set_label(r"$\log_{10}(2 + \delta)$", fontsize=fs)
cb.ax.tick_params(labelsize=fs)
cax.xaxis.set_ticks_position("bottom")
cax.xaxis.set_label_position("bottom")
fig.savefig(
simdir + f"{figname}.png",
bbox_inches="tight",
dpi=300,
transparent=True,
)
fig.savefig(
simdir + f"{figname}.pdf",
bbox_inches="tight",
dpi=300,
)
plt.show()
In [14]:
full_field_p3m = np.log10(2+read_field(simdir + f"nsteps{nsteps_p3m}_final_density_p3m.h5").data)
if N <= 128:
fig = plotly_3d(full_field_p3m, size=N, L=L, colormap=thermal_plotly, limits="default")
else:
# Downsample the grid for visualisation
downsample_factor = N // 128
downsampled_field = full_field_p3m[
::downsample_factor, ::downsample_factor, ::downsample_factor
]
fig = plotly_3d(downsampled_field, size=N, L=L, colormap=thermal_plotly, limits="default")
fig.show()
clear_large_plot(fig) # Uncomment to clear the Plotly figure to avoid memory issues
Compute and plot the power spectra of the evolved dark matter fields¶
In [15]:
G = read_FourierGrid(simdir + "input_ss_k_grid.h5")
k = G.k_modes[1:]
AliasingCorr = False
DELTA = read_field(simdir + "initial_density.h5")
Pk_INI, Vk_INI = get_autocorrelation(DELTA, G, AliasingCorr)
Pk_INI, Vk_INI = Pk_INI[1:], Vk_INI[1:]
Sk_INI = np.sqrt(Vk_INI)
DELTA = read_field(simdir + "lpt_density.h5")
Pk_LPT, Vk_LPT = get_autocorrelation(DELTA, G, AliasingCorr)
Pk_LPT, Vk_LPT = Pk_LPT[1:], Vk_LPT[1:]
Sk_LPT = np.sqrt(Vk_LPT)
DELTA = read_field(simdir + f"nsteps{nsteps_pm}_final_density_pm.h5")
Pk_PM, Vk_PM = get_autocorrelation(DELTA, G, AliasingCorr)
Pk_PM, Vk_PM = Pk_PM[1:], Vk_PM[1:]
Sk_PM = np.sqrt(Vk_PM)
DELTA = read_field(simdir + f"nsteps{nsteps_cola}_final_density_cola.h5")
Pk_COLA, Vk_COLA = get_autocorrelation(DELTA, G, AliasingCorr)
Pk_COLA, Vk_COLA = Pk_COLA[1:], Vk_COLA[1:]
Sk_COLA = np.sqrt(Vk_COLA)
DELTA = read_field(simdir + f"nsteps{nsteps_spm}_final_density_spm.h5")
Pk_sPM, Vk_sPM = get_autocorrelation(DELTA, G, AliasingCorr)
Pk_sPM, Vk_sPM = Pk_sPM[1:], Vk_sPM[1:]
Sk_sPM = np.sqrt(Vk_sPM)
DELTA = read_field(simdir + f"nsteps{nsteps_p3m}_final_density_p3m.h5")
Pk_P3M, Vk_p3m = get_autocorrelation(DELTA, G, AliasingCorr)
Pk_P3M, Vk_p3m = Pk_P3M[1:], Vk_p3m[1:]
Sk_p3m = np.sqrt(Vk_p3m)
In [16]:
Pk_ref = Pk_PM
fig, ax = plt.subplots(figsize=(7, 4))
ax.set_xscale("log")
k = G.k_modes[1:]
kmin, kmax = k.min(), k.max()
print(f"kmin = {kmin}, kmax = {kmax}")
log_pad = 0.02
log_k_min = np.log10(kmin)
log_k_max = np.log10(kmax)
log_range = log_k_max - log_k_min
xlim_min = 10 ** (log_k_min - log_pad * log_range)
xlim_max = 10 ** (log_k_max + log_pad * log_range)
plt.xlim(xlim_min, xlim_max)
ax.set_ylim([0.2, 1.8])
# ax.set_ylim([0.5, 1.5])
dark_grey_bnd = 0.05
light_grey_bnd = 0.1
label_ref = f"PM with $n_\\mathrm{{steps}}={nsteps_pm}$"
line1 = ax.plot([kmin, kmax], [1, 1], color="black", linestyle="-", label=label_ref)
# line2 = ax.plot(k, Pk_LPT / Pk_ref, label="2LPT", color=cols[0], linestyle="-")
ax.plot(
k,
Pk_COLA / Pk_ref,
label=f"COLA with $n_\\mathrm{{steps}}={nsteps_cola}$",
linestyle="-",
color=cols[2],
)
# ax.plot(
# k,
# Pk_PM / Pk_ref,
# label=f"PM with $n_\\mathrm{{steps}}={nsteps_pm}$",
# linestyle="-",
# color=cols[3],
# )
ax.plot(
k,
Pk_sPM / Pk_ref,
label=f"sPM with $n_\\mathrm{{steps}}={nsteps_p3m}$",
linestyle="--",
color=cols[4],
)
ax.plot(
k,
Pk_P3M / Pk_ref,
label=f"P3M with $n_\\mathrm{{steps}}={nsteps_p3m}$",
linestyle="--",
color=cols[5],
)
ax.axhspan(1 - dark_grey_bnd, 1 + dark_grey_bnd, color="grey", alpha=0.2)
ax.axhspan(1 - light_grey_bnd, 1 + light_grey_bnd, color="grey", alpha=0.1)
for i in range(1, len(k)):
ax.axvline(k[i], color="black", linestyle=":", linewidth=1, alpha=0.1)
ax.axvline(np.pi * N / L, color="grey", linestyle="--", linewidth=1) # 1D Nyquist
ax.yaxis.set_major_locator(plt.MaxNLocator(6))
ax.yaxis.get_major_ticks()[0].label1.set_visible(False)
ax.set_xlabel("$k$ [$h/\\mathrm{Mpc}$]", fontsize=fs)
ax.set_ylabel("$P(k)/P_\\mathrm{ref}(k)$", fontsize=fs)
ax.tick_params(which="both", direction="in")
ax.tick_params(axis="both", which="major", labelsize=fs)
ax.tick_params(axis="both", which="minor", labelsize=fs)
empty_patch = mpatches.Patch(color="none", label="")
handles, labels = plt.gca().get_legend_handles_labels()
# handles = [empty_patch, *handles]
# labels = ["", *labels]
plt.legend(
handles,
labels,
loc="upper center",
ncol=2,
bbox_to_anchor=(0.5, -0.2),
fontsize=fs,
)
fig.savefig(
simdir + "power_spectrum.png",
bbox_inches="tight",
dpi=300,
transparent=True,
)
fig.savefig(
simdir + "power_spectrum.pdf",
bbox_inches="tight",
dpi=300,
)
plt.show()
Force exerted by particles on other particles¶
In [17]:
r1, fmag1, _ = load_force_diagnostic(OutputForceDiagnostic_pm)
r4, fmag4, _ = load_force_diagnostic(OutputForceDiagnostic_cola)
r2, fmag2, _ = load_force_diagnostic(OutputForceDiagnostic_spm)
r3, fmag3, _ = load_force_diagnostic(OutputForceDiagnostic_p3m)
rr = np.array([r1, r4, r2, r3], dtype=object)
ff = np.array([fmag1, fmag4, fmag2, fmag3], dtype=object)
ll = np.array(["PM", "COLA", "sPM", "P3M"])
ix = [0, 1, 2, 3]
Newton_prefactor = (L / Np)**3 / (4*np.pi)
print(f"Newton prefactor = {Newton_prefactor:.2e}")
plot_force_distance_comparison(rr=rr[ix], ff=ff[ix], ll=ll[ix], L=L, Np=Np, Npm=Npm, a=Newton_prefactor, title="Particle's contributions to total force", save_path=simdir + "force_diagnostic_comparison.png")
In [ ]: