From 17f9099e556b3bfe54b3255c3ead5c6a9449e0fd Mon Sep 17 00:00:00 2001 From: Deaglan Bartlett Date: Wed, 24 Apr 2024 00:04:01 +0200 Subject: [PATCH] Minimum example to give MCMC samples --- borg_velocity/forwards.py | 1 + borg_velocity/likelihood.py | 168 +++++++++++++++++++++++++++++++++--- borg_velocity/mock_maker.py | 12 +-- borg_velocity/projection.py | 2 +- conf/basic_ini.ini | 3 +- 5 files changed, 168 insertions(+), 18 deletions(-) diff --git a/borg_velocity/forwards.py b/borg_velocity/forwards.py index 6cda0f3..d8459e8 100644 --- a/borg_velocity/forwards.py +++ b/borg_velocity/forwards.py @@ -3,6 +3,7 @@ import jax import jax.numpy as jnp from functools import partial +from borg_velocity.utils import myprint class NullForward(borg.forward.BaseForwardModel): """ diff --git a/borg_velocity/likelihood.py b/borg_velocity/likelihood.py index 3e2cc13..0f17a32 100644 --- a/borg_velocity/likelihood.py +++ b/borg_velocity/likelihood.py @@ -4,6 +4,9 @@ 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 @@ -48,6 +51,8 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood): 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 @@ -55,7 +60,16 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood): 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) @@ -64,6 +78,25 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood): 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: """ @@ -120,6 +153,12 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood): 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: @@ -133,7 +172,7 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood): """ self.MB_pos = [None] * self.nsamp - self.r = [None] * self.nsamp + self.r_integration = [None] * self.nsamp self.RA = [None] * self.nsamp self.DEC = [None] * self.nsamp @@ -152,10 +191,11 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood): rmax = rmax.at[rmax > self.R_lim].set(self.R_lim) # Compute coordinates of integration points - self.r[i] = np.linspace(rmin, rmax, self.Nint_points) - cartesian_pos_MB = np.expand_dims(self.r[i], axis=2) * r_hat + 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, 0, 1)) + 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: @@ -180,9 +220,59 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood): 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() - quit() + + 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): """ @@ -207,8 +297,7 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood): self.delta = output_density - # L = self.dens2like(output_density) - L = 1. + L = self.dens2like(output_density) myprint(f"var(s_hat): {np.var(s_hat)}, Call to logLike: {L}") return L @@ -234,10 +323,10 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood): self.fwd.forwardModel_v2(s_hat) self.fwd.getDensityFinal(output_density) -# mygradient = self.grad_like(output_density) -# mygradient = np.array(mygradient, dtype=np.float64) + mygradient = self.grad_like(output_density) + mygradient = np.array(mygradient, dtype=np.float64) -# self.fwd.adjointModel_v2(mygradient) + self.fwd.adjointModel_v2(mygradient) mygrad_hat = np.zeros(s_hat.shape, dtype=np.complex128) self.fwd.getAdjointModel(mygrad_hat) @@ -253,6 +342,65 @@ class VelocityBORGLikelihood(borg.likelihood.BaseLikelihood): 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 diff --git a/borg_velocity/mock_maker.py b/borg_velocity/mock_maker.py index 6c15243..a604df2 100644 --- a/borg_velocity/mock_maker.py +++ b/borg_velocity/mock_maker.py @@ -71,7 +71,7 @@ def borg_mock(s_hat, state, fwd_model, ini_file): fwd_model.getOutputBoxModel().L[0], smooth_R) # Add bulk flow - bulk_flow = np.array(ast.literal_eval(config['mock']['bulk_flow'])) + bulk_flow = np.array(ast.literal_eval(config['model']['bulk_flow'])) output_vel = output_vel + bulk_flow.reshape((3, 1, 1, 1)) # Sample positions according to bias model @@ -117,20 +117,20 @@ def borg_mock(s_hat, state, fwd_model, ini_file): np.array(fwd_model.getOutputBoxModel().xmin), interp_order ) - vr_true[i] = projection.project_radial( + vr_true[i] = np.squeeze(projection.project_radial( tracer_vel, np.expand_dims(coord_true[i], axis=2), np.zeros(3,) - ) + )) # Compute observed redshifts (including noise) sig_v = float(config['model']['sig_v']) cz_obs = [None] * nsamp for i in range(nsamp): - rtrue = np.sqrt(np.sum(coord_true[i] ** 2, axis=1)) + rtrue = np.sqrt(np.sum(coord_true[i] ** 2, axis=0)) zco = utils.z_cos(rtrue, cosmo.omega_m) - cz_obs = utils.speed_of_light * zco + (1 + zco) * vr_true[i] - cz_obs += np.random.normal(size=cz_obs.shape) * sig_v # CHECK THIS LINE!!!!! + cz_obs[i] = utils.speed_of_light * zco + (1 + zco) * vr_true[i] + cz_obs[i] += np.random.normal(size=cz_obs[i].shape) * sig_v # CHECK THIS LINE!!!!! # Add observational systematic due to incorrect distance estimate # \mu -> \mu + 5 log10(A), or equivalently d -> A d diff --git a/borg_velocity/projection.py b/borg_velocity/projection.py index 8d657e5..c9b351d 100644 --- a/borg_velocity/projection.py +++ b/borg_velocity/projection.py @@ -91,6 +91,6 @@ def get_radial_vectors(coord_meas): # Get unit vectors along line of sight r_hat = jnp.array(SkyCoord(ra=RA*apu.deg, dec=DEC*apu.deg).cartesian.xyz) - r_hat = jnp.expand_dims(r_hat, axis=0) + r_hat = jnp.expand_dims(r_hat.T, axis=0) return r_hat \ No newline at end of file diff --git a/conf/basic_ini.ini b/conf/basic_ini.ini index 6872689..91db0e4 100644 --- a/conf/basic_ini.ini +++ b/conf/basic_ini.ini @@ -29,6 +29,7 @@ ares_heat = 1.0 number_to_generate = 10 random_ic = false init_random_scaling = 0.1 +bignum = 1e300 [model] gravity = lpt @@ -42,6 +43,7 @@ sig_v = 150. R_lim = none Nint_points = 201 Nsig = 10 +bulk_flow = [0.0, 0.0, 0.0] [prior] omega_m = [0.1, 0.8] @@ -73,7 +75,6 @@ NSAMP = 2 [mock] R_max = 100 -bulk_flow = [0.0, 0.0, 0.0] [python] likelihood_path = /home/bartlett/fsigma8/borg_velocity/borg_velocity/likelihood.py