class PhaseLPF(BaseLPF):
    def __init__(self, name: str):
        times, fluxes, _, _, _, _ = read_tess(tess_file,
                                              zero_epoch,
                                              period,
                                              use_pdc=True,
                                              baseline_duration_d=period)
        phase, time, flux = [], [], []
        for t, f in zip(times, fluxes):
            ph = (fold(t, period.n, zero_epoch.n, 0.5) - 0.5) * period.n
            mask = abs(ph) > 0.03
            phase.append(ph[mask])
            time.append(t[mask])
            flux.append(f[mask])
        super().__init__(name, ['TESS'],
                         time,
                         flux,
                         wnids=arange(len(flux)),
                         lnlikelihood='celerite')

    def _post_initialisation(self):
        self.t0 = 2458491.8771169
        self.period = 1.2652328
        self.k2 = 0.3**2
        self.a = 10.
        self.phase = fold(self.timea, self.period, self.t0)
        self.em = UniformModel()
        self.em.set_data(self.timea)

        self._mec = self.em.evaluate(sqrt(
            self.k2), self.t0 + 0.5 * self.period, self.period, self.a,
                                     0.5 * pi) - 1
        self._mec = 1 + self._mec / self._mec.ptp()

        phi = 2 * pi * self.phase
        self.alpha = a = abs(phi - pi)
        self._rff = (sin(a) +
                     (pi - a) * cos(a)) / pi * self._mec  # Reflected light
        self._dbf = sin(phi)  # Doppler boosting
        self._evf = -cos(2 * phi)  # Ellipsoidal variations

        self._cwl = 1e-9 * tess.wl
        self._ctm = tess.tm

        self.set_prior('ms', 'NP', star_m.n, star_m.s)
        self.set_prior('teffh', 'NP', 3300, 100)
        self.set_prior('teffc', 'UP', 100, 3300)
        self.set_prior('gp_log10_wn', 'NP', -1.74, 0.15)

    def _init_parameters(self):
        self.ps = ParameterSet()
        ppc = [
            GParameter('ab', 'Bond albedo', '', UP(0, 1), (0, 1)),
            GParameter('mp', 'log10 planet mass', 'MJup',
                       UP(log10(0.1), log10(300)), (0, inf)),
            GParameter('ms', 'Star mass', 'MSun', NP(1.0, 0.1), (0, inf)),
            GParameter('teffh', 'Host effective temperature', '',
                       UP(2000, 10000), (0, inf)),
            GParameter('teffc', 'Companion effective temperature', '',
                       UP(500, 10000), (0, inf)),
            GParameter('bl', 'Baseline level', '', NP(1, 0.005), (0, inf))
        ]
        self.ps.add_global_block('phase_curve', ppc)
        self.ps.freeze()

    def baseline(self, pv):
        pv = atleast_2d(pv)
        return pv[:, 5:6]

    def emitted_flux_ratio(self, pv):
        pv = atleast_2d(pv)
        return summed_planck(pv[:, 4:5], self._cwl, self._ctm) / summed_planck(
            pv[:, 3:4], self._cwl, self._ctm)

    def emitted_light(self, pv):
        pv = atleast_2d(pv)
        return squeeze(
            self.emitted_flux_ratio(pv)[:, newaxis] * self.k2 * self._mec)

    def reflected_light(self, pv):
        pv = atleast_2d(pv)
        return squeeze(reflected_fr(self.a, pv[:, 0:1]) * self.k2 * self._rff)

    def boosting(self, pv):
        pv = atleast_2d(pv)
        a = boosting_amplitude(10**pv[:, 1:2],
                               pv[:, 2:3],
                               self.period,
                               alpha=8.5)
        return squeeze(a * self._dbf)

    def ellipsoidal_variation(self, pv):
        pv = atleast_2d(pv)
        a = ev_amplitude(10**pv[:, 1:2], pv[:, 2:3], 10, 0.55, 0.3)
        return squeeze(a * self._evf)

    def phase_model(self, pv):
        pv = atleast_2d(pv)
        return squeeze(
            self.ellipsoidal_variation(pv) + self.boosting(pv) +
            self.reflected_light(pv) + self.emitted_light(pv))

    def flux_model(self, pv):
        pv = atleast_2d(pv)
        return squeeze(self.baseline(pv) + self.phase_model(pv))

    def create_pv_population(self, npop: int = 50):
        return self.ps.sample_from_prior(npop)
Example #2
0
class LogPosteriorFunction:
    _lpf_name = 'LogPosteriorFunction'

    def __init__(self, name: str, result_dir: Union[Path, str] = '.'):
        """The Log Posterior Function class.

        Parameters
        ----------
        name: str
            Name of the log posterior function instance.
        """
        self.name = name
        self.result_dir = Path(result_dir if result_dir is not None else '.')

        # Declare high-level objects
        # --------------------------
        self.ps = None  # Parametrisation
        self.de = None  # Differential evolution optimiser
        self.sampler = None  # MCMC sampler
        self._local_minimization = None

        # Initialise the additional lnprior list
        # --------------------------------------
        self._additional_log_priors = []

        self._old_de_fitness = None
        self._old_de_population = None

    def print_parameters(self, columns: int = 2):
        columns = max(1, columns)
        for i, p in enumerate(self.ps):
            print(p.__repr__(), end=('\n' if i % columns == columns - 1 else '\t'))

    def _init_parameters(self):
        self.ps = ParameterSet()
        self.ps.freeze()

    def create_pv_population(self, npop=50):
        return self.ps.sample_from_prior(npop)

    def set_prior(self, parameter, prior, *nargs) -> None:
        if isinstance(parameter, str):
            descriptions = self.ps.descriptions
            names = self.ps.names
            if parameter in descriptions:
                parameter = descriptions.index(parameter)
            elif parameter in names:
                parameter = names.index(parameter)
            else:
                params = ', '.join([f"{ln} ({sn})" for ln, sn in zip(self.ps.descriptions, self.ps.names)])
                raise ValueError(f'Parameter "{parameter}" not found from the parameter set: {params}')

        if isinstance(prior, str):
            if prior.lower() in ['n', 'np', 'normal']:
                prior = NP(nargs[0], nargs[1])
            elif prior.lower() in ['u', 'up', 'uniform']:
                prior = UP(nargs[0], nargs[1])
            else:
                raise ValueError(f'Unknown prior "{prior}". Allowed values are (N)ormal and (U)niform.')

        self.ps[parameter].prior = prior

    def lnprior(self, pv: ndarray) -> Union[Iterable, float]:
        """Log prior density for a 1D or 2D array of model parameters.

        Parameters
        ----------
        pv: ndarray
            Either a 1D parameter vector or a 2D parameter array.

        Returns
        -------
            Log prior density for the given parameter vector(s).
        """
        return self.ps.lnprior(pv) + self.additional_priors(pv)

    def additional_priors(self, pv):
        pv = atleast_2d(pv)
        return sum([f(pv) for f in self._additional_log_priors], 0)

    def lnlikelihood(self, pv):
        raise NotImplementedError

    def lnposterior(self, pv):
        lnp = self.lnprior(pv) + self.lnlikelihood(pv)
        return where(isfinite(lnp), lnp, -inf)

    def __call__(self, pv):
        return self.lnposterior(pv)

    def optimize_local(self, pv0=None, method='powell'):
        if pv0 is None:
            if self.de is not None:
                pv0 = self.de.minimum_location
            else:
                pv0 = self.ps.mean_pv
        res = minimize(lambda pv: -self.lnposterior(pv), pv0, method=method)
        self._local_minimization = res

    def optimize_global(self, niter=200, npop=50, population=None, pool=None, lnpost=None, vectorize=True,
                        label='Global optimisation', leave=False, plot_convergence: bool = True, use_tqdm: bool = True,
                        plot_parameters: tuple = (0, 2, 3, 4)):

        lnpost = lnpost or self.lnposterior
        if self.de is None:
            self.de = DiffEvol(lnpost, clip(self.ps.bounds, -1, 1), npop, maximize=True, vectorize=vectorize, pool=pool)
            if population is None:
                self.de._population[:, :] = self.create_pv_population(npop)
            else:
                self.de._population[:, :] = population
        for _ in tqdm(self.de(niter), total=niter, desc=label, leave=leave, disable=(not use_tqdm)):
            pass

        if plot_convergence:
            fig, axs = subplots(1, 1 + len(plot_parameters), figsize=(13, 2), constrained_layout=True)
            rfit = self.de._fitness
            mfit = isfinite(rfit)

            if self._old_de_fitness is not None:
                m = isfinite(self._old_de_fitness)
                axs[0].hist(-self._old_de_fitness[m], facecolor='midnightblue', bins=25, alpha=0.25)
            axs[0].hist(-rfit[mfit], facecolor='midnightblue', bins=25)

            for i, ax in zip(plot_parameters, axs[1:]):
                if self._old_de_fitness is not None:
                    m = isfinite(self._old_de_fitness)
                    ax.plot(self._old_de_population[m, i], -self._old_de_fitness[m], 'kx', alpha=0.25)
                ax.plot(self.de.population[mfit, i], -rfit[mfit], 'k.')
                ax.set_xlabel(self.ps.descriptions[i])
            setp(axs, yticks=[])
            setp(axs[1], ylabel='Log posterior')
            setp(axs[0], xlabel='Log posterior')
            sb.despine(fig, offset=5)
        self._old_de_population = self.de.population.copy()
        self._old_de_fitness = self.de._fitness.copy()

    def sample_mcmc(self, niter: int = 500, thin: int = 5, repeats: int = 1, npop: int = None, population=None,
                    label='MCMC sampling', reset=True, leave=True, save=False, use_tqdm: bool = True, pool=None,
                    lnpost=None, vectorize: bool = True):

        if save and self.result_dir is None:
            raise ValueError('The MCMC sampler is set to save the results, but the result directory is not set.')

        lnpost = lnpost or self.lnposterior
        if self.sampler is None:
            if population is not None:
                pop0 = population
            elif hasattr(self, '_local_minimization') and self._local_minimization is not None:
                pop0 = multivariate_normal(self._local_minimization.x, diag(full(len(self.ps), 0.001 ** 2)), size=npop)
            elif self.de is not None:
                pop0 = self.de.population.copy()
            else:
                raise ValueError('Sample MCMC needs an initial population.')
            self.sampler = EnsembleSampler(pop0.shape[0], pop0.shape[1], lnpost, vectorize=vectorize, pool=pool)
        else:
            pop0 = self.sampler.chain[:, -1, :].copy()

        for i in tqdm(range(repeats), desc=label, disable=(not use_tqdm), leave=leave):
            if reset or i > 0:
                self.sampler.reset()
            for _ in tqdm(self.sampler.sample(pop0, iterations=niter, thin=thin), total=niter,
                          desc='Run {:d}/{:d}'.format(i + 1, repeats), leave=False, disable=(not use_tqdm)):
                pass
            if save:
                self.save(self.result_dir)
            pop0 = self.sampler.chain[:, -1, :].copy()

    def posterior_samples(self, burn: int = 0, thin: int = 1):
        fc = self.sampler.chain[:, burn::thin, :].reshape([-1, len(self.ps)])
        df = pd.DataFrame(fc, columns=self.ps.names)
        return df

    def plot_mcmc_chains(self, pid: int = 0, alpha: float = 0.1, thin: int = 1, ax=None):
        fig, ax = (None, ax) if ax is not None else subplots()
        ax.plot(self.sampler.chain[:, ::thin, pid].T, 'k', alpha=alpha)
        fig.tight_layout()
        return fig

    def save(self, save_path: Path = '.'):
        save_path = Path(save_path)
        npar = len(self.ps)

        if self.de:
            de = xa.DataArray(self.de.population, dims='pvector parameter'.split(), coords={'parameter': self.ps.names})
        else:
            de = None

        if self.sampler is not None:
            mc = xa.DataArray(self.sampler.chain, dims='pvector step parameter'.split(),
                              coords={'parameter': self.ps.names}, attrs={'ndim': npar, 'npop': self.sampler.nwalkers})
        else:
            mc = None

        ds = xa.Dataset(data_vars={'de_population': de, 'mcmc_samples': mc},
                        attrs={'created': strftime('%Y-%m-%d %H:%M:%S'), 'name': self.name})
        ds.to_netcdf(save_path.joinpath(f'{self.name}.nc'))

        try:
            if self.sampler is not None:
                fname = save_path / f'{self.name}.fits'
                chains = self.sampler.chain
                nchains = chains.shape[0]
                nsteps = chains.shape[1]
                idch = repeat(arange(nchains), nsteps)
                idst = tile(arange(nsteps), nchains)
                flc = chains.reshape([-1, chains.shape[2]])
                tb1 = Table([idch, idst], names=['chain', 'step'])
                tb1.add_columns(flc.T, names=self.ps.names)
                tb2 = Table([idch, idst], names=['chain', 'step'])
                tb2.add_column(self.sampler.lnprobability.ravel(), name='lnp')
                tbhdu1 = pf.BinTableHDU(tb1, name='posterior')
                tbhdu2 = pf.BinTableHDU(tb2, name='sample_stats')
                hdul = pf.HDUList([pf.PrimaryHDU(), tbhdu1, tbhdu2])
                hdul.writeto(fname, overwrite=True)
        except ValueError:
            print('Could not save the samples in fits format.')

    def __repr__(self):
        return f"Target: {self.name}\nLPF: {self._lpf_name}"