Add cosmology sampler (example2)
This commit is contained in:
Deaglan Bartlett
parent
d1212072b7
commit
174acf7fc1
3 changed files with 528 additions and 0 deletions
430
example2.py
Normal file
430
example2.py
Normal file
|
|
@ -0,0 +1,430 @@
|
|||
import aquila_borg as borg
|
||||
import numpy as np
|
||||
import numbers
|
||||
import jaxlib
|
||||
import jax.numpy as jnp
|
||||
import jax
|
||||
import configparser
|
||||
import ast
|
||||
import warnings
|
||||
from functools import partial
|
||||
|
||||
# Output stream management
|
||||
cons = borg.console()
|
||||
def myprint(x):
|
||||
if isinstance(x, str):
|
||||
cons.print_std(x)
|
||||
else:
|
||||
cons.print_std(repr(x))
|
||||
|
||||
def get_cosmopar(ini_file):
|
||||
"""
|
||||
Extract cosmological parameters from an ini file
|
||||
|
||||
Args:
|
||||
:ini_file (str): Path to the ini file
|
||||
|
||||
Returns:
|
||||
:cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters
|
||||
"""
|
||||
|
||||
config = configparser.ConfigParser()
|
||||
config.read(ini_file)
|
||||
|
||||
cpar = borg.cosmo.CosmologicalParameters()
|
||||
cpar.default()
|
||||
cpar.fnl = float(config['cosmology']['fnl'])
|
||||
cpar.omega_k = float(config['cosmology']['omega_k'])
|
||||
cpar.omega_m = float(config['cosmology']['omega_m'])
|
||||
cpar.omega_b = float(config['cosmology']['omega_b'])
|
||||
cpar.omega_q = float(config['cosmology']['omega_q'])
|
||||
cpar.h = float(config['cosmology']['h100'])
|
||||
cpar.sigma8 = float(config['cosmology']['sigma8'])
|
||||
cpar.n_s = float(config['cosmology']['n_s'])
|
||||
cpar.w = float(config['cosmology']['w'])
|
||||
cpar.wprime = float(config['cosmology']['wprime'])
|
||||
|
||||
cpar = compute_As(cpar)
|
||||
|
||||
return cpar
|
||||
|
||||
|
||||
def compute_As(cpar):
|
||||
"""
|
||||
Compute As given values of sigma8
|
||||
|
||||
Args:
|
||||
:cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters with wrong As
|
||||
|
||||
Returns:
|
||||
:cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters with updated As
|
||||
|
||||
"""
|
||||
|
||||
# requires BORG-CLASS
|
||||
if not hasattr(borg.cosmo, 'ClassCosmo'):
|
||||
raise ImportError(
|
||||
"BORG-CLASS is required to compute As, but is not installed.")
|
||||
|
||||
sigma8_true = jnp.copy(cpar.sigma8)
|
||||
cpar.sigma8 = 0
|
||||
cpar.A_s = 2.3e-9
|
||||
k_max, k_per_decade = 10, 100
|
||||
extra_class = {}
|
||||
extra_class['YHe'] = '0.24'
|
||||
cosmo = borg.cosmo.ClassCosmo(cpar, k_per_decade, k_max, extra=extra_class)
|
||||
cosmo.computeSigma8()
|
||||
cos = cosmo.getCosmology()
|
||||
cpar.A_s = (sigma8_true/cos['sigma_8'])**2*cpar.A_s
|
||||
cpar.sigma8 = sigma8_true
|
||||
|
||||
print('Updated cosmology:', cpar)
|
||||
|
||||
return cpar
|
||||
|
||||
|
||||
class MyLikelihood(borg.likelihood.BaseLikelihood):
|
||||
"""
|
||||
HADES likelihood class
|
||||
"""
|
||||
|
||||
def __init__(self, fwd: borg.forward.BaseForwardModel, ini_fname: str):
|
||||
|
||||
self.fwd = fwd
|
||||
|
||||
# Read the ini file
|
||||
self.ini_fname = ini_fname
|
||||
self.config = configparser.ConfigParser()
|
||||
self.config.read(ini_fname)
|
||||
self.N = [int(self.config['system'][f'N{i}']) for i in range(3)] # Number of grid points per side
|
||||
self.L = [float(self.config['system'][f'L{i}']) for i in range(3)] # Box size lenght Mpc/h
|
||||
|
||||
self.sigma_dens = float(self.config['mock']['sigma_dens']) # Density scatter
|
||||
|
||||
myprint(f"Likelihood initialized with {self.N} grid points and box size {self.L} Mpc/h")
|
||||
super().__init__(fwd, self.N, self.L)
|
||||
|
||||
# Set up cosmoligical parameters
|
||||
cpar = get_cosmopar(ini_fname)
|
||||
self.updateCosmology(cpar)
|
||||
|
||||
# Gradient of the likelihood
|
||||
self.grad_like = jax.grad(self.dens2like)
|
||||
|
||||
def updateCosmology(self, cosmo: borg.cosmo.CosmologicalParameters) -> None:
|
||||
cpar = compute_As(cosmo)
|
||||
self.fwd.setCosmoParams(cpar)
|
||||
|
||||
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.
|
||||
|
||||
"""
|
||||
cosmo = state['cosmology']
|
||||
cpar = compute_As(cosmo)
|
||||
self.fwd.setCosmoParams(cpar)
|
||||
|
||||
def initializeLikelihood(self, state: borg.likelihood.MarkovState) -> None:
|
||||
"""
|
||||
Initialize the likelihood function.
|
||||
|
||||
Args:
|
||||
- state (borg.likelihood.MarkovState): The state object to be used in the likelihood.
|
||||
|
||||
"""
|
||||
myprint("Init likelihood")
|
||||
state.newArray3d("BORG_final_density", *self.fwd.getOutputBoxModel().N, True)
|
||||
|
||||
# Could load real data
|
||||
# We'll generate mock data which has its own function
|
||||
|
||||
def generateMockData(self, s_hat:np.ndarray, state: borg.likelihood.MarkovState) -> None:
|
||||
"""
|
||||
Generates mock data by simulating the forward model with the given white noise
|
||||
|
||||
Args:
|
||||
- s_hat (np.ndarray): The input (initial) white noise field.
|
||||
- state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
|
||||
"""
|
||||
myprint('Making mock from BORG')
|
||||
|
||||
# Get density field from the initial conditions
|
||||
# Could replace with any (better) simulation here
|
||||
# This version is self-consistnet
|
||||
dens = np.zeros(self.fwd.getOutputBoxModel().N)
|
||||
myprint('Running forward model')
|
||||
myprint(self.fwd.getCosmoParams())
|
||||
self.fwd.forwardModel_v2(s_hat)
|
||||
self.fwd.getDensityFinal(dens)
|
||||
state["BORG_final_density"][:] = dens
|
||||
self.true_dens = dens.copy()
|
||||
|
||||
# Add some scatter
|
||||
myprint('Adding scatter')
|
||||
self.obs_dens = self.true_dens + np.random.randn(*self.true_dens.shape) * self.sigma_dens
|
||||
|
||||
# Compute the likelihood and print it
|
||||
myprint('From mock')
|
||||
self.saved_s_hat = s_hat.copy()
|
||||
self.logLikelihoodComplex(s_hat, False)
|
||||
self.commitAuxiliaryFields(state)
|
||||
myprint('Done')
|
||||
|
||||
|
||||
def dens2like(self, output_density: np.ndarray):
|
||||
"""
|
||||
Compute the likelihood from the density field
|
||||
Args:
|
||||
- output_density (np.ndarray): The density field to be used in the likelihood.
|
||||
Returns:
|
||||
- float: The likelihood value.
|
||||
"""
|
||||
# Compute the likelihood from the density field
|
||||
# This is a simple Gaussian likelihood
|
||||
# Could be replaced with any other likelihood
|
||||
diff = output_density - self.obs_dens
|
||||
like = 0.5 * jnp.sum(diff**2) / (self.sigma_dens**2)
|
||||
|
||||
return like
|
||||
|
||||
|
||||
def logLikelihoodComplex(self, s_hat:np.ndarray, gradientIsNext:bool):
|
||||
|
||||
# myprint('Getting density field now')
|
||||
# Get the density field from the forward model
|
||||
dens = np.zeros(self.fwd.getOutputBoxModel().N)
|
||||
self.fwd.forwardModel_v2(s_hat)
|
||||
self.fwd.getDensityFinal(dens)
|
||||
|
||||
L = self.dens2like(dens)
|
||||
|
||||
if isinstance(L, numbers.Number) or isinstance(L, jaxlib.xla_extension.ArrayImpl):
|
||||
myprint(f"var(s_hat): {np.var(s_hat)}, Call to logLike: {L}")
|
||||
|
||||
self.delta = dens.copy()
|
||||
|
||||
return L
|
||||
|
||||
|
||||
def gradientLikelihoodComplex(self, s_hat:np.ndarray):
|
||||
|
||||
# Run BORG density field
|
||||
output_density = np.zeros(self.N)
|
||||
self.fwd.forwardModel_v2(s_hat)
|
||||
self.fwd.getDensityFinal(output_density)
|
||||
|
||||
# Compute the gradient of the likelihood
|
||||
# d logL / d dens
|
||||
mygradient = self.grad_like(output_density)
|
||||
|
||||
# Now get d logL / d s_hat
|
||||
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)
|
||||
self.fwd.clearAdjointGradient()
|
||||
|
||||
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.copy()
|
||||
|
||||
|
||||
@borg.registerGravityBuilder
|
||||
def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.BoxModel, ini_fname=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
|
||||
myprint("Building gravity model")
|
||||
|
||||
if ini_fname is None:
|
||||
ini_fname=borg.getIniConfigurationFilename()
|
||||
config = configparser.ConfigParser()
|
||||
config.read(ini_fname)
|
||||
# READ FROM INI FILE
|
||||
which_model = config['gravity']['which_model']
|
||||
ai = float(config['gravity']['ai']) # Initial scale factor
|
||||
af = float(config['gravity']['af']) # Final scale factor
|
||||
supersampling = int(config['gravity']['supersampling'])
|
||||
forcesampling = int(config['gravity']['forcesampling'])
|
||||
nsteps = int(config['gravity']['nsteps']) # Number of steps in the PM solver
|
||||
|
||||
chain = borg.forward.ChainForwardModel(box)
|
||||
|
||||
# Make sure that the initial conditions are real in position space
|
||||
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'}}) # helps deals with errors with primordial physics in CLASS for weird cosmologies
|
||||
chain @= transfer_class
|
||||
|
||||
# This one works
|
||||
# chain @= borg.forward.models.Primordial(box, ai)
|
||||
# chain @= borg.forward.models.EisensteinHu(box)
|
||||
|
||||
# Gravity model
|
||||
if which_model == '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 which_model == '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 which_model == 'pm':
|
||||
mod = borg.forward.model_lib.M_PM_CIC(
|
||||
box,
|
||||
opts=dict(a_initial=af,
|
||||
a_final=af,
|
||||
pm_start_z=1/ai - 1,
|
||||
do_rsd=False,
|
||||
supersampling=supersampling,
|
||||
forcesampling=forcesampling,
|
||||
lightcone=False,
|
||||
part_factor=1.01,
|
||||
pm_nsteps=nsteps, # Number of steps in the PM solver
|
||||
tcola=False
|
||||
))
|
||||
elif which_model == 'cola':
|
||||
mod = borg.forward.model_lib.M_PM_CIC(
|
||||
box,
|
||||
opts=dict(a_initial=af,
|
||||
a_final=af,
|
||||
pm_start_z=1/ai - 1,
|
||||
do_rsd=False,
|
||||
supersampling=supersampling,
|
||||
forcesampling=forcesampling,
|
||||
lightcone=False,
|
||||
part_factor=1.01,
|
||||
pm_nsteps=nsteps, # Number of steps in the PM solver
|
||||
tcola=True
|
||||
))
|
||||
else:
|
||||
raise ValueError(f"Unknown model {which_model}")
|
||||
|
||||
mod.accumulateAdjoint(True)
|
||||
chain @= mod
|
||||
|
||||
# Cosmological parameters
|
||||
cpar = get_cosmopar(borg.getIniConfigurationFilename())
|
||||
print('Setting cosmo params', cpar)
|
||||
chain.setCosmoParams(cpar)
|
||||
|
||||
return chain
|
||||
|
||||
|
||||
def check_cosmo_sampling(loop):
|
||||
return loop.getStepID() > begin_cosmo
|
||||
|
||||
|
||||
@borg.registerSamplerBuilder
|
||||
def build_sampler(state: borg.likelihood.MarkovState, info: borg.likelihood.LikelihoodInfo, loop: borg.samplers.MainLoop):
|
||||
"""
|
||||
Builds the sampler and returns the main loop.
|
||||
|
||||
Args:
|
||||
- state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
|
||||
- info (borg.likelihood.LikelihoodInfo): The likelihood info object to be used in the likelihood.
|
||||
- loop (borg.samplers.MainLoop): The main loop object to be used in the likelihood.
|
||||
|
||||
Returns:
|
||||
borg.samplers.MainLoop: The main loop.
|
||||
"""
|
||||
|
||||
global begin_cosmo
|
||||
|
||||
myprint("Building sampler")
|
||||
|
||||
# Read from config file what to sample
|
||||
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.endswith(end) and v.lower() == 'false']
|
||||
myprint(f"Sampling {to_sample}")
|
||||
|
||||
all_sampler = []
|
||||
|
||||
# Add cosmology samplers
|
||||
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)
|
||||
|
||||
begin_cosmo = int(config['mcmc']['warmup_cosmo'])
|
||||
|
||||
if len(params) > 0:
|
||||
myprint(f"Sampling cosmology parameters {params}")
|
||||
cosmo_sampler = borg.samplers.ModelParamsSampler("", params, likelihood, chain, initial_values, prior)
|
||||
cosmo_sampler.setName("cosmo_sampler")
|
||||
all_sampler.append(cosmo_sampler)
|
||||
loop.push(cosmo_sampler)
|
||||
loop.addToConditionGroup("warmup_cosmo", "cosmo_sampler")
|
||||
loop.addConditionToConditionGroup("warmup_cosmo", partial(check_cosmo_sampling, loop))
|
||||
|
||||
# Here you can add model parameter sampling, etc.
|
||||
|
||||
return []
|
||||
|
||||
|
||||
@borg.registerLikelihoodBuilder
|
||||
def build_likelihood(state: borg.likelihood.MarkovState, info: borg.likelihood.LikelihoodInfo) -> borg.likelihood.BaseLikelihood:
|
||||
"""
|
||||
Builds the likelihood and returns the likelihood object.
|
||||
|
||||
Args:
|
||||
- state (borg.likelihood.MarkovState): The Markov state object to be used in the likelihood.
|
||||
- info (borg.likelihood.LikelihoodInfo): The likelihood info object to be used in the likelihood.
|
||||
|
||||
Returns:
|
||||
borg.likelihood.BaseLikelihood: The likelihood object.
|
||||
"""
|
||||
|
||||
myprint("Building likelihood")
|
||||
global likelihood
|
||||
likelihood = MyLikelihood(chain, borg.getIniConfigurationFilename())
|
||||
return likelihood
|
||||
71
ini_example2.ini
Normal file
71
ini_example2.ini
Normal file
|
|
@ -0,0 +1,71 @@
|
|||
[system]
|
||||
console_output = borg_log
|
||||
VERBOSE_LEVEL = 2
|
||||
N0 = 64
|
||||
N1 = 64
|
||||
N2 = 64
|
||||
L0 = 500.0
|
||||
L1 = 500.0
|
||||
L2 = 500.0
|
||||
corner0 = -250.0
|
||||
corner1 = -250.0
|
||||
corner2 = -250.0
|
||||
NUM_MODES = 100
|
||||
test_mode = true
|
||||
seed_cpower = true
|
||||
|
||||
[block_loop]
|
||||
hades_sampler_blocked = false
|
||||
sigma8_sampler_blocked = false
|
||||
omega_m_sampler_blocked = false
|
||||
bias_sampler_blocked= true
|
||||
ares_heat = 1.0
|
||||
|
||||
[prior]
|
||||
sigma8 = [0.5, 1.2]
|
||||
omega_m = [0.1, 0.5]
|
||||
|
||||
[mcmc]
|
||||
number_to_generate = 5
|
||||
random_ic = false
|
||||
init_random_scaling = 0.1
|
||||
warmup_cosmo = 2
|
||||
|
||||
[hades]
|
||||
algorithm = HMC
|
||||
max_epsilon = 0.01
|
||||
max_timesteps = 50
|
||||
mixing = 1
|
||||
|
||||
[python]
|
||||
likelihood_path = /home/bartlett/borg_examples/example2.py
|
||||
|
||||
[run]
|
||||
run_type = mock
|
||||
NCAT = 0
|
||||
|
||||
[cosmology]
|
||||
omega_r = 0
|
||||
fnl = 0
|
||||
omega_k = 0
|
||||
omega_m = 0.315
|
||||
omega_b = 0.049
|
||||
omega_q = 0.685
|
||||
h100 = 0.68
|
||||
sigma8 = 0.81
|
||||
n_s = 0.97
|
||||
w = -1
|
||||
wprime = 0
|
||||
beta = 1.5
|
||||
z0 = 0
|
||||
|
||||
[mock]
|
||||
sigma_dens = 1.
|
||||
|
||||
[gravity]
|
||||
which_model = lpt
|
||||
ai = 0.05
|
||||
af = 1.0
|
||||
supersampling = 2
|
||||
forcesampling = 2
|
||||
nsteps = 20
|
||||
27
run_example2.sh
Executable file
27
run_example2.sh
Executable file
|
|
@ -0,0 +1,27 @@
|
|||
#!/bin/sh
|
||||
|
||||
# Modules
|
||||
module purge
|
||||
module restore myborg
|
||||
module load cuda/12.6
|
||||
|
||||
# Environment
|
||||
source /home/bartlett/.bashrc
|
||||
source /home/bartlett/anaconda3/etc/profile.d/conda.sh
|
||||
conda deactivate
|
||||
conda activate borg_new
|
||||
|
||||
# Kill job if there are any errors
|
||||
set -e
|
||||
|
||||
# Path variables
|
||||
BORG=/data101/bartlett/build_borg/tools/hades_python/hades_python
|
||||
RUN_DIR=/data101/bartlett/borg_examples/example2
|
||||
|
||||
mkdir -p $RUN_DIR
|
||||
cd $RUN_DIR
|
||||
|
||||
INI_FILE=/home/bartlett/borg_examples/ini_example2.ini
|
||||
|
||||
cp $INI_FILE ini_file.ini
|
||||
$BORG INIT ini_file.ini
|
||||
Loading…
Add table
Reference in a new issue