Aquila-Consortium/csiborgtools
0
0
Fork 0
mirror of https://github.com/Richard-Sti/csiborgtools_public.git synced 2025-05-14 06:31:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

More flow preparation & Olympics (#143)

Browse source 
* Add more comments

* Add flow paths

* Simplify paths

* Update default arguemnts

* Update paths

* Update param names

* Update some of scipts for reading files

* Add the Mike method option

* Update plotting

* Update fnames

* Simplify things

* Make more default options

* Add print

* Update

* Downsample CF4

* Update numpyro selection

* Add selection fitting nb

* Add coeffs

* Update script

* Add nb

* Add label

* Increase number of steps

* Update default params

* Add more labels

* Improve file name

* Update nb

* Fix little bug

* Remove import

* Update scales

* Update labels

* Add script

* Update script

* Add more

* Add more labels

* Add script

* Add submit

* Update spacing

* Update submit scrips

* Update script

* Update defaults

* Update defaults

* Update nb

* Update test

* Update imports

* Add script

* Add support for Indranil void

* Add a dipole

* Update nb

* Update submit

* Update Om0

* Add final

* Update default params

* Fix bug

* Add option to fix to LG frame

* Add Vext label

* Add Vext label

* Update script

* Rm fixed LG

* rm LG stuff

* Update script

* Update bulk flow plotting

* Update nb

* Add no field option

* Update defaults

* Update nb

* Update script

* Update nb

* Update nb

* Add names to plots

* Update nb

* Update plot

* Add more latex names

* Update default

* Update nb

* Update np

* Add plane slicing

* Add nb with slices

* Update nb

* Update script

* Upddate nb

* Update nb
This commit is contained in:
Richard Stiskalek 2024-09-11 08:45:42 +02:00 • committed by GitHub
parent 3d1e1c0ae3
commit 2b938c112c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 3106 additions and 379 deletions
Download patch file Download diff file Expand all files Collapse all files

248
notebooks/flow/reconstruction_comparison.py
View file

@ -13,8 +13,7 @@
# 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
from copy import copy, deepcopy
import numpy as np
from jax import numpy as jnp
@ -41,25 +40,6 @@ def cartesian_to_radec(x, y, z):
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 #
###############################################################################
@ -69,30 +49,96 @@ 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",
"Vmag": "V_{\\rm ext} ~ [\\mathrm{km} / \\mathrm{s}]",
"Vx": "V_x ~ [\\mathrm{km} / \\mathrm{s}]",
"Vy": "V_y ~ [\\mathrm{km} / \\mathrm{s}]",
"Vz": "V_z ~ [\\mathrm{km} / \\mathrm{s}]",
"sigma_v": "\\sigma_v ~ [\\mathrm{km} / \\mathrm{s}]",
"alpha_cal": "\\mathcal{A}",
"beta_cal": "\\mathcal{B}",
"mag_cal": "\\mathcal{M}",
"e_mu": "\\sigma_\\mu",
"aTF": "a_{\\rm TF}",
"bTF": "b_{\\rm TF}",
"l": "\\ell ~ [\\mathrm{deg}]",
"b": "b ~ [\\mathrm{deg}]",
}
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}}$",
"l": r"$\ell$",
"b": r"$b$",
"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}$",
}
names = copy(names)
for i, name in enumerate(names):
if "SFI_gals" in name:
names[i] = names[i].replace("SFI_gals", "SFI")
if "CF4_GroupAll" in name:
names[i] = names[i].replace("CF4_GroupAll", "CF4Group")
if "CF4_TFR_i" in name:
names[i] = names[i].replace("CF4_TFR_i", "CF4,TFR")
for cat in ["2MTF", "SFI", "CF4,TFR"]:
ltx[f"a_{cat}"] = f"a_{{\\rm TF}}^{{\\rm {cat}}}"
ltx[f"b_{cat}"] = f"b_{{\\rm TF}}^{{\\rm {cat}}}"
ltx[f"c_{cat}"] = f"c_{{\\rm TF}}^{{\\rm {cat}}}"
ltx[f"corr_mag_eta_{cat}"] = f"\\rho_{{m,\\eta}}^{{\\rm {cat}}}"
ltx[f"eta_mean_{cat}"] = f"\\widehat{{\\eta}}^{{\\rm {cat}}}"
ltx[f"eta_std_{cat}"] = f"\\widehat{{\\sigma}}_\\eta^{{\\rm {cat}}}"
ltx[f"mag_mean_{cat}"] = f"\\widehat{{m}}^{{\\rm {cat}}}"
ltx[f"mag_std_{cat}"] = f"\\widehat{{\\sigma}}_m^{{\\rm {cat}}}"
ltx_corner[f"a_{cat}"] = rf"$a_{{\rm TF}}^{{\rm {cat}}}$"
ltx_corner[f"b_{cat}"] = rf"$b_{{\rm TF}}^{{\rm {cat}}}$"
ltx_corner[f"c_{cat}"] = rf"$c_{{\rm TF}}^{{\rm {cat}}}$"
ltx_corner[f"corr_mag_eta_{cat}"] = rf"$\rho_{{m,\eta}}^{{\rm {cat}}}$"
ltx_corner[f"eta_mean_{cat}"] = rf"$\widehat{{\eta}}^{{\rm {cat}}}$"
ltx_corner[f"eta_std_{cat}"] = rf"$\widehat{{\sigma}}_\eta^{{\rm {cat}}}$" # noqa
ltx_corner[f"mag_mean_{cat}"] = rf"$\widehat{{m}}^{{\rm {cat}}}$"
ltx_corner[f"mag_std_{cat}"] = rf"$\widehat{{\sigma}}_m^{{\rm {cat}}}$"
for cat in ["2MTF", "SFI", "Foundation", "LOSS", "CF4Group", "CF4,TFR"]:
ltx[f"alpha_{cat}"] = f"\\alpha^{{\\rm {cat}}}"
ltx[f"e_mu_{cat}"] = f"\\sigma_{{\\mu}}^{{\\rm {cat}}}"
ltx[f"a_dipole_mag_{cat}"] = f"\\epsilon_{{\\rm mag}}^{{\\rm {cat}}}"
ltx[f"a_dipole_l_{cat}"] = f"\\epsilon_{{\\ell}}^{{\\rm {cat}}} ~ [\\mathrm{{deg}}]" # noqa
ltx[f"a_dipole_b_{cat}"] = f"\\epsilon_{{b}}^{{\\rm {cat}}} ~ [\\mathrm{{deg}}]" # noqa
ltx["a_dipole_mag"] = "\\epsilon_{{\\rm mag}}"
ltx["a_dipole_l"] = "\\epsilon_{{\\ell}} ~ [\\mathrm{{deg}}]"
ltx["a_dipole_b"] = "\\epsilon_{{b}} ~ [\\mathrm{{deg}}]"
ltx_corner[f"alpha_{cat}"] = rf"$\alpha^{{\rm {cat}}}$"
ltx_corner[f"e_mu_{cat}"] = rf"$\sigma_{{\mu}}^{{\rm {cat}}}$"
ltx_corner[f"a_dipole_mag_{cat}"] = rf"$\epsilon_{{\rm mag}}^{{\rm {cat}}}$" # noqa
ltx_corner[f"a_dipole_l_{cat}"] = rf"$\epsilon_{{\ell}}^{{\rm {cat}}}$"
ltx_corner[f"a_dipole_b_{cat}"] = rf"$\epsilon_{{b}}^{{\rm {cat}}}$"
for cat in ["Foundation", "LOSS"]:
ltx[f"alpha_cal_{cat}"] = f"\\mathcal{{A}}^{{\\rm {cat}}}"
ltx[f"beta_cal_{cat}"] = f"\\mathcal{{B}}^{{\\rm {cat}}}"
ltx[f"mag_cal_{cat}"] = f"\\mathcal{{M}}^{{\\rm {cat}}}"
ltx_corner[f"alpha_cal_{cat}"] = rf"$\mathcal{{A}}^{{\rm {cat}}}$"
ltx_corner[f"beta_cal_{cat}"] = rf"$\mathcal{{B}}^{{\rm {cat}}}$"
ltx_corner[f"mag_cal_{cat}"] = rf"$\mathcal{{M}}^{{\rm {cat}}}$"
for cat in ["CF4Group"]:
ltx[f"dmu_{cat}"] = f"\\Delta\\mu^{{\\rm {cat}}}"
ltx[f"dmu_dipole_mag_{cat}"] = f"\\epsilon_\\mu_{{\\rm mag}}^{{\\rm {cat}}}" # noqa
ltx[f"dmu_dipole_l_{cat}"] = f"\\epsilon_\\mu_{{\\ell}}^{{\\rm {cat}}} ~ [\\mathrm{{deg}}]" # noqa
ltx[f"dmu_dipole_b_{cat}"] = f"\\epsilon_\\mu_{{b}}^{{\\rm {cat}}} ~ [\\mathrm{{deg}}]" # noqa
ltx_corner[f"dmu_{cat}"] = rf"$\Delta\mu_{{0}}^{{\rm {cat}}}$"
ltx_corner[f"dmu_dipole_mag_{cat}"] = rf"$\epsilon_{{\rm mag}}^{{\rm {cat}}}$" # noqa
ltx_corner[f"dmu_dipole_l_{cat}"] = rf"$\epsilon_{{\ell}}^{{\rm {cat}}}$" # noqa
ltx_corner[f"dmu_dipole_b_{cat}"] = rf"$\epsilon_{{b}}^{{\rm {cat}}}$" # noqa
labels = copy(names)
for i, label in enumerate(names):
if for_corner:
@ -113,21 +159,35 @@ def simname_to_pretty(simname):
}
if isinstance(simname, list):
return [ltx[s] if s in ltx else s for s in simname]
names = [ltx[s] if s in ltx else s for s in simname]
return "".join([f"{n}, " for n in names]).rstrip(", ")
return ltx[simname] if simname in ltx else simname
def catalogue_to_pretty(catalogue):
ltx = {"SFI_gals": "SFI",
"CF4_TFR_not2MTForSFI_i": r"CF4 $i$-band",
"CF4_TFR_i": r"CF4 TFR $i$",
"CF4_TFR_w1": r"CF4 TFR W1",
}
if isinstance(catalogue, list):
names = [ltx[s] if s in ltx else s for s in catalogue]
return "".join([f"{n}, " for n in names]).rstrip(", ")
return ltx[catalogue] if catalogue in ltx else catalogue
###############################################################################
# Read in goodness-of-fit #
###############################################################################
def get_gof(kind, simname, catalogue, ksmooth=0, nsim=None, sample_beta=True):
def get_gof(kind, fname):
"""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'")
if kind not in ["BIC", "AIC", "neg_lnZ_harmonic"]:
raise ValueError("`kind` must be one of 'BIC', 'AIC', 'neg_lnZ_harmonic'.") # noqa
fname = get_fname(simname, catalogue, ksmooth, nsim, sample_beta)
with File(fname, 'r') as f:
return f[f"gof/{kind}"][()]
@ -136,29 +196,48 @@ def get_gof(kind, simname, catalogue, ksmooth=0, nsim=None, sample_beta=True):
# Read in samples #
###############################################################################
def get_samples(simname, catalogue, ksmooth=0, nsim=None, sample_beta=True,
convert_Vext_to_galactic=True):
def get_samples(fname, 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])
else:
Vext = samples.pop("Vext")
samples["Vx"] = Vext[:, 0]
samples["Vy"] = Vext[:, 1]
samples["Vz"] = Vext[:, 2]
keys = list(samples.keys())
for key in keys:
if "dmu_dipole_" in key:
dmu = samples.pop(key)
dmu = csiborgtools.cartesian_to_radec(dmu)
dmu_mag = dmu[:, 0]
l, b = csiborgtools.radec_to_galactic(dmu[:, 1], dmu[:, 2])
samples[key.replace("dmu_dipole_", "dmu_dipole_mag_")] = dmu_mag
samples[key.replace("dmu_dipole_", "dmu_dipole_l_")] = l
samples[key.replace("dmu_dipole_", "dmu_dipole_b_")] = b
if "a_dipole" in key:
adipole = samples.pop(key)
adipole = csiborgtools.cartesian_to_radec(adipole)
adipole_mag = adipole[:, 0]
l, b = csiborgtools.radec_to_galactic(adipole[:, 1], adipole[:, 2])
samples[key.replace("a_dipole", "a_dipole_mag")] = adipole_mag
samples[key.replace("a_dipole", "a_dipole_l")] = l
samples[key.replace("a_dipole", "a_dipole_b")] = b
return samples
@ -180,12 +259,20 @@ def get_bulkflow_simulation(simname, convert_to_galactic=True):
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):
def get_bulkflow(fname, simname, convert_to_galactic=True, downsample=1,
Rmax=125):
# Read in the samples
with File(fname, "r") as f:
Vext = f["samples/Vext"][...]
try:
beta = f["samples/beta"][...]
except KeyError:
beta = jnp.ones(len(Vext))
# 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))
@ -197,38 +284,18 @@ def get_bulkflow(simname, catalogue, ksmooth=0, nsim=None, sample_beta=True,
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)
Vext = Vext[::downsample]
beta = beta[::downsample]
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
# 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)
# Add V_ext, shape of B_i is `(nsims, nradial, nsamples)``
Bx = Bx + Vext[:, 0]
By = By + Vext[:, 1]
Bz = Bz + Vext[:, 2]
if convert_to_galactic:
Bmag, Bl, Bb = cartesian_to_radec(Bx, By, Bz)
@ -237,6 +304,8 @@ def get_bulkflow(simname, catalogue, ksmooth=0, nsim=None, sample_beta=True,
else:
B = np.stack([Bx, By, Bz], axis=-1)
# Stack over the simulations
B = np.hstack([B[i] for i in range(len(B))])
return r, B
###############################################################################
@ -245,25 +314,30 @@ def get_bulkflow(simname, catalogue, ksmooth=0, nsim=None, sample_beta=True,
def samples_for_corner(samples):
samples = deepcopy(samples)
# Remove the true parameters of each galaxy.
keys = list(samples.keys())
for key in keys:
# Generally don't want to plot the true latent parameters..
if "x_TFR" in key or "_true_" in key:
samples.pop(key)
keys = list(samples.keys())
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
return data, labels, keys
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
def samples_to_getdist(samples, label):
data, __, keys = samples_for_corner(samples)
return MCSamples(
samples=data, names=names,
labels=names_to_latex(names, for_corner=False),
samples=data, names=keys,
labels=names_to_latex(keys, for_corner=False),
label=label)