mirror of
https://github.com/Richard-Sti/csiborgtools.git
synced 2024-10-18 20:15:05 +02:00
43ebf660ee
* Update priors * Update submit * Update README * Relax width of calibration priors * Relax width of calibration priors * Update smooth scales * Update settings * Update README * Update submit * Add names * Update nb * Fix bug * Update script * Move files * Small bug fix * Add script * Add script * Add script * Add more power * Update script * Quick fix * Rm blank line * Update nb * Update nb * Update nb
270 lines
9.3 KiB
Python
270 lines
9.3 KiB
Python
# 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",
|
|
"CF4": "CF4",
|
|
"CF4gp": "CF4group",
|
|
}
|
|
|
|
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
|
|
|
|
|
|
###############################################################################
|
|
# Bulk flow plotting #
|
|
###############################################################################
|
|
|
|
|
|
def get_bulkflow_simulation(simname, convert_to_galactic=True):
|
|
f = np.load(f"/mnt/extraspace/rstiskalek/csiborg_postprocessing/field_shells/enclosed_mass_{simname}.npz") # noqa
|
|
r, B = f["distances"], f["cumulative_velocity"]
|
|
|
|
if convert_to_galactic:
|
|
Bmag, Bl, Bb = cartesian_to_radec(B[..., 0], B[..., 1], B[..., 2])
|
|
Bl, Bb = csiborgtools.radec_to_galactic(Bl, Bb)
|
|
B = np.stack([Bmag, Bl, Bb], axis=-1)
|
|
|
|
return r, B
|
|
|
|
|
|
def get_bulkflow(simname, catalogue, ksmooth=0, nsim=None, sample_beta=True,
|
|
convert_to_galactic=True, weight_simulations=True,
|
|
downsample=1, Rmax=125):
|
|
# Read in the bulk flow
|
|
f = np.load(f"/mnt/extraspace/rstiskalek/csiborg_postprocessing/field_shells/enclosed_mass_{simname}.npz") # noqa
|
|
r = f["distances"]
|
|
# Shape of B_i is (nsims, nradial)
|
|
Bx, By, Bz = (f["cumulative_velocity"][..., i] for i in range(3))
|
|
|
|
# Mask out the unconstrained large scales
|
|
Rmax = Rmax # Mpc/h
|
|
mask = r < Rmax
|
|
r = r[mask]
|
|
Bx = Bx[:, mask]
|
|
By = By[:, mask]
|
|
Bz = Bz[:, mask]
|
|
|
|
# Read in the samples
|
|
fname_samples = get_fname(simname, catalogue, ksmooth, nsim, sample_beta)
|
|
with File(fname_samples, 'r') as f:
|
|
# Shape of Vext_i is (nsamples,)
|
|
Vext_x, Vext_y, Vext_z = (f["samples/Vext"][...][::downsample, i] for i in range(3)) # noqa
|
|
nsamples = len(Vext_x)
|
|
|
|
if weight_simulations:
|
|
simulation_weights = jnp.exp(f["simulation_weights"][...])[::downsample] # noqa
|
|
else:
|
|
nsims = len(Bx)
|
|
simulation_weights = jnp.ones((nsamples, nsims)) / nsims
|
|
|
|
if sample_beta:
|
|
beta = f["samples/beta"][...][::downsample]
|
|
else:
|
|
beta = jnp.ones(nsamples)
|
|
|
|
# Multiply the simulation velocities by beta
|
|
Bx = Bx[..., None] * beta
|
|
By = By[..., None] * beta
|
|
Bz = Bz[..., None] * beta
|
|
|
|
# Shape of B_i is (nsims, nradial, nsamples)
|
|
Bx = Bx + Vext_x
|
|
By = By + Vext_y
|
|
Bz = Bz + Vext_z
|
|
|
|
simulation_weights = simulation_weights.T[:, None, :]
|
|
Bx = jnp.sum(Bx * simulation_weights, axis=0)
|
|
By = jnp.sum(By * simulation_weights, axis=0)
|
|
Bz = jnp.sum(Bz * simulation_weights, axis=0)
|
|
|
|
if convert_to_galactic:
|
|
Bmag, Bl, Bb = cartesian_to_radec(Bx, By, Bz)
|
|
Bl, Bb = csiborgtools.radec_to_galactic(Bl, Bb)
|
|
B = np.stack([Bmag, Bl, Bb], axis=-1)
|
|
else:
|
|
B = np.stack([Bx, By, Bz], 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)
|