def mcmc_emcee(self, n_walkers, n_run, n_burn, mean_start, sigma_start):
     """
     returns the mcmc analysis of the parameter space
     """
     sampler = emcee.EnsembleSampler(n_walkers, self.cosmoParam.numParam, self.chain.likelihood)
     p0 = emcee.utils.sample_ball(mean_start, sigma_start, n_walkers)
     new_pos, _, _, _ = sampler.run_mcmc(p0, n_burn)
     sampler.reset()
     store = InMemoryStorageUtil()
     for pos, prob, _, _ in sampler.sample(new_pos, iterations=n_run):
         store.persistSamplingValues(pos, prob, None)
     return store.samples
    def mcmc_CH(self, walkerRatio, n_run, n_burn, mean_start, sigma_start, threadCount=1, init_pos=None, mpi_monch=False):
        """
        runs mcmc on the parameter space given parameter bounds with CosmoHammerSampler
        returns the chain
        """
        lowerLimit, upperLimit = self.cosmoParam.param_bounds
        params = np.array([mean_start, lowerLimit, upperLimit, sigma_start]).T

        chain = LikelihoodComputationChain(
            min=lowerLimit,
            max=upperLimit)

        temp_dir = tempfile.mkdtemp("Hammer")
        file_prefix = os.path.join(temp_dir, "logs")

        # chain.addCoreModule(CambCoreModule())
        chain.addLikelihoodModule(self.chain)
        chain.setup()

        store = InMemoryStorageUtil()
        if mpi_monch is True:
            sampler = MpiCosmoHammerSampler(
            params=params,
            likelihoodComputationChain=chain,
            filePrefix=file_prefix,
            walkersRatio=walkerRatio,
            burninIterations=n_burn,
            sampleIterations=n_run,
            threadCount=1,
            initPositionGenerator=init_pos,
            storageUtil=store)
        else:
            sampler = CosmoHammerSampler(
                params=params,
                likelihoodComputationChain=chain,
                filePrefix=file_prefix,
                walkersRatio=walkerRatio,
                burninIterations=n_burn,
                sampleIterations=n_run,
                threadCount=threadCount,
                initPositionGenerator=init_pos,
                storageUtil=store)
        time_start = time.time()
        if sampler.isMaster():
            print('Computing the MCMC...')
            print('Number of walkers = ', len(mean_start)*walkerRatio)
            print('Burn-in itterations: ', n_burn)
            print('Sampling itterations:', n_run)
        sampler.startSampling()
        if sampler.isMaster():
            time_end = time.time()
            print(time_end - time_start, 'time taken for MCMC sampling')
        # if sampler._sampler.pool is not None:
        #     sampler._sampler.pool.close()
        try:
            shutil.rmtree(temp_dir)
        except Exception as ex:
            print(ex)
            pass
        return store.samples
class TestCosmoHammerSampler(object):
    storageUtil = None
    samples = np.array([[1, 2, 3], [4, 5, 6]])
    prob = np.array([1, 1])

    def setup(self):
        self.storageUtil = InMemoryStorageUtil()

    def test_no_data(self):

        assert self.storageUtil.samplesBurnin == None
        assert self.storageUtil.probBurnin == None

    def test_persistBurninValues(self):
        self.storageUtil.persistBurninValues(self.samples, self.prob, [])

        assert self.storageUtil.samplesBurnin is not None
        assert self.storageUtil.probBurnin is not None
        assert self.storageUtil.samplesBurnin.shape == (2, 3)
        assert self.storageUtil.probBurnin.shape == (2, )

        self.storageUtil.persistBurninValues(self.samples, self.prob, [])

        assert self.storageUtil.samplesBurnin is not None
        assert self.storageUtil.probBurnin is not None
        assert self.storageUtil.samplesBurnin.shape == (4, 3)
        assert self.storageUtil.probBurnin.shape == (4, )

    def test_persistSamplingValues(self):
        self.storageUtil.persistSamplingValues(self.samples, self.prob, [])

        assert self.storageUtil.samples is not None
        assert self.storageUtil.prob is not None
        assert self.storageUtil.samples.shape == (2, 3)
        assert self.storageUtil.prob.shape == (2, )

        self.storageUtil.persistSamplingValues(self.samples, self.prob, [])

        assert self.storageUtil.samples is not None
        assert self.storageUtil.prob is not None
        assert self.storageUtil.samples.shape == (4, 3)
        assert self.storageUtil.prob.shape == (4, )
class TestCosmoHammerSampler(object):
    storageUtil = None
    samples = np.array([[1, 2, 3], [4, 5, 6]])
    prob = np.array([1, 1])

    def setup(self):
        self.storageUtil = InMemoryStorageUtil()

    def test_no_data(self):

        assert self.storageUtil.samplesBurnin == None
        assert self.storageUtil.probBurnin == None

    def test_persistBurninValues(self):
        self.storageUtil.persistBurninValues(self.samples, self.prob, [])

        assert self.storageUtil.samplesBurnin is not None
        assert self.storageUtil.probBurnin is not None
        assert self.storageUtil.samplesBurnin.shape == (2, 3)
        assert self.storageUtil.probBurnin.shape == (2,)

        self.storageUtil.persistBurninValues(self.samples, self.prob, [])

        assert self.storageUtil.samplesBurnin is not None
        assert self.storageUtil.probBurnin is not None
        assert self.storageUtil.samplesBurnin.shape == (4, 3)
        assert self.storageUtil.probBurnin.shape == (4,)

    def test_persistSamplingValues(self):
        self.storageUtil.persistSamplingValues(self.samples, self.prob, [])

        assert self.storageUtil.samples is not None
        assert self.storageUtil.prob is not None
        assert self.storageUtil.samples.shape == (2, 3)
        assert self.storageUtil.prob.shape == (2,)

        self.storageUtil.persistSamplingValues(self.samples, self.prob, [])

        assert self.storageUtil.samples is not None
        assert self.storageUtil.prob is not None
        assert self.storageUtil.samples.shape == (4, 3)
        assert self.storageUtil.prob.shape == (4,)
 def setup(self):
     self.storageUtil = InMemoryStorageUtil()
 def setup(self):
     self.storageUtil = InMemoryStorageUtil()