# 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, deepcopy 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 ############################################################################### # 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} ~ [\\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}", "l": "\\ell ~ [\\mathrm{deg}]", "b": "b ~ [\\mathrm{deg}]", } ltx_corner = {"alpha": r"$\alpha$", "beta": r"$\beta$", "Vmag": r"$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}$", } 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: 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": "Courtois+23", "CF4gp": "CF4group", "CLONES": "Sorce+2018", "IndranilVoid_exp": "Exponential", "IndranilVoid_gauss": "Gaussian", "IndranilVoid_mb": "Maxwell-Boltzmann", } if isinstance(simname, list): 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": r"SFI\texttt{++}", "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, fname): """Read in the goodness-of-fit statistics `kind`.""" if kind not in ["BIC", "AIC", "neg_lnZ_harmonic"]: raise ValueError("`kind` must be one of 'BIC', 'AIC', 'neg_lnZ_harmonic'.") # noqa with File(fname, 'r') as f: return f[f"gof/{kind}"][()] ############################################################################### # Read in samples # ############################################################################### def get_samples(fname, convert_Vext_to_galactic=True): """Read in the samples from the HDF5 file.""" samples = {} with File(fname, 'r') as f: grp = f["samples"] for key in grp.keys(): samples[key] = grp[key][...] 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 ############################################################################### # 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(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)) # 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] Vext = Vext[::downsample] beta = beta[::downsample] # Multiply the simulation velocities by beta. Bx = Bx[..., None] * beta By = By[..., None] * beta Bz = Bz[..., None] * beta # 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) 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) # Stack over the simulations B = np.hstack([B[i] for i in range(len(B))]) return r, B ############################################################################### # Prepare samples for plotting # ############################################################################### 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, keys def samples_to_getdist(samples, label): data, __, keys = samples_for_corner(samples) return MCSamples( samples=data, names=keys, labels=names_to_latex(keys, for_corner=False), label=label)