mirror of
https://github.com/Richard-Sti/csiborgtools.git
synced 2024-12-22 17:38:02 +00:00
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
GitHub
parent
7dad6885e8
commit
c6f49790bf
13 changed files with 1208 additions and 2680 deletions
|
|
@ -666,7 +666,8 @@ class PV_validation_model(BaseFlowValidationModel):
|
||||||
self.Omega_m = Omega_m
|
self.Omega_m = Omega_m
|
||||||
self.norm = - self.ndata * jnp.log(self.num_sims)
|
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."""
|
"""NumPyro PV validation model."""
|
||||||
Vext, sigma_v, alpha, beta = sample_calibration(**calibration_hyperparams) # noqa
|
Vext, sigma_v, alpha, beta = sample_calibration(**calibration_hyperparams) # noqa
|
||||||
cz_err = jnp.sqrt(sigma_v**2 + self.e2_cz_obs)
|
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)
|
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.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.deterministic("ll_values", ll)
|
||||||
numpyro.factor("ll", ll)
|
numpyro.factor("ll", ll)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -18,7 +18,8 @@ from .catalogue import (CSiBORG1Catalogue, CSiBORG2Catalogue,
|
||||||
MDPL2Catalogue, fiducial_observers) # noqa
|
MDPL2Catalogue, fiducial_observers) # noqa
|
||||||
from .snapshot import (CSiBORG1Snapshot, CSiBORG2Snapshot, QuijoteSnapshot, # noqa
|
from .snapshot import (CSiBORG1Snapshot, CSiBORG2Snapshot, QuijoteSnapshot, # noqa
|
||||||
CSiBORG1Field, CSiBORG2Field, CSiBORG2XField, # 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
|
from .obs import (SDSS, MCXCClusters, PlanckClusters, TwoMPPGalaxies, # noqa
|
||||||
TwoMPPGroups, ObservedCluster, match_array_to_no_masking, # noqa
|
TwoMPPGroups, ObservedCluster, match_array_to_no_masking, # noqa
|
||||||
cols_to_structured) # noqa
|
cols_to_structured) # noqa
|
||||||
|
|
|
||||||
|
|
@ -372,6 +372,24 @@ class Paths:
|
||||||
else:
|
else:
|
||||||
raise ValueError(f"Unsupported CSiBORG2X field: `{kind}`.")
|
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 MAS == "SPH" and kind in ["density", "velocity"]:
|
||||||
if simname == "csiborg1":
|
if simname == "csiborg1":
|
||||||
return join(self.csiborg1_srcdir, "field",
|
return join(self.csiborg1_srcdir, "field",
|
||||||
|
|
|
||||||
|
|
@ -20,7 +20,7 @@ observed RA-dec can be mapped into the simulation box.
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from os.path import join
|
from os.path import join
|
||||||
|
|
||||||
import numpy
|
import numpy as np
|
||||||
from h5py import File
|
from h5py import File
|
||||||
|
|
||||||
from ..params import paths_glamdring, simname2boxsize, simname2Omega_m
|
from ..params import paths_glamdring, simname2boxsize, simname2Omega_m
|
||||||
|
|
@ -40,11 +40,11 @@ class BaseSnapshot(ABC):
|
||||||
flip_xz=False):
|
flip_xz=False):
|
||||||
self._keep_snapshot_open = None
|
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")
|
raise TypeError("`nsim` must be an integer")
|
||||||
self._nsim = int(nsim)
|
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")
|
raise TypeError("`nsnap` must be an integer")
|
||||||
self._nsnap = int(nsnap)
|
self._nsnap = int(nsnap)
|
||||||
|
|
||||||
|
|
@ -409,16 +409,16 @@ class CSiBORG2Snapshot(BaseSnapshot):
|
||||||
with File(self._snapshot_path, "r") as f:
|
with File(self._snapshot_path, "r") as f:
|
||||||
if kind == "Masses":
|
if kind == "Masses":
|
||||||
npart = f["Header"].attrs["NumPart_Total"][1]
|
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]
|
x *= f["Header"].attrs["MassTable"][1]
|
||||||
else:
|
else:
|
||||||
x = f[f"PartType1/{kind}"][...]
|
x = f[f"PartType1/{kind}"][...]
|
||||||
|
|
||||||
if not high_resolution_only:
|
if not high_resolution_only:
|
||||||
if x.ndim == 1:
|
if x.ndim == 1:
|
||||||
x = numpy.hstack([x, f[f"PartType5/{kind}"][...]])
|
x = np.hstack([x, f[f"PartType5/{kind}"][...]])
|
||||||
else:
|
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"]:
|
if self.flip_xz and kind in ["Coordinates", "Velocities"]:
|
||||||
x[:, [0, 2]] = x[:, [2, 0]]
|
x[:, [0, 2]] = x[:, [2, 0]]
|
||||||
|
|
@ -453,7 +453,7 @@ class CSiBORG2Snapshot(BaseSnapshot):
|
||||||
|
|
||||||
if i1 is not None and j1 - i1 > 0:
|
if i1 is not None and j1 - i1 > 0:
|
||||||
if kind == "Masses":
|
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]
|
x1 *= f["Header"].attrs["MassTable"][1]
|
||||||
else:
|
else:
|
||||||
x1 = f[f"PartType1/{kind}"][i1:j1]
|
x1 = f[f"PartType1/{kind}"][i1:j1]
|
||||||
|
|
@ -481,9 +481,9 @@ class CSiBORG2Snapshot(BaseSnapshot):
|
||||||
return x5
|
return x5
|
||||||
|
|
||||||
if x1.ndim > 1:
|
if x1.ndim > 1:
|
||||||
x1 = numpy.vstack([x1, x5])
|
x1 = np.vstack([x1, x5])
|
||||||
else:
|
else:
|
||||||
x1 = numpy.hstack([x1, x5])
|
x1 = np.hstack([x1, x5])
|
||||||
|
|
||||||
return x1
|
return x1
|
||||||
|
|
||||||
|
|
@ -565,7 +565,7 @@ class BaseField(ABC):
|
||||||
Base class for reading fields such as density or velocity fields.
|
Base class for reading fields such as density or velocity fields.
|
||||||
"""
|
"""
|
||||||
def __init__(self, nsim, paths, flip_xz=False):
|
def __init__(self, nsim, paths, flip_xz=False):
|
||||||
if isinstance(nsim, numpy.integer):
|
if isinstance(nsim, np.integer):
|
||||||
nsim = int(nsim)
|
nsim = int(nsim)
|
||||||
if not isinstance(nsim, int):
|
if not isinstance(nsim, int):
|
||||||
raise TypeError(f"`nsim` must be an integer. Received `{type(nsim)}`.") # noqa
|
raise TypeError(f"`nsim` must be an integer. Received `{type(nsim)}`.") # noqa
|
||||||
|
|
@ -644,7 +644,7 @@ class CSiBORG1Field(BaseField):
|
||||||
field = f["density"][:]
|
field = f["density"][:]
|
||||||
field /= (677.7 * 1e3 / grid)**3 # Convert to h^2 Msun / kpc^3
|
field /= (677.7 * 1e3 / grid)**3 # Convert to h^2 Msun / kpc^3
|
||||||
else:
|
else:
|
||||||
field = numpy.load(fpath)
|
field = np.load(fpath)
|
||||||
|
|
||||||
if self.flip_xz:
|
if self.flip_xz:
|
||||||
field = field.T
|
field = field.T
|
||||||
|
|
@ -660,9 +660,9 @@ class CSiBORG1Field(BaseField):
|
||||||
v0 = f["p0"][:] / density
|
v0 = f["p0"][:] / density
|
||||||
v1 = f["p1"][:] / density
|
v1 = f["p1"][:] / density
|
||||||
v2 = f["p2"][:] / density
|
v2 = f["p2"][:] / density
|
||||||
field = numpy.array([v0, v1, v2])
|
field = np.array([v0, v1, v2])
|
||||||
else:
|
else:
|
||||||
field = numpy.load(fpath)
|
field = np.load(fpath)
|
||||||
|
|
||||||
if self.flip_xz:
|
if self.flip_xz:
|
||||||
field[0, ...] = field[0, ...].T
|
field[0, ...] = field[0, ...].T
|
||||||
|
|
@ -678,7 +678,7 @@ class CSiBORG1Field(BaseField):
|
||||||
"for the flipped x- and z-axes.")
|
"for the flipped x- and z-axes.")
|
||||||
|
|
||||||
fpath = self.paths.field("radvel", MAS, grid, self.nsim, "csiborg1")
|
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 *= 1e10 # Convert to Msun / h
|
||||||
field /= (676.6 * 1e3 / grid)**3 # Convert to h^2 Msun / kpc^3
|
field /= (676.6 * 1e3 / grid)**3 # Convert to h^2 Msun / kpc^3
|
||||||
else:
|
else:
|
||||||
field = numpy.load(fpath)
|
field = np.load(fpath)
|
||||||
|
|
||||||
if self.flip_xz:
|
if self.flip_xz:
|
||||||
field = field.T
|
field = field.T
|
||||||
|
|
@ -744,9 +744,9 @@ class CSiBORG2Field(BaseField):
|
||||||
v0 = f["p0"][:] / density
|
v0 = f["p0"][:] / density
|
||||||
v1 = f["p1"][:] / density
|
v1 = f["p1"][:] / density
|
||||||
v2 = f["p2"][:] / density
|
v2 = f["p2"][:] / density
|
||||||
field = numpy.array([v0, v1, v2])
|
field = np.array([v0, v1, v2])
|
||||||
else:
|
else:
|
||||||
field = numpy.load(fpath)
|
field = np.load(fpath)
|
||||||
|
|
||||||
if self.flip_xz:
|
if self.flip_xz:
|
||||||
field[0, ...] = field[0, ...].T
|
field[0, ...] = field[0, ...].T
|
||||||
|
|
@ -763,7 +763,7 @@ class CSiBORG2Field(BaseField):
|
||||||
|
|
||||||
fpath = self.paths.field("radvel", MAS, grid, self.nsim,
|
fpath = self.paths.field("radvel", MAS, grid, self.nsim,
|
||||||
f"csiborg2_{self.kind}")
|
f"csiborg2_{self.kind}")
|
||||||
return numpy.load(fpath)
|
return np.load(fpath)
|
||||||
|
|
||||||
|
|
||||||
class CSiBORG2XField(BaseField):
|
class CSiBORG2XField(BaseField):
|
||||||
|
|
@ -784,7 +784,7 @@ class CSiBORG2XField(BaseField):
|
||||||
fpath = self.paths.field(
|
fpath = self.paths.field(
|
||||||
"overdensity", None, None, self.nsim, "csiborg2X")
|
"overdensity", None, None, self.nsim, "csiborg2X")
|
||||||
with File(fpath, "r") as f:
|
with File(fpath, "r") as f:
|
||||||
field = f["delta_cic"][...].astype(numpy.float32)
|
field = f["delta_cic"][...].astype(np.float32)
|
||||||
|
|
||||||
return field
|
return field
|
||||||
|
|
||||||
|
|
@ -803,7 +803,7 @@ class CSiBORG2XField(BaseField):
|
||||||
v0 = f["v_0"][...]
|
v0 = f["v_0"][...]
|
||||||
v1 = f["v_1"][...]
|
v1 = f["v_1"][...]
|
||||||
v2 = f["v_2"][...]
|
v2 = f["v_2"][...]
|
||||||
field = numpy.array([v0, v1, v2])
|
field = np.array([v0, v1, v2])
|
||||||
|
|
||||||
return field
|
return field
|
||||||
|
|
||||||
|
|
@ -833,7 +833,7 @@ class BORG1Field(BaseField):
|
||||||
def overdensity_field(self):
|
def overdensity_field(self):
|
||||||
fpath = self.paths.field(None, None, None, self.nsim, "borg1")
|
fpath = self.paths.field(None, None, None, self.nsim, "borg1")
|
||||||
with File(fpath, "r") as f:
|
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
|
return field
|
||||||
|
|
||||||
|
|
@ -874,7 +874,7 @@ class BORG2Field(BaseField):
|
||||||
def overdensity_field(self):
|
def overdensity_field(self):
|
||||||
fpath = self.paths.field(None, None, None, self.nsim, "borg2")
|
fpath = self.paths.field(None, None, None, self.nsim, "borg2")
|
||||||
with File(fpath, "r") as f:
|
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
|
return field
|
||||||
|
|
||||||
|
|
@ -894,7 +894,7 @@ class BORG2Field(BaseField):
|
||||||
|
|
||||||
|
|
||||||
###############################################################################
|
###############################################################################
|
||||||
# TNG300-1 field #
|
# Additional fields #
|
||||||
###############################################################################
|
###############################################################################
|
||||||
|
|
||||||
class TNG300_1Field(BaseField):
|
class TNG300_1Field(BaseField):
|
||||||
|
|
@ -922,7 +922,7 @@ class TNG300_1Field(BaseField):
|
||||||
def density_field(self, MAS, grid):
|
def density_field(self, MAS, grid):
|
||||||
fpath = join(self.paths.tng300_1(), "postprocessing", "density_field",
|
fpath = join(self.paths.tng300_1(), "postprocessing", "density_field",
|
||||||
f"rho_dm_099_{grid}_{MAS}.npy")
|
f"rho_dm_099_{grid}_{MAS}.npy")
|
||||||
return numpy.load(fpath)
|
return np.load(fpath)
|
||||||
|
|
||||||
def velocity_field(self, MAS, grid):
|
def velocity_field(self, MAS, grid):
|
||||||
raise RuntimeError("The velocity field is not available.")
|
raise RuntimeError("The velocity field is not available.")
|
||||||
|
|
@ -930,10 +930,6 @@ class TNG300_1Field(BaseField):
|
||||||
def radial_velocity_field(self, MAS, grid):
|
def radial_velocity_field(self, MAS, grid):
|
||||||
raise RuntimeError("The radial velocity field is not available.")
|
raise RuntimeError("The radial velocity field is not available.")
|
||||||
|
|
||||||
###############################################################################
|
|
||||||
# Quijote field class #
|
|
||||||
###############################################################################
|
|
||||||
|
|
||||||
|
|
||||||
class QuijoteField(CSiBORG1Field):
|
class QuijoteField(CSiBORG1Field):
|
||||||
"""
|
"""
|
||||||
|
|
@ -951,6 +947,90 @@ class QuijoteField(CSiBORG1Field):
|
||||||
self._simname = "quijote"
|
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 #
|
# Supplementary functions #
|
||||||
###############################################################################
|
###############################################################################
|
||||||
|
|
|
||||||
File diff suppressed because one or more lines are too long
|
|
@ -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
|
|
||||||
293
notebooks/flow/flow_los.ipynb
Normal file
293
notebooks/flow/flow_los.ipynb
Normal 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
|
||||||
|
}
|
||||||
481
notebooks/flow/reconstruction_comparison.ipynb
Normal file
481
notebooks/flow/reconstruction_comparison.ipynb
Normal file
File diff suppressed because one or more lines are too long
230
notebooks/flow/reconstruction_comparison.py
Normal file
230
notebooks/flow/reconstruction_comparison.py
Normal 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)
|
||||||
|
|
@ -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
|
|
||||||
}
|
|
||||||
|
|
@ -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.
|
main bottleneck is reading the data from disk.
|
||||||
"""
|
"""
|
||||||
from argparse import ArgumentParser
|
from argparse import ArgumentParser
|
||||||
from os.path import join
|
from datetime import datetime
|
||||||
from gc import collect
|
from gc import collect
|
||||||
|
from os.path import join
|
||||||
|
|
||||||
import csiborgtools
|
import csiborgtools
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
from astropy import units as u
|
||||||
|
from astropy.coordinates import CartesianRepresentation, SkyCoord
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
|
|
||||||
from datetime import datetime
|
|
||||||
|
|
||||||
|
|
||||||
###############################################################################
|
###############################################################################
|
||||||
# Read in information about the simulation #
|
# 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)
|
cumulative_velocity=cumulative_velocity)
|
||||||
|
|
||||||
|
|
||||||
def main_csiborg2X(args, folder):
|
def main_from_field(args, folder):
|
||||||
"""Bulk flow in the Manticore boxes provided by Stuart."""
|
"""Bulk flow in the Manticore boxes provided by Stuart."""
|
||||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||||
boxsize = csiborgtools.simname2boxsize(args.simname)
|
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_y = np.zeros_like(cumulative_vel_x)
|
||||||
cumulative_vel_z = np.zeros_like(cumulative_vel_x)
|
cumulative_vel_z = np.zeros_like(cumulative_vel_x)
|
||||||
for i, nsim in enumerate(tqdm(nsims, desc="Simulations")):
|
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()
|
density_field = reader.density_field()
|
||||||
velocity_field = reader.velocity_field()
|
velocity_field = reader.velocity_field()
|
||||||
|
|
@ -229,6 +236,20 @@ def main_csiborg2X(args, folder):
|
||||||
cumulative_vel_z[i, :], __ = csiborgtools.field.field_enclosed_mass(
|
cumulative_vel_z[i, :], __ = csiborgtools.field.field_enclosed_mass(
|
||||||
velocity_field[2], distances, boxsize, verbose=False)
|
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 = np.stack(
|
||||||
[cumulative_vel_x, cumulative_vel_y, cumulative_vel_z], axis=-1)
|
[cumulative_vel_x, cumulative_vel_y, cumulative_vel_z], axis=-1)
|
||||||
cumulative_vel /= cumulative_volume[..., None]
|
cumulative_vel /= cumulative_volume[..., None]
|
||||||
|
|
@ -236,20 +257,28 @@ def main_csiborg2X(args, folder):
|
||||||
# Finally save the output
|
# Finally save the output
|
||||||
fname = f"enclosed_mass_{args.simname}.npz"
|
fname = f"enclosed_mass_{args.simname}.npz"
|
||||||
fname = join(folder, fname)
|
fname = join(folder, fname)
|
||||||
|
print(f"Saving to `{fname}`.")
|
||||||
np.savez(fname, enclosed_mass=cumulative_mass, distances=distances,
|
np.savez(fname, enclosed_mass=cumulative_mass, distances=distances,
|
||||||
cumulative_velocity=cumulative_vel,
|
cumulative_velocity=cumulative_vel,
|
||||||
enclosed_volume=cumulative_volume)
|
enclosed_volume=cumulative_volume)
|
||||||
|
|
||||||
|
|
||||||
|
###############################################################################
|
||||||
|
# Command line interface #
|
||||||
|
###############################################################################
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
parser = ArgumentParser()
|
parser = ArgumentParser()
|
||||||
parser.add_argument("--simname", type=str, help="Simulation name.",
|
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()
|
args = parser.parse_args()
|
||||||
|
|
||||||
folder = "/mnt/extraspace/rstiskalek/csiborg_postprocessing/field_shells"
|
folder = "/mnt/extraspace/rstiskalek/csiborg_postprocessing/field_shells"
|
||||||
if args.simname == "csiborg2X":
|
if args.simname in ["csiborg2X", "Carrick2015", "Lilow2024"]:
|
||||||
main_csiborg2X(args, folder)
|
main_from_field(args, folder)
|
||||||
elif "csiborg" in args.simname:
|
elif "csiborg" in args.simname:
|
||||||
main_csiborg(args, folder)
|
main_csiborg(args, folder)
|
||||||
elif "borg" in args.simname:
|
elif "borg" in args.simname:
|
||||||
|
|
|
||||||
|
|
@ -48,18 +48,18 @@ def parse_args():
|
||||||
ARGS = parse_args()
|
ARGS = parse_args()
|
||||||
# This must be done before we import JAX etc.
|
# This must be done before we import JAX etc.
|
||||||
from numpyro import set_host_device_count, set_platform # noqa
|
from numpyro import set_host_device_count, set_platform # noqa
|
||||||
|
|
||||||
set_platform(ARGS.device) # noqa
|
set_platform(ARGS.device) # noqa
|
||||||
set_host_device_count(ARGS.ndevice) # noqa
|
set_host_device_count(ARGS.ndevice) # noqa
|
||||||
|
|
||||||
import sys # noqa
|
import sys # noqa
|
||||||
from os.path import join # noqa
|
from os.path import join # noqa
|
||||||
|
|
||||||
|
import csiborgtools # noqa
|
||||||
import jax # noqa
|
import jax # noqa
|
||||||
from h5py import File # noqa
|
from h5py import File # noqa
|
||||||
from mpi4py import MPI # noqa
|
from mpi4py import MPI # noqa
|
||||||
from numpyro.infer import MCMC, NUTS, init_to_median # noqa
|
from numpyro.infer import MCMC, NUTS, Predictive, init_to_median # noqa
|
||||||
|
|
||||||
import csiborgtools # noqa
|
|
||||||
|
|
||||||
|
|
||||||
def print_variables(names, variables):
|
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)
|
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,
|
def run_model(model, nsteps, nburn, model_kwargs, out_folder, sample_beta,
|
||||||
calculate_evidence, nchains_harmonic, epoch_num, kwargs_print):
|
calculate_evidence, nchains_harmonic, epoch_num, kwargs_print):
|
||||||
"""Run the NumPyro model and save output to a file."""
|
"""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)
|
mcmc.run(rng_key, extra_fields=("potential_energy",), **model_kwargs)
|
||||||
samples = mcmc.get_samples()
|
samples = mcmc.get_samples()
|
||||||
|
simulation_weights = get_simulation_weights(samples, model, model_kwargs)
|
||||||
|
|
||||||
log_posterior = -mcmc.get_extra_fields()["potential_energy"]
|
log_posterior = -mcmc.get_extra_fields()["potential_energy"]
|
||||||
log_likelihood = samples.pop("ll_values")
|
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:
|
if calculate_evidence:
|
||||||
print("Calculating the evidence using `harmonic`.", flush=True)
|
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)
|
samples, log_posterior, nchains_harmonic, epoch_num)
|
||||||
print(f"{'ln(Z)':<20} {ln_evidence}")
|
print(f"{'-ln(Z)':<20} {neg_ln_evidence}")
|
||||||
print(f"{'ln(Z) error':<20} {ln_evidence_err}")
|
print(f"{'-ln(Z) error':<20} {neg_ln_evidence_err}")
|
||||||
else:
|
else:
|
||||||
ln_evidence = jax.numpy.nan
|
neg_ln_evidence = jax.numpy.nan
|
||||||
ln_evidence_err = (jax.numpy.nan, 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
|
fname = f"samples_{ARGS.simname}_{ARGS.catalogue}_ksmooth{ARGS.ksmooth}.hdf5" # noqa
|
||||||
if ARGS.ksim is not None:
|
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
|
# Write log likelihood and posterior
|
||||||
f.create_dataset("log_likelihood", data=log_likelihood)
|
f.create_dataset("log_likelihood", data=log_likelihood)
|
||||||
f.create_dataset("log_posterior", data=log_posterior)
|
f.create_dataset("log_posterior", data=log_posterior)
|
||||||
|
f.create_dataset("simulation_weights", data=simulation_weights)
|
||||||
|
|
||||||
# Write goodness of fit
|
# Write goodness of fit
|
||||||
grp = f.create_group("gof")
|
grp = f.create_group("gof")
|
||||||
grp.create_dataset("BIC", data=BIC)
|
grp.create_dataset("BIC", data=BIC)
|
||||||
grp.create_dataset("AIC", data=AIC)
|
grp.create_dataset("AIC", data=AIC)
|
||||||
grp.create_dataset("lnZ", data=ln_evidence)
|
grp.create_dataset("neg_lnZ", data=neg_ln_evidence)
|
||||||
grp.create_dataset("lnZ_err", data=ln_evidence_err)
|
grp.create_dataset("neg_lnZ_err", data=neg_ln_evidence_err)
|
||||||
|
|
||||||
fname_summary = fname.replace(".hdf5", ".txt")
|
fname_summary = fname.replace(".hdf5", ".txt")
|
||||||
print(f"Saving summary to `{fname_summary}`.")
|
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("HMC summary:")
|
||||||
print(f"{'BIC':<20} {BIC}")
|
print(f"{'BIC':<20} {BIC}")
|
||||||
print(f"{'AIC':<20} {AIC}")
|
print(f"{'AIC':<20} {AIC}")
|
||||||
print(f"{'ln(Z)':<20} {ln_evidence}")
|
print(f"{'-ln(Z)':<20} {neg_ln_evidence}")
|
||||||
print(f"{'ln(Z) error':<20} {ln_evidence_err}")
|
print(f"{'-ln(Z) error':<20} {neg_ln_evidence_err}")
|
||||||
mcmc.print_summary(exclude_deterministic=False)
|
mcmc.print_summary(exclude_deterministic=False)
|
||||||
sys.stdout = original_stdout
|
sys.stdout = original_stdout
|
||||||
|
|
||||||
|
|
@ -238,7 +253,7 @@ if __name__ == "__main__":
|
||||||
print_variables(
|
print_variables(
|
||||||
calibration_hyperparams.keys(), calibration_hyperparams.values())
|
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,
|
distmod_hyperparams = {"e_mu_mean": 0.1, "e_mu_std": 0.05,
|
||||||
"mag_cal_mean": -18.25, "mag_cal_std": 0.5,
|
"mag_cal_mean": -18.25, "mag_cal_std": 0.5,
|
||||||
"alpha_cal_mean": 0.148, "alpha_cal_std": 0.05,
|
"alpha_cal_mean": 0.148, "alpha_cal_std": 0.05,
|
||||||
|
|
|
||||||
|
|
@ -44,10 +44,10 @@ COMMENT
|
||||||
|
|
||||||
# Submit a job for each combination of simname, catalogue, ksim
|
# Submit a job for each combination of simname, catalogue, ksim
|
||||||
# for simname in "Lilow2024" "CF4" "CF4gp" "csiborg2_main" "csiborg2X"; do
|
# 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 simname in "csiborg1" "csiborg2_main" "csiborg2X"; do
|
||||||
for catalogue in "LOSS" "Foundation" "2MTF" "Pantheon+" "Pantheon+_groups" "Pantheon+_zSN" "SFI_gals"; 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"; do
|
||||||
# for ksim in 0 1 2; 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 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
|
for ksim in "none"; do
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue