species.analysis.fit_spectrum — species 0.1.2 documentation

# Source code for species.analysis.fit_spectrum

    """
    Module with functionalities for fitting a calibration spectrum.
    """

    import math

    from multiprocessing import Pool, cpu_count

    import emcee
    import numpy as np

    from species.analysis import photometry
    from species.data import database
    from species.read import read_object, read_calibration

    def lnprob(param,
               bounds,
               modelpar,
               objphot,
               specphot,
               bands):
        """
        Internal function for the posterior probability.

        Parameters
        ----------
        param : numpy.ndarray
            Values of the main scaling parameter and optionally additional band-dependent scaling
            parameters.
        bounds : dict
            Boundaries of the main scaling parameter.
        modelpar : list(str, )
            Parameter names.
        objphot : list(tuple(float, float), )
            Photometry of the object.
        specphot : list(float, )
            Synthetic photometry of the calibration spectrum for the same filters as the photometry
            of the object.
        bands : bool
            Use band-dependent scaling parameters in addition to the main scaling parameter which
            is used for the full spectrum.

        Returns
        -------
        float
            Log posterior probability.
        """

        for i, item in enumerate(modelpar):

            if bounds[item][0] <= param[i] <= bounds[item][1]:
                ln_prior = 0.

            else:
                ln_prior = -np.inf
                break

        if math.isinf(ln_prior):
            ln_prob = -np.inf

        else:
            chisq = 0.
            for i, _ in enumerate(objphot):
                if bands:
                    chisq += (objphot[i][0] - param[0]*param[i+1]*specphot[i])**2 / objphot[i][1]**2

                else:
                    chisq += (objphot[i][0] - param[0]*specphot[i])**2 / objphot[i][1]**2

            ln_prob = ln_prior - 0.5*chisq

        return ln_prob

    class FitSpectrum:
        """
        Class for fitting a calibration spectrum to photometric data.
        """

        def __init__(self,
                     object_name,
                     filters,
                     spectrum,
                     bounds):
            """
            Parameters
            ----------
            object_name : str
                Object name in the database.
            filters : tuple(str, )
                Filter IDs for which the photometry is selected. All available photometry of the
                object is selected if set to None.
            spectrum : str
                Calibration spectrum.
            bounds : dict
                Boundaries of the scaling parameter, as {'scaling':(min, max)}.

            Returns
            -------
            None
            """

            self.object = read_object.ReadObject(object_name)

            self.spectrum = spectrum
            self.bounds = bounds

            self.objphot = []
            self.specphot = []

            if filters is None:
                species_db = database.Database()

                objectbox = species_db.get_object(object_name,
                                                  filters=None,
                                                  inc_phot=True,
                                                  inc_spec=False)

                filters = objectbox.filters

            for item in filters:
                readcalib = read_calibration.ReadCalibration(self.spectrum, item)
                calibspec = readcalib.get_spectrum()

                synphot = photometry.SyntheticPhotometry(item)
                spec_phot = synphot.spectrum_to_flux(calibspec.wavelength, calibspec.flux)
                self.specphot.append(spec_phot)

                obj_phot = self.object.get_photometry(item)
                self.objphot.append((obj_phot[2], obj_phot[3]))

            self.modelpar = ['scaling']

        def run_mcmc(self,
                     nwalkers,
                     nsteps,
                     guess,
                     tag,
                     bands=False):
            """
            Function to run the MCMC sampler.

            Parameters
            ----------
            nwalkers : int
                Number of walkers.
            nsteps : int
                Number of steps per walker.
            guess : dict
                Guess of the scaling parameter.
            tag : str
                Database tag where the MCMC samples are stored.
            bands : bool
                Use band-dependent scaling parameters in addition to the main scaling parameter which
                is used for the full spectrum.

            Returns
            -------
            None
            """

            print('Running MCMC...')

            if bands:
                ndim = 1 + len(self.objphot)

            else:
                ndim = 1

            initial = np.zeros((nwalkers, ndim))
            initial[:, 0] = guess['scaling'] + np.random.normal(0, 1e-1*guess['scaling'], nwalkers)

            if ndim > 1:
                for i in range(1, ndim):
                    initial[:, i] = 1. + np.random.normal(0, 0.1, nwalkers)
                    self.modelpar.append('scaling'+str(i))
                    self.bounds['scaling'+str(i)] = (0., 1e2)

            with Pool(processes=cpu_count()) as pool:
                sampler = emcee.EnsembleSampler(nwalkers,
                                                ndim,
                                                lnprob,
                                                args=([self.bounds,
                                                       self.modelpar,
                                                       self.objphot,
                                                       self.specphot,
                                                       bands]))

                sampler.run_mcmc(initial, nsteps, progress=True)

            species_db = database.Database()

            species_db.add_samples(sampler=sampler,
                                   spectrum=('calibration', self.spectrum),
                                   tag=tag,
                                   modelpar=self.modelpar,
                                   distance=None)