csiborgtools/notebooks/flow/quijote_Cell.ipynb
Richard Stiskalek 43ebf660ee
Update prior in VF (#135)
* Update priors

* Update submit

* Update README

* Relax width of calibration priors

* Relax width of calibration priors

* Update smooth scales

* Update settings

* Update README

* Update submit

* Add names

* Update nb

* Fix bug

* Update script

* Move files

* Small bug fix

* Add script

* Add script

* Add script

* Add more power

* Update script

* Quick fix

* Rm blank line

* Update nb

* Update nb

* Update nb
2024-07-12 15:46:45 +01:00

52 KiB

Density & velocity fields alond a LOS

In [128]:
# Copyright (C) 2024 Richard Stiskalek
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
import numpy as np
import matplotlib.pyplot as plt
import jax
from jax import numpy as jnp
from numpyro.infer import MCMC, NUTS, init_to_median

import csiborgtools


%load_ext autoreload
%autoreload 2
%matplotlib inline

paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

LOS density & radial velocity plots

In [129]:
fpath = "/mnt/extraspace/rstiskalek/catalogs/PV_compilation.hdf5"

loader_carrick = csiborgtools.flow.DataLoader("Carrick2015", [0], "Pantheon+", fpath, paths, ksmooth=0, )
# loaders_csiborg2X = [csiborgtools.flow.DataLoader("csiborg2X", i, "LOSS", fpath, paths, ksmooth=1, verbose=False) for i in range(20)]
# loaders_csiborg2 = [csiborgtools.flow.DataLoader("csiborg2_main", i, "LOSS", fpath, paths, ksmooth=1, verbose=False) for i in range(20)]

# loader_CF4 = csiborgtools.flow.DataLoader("CF4gp", [0], "LOSS", fpath, paths, ksmooth=0, )
# loader_lilow = csiborgtools.flow.DataLoader("Lilow2024", [0], "LOSS", fpath, paths, ksmooth=0, )
2024-07-17 10:04:21.726937:   reading the catalogue,
2024-07-17 10:04:21.744932:   reading the interpolated field,
2024-07-17 10:04:21.770393:   calculating the radial velocity.
/mnt/users/rstiskalek/csiborgtools/csiborgtools/flow/flow_model.py:91: UserWarning: The number of radial steps is even. Skipping the first step at 0.0 because Simpson's rule requires an odd number of steps.
  warn(f"The number of radial steps is even. Skipping the first "
In [130]:
# ks = [i for i in range(50)]
# ks = [30, 50,]
ks = np.random.choice(50, 2, replace=False)

for k in ks:
    fig, axs = plt.subplots(2, 1, figsize=(7, 7), sharex=True)
    fig.subplots_adjust(wspace=0)
    cols = plt.rcParams['axes.prop_cycle'].by_key()['color']

    # # CSiBORG2
    # x = loaders_csiborg2X[0].rdist
    # y = np.asarray([loaders_csiborg2[i].los_density[k, :] for i in range(len(loaders_csiborg2X))])
    # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)
    # axs[0].fill_between(x, ylow, yhigh, color=cols[0], alpha=0.25)
    # axs[0].plot(x, ymed, color=cols[0], label="CSiBORG2")

    # y = np.asarray([loaders_csiborg2[i].los_radial_velocity[k, :] for i in range(len(loaders_csiborg2X))])
    # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)
    # axs[1].fill_between(x, ylow, yhigh, color=cols[0], alpha=0.25)
    # axs[1].plot(x, ymed, color=cols[0], label="CSiBORG2")

    # # CSiBORG2X
    # x = loaders_csiborg2X[0].rdist
    # y = np.asarray([loaders_csiborg2X[i].los_density[k, :] for i in range(len(loaders_csiborg2X))])
    # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)
    # axs[0].fill_between(x, ylow, yhigh, color=cols[1], alpha=0.25)
    # axs[0].plot(x, ymed, color=cols[1], label="CSiBORG2X")

    # y = np.asarray([loaders_csiborg2X[i].los_radial_velocity[k, :] for i in range(len(loaders_csiborg2X))])
    # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)
    # axs[1].fill_between(x, ylow, yhigh, color=cols[1], alpha=0.25)
    # axs[1].plot(x, ymed, color=cols[1], label="CSiBORG2X")

    # Plot Carrick+2015
    axs[0].plot(loader_carrick.rdist, loader_carrick.los_density[0, k, :], color="red", label="Carrick+2015")
    axs[1].plot(loader_carrick.rdist, loader_carrick.los_radial_velocity[0, k, :] * 0.43, color="red")
    axs[1].axvline(loader_carrick._rmax[0, k], c="black", linestyle="--")

    # # Plot CF4
    # c = cols[4]
    # axs[0].plot(loader_CF4.rdist, loader_CF4.los_density[0, k, :], color=c, label="CF4")
    # axs[1].plot(loader_CF4.rdist, loader_CF4.los_radial_velocity[0, k, :], color=c)

    # # Plot Lilow2024
    # c = cols[5]
    # axs[0].plot(loader_lilow.rdist, loader_lilow.los_density[0, k, :], color=c, label="Lilow+2024")
    # axs[1].plot(loader_lilow.rdist, loader_lilow.los_radial_velocity[0, k, :], color=c)
    axs[1].axhline(0, color="black", linestyle="--")


    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()
No description has been provided for this image
No description has been provided for this image

Test running a model

In [84]:
fpath_data = "/mnt/extraspace/rstiskalek/catalogs/PV_compilation.hdf5"

simname = "Carrick2015"
catalogue = "LOSS"
loader = csiborgtools.flow.DataLoader(simname, [0, 0], catalogue, fpath_data, paths, ksmooth=0, )

SN_hyperparams = {"e_mu_mean": 0.1, "e_mu_std": 0.05,
                  "mag_cal_mean": -18.25, "mag_cal_std": 0.5,
                  "alpha_cal_mean": 0.148, "alpha_cal_std": 0.05,
                  "beta_cal_mean": 3.112, "beta_cal_std": 1.0,
                  }
calibration_hyperparams = {"Vext_std": 250,
                           "alpha_mean": 1.0, "alpha_std": 0.5,
                           "beta_mean": 1.0, "beta_std": 0.5,
                           "sigma_v_mean": 150., "sigma_v_std": 100.,
                           "sample_alpha": True, "sample_beta": True,
                           }
get_model_kwargs = {"zcmb_max": 0.05}
2024-06-29 19:40:25.229961:   reading the catalogue,
2024-06-29 19:40:25.243502:   reading the interpolated field,
2024-06-29 19:40:25.261423:   calculating the radial velocity.
/mnt/users/rstiskalek/csiborgtools/csiborgtools/flow/flow_model.py:91: UserWarning: The number of radial steps is even. Skipping the first step at 0.0 because Simpson's rule requires an odd number of steps.
  warn(f"The number of radial steps is even. Skipping the first "

Running HMC

In [85]:
model = csiborgtools.flow.get_model(loader, **get_model_kwargs)
model_kwargs = {"distmod_hyperparams": SN_hyperparams, "calibration_hyperparams": calibration_hyperparams,}
Selected 50/50 galaxies.
In [86]:
kernel = NUTS(model, init_strategy=init_to_median(num_samples=100))
mcmc = MCMC(kernel, num_warmup=500, num_samples=500)

rng_key = jax.random.PRNGKey(5)
mcmc.run(rng_key, extra_fields=("potential_energy",), **model_kwargs)
mcmc.print_summary()
samples = mcmc.get_samples()
sample: 100%|██████████| 1000/1000 [02:10<00:00,  7.68it/s, 7 steps of size 4.49e-01. acc. prob=0.90]  
                 mean       std    median      5.0%     95.0%     n_eff     r_hat
    Vext[0]     -3.71     69.92     -3.04   -123.73    103.87    469.72      1.00
    Vext[1]    -27.47     95.52    -30.48   -151.20    172.63    308.02      1.00
    Vext[2]    -59.27    131.26    -57.79   -273.64    137.55    456.29      1.00
      alpha      1.09      0.38      1.10      0.50      1.69    400.05      1.00
  alpha_cal      0.13      0.03      0.13      0.09      0.17    558.81      1.00
       beta      0.43      0.11      0.44      0.27      0.61    341.86      1.00
   beta_cal      3.54      0.18      3.54      3.23      3.81    606.77      1.00
       e_mu      0.08      0.03      0.08      0.04      0.12    330.71      1.00
    mag_cal    -18.19      0.04    -18.19    -18.25    -18.13    389.94      1.00
    sigma_v    176.93     52.05    169.93    102.74    267.56    315.30      1.00

Number of divergences: 0
In [ ]: