623 lines
No EOL
24 KiB
Python
623 lines
No EOL
24 KiB
Python
import numpy as np
|
||
import jax.numpy as jnp
|
||
import configparser
|
||
import warnings
|
||
import aquila_borg as borg
|
||
import symbolic_pofk.linear
|
||
import jax
|
||
from functools import partial
|
||
import ast
|
||
|
||
import borg_velocity.utils as utils
|
||
from borg_velocity.utils import myprint
|
||
import borg_velocity.forwards as forwards
|
||
import borg_velocity.mock_maker as mock_maker
|
||
import borg_velocity.projection as projection
|
||
|
||
class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood):
|
||
"""
|
||
HADES likelihood for distance-tracers
|
||
"""
|
||
|
||
def __init__(self, fwd: borg.forward.BaseForwardModel, param_model: forwards.NullForward, ini_file: str) -> None:
|
||
"""
|
||
Initialises the VelocityBORGLikelihood class
|
||
|
||
Args:
|
||
- fwd (borg.forward.BaseForwardModel): The forward model to be used in the likelihood.
|
||
- param_model (forwards.NullForward): An empty forward model for storing model parameters.
|
||
- ini_file (str): The name of the ini file containing the model and borg parameters.
|
||
"""
|
||
|
||
self.ini_file = ini_file
|
||
|
||
myprint("Reading from configuration file: " + ini_file)
|
||
config = configparser.ConfigParser()
|
||
config.read(ini_file)
|
||
|
||
# Grid parameters
|
||
self.N = [int(config['system'][f'N{i}']) for i in range(3)]
|
||
self.L = [float(config['system'][f'L{i}']) for i in range(3)]
|
||
|
||
# Catalogues
|
||
self.nsamp = int(config['run']['nsamp'])
|
||
assert self.nsamp > 0, "Need at least one sample to run"
|
||
self.Nt = [int(config[f'sample_{i}']['Nt']) for i in range(self.nsamp)]
|
||
self.alpha = [float(config[f'sample_{i}']['alpha']) for i in range(self.nsamp)]
|
||
self.muA = [float(config[f'sample_{i}']['muA']) for i in range(self.nsamp)]
|
||
self.frac_sig_rhMpc = [float(config[f'sample_{i}']['frac_sig_rhMpc']) for i in range(self.nsamp)]
|
||
|
||
# What type of run we're doing
|
||
self.run_type = config['run']['run_type']
|
||
|
||
# Model parameters
|
||
self.ai = float(config['model']['ai'])
|
||
self.af = float(config['model']['af'])
|
||
self.sig_v = float(config['model']['sig_v'])
|
||
if config['model']['R_lim'] == 'none':
|
||
self.R_lim = self.L[0]/2
|
||
else:
|
||
self.R_lim = float(config['model']['R_lim'])
|
||
self.Nint_points = int(config['model']['nint_points'])
|
||
self.Nsig = float(config['model']['Nsig'])
|
||
self.sig_v = float(config['model']['sig_v'])
|
||
self.smooth_R = float(config['model']['smooth_R'])
|
||
self.bias_epsilon = float(config['model']['bias_epsilon'])
|
||
self.interp_order = int(config['model']['interp_order'])
|
||
|
||
# Load bulk flow
|
||
self.bulk_flow = np.array(ast.literal_eval(config['model']['bulk_flow']))
|
||
|
||
# For log-likelihood values
|
||
self.bignum = float(config['mcmc']['bignum'])
|
||
|
||
myprint(f" Init {self.N}, {self.L}")
|
||
super().__init__(fwd, self.N, self.L)
|
||
|
||
# Define the forward models
|
||
self.fwd = fwd
|
||
self.fwd_param = param_model
|
||
|
||
# Initialise model parameters
|
||
model_params = {
|
||
**{f'mua{i}':self.muA[i] for i in range(self.nsamp)},
|
||
**{f'alpha{i}':self.alpha[i] for i in range(self.nsamp)},
|
||
**{f'bulk_flow_{d}':self.bulk_flow[i] for i, d in enumerate(['x', 'y', 'z'])},
|
||
'sig_v':self.sig_v
|
||
}
|
||
self.fwd_param.setModelParams(model_params)
|
||
|
||
# Initialise cosmological parameters
|
||
cpar = utils.get_cosmopar(self.ini_file)
|
||
self.fwd.setCosmoParams(cpar)
|
||
self.fwd_param.setCosmoParams(cpar)
|
||
self.updateCosmology(cpar)
|
||
myprint(f"Original cosmological parameters: {self.fwd.getCosmoParams()}")
|
||
|
||
# Initialise derivative
|
||
self.grad_like = jax.grad(self.dens2like)
|
||
|
||
|
||
def initializeLikelihood(self, state: borg.likelihood.MarkovState) -> None:
|
||
"""
|
||
Initialise the likelihood internal variables and MarkovState variables.
|
||
|
||
Args:
|
||
- state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
|
||
"""
|
||
|
||
myprint("Init likelihood")
|
||
|
||
# for i in range(self.nsamp):
|
||
# state.newArray1d(f"tracer_vr_{i}", self.Nt[i], True)
|
||
# if self.run_type != 'data':
|
||
# state.newArray1d(f"data_x_{i}", self.Nt[i], True)
|
||
# state.newArray1d(f"data_y_{i}", self.Nt[i], True)
|
||
# state.newArray1d(f"data_z_{i}", self.Nt[i], True)
|
||
# state.newArray1d(f"true_x_{i}", self.Nt[i], True)
|
||
# state.newArray1d(f"true_y_{i}", self.Nt[i], True)
|
||
# state.newArray1d(f"true_z_{i}", self.Nt[i], True)
|
||
# self.data = [state[f"tracer_vr_{i}"] for i in range(self.nsamp)]
|
||
state.newArray3d("BORG_final_density", *self.fwd.getOutputBoxModel().N, True)
|
||
|
||
|
||
def updateMetaParameters(self, state: borg.likelihood.MarkovState) -> None:
|
||
"""
|
||
Update the meta parameters of the sampler (not sampled) from the MarkovState.
|
||
|
||
Args:
|
||
- state (borg.likelihood.MarkovState): The state object to be used in the likelihood.
|
||
|
||
"""
|
||
cpar = state['cosmology']
|
||
cpar.omega_q = 1. - cpar.omega_m - cpar.omega_k
|
||
self.fwd.setCosmoParams(cpar)
|
||
self.fwd_param.setCosmoParams(cpar)
|
||
|
||
|
||
def updateCosmology(self, cosmo: borg.cosmo.CosmologicalParameters) -> None:
|
||
"""
|
||
Updates the forward model's cosmological parameters with the given values.
|
||
|
||
Args:
|
||
- cosmo (borg.cosmo.CosmologicalParameters): The cosmological parameters.
|
||
|
||
"""
|
||
cpar = cosmo
|
||
|
||
# Convert sigma8 to As
|
||
cpar.A_s = 1.e-9 * symbolic_pofk.linear.sigma8_to_As(
|
||
cpar.sigma8, cpar.omega_m, cpar.omega_b, cpar.h, cpar.n_s)
|
||
myprint(f"Updating cosmology Om = {cosmo.omega_m}, sig8 = {cosmo.sigma8}, As = {cosmo.A_s}")
|
||
|
||
cpar.omega_q = 1. - cpar.omega_m - cpar.omega_k
|
||
self.fwd.setCosmoParams(cpar)
|
||
self.fwd_param.setCosmoParams(cpar)
|
||
|
||
# Compute growth rate
|
||
cosmology = borg.cosmo.Cosmology(cosmo)
|
||
f = cosmology.gplus(self.af) # dD / da
|
||
f *= self.af / cosmology.d_plus(self.af) # f = dlnD / dlna
|
||
self.f = f
|
||
|
||
|
||
def generateMBData(self) -> None:
|
||
"""
|
||
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
|
||
"""
|
||
|
||
self.MB_pos = [None] * self.nsamp
|
||
self.r_integration = [None] * self.nsamp
|
||
self.RA = [None] * self.nsamp
|
||
self.DEC = [None] * self.nsamp
|
||
|
||
for i in range(self.nsamp):
|
||
myprint(f"Making MB data for sample {i}")
|
||
|
||
# Get angular coordinates of all points
|
||
r_hat = projection.get_radial_vectors(self.coord_meas[i])
|
||
|
||
# Get min and max distance to integrate over
|
||
robs = self.cz_obs[i] / 100
|
||
sigr = np.sqrt((self.sig_v / 100) ** 2 + (self.frac_sig_rhMpc[i] * robs)**2)
|
||
rmin = robs - self.Nsig * sigr
|
||
rmin = rmin.at[rmin <= 0].set(self.L[0] / self.N[0] / 100.)
|
||
rmax = robs + self.Nsig * sigr
|
||
rmax = rmax.at[rmax > self.R_lim].set(self.R_lim)
|
||
|
||
# Compute coordinates of integration points
|
||
self.r_integration[i] = np.linspace(rmin, rmax, self.Nint_points)
|
||
cartesian_pos_MB = np.expand_dims(self.r_integration[i], axis=2) * r_hat
|
||
self.r_integration[i] = self.r_integration[i].T
|
||
self.MB_pos[i] = cartesian_pos_MB
|
||
self.MB_pos[i] = jnp.transpose(self.MB_pos[i], (2, 1, 0))
|
||
|
||
|
||
def generateMockData(self, s_hat: np.ndarray, state: borg.likelihood.MarkovState, make_plot: bool=True) -> None:
|
||
"""
|
||
Generates mock data by simulating the forward model with the given white noise,
|
||
drawing distance tracers from the density field, computing their distance
|
||
moduli and radial velocities, and adding Gaussian noise to the appropriate
|
||
variables. Also calculates the initial negative log-likelihood of the data.
|
||
|
||
Args:
|
||
- s_hat (np.ndarray): The input (initial) density field.
|
||
- state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
|
||
- make_plot (bool, default=True): Whether to make diagnostic plots for the mock data generation
|
||
"""
|
||
if self.run_type == 'data':
|
||
raise NotImplementedError
|
||
elif self.run_type == 'velmass':
|
||
raise NotImplementedError
|
||
elif self.run_type == 'mock':
|
||
self.coord_true, self.coord_meas, self.sig_mu, self.vr_true, self.cz_obs = \
|
||
mock_maker.borg_mock(s_hat, state, self.fwd, self.ini_file)
|
||
else:
|
||
raise NotImplementedError
|
||
|
||
self.r_hMpc = [np.sqrt(np.sum(self.coord_meas[i] ** 2, axis=0)) for i in range(self.nsamp)]
|
||
self.generateMBData()
|
||
|
||
|
||
def dens2like(self, output_density: np.ndarray):
|
||
"""
|
||
Given stored distance tracer data, computes the negative log-likelihood of the data
|
||
for this final density field.
|
||
|
||
Args:
|
||
output_density (np.ndarray): The z=0 density field.
|
||
Return:
|
||
lkl (float): The negative log-likelihood of the data.
|
||
"""
|
||
|
||
lkl = 0
|
||
sig_v = self.fwd_param.getModelParam('nullforward', 'sig_v')
|
||
|
||
# Compute velocity field
|
||
bulk_flow = jnp.array([self.fwd_param.getModelParam('nullforward', 'bulk_flow_x'),
|
||
self.fwd_param.getModelParam('nullforward', 'bulk_flow_y'),
|
||
self.fwd_param.getModelParam('nullforward', 'bulk_flow_z')])
|
||
v = forwards.dens2vel_linear(output_density, self.f,
|
||
self.fwd.getOutputBoxModel().L[0], self.smooth_R)
|
||
v = v + self.bulk_flow.reshape((3, 1, 1, 1))
|
||
|
||
omega_m = self.fwd.getCosmoParams().omega_m
|
||
|
||
for i in range(self.nsamp):
|
||
muA = self.fwd_param.getModelParam('nullforward', f'mua{i}')
|
||
alpha = self.fwd_param.getModelParam('nullforward', f'alpha{i}')
|
||
lkl += vel2like(
|
||
self.cz_obs[i],
|
||
v,
|
||
output_density,
|
||
self.MB_pos[i],
|
||
self.r_integration[i],
|
||
self.r_hMpc[i],
|
||
self.sig_mu[i],
|
||
sig_v,
|
||
omega_m,
|
||
muA,
|
||
alpha,
|
||
self.L[0],
|
||
jnp.array(self.fwd.getOutputBoxModel().xmin),
|
||
self.interp_order,
|
||
self.bias_epsilon
|
||
)
|
||
|
||
if not jnp.isfinite(lkl):
|
||
lkl = self.bignum
|
||
|
||
return lkl
|
||
|
||
def logLikelihoodComplex(self, s_hat: np.ndarray, gradientIsNext: bool):
|
||
"""
|
||
Calculates the negative log-likelihood of the data.
|
||
|
||
Args:
|
||
- s_hat (np.ndarray): The input white noise.
|
||
- gradientIsNext (bool): If True, prepares the forward model for gradient calculations.
|
||
|
||
Returns:
|
||
The negative log-likelihood value.
|
||
|
||
"""
|
||
|
||
N = self.fwd.getBoxModel().N[0]
|
||
L = self.fwd.getOutputBoxModel().L[0]
|
||
|
||
# Run BORG density field
|
||
output_density = np.zeros((N,N,N))
|
||
self.fwd.forwardModel_v2(s_hat)
|
||
self.fwd.getDensityFinal(output_density)
|
||
|
||
self.delta = output_density
|
||
|
||
L = self.dens2like(output_density)
|
||
myprint(f"var(s_hat): {np.var(s_hat)}, Call to logLike: {L}")
|
||
|
||
return L
|
||
|
||
|
||
def gradientLikelihoodComplex(self, s_hat: np.ndarray):
|
||
"""
|
||
Calculates the adjoint negative log-likelihood of the data.
|
||
|
||
Args:
|
||
- s_hat (np.ndarray): The input density field.
|
||
|
||
Returns:
|
||
The adjoint negative log-likelihood gradient.
|
||
|
||
"""
|
||
|
||
N = self.fwd.getBoxModel().N[0]
|
||
L = self.fwd.getOutputBoxModel().L[0]
|
||
|
||
# Run BORG density field
|
||
output_density = np.zeros((N,N,N))
|
||
self.fwd.forwardModel_v2(s_hat)
|
||
self.fwd.getDensityFinal(output_density)
|
||
|
||
mygradient = self.grad_like(output_density)
|
||
mygradient = np.array(mygradient, dtype=np.float64)
|
||
|
||
self.fwd.adjointModel_v2(mygradient)
|
||
mygrad_hat = np.zeros(s_hat.shape, dtype=np.complex128)
|
||
self.fwd.getAdjointModel(mygrad_hat)
|
||
|
||
return mygrad_hat
|
||
|
||
|
||
def commitAuxiliaryFields(self, state: borg.likelihood.MarkovState) -> None:
|
||
"""
|
||
Commits the final density field to the Markov state.
|
||
Args:
|
||
- state (borg.state.State): The state object containing the final density field.
|
||
"""
|
||
self.updateCosmology(self.fwd.getCosmoParams())
|
||
self.dens2like(self.delta)
|
||
state["BORG_final_density"][:] = self.delta
|
||
|
||
|
||
@partial(jax.jit, static_argnames=['L_BOX', 'interp_order', 'bias_epsilon'])
|
||
def vel2like(cz_obs, v, MB_field, MB_pos, r, r_hMpc, sig_mu, sig_v, omega_m, muA, alpha, L_BOX, X_MIN, interp_order, bias_epsilon):
|
||
"""
|
||
Jitted part of dens2like
|
||
"""
|
||
|
||
tracer_vel = projection.interp_field(v,
|
||
MB_pos,
|
||
L_BOX,
|
||
X_MIN,
|
||
interp_order,
|
||
use_jitted=True,
|
||
)
|
||
|
||
tracer_vr = projection.project_radial(
|
||
tracer_vel,
|
||
MB_pos,
|
||
jnp.zeros(3,)
|
||
)
|
||
|
||
# Convert velocities to redshifts
|
||
zco = utils.z_cos(r, omega_m)
|
||
cz_pred = utils.speed_of_light * zco + (1 + zco) * tracer_vr
|
||
delta_cz_sigv = (cz_pred - jnp.expand_dims(cz_obs, axis=1)) / sig_v
|
||
|
||
# Get p_r
|
||
delta_mu = 5. * jnp.log10(r / jnp.expand_dims(r_hMpc,axis=1) * muA)
|
||
|
||
# Get los biased density
|
||
los_density = projection.interp_field(
|
||
MB_field,
|
||
MB_pos,
|
||
L_BOX,
|
||
X_MIN,
|
||
interp_order,
|
||
use_jitted=True,
|
||
)
|
||
los_density = jnp.power(1. + los_density + bias_epsilon, alpha)
|
||
|
||
d2 = (delta_mu / sig_mu) ** 2
|
||
best = jnp.amin(jnp.abs(d2), axis=1)
|
||
|
||
# Multiply p_r by arbitrary number for numerical stability (cancels in p_r / p_r_norm)
|
||
d2 = d2 - jnp.expand_dims(jnp.nanmin(d2, axis=1), axis=1)
|
||
p_r = r ** 2 * jnp.exp(-0.5 * d2) * los_density
|
||
p_r_norm = jnp.expand_dims(jnp.trapz(p_r, r, axis=1), axis=1)
|
||
|
||
# Integrate to get likelihood
|
||
d2 = delta_cz_sigv**2
|
||
scale = jnp.nanmin(d2, axis=1)
|
||
d2 = d2 - jnp.expand_dims(scale, axis=1)
|
||
exp_delta_cz = jnp.exp(-0.5 * d2)
|
||
p_cz = jnp.trapz(exp_delta_cz * p_r / p_r_norm, r, axis=1)
|
||
lkl_ind = jnp.log(p_cz) - scale / 2 - 0.5 * jnp.log(2 * np.pi * sig_v**2)
|
||
lkl = - lkl_ind.sum()
|
||
|
||
return lkl
|
||
|
||
|
||
@borg.registerGravityBuilder
|
||
def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.BoxModel, ini_file=None) -> borg.forward.BaseForwardModel:
|
||
"""
|
||
Builds the gravity model and returns the forward model chain.
|
||
|
||
Args:
|
||
- state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
|
||
- box (borg.forward.BoxModel): The input box model.
|
||
- ini_file (str, default=None): The location of the ini file. If None, use borg.getIniConfigurationFilename()
|
||
|
||
Returns:
|
||
borg.forward.BaseForwardModel: The forward model.
|
||
|
||
"""
|
||
|
||
global chain, fwd_param
|
||
myprint("Building gravity model")
|
||
|
||
if ini_file is None:
|
||
myprint("Reading from configuration file: " + borg.getIniConfigurationFilename())
|
||
config = configparser.ConfigParser()
|
||
config.read(borg.getIniConfigurationFilename())
|
||
else:
|
||
myprint("Reading from configuration file: " + ini_file)
|
||
config = configparser.ConfigParser()
|
||
config.read(ini_file)
|
||
ai = float(config['model']['ai'])
|
||
af = float(config['model']['af'])
|
||
|
||
# Setup forward model
|
||
chain = borg.forward.ChainForwardModel(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 config['model']['gravity'] == 'linear':
|
||
raise NotImplementedError(config['model']['gravity'])
|
||
elif config['model']['gravity'] == 'lpt':
|
||
chain @= borg.forward.model_lib.M_LPT_CIC(
|
||
box,
|
||
opts=dict(a_initial=af,
|
||
a_final=af,
|
||
do_rsd=False,
|
||
supersampling=1,
|
||
lightcone=False,
|
||
part_factor=1.01,))
|
||
elif config['model']['gravity'] == '2lpt':
|
||
chain @= borg.forward.model_lib.M_2LPT_CIC(
|
||
box,
|
||
opts=dict(a_initial=af,
|
||
a_final=af,
|
||
do_rsd=False,
|
||
supersampling=1,
|
||
lightcone=False,
|
||
part_factor=1.01,))
|
||
elif config['model']['gravity'] == 'pm':
|
||
chain @= borg.forward.model_lib.M_PM_CIC(
|
||
box,
|
||
opts=dict(a_initial=af,
|
||
a_final=af,
|
||
do_rsd=False,
|
||
supersampling=1,
|
||
part_factor=1.01,
|
||
forcesampling=2,
|
||
pm_start_z=1/ai - 1,
|
||
pm_nsteps=int(config['model']['nsteps']),
|
||
tcola=False))
|
||
elif config['model']['gravity'] == 'cola':
|
||
chain @= borg.forward.model_lib.M_PM_CIC(
|
||
box,
|
||
opts=dict(a_initial=af,
|
||
a_final=af,
|
||
do_rsd=False,
|
||
supersampling=1,
|
||
part_factor=1.01,
|
||
forcesampling=2,
|
||
pm_start_z=1/ai - 1,
|
||
pm_nsteps=int(config['model']['nsteps']),
|
||
tcola=True))
|
||
else:
|
||
raise NotImplementedError(config['model']['gravity'])
|
||
|
||
# Cosmological parameters
|
||
if ini_file is None:
|
||
cpar = utils.get_cosmopar(borg.getIniConfigurationFilename())
|
||
else:
|
||
cpar = utils.get_cosmopar(ini_file)
|
||
chain.setCosmoParams(cpar)
|
||
|
||
# This is the forward model for the model parameters
|
||
fwd_param = borg.forward.ChainForwardModel(box)
|
||
mod = forwards.NullForward(box)
|
||
fwd_param.addModel(mod)
|
||
fwd_param.setCosmoParams(cpar)
|
||
|
||
return chain
|
||
|
||
|
||
@borg.registerSamplerBuilder
|
||
def build_sampler(
|
||
state: borg.likelihood.MarkovState,
|
||
info: borg.likelihood.LikelihoodInfo,
|
||
):
|
||
"""
|
||
Builds the sampler and returns it.
|
||
Which parameters to sample are given in the ini file.
|
||
We assume all parameters are NOT meant to be sampled, unless we find "XX_sampler_blocked = false" in the ini file
|
||
|
||
Args:
|
||
- state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
|
||
- info (borg.likelihood.LikelihoodInfo): The likelihood information.
|
||
|
||
Returns:
|
||
List of samplers to use.
|
||
|
||
"""
|
||
myprint("Building sampler")
|
||
|
||
myprint("Reading from configuration file: " + borg.getIniConfigurationFilename())
|
||
config = configparser.ConfigParser()
|
||
config.read(borg.getIniConfigurationFilename())
|
||
end = '_sampler_blocked'
|
||
to_sample = [k[:-len(end)] for (k, v) in config['block_loop'].items() if k[-len(end):] == end and v.lower() == 'false']
|
||
myprint(f'Parameters to sample: {to_sample}')
|
||
nsamp = int(config['run']['nsamp'])
|
||
|
||
all_sampler = []
|
||
|
||
# Cosmology sampler arguments
|
||
prefix = ""
|
||
params = []
|
||
initial_values = {}
|
||
prior = {}
|
||
for p in ["omega_m", "sigma8"]:
|
||
if p not in to_sample:
|
||
continue
|
||
if p in config['prior'].keys() and p in config['cosmology'].keys():
|
||
myprint(f'Adding {p} sampler')
|
||
params.append(f"cosmology.{p}")
|
||
initial_values[f"cosmology.{p}"] = float(config['cosmology'][p])
|
||
prior[f"cosmology.{p}"] = np.array(ast.literal_eval(config['prior'][p]))
|
||
else:
|
||
s = f'Could not find {p} prior and/or default, so will not sample'
|
||
warnings.warn(s, stacklevel=2)
|
||
# Remove for later to prevent duplication
|
||
to_sample.remove(p)
|
||
if len(params) > 0:
|
||
myprint('Adding cosmological parameter sampler')
|
||
all_sampler.append(borg.samplers.ModelParamsSampler(prefix, params, likelihood, chain, initial_values, prior))
|
||
|
||
# Model parameter sampler
|
||
prefix = ""
|
||
params = []
|
||
initial_values = {}
|
||
prior = {}
|
||
for p in to_sample:
|
||
if p in config['prior'].keys():
|
||
if p == 'sig_v':
|
||
myprint(f'Adding {p} sampler')
|
||
params.append(p)
|
||
initial_values[p] = float(config['model'][p])
|
||
if 'inf' in config['prior'][p]:
|
||
x = ast.literal_eval(config['prior'][p].replace('inf', '"inf"'))
|
||
prior[p] = np.array([xx if xx != 'inf' else np.inf for xx in x])
|
||
else:
|
||
prior[p] = np.array(ast.literal_eval(config['prior'][p]))
|
||
elif p == 'bulk_flow':
|
||
for i, d in enumerate(['_x', '_y', '_z']):
|
||
myprint(f'Adding {p}{d} sampler')
|
||
params.append(f'{p}{d}')
|
||
initial_values[f'{p}{d}'] = np.array(ast.literal_eval(config['model']['bulk_flow']))[i]
|
||
if 'inf' in config['prior'][p]:
|
||
x = ast.literal_eval(config['prior'][p].replace('inf', '"inf"'))
|
||
prior[f'{p}{d}'] = np.array([xx if xx != 'inf' else np.inf for xx in x])
|
||
else:
|
||
prior[f'{p}{d}'] = np.array(ast.literal_eval(config['prior'][p]))
|
||
else:
|
||
for i in range(nsamp):
|
||
myprint(f'Adding {p}{i} sampler')
|
||
params.append(f'{p}{i}')
|
||
initial_values[f'{p}{i}'] = float(config[f'sample_{i}'][p])
|
||
if 'inf' in config['prior'][p]:
|
||
x = ast.literal_eval(config['prior'][p].replace('inf', '"inf"'))
|
||
prior[f'{p}{i}'] = np.array([xx if xx != 'inf' else np.inf for xx in x])
|
||
else:
|
||
prior[f'{p}{i}'] = np.array(ast.literal_eval(config['prior'][p]))
|
||
else:
|
||
s = f'Could not find {p} prior, so will not sample'
|
||
warnings.warn(s, stacklevel=2)
|
||
if len(params) > 0:
|
||
myprint('Adding model parameter sampler')
|
||
all_sampler.append(borg.samplers.ModelParamsSampler(prefix, params, likelihood, fwd_param, initial_values, prior))
|
||
|
||
return all_sampler
|
||
|
||
|
||
@borg.registerLikelihoodBuilder
|
||
def build_likelihood(state: borg.likelihood.MarkovState, info: borg.likelihood.LikelihoodInfo) -> borg.likelihood.BaseLikelihood:
|
||
"""
|
||
Builds the likelihood object and returns it.
|
||
|
||
Args:
|
||
- state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
|
||
- info (borg.likelihood.LikelihoodInfo): The likelihood information.
|
||
|
||
Returns:
|
||
borg.likelihood.BaseLikelihood: The likelihood object.
|
||
|
||
"""
|
||
global likelihood, fwd_param
|
||
myprint("Building likelihood")
|
||
myprint(chain.getCosmoParams())
|
||
boxm = chain.getBoxModel()
|
||
likelihood = VelocityBORGLikelihood(chain, fwd_param, borg.getIniConfigurationFilename())
|
||
return likelihood |