Plots of VF (#134)

* Add VF plots

* Update nb

* Add CMB velocity note

* rm nb

* Add option to return alllikelihood

* Add simulation weights

* Update nb

* Add bulkflow

* Update nb

* Add values of beta

* Update imports

* Update imports

* Add paths to Carrick and Lilow fiels

* Add Carrick and Lilow fields

* Add support for more fields

* Update bulkflow comp

* Update nb

* Update script
This commit is contained in:
Richard Stiskalek 2024-07-01 11:48:50 +01:00 committed by GitHub
parent 7dad6885e8
commit c6f49790bf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 1208 additions and 2680 deletions

View file

@ -666,7 +666,8 @@ class PV_validation_model(BaseFlowValidationModel):
self.Omega_m = Omega_m
self.norm = - self.ndata * jnp.log(self.num_sims)
def __call__(self, calibration_hyperparams, distmod_hyperparams):
def __call__(self, calibration_hyperparams, distmod_hyperparams,
store_ll_all=False):
"""NumPyro PV validation model."""
Vext, sigma_v, alpha, beta = sample_calibration(**calibration_hyperparams) # noqa
cz_err = jnp.sqrt(sigma_v**2 + self.e2_cz_obs)
@ -698,8 +699,11 @@ class PV_validation_model(BaseFlowValidationModel):
ptilde *= calculate_likelihood_zobs(self.z_obs, zobs, cz_err)
ll = jnp.log(simpson(ptilde, dx=self.dr, axis=-1)) - jnp.log(pnorm)
ll = jnp.sum(logsumexp(ll, axis=0)) + self.norm
if store_ll_all:
numpyro.deterministic("ll_all", ll)
ll = jnp.sum(logsumexp(ll, axis=0)) + self.norm
numpyro.deterministic("ll_values", ll)
numpyro.factor("ll", ll)

View file

@ -18,7 +18,8 @@ from .catalogue import (CSiBORG1Catalogue, CSiBORG2Catalogue,
MDPL2Catalogue, fiducial_observers) # noqa
from .snapshot import (CSiBORG1Snapshot, CSiBORG2Snapshot, QuijoteSnapshot, # noqa
CSiBORG1Field, CSiBORG2Field, CSiBORG2XField, # noqa
QuijoteField, BORG2Field, BORG1Field, TNG300_1Field) # noqa
QuijoteField, BORG2Field, BORG1Field, TNG300_1Field, # noqa
Carrick2015Field, Lilow2024Field) # noqa
from .obs import (SDSS, MCXCClusters, PlanckClusters, TwoMPPGalaxies, # noqa
TwoMPPGroups, ObservedCluster, match_array_to_no_masking, # noqa
cols_to_structured) # noqa

View file

@ -372,6 +372,24 @@ class Paths:
else:
raise ValueError(f"Unsupported CSiBORG2X field: `{kind}`.")
if simname == "Carrick2015":
basedir = "/mnt/extraspace/rstiskalek/catalogs"
if kind == "overdensity":
return join(basedir, "twompp_density_carrick2015.npy")
elif kind == "velocity":
return join(basedir, "twompp_velocity_carrick2015.npy")
else:
raise ValueError(f"Unsupported Carrick2015 field: `{kind}`.")
if simname == "Lilow2024":
basedir = "/mnt/extraspace/rstiskalek/catalogs"
if kind == "overdensity":
return join(basedir, "Lilow2024_density.npy")
elif kind == "velocity":
return join(basedir, "Lilow2024_xVelocity.npy")
else:
raise ValueError(f"Unsupported Lilow2024 field: `{kind}`.")
if MAS == "SPH" and kind in ["density", "velocity"]:
if simname == "csiborg1":
return join(self.csiborg1_srcdir, "field",

View file

@ -20,7 +20,7 @@ observed RA-dec can be mapped into the simulation box.
from abc import ABC, abstractmethod
from os.path import join
import numpy
import numpy as np
from h5py import File
from ..params import paths_glamdring, simname2boxsize, simname2Omega_m
@ -40,11 +40,11 @@ class BaseSnapshot(ABC):
flip_xz=False):
self._keep_snapshot_open = None
if not isinstance(nsim, (int, numpy.integer)):
if not isinstance(nsim, (int, np.integer)):
raise TypeError("`nsim` must be an integer")
self._nsim = int(nsim)
if not isinstance(nsnap, (int, numpy.integer)):
if not isinstance(nsnap, (int, np.integer)):
raise TypeError("`nsnap` must be an integer")
self._nsnap = int(nsnap)
@ -409,16 +409,16 @@ class CSiBORG2Snapshot(BaseSnapshot):
with File(self._snapshot_path, "r") as f:
if kind == "Masses":
npart = f["Header"].attrs["NumPart_Total"][1]
x = numpy.ones(npart, dtype=numpy.float32)
x = np.ones(npart, dtype=np.float32)
x *= f["Header"].attrs["MassTable"][1]
else:
x = f[f"PartType1/{kind}"][...]
if not high_resolution_only:
if x.ndim == 1:
x = numpy.hstack([x, f[f"PartType5/{kind}"][...]])
x = np.hstack([x, f[f"PartType5/{kind}"][...]])
else:
x = numpy.vstack([x, f[f"PartType5/{kind}"][...]])
x = np.vstack([x, f[f"PartType5/{kind}"][...]])
if self.flip_xz and kind in ["Coordinates", "Velocities"]:
x[:, [0, 2]] = x[:, [2, 0]]
@ -453,7 +453,7 @@ class CSiBORG2Snapshot(BaseSnapshot):
if i1 is not None and j1 - i1 > 0:
if kind == "Masses":
x1 = numpy.ones(j1 - i1, dtype=numpy.float32)
x1 = np.ones(j1 - i1, dtype=np.float32)
x1 *= f["Header"].attrs["MassTable"][1]
else:
x1 = f[f"PartType1/{kind}"][i1:j1]
@ -481,9 +481,9 @@ class CSiBORG2Snapshot(BaseSnapshot):
return x5
if x1.ndim > 1:
x1 = numpy.vstack([x1, x5])
x1 = np.vstack([x1, x5])
else:
x1 = numpy.hstack([x1, x5])
x1 = np.hstack([x1, x5])
return x1
@ -565,7 +565,7 @@ class BaseField(ABC):
Base class for reading fields such as density or velocity fields.
"""
def __init__(self, nsim, paths, flip_xz=False):
if isinstance(nsim, numpy.integer):
if isinstance(nsim, np.integer):
nsim = int(nsim)
if not isinstance(nsim, int):
raise TypeError(f"`nsim` must be an integer. Received `{type(nsim)}`.") # noqa
@ -644,7 +644,7 @@ class CSiBORG1Field(BaseField):
field = f["density"][:]
field /= (677.7 * 1e3 / grid)**3 # Convert to h^2 Msun / kpc^3
else:
field = numpy.load(fpath)
field = np.load(fpath)
if self.flip_xz:
field = field.T
@ -660,9 +660,9 @@ class CSiBORG1Field(BaseField):
v0 = f["p0"][:] / density
v1 = f["p1"][:] / density
v2 = f["p2"][:] / density
field = numpy.array([v0, v1, v2])
field = np.array([v0, v1, v2])
else:
field = numpy.load(fpath)
field = np.load(fpath)
if self.flip_xz:
field[0, ...] = field[0, ...].T
@ -678,7 +678,7 @@ class CSiBORG1Field(BaseField):
"for the flipped x- and z-axes.")
fpath = self.paths.field("radvel", MAS, grid, self.nsim, "csiborg1")
return numpy.load(fpath)
return np.load(fpath)
###############################################################################
@ -727,7 +727,7 @@ class CSiBORG2Field(BaseField):
field *= 1e10 # Convert to Msun / h
field /= (676.6 * 1e3 / grid)**3 # Convert to h^2 Msun / kpc^3
else:
field = numpy.load(fpath)
field = np.load(fpath)
if self.flip_xz:
field = field.T
@ -744,9 +744,9 @@ class CSiBORG2Field(BaseField):
v0 = f["p0"][:] / density
v1 = f["p1"][:] / density
v2 = f["p2"][:] / density
field = numpy.array([v0, v1, v2])
field = np.array([v0, v1, v2])
else:
field = numpy.load(fpath)
field = np.load(fpath)
if self.flip_xz:
field[0, ...] = field[0, ...].T
@ -763,7 +763,7 @@ class CSiBORG2Field(BaseField):
fpath = self.paths.field("radvel", MAS, grid, self.nsim,
f"csiborg2_{self.kind}")
return numpy.load(fpath)
return np.load(fpath)
class CSiBORG2XField(BaseField):
@ -784,7 +784,7 @@ class CSiBORG2XField(BaseField):
fpath = self.paths.field(
"overdensity", None, None, self.nsim, "csiborg2X")
with File(fpath, "r") as f:
field = f["delta_cic"][...].astype(numpy.float32)
field = f["delta_cic"][...].astype(np.float32)
return field
@ -803,7 +803,7 @@ class CSiBORG2XField(BaseField):
v0 = f["v_0"][...]
v1 = f["v_1"][...]
v2 = f["v_2"][...]
field = numpy.array([v0, v1, v2])
field = np.array([v0, v1, v2])
return field
@ -833,7 +833,7 @@ class BORG1Field(BaseField):
def overdensity_field(self):
fpath = self.paths.field(None, None, None, self.nsim, "borg1")
with File(fpath, "r") as f:
field = f["scalars/BORG_final_density"][:].astype(numpy.float32)
field = f["scalars/BORG_final_density"][:].astype(np.float32)
return field
@ -874,7 +874,7 @@ class BORG2Field(BaseField):
def overdensity_field(self):
fpath = self.paths.field(None, None, None, self.nsim, "borg2")
with File(fpath, "r") as f:
field = f["scalars/BORG_final_density"][:].astype(numpy.float32)
field = f["scalars/BORG_final_density"][:].astype(np.float32)
return field
@ -894,7 +894,7 @@ class BORG2Field(BaseField):
###############################################################################
# TNG300-1 field #
# Additional fields #
###############################################################################
class TNG300_1Field(BaseField):
@ -922,7 +922,7 @@ class TNG300_1Field(BaseField):
def density_field(self, MAS, grid):
fpath = join(self.paths.tng300_1(), "postprocessing", "density_field",
f"rho_dm_099_{grid}_{MAS}.npy")
return numpy.load(fpath)
return np.load(fpath)
def velocity_field(self, MAS, grid):
raise RuntimeError("The velocity field is not available.")
@ -930,10 +930,6 @@ class TNG300_1Field(BaseField):
def radial_velocity_field(self, MAS, grid):
raise RuntimeError("The radial velocity field is not available.")
###############################################################################
# Quijote field class #
###############################################################################
class QuijoteField(CSiBORG1Field):
"""
@ -951,6 +947,90 @@ class QuijoteField(CSiBORG1Field):
self._simname = "quijote"
class Carrick2015Field(BaseField):
"""
Carrick+2015 `z = 0` field class. The field is in galactic coordinates.
Parameters
----------
paths : Paths, optional
Paths object. By default, the paths are set to the `glamdring` paths.
"""
def __init__(self, paths=None):
super().__init__(0, paths, False)
def overdensity_field(self, **kwargs):
fpath = self.paths.field(
"overdensity", None, None, self.nsim, "Carrick2015")
return np.load(fpath)
def density_field(self, **kwargs):
field = self.overdensity_field()
omega0 = simname2Omega_m("Carrick2015")
rho_mean = omega0 * 277.53662724583074 # Msun / kpc^3
field += 1
field *= rho_mean
return field
def velocity_field(self, **kwargs):
fpath = self.paths.field(
"velocity", None, None, self.nsim, "Carrick2015")
field = np.load(fpath)
# 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
def radial_velocity_field(self, **kwargs):
raise RuntimeError("The radial velocity field is not available.")
class Lilow2024Field(BaseField):
"""
Lilow+2024 `z = 0` field class. The field is in galactic coordinates.
Parameters
----------
paths : Paths, optional
Paths object. By default, the paths are set to the `glamdring` paths.
"""
def __init__(self, paths=None):
super().__init__(0, paths, False)
def overdensity_field(self, **kwargs):
fpath = self.paths.field(
"overdensity", None, None, self.nsim, "Lilow2024")
return np.load(fpath) - 1
def density_field(self, **kwargs):
field = self.overdensity_field()
omega0 = simname2Omega_m("Lilow2024")
rho_mean = omega0 * 277.53662724583074 # Msun / kpc^3
field += 1
field *= rho_mean
return field
def velocity_field(self, **kwargs):
fpath = self.paths.field(
"velocity", None, None, self.nsim, "Lilow2024")
vx = np.load(fpath)
vy = np.load(fpath.replace("xVelocity", "yVelocity"))
vz = np.load(fpath.replace("xVelocity", "zVelocity"))
return np.stack([vx, vy, vz])
def radial_velocity_field(self, **kwargs):
raise RuntimeError("The radial velocity field is not available.")
###############################################################################
# Supplementary functions #
###############################################################################

File diff suppressed because one or more lines are too long

View file

@ -1,194 +0,0 @@
# 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.
"""Script to help with plots in `flow_calibration.ipynb`."""
from copy import copy
from os.path import join, exists
import numpy as np
from getdist import MCSamples
from h5py import File
import csiborgtools
def read_samples(catalogue, simname, ksmooth, include_calibration=False,
return_MCsamples=False, subtract_LG_velocity=-1):
print(f"\nReading {catalogue} fitted to {simname} with ksmooth = {ksmooth}.", flush=True) # noqa
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsims = paths.get_ics(simname)
# The last simulation was used to draw the mocks.
if catalogue in ["CB2_small", "CB2_large"]:
nsims = nsims[:-5]
FDIR_LG = "/mnt/extraspace/rstiskalek/csiborg_postprocessing/peculiar_velocity/observer" # noqa
Vx, Vy, Vz, beta, sigma_v, alpha = [], [], [], [], [], []
BIC, AIC, logZ, chi2 = [], [], [], []
if catalogue in ["LOSS", "Foundation"] or "Pantheon+" in catalogue:
alpha_cal, beta_cal, mag_cal, e_mu_intrinsic = [], [], [], []
elif catalogue in ["2MTF", "SFI_gals", "SFI_gals_masked"]:
a, b, e_mu_intrinsic = [], [], []
elif catalogue == "SFI_groups":
h = []
elif catalogue in ["CB2_small", "CB2_large"]:
pass
else:
raise ValueError(f"Catalogue {catalogue} not recognized.")
fname = f"/mnt/extraspace/rstiskalek/csiborg_postprocessing/peculiar_velocity/flow_samples_{catalogue}_{simname}_smooth_{ksmooth}.hdf5" # noqa
with File(fname, 'r') as f:
for i, nsim in enumerate(nsims):
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"][:])
if subtract_LG_velocity >= 0:
fname = join(FDIR_LG, f"{simname}_{nsim}_observer_velocity.npz") # noqa
if not exists(fname):
raise FileNotFoundError(f"File {fname} not found.")
d = np.load(fname)
R = d["smooth_scales"][subtract_LG_velocity]
if i == 0:
print(f"Subtracting LG velocity with kernel {R} Mpc / h.", flush=True) # noqa
Vx_LG, Vy_LG, Vz_LG = d["vobs"][subtract_LG_velocity]
if simname == "Carrick2015":
Vx[-1] += beta[-1] * Vx_LG
Vy[-1] += beta[-1] * Vy_LG
Vz[-1] += beta[-1] * Vz_LG
else:
Vx[-1] += Vx_LG
Vy[-1] += Vy_LG
Vz[-1] += Vz_LG
BIC.append(f[f"sim_{nsim}/BIC"][...])
AIC.append(f[f"sim_{nsim}/AIC"][...])
logZ.append(f[f"sim_{nsim}/logZ"][...])
try:
chi2.append(f[f"sim_{nsim}/chi2"][...])
except KeyError:
chi2.append([0.])
if catalogue in ["LOSS", "Foundation"] or "Pantheon+" in catalogue: # noqa
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"][:])
elif catalogue in ["2MTF", "SFI_gals"]:
a.append(f[f"sim_{nsim}/a"][:])
b.append(f[f"sim_{nsim}/b"][:])
e_mu_intrinsic.append(f[f"sim_{nsim}/e_mu_intrinsic"][:])
elif catalogue == "SFI_groups":
h.append(f[f"sim_{nsim}/h"][:])
elif catalogue in ["CB2_small", "CB2_large"]:
pass
else:
raise ValueError(f"Catalogue {catalogue} not recognized.")
Vx, Vy, Vz, alpha, beta, sigma_v = np.hstack(Vx), np.hstack(Vy), np.hstack(Vz), np.hstack(alpha), np.hstack(beta), np.hstack(sigma_v) # noqa
gof = np.hstack(BIC), np.hstack(AIC), np.hstack(logZ), np.hstack(chi2)
if catalogue in ["LOSS", "Foundation"] or "Pantheon+" in catalogue:
alpha_cal, beta_cal, mag_cal, e_mu_intrinsic = np.hstack(alpha_cal), np.hstack(beta_cal), np.hstack(mag_cal), np.hstack(e_mu_intrinsic) # noqa
elif catalogue in ["2MTF", "SFI_gals", "SFI_gals_masked"]:
a, b, e_mu_intrinsic = np.hstack(a), np.hstack(b), np.hstack(e_mu_intrinsic) # noqa
elif catalogue == "SFI_groups":
h = np.hstack(h)
elif catalogue in ["CB2_small", "CB2_large"]:
pass
else:
raise ValueError(f"Catalogue {catalogue} not recognized.")
# Calculate magnitude of V_ext
Vmag = np.sqrt(Vx**2 + Vy**2 + Vz**2)
# Calculate direction in galactic coordinates of V_ext
V = np.vstack([Vx, Vy, Vz]).T
V = csiborgtools.cartesian_to_radec(V)
l, b = csiborgtools.radec_to_galactic(V[:, 1], V[:, 2])
data = [alpha, beta, Vmag, l, b, sigma_v]
names = ["alpha", "beta", "Vmag", "l", "b", "sigma_v"]
if include_calibration:
if catalogue in ["LOSS", "Foundation"] or "Pantheon+" in catalogue:
data += [alpha_cal, beta_cal, mag_cal, e_mu_intrinsic]
names += ["alpha_cal", "beta_cal", "mag_cal", "e_mu_intrinsic"]
elif catalogue in ["2MTF", "SFI_gals", "SFI_gals_masked"]:
data += [a, b, e_mu_intrinsic]
names += ["a", "b", "e_mu_intrinsic"]
elif catalogue == "SFI_groups":
data += [h]
names += ["h"]
else:
raise ValueError(f"Catalogue {catalogue} not recognized.")
print("BIC = {:4f} +- {:4f}".format(np.mean(gof[0]), np.std(gof[0])))
print("AIC = {:4f} +- {:4f}".format(np.mean(gof[1]), np.std(gof[1])))
print("logZ = {:4f} +- {:4f}".format(np.mean(gof[2]), np.std(gof[2])))
print("chi2 = {:4f} +- {:4f}".format(np.mean(gof[3]), np.std(gof[3])))
data = np.vstack(data).T
if return_MCsamples:
simname = simname_to_pretty(simname)
if ksmooth == 1:
simname = fr"{simname} (2)"
if subtract_LG_velocity >= 0:
simname += " (LG)"
label = fr"{catalogue}, {simname}, $\log \mathcal{{Z}} = {np.mean(gof[2]):.1f}$" # noqa
return MCSamples(samples=data, names=names,
labels=names_to_latex(names), label=label)
return data, names, gof
def simname_to_pretty(simname):
ltx = {"Carrick2015": "C+15",
"csiborg1": "CB1",
"csiborg2_main": "CB2",
}
return ltx[simname] if simname in ltx else simname
def names_to_latex(names, for_corner=False):
ltx = {"alpha": "\\alpha",
"beta": "\\beta",
"Vmag": "V_{\\rm ext} ~ [\\mathrm{km} / \\mathrm{s}]",
"sigma_v": "\\sigma_v ~ [\\mathrm{km} / \\mathrm{s}]",
}
ltx_corner = {"alpha": r"$\alpha$",
"beta": r"$\beta$",
"Vmag": r"$V_{\rm ext}$",
"sigma_v": r"$\sigma_v$",
"h": r"$h$",
}
labels = copy(names)
for i, label in enumerate(names):
if label in ltx:
labels[i] = ltx_corner[label] if for_corner else ltx[label]
return labels

View file

@ -0,0 +1,293 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Density & velocity fields alond a LOS"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"The autoreload extension is already loaded. To reload it, use:\n",
" %reload_ext autoreload\n"
]
}
],
"source": [
"# Copyright (C) 2024 Richard Stiskalek\n",
"# This program is free software; you can redistribute it and/or modify it\n",
"# under the terms of the GNU General Public License as published by the\n",
"# Free Software Foundation; either version 3 of the License, or (at your\n",
"# option) any later version.\n",
"#\n",
"# This program is distributed in the hope that it will be useful, but\n",
"# WITHOUT ANY WARRANTY; without even the implied warranty of\n",
"# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General\n",
"# Public License for more details.\n",
"#\n",
"# You should have received a copy of the GNU General Public License along\n",
"# with this program; if not, write to the Free Software Foundation, Inc.,\n",
"# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.\n",
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
"import jax\n",
"from jax import numpy as jnp\n",
"from numpyro.infer import MCMC, NUTS, init_to_median\n",
"\n",
"import csiborgtools\n",
"\n",
"\n",
"%load_ext autoreload\n",
"%autoreload 2\n",
"%matplotlib inline\n",
"\n",
"paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## LOS density & radial velocity plots "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"fpath = \"/mnt/extraspace/rstiskalek/catalogs/PV_compilation.hdf5\"\n",
"\n",
"loader_carrick = csiborgtools.flow.DataLoader(\"Carrick2015\", [0], \"LOSS\", fpath, paths, ksmooth=0, )\n",
"# loaders_csiborg2X = [csiborgtools.flow.DataLoader(\"csiborg2X\", i, \"LOSS\", fpath, paths, ksmooth=1, verbose=False) for i in range(20)]\n",
"# loaders_csiborg2 = [csiborgtools.flow.DataLoader(\"csiborg2_main\", i, \"LOSS\", fpath, paths, ksmooth=1, verbose=False) for i in range(20)]\n",
"\n",
"loader_CF4 = csiborgtools.flow.DataLoader(\"CF4gp\", [0], \"LOSS\", fpath, paths, ksmooth=0, )\n",
"loader_lilow = csiborgtools.flow.DataLoader(\"Lilow2024\", [0], \"LOSS\", fpath, paths, ksmooth=0, )"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# ks = [i for i in range(50)]\n",
"ks = [0,]\n",
"\n",
"for k in ks:\n",
" fig, axs = plt.subplots(2, 1, figsize=(7, 7), sharex=True)\n",
" fig.subplots_adjust(wspace=0)\n",
" cols = plt.rcParams['axes.prop_cycle'].by_key()['color']\n",
"\n",
" # # CSiBORG2\n",
" # x = loaders_csiborg2X[0].rdist\n",
" # y = np.asarray([loaders_csiborg2[i].los_density[k, :] for i in range(len(loaders_csiborg2X))])\n",
" # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)\n",
" # axs[0].fill_between(x, ylow, yhigh, color=cols[0], alpha=0.25)\n",
" # axs[0].plot(x, ymed, color=cols[0], label=\"CSiBORG2\")\n",
"\n",
" # y = np.asarray([loaders_csiborg2[i].los_radial_velocity[k, :] for i in range(len(loaders_csiborg2X))])\n",
" # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)\n",
" # axs[1].fill_between(x, ylow, yhigh, color=cols[0], alpha=0.25)\n",
" # axs[1].plot(x, ymed, color=cols[0], label=\"CSiBORG2\")\n",
"\n",
" # # CSiBORG2X\n",
" # x = loaders_csiborg2X[0].rdist\n",
" # y = np.asarray([loaders_csiborg2X[i].los_density[k, :] for i in range(len(loaders_csiborg2X))])\n",
" # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)\n",
" # axs[0].fill_between(x, ylow, yhigh, color=cols[1], alpha=0.25)\n",
" # axs[0].plot(x, ymed, color=cols[1], label=\"CSiBORG2X\")\n",
"\n",
" # y = np.asarray([loaders_csiborg2X[i].los_radial_velocity[k, :] for i in range(len(loaders_csiborg2X))])\n",
" # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)\n",
" # axs[1].fill_between(x, ylow, yhigh, color=cols[1], alpha=0.25)\n",
" # axs[1].plot(x, ymed, color=cols[1], label=\"CSiBORG2X\")\n",
"\n",
" # Plot Carrick+2015\n",
" axs[0].plot(loader_carrick.rdist, loader_carrick.los_density[0, k, :], color=\"red\", label=\"Carrick+2015\")\n",
" axs[1].plot(loader_carrick.rdist, loader_carrick.los_radial_velocity[0, k, :] * 0.43, color=\"red\")\n",
"\n",
" # Plot CF4\n",
" c = cols[4]\n",
" axs[0].plot(loader_CF4.rdist, loader_CF4.los_density[0, k, :], color=c, label=\"CF4\")\n",
" axs[1].plot(loader_CF4.rdist, loader_CF4.los_radial_velocity[0, k, :], color=c)\n",
"\n",
" # Plot Lilow2024\n",
" c = cols[5]\n",
" axs[0].plot(loader_lilow.rdist, loader_lilow.los_density[0, k, :], color=c, label=\"Lilow+2024\")\n",
" axs[1].plot(loader_lilow.rdist, loader_lilow.los_radial_velocity[0, k, :], color=c)\n",
"\n",
"\n",
" axs[1].set_xlabel(r\"$r ~ [\\mathrm{Mpc} / h]$\")\n",
" axs[0].set_ylabel(r\"$\\rho_{\\rm LOS} / \\langle \\rho_{\\rm matter} \\rangle$\")\n",
" axs[1].set_ylabel(r\"$v_{\\rm LOS} ~ [\\mathrm{km/s}]$\")\n",
" axs[0].set_yscale(\"log\")\n",
"\n",
" axs[0].legend(loc=\"upper right\")\n",
" axs[0].set_xlim(0, 200)\n",
"\n",
" fig.tight_layout(w_pad=0, h_pad=0)\n",
" fig.savefig(f\"../../plots/LOSS_los_{k}.png\", dpi=500, bbox_inches=\"tight\")\n",
"\n",
" fig.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Test running a model"
]
},
{
"cell_type": "code",
"execution_count": 84,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"2024-06-29 19:40:25.229961: reading the catalogue,\n",
"2024-06-29 19:40:25.243502: reading the interpolated field,\n",
"2024-06-29 19:40:25.261423: calculating the radial velocity.\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"/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.\n",
" warn(f\"The number of radial steps is even. Skipping the first \"\n"
]
}
],
"source": [
"fpath_data = \"/mnt/extraspace/rstiskalek/catalogs/PV_compilation.hdf5\"\n",
"\n",
"simname = \"Carrick2015\"\n",
"catalogue = \"LOSS\"\n",
"loader = csiborgtools.flow.DataLoader(simname, [0, 0], catalogue, fpath_data, paths, ksmooth=0, )\n",
"\n",
"SN_hyperparams = {\"e_mu_mean\": 0.1, \"e_mu_std\": 0.05,\n",
" \"mag_cal_mean\": -18.25, \"mag_cal_std\": 0.5,\n",
" \"alpha_cal_mean\": 0.148, \"alpha_cal_std\": 0.05,\n",
" \"beta_cal_mean\": 3.112, \"beta_cal_std\": 1.0,\n",
" }\n",
"calibration_hyperparams = {\"Vext_std\": 250,\n",
" \"alpha_mean\": 1.0, \"alpha_std\": 0.5,\n",
" \"beta_mean\": 1.0, \"beta_std\": 0.5,\n",
" \"sigma_v_mean\": 150., \"sigma_v_std\": 100.,\n",
" \"sample_alpha\": True, \"sample_beta\": True,\n",
" }\n",
"get_model_kwargs = {\"zcmb_max\": 0.05}"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Running HMC"
]
},
{
"cell_type": "code",
"execution_count": 85,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Selected 50/50 galaxies.\n"
]
}
],
"source": [
"model = csiborgtools.flow.get_model(loader, **get_model_kwargs)\n",
"model_kwargs = {\"distmod_hyperparams\": SN_hyperparams, \"calibration_hyperparams\": calibration_hyperparams,}"
]
},
{
"cell_type": "code",
"execution_count": 86,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"sample: 100%|██████████| 1000/1000 [02:10<00:00, 7.68it/s, 7 steps of size 4.49e-01. acc. prob=0.90] \n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" mean std median 5.0% 95.0% n_eff r_hat\n",
" Vext[0] -3.71 69.92 -3.04 -123.73 103.87 469.72 1.00\n",
" Vext[1] -27.47 95.52 -30.48 -151.20 172.63 308.02 1.00\n",
" Vext[2] -59.27 131.26 -57.79 -273.64 137.55 456.29 1.00\n",
" alpha 1.09 0.38 1.10 0.50 1.69 400.05 1.00\n",
" alpha_cal 0.13 0.03 0.13 0.09 0.17 558.81 1.00\n",
" beta 0.43 0.11 0.44 0.27 0.61 341.86 1.00\n",
" beta_cal 3.54 0.18 3.54 3.23 3.81 606.77 1.00\n",
" e_mu 0.08 0.03 0.08 0.04 0.12 330.71 1.00\n",
" mag_cal -18.19 0.04 -18.19 -18.25 -18.13 389.94 1.00\n",
" sigma_v 176.93 52.05 169.93 102.74 267.56 315.30 1.00\n",
"\n",
"Number of divergences: 0\n"
]
}
],
"source": [
"kernel = NUTS(model, init_strategy=init_to_median(num_samples=100))\n",
"mcmc = MCMC(kernel, num_warmup=500, num_samples=500)\n",
"\n",
"rng_key = jax.random.PRNGKey(5)\n",
"mcmc.run(rng_key, extra_fields=(\"potential_energy\",), **model_kwargs)\n",
"mcmc.print_summary()\n",
"samples = mcmc.get_samples()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "venv_csiborg",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.7"
}
},
"nbformat": 4,
"nbformat_minor": 2
}

File diff suppressed because one or more lines are too long

View file

@ -0,0 +1,230 @@
# 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.
"""Script to help with plots in `flow_calibration.ipynb`."""
from copy import copy
from os.path import join
import numpy as np
from jax import numpy as jnp
from getdist import MCSamples
from h5py import File
import csiborgtools
###############################################################################
# Convert between coordinate systems #
###############################################################################
def cartesian_to_radec(x, y, z):
d = (x**2 + y**2 + z**2)**0.5
dec = np.arcsin(z / d)
ra = np.arctan2(y, x)
ra[ra < 0] += 2 * np.pi
ra *= 180 / np.pi
dec *= 180 / np.pi
return d, ra, dec
###############################################################################
# Get the filename of the samples #
###############################################################################
def get_fname(simname, catalogue, ksmooth=0, nsim=None, sample_beta=True):
"""Get the filename of the HDF5 file containing the posterior samples."""
FDIR = "/mnt/extraspace/rstiskalek/csiborg_postprocessing/peculiar_velocity/" # noqa
fname = join(FDIR, f"samples_{simname}_{catalogue}_ksmooth{ksmooth}.hdf5")
if nsim is not None:
fname = fname.replace(".hdf5", f"_nsim{nsim}.hdf5")
if sample_beta:
fname = fname.replace(".hdf5", "_sample_beta.hdf5")
return fname
###############################################################################
# Convert names to LaTeX #
###############################################################################
def names_to_latex(names, for_corner=False):
"""Convert the names of the parameters to LaTeX."""
ltx = {"alpha": "\\alpha",
"beta": "\\beta",
"Vmag": "V_{\\rm ext}",
"sigma_v": "\\sigma_v",
"alpha_cal": "\\mathcal{A}",
"beta_cal": "\\mathcal{B}",
"mag_cal": "\\mathcal{M}",
"e_mu": "\\sigma_\\mu",
"aTF": "a_{\\rm TF}",
"bTF": "b_{\\rm TF}",
}
ltx_corner = {"alpha": r"$\alpha$",
"beta": r"$\beta$",
"Vmag": r"$V_{\rm ext}$",
"l": r"$\ell_{V_{\rm ext}}$",
"b": r"$b_{V_{\rm ext}}$",
"sigma_v": r"$\sigma_v$",
"alpha_cal": r"$\mathcal{A}$",
"beta_cal": r"$\mathcal{B}$",
"mag_cal": r"$\mathcal{M}$",
"e_mu": r"$\sigma_\mu$",
"aTF": r"$a_{\rm TF}$",
"bTF": r"$b_{\rm TF}$",
}
labels = copy(names)
for i, label in enumerate(names):
if for_corner:
labels[i] = ltx_corner[label] if label in ltx_corner else label
else:
labels[i] = ltx[label] if label in ltx else label
return labels
def simname_to_pretty(simname):
ltx = {"Carrick2015": "Carrick+15",
"Lilow2024": "Lilow+24",
"csiborg1": "CB1",
"csiborg2_main": "CB2",
"csiborg2X": "Manticore",
}
if isinstance(simname, list):
return [ltx[s] if s in ltx else s for s in simname]
return ltx[simname] if simname in ltx else simname
###############################################################################
# Read in goodness-of-fit #
###############################################################################
def get_gof(kind, simname, catalogue, ksmooth=0, nsim=None, sample_beta=True):
"""Read in the goodness-of-fit statistics `kind`."""
if kind not in ["BIC", "AIC", "lnZ"]:
raise ValueError("`kind` must be one of 'BIC', 'AIC', 'lnZ'")
fname = get_fname(simname, catalogue, ksmooth, nsim, sample_beta)
with File(fname, 'r') as f:
return f[f"gof/{kind}"][()]
###############################################################################
# Read in samples #
###############################################################################
def get_samples(simname, catalogue, ksmooth=0, nsim=None, sample_beta=True,
convert_Vext_to_galactic=True):
"""Read in the samples from the HDF5 file."""
fname = get_fname(simname, catalogue, ksmooth, nsim, sample_beta)
samples = {}
with File(fname, 'r') as f:
grp = f["samples"]
for key in grp.keys():
samples[key] = grp[key][...]
# Rename TF parameters
if "a" in samples:
samples["aTF"] = samples.pop("a")
if "b" in samples:
samples["bTF"] = samples.pop("b")
if convert_Vext_to_galactic:
Vext = samples.pop("Vext")
samples["Vmag"] = np.linalg.norm(Vext, axis=1)
Vext = csiborgtools.cartesian_to_radec(Vext)
samples["l"], samples["b"] = csiborgtools.radec_to_galactic(
Vext[:, 1], Vext[:, 2])
return samples
###############################################################################
def get_bulkflow(simname, catalogue, ksmooth=0, nsim=None, sample_beta=True,
convert_to_galactic=True, simulation_only=False):
# Read in the samples
fname_samples = get_fname(simname, catalogue, ksmooth, nsim, sample_beta)
with File(fname_samples, 'r') as f:
simulation_weights = jnp.exp(f["simulation_weights"][...])
Vext = f["samples/Vext"][...]
try:
beta = f["samples/beta"][...]
except KeyError:
beta = jnp.ones_like(simulation_weights)
# Read in the bulk flow
f = np.load(f"/mnt/extraspace/rstiskalek/csiborg_postprocessing/field_shells/enclosed_mass_{simname}.npz") # noqa
r, Bsim = f["distances"], f["cumulative_velocity"]
# Mask out the unconstrained large scales
Rmax = 150 # Mpc/h
mask = r < Rmax
r = r[mask]
Bsim = Bsim[:, mask, :]
# Add the external flow to the bulk flow and weight it
B = Bsim[:, :, None, :] * beta[None, None, :, None] # noqa
if not simulation_only:
B += Vext[None, None, :, :]
B = jnp.sum(B * simulation_weights.T[:, None, :, None], axis=0)
if convert_to_galactic:
Bmag, Bl, Bb = cartesian_to_radec(B[..., 0], B[..., 1], B[..., 2])
Bl, Bb = csiborgtools.radec_to_galactic(Bl, Bb)
return r, np.stack([Bmag, Bl, Bb], axis=-1)
return r, B
###############################################################################
# Prepare samples for plotting #
###############################################################################
def samples_for_corner(samples):
if any(x.ndim > 1 for x in samples.values()):
raise ValueError("All samples must be 1D arrays.")
data = np.vstack([x for x in samples.values()]).T
labels = names_to_latex(list(samples.keys()), for_corner=True)
return data, labels
def samples_to_getdist(samples, simname, catalogue=None):
data, __ = samples_for_corner(samples)
names = list(samples.keys())
if catalogue is None:
label = simname_to_pretty(simname)
else:
label = catalogue
return MCSamples(
samples=data, names=names,
labels=names_to_latex(names, for_corner=False),
label=label)

View file

@ -1,168 +0,0 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import scipy.integrate\n",
"import symbolic_pofk.linear\n",
"import symbolic_pofk.syrenhalofit as syrenhalofit"
]
},
{
"cell_type": "code",
"execution_count": 26,
"metadata": {},
"outputs": [],
"source": [
"def compute_sigma8_from_pk(k, pk):\n",
" \"\"\"Given a power spectrum P(k), compute sigma8.\"\"\"\n",
" R = 8.0\n",
" x = k * R\n",
" W = np.zeros(x.shape)\n",
" m = x < 1.e-3\n",
" W[m] = 1.0\n",
" W[~m] =3.0 / x[~m]**3 * (np.sin(x[~m]) - x[~m] * np.cos(x[~m]))\n",
" y = pk * W**2 * k**3\n",
" sigma2 = scipy.integrate.simpson(y, x=np.log(x))\n",
" sigma = np.sqrt(sigma2 / (2.0 * np.pi**2))\n",
"\n",
" return sigma\n",
"\n",
"# Cosmological parameters\n",
"As = 2.105 # 10^9 A_s\n",
"h = 0.6766\n",
"Om = 0.3111\n",
"Ob = 0.02242 / h ** 2\n",
"ns = 0.9665\n",
"tau = 0.0561\n",
"\n",
"# Define k integration range\n",
"\n",
"\n",
"def linear_sigma8(As, Om, Ob, h, ns):\n",
" \"\"\"Calculated from Deaglan's emulator.\"\"\"\n",
" return symbolic_pofk.linear.As_to_sigma8(As, Om, Ob, h, ns)\n",
" # print('sigma8 from emulator', sigma8)\n",
"\n",
"# # Test linear sigma8\n",
"# pk_lin = symbolic_pofk.linear.plin_emulated(k, sigma8, Om, Ob, h, ns,\n",
"# emulator='fiducial', extrapolate=True)\n",
"# new_sigma8 = compute_sigma8(k, pk_lin)\n",
"# print('sigma8 from integral:', new_sigma8)\n",
"\n",
"# Get non-linear sigma8\n",
"def nonlinear_sigma8(As, Om0, Ob, h, ns, ks):\n",
" a = 1.\n",
" sigma8 = linear_sigma8(As, Om0, Ob, h, ns) # Linear sigma8\n",
" pk_nl = syrenhalofit.run_halofit(\n",
" ks, sigma8, Om, Ob, h, ns, a, emulator='fiducial', extrapolate=True,\n",
" which_params='Bartlett', add_correction=True)\n",
" return compute_sigma8_from_pk(ks, pk_nl)\n",
"# print('non-linear sigma8:', sigma8_nl)\n",
"# print('sigma8 non-linear bigger by a factor', sigma8_nl / sigma8)\n"
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {},
"outputs": [],
"source": [
"kmin, kmax, nk = 1e-4, 1e1, 256\n",
"ks = np.logspace(np.log10(kmin), np.log10(kmax), nk) # Wavenumber"
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"0.825706107176727"
]
},
"execution_count": 21,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"linear_sigma8(As, Om, Ob, h, ns)"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"0.9203769215120938"
]
},
"execution_count": 20,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"nonlinear_sigma8(As, Om, Ob, h, ns, ks)"
]
},
{
"cell_type": "code",
"execution_count": 25,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"2.1005829616811546e-09"
]
},
"execution_count": 25,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
" 1e-10 * np.exp(3.0448)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "venv_csiborg",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.4"
}
},
"nbformat": 4,
"nbformat_minor": 2
}

View file

@ -21,16 +21,16 @@ The script is not parallelized in any way but it should not take very long, the
main bottleneck is reading the data from disk.
"""
from argparse import ArgumentParser
from os.path import join
from datetime import datetime
from gc import collect
from os.path import join
import csiborgtools
import numpy as np
from astropy import units as u
from astropy.coordinates import CartesianRepresentation, SkyCoord
from tqdm import tqdm
from datetime import datetime
###############################################################################
# Read in information about the simulation #
###############################################################################
@ -132,7 +132,7 @@ def get_particles(reader, boxsize, get_velocity=True, verbose=True):
###############################################################################
# Main & command line interface #
# Main #
###############################################################################
@ -201,7 +201,7 @@ def main_csiborg(args, folder):
cumulative_velocity=cumulative_velocity)
def main_csiborg2X(args, folder):
def main_from_field(args, folder):
"""Bulk flow in the Manticore boxes provided by Stuart."""
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
boxsize = csiborgtools.simname2boxsize(args.simname)
@ -214,7 +214,14 @@ def main_csiborg2X(args, folder):
cumulative_vel_y = np.zeros_like(cumulative_vel_x)
cumulative_vel_z = np.zeros_like(cumulative_vel_x)
for i, nsim in enumerate(tqdm(nsims, desc="Simulations")):
reader = csiborgtools.read.CSiBORG2XField(nsim, paths)
if args.simname == "csiborg2X":
reader = csiborgtools.read.CSiBORG2XField(nsim, paths)
elif args.simname == "Carrick2015":
reader = csiborgtools.read.Carrick2015Field(paths)
elif args.simname == "Lilow2024":
reader = csiborgtools.read.Lilow2024Field(paths)
else:
raise ValueError(f"Unknown simname: `{args.simname}`.")
density_field = reader.density_field()
velocity_field = reader.velocity_field()
@ -229,6 +236,20 @@ def main_csiborg2X(args, folder):
cumulative_vel_z[i, :], __ = csiborgtools.field.field_enclosed_mass(
velocity_field[2], distances, boxsize, verbose=False)
if args.simname in ["Carrick2015", "Lilow2024"]:
# Carrick+2015 and Lilow+2024 box is in galactic coordinates, so we
# need to convert the bulk flow vector to RA/dec Cartesian
# representation.
galactic_cartesian = CartesianRepresentation(
cumulative_vel_x, cumulative_vel_y, cumulative_vel_z,
unit=u.km/u.s)
galactic_coord = SkyCoord(galactic_cartesian, frame='galactic')
icrs_cartesian = galactic_coord.icrs.cartesian
cumulative_vel_x = icrs_cartesian.x.to(u.km/u.s).value
cumulative_vel_y = icrs_cartesian.y.to(u.km/u.s).value
cumulative_vel_z = icrs_cartesian.z.to(u.km/u.s).value
cumulative_vel = np.stack(
[cumulative_vel_x, cumulative_vel_y, cumulative_vel_z], axis=-1)
cumulative_vel /= cumulative_volume[..., None]
@ -236,20 +257,28 @@ def main_csiborg2X(args, folder):
# Finally save the output
fname = f"enclosed_mass_{args.simname}.npz"
fname = join(folder, fname)
print(f"Saving to `{fname}`.")
np.savez(fname, enclosed_mass=cumulative_mass, distances=distances,
cumulative_velocity=cumulative_vel,
enclosed_volume=cumulative_volume)
###############################################################################
# Command line interface #
###############################################################################
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("--simname", type=str, help="Simulation name.",
choices=["csiborg1", "csiborg2_main", "csiborg2_varysmall", "csiborg2_random", "borg1", "borg2", "borg2_all", "csiborg2X"]) # noqa
choices=["csiborg1", "csiborg2_main", "csiborg2_varysmall", "csiborg2_random", # noqa
"borg1", "borg2", "borg2_all", "csiborg2X", "Carrick2015", # noqa
"Lilow2024"]) # noqa
args = parser.parse_args()
folder = "/mnt/extraspace/rstiskalek/csiborg_postprocessing/field_shells"
if args.simname == "csiborg2X":
main_csiborg2X(args, folder)
if args.simname in ["csiborg2X", "Carrick2015", "Lilow2024"]:
main_from_field(args, folder)
elif "csiborg" in args.simname:
main_csiborg(args, folder)
elif "borg" in args.simname:

View file

@ -48,18 +48,18 @@ def parse_args():
ARGS = parse_args()
# This must be done before we import JAX etc.
from numpyro import set_host_device_count, set_platform # noqa
set_platform(ARGS.device) # noqa
set_host_device_count(ARGS.ndevice) # noqa
import sys # noqa
from os.path import join # noqa
import csiborgtools # noqa
import jax # noqa
from h5py import File # noqa
from mpi4py import MPI # noqa
from numpyro.infer import MCMC, NUTS, init_to_median # noqa
import csiborgtools # noqa
from numpyro.infer import MCMC, NUTS, Predictive, init_to_median # noqa
def print_variables(names, variables):
@ -113,6 +113,19 @@ def get_harmonic_evidence(samples, log_posterior, nchains_harmonic, epoch_num):
data, log_posterior, return_flow_samples=False, epochs_num=epoch_num)
def get_simulation_weights(samples, model, model_kwargs):
"""Get the weights per posterior samples for each simulation."""
predictive = Predictive(model, samples)
ll_all = predictive(
jax.random.PRNGKey(1), store_ll_all=True, **model_kwargs)["ll_all"]
# Multiply the likelihood of galaxies
ll_per_simulation = jax.numpy.sum(ll_all, axis=-1)
# Normalization by summing the likelihood over simulations
norm = jax.scipy.special.logsumexp(ll_per_simulation, axis=-1)
return ll_per_simulation - norm[:, None]
def run_model(model, nsteps, nburn, model_kwargs, out_folder, sample_beta,
calculate_evidence, nchains_harmonic, epoch_num, kwargs_print):
"""Run the NumPyro model and save output to a file."""
@ -128,6 +141,7 @@ def run_model(model, nsteps, nburn, model_kwargs, out_folder, sample_beta,
mcmc.run(rng_key, extra_fields=("potential_energy",), **model_kwargs)
samples = mcmc.get_samples()
simulation_weights = get_simulation_weights(samples, model, model_kwargs)
log_posterior = -mcmc.get_extra_fields()["potential_energy"]
log_likelihood = samples.pop("ll_values")
@ -141,13 +155,13 @@ def run_model(model, nsteps, nburn, model_kwargs, out_folder, sample_beta,
if calculate_evidence:
print("Calculating the evidence using `harmonic`.", flush=True)
ln_evidence, ln_evidence_err = get_harmonic_evidence(
neg_ln_evidence, neg_ln_evidence_err = get_harmonic_evidence(
samples, log_posterior, nchains_harmonic, epoch_num)
print(f"{'ln(Z)':<20} {ln_evidence}")
print(f"{'ln(Z) error':<20} {ln_evidence_err}")
print(f"{'-ln(Z)':<20} {neg_ln_evidence}")
print(f"{'-ln(Z) error':<20} {neg_ln_evidence_err}")
else:
ln_evidence = jax.numpy.nan
ln_evidence_err = (jax.numpy.nan, jax.numpy.nan)
neg_ln_evidence = jax.numpy.nan
neg_ln_evidence_err = (jax.numpy.nan, jax.numpy.nan)
fname = f"samples_{ARGS.simname}_{ARGS.catalogue}_ksmooth{ARGS.ksmooth}.hdf5" # noqa
if ARGS.ksim is not None:
@ -167,13 +181,14 @@ def run_model(model, nsteps, nburn, model_kwargs, out_folder, sample_beta,
# Write log likelihood and posterior
f.create_dataset("log_likelihood", data=log_likelihood)
f.create_dataset("log_posterior", data=log_posterior)
f.create_dataset("simulation_weights", data=simulation_weights)
# Write goodness of fit
grp = f.create_group("gof")
grp.create_dataset("BIC", data=BIC)
grp.create_dataset("AIC", data=AIC)
grp.create_dataset("lnZ", data=ln_evidence)
grp.create_dataset("lnZ_err", data=ln_evidence_err)
grp.create_dataset("neg_lnZ", data=neg_ln_evidence)
grp.create_dataset("neg_lnZ_err", data=neg_ln_evidence_err)
fname_summary = fname.replace(".hdf5", ".txt")
print(f"Saving summary to `{fname_summary}`.")
@ -188,8 +203,8 @@ def run_model(model, nsteps, nburn, model_kwargs, out_folder, sample_beta,
print("HMC summary:")
print(f"{'BIC':<20} {BIC}")
print(f"{'AIC':<20} {AIC}")
print(f"{'ln(Z)':<20} {ln_evidence}")
print(f"{'ln(Z) error':<20} {ln_evidence_err}")
print(f"{'-ln(Z)':<20} {neg_ln_evidence}")
print(f"{'-ln(Z) error':<20} {neg_ln_evidence_err}")
mcmc.print_summary(exclude_deterministic=False)
sys.stdout = original_stdout
@ -238,7 +253,7 @@ if __name__ == "__main__":
print_variables(
calibration_hyperparams.keys(), calibration_hyperparams.values())
if ARGS.catalogue in ["LOSS", "Foundation", "Pantheon+", "Pantheon+_groups"]: # noqa
if ARGS.catalogue in ["LOSS", "Foundation", "Pantheon+", "Pantheon+_groups", "Pantheon+_zSN"]: # noqa
distmod_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,

View file

@ -44,10 +44,10 @@ COMMENT
# Submit a job for each combination of simname, catalogue, ksim
# for simname in "Lilow2024" "CF4" "CF4gp" "csiborg2_main" "csiborg2X"; do
for simname in "Lilow2024" "Carrick2015" "CF4" "CF4gp"; do
for simname in "Lilow2024"; do
# for simname in "csiborg1" "csiborg2_main" "csiborg2X"; do
for catalogue in "LOSS" "Foundation" "2MTF" "Pantheon+" "Pantheon+_groups" "Pantheon+_zSN" "SFI_gals"; do
# for catalogue in "2MTF" "Pantheon+"; do
# for catalogue in "LOSS" "Foundation" "2MTF" "Pantheon+" "Pantheon+_groups" "Pantheon+_zSN" "SFI_gals"; do
for catalogue in "LOSS"; do
# for ksim in 0 1 2; do
# for ksim in 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20; do
for ksim in "none"; do