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)