diff --git a/csiborgtools/flow/void_model.py b/csiborgtools/flow/void_model.py index 0ff43a8..794bb95 100644 --- a/csiborgtools/flow/void_model.py +++ b/csiborgtools/flow/void_model.py @@ -19,11 +19,44 @@ from os.path import join from re import search import numpy as np +from astropy.coordinates import SkyCoord, angular_separation from jax import numpy as jnp from jax import vmap from jax.scipy.ndimage import map_coordinates +from scipy.interpolate import RegularGridInterpolator from tqdm import tqdm +from ..utils import galactic_to_radec +from ..params import SPEED_OF_LIGHT +from .cosmography import distmod2dist, distmod2redshift + +############################################################################### +# Basic void computations # +############################################################################### + + +def angular_distance_from_void_axis(RA, dec): + """ + Calculate the angular distance of a galaxy from the void axis, all in + degrees. + """ + # Calculate the separation angle between the galaxy and the model axis. + model_axis = SkyCoord(l=117, b=4, frame='galactic', unit='deg').icrs + coords = SkyCoord(ra=RA, dec=dec, unit='deg').icrs + return angular_separation( + coords.ra.rad, coords.dec.rad, + model_axis.ra.rad, model_axis.dec.rad) * 180 / np.pi + + +def select_void_h(kind): + """Select 'little h' for void profile `kind`.""" + hs = {"mb": 0.7615, "gauss": 0.7724, "exp": 0.7725} + try: + return hs[kind] + except KeyError: + raise ValueError(f"Unknown void kind: `{kind}`.") + + ############################################################################### # I/O of the void data # ############################################################################### @@ -97,9 +130,9 @@ def interpolate_void(rLG, r, phi, data, rgrid_min, rgrid_max, rLG_min, rLG_max, ---------- rLG : float The observer's distance from the center of the void. - r : 1-dimensional array + r : 1-dimensional array of shape `(nsteps,) The radial distances at which to interpolate the velocities. - phi : 1-dimensional array + phi : 1-dimensional array of shape `(ngal,)` The angles at which to interpolate the velocities, in degrees, defining the galaxy position. data : 3-dimensional array of shape (nLG, nrad, nphi) @@ -114,7 +147,7 @@ def interpolate_void(rLG, r, phi, data, rgrid_min, rgrid_max, rLG_min, rLG_max, Returns ------- - vel : 2-dimensional array of shape (len(phi), len(r)) + vel : 2-dimensional array of shape `(ngal, nsteps)` """ nLG, nrad, nphi = data.shape @@ -139,3 +172,106 @@ def interpolate_void(rLG, r, phi, data, rgrid_min, rgrid_max, rLG_min, rLG_max, return map_coordinates(data, X, order=order, mode='nearest') return vmap(interpolate_single_phi)(phi) + + +############################################################################### +# Mock void data # +############################################################################### + + +def mock_void(vrad_data, rLG_index, profile, + a_TF=-22.8, b_TF=-7.2, sigma_TF=0.1, sigma_v=100., + mean_eta=0.069, std_eta=0.078, mean_e_eta=0.012, + mean_mag=10.31, std_mag=0.83, mean_e_mag=0.044, + bmin=None, add_malmquist=False, nsamples=2000, seed=42, + Om0=0.3175, verbose=False, **kwargs): + """Mock 2MTF-like TFR data with void velocities.""" + truths = {"a": a_TF, "b": b_TF, "e_mu": sigma_TF, "sigma_v": sigma_v, + "mean_eta": mean_eta, "std_eta": std_eta, + "mean_mag": mean_mag, "std_mag": std_mag, + } + + gen = np.random.default_rng(seed) + + # Sample the sky-distribution, either full-sky or mask out the Galactic + # plane. + l = gen.uniform(0, 360, size=nsamples) # noqa + if bmin is None: + b = np.arcsin(gen.uniform(-1, 1, size=nsamples)) + else: + b = np.arcsin(gen.uniform(np.sin(np.deg2rad(bmin)), 1, + size=nsamples)) + b[gen.rand(nsamples) < 0.5] *= -1 + + b = np.rad2deg(b) + + RA, DEC = galactic_to_radec(l, b) + # Calculate the angular separation from the void axis, in degrees. + phi = angular_distance_from_void_axis(RA, DEC) + + # Sample the linewidth of each galaxy from a Gaussian distribution to mimic + # the MNR procedure. + eta_true = gen.normal(mean_eta, std_eta, nsamples) + eta_obs = gen.normal(eta_true, mean_e_eta) + + # Subtract the mean of the observed linewidths, so that they are + # centered around zero. For consistency subtract from both observed + # and true values. + eta_mean_sampled = np.mean(eta_obs) + eta_true -= eta_mean_sampled + eta_obs -= eta_mean_sampled + + # Sample the magnitude from some Gaussian distribution to replicate MNR. + mag_true = gen.normal(mean_mag, std_mag, nsamples) + mag_obs = gen.normal(mag_true, mean_e_mag) + + # Calculate the 'true' distance modulus and redshift from the TFR distance. + mu_TFR = mag_true - (a_TF + b_TF * eta_true) + if add_malmquist: + raise NotImplementedError("Malmquist bias not implemented yet.") + else: + mu_true = gen.normal(mu_TFR, sigma_TF) + + # Convert the true distance modulus to true distance and cosmological + # redshift. + r = distmod2dist(mu_true, Om0) + zcosmo = distmod2redshift(mu_true, Om0) + + # Little h of this void profile + h = select_void_h(profile) + + # Extract the velocities for the galaxies from the grid for this LG + # index. + vrad_data_rLG = vrad_data[rLG_index] + + r_grid = np.arange(0, 251) * h + phi_grid = np.arange(0, 181) + Vr = RegularGridInterpolator((r_grid, phi_grid), vrad_data_rLG, + fill_value=np.nan, bounds_error=False, + method="cubic")(np.vstack([r, phi]).T) + + # The true redshift of the source. + zCMB_true = (1 + zcosmo) * (1 + Vr / SPEED_OF_LIGHT) - 1 + zCMB_obs = gen.normal(zCMB_true, sigma_v / SPEED_OF_LIGHT) + + sample = {"RA": RA, + "DEC": DEC, + "z_CMB": zCMB_obs, + "eta": eta_obs, + "mag": mag_obs, + "e_eta": np.ones(nsamples) * mean_e_eta, + "e_mag": np.ones(nsamples) * mean_e_mag, + "r": r, + "distmod_true": mu_true, + "distmod_TFR": mu_TFR} + + # Apply a true distance cut to the mocks. + mask = r < np.max(r_grid) + for key in sample: + sample[key] = sample[key][mask] + + if verbose and np.any(~mask): + print(f"Removed {(~mask).sum()} out of {mask.size} samples " + "due to the true distance cutoff.") + + return sample, truths