def test_choose_pool(self): import schwimmbad pool = choose_pool(mpi=False, processes=1, use_dill=True) assert pool.is_master() is True assert isinstance(pool, schwimmbad.serial.SerialPool) pool = choose_pool(mpi=False, processes=2, use_dill=True) assert pool.is_master() is True assert isinstance(pool, MultiPool)
def optimize(self, n_particles=50, n_iterations=250, verbose=False, threadCount=1): """ :param n_particles: number of PSO particles, will be ignored if self._particle_swarm is False :param n_iterations: number of PSO iterations, will be ignored if self._particle_swarm is False :param verbose: whether to print stuff :param threadCount: integer; number of threads in multi-threading mode :return: keyword arguments that map (x_image, y_image) to the same source coordinate (source_x, source_y) """ if self._particle_swarm: if threadCount > 1: pool = choose_pool(mpi=False, processes=threadCount) else: pool = None kwargs = self._fit_pso(n_particles, n_iterations, pool, verbose) else: kwargs = self._param_class.kwargs_lens kwargs_lens_final, source_penalty = self._fit_amoeba(kwargs, verbose) args_lens_final = self._param_class.kwargs_to_args(kwargs_lens_final) source_x_array, source_y_array = self.fast_rayshooting.ray_shooting_fast(args_lens_final) source_x, source_y = np.mean(source_x_array), np.mean(source_y_array) if verbose: print('optimization done.') print('Recovered source position: ', (source_x_array, source_y_array)) return kwargs_lens_final, [source_x, source_y]
def pso(self, n_particles, n_iterations, lower_start=None, upper_start=None, threadCount=1, init_pos=None, mpi=False, print_key='PSO'): """ Return the best fit for the lens model on catalogue basis with particle swarm optimizer. :param n_particles: number of particles in the sampling process :param n_iterations: number of iterations of the swarm :param lower_start: numpy array, lower end parameter of the values of the starting particles :param upper_start: numpy array, upper end parameter of the values of the starting particles :param threadCount: number of threads in the computation (only applied if mpi=False) :param init_pos: numpy array, position of the initial best guess model :param mpi: bool, if True, makes instance of MPIPool to allow for MPI execution :param print_key: string, prints the process name in the progress bar (optional) :return: kwargs_result (of best fit), [lnlikelihood of samples, positions of samples, velocity of sampels) """ if lower_start is None or upper_start is None: lower_start, upper_start = np.array(self.lower_limit), np.array(self.upper_limit) print("PSO initialises its particles with default values") else: lower_start = np.maximum(lower_start, self.lower_limit) upper_start = np.minimum(upper_start, self.upper_limit) pool = choose_pool(mpi=mpi, processes=threadCount, use_dill=True) if mpi is True and pool.is_master(): print('MPI option chosen for PSO.') pso = ParticleSwarmOptimizer(self.chain.logL, lower_start, upper_start, n_particles, pool=pool) if init_pos is None: init_pos = (upper_start - lower_start) / 2 + lower_start pso.set_global_best(init_pos, [0]*len(init_pos), self.chain.logL(init_pos)) if pool.is_master(): print('Computing the %s ...' % print_key) time_start = time.time() result, [chi2_list, pos_list, vel_list] = pso.optimize(n_iterations) if pool.is_master(): kwargs_return = self.chain.param.args2kwargs(result) print(pso.global_best.fitness * 2 / (max( self.chain.effective_num_data_points(**kwargs_return), 1)), 'reduced X^2 of best position') print(pso.global_best.fitness, 'logL') print(self.chain.effective_num_data_points(**kwargs_return), 'effective number of data points') print(kwargs_return.get('kwargs_lens', None), 'lens result') print(kwargs_return.get('kwargs_source', None), 'source result') print(kwargs_return.get('kwargs_lens_light', None), 'lens light result') print(kwargs_return.get('kwargs_ps', None), 'point source result') print(kwargs_return.get('kwargs_special', None), 'special param result') time_end = time.time() print(time_end - time_start, 'time used for ', print_key) print('===================') return result, [chi2_list, pos_list, vel_list]
def pso(self, n_particles=10, n_iterations=10, lowerLimit=-0.2, upperLimit=0.2, threadCount=1, mpi=False, print_key='default'): """ returns the best fit for the lense model on catalogue basis with particle swarm optimizer """ init_pos = self.chain.get_args(self.chain.kwargs_data_init) num_param = self.chain.num_param lowerLimit = [lowerLimit] * num_param upperLimit = [upperLimit] * num_param pool = choose_pool(mpi=mpi, processes=threadCount, use_dill=True) pso = ParticleSwarmOptimizer(self.chain, lowerLimit, upperLimit, n_particles, pool=pool) if init_pos is not None: pso.set_global_best(init_pos, [0]*len(init_pos), self.chain.likelihood(init_pos)) if pool.is_master(): print('Computing the %s ...' % print_key) time_start = time.time() result, [chi2_list, pos_list, vel_list] = pso.optimize(n_iterations) kwargs_data = self.chain.update_data(result) if pool.is_master(): time_end = time.time() print("Shifts found: ", result) print(time_end - time_start, 'time used for ', print_key) return kwargs_data, [chi2_list, pos_list, vel_list]
def nautilus_sampling(self, prior_type='uniform', mpi=False, thread_count=1, verbose=True, one_step=False, **kwargs_nautilus): """ :param prior_type: string; prior type. Currently only 'uniform' supported (in addition to Prior class in Likelihood module) :param mpi: MPI option (currently not supported) :param thread_count: integer; multi-threading option (currently not supported) :param verbose: verbose statements of Nautilus :param one_step: boolean, if True, only runs one iteration of filling the sampler and re-training. This is meant for test purposes of the sampler to operate with little computational effort :param kwargs_nautilus: additional keyword arguments for Nautilus :return: points, log_w, log_l, log_z """ from nautilus import Prior, Sampler prior = Prior() # TODO better prior integration with Nautilus if prior_type == 'uniform': for i in range(self._num_param): prior.add_parameter(dist=(self._lower_limit[i], self._upper_limit[i])) # assert self._num_param == prior.dimensionality() # print(self._num_param, prior.dimensionality(), 'number of param, dimensionality') else: raise ValueError( 'prior_type %s is not supported for Nautilus wrapper.' % prior_type) # loop through prior pool = choose_pool(mpi=mpi, processes=thread_count, use_dill=True) sampler = Sampler(prior, likelihood=self.likelihood, pool=pool, **kwargs_nautilus) time_start = time.time() if one_step is True: sampler.add_bound() sampler.fill_bound() else: sampler.run(verbose=verbose) points, log_w, log_l = sampler.posterior() log_z = sampler.evidence() time_end = time.time() if pool.is_master(): print(time_end - time_start, 'time taken for MCMC sampling') return points, log_w, log_l, log_z
def mcmc_emcee(self, n_walkers, n_run, n_burn, mean_start, sigma_start, mpi=False, progress=False, threadCount=1, initpos=None, backup_filename=None, start_from_backup=False): """ Run MCMC with emcee. For details, please have a look at the documentation of the emcee packager. :param n_walkers: number of walkers in the emcee process :type n_walkers: integer :param n_run: number of sampling (after burn-in) of the emcee :type n_run: integer :param n_burn: number of burn-in iterations (those will not be saved in the output sample) :type n_burn: integer :param mean_start: mean of the parameter position of the initialising sample :type mean_start: numpy array of length the number of parameters :param sigma_start: spread of the parameter values (uncorrelated in each dimension) of the initialising sample :type sigma_start: numpy array of length the number of parameters :param mpi: if True, initializes an MPIPool to allow for MPI execution of the sampler :type mpi: bool :param progress: if True, prints the progress bar :type progress: bool :param threadCount: number of threats in multi-processing (not applicable for MPI) :type threadCount: integer :param initpos: initial walker position to start sampling (optional) :type initpos: numpy array of size num param x num walkser :param backup_filename: name of the HDF5 file where sampling state is saved (through emcee backend engine) :type backup_filename: string :param start_from_backup: if True, start from the state saved in `backup_filename`. Otherwise, create a new backup file with name `backup_filename` (any already existing file is overwritten!). :type start_from_backup: bool :return: samples, ln likelihood value of samples :rtype: numpy 2d array, numpy 1d array """ num_param, _ = self.chain.param.num_param() if initpos is None: initpos = sampling_util.sample_ball(mean_start, sigma_start, n_walkers, dist='normal') pool = choose_pool(mpi=mpi, processes=threadCount, use_dill=True) if backup_filename is not None: backend = emcee.backends.HDFBackend(backup_filename, name="lenstronomy_mcmc_emcee") if pool.is_master(): print("Warning: All samples (including burn-in) will be saved in backup file '{}'.".format(backup_filename)) if start_from_backup: initpos = None n_run_eff = n_run else: n_run_eff = n_burn + n_run backend.reset(n_walkers, num_param) if pool.is_master(): print("Warning: backup file '{}' has been reset!".format(backup_filename)) else: backend = None n_run_eff = n_burn + n_run time_start = time.time() sampler = emcee.EnsembleSampler(n_walkers, num_param, self.chain.logL, pool=pool, backend=backend) sampler.run_mcmc(initpos, n_run_eff, progress=progress) flat_samples = sampler.get_chain(discard=n_burn, thin=1, flat=True) dist = sampler.get_log_prob(flat=True, discard=n_burn, thin=1) if pool.is_master(): print('Computing the MCMC...') print('Number of walkers = ', n_walkers) print('Burn-in iterations: ', n_burn) print('Sampling iterations (in current run):', n_run_eff) time_end = time.time() print(time_end - time_start, 'time taken for MCMC sampling') return flat_samples, dist
def mcmc_zeus(self, n_walkers, n_run, n_burn, mean_start, sigma_start, mpi=False, threadCount=1, progress=False, initpos=None, backend_filename=None, moves=None, tune=True, tolerance=0.05, patience=5, maxsteps=10000, mu=1.0, maxiter=10000, pool=None, vectorize=False, blobs_dtype=None, verbose=True, check_walkers=True, shuffle_ensemble=True, light_mode=False): """ Lightning fast MCMC with zeus: https://github.com/minaskar/zeus For the full list of arguments for the EnsembleSampler, see see `the zeus docs <https://zeus-mcmc.readthedocs.io/en/latest/api/sampler.html>`_. If you use the zeus sampler, you should cite the following papers: 2105.03468, 2002.06212. :param n_walkers: number of walkers per parameter :type n_walkers: integer :param n_run: number of sampling steps :type n_run: integer :param n_burn: number of burn-in steps :type n_burn: integer :param mean_start: mean of the parameter position of the initialising sample :type mean_start: numpy array of length the number of parameters :param sigma_start: spread of the parameter values (uncorrelated in each dimension) of the initialising sample :type sigma_start: numpy array of length the number of parameters :param progress: :type progress: bool :param initpos: initial walker position to start sampling (optional) :type initpos: numpy array of size num param x num walkser :param backup_filename: name of the HDF5 file where sampling state is saved (through zeus callback function) :type backup_filename: string :return: samples, ln likelihood value of samples :rtype: numpy 2d array, numpy 1d array """ import zeus print('Using zeus to perform the MCMC.') num_param, _ = self.chain.param.num_param() if initpos is None: initpos = sampling_util.sample_ball_truncated(mean_start, sigma_start, self.lower_limit, self.upper_limit, size=n_walkers) if backend_filename is not None: backend = zeus.callbacks.SaveProgressCallback( filename=backend_filename, ncheck=1) n_run_eff = n_burn + n_run else: backend = None n_run_eff = n_burn + n_run pool = choose_pool(mpi=mpi, processes=threadCount, use_dill=True) sampler = zeus.EnsembleSampler(nwalkers=n_walkers, ndim=num_param, logprob_fn=self.chain.logL, moves=moves, tune=tune, tolerance=tolerance, patience=patience, maxsteps=maxsteps, mu=mu, maxiter=maxiter, pool=pool, vectorize=vectorize, blobs_dtype=blobs_dtype, verbose=verbose, check_walkers=check_walkers, shuffle_ensemble=shuffle_ensemble, light_mode=light_mode) sampler.run_mcmc(initpos, n_run_eff, progress=progress, callbacks=backend) flat_samples = sampler.get_chain(flat=True, thin=1, discard=n_burn) dist = sampler.get_log_prob(flat=True, thin=1, discard=n_burn) return flat_samples, dist