Exemplo n.º 1
0
    def mcmc(
            self,
            n_burn,
            n_run,
            walkerRatio=None,
            n_walkers=None,
            sigma_scale=1,
            threadCount=1,
            init_samples=None,
            re_use_samples=True,
            sampler_type='EMCEE',
            progress=True,
            backend_filename=None,
            start_from_backend=False,
            # zeus specific kwargs
            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):
        """
        MCMC routine

        :param n_burn: number of burn in iterations (will not be saved)
        :param n_run: number of MCMC iterations that are saved
        :param walkerRatio: ratio of walkers/number of free parameters
        :param n_walkers: integer, number of walkers of emcee (optional, if set, overwrites the walkerRatio input
        :param sigma_scale: scaling of the initial parameter spread relative to the width in the initial settings
        :param threadCount: number of CPU threads. If MPI option is set, threadCount=1
        :param init_samples: initial sample from where to start the MCMC process
        :param re_use_samples: bool, if True, re-uses the samples described in init_samples.nOtherwise starts from
         scratch.
        :param sampler_type: string, which MCMC sampler to be used. Options are: 'EMCEE', 'ZEUS'
        :param progress: boolean, if True shows progress bar in EMCEE
        :param backend_filename: name of the HDF5 file where sampling state is saved (through emcee backend engine)
        :type backend_filename: string
        :param start_from_backend: 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_backend: bool
        :return: list of output arguments, e.g. MCMC samples, parameter names, logL distances of all samples specified
         by the specific sampler used
        """

        param_class = self.param_class
        # run PSO
        mcmc_class = Sampler(likelihoodModule=self.likelihoodModule)
        kwargs_temp = self._updateManager.parameter_state
        mean_start = param_class.kwargs2args(**kwargs_temp)
        kwargs_sigma = self._updateManager.sigma_kwargs
        sigma_start = np.array(
            param_class.kwargs2args(**kwargs_sigma)) * sigma_scale
        num_param, param_list = param_class.num_param()
        if n_walkers is None:
            if walkerRatio is None:
                raise ValueError(
                    'MCMC sampler needs either n_walkers or walkerRatio as input argument'
                )
            n_walkers = num_param * walkerRatio
        # run MCMC
        if init_samples is not None and re_use_samples is True:
            num_samples, num_param_prev = np.shape(init_samples)
            if num_param_prev == num_param:
                print(
                    "re-using previous samples to initialize the next MCMC run."
                )
                idxs = np.random.choice(len(init_samples), n_walkers)
                initpos = init_samples[idxs]
            else:
                raise ValueError(
                    "Can not re-use previous MCMC samples as number of parameters have changed!"
                )
        else:
            initpos = None

        if sampler_type == 'EMCEE':
            samples, dist = mcmc_class.mcmc_emcee(
                n_walkers,
                n_run,
                n_burn,
                mean_start,
                sigma_start,
                mpi=self._mpi,
                threadCount=threadCount,
                progress=progress,
                initpos=initpos,
                backend_filename=backend_filename,
                start_from_backend=start_from_backend)
            output = [sampler_type, samples, param_list, dist]

        elif sampler_type == 'ZEUS':
            samples, dist = mcmc_class.mcmc_zeus(
                n_walkers,
                n_run,
                n_burn,
                mean_start,
                sigma_start,
                mpi=self._mpi,
                threadCount=threadCount,
                progress=progress,
                initpos=initpos,
                backend_filename=backend_filename,
                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)
            output = [sampler_type, samples, param_list, dist]
        else:
            raise ValueError('sampler_type %s not supported!' % sampler_type)
        self._mcmc_init_samples = samples  # overwrites previous samples to continue from there in the next MCMC run
        return output
Exemplo n.º 2
0
class TestSampler(object):
    """
    test the fitting sequences
    """
    def setup(self):

        # data specifics
        sigma_bkg = 0.05  # background noise per pixel
        exp_time = 100  # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit)
        numPix = 10  # cutout pixel size
        deltaPix = 0.1  # pixel size in arcsec (area per pixel = deltaPix**2)
        fwhm = 0.5  # full width half max of PSF

        # PSF specification

        kwargs_data = sim_util.data_configure_simple(numPix, deltaPix,
                                                     exp_time, sigma_bkg)
        data_class = ImageData(**kwargs_data)
        kwargs_psf_gaussian = {
            'psf_type': 'GAUSSIAN',
            'fwhm': fwhm,
            'pixel_size': deltaPix
        }
        psf = PSF(**kwargs_psf_gaussian)
        kwargs_psf = {
            'psf_type': 'PIXEL',
            'kernel_point_source': psf.kernel_point_source
        }
        psf_class = PSF(**kwargs_psf)
        kwargs_spemd = {
            'theta_E': 1.,
            'gamma': 1.8,
            'center_x': 0,
            'center_y': 0,
            'e1': 0.1,
            'e2': 0.1
        }

        lens_model_list = ['SPEP']
        self.kwargs_lens = [kwargs_spemd]
        lens_model_class = LensModel(lens_model_list=lens_model_list)
        kwargs_sersic = {
            'amp': 1.,
            'R_sersic': 0.1,
            'n_sersic': 2,
            'center_x': 0,
            'center_y': 0
        }
        # 'SERSIC_ELLIPSE': elliptical Sersic profile
        kwargs_sersic_ellipse = {
            'amp': 1.,
            'R_sersic': .6,
            'n_sersic': 3,
            'center_x': 0,
            'center_y': 0,
            'e1': 0.1,
            'e2': 0.1
        }

        lens_light_model_list = ['SERSIC']
        self.kwargs_lens_light = [kwargs_sersic]
        lens_light_model_class = LightModel(
            light_model_list=lens_light_model_list)
        source_model_list = ['SERSIC_ELLIPSE']
        self.kwargs_source = [kwargs_sersic_ellipse]
        source_model_class = LightModel(light_model_list=source_model_list)

        kwargs_numerics = {
            'supersampling_factor': 1,
            'supersampling_convolution': False,
            'compute_mode': 'regular'
        }
        imageModel = ImageModel(data_class,
                                psf_class,
                                lens_model_class,
                                source_model_class,
                                lens_light_model_class,
                                kwargs_numerics=kwargs_numerics)
        image_sim = sim_util.simulate_simple(imageModel, self.kwargs_lens,
                                             self.kwargs_source,
                                             self.kwargs_lens_light)

        data_class.update_data(image_sim)
        kwargs_data['image_data'] = image_sim
        kwargs_data_joint = {
            'multi_band_list': [[kwargs_data, kwargs_psf, kwargs_numerics]],
            'multi_band_type': 'single-band'
        }
        self.data_class = data_class
        self.psf_class = psf_class

        kwargs_model = {
            'lens_model_list': lens_model_list,
            'source_light_model_list': source_model_list,
            'lens_light_model_list': lens_light_model_list,
            'fixed_magnification_list': [False],
        }
        self.kwargs_numerics = {'subgrid_res': 1, 'psf_subgrid': False}

        kwargs_constraints = {
            'image_plane_source_list': [False] * len(source_model_list)
        }

        kwargs_likelihood = {
            'source_marg': False,
            'image_position_uncertainty': 0.004,
            'check_matched_source_position': False,
            'source_position_tolerance': 0.001,
            'source_position_sigma': 0.001,
        }
        self.param_class = Param(kwargs_model, **kwargs_constraints)
        self.Likelihood = LikelihoodModule(kwargs_data_joint=kwargs_data_joint,
                                           kwargs_model=kwargs_model,
                                           param_class=self.param_class,
                                           **kwargs_likelihood)
        self.sampler = Sampler(likelihoodModule=self.Likelihood)

    def test_pso(self):
        n_particles = 2
        n_iterations = 2
        result, chain = self.sampler.pso(n_particles,
                                         n_iterations,
                                         lower_start=None,
                                         upper_start=None,
                                         threadCount=1,
                                         init_pos=None,
                                         mpi=False,
                                         print_key='PSO')

        assert len(result) == 16

    def test_mcmc_emcee(self):
        n_walkers = 36
        n_run = 2
        n_burn = 2
        mean_start = self.param_class.kwargs2args(
            kwargs_lens=self.kwargs_lens,
            kwargs_source=self.kwargs_source,
            kwargs_lens_light=self.kwargs_lens_light)
        sigma_start = np.ones_like(mean_start) * 0.1
        samples, dist = self.sampler.mcmc_emcee(n_walkers,
                                                n_run,
                                                n_burn,
                                                mean_start,
                                                sigma_start,
                                                mpi=False)
        assert len(samples) == n_walkers * n_run
        assert len(dist) == len(samples)

        # test of backup file
        # 1) run a chain specifiying a backup file name
        backup_filename = 'test_mcmc_emcee.h5'
        samples_1, dist_1 = self.sampler.mcmc_emcee(
            n_walkers,
            n_run,
            n_burn,
            mean_start,
            sigma_start,
            mpi=False,
            backend_filename=backup_filename)
        assert len(samples_1) == n_walkers * n_run
        # 2) run a chain starting from the backup of previous run
        samples_2, dist_2 = self.sampler.mcmc_emcee(
            n_walkers,
            n_run,
            n_burn,
            mean_start,
            sigma_start,
            mpi=False,
            backend_filename=backup_filename,
            start_from_backend=True)
        assert len(samples_2) == len(samples_1) + n_walkers * n_run
        assert len(dist_2) == len(samples_2)

        os.remove(backup_filename)  # just remove the backup file created above

    def test_mcmc_zeus(self):
        n_walkers = 36
        n_run = 2
        n_burn = 2
        mean_start = self.param_class.kwargs2args(
            kwargs_lens=self.kwargs_lens,
            kwargs_source=self.kwargs_source,
            kwargs_lens_light=self.kwargs_lens_light)
        sigma_start = np.ones_like(mean_start) * 0.1
        samples, dist = self.sampler.mcmc_zeus(n_walkers, n_run, n_burn,
                                               mean_start, sigma_start)
        assert len(samples) == n_walkers * n_run
        assert len(dist) == len(samples)

        # test of backup file
        backup_filename = 'test_mcmc_zeus.h5'
        samples_1, dist_1 = self.sampler.mcmc_zeus(
            n_walkers,
            n_run,
            n_burn,
            mean_start,
            sigma_start,
            backend_filename=backup_filename)
        assert len(samples_1) == n_walkers * n_run

        os.remove(backup_filename)