import aquila_borg as borg import numpy as np from astropy.coordinates import SkyCoord import astropy.units as apu import pandas as pd import jax.numpy as jnp import jax.scipy.special import matplotlib.pyplot as plt import corner import numpyro import numpyro.distributions as dist from jax import lax, random import borg_velocity.poisson_process as poisson_process import borg_velocity.projection as projection import borg_velocity.utils as utils # Output stream management cons = borg.console() myprint = lambda x: cons.print_std(x) if type(x) == str else cons.print_std(repr(x)) def build_gravity_model(box, cpar, ai=0.05, af=1.0, nsteps=20, forcesampling=4, supersampling=2, rsmooth=4.0, gravity='lpt', velmodel_name='CICModel'): """ Builds the gravity model and returns the forward model chain. Args: - box (borg.forward.BoxModel): Box within which to run simulation - cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters to use - ai (float, default=0.05): Scale factor to begin simulation - af (float, default=1.0): Scale factor to end simulation - nsteps (int, default=20): Number of steps to use in the simulation - forcesampling (int, default=4): Sampling factor for force evaluations - supersampling (int, default=2): Supersampling factor of particles - rsmooth (float, default=4.0): Smoothing scale for velocity field (Mpc/h) - gravity (str, default='lpt'): Which gravity model to use - velmodel_name (str, default='CICModel'): Which velocity estimator to use Returns: - chain (borg.forward.BaseForwardModel): The forward model for density - fwd_vel (borg.forward.VelocityBase): The forward model for velocity """ myprint(f"Building gravity model {gravity}") # Setup forward model chain = borg.forward.ChainForwardModel(box) chain.addModel(borg.forward.models.HermiticEnforcer(box)) # CLASS transfer function chain @= borg.forward.model_lib.M_PRIMORDIAL_AS(box) transfer_class = borg.forward.model_lib.M_TRANSFER_CLASS(box, opts=dict(a_transfer=1.0)) transfer_class.setModelParams({"extra_class_arguments":{'YHe':'0.24'}}) chain @= transfer_class if gravity == 'linear': raise NotImplementedError(gravity) elif gravity == 'lpt': mod = borg.forward.model_lib.M_LPT_CIC( box, opts=dict(a_initial=af, a_final=af, do_rsd=False, supersampling=supersampling, lightcone=False, part_factor=1.01,)) elif gravity == '2lpt': mod = borg.forward.model_lib.M_2LPT_CIC( box, opts=dict(a_initial=af, a_final=af, do_rsd=False, supersampling=supersampling, lightcone=False, part_factor=1.01,)) elif gravity == 'pm': mod = borg.forward.model_lib.M_PM_CIC( box, opts=dict(a_initial=af, a_final=af, do_rsd=False, supersampling=supersampling, part_factor=1.01, forcesampling=forcesampling, pm_start_z=1/ai - 1, pm_nsteps=nsteps, tcola=False)) elif gravity == 'cola': mod = borg.forward.model_lib.M_PM_CIC( box, opts=dict(a_initial=af, a_final=af, do_rsd=False, supersampling=supersampling, part_factor=1.01, forcesampling=forcesampling, pm_start_z=1/ai - 1, pm_nsteps=nsteps, tcola=True)) else: raise NotImplementedError(gravity) mod.accumulateAdjoint(True) chain @= mod # Cosmological parameters chain.setCosmoParams(cpar) # This is the forward model for velocity velmodel = getattr(borg.forward.velocity, velmodel_name) if velmodel_name == 'LinearModel': fwd_vel = velmodel(box, mod, af) elif velmodel_name == 'CICModel': fwd_vel = velmodel(box, mod, rsmooth) else: fwd_vel = velmodel(box, mod) return chain, fwd_vel def get_fields(L, N, xmin, gravity='lpt', velmodel_name='CICModel'): """ Obtain a density and velocity field to use for mock Args: - L (float): Box length (Mpc/h) - N (int): Number of grid cells per side - xmin (float): Coordinate of corner of the box (Mpc/h) - gravity (str, default='lpt'): Which gravity model to use - velmodel_name (str, default='CICModel'): Which velocity estimator to use Returns: - cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters to use - output_density (np.ndarray): Over-density field - output_velocity (np.ndarray): Velocity field """ # Setup box and cosmology cpar = borg.cosmo.CosmologicalParameters() cpar.default() box = borg.forward.BoxModel() box.L = (L, L, L) box.N = (N, N, N) box.xmin = (xmin, xmin, xmin) # Get forward models fwd_model, fwd_vel = build_gravity_model(box, cpar, gravity=gravity, velmodel_name=velmodel_name) # Make some initial conditions s_hat = np.fft.rfftn(np.random.randn(*box.N)) / box.Ntot ** (0.5) # Obtain density and velocity fields output_density = np.zeros(box.N) fwd_model.forwardModel_v2(s_hat) fwd_model.getDensityFinal(output_density) output_velocity = fwd_vel.getVelocityField() return cpar, output_density, output_velocity def create_mock(Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh, a_TFR, b_TFR, sigma_TFR, sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma, sigma_v, interp_order=1, bias_epsilon=1e-7): """ Create mock TFR catalogue from a density and velocity field Args: - Nt (int): Number of tracers to produce - L (float): Box length (Mpc/h) - xmin (float): Coordinate of corner of the box (Mpc/h) - cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters to use - dens (np.ndarray): Over-density field (shape = (N, N, N)) - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N)) - Rmax (float): Maximum allowed comoving radius of a tracer (Mpc/h) - alpha (float): Exponent for bias model - mthresh (float): Threshold absolute magnitude in selection - a_TFR (float): TFR relation intercept - b_TFR (float): TFR relation slope - sigma_TFR (float): Intrinsic scatter in the TFR - sigma_m (float): Uncertainty on the apparent magnitude measurements - sigma_eta (float): Uncertainty on the linewidth measurements - hyper_eta_mu (float): Mean of the Gaussian hyper-prior for the true eta values - hyper_eta_sigma (float): Std deviation of the Gaussian hyper-prior for the true eta values - sigma_v (float): Uncertainty on the velocity field (km/s) - interp_order (int, default=1): Order of interpolation from grid points to the line of sight - bias_epsilon (float, default=1e-7): Small number to add to 1 + delta to prevent 0^# Returns: - all_RA (np.ndarrary): Right Ascension (degrees) of the tracers (shape = (Nt,)) - all_Dec (np.ndarrary): Dec (np.ndarray): Delination (degrees) of the tracers (shape = (Nt,)) - czCMB (np.ndarrary): Observed redshifts (km/s) of the tracers (shape = (Nt,)) - all_mtrue (np.ndarrary): True apparent magnitudes of the tracers (shape = (Nt,)) - all_etatrue (np.ndarrary): True linewidths of the tracers (shape = (Nt,)) - all_mobs (np.ndarrary): Observed apparent magnitudes of the tracers (shape = (Nt,)) - all_etaobs (np.ndarrary): Observed linewidths of the tracers (shape = (Nt,)) - all_xtrue (np.ndarrary): True comoving coordinates of the tracers (Mpc/h) (shape = (3, Nt)) - vbulk (np.ndarray): The bulk velocity of the box (km/s) """ # Initialize lists to store valid positions and corresponding sig_mu values all_xtrue = np.empty((3, Nt)) all_mtrue = np.empty(Nt) all_etatrue = np.empty(Nt) all_mobs = np.empty(Nt) all_etaobs = np.empty(Nt) all_RA = np.empty(Nt) all_Dec = np.empty(Nt) # Counter for accepted positions accepted_count = 0 # Bias model phi = (1. + dens + bias_epsilon) ** alpha # Only use centre of box x = np.linspace(xmin, xmin + L, dens.shape[0]+1) i0 = np.argmin(np.abs(x + Rmax)) i1 = np.argmin(np.abs(x - Rmax)) L_small = x[i1] - x[i0] xmin_small = x[i0] phi_small = phi[i0:i1, i0:i1, i0:i1] # Loop until we have Nt valid positions while accepted_count < Nt: # Generate positions (comoving) xtrue = poisson_process.sample_3d(phi_small, Nt, L_small, (xmin_small, xmin_small, xmin_small)) # Convert to RA, Dec, Distance (comoving) rtrue = np.sqrt(np.sum(xtrue** 2, axis=0)) # Mpc/h c = SkyCoord(x=xtrue[0], y=xtrue[1], z=xtrue[2], representation_type='cartesian') RA = c.spherical.lon.degree Dec = c.spherical.lat.degree r_hat = np.array(SkyCoord(ra=RA*apu.deg, dec=Dec*apu.deg).cartesian.xyz) # Compute cosmological redshift zcosmo = utils.z_cos(rtrue, cpar.omega_m) # Compute luminosity distance # DO I NEED TO DO /h??? dL = (1 + zcosmo) * rtrue / cpar.h # Mpc # Compute true distance modulus mutrue = 5 * np.log10(dL) + 25 # Sample true linewidth (eta) from its prior etatrue = hyper_eta_mu + hyper_eta_sigma * np.random.randn(Nt) # Obtain muTFR from mutrue using the intrinsic scatter muTFR = mutrue + sigma_TFR * np.random.randn(Nt) # Obtain apparent magnitude from the TFR mtrue = muTFR + (a_TFR + b_TFR * etatrue) # Scatter true observed apparent magnitudes and linewidths mobs = mtrue + sigma_m * np.random.randn(Nt) etaobs = etatrue + sigma_eta * np.random.randn(Nt) # Apply apparement magnitude cut m = mobs <= mthresh # m = np.ones(mobs.shape, dtype=bool) mtrue = mtrue[m] etatrue = etatrue[m] mobs = mobs[m] etaobs = etaobs[m] xtrue = xtrue[:,m] RA = RA[m] Dec = Dec[m] # Calculate how many valid positions we need to reach Nt remaining_needed = Nt - accepted_count selected_count = min(xtrue.shape[1], remaining_needed) # Append only the needed number of valid positions imin = accepted_count imax = accepted_count + selected_count all_xtrue[:,imin:imax] = xtrue[:,:selected_count] all_mtrue[imin:imax] = mtrue[:selected_count] all_etatrue[imin:imax] = etatrue[:selected_count] all_mobs[imin:imax] = mobs[:selected_count] all_etaobs[imin:imax] = etaobs[:selected_count] all_RA[imin:imax] = RA[:selected_count] all_Dec[imin:imax] = Dec[:selected_count] # Update the accepted count accepted_count += selected_count myprint(f'\tMade {accepted_count} of {Nt}') # Obtain a bulk velocity vhat = np.random.randn(3) vhat = vhat / np.linalg.norm(vhat) vbulk = np.random.randn() * utils.get_sigma_bulk(L, cpar) vbulk = vhat * vbulk # Get the radial component of the peculiar velocity at the positions of the objects myprint('Obtaining peculiar velocities') tracer_vel = projection.interp_field( vel, np.expand_dims(all_xtrue, axis=2), L, np.array([xmin, xmin, xmin]), interp_order ) # km/s myprint('Adding bulk velocity') tracer_vel = tracer_vel + vbulk[:,None,None] myprint('Radial projection') vr_true = np.squeeze(projection.project_radial( tracer_vel, np.expand_dims(all_xtrue, axis=2), np.zeros(3,) )) # km/s # Recompute cosmological redshift rtrue = jnp.sqrt(jnp.sum(all_xtrue ** 2, axis=0)) zcosmo = utils.z_cos(rtrue, cpar.omega_m) # Obtain total redshift vr_noised = vr_true + sigma_v * np.random.randn(Nt) czCMB = ((1 + zcosmo) * (1 + vr_noised / utils.speed_of_light) - 1) * utils.speed_of_light return all_RA, all_Dec, czCMB, all_mtrue, all_etatrue, all_mobs, all_etaobs, all_xtrue, vbulk def estimate_data_parameters(): """ Using the 2MASS catalogue, estimate some parameters to use in mock generation. The file contains the following columns: ID 2MASS XSC ID name (HHMMSSss+DDMMSSs) RAdeg Right ascension (J2000) DEdeg Declination (J2000) cz2mrs Heliocentric redshift from the 2MRS (km/s) Kmag NIR magnitudes in the K band from the 2MRS (mag) Hmag NIR magnitudes in the H band from the 2MRS (mag) Jmag NIR magnitudes in the J band from the 2MRS (mag) e_Kmag Error of the NIR magnitudes in K band from the (mag) e_Hmag Error of the NIR magnitudes in H band from the (mag) e_Jmag Error of the NIR magnitudes in J band from the (mag) WHIc Corrected HI width (km/s) e_WHIc Error of corrected HI width (km/s) Returns: - sigma_m (float): Estimate of the uncertainty on the apparent magnitude measurements - sigma_eta (float): Estimate of the uncertainty on the linewidth measurements - hyper_eta_mu (float): Estimate of the mean of the Gaussian hyper-prior for the true eta values - hyper_eta_sigma (float): Estimate of the std deviation of the Gaussian hyper-prior for the true eta values - hyper_m_sigma (float): Estimate of the std deviation of the Gaussian hyper-prior for the true m values """ columns = ['ID', 'RAdeg', 'DEdeg', 'cz2mrs', 'Kmag', 'Hmag', 'Jmag', 'e_Kmag', 'e_Hmah', 'e_Jmag', 'WHIc', 'e_WHIc'] fname = '/data101/bartlett/fsigma8/PV_data/2MASS/table1.dat' df = pd.read_csv(fname, sep='\s+', names=columns) sigma_m = np.median(df['e_Kmag']) eta = np.log10(df['WHIc']) - 2.5 sigma_eta = np.median(df['e_WHIc'] / df['WHIc'] / np.log(10)) hyper_eta_mu = np.median(eta) hyper_eta_sigma = (np.percentile(eta, 84) - np.percentile(eta, 16)) / 2 hyper_m_sigma = np.amax(df['Kmag']) - np.percentile(df['Kmag'], 16) return sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma, hyper_m_sigma def generateMBData(RA, Dec, cz_obs, L, N, R_lim, Nsig, Nint_points, sigma_v, frac_sigma_r): """ Generate points along the line of sight of each tracer to enable marginalisation over distance uncertainty. The central distance is given such that the observed redshift equals the cosmological redshift at this distance. The range is then +/- Nsig * sig, where sig^2 = (sig_v/100)^2 + sig_r^2 and sig_v is the velocity uncertainty in km/s Args: - RA (np.ndarray): Right Ascension (degrees) of the tracers (shape = (Nt,)) - Dec (np.ndarray): Delination (degrees) of the tracers (shape = (Nt,)) - cz_obs (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,)) - L (float): Box length (Mpc/h) - N (int): Number of grid cells per side - R_lim (float): Maximum allowed (true) comoving distance of a tracer (Mpc/h) - Nsig (float): How many standard deviations away from the predicted radius to use - Nint_points (int): Number of radii over which to integrate the likelihood - sigma_v (float): Uncertainty on the velocity field (km/s) - frac_sigma_r (float): An estimate of the fractional uncertainty on the positions of tracers Returns: - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h). The shape is (3, Nt, Nsig) """ myprint(f"Making MB data") # Convert RA, DEC to radial vector r_hat = np.array(SkyCoord(ra=RA*apu.deg, dec=Dec*apu.deg).cartesian.xyz).T # Get min and max distance to integrate over # cz = 100 h r, so sigma_v corresponds to a sigma_r of ~ sigma_v / 100 robs = cz_obs / 100 sigr = np.sqrt((sigma_v / 100) ** 2 + (frac_sigma_r * robs)**2) rmin = robs - Nsig * sigr rmin = rmin.at[rmin <= 0].set(L / N / 100.) rmax = robs + Nsig * sigr rmax = rmax.at[rmax > R_lim].set(R_lim) # Compute coordinates of integration points r_integration = np.linspace(rmin, rmax, Nint_points) MB_pos = np.expand_dims(r_integration, axis=2) * r_hat[None,...] MB_pos = jnp.transpose(MB_pos, (2, 1, 0)) return MB_pos def likelihood_vel(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true, vbulk, dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, cz_obs, MB_pos, mthresh): """ Evaluate the terms in the likelihood from the velocity and malmquist bias Args: - alpha (float): Exponent for bias model - a_TFR (float): TFR relation intercept - b_TFR (float): TFR relation slope - sigma_TFR (float): Intrinsic scatter in the TFR - sigma_v (float): Uncertainty on the velocity field (km/s) - m_true (np.ndarray): True apparent magnitudes of the tracers (shape = (Nt,)) - eta_true (np.ndarray): True linewidths of the tracers (shape = (Nt,)) - vbulk (np.ndarray): Bulk velocity of the box (km/s) (shape=(3,)) - dens (np.ndarray): Over-density field (shape = (N, N, N)) - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N)) - omega_m (float): Matter density parameter Om - h (float): Hubble constant H0 = 100 h km/s/Mpc - L (float): Comoving box size (Mpc/h) - xmin (float): Coordinate of corner of the box (Mpc/h) - interp_order (int): Order of interpolation from grid points to the line of sight - bias_epsilon (float): Small number to add to 1 + delta to prevent 0^# - cz_obs (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,)) - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h). The shape is (3, Nt, Nsig) - mthresh (float): Threshold absolute magnitude in selection Returns: - loglike (float): The log-likelihood of the data """ # Comoving radii of integration points (Mpc/h) r = jnp.sqrt(jnp.sum(MB_pos ** 2, axis=0)) # p_r = r^2 n(r) N(mutrue; muTFR, sigmaTFR) # Multiply by arbitrary number for numerical stability (cancels in p_r / p_r_norm) number_density = projection.interp_field( dens, MB_pos, L, jnp.array([xmin, xmin, xmin]), interp_order, use_jitted=True, ) number_density = jax.nn.relu(1. + number_density) number_density = jnp.power(number_density + bias_epsilon, alpha) zcosmo = utils.z_cos(r, omega_m) mutrue = 5 * jnp.log10((1 + zcosmo) * r / h) + 25 muTFR = m_true - (a_TFR + b_TFR * eta_true) d2 = ((mutrue - muTFR[:,None]) / sigma_TFR) ** 2 best = jnp.amin(jnp.abs(d2), axis=1) d2 = d2 - jnp.expand_dims(jnp.nanmin(d2, axis=1), axis=1) p_r = r ** 2 * jnp.exp(-0.5 * d2) * number_density p_r_norm = jnp.expand_dims(jnp.trapezoid(p_r, r, axis=1), axis=1) # Peculiar velocity term tracer_vel = projection.interp_field( vel, MB_pos, L, jnp.array([xmin, xmin, xmin]), interp_order, use_jitted=True, ) tracer_vel = tracer_vel + jnp.squeeze(vbulk)[...,None,None] tracer_vr = projection.project_radial( tracer_vel, MB_pos, jnp.zeros(3,) ) cz_pred = ((1 + zcosmo) * (1 + tracer_vr / utils.speed_of_light) - 1) * utils.speed_of_light d2 = ((cz_pred - jnp.expand_dims(cz_obs, axis=1)) / sigma_v)**2 scale = jnp.nanmin(d2, axis=1) d2 = d2 - jnp.expand_dims(scale, axis=1) # Integrate to get likelihood p_cz = jnp.trapezoid(jnp.exp(-0.5 * d2) * p_r / p_r_norm, r, axis=1) lkl_ind = jnp.log(p_cz) - scale / 2 - 0.5 * jnp.log(2 * np.pi * sigma_v**2) loglike = lkl_ind.sum() return loglike def likelihood_m(m_true, m_obs, sigma_m, mthresh): """ Evaluate the terms in the likelihood from apparent magnitude Args: - m_true (np.ndarray): True apparent magnitudes of the tracers (shape = (Nt,)) - m_obs (np.ndarray): Observed apparent magnitudes of the tracers (shape = (Nt,)) - sigma_m (float or np.ndarray): Uncertainty on the apparent magnitude measurements - mthresh (float): Threshold absolute magnitude in selection Returns: - loglike (float): The log-likelihood of the data """ Nt = m_obs.shape[0] norm0 = ( jnp.log(2) - jnp.log(jax.scipy.special.erfc(- (mthresh - m_true) / (jnp.sqrt(2) * sigma_m))) - 0.5 * jnp.log(2 * jnp.pi * sigma_m ** 2) ) norm1 = jnp.where( jnp.ndim(sigma_m) == 0, Nt * 0.5 * jnp.log(2 * jnp.pi * sigma_m ** 2), jnp.sum(0.5 * jnp.log(2 * jnp.pi * sigma_m ** 2)) ) loglike = - ( 0.5 * jnp.sum((m_obs - m_true) ** 2 / sigma_m ** 2) - jnp.sum(norm0) + norm1 ) return loglike def likelihood_eta(eta_true, eta_obs, sigma_eta): """ Evaluate the terms in the likelihood from linewidth Args: - eta_true (np.ndarray): True linewidths of the tracers (shape = (Nt,)) - eta_obs (np.ndarray): Observed linewidths of the tracers (shape = (Nt,)) - sigma_eta (float or np.ndarray): Uncertainty on the linewidth measurements Returns: - loglike (float): The log-likelihood of the data """ Nt = eta_obs.shape[0] norm = jnp.where( jnp.ndim(sigma_eta) == 0, Nt * 0.5 * jnp.log(2 * jnp.pi * sigma_eta ** 2), jnp.sum(0.5 * jnp.log(2 * jnp.pi * sigma_eta ** 2)) ) loglike = - ( 0.5 * jnp.sum((eta_obs - eta_true) ** 2 / sigma_eta ** 2) + norm ) return loglike def likelihood(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true, vbulk, dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, cz_obs, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh): """ Evaluate the likelihood for TFR sample Args: - alpha (float): Exponent for bias model - a_TFR (float): TFR relation intercept - b_TFR (float): TFR relation slope - sigma_TFR (float): Intrinsic scatter in the TFR - sigma_v (float): Uncertainty on the velocity field (km/s) - m_true (np.ndarray): True apparent magnitudes of the tracers (shape = (Nt,)) - eta_true (np.ndarray): True linewidths of the tracers (shape = (Nt,)) - vbulk (np.ndarray): Bulk velocity of the box (km/s) (shape=(3,)) - dens (np.ndarray): Over-density field (shape = (N, N, N)) - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N)) - omega_m (float): Matter density parameter Om - h (float): Hubble constant H0 = 100 h km/s/Mpc - L (float): Comoving box size (Mpc/h) - xmin (float): Coordinate of corner of the box (Mpc/h) - interp_order (int): Order of interpolation from grid points to the line of sight - bias_epsilon (float): Small number to add to 1 + delta to prevent 0^# - cz_obs (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,)) - m_obs (np.ndarray): Observed apparent magnitudes of the tracers (shape = (Nt,)) - eta_obs (np.ndarray): Observed linewidths of the tracers (shape = (Nt,)) - sigma_m (float or np.ndarray): Uncertainty on the apparent magnitude measurements - sigma_eta (float or np.ndarray): Uncertainty on the linewidth measurements - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h). The shape is (3, Nt, Nsig) - mthresh (float): Threshold absolute magnitude in selection Returns: - loglike (float): The log-likelihood of the data """ loglike_vel = likelihood_vel(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true, vbulk, dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, cz_obs, MB_pos, mthresh) loglike_m = likelihood_m(m_true, m_obs, sigma_m, mthresh) loglike_eta = likelihood_eta(eta_true, eta_obs, sigma_eta) loglike = (loglike_vel + loglike_m + loglike_eta) return loglike def test_likelihood_scan(prior, alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true, vbulk, dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh): """ Plot likelihood as we scan through the paramaters [alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, mthresh] to verify that the likelihood shape looks reasonable Args: - prior (dict): Upper and lower bounds for a uniform prior for the parameters - alpha (float): Exponent for bias model - a_TFR (float): TFR relation intercept - b_TFR (float): TFR relation slope - sigma_TFR (float): Intrinsic scatter in the TFR - sigma_v (float): Uncertainty on the velocity field (km/s) - m_true (np.ndarray): True apparent magnitudes of the tracers (shape = (Nt,)) - eta_true (np.ndarray): True linewidths of the tracers (shape = (Nt,)) - dens (np.ndarray): Over-density field (shape = (N, N, N)) - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N)) - omega_m (float): Matter density parameter Om - h (float): Hubble constant H0 = 100 h km/s/Mpc - L (float): Comoving box size (Mpc/h) - xmin (float): Coordinate of corner of the box (Mpc/h) - interp_order (int): Order of interpolation from grid points to the line of sight - bias_epsilon (float): Small number to add to 1 + delta to prevent 0^# - czCMB (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,)) - m_obs (np.ndarray): Observed apparent magnitudes of the tracers (shape = (Nt,)) - eta_obs (np.ndarray): Observed linewidths of the tracers (shape = (Nt,)) - sigma_m (float or np.ndarray): Uncertainty on the apparent magnitude measurements - sigma_eta (float or np.ndarray): Uncertainty on the apparent linewidth measurements - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h). The shape is (3, Nt, Nsig) - mthresh (float): Threshold absolute magnitude in selection """ pars = [alpha, a_TFR, b_TFR, sigma_TFR, sigma_v] par_names = ['alpha', 'a_TFR', 'b_TFR', 'sigma_TFR', 'sigma_v'] orig_ll = - likelihood(*pars, m_true, eta_true, vbulk, dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh) for i, name in enumerate(par_names): myprint(f'Scanning {name}') if name in prior: x = np.linspace(*prior[name], 20) else: pmin = pars[i] * 0.2 pmax = pars[i] * 2.0 x = np.linspace(pmin, pmax, 20) all_ll = np.empty(x.shape) orig_x = pars[i] for j, xx in enumerate(x): pars[i] = xx all_ll[j] = - likelihood(*pars, m_true, eta_true, vbulk, dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh) pars[i] = orig_x plt.figure() plt.plot(x, all_ll, '.') plt.axvline(orig_x, ls='--', color='k') plt.axhline(orig_ll, ls='--', color='k') plt.xlabel(name) plt.ylabel('Negative log-likelihood') plt.savefig(f'tfr_likelihood_scan_{name}.png') fig = plt.gcf() plt.clf() plt.close(fig) # Now check the effect of varying mthresh on the likelihood myprint(f'Scanning mthresh') x = np.linspace(mthresh - 0.5, mthresh + 0.5, 20) all_ll = np.empty(x.shape) for j, xx in enumerate(x): all_ll[j] = - likelihood(*pars, m_true, eta_true, vbulk, dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, xx) plt.figure() plt.plot(x, all_ll, '.') plt.axvline(mthresh, ls='--', color='k') plt.axhline(orig_ll, ls='--', color='k') plt.xlabel('mthresh') plt.ylabel('Negative log-likelihood') plt.savefig(f'tfr_likelihood_scan_mthresh.png') fig = plt.gcf() plt.clf() plt.close(fig) return def run_mcmc(num_warmup, num_samples, prior, initial, dens, vel, cpar, L, xmin, interp_order, bias_epsilon, czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh,): """ Run MCMC over the model parameters Args: - num_warmup (int): Number of warmup steps to take in the MCMC - num_samples (int): Number of samples to take in the MCMC - prior (dict): Upper and lower bounds for a uniform prior for the parameters - initial (dict): Initial values for the MCMC - dens (np.ndarray): Over-density field (shape = (N, N, N)) - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N)) - cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters to use - L (float): Comoving box size (Mpc/h) - xmin (float): Coordinate of corner of the box (Mpc/h) - interp_order (int): Order of interpolation from grid points to the line of sight - bias_epsilon (float): Small number to add to 1 + delta to prevent 0^# - czCMB (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,)) - m_obs (np.ndarray): Observed apparent magnitudes of the tracers (shape = (Nt,)) - eta_obs (np.ndarray): Observed linewidths of the tracers (shape = (Nt,)) - sigma_m (float or np.ndarray): Uncertainty on the apparent magnitude measurements - sigma_eta (float or np.ndarray): Uncertainty on the apparent linewidth measurements - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h). The shape is (3, Nt, Nsig) - mthresh (float): Threshold absolute magnitude in selection Returns: - mcmc (numpyro.infer.MCMC): MCMC object which has been run """ Nt = eta_obs.shape[0] omega_m = cpar.omega_m h = cpar.h sigma_bulk = utils.get_sigma_bulk(L, cpar) def tfr_model(): alpha = numpyro.sample("alpha", dist.Uniform(*prior['alpha'])) a_TFR = numpyro.sample("a_TFR", dist.Uniform(*prior['a_TFR'])) b_TFR = numpyro.sample("b_TFR", dist.Uniform(*prior['b_TFR'])) sigma_TFR = numpyro.sample("sigma_TFR", dist.HalfCauchy(1.0)) # sigma_v = numpyro.sample("sigma_v", dist.HalfCauchy(1.0)) sigma_v = numpyro.sample("sigma_v", dist.Uniform(*prior['sigma_v'])) hyper_mean_m = numpyro.sample("hyper_mean_m", dist.Uniform(*prior['hyper_mean_m'])) hyper_sigma_m = numpyro.sample("hyper_sigma_m", dist.HalfCauchy(1.0)) # Equivalent to 1/sigma prior hyper_mean_eta = numpyro.sample("hyper_mean_eta", dist.Uniform(*prior['hyper_mean_eta'])) hyper_sigma_eta = numpyro.sample("hyper_sigma_eta", dist.HalfCauchy(1.0)) # Equivalent to 1/sigma prior # Sample correlation matrix using LKJ prior L_corr = numpyro.sample("L_corr", dist.LKJCholesky(2, concentration=1.0)) # Cholesky factor of correlation matrix corr_matrix = L_corr @ L_corr.T # Convert to full correlation matrix # Construct full covariance matrix: Σ = D * Corr * D hyper_mean = jnp.array([hyper_mean_m, hyper_mean_eta]) hyper_sigma = jnp.array([hyper_sigma_m, hyper_sigma_eta]) hyper_cov = jnp.diag(hyper_sigma) @ corr_matrix @ jnp.diag(hyper_sigma) # Sample m_true and eta_true x = numpyro.sample("true_vars", dist.MultivariateNormal(hyper_mean, hyper_cov), sample_shape=(Nt,)) m_true = numpyro.deterministic("m_true", x[:, 0]) eta_true = numpyro.deterministic("eta_true", x[:, 1]) # Sample bulk velocity vbulk_x = numpyro.sample("vbulk_x", dist.Normal(0, sigma_bulk / jnp.sqrt(3))) vbulk_y = numpyro.sample("vbulk_y", dist.Normal(0, sigma_bulk / jnp.sqrt(3))) vbulk_z = numpyro.sample("vbulk_z", dist.Normal(0, sigma_bulk / jnp.sqrt(3))) # Evaluate the likelihood numpyro.sample("obs", TFRLikelihood(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true, vbulk_x, vbulk_y, vbulk_z), obs=jnp.array([m_obs, eta_obs])) class TFRLikelihood(dist.Distribution): support = dist.constraints.real def __init__(self, alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true, vbulk_x, vbulk_y, vbulk_z): self.alpha, self.a_TFR, self.b_TFR, self.sigma_TFR, self.sigma_v, self.m_true, self.eta_true, self.vbulk_x, self.vbulk_y, self.vbulk_z = dist.util.promote_shapes(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true, vbulk_x, vbulk_y, vbulk_z) batch_shape = lax.broadcast_shapes( jnp.shape(alpha), jnp.shape(a_TFR), jnp.shape(b_TFR), jnp.shape(sigma_TFR), jnp.shape(sigma_v), jnp.shape(m_true), jnp.shape(eta_true), jnp.shape(vbulk_x), jnp.shape(vbulk_y), jnp.shape(vbulk_z), ) super(TFRLikelihood, self).__init__(batch_shape = batch_shape) def sample(self, key, sample_shape=()): raise NotImplementedError def log_prob(self, value): vbulk = jnp.array([self.vbulk_x, self.vbulk_y, self.vbulk_z]) loglike = likelihood(self.alpha, self.a_TFR, self.b_TFR, self.sigma_TFR, self.sigma_v, self.m_true, self.eta_true, vbulk, dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh) return loglike rng_key = random.PRNGKey(6) rng_key, rng_key_ = random.split(rng_key) values = initial values['true_vars'] = jnp.array([m_obs, eta_obs]).T values['L_corr'] = jnp.identity(2) values['vbulk_x'] = 0. values['vbulk_y'] = 0. values['vbulk_z'] = 0. myprint('Preparing MCMC kernel') kernel = numpyro.infer.NUTS(tfr_model, init_strategy=numpyro.infer.initialization.init_to_value(values=initial) ) mcmc = numpyro.infer.MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples) myprint('Running MCMC') mcmc.run(rng_key_) mcmc.print_summary() return mcmc def process_mcmc_run(mcmc, param_labels, truths, true_vars): """ Make summary plots from the MCMC and save these to file Args: - mcmc (numpyro.infer.MCMC): MCMC object which has been run - param_labels (list[str]): Names of the parameters to plot - truths (list[float]): True values of the parameters to plot. If unknown, then entry is None - true_vars (dict): True values of the observables to compare against inferred ones """ # Convert samples into a single array samples = mcmc.get_samples() samps = jnp.empty((len(samples[param_labels[0]]), len(param_labels))) for i, p in enumerate(param_labels): if p == 'hyper_corr': L_corr = samples['L_corr'] corr_matrix = jnp.matmul(L_corr, jnp.transpose(L_corr, (0, 2, 1))) samps = samps.at[:,i].set(corr_matrix[:,0,1]) else: samps = samps.at[:,i].set(samples[p]) # Trace plot of non-redshift quantities fig1, axs1 = plt.subplots(samps.shape[1], 1, figsize=(6,3*samps.shape[1]), sharex=True) axs1 = np.atleast_1d(axs1) for i in range(samps.shape[1]): axs1[i].plot(samps[:,i]) axs1[i].set_ylabel(param_labels[i]) if truths[i] is not None: axs1[i].axhline(truths[i], color='k') axs1[-1].set_xlabel('Step Number') fig1.tight_layout() fig1.savefig('tfr_trace.png') # Corner plot fig2, axs2 = plt.subplots(samps.shape[1], samps.shape[1], figsize=(20,20)) corner.corner( np.array(samps), labels=param_labels, fig=fig2, truths=truths ) fig2.savefig('tfr_corner.png') # True vs predicted for var in ['eta', 'm']: vname = var + '_true' if vname in samples.keys(): xtrue = true_vars[var] xpred_median = np.median(samples[vname], axis=0) xpred_plus = np.percentile(samples[vname], 84, axis=0) - xpred_median xpred_minus = xpred_median - np.percentile(samples[vname], 16, axis=0) fig3, axs3 = plt.subplots(2, 1, figsize=(10,8), sharex=True) plot_kwargs = {'fmt':'.', 'markersize':3, 'zorder':10, 'capsize':1, 'elinewidth':1, 'alpha':1} axs3[0].errorbar(xtrue, xpred_median, yerr=[xpred_minus, xpred_plus], **plot_kwargs) axs3[1].errorbar(xtrue, xpred_median - xtrue, yerr=[xpred_minus, xpred_plus], **plot_kwargs) axs3[1].set_xlabel('True') axs3[0].set_ylabel('Predicted') axs3[1].set_ylabel('Predicted - True') xlim = axs3[0].get_xlim() ylim = axs3[0].get_ylim() axs3[0].plot(xlim, xlim, color='k', zorder=0) axs3[0].set_xlim(xlim) axs3[0].set_ylim(ylim) axs3[1].axhline(0, color='k', zorder=0) fig3.suptitle(var) fig3.align_labels() fig3.tight_layout() fig3.savefig(f'tfr_true_predicted_{var}.png') return def main(): myprint('Beginning') # Get some parameters from the data sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma, hyper_m_sigma = estimate_data_parameters() # Other parameters to use L = 500.0 N = 64 xmin = -L/2 R_lim = L / 2 Rmax = 100 Nt = 100 alpha = 1.4 mthresh = 11.25 a_TFR = -23 b_TFR = -8.2 sigma_TFR = 0.3 sigma_v = 150 Nint_points = 201 Nsig = 10 frac_sigma_r = 0.07 # WANT A BETTER WAY OF DOING THIS - ESTIMATE THROUGH SIGMAS FROM TFR interp_order = 1 bias_epsilon = 1.e-7 num_warmup = 1000 num_samples = 2000 prior = { 'alpha': [0.5, 2.5], 'a_TFR': [-25, -20], 'b_TFR': [-10, -5], 'hyper_mean_eta': [hyper_eta_mu - 0.5, hyper_eta_mu + 0.5], 'hyper_mean_m':[mthresh - 5, mthresh + 5], 'sigma_v': [10, 3000], } initial = { 'alpha': alpha, 'a_TFR': a_TFR, 'b_TFR': b_TFR, 'hyper_mean_eta': hyper_eta_mu, 'hyper_sigma_eta': hyper_eta_sigma, 'hyper_mean_m': mthresh, 'hyper_sigma_m': hyper_m_sigma, 'sigma_TFR': sigma_TFR, 'sigma_v': sigma_v, } # Make mock np.random.seed(123) cpar, dens, vel = get_fields(L, N, xmin) RA, Dec, czCMB, m_true, eta_true, m_obs, eta_obs, xtrue, vbulk = create_mock( Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh, a_TFR, b_TFR, sigma_TFR, sigma_m, sigma_eta, hyper_eta_mu, hyper_eta_sigma, sigma_v, interp_order=interp_order, bias_epsilon=bias_epsilon) MB_pos = generateMBData(RA, Dec, czCMB, L, N, R_lim, Nsig, Nint_points, sigma_v, frac_sigma_r) # Test likelihood loglike = likelihood(alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true, vbulk, dens, vel, cpar.omega_m, cpar.h, L, xmin, interp_order, bias_epsilon, czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh) myprint(f'loglike {loglike}') # Scan over parameters to make plots verifying behaviour test_likelihood_scan(prior, alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, eta_true, vbulk, dens, vel, cpar.omega_m, cpar.h, L, xmin, interp_order, bias_epsilon, czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh) # Run a MCMC mcmc = run_mcmc(num_warmup, num_samples, prior, initial, dens, vel, cpar, L, xmin, interp_order, bias_epsilon, czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh,) param_labels = ['alpha', 'a_TFR', 'b_TFR', 'sigma_TFR', 'sigma_v', 'hyper_mean_m', 'hyper_sigma_m', 'hyper_mean_eta', 'hyper_sigma_eta', 'hyper_corr', 'vbulk_x', 'vbulk_y', 'vbulk_z'] truths = [alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, None, None, hyper_eta_mu, hyper_eta_sigma, None, vbulk[0], vbulk[1], vbulk[2]] true_vars = {'m':m_true, 'eta':eta_true} process_mcmc_run(mcmc, param_labels, truths, true_vars) if __name__ == "__main__": main()