Samplers and Minimisers

Samplers and Minimisers#

We provide tools to interface the likelihood with the most commonly used MCMC samplers in cosmology: Cobaya, MontePython, and CosmoSIS. Moreover, we show how the likelihood can be interfaced with tools that explicitly exploit its differentiability: BlackJAX and Optax.

Warning

When using MCMC samplers, be careful to only apply priors once. For some data sets candl includes Gaussian priors on cosmological and nuisance parameters, to ensure a fully differentiable pipeline. When interfacing candl with Cobaya and MontePython these will be turned off by default and you are expected to specify them within the relevant framework. However, this may not be available for all samplers (e.g. BlackJAX). Look out for clear_internal_priors flags in the interface functions and take care not to double count information. Also see clear_specific_priors for more fine-grained control.

Tip

See Data Selection for info on how to select subsets of the data.

The examples in the section below use the candl likelihoods made available through the candl_data repository. Simply clone the repo to get the data and run pip install . to install the shortcuts. More details can be found in the candl_data readme.

Cobaya#

Interactive#

To interface with Cobaya interactively, you can use the supplied interface code. For example:

import candl
import candl.interface
import candl_data

candl_like = candl.Like(candl_data.SPT3G_2018_TTTEEE)
cobaya_dict = {"likelihood": candl.interface.get_cobaya_info_dict_for_like(candl_like)}

This will interface the likelihood with Cobaya and register all of its requirements. You can then proceed to populate cobaya_dict with the the parameters to be sampled etc. and run Cobaya as usual. By default the candl internal priors are not applied, add "clear_internal_priors": False to the relevant likelihood block in the dictionary if you want to use them. If you only want to clear some internal priors, you can supply a string or list of strings to clear_specific_priors, e.g. "clear_internal_priors": "tau". Note that any priors that feature the specified parameters will be ignored, i.e. an entire multi-dimensional prior is ignored if any of its parameters are specified in clear_specific_priors. Note that Cobaya prefers to initialise the likelihood itself, hence any modifications of candl_like won’t be reflected in the Cobaya likelihood. The exception are data selections, which are usually propagated correctly. It doesn’t hurt checking the output of the initialisation though to be sure. To use a wrapper likelihood, such as clipy, be sure to set the "warpper" keyword.

candl.interface.get_cobaya_info_dict_for_like(like, name='candl_like', data_selection=Ellipsis, clear_internal_priors=True, clear_specific_priors=[], feedback=True, wrapper=None, additional_args={})#

Thin wrapper for CandlCobayaLikelihood that returns the class with the requested data set in the format expected by Cobaya. Note that since Cobaya prefers to instantiate likelihoods itself, this will instantiate a new instance of the likelihood. That means, any modifications to the likelihood object you pass will not be reflected in the likelihood that Cobaya uses. Still, the wrapper will try to copy over any data selection - it’s worth double checking this with the feedback printout. Use the optional keywords to pass arguments to the new Cobaya likelihood object during initialisation. Note that candl internal priors are cleared by default, be sure to either apply these in the Cobaya parameter block or override the default behaviour.

Parameters:

like: candl.Like, candl.LensLike: A candl likelihood to be used.
name: str: Name to give the likelihood in Cobaya
data_selectionany: Data selection to be used. None, string, list of string, binary mask, or path to a mask are supported.
clear_internal_priorsbool: Whether to clear internal priors.
clear_specific_priorsstr or list: String or list of strings of specific priors to clear. If a given parameter appears in an internal prior, the full prior will be cleared (beware for multi-dimensional priors). Doesn’t do anything if clear_internal_priors is True.
feedbackbool: Whether to print feedback when initialising the likelihood.
wrapperstr: Identifier of a supported external wrapper likelihood.
additional_args: dict: Additional arguments to be passed to the likelihood at initialisation.

Returns:

dict:: Dictionary to use in Cobaya’s ‘likelihood’ entry.

Command-Line#

In order to run Cobaya from the command line it sufficies to include the following lines in the likelihood block.

likelihood:
    candl.interface.CandlCobayaLikelihood:
        data_set_file: candl_data.SPT3G_2018_TTTEEE # Data set or path to .yaml file
        variant: None # Select a variant of the data set if pointing to an index file
        lensing: False # Switch on for lensing likelihoods
        feedback: True # Switch off to hide feedback from candl initialisation
        data_selection: ... # Select a subset of the data set
        clear_internal_priors: True # Switch off to use candl internal priors
        clear_specific_priors: [] # Specify a string or list of strings to clear specific internal priors, e.g. 'tau'
        wrapper: None # Use a wrapper likelihood, e.g. 'clipy'
        additional_args: {} # Additional arguments to pass to the likelihood at initialisation
        add_logdet: False # Whether to add the (constant) logdet term to the likelihood (applies only if the likelihood is a simple Gaussian)

Only data_set_file is required, the other arguments are optional. Again, by default the candl internal priors are not applied, set clear_internal_priors: False if you want to use them. If you only want to clear some internal priors, you can supply a string or list of strings to clear_specific_priors, e.g. "clear_internal_priors": "tau". Note that any priors that feature the specified parameters will be ignored, i.e. an entire multi-dimensional prior is ignored if any of its parameters are specified in clear_specific_priors. If you are pointing to an index file with data_set_file, use variant to select the desired variant.

Tip

If you are encountering issues pointing to candl.interface.CandlCobayaLikelihood this way, try adding external: !!python/name:candl.interface.CandlCobayaLikelihood '' to the block and replacing candl.interface.CandlCobayaLikelihood by a name of your choice, e.g. candl_like.

Tip

If you are resuming a Cobaya run and encounter temperature mismatch errors or inconsistencies - especially when using lensing likelihoods - consider setting add_logdet: True. This can help avoid numerical issues in Cobaya’s internal checks.

Connecting Theory Codes to Cobaya#

In the interface module we also supply an interface to allow for the use of CosmoPower, CosmoPower-JAX, and Capse.jl with Cobaya. This works through custom cobaya.theory.Theory subclasses. For example, to use CosmoPowerJAX in an interactive session with Cobaya, you can use the following code:

import candl.interface

cp_emulator_filenames = {"TT": "your_desired_emu_model"}
cobaya_dict = {"theory": {"CosmoPowerJAX": {"external": candl.interface.CobayaTheoryCosmoPowerJAX,
                                            "emulator_filenames": cp_emulator_filenames}}}

You can then add the other info needed to cobaya_dict and initialise Cobaya. You can similarly use this interface when launching Cobaya from the command line.

MontePython#

In order to interface candl with Montepython, copy the candl_mp folder into your montepython/likelihoods directory. The folder contains the wrapper in the __init__.py file and options can be set in the candl_mp.data file. The provided template for the latter looks like this:

candl_mp.data_set_file = #"candl_data.the_data_set_you_want" or "path/to/info.yaml"
candl_mp.lensing = False
candl_mp.feedback = True
candl_mp.data_selection = ...
candl_mp.clear_internal_priors = True

You have to insert the data set you want in the first line: either the name of a released data set or a path to a .yaml file. The lensing flag must be set to True for lensing data sets.

Note

By default the internal priors declared in candl’s data set .yaml file are ignored. It is assumed that Montepython handles these. If you want to use the internal priors, set clear_internal_priors = False in candl_mp.data.

You can then add candl_mp to the likelihoods in your .param file. Here, MontePython also requires you to add all nuisance parameters for the likelihood. This can be a little tedious, so we provide a helper function to do this for you. For example, for the SPT-3G 2018 TTTEEE likelihood, run the following python code:

import candl
import candl_data
import candl.interface

candl_like = candl.Like(candl_data.SPT3G_2018_TTTEEE)
candl.interface.get_montepython_nuisance_param_block_for_like(candl_like)

This will print the nuisance parameter block to the terminal, which you can then copy over to your .param file.

candl.interface.get_montepython_nuisance_param_block_for_like(like)#

Prints out info needed by Montepython .param file for nuisance parameters.

Parameters:

like: candl.Like: Likelihood to be used.

Returns:

None

Note that the MontePython interface does not currently accept clipy wrapper likelihoods.

CosmoSIS#

A candl interface is included automatically in the latest version of CosmoSIS on GitHub. If you are not able to use this version, you can alternatively copy the candl_cosmosis folder from the candl GitHub into your cosmosis-standard-library/likelihood directory. The folder contains the wrapper (candl_cosmosis_interface.py) as well as the module file required by CosmoSIS. Either way, in order run chains with e.g. the SPT-3G 2018 lensing likelihood, include the following block in your .ini file:

file = ./likelihood/candl/candl_cosmosis_interface.py ; Location of interface code - change depending on the location of your .ini file
data_set = 'candl_data.SPT3G_2018_Lens' ; Data set or path to .yaml file
variant = 'use_CMB' ; Select a variant of the data set if pointing to an index file
lensing = T ; Switch on for lensing likelihoods
feedback = T ; Switch on to request feedback from candl initialisation
data_selection = "..." ; Select a subset of the data set
clear_1d_internal_priors = T ; Switch off to use candl internal 1d priors
clear_nd_internal_priors = F ; Switch on to ignore candl internal higher dimensional priors. Careful: higher-dimensional priors are not implemented in CosmoSIS itself.
force_ignore_transformations = '' ; Backdoor if you want to ignore certain transformations in the data model.

This wrapper was written by Y. Omori and L. Balkenhol, with help from J. Zuntz.

Note

By default the 1-dimensional internal priors declared in candl’s data set .yaml file are ignored, while the multi-dimensional priors are applied. If you want to modify this behaviour, set clear_1d_internal_priors and clear_nd_internal_priors.

Note that the CosmoSIS interface does not currently accept clipy wrapper likelihoods.

BlackJAX#

BlackJAX gradient-based samplers for JAX. We show how to interface with BlackJAX and run NUTS chains in the differentiability tutorial. Below is a quick example using the ACT DR4 likelihood and CosmoPower-JAX. Note that you need to have downloaded the CosmoPower SPT high-accuracy models (available here) and placed them in your cosmopower_jax/trained_models folder. First, we initialise the likelihood and the theory code:

import candl
import candl.interface
import candl.tools
import candl_data

candl_like = candl.Like(candl_data.ACT_DR4_TTTEEE)
cp_emulator_filenames = {"TT": "cmb_spt_TT_NN",
                        "TE": "cmb_spt_TE_PCAplusNN",
                        "EE": "cmb_spt_EE_NN"}
pars_to_theory_specs = candl.interface.get_CosmoPowerJAX_pars_to_theory_specs_func(cp_emulator_filenames)
like = candl.tools.get_params_to_logl_func(candl_like, pars_to_theory_specs)

In the last line we obtain a function that moves from a dictionary of cosmological parameters to the log likelihood value. For BlackJAX, it’s easier to work with a normalised parameter vector as an input. So next, we define the Planck 18 best-fit point as our reference and use the Fisher matrix (from this point) to normalise the parameters.

import jax
import numpy as np

fid_pars = {'H0': 67.37, 'ombh2': 0.02233, 'omch2': 0.1198, 'logA': 3.043, 'ns': 0.9652, 'tau': 0.054, 'yp': 1.0}
fid_pars_vec = np.array([fid_pars[p] for p in par_order])
par_cov, par_order = candl.tools.get_fisher_matrix(pars_to_theory_specs, candl_like, fid_pars, par_order=None, return_par_order=True)

def norm_pars(pars):
    return (pars-fid_pars_vec)/np.sqrt(np.diag(par_cov))

def denorm_pars(pars):
    return pars*np.sqrt(np.diag(par_cov))+fid_pars_vec

@jax.jit
def like_normed_vec(input_vec):
    denormed_vec = denorm_pars(input_vec)
    return like({par_order[i]: denormed_vec[i] for i in range(len(par_order))})

Now we have a function that takes a normalised parameter vector and returns the log likelihood value. Following the NUTS example provided by BlackJAX (here) we define the inference loop.

# This code is taken from the BlackJAX tutorial
def inference_loop(rng_key, kernel, initial_state, num_samples):
    @jax.jit
    def one_step(state, rng_key):
        state, _ = kernel(rng_key, state)
        return state, state
    keys = jax.random.split(rng_key, num_samples)
    _, states = jax.lax.scan(one_step, initial_state, keys)
    return states

Finally we feed this into BlackJAX and run the sampler.

import blackjax
from datetime import date

rng_key = jax.random.key(int(date.today().strftime("%Y%m%d")))# grab random number for starting point
nuts = blackjax.nuts(like_normed_vec, step_size=0.1, inverse_mass_matrix=np.ones(len(par_order)))
initial_state = nuts.init(np.zeros(len(par_order)))
states = inference_loop(rng_key,
                        nuts.step,
                        initial_state,
                        100)# Number of desired samples
NUTS_samples = jax.numpy.apply_along_axis(denorm_pars, 1, states.position)

Where in the last line we move from the normalised parameter space back to our familiar one (i.e. get the right units).

Optax#

Optax is a library for gradient-based optimisation written for JAX. We show how to interface with Optax and run the ADAM minimiser in the differentiability tutorial. Below is a quick example using the ACT DR4 likelihood and CosmoPower-JAX. Note that you need to have downloaded the CosmoPower SPT high-accuracy models (available here) and placed them in your cosmopower_jax/trained_models folder. Like before, we initialise the likelihood and the theory code:

import candl
import candl.interface
import candl.tools
import candl_data

candl_like = candl.Like(candl_data.ACT_DR4_TTTEEE)
cp_emulator_filenames = {"TT": "cmb_spt_TT_NN",
                        "TE": "cmb_spt_TE_PCAplusNN",
                        "EE": "cmb_spt_EE_NN"}
pars_to_theory_specs = candl.interface.get_CosmoPowerJAX_pars_to_theory_specs_func(cp_emulator_filenames)
like = candl.tools.get_params_to_logl_func(candl_like, pars_to_theory_specs)

Optax also prefers normalised input parameters, so as before we define the Planck 18 best-fit point as our reference and use the Fisher matrix (from this point) to normalise the parameters.

import jax
import numpy as np

fid_pars = {'H0': 67.37, 'ombh2': 0.02233, 'omch2': 0.1198, 'logA': 3.043, 'ns': 0.9652, 'tau': 0.054, 'yp': 1.0}
par_cov, par_order = candl.tools.get_fisher_matrix(pars_to_theory_specs, candl_like, fid_pars, par_order=None, return_par_order=True)
par_scales = {par_order[i]: np.sqrt(par_cov[i,i]) for i in range(len(par_order))}

def transform_to_zero_mean_unit_var(par_dict):
    new_par_dict = {}
    for p in par_dict:
        if p in fid_pars and p in par_order:
            new_par_dict[p] = (par_dict[p] - fid_pars[p])/par_scales[p]
    return new_par_dict

def transform_from_zero_mean_unit_var(par_dict):
    new_par_dict = {}
    for p in par_dict:
        if p in fid_pars and p in par_order:
            new_par_dict[p] = fid_pars[p] + par_dict[p]*par_scales[p]
    return new_par_dict

like_normed = jax.jit(lambda p: -1.0*like(transform_from_zero_mean_unit_var(p)))
like_normed_deriv = jax.jit(jax.jacfwd(like_normed))

Note the minus sign in the normalised likelihood function; Optax uses a different sign convention (expeting the logl to be positive). We are now ready to initialise and run the ADAM minimiser.

import optax
from copy import deepcopy

# Initialise the ADAM minimiser and starting point
adam_optimiser = optax.adam(learning_rate = 0.75)
starting_pars = transform_to_zero_mean_unit_var(fid_pars)
opt_state = adam_optimiser.init(starting_pars)

# Minimise!
this_pars = deepcopy(starting_pars)
for i_adam in range(100):

    # Get like value and gradient
    this_logl = like_normed(this_pars)
    this_grad = like_normed_deriv(this_pars)

    # Pass information to optax and update
    updates, opt_state = adam_optimiser.update(this_grad,
                                            opt_state,
                                            this_pars)
    this_pars = optax.apply_updates(this_pars,
                                    updates)

bf_point = transform_from_zero_mean_unit_var(this_pars)