def pool(request): multimode = 'None' # multimode = 'Serial' # multimode = 'Multi' # multimode = 'MPI' # setup code pool = None if multimode == 'Serial': from schwimmbad import SerialPool pool = SerialPool() elif multimode == 'Multi': from schwimmbad import MultiPool pool = MultiPool() elif multimode == 'MPI': from schwimmbad import MPIPool pool = MPIPool() if not pool.is_master(): pool.wait() import sys sys.exit(0) # inject class variables request.cls.pool = pool yield # tear down if multimode == 'Multi' or multimode == 'MPI': pool.close()
def sample(self, n_samples): if self.pool is None or _GPU_ENABLED: pool = SerialPool() else: if isinstance(self.pool, int): pool = MultiPool(self.pool) elif isinstance(self.pool, (SerialPool, MultiPool)): pool = self.pool else: raise TypeError( "Does not understand the given multiprocessing pool.") drawn_samples = list( pool.map(self.draw_one_joint_posterior_sample_map, range(n_samples))) pool.close() drawn_zs = [drawn_samples[i][0] for i in range(n_samples)] drawn_inference_posteriors = [ drawn_samples[i][1] for i in range(n_samples) ] drawn_joint_posterior_samples = pd.DataFrame( drawn_inference_posteriors) drawn_joint_posterior_samples["redshift"] = drawn_zs return drawn_joint_posterior_samples
def compute_mean_selection_function(selection_function, N_avg, pool=None): if pool is None: pool = SerialPool() elif isinstance(pool, int): pool = MultiPool(pool) elif isinstance(pool, (SerialPool, MultiPool)): pool = pool else: raise TypeError("Does not understand the given multiprocessing pool.") out = list( pool.starmap(selection_function.evaluate, [() for _ in range(N_avg)])) avg = np.average(out) pool.close() return avg