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)
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}"