def sample_emcee(model,
data,
nwalkers,
nsamples,
walker_initial_pos,
parallel='auto',
seed=None):
if _EMCEE_MISSING:
raise DependencyMissing(
'emcee', "Install it with \'conda install -c conda-forge emcee\'.")
obj_func = LnpostWrapper(model, data)
pool = choose_pool(parallel)
sampler = emcee.EnsembleSampler(nwalkers,
len(model._parameters),
obj_func.evaluate,
pool=pool)
if seed is not None:
np.random.seed(seed)
seed_state = np.random.mtrand.RandomState(seed).get_state()
sampler.random_state = seed_state
sampler.run_mcmc(walker_initial_pos, nsamples)
if pool is not parallel:
pool.close()
return sampler
def test_MPI(self):
if MPIPool.enabled():
# mpi should work
mpi_pool = choose_pool('mpi')
self.assertTrue(isinstance(mpi_pool, MPIPool))
elif find_spec('mpi4py') is None:
# mpi4py not installed
self.assertRaises(ImportError, choose_pool, 'mpi')
else:
# mpi4py installed but only one process available
self.assertRaises(ValueError, choose_pool, 'mpi')
def run_cma(obj_func, parameters, initial_population, weight_function,
tols={}, seed=None, parallel='auto'):
"""
instantiate and run a CMAEvolutionStrategy object
Parameters
----------
obj_func : Function
function to be minimized (not maximized like posterior)
parameters : list of Prior objects
parameters to fit
initial_population: array
starting population with shape = (popsize, len(parameters))
weight_function: function
takes arguments (i, popsize), i in range(popsize); returns weight of i
tols: dict, optional
tolerance values to overwrite the cma defaults
seed: int, optional
random seed to use
parallel: optional
number of threads to use or pool object or one of {None, 'all', 'mpi'}.
Default tries 'mpi' then 'all'.
"""
if _CMA_MISSING:
raise DependencyMissing('cma', "Install it with \'pip install cma\'.")
popsize = len(initial_population)
stds = [par.sd if isinstance(par, prior.Gaussian)
else par.interval/4 for par in parameters]
weights = [weight_function(i, popsize) for i in range(popsize)]
if weights[-1] > 0:
weights[-1] = 0
warnings.warn('Setting weight of worst parent to 0')
with tempfile.TemporaryDirectory() as tempdir:
cmaoptions = {'CMA_stds': stds, 'CMA_recombination_weights': weights,
'verb_filenameprefix': tempdir, 'verbose': -3}
cmaoptions.update(tols)
if seed is not None:
cmaoptions.update({'seed': seed})
guess = [par.guess for par in parameters]
cma_strategy = cma.CMAEvolutionStrategy(guess, 1, cmaoptions)
cma_strategy.inject(initial_population, force=True)
solutions = np.zeros((popsize, len(parameters)))
func_vals = np.zeros(popsize)
pool = choose_pool(parallel)
while not cma_strategy.stop():
invalid = np.ones(popsize, dtype=bool)
inf_replace_counter = 0
while invalid.any() and inf_replace_counter < 10:
attempts = cma_strategy.ask(np.sum(invalid))
solutions[invalid, :] = attempts
func_vals[invalid] = list(pool.map(obj_func, attempts))
invalid = ~np.isfinite(func_vals)
inf_replace_counter += 1 # catches case where all are inf
cma_strategy.tell(solutions, func_vals)
cma_strategy.logger.add()
cma_strategy.logger.load()
if pool is not parallel:
# I made pool, responsible for closing it.
pool.close()
return cma_strategy
def test_auto(self):
auto_pool = choose_pool('auto')
self.assertTrue(isinstance(auto_pool, (pool.BasePool, mp.pool.Pool)))
def test_schwimmbad_multipool(self):
multi_pool = choose_pool(5)
self.assertTrue(isinstance(multi_pool, MultiPool))
self.assertEqual(multi_pool._processes, 5)
def test_counting_all_cores(self):
all_pool = choose_pool('all')
self.assertTrue(isinstance(all_pool, MultiPool))
self.assertEqual(all_pool._processes, mp.cpu_count())
def test_nonepool(self):
none_pool = choose_pool(None)
self.assertFalse(isinstance(none_pool, (pool.BasePool, mp.pool.Pool)))
self.assertEqual(list(none_pool.map(len, [[0, 1, 2], 'asdf'])), [3, 4])
self.assertTrue(hasattr(none_pool, "close"))
def test_multiprocessing_pool(self):
mp_pool = mp.pool.Pool(5)
self.assertTrue(choose_pool(mp_pool) is mp_pool)
def test_custom_pool(self):
custom_pool = DummyPool(17)
chosen_pool = choose_pool(custom_pool)
self.assertTrue(choose_pool(custom_pool) is custom_pool)
