Source code for threeML.bayesian.emcee_sampler

from typing import Dict, List, Optional

import emcee
import numpy as np
from astromodels import ModelAssertionViolation, use_astromodels_memoization

from threeML.bayesian.sampler_base import MCMCSampler
from threeML.config import threeML_config
from threeML.config.config import threeML_config
from import is_inside_notebook
from import setup_logger
from threeML.parallel.parallel_client import ParallelClient

log = setup_logger(__name__)

[docs]class EmceeSampler(MCMCSampler): def __init__(self, likelihood_model=None, data_list=None, **kwargs): """ Sample using the emcee sampler. For details: :param likelihood_model: :param data_list: :returns: :rtype: """ super(EmceeSampler, self).__init__( likelihood_model, data_list, **kwargs)
[docs] def setup(self, n_iterations: int, n_burn_in: Optional[int] = None, n_walkers: int = 20, seed=None, **kwargs): """TODO describe function :param n_iterations: :type n_iterations: int :param n_burn_in: :type n_burn_in: Optional[int] :param n_walkers: :type n_walkers: int :param seed: :type seed: :returns: """ log.debug(f"Setup for Emcee sampler: n_iterations:{n_iterations}, n_burn_in:{n_burn_in}," f"n_walkers: {n_walkers}, seed: {seed}.") self._n_iterations = int(n_iterations) if n_burn_in is None: self._n_burn_in = int(np.floor(n_iterations / 4.0)) else: self._n_burn_in = n_burn_in self._n_walkers = int(n_walkers) self._seed = seed self._kwargs = kwargs # we control progress with the config if "progress" in self._kwargs: _ = self._kwargs.pop("progress") self._is_setup = True
[docs] def sample(self, quiet=False): if not self._is_setup:"You forgot to setup the sampler!") return loud = not quiet self._update_free_parameters() n_dim = len(list(self._free_parameters.keys())) # Get starting point p0 = emcee.State(self._get_starting_points(self._n_walkers)) # Deactivate memoization in astromodels, which is useless in this case since we will never use twice the # same set of parameters with use_astromodels_memoization(False): if threeML_config["parallel"]["use_parallel"]: c = ParallelClient() view = c[:] sampler = emcee.EnsembleSampler( self._n_walkers, n_dim, self.get_posterior, pool=view ) else: sampler = emcee.EnsembleSampler( self._n_walkers, n_dim, self.get_posterior ) # If a seed is provided, set the random number seed if self._seed is not None: sampler._random.seed(self._seed) log.debug("Start emcee run") # Sample the burn-in if threeML_config.interface.progress_bars: if is_inside_notebook(): progress = "notebook" else: progress = True else: progress = False pos, prob, state = sampler.run_mcmc( initial_state=p0, nsteps=self._n_burn_in, progress=progress ) log.debug("Emcee run done") # Reset sampler sampler.reset() state = emcee.State(pos, prob, random_state=state) # Run the true sampling _ = sampler.run_mcmc( initial_state=state, nsteps=self._n_iterations, progress=progress) acc=np.mean(sampler.acceptance_fraction)"Mean acceptance fraction: {acc}") self._sampler=sampler self._raw_samples=sampler.get_chain(flat=True) # Compute the corresponding values of the likelihood # First we need the prior log_prior=[self._log_prior(x) for x in self._raw_samples] # Now we get the log posterior and we remove the log prior self._log_like_values=sampler.get_log_prob(flat=True) - log_prior # we also want to store the log probability self._log_probability_values=sampler.get_log_prob(flat=True) self._marginal_likelihood=None self._build_samples_dictionary() self._build_results() # Display results if loud: self._results.display() return self.samples