mirror of
https://github.com/DifferentiableUniverseInitiative/JaxPM.git
synced 2025-02-23 18:10:55 +00:00
330 KiB
330 KiB
In [1]:
import numpy as np
import matplotlib.pyplot as plt
from h5py import File
%matplotlib inline
from jax import numpy as jnp
import jax
from tqdm import tqdm
from numpyro.infer import MCMC, NUTS, util, init_to_sample
import csiborgtools
import corner
%load_ext autoreload
%autoreload 2
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
LOS density & radial velocity plots¶
In [14]:
# fpath = "/mnt/extraspace/rstiskalek/catalogs/A2.h5"
fpath = "/mnt/extraspace/rstiskalek/catalogs/PV_compilation_Supranta2019.hdf5"
loader_carrick = csiborgtools.flow.DataLoader("Carrick2015", "LOSS", fpath, paths, ksmooth=0)
loader_csiborg = csiborgtools.flow.DataLoader("csiborg1", "LOSS", fpath, paths, ksmooth=0)
loader_csiborg2 = csiborgtools.flow.DataLoader("csiborg2_main", "LOSS", fpath, paths, ksmooth=0)
In [19]:
# ks = [115, 53, 77, 105, 26, 61, 86, 29, 80, 21]
ks = [19, 8, 15, 0, 16, 6, 48, 38, 26, 44]
# ks = [19]
# ks = np.random.choice(50, 10, replace=False)
# k = 6
for k in ks:
fig, axs = plt.subplots(2, 1, figsize=(7, 7), sharex=True)
# Get rid of vertical spacing
fig.subplots_adjust(wspace=0)
# Plot CSiBORG
for i in range(loader_csiborg.los_density.shape[1]):
axs[0].plot(loader_csiborg.rdist, loader_csiborg.los_density[k, i, :], alpha=0.1, color="black")
axs[1].plot(loader_csiborg.rdist, loader_csiborg.los_radial_velocity[k, i, :], alpha=0.1, color="black")
# CSiBORG1
axs[0].plot(loader_csiborg.rdist, loader_csiborg.los_density[k, :, :].mean(axis=0), color="red", label="CSiBORG1")
axs[1].plot(loader_csiborg.rdist, loader_csiborg.los_radial_velocity[k, :, :].mean(axis=0), color="red")
# CSiBORG2
axs[0].plot(loader_csiborg2.rdist, loader_csiborg2.los_density[k, :, :].mean(axis=0), color="violet", label="CSiBORG2")
axs[1].plot(loader_csiborg2.rdist, loader_csiborg2.los_radial_velocity[k, :, :].mean(axis=0), color="violet")
# Plot Carrick+2015
axs[0].plot(loader_carrick.rdist, loader_carrick.los_density[k, 0, :], color="blue", label="Carrick+2015")
axs[1].plot(loader_carrick.rdist, loader_carrick.los_radial_velocity[k, 0, :] * 0.43, color="blue")
# for i in range(2):
# label = "SN"
# rdist = loader_csiborg.cat["r_hMpc"][k]
# axs[i].axvline(rdist, color="violet", linestyle="--",
# zorder=0, label=label)
axs[1].set_xlabel(r"$r ~ [\mathrm{Mpc} / h]$")
axs[0].set_ylabel(r"$\rho_{\rm LOS} / \langle \rho_{\rm matter} \rangle$")
axs[1].set_ylabel(r"$v_{\rm LOS} ~ [\mathrm{km/s}]$")
axs[0].set_yscale("log")
axs[0].legend(loc="upper right")
axs[0].set_xlim(0, 200)
fig.tight_layout(w_pad=0, h_pad=0)
fig.savefig(f"../plots/LOSS_los_{k}.png", dpi=500, bbox_inches="tight")
fig.show()
Test running a model¶
In [15]:
fpath_data = "/mnt/extraspace/rstiskalek/catalogs/PV_compilation_Supranta2019.hdf5"
# fpath_data = "/mnt/extraspace/rstiskalek/catalogs/A2.h5"
simname = "csiborg1"
catalogue = "2MTF"
loader = csiborgtools.flow.DataLoader(simname, catalogue, fpath_data, paths, ksmooth=0)
Omega_m = csiborgtools.simname2Omega_m(simname)
In [16]:
print(f"Length of the catalogue is {len(loader.cat['RA'])}.")
# mask = np.random.choice(len(loader.cat["RA"]), 300, replace=False)
mask = np.ones(len(loader.cat["RA"]), dtype=bool)
# mask = loader.cat["r_hMpc"] < 1000
print(f"Selected {sum(mask)} galaxies.")
los_overdensity = loader.los_density[mask, -1, :]
los_velocity = loader.los_radial_velocity[mask, -1, :] #* 0.
# PV calibration
if catalogue in ["LOSS", "Foundation"]:
RA = loader.cat["RA"][mask]
dec = loader.cat["DEC"][mask]
zCMB = loader.cat["z_CMB"][mask]
mB = loader.cat["mB"][mask]
x1 = loader.cat["x1"][mask]
c = loader.cat["c"][mask]
e_mB = loader.cat["e_mB"][mask]
e_x1 = loader.cat["e_x1"][mask]
e_c = loader.cat["e_c"][mask]
model = csiborgtools.flow.SN_PV_validation_model(los_overdensity, los_velocity, RA, dec, zCMB, mB, x1, c, e_mB, e_x1, e_c, loader.rdist, Omega_m)
elif catalogue in ["SFI_gals", "2MTF"]:
RA = loader.cat["RA"][mask]
dec = loader.cat["DEC"][mask]
zCMB = loader.cat["z_CMB"][mask]
mag = loader.cat["mag"][mask]
eta = loader.cat["eta"][mask]
e_mag = loader.cat["e_mag"][mask]
e_eta = loader.cat["e_eta"][mask]
model = csiborgtools.flow.TF_PV_validation_model(los_overdensity, los_velocity, RA, dec, zCMB, mag, eta, e_mag, e_eta, loader.rdist, Omega_m)
elif catalogue == "A2" or "csiborg2" in catalogue:
RA = loader.cat["RA"][mask]
dec = loader.cat["DEC"][mask]
z_obs = loader.cat["z_obs"][mask]
r_hMpc = loader.cat["r_hMpc"][mask]
e_r_hMpc = loader.cat["e_rhMpc"][mask]
model = csiborgtools.flow.SD_PV_validation_model(los_overdensity, los_velocity, RA, dec, z_obs, r_hMpc, e_r_hMpc, loader.rdist, Omega_m)
else:
raise ValueError(f"Catalogue {catalogue} not recognized.")
print("RA ", RA[:20])
print("DEC ", dec[:20])
model
Out[16]:
In [19]:
true_samples = {'Vext_x': jnp.array(0.0, dtype=jnp.float32),
'Vext_y': jnp.array(0.0, dtype=jnp.float32),
'Vext_z': jnp.array(0.0, dtype=jnp.float32),
'alpha': jnp.array(0.0, dtype=jnp.float32),
'beta': jnp.array(1.0, dtype=jnp.float32),
'sigma_v': jnp.array(112, dtype=jnp.float32),
"e_mu_intrinsic": jnp.array(0.1, dtype=jnp.float32),
"mag_cal": jnp.array(-19, dtype=jnp.float32),
"alpha_cal": jnp.array(0.14, dtype=jnp.float32),
"beta_cal": jnp.array(3.1, dtype=jnp.float32),
"a" : jnp.array(-21., dtype=jnp.float32),
"b": jnp.array(-6.0, dtype=jnp.float32),
"h": jnp.array(0.7, dtype=jnp.float32),
}
# true_samples = {'beta': jnp.array(1.0, dtype=jnp.float32),}
util.log_likelihood(model, true_samples)
Out[19]:
In [20]:
nuts_kernel = NUTS(model, init_strategy=init_to_sample())
# nuts_kernel = NUTS(model)
num_chains = 1
mcmc = MCMC(nuts_kernel, num_warmup=250, num_samples=500)
rng_key = jax.random.PRNGKey(5)
In [22]:
mcmc.run(rng_key, sample_alpha=False)
In [11]:
mcmc.print_summary()
samples = mcmc.get_samples(group_by_chain=False)
In [12]:
plt.figure()
plt.hist(samples["beta"], bins="auto")
plt.show()
In [ ]:
In [566]:
Vmag = np.sqrt(samples["Vext_x"]**2 + samples["Vext_y"]**2 + samples["Vext_z"]**2)
V = np.vstack([samples["Vext_x"], samples["Vext_y"], samples["Vext_z"]]).T
V = csiborgtools.cartesian_to_radec(V)
l, b = csiborgtools.flow.radec_to_galactic(V[:, 1], V[:, 2])
In [543]:
print(f"|V| = {np.mean(Vmag)} +- {np.std(Vmag)}")
print(f"l = {np.mean(l)} +- {np.std(l)}")
print(f"b = {np.mean(b)} +- {np.std(b)}")
print(f"beta = {np.mean(samples['beta'])} +- {np.std(samples['beta'])}")
In [544]:
plt.figure()
plt.scatter(l, b)
plt.xlabel(r"$l$")
plt.ylabel(r"$b$")
plt.show()
In [545]:
plt.figure()
plt.hist(V[:, 0], bins="auto")
plt.show()
In [337]:
data = np.vstack([samples["e_mu_intrinsic"], samples["beta"], l, b, Vmag]).T
labels = [r"$\sigma_\mu$", r"$\beta$", r"$l$", r"$b$", r"$|\bf{V}_{\rm ext}|$"]
In [184]:
fig = corner.corner(data, labels=labels, show_titles=True, title_fmt=".3f", title_kwargs={"fontsize": 12})
# fig.savefig(f"../plots/corner.png", dpi=300, bbox_inches="tight")
Vizualize the results¶
In [17]:
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
catalogue = "LOSS"
simname = "csiborg2_random"
nsims = paths.get_ics(simname)
Vx = []
Vy = []
Vz = []
beta = []
sigma_v = []
alpha = []
alpha_cal = []
beta_cal = []
mag_cal = []
e_mu_intrinsic = []
with File(f"/mnt/extraspace/rstiskalek/csiborg_postprocessing/peculiar_velocity/flow_samples_{catalogue}_{simname}_smooth_0.hdf5", 'r') as f:
for i, nsim in enumerate(tqdm(nsims)):
if i == 0:
print(f[f"sim_{nsim}"].keys())
Vx.append(f[f"sim_{nsim}/Vext_x"][:])
Vy.append(f[f"sim_{nsim}/Vext_y"][:])
Vz.append(f[f"sim_{nsim}/Vext_z"][:])
alpha.append(f[f"sim_{nsim}/alpha"][:])
beta.append(f[f"sim_{nsim}/beta"][:])
sigma_v.append(f[f"sim_{nsim}/sigma_v"][:])
alpha_cal.append(f[f"sim_{nsim}/alpha_cal"][:])
beta_cal.append(f[f"sim_{nsim}/beta_cal"][:])
mag_cal.append(f[f"sim_{nsim}/mag_cal"][:])
e_mu_intrinsic.append(f[f"sim_{nsim}/e_mu_intrinsic"][:])
Vx = np.hstack(Vx)
Vy = np.hstack(Vy)
Vz = np.hstack(Vz)
alpha = np.hstack(alpha)
beta = np.hstack(beta)
sigma_v = np.hstack(sigma_v)
alpha_cal = np.hstack(alpha_cal)
beta_cal = np.hstack(beta_cal)
mag_cal = np.hstack(mag_cal)
e_mu_intrinsic = np.hstack(e_mu_intrinsic)
Vmag = np.sqrt(Vx**2 + Vy**2 + Vz**2)
V = np.vstack([Vx, Vy, Vz]).T
V = csiborgtools.cartesian_to_radec(V)
l, b = csiborgtools.flow.radec_to_galactic(V[:, 1], V[:, 2])
data = np.vstack([alpha, beta, Vmag, l, b, sigma_v, alpha_cal, beta_cal, mag_cal, e_mu_intrinsic]).T
labels = [r"$\alpha$", r"$\beta$", r"$|\bf{V}_{\rm ext}|$", r"$l$", r"$b$", r"$\sigma_v$", r"$\alpha_{\rm cal}$", r"$\beta_{\rm cal}$", r"$M$", r"$\sigma_\mu$"]
In [18]:
fig = corner.corner(data, labels=labels, show_titles=True, title_fmt=".3f", title_kwargs={"fontsize": 12})
fig.suptitle(f"{catalogue}, {simname}")
fig.savefig(f"../plots/corner_{catalogue}_{simname}.png", dpi=300, bbox_inches="tight")
In [210]:
np.mean(l)
Out[210]:
In [299]:
0.906 * 0.43
Out[299]:
In [214]:
plt.figure()
plt.hist(V[:, 0], bins="auto", density=True)
plt.show()
In [219]:
print(beta.mean(), beta.std())
print(np.mean(V[:, 0]), np.std(V[:, 0]))
print(l.mean(), b.mean())
In [213]:
plt.figure()
plt.hexbin(l, b, gridsize=50, mincnt=1)
plt.xlabel(r"$l$")
plt.ylabel(r"$b$")
plt.show()
In [ ]:
In [130]:
plt.figure()
for i in range(len(x)):
plt.hist(x[i], bins="auto", density=True, histtype="step")
plt.show()
Test bias¶
In [437]:
from jax.lax import cond
In [488]:
def bias(rho_norm, b):
def bias_small(rho_norm, b):
return jnp.log((1 + jnp.exp((1 - b + b * rho_norm))) / (1 + jnp.exp((1 - b))))
def bias_large(rho_norm, b):
return 1 - b + b * rho_norm - jnp.log(1 + jnp.exp(1 - b))
return jnp.where(rho_norm < 3,
bias_small(rho_norm, b),
bias_large(rho_norm, b))
In [489]:
rho_norm = jnp.linspace(0, 100, 100)
b = 2
y = bias(rho_norm, b)
In [490]:
y
Out[490]:
In [485]:
bias_small(rho_norm, b)
Out[485]:
In [487]:
plt.figure()
plt.plot(rho_norm, y)
plt.plot(rho_norm, bias_small(rho_norm, b), ls="--")
plt.plot(rho_norm, bias_large(rho_norm, b), ls="dotted")
plt.show()
In [ ]:
In [462]:
bias(loader._los_density[0], 2.)
Out[462]:
In [ ]:
In [ ]:
In [436]:
b = 1.4
rho = np.linspace(0, 3, 100)
def bias(rho, b):
return np.log((1 + np.exp((1 - b + b * rho))) / (1 + np.exp((1 - b))))
def linear_bias(rho, b):
return (1 - b + b * rho)
plt.figure()
cols = plt.rcParams['axes.prop_cycle'].by_key()['color']
for i, b in enumerate([0.5, 1, 1.5, 3]):
plt.plot(rho, bias(rho, b), label=f"b = {b}", c=cols[i])
plt.plot(rho, 1 - b + b * rho - np.log(1 + np.exp(1 - b)), ls="--", c=cols[i])
# plt.plot(rho, linear_bias(rho, b), ls="--", c=cols[i])
plt.axline((1, 1), slope=1, color="red", linestyle="dotted")
plt.axhline(0, color="black", linestyle="--")
plt.xlabel(r"$\rho / \langle \rho \rangle$")
plt.ylabel(r"$1 + \delta_g$")
plt.legend()
plt.show()
In [ ]:
In [133]:
def get_field(kind):
folder = "/mnt/extraspace/rstiskalek/catalogs"
from os.path import join
if kind == "density":
fpath = join(folder, "twompp_density_carrick2015.npy")
return np.load(fpath).astype(np.float32)
elif kind == "velocity":
fpath = join(folder, "twompp_velocity_carrick2015.npy")
field = np.load(fpath).astype(np.float32)
# Because the Carrick+2015 data is in the following form:
# "The velocities are predicted peculiar velocities in the CMB
# frame in Galactic Cartesian coordinates, generated from the
# \(\delta_g^*\) field with \(\beta^* = 0.43\) and an external
# dipole \(V_\mathrm{ext} = [89,-131,17]\) (Carrick et al Table 3)
# has already been added.""
field[0] -= 89
field[1] -= -131
field[2] -= 17
field /= 0.43
return field
In [135]:
density = get_field("density")
velocity = get_field("velocity")
In [136]:
from h5py import File
In [ ]:
In [137]:
with File("/mnt/extraspace/rstiskalek/catalogs/Carrick2015_unscaled.hdf5", 'w') as f:
f.create_dataset("density", data=density)
f.create_dataset("velocity", data=velocity)
In [ ]: