426 lines
18 KiB
Python
426 lines
18 KiB
Python
import aquila_borg as borg
|
||
import numpy as np
|
||
from astropy.coordinates import SkyCoord
|
||
import astropy.units as apu
|
||
import jax.numpy as jnp
|
||
import jax
|
||
import configparser
|
||
import ast
|
||
|
||
import borg_velocity.poisson_process as poisson_process
|
||
import borg_velocity.projection as projection
|
||
import borg_velocity.utils as utils
|
||
from borg_velocity.utils import myprint
|
||
|
||
|
||
def tfr_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))
|
||
|
||
"""
|
||
|
||
# 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
|
||
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}')
|
||
|
||
# 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('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
|
||
|
||
|
||
|
||
def sn_create_mock(Nt, L, xmin, cpar, dens, vel, Rmax, alpha,
|
||
a_tripp, b_tripp, M_SN, sigma_SN, sigma_m, sigma_stretch, sigma_c,
|
||
hyper_stretch_mu, hyper_stretch_sigma, hyper_c_mu, hyper_c_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
|
||
- a_tripp (float): Coefficient of stretch in the Tripp relation
|
||
- b_tripp (float): Coefficient of colour in the Tripp relation
|
||
- M_SN (float): Absolute magnitude of supernovae
|
||
- sigma_SN (float): Intrinsic scatter in the Tripp relation
|
||
- sigma_m (float): Uncertainty on the apparent magnitude measurements
|
||
- sigma_stretch (float): Uncertainty on the stretch measurements
|
||
- sigma_c (float): Uncertainty on the colour measurements
|
||
- hyper_stretch_mu (float): Mean of Gaussian hyper prior for the true stretch values
|
||
- hyper_stretch_sigma (float): Std of Gaussian hyper prior for the true stretch values
|
||
- hyper_c_mu (float): Mean of hyper Gaussian prior for the true colour values
|
||
- hyper_c_sigma (float): Std of Gaussian hyper prior for the true colour 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_mobs (np.ndarrary): Observed apparent magnitudes of the tracers (shape = (Nt,))
|
||
|
||
- all_xtrue (np.ndarrary): True comoving coordinates of the tracers (Mpc/h) (shape = (3, Nt))
|
||
|
||
"""
|
||
|
||
# Initialize lists to store valid positions and corresponding sig_mu values
|
||
all_xtrue = np.empty((3, Nt))
|
||
all_mtrue = np.empty(Nt)
|
||
all_stretchtrue = np.empty(Nt)
|
||
all_ctrue = np.empty(Nt)
|
||
all_mobs = np.empty(Nt)
|
||
all_stretchobs = np.empty(Nt)
|
||
all_cobs = 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 stretch and colour (c) from its prior
|
||
stretchtrue = hyper_stretch_mu + hyper_stretch_sigma * np.random.randn(Nt)
|
||
ctrue = hyper_c_mu + hyper_c_sigma * np.random.randn(Nt)
|
||
|
||
# Obtain muSN from mutrue using the intrinsic scatter
|
||
muSN = mutrue + sigma_SN * np.random.randn(Nt)
|
||
|
||
# Obtain apparent magnitude from the TFR
|
||
mtrue = muSN - (a_tripp * stretchtrue - b_tripp * ctrue) + M_SN
|
||
|
||
# Scatter true observed apparent magnitudes and linewidths
|
||
mobs = mtrue + sigma_m * np.random.randn(Nt)
|
||
stretchobs = stretchtrue + sigma_stretch * np.random.randn(Nt)
|
||
cobs = ctrue + sigma_c * np.random.randn(Nt)
|
||
|
||
# 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_stretchtrue[imin:imax] = stretchtrue[:selected_count]
|
||
all_ctrue[imin:imax] = ctrue[:selected_count]
|
||
all_mobs[imin:imax] = mobs[:selected_count]
|
||
all_stretchobs[imin:imax] = stretchobs[:selected_count]
|
||
all_cobs[imin:imax] = cobs[: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}')
|
||
|
||
# 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('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_stretchtrue, all_ctrue, all_mobs, all_stretchobs, all_cobs, all_xtrue
|
||
|
||
|
||
|
||
def borg_mock(s_hat, state, fwd_model, fwd_vel, ini_file, seed=None):
|
||
"""
|
||
WRITE DOCSTRING
|
||
"""
|
||
|
||
myprint('Making mock from BORG')
|
||
|
||
config = configparser.ConfigParser()
|
||
config.read(ini_file)
|
||
|
||
# Ensure cosmology is correct
|
||
cosmo = utils.get_cosmopar(ini_file)
|
||
fwd_model.setCosmoParams(cosmo)
|
||
|
||
# Run BORG density field
|
||
dens = np.zeros(fwd_model.getOutputBoxModel().N)
|
||
fwd_model.forwardModel_v2(s_hat)
|
||
fwd_model.getDensityFinal(dens)
|
||
state["BORG_final_density"][:] = dens
|
||
|
||
# Get velocity field
|
||
cosmo = utils.get_cosmopar(ini_file)
|
||
vel = fwd_vel.getVelocityField()
|
||
|
||
# Obtain a bulk velocity
|
||
vbulk = np.array(ast.literal_eval(config['model']['bulk_flow'])).reshape((3, 1, 1, 1))
|
||
|
||
if seed is not None:
|
||
np.random.seed(seed)
|
||
|
||
L = float(config['system']['L0'])
|
||
xmin = float(config['system']['corner0'])
|
||
nsamp = int(config['run']['nsamp'])
|
||
Rmax = float(config['mock']['R_max'])
|
||
sigma_v = float(config['model']['sig_v'])
|
||
interp_order = int(config['model']['interp_order'])
|
||
bias_epsilon = float(config['model']['bias_epsilon'])
|
||
|
||
tracer_type = [None] * nsamp
|
||
RA = [None] * nsamp
|
||
Dec = [None] * nsamp
|
||
czCMB = [None] * nsamp
|
||
m_true = [None] * nsamp
|
||
m_obs = [None] * nsamp
|
||
eta_true = [None] * nsamp
|
||
eta_obs = [None] * nsamp
|
||
stretch_true = [None] * nsamp
|
||
stretch_obs = [None] * nsamp
|
||
c_true = [None] * nsamp
|
||
c_obs = [None] * nsamp
|
||
xtrue = [None] * nsamp
|
||
|
||
for i in range(nsamp):
|
||
|
||
myprint(f'Making mock for sample {i}')
|
||
|
||
Nt = int(config[f'sample_{i}']['Nt'])
|
||
alpha = float(config[f'sample_{i}']['alpha'])
|
||
tracer_type[i] = config[f'sample_{i}']['tracer_type']
|
||
|
||
myprint(f'Tracer type: {tracer_type[i]}')
|
||
|
||
if tracer_type[i] == 'tfr':
|
||
mthresh = float(config[f'sample_{i}']['mthresh'])
|
||
a_TFR = float(config[f'sample_{i}']['a_tfr'])
|
||
b_TFR = float(config[f'sample_{i}']['b_tfr'])
|
||
sigma_TFR = float(config[f'sample_{i}']['sigma_tfr'])
|
||
sigma_m = float(config[f'sample_{i}']['sigma_m'])
|
||
sigma_eta = float(config[f'sample_{i}']['sigma_eta'])
|
||
hyper_eta_mu = float(config[f'sample_{i}']['hyper_eta_mu'])
|
||
hyper_eta_sigma = float(config[f'sample_{i}']['hyper_eta_sigma'])
|
||
mthresh = float(config[f'sample_{i}']['mthresh'])
|
||
RA[i], Dec[i], czCMB[i], m_true[i], eta_true[i], m_obs[i], eta_obs[i], xtrue[i] = tfr_create_mock(
|
||
Nt, L, xmin, cosmo, dens, vel + vbulk,
|
||
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)
|
||
elif tracer_type[i] == 'sn':
|
||
a_tripp = float(config[f'sample_{i}']['a_tripp'])
|
||
b_tripp = float(config[f'sample_{i}']['b_tripp'])
|
||
M_SN = float(config[f'sample_{i}']['m_sn'])
|
||
sigma_SN = float(config[f'sample_{i}']['sigma_sn'])
|
||
sigma_m = float(config[f'sample_{i}']['sigma_m'])
|
||
sigma_stretch = float(config[f'sample_{i}']['sigma_stretch'])
|
||
sigma_c = float(config[f'sample_{i}']['sigma_c'])
|
||
hyper_stretch_mu = float(config[f'sample_{i}']['hyper_stretch_mu'])
|
||
hyper_stretch_sigma = float(config[f'sample_{i}']['hyper_stretch_sigma'])
|
||
hyper_c_mu = float(config[f'sample_{i}']['hyper_c_mu'])
|
||
hyper_c_sigma = float(config[f'sample_{i}']['hyper_c_sigma'])
|
||
RA[i], Dec[i], czCMB[i], m_true[i], stretch_true[i], c_true[i], m_obs[i], stretch_obs[i], c_obs[i], xtrue[i] = sn_create_mock(
|
||
Nt, L, xmin, cosmo, dens, vel + vbulk, Rmax, alpha,
|
||
a_tripp, b_tripp, M_SN, sigma_SN, sigma_m, sigma_stretch, sigma_c,
|
||
hyper_stretch_mu, hyper_stretch_sigma, hyper_c_mu, hyper_c_sigma,
|
||
sigma_v, interp_order=1, bias_epsilon=1e-7)
|
||
else:
|
||
raise NotImplementedError
|
||
|
||
return tracer_type, RA, Dec, czCMB, m_true, m_obs, eta_true, eta_obs, stretch_true, stretch_obs, c_true, c_obs, xtrue
|