def plotsample_QP(p, t, y, tsel, ysel):
kern, sig = kern_QP(p)
gp = GP(kern)
yerr = np.ones(len(ysel)) * sig
gp.compute(tsel, yerr)
mu = gp.sample_conditional(ysel, t)
pl.plot(t, mu, color='c', alpha=0.3)
return
class model_transit_lightcurve(object):
per = 0
ldmethod = 'quad'
def __init__(self,
data=None,
param={},
bounds={},
fixed=[],
group=None,
**kwargs):
"""
:param data:
:param param:
:param bounds:
:param group:
"""
self.data = []
self.param, self.param_name, self.param_bounds = np.array(
[]), np.array([]), []
self.Ndata = 0
self.index_param_indiv = []
self.index_bounds_indiv = []
self.index_param_common = []
self.index_bounds_common = []
self.group_indiv = []
self.fixed = []
self.group_common = []
self.detrendvar = {}
self._varind = 0
self.gp = GP(solver=BasicSolver)
self._mcmc_parinit_optimized = False
self.add_data(data, param, bounds, fixed, group, **kwargs)
def add_data(self, data, param, bounds, fixed=[], group=None, **kwargs):
if data:
if type(data) == str:
data = np.loadtxt(data, unpack=True, **kwargs)
if self.Ndata == 0: self.data = data
if self.Ndata == 1: self.data = [self.data]
if self.Ndata > 0: self.data.append(data)
parlen = len(self.param)
bounlen = len(self.param_bounds)
if param:
if isinstance(param, dict) and type(param) != NamedParam:
param = NamedParam(param)
self.param_name = np.append(self.param_name, param.keys())
self.param = np.append(self.param,
[param[key] for key in param])
fixedkeys = [
param[fx] if type(fx) == int else fx for fx in fixed
]
self.param_bounds += [
bounds[key] for key in bounds if key not in fixedkeys
]
self.group_indiv.append(group)
self.fixed += [
(fx if type(fx) == int else list(param).index(fx)) + parlen
for fx in fixed
]
self.index_param_indiv.append(np.arange(len(param)) + parlen)
self.index_bounds_indiv.append(np.arange(len(bounds)) + bounlen)
if self.Ndata == 1: self.detrendvar = [self.detrendvar]
if self.Ndata > 0: self.detrendvar.append({})
self.Ndata += 1
def add_indiv_param(self, param, bounds, fixed=[], dataindex=-1):
if self.Ndata < 1: return
if param:
if isinstance(param, dict) and type(param) != NamedParam:
param = NamedParam(param)
parlen = len(self.param)
bounlen = len(self.param_bounds)
self.param_name = np.append(self.param_name, param.keys())
self.param = np.append(self.param, [param[key] for key in param])
fixedkeys = [param[fx] if type(fx) == int else fx for fx in fixed]
self.param_bounds += [
bounds[key] for key in bounds if key not in fixedkeys
]
self.index_param_indiv[dataindex] = np.append(
self.index_param_indiv[dataindex],
np.arange(len(param)) + parlen)
self.index_bounds_indiv[dataindex] = np.append(
self.index_bounds_indiv[dataindex],
np.arange(len(bounds)) + bounlen)
self.fixed += [parlen + fx for fx in fixed]
def add_common_param(
self,
param,
bounds={},
fixed=[],
group=None): # TODO: ignore bounds keys if fixed is not []
if self.Ndata < 1: return
if param:
if isinstance(param, dict) and type(param) != NamedParam:
param = NamedParam(param)
parlen = len(self.param)
bounlen = len(self.param_bounds)
self.param_name = np.append(self.param_name, param.keys())
self.param = np.append(self.param, [param[key] for key in param])
fixedkeys = [param[fx] if type(fx) == int else fx for fx in fixed]
self.param_bounds += [
bounds[key] for key in bounds if key not in fixedkeys
]
self.group_common.append(group)
self.fixed += [
(fx if type(fx) == int else list(param.keys()).index(fx)) +
parlen for fx in fixed
]
self.index_param_common.append(np.arange(len(param)) + parlen)
self.index_bounds_common.append(np.arange(len(bounds)) + bounlen)
def add_detrend_param(self,
variable=0,
name='',
dataindex=-1,
coeff=[],
bounds=[],
fixed=[]):
if type(variable) == int:
varnames = ['t', 'flux', 'err']
if self.Ndata == 1 and dataindex in [0, -1]:
var = self.data[variable]
elif self.Ndata > 1:
var = self.data[dataindex][variable]
if not name: name = varnames[variable]
else:
var = variable
if not name:
name = 'var' + str(self._varind)
self._varind += 1
name = 'det_' + name
if coeff:
parlen = len(self.param)
bounlen = len(self.param_bounds)
self.param = np.append(self.param, coeff)
self.param_name = np.append(
self.param_name, [name + f'_{i}' for i in range(len(coeff))])
self.param_bounds += bounds
if self.Ndata == 1 and dataindex in [0, -1]:
self.detrendvar.update({name: (var, len(coeff))})
elif self.Ndata > 1:
self.detrendvar[dataindex].update({name: (var, len(coeff))})
self.index_param_indiv[dataindex] = np.append(
self.index_param_indiv[dataindex],
np.arange(len(coeff)) + parlen)
self.index_bounds_indiv[dataindex] = np.append(
self.index_bounds_indiv[dataindex],
np.arange(len(bounds)) + bounlen)
self.fixed += [parlen + fx for fx in fixed]
def get_named_param(self, param):
if type(param) != np.ndarray: param = np.array(param)
if self.Ndata == 1:
nparam = NamedParam(
zip(self.param_name[self.index_param_indiv[0]],
param[self.index_param_indiv[0]]))
for j, icarr in enumerate(self.index_param_common):
if not self.group_indiv[0] or not self.group_common[
j] or self.group_indiv[0] == self.group_common[j]:
nparam.update(zip(self.param_name[icarr], param[icarr]))
return nparam
paramseg = [NamedParam() for _ in self.index_param_indiv]
for i, iiarr in enumerate(self.index_param_indiv):
for j, icarr in enumerate(self.index_param_common):
if not self.group_indiv[i] or not self.group_common[
j] or self.group_indiv[i] == self.group_common[j]:
paramseg[i].update(
zip(self.param_name[icarr], param[icarr]))
paramseg[i].update(zip(self.param_name[iiarr], param[iiarr]))
return paramseg
def _model_function(self, param, t, named=False, detrend=True):
if not named: param = self.get_named_param(param)
if detrend:
return direct_tfm_with_detrend(param, t, self.per, self.detrendvar,
self.ldmethod)
return direct_tfm(param, t, self.per, self.ldmethod)
@staticmethod
def _contains_gppar(param):
if isinstance(param, (dict, NamedParam)):
return 'gpa' in param and 'gptau' in param
return ['gpa' in par and 'gptau' in par for par in param]
@staticmethod
def _contais_detrendvar(detrendvar):
if isinstance(detrendvar, dict):
return len(detrendvar) == 0
return [len(dtv) == 0 for dtv in detrendvar]
def log_likelihood_gp(self, param):
if self.Ndata == 0: return
if not self._calcgp: return
nparam = self.get_named_param(param)
if self.Ndata == 1:
detfac = get_detrend_factor(
nparam, self.detrendvar) if self._calcdetrend else 1
self.gp.kernel = kernel(nparam['gpa'], nparam['gptau'])
if np.any(np.isnan(detfac)): print('detfac nan')
if np.any(np.isinf(detfac)): print('detfac inf')
if np.any(np.isnan(nparam['gpa'])): print('gpa nan')
if np.any(np.isinf(nparam['gpa'])): print('gpa inf')
if np.any(np.isnan(nparam['gptau'])): print('gptau nan')
if np.any(np.isinf(nparam['gptau'])): print('gptau inf')
if np.any(np.isnan(self.data[0])): print('t nan')
if np.any(np.isinf(self.data[0])): print('t inf')
if np.any(np.isnan(self.data[2])): print('err nan')
if np.any(np.isinf(self.data[2])): print('err inf')
try:
self.gp.compute(self.data[0], self.data[2] / detfac)
except:
print(nparam)
print(np.any(detfac == 0))
raise
return self.gp.log_likelihood(
self.data[1] / detfac - self._model_function(
nparam, self.data[0], named=True, detrend=False))
detfac = get_detrend_factor(
nparam, self.detrendvar) if self._calcdetrend else np.ones(
self.Ndata)
llhood = []
for i in range(self.Ndata):
if self._calcgp[i]:
self.gp.kernel = kernel(nparam[i]['gpa'], nparam[i]['gptau'])
self.gp.compute(self.data[i][0], self.data[i][2] / detfac[i])
llhood.append(
self.gp.log_likelihood(
self.data[i][1] / detfac[i] -
self._model_function(nparam[i],
self.data[i][0],
named=True,
detrend=False)))
else:
llhood.append(None)
return llhood
def run_mcmc(self, *args, **kwargs):
data = self.data if self.Ndata == 1 else list(
map(list, zip(*self.data)))
param = kwargs.pop('param_init', self.param)
bounds = kwargs.pop('param_bounds', self.param_bounds)
priorpdf_type = kwargs.pop('priorpdftype', 'uniform')
priorpdf_args = kwargs.pop('priorpdfargs', ())
self._calcgp = self._contains_gppar(self.get_named_param(param))
if not np.any(self._calcgp): self._calcgp = False
self._calcdetrend = self._contais_detrendvar(self.detrendvar)
if not np.any(self._calcdetrend): self._calcdetrend = False
llhood = self.log_likelihood_gp if self._calcgp else None
self.mcmc = MCMC(np.array(data),
param,
bounds,
fixpos=self.fixed,
modelfunc=self._model_function,
loglikelihood=llhood,
ignorex_loglikelihood=True,
priorpdf_type=priorpdf_type,
priorpdf_args=priorpdf_args)
self.mcmc.mcmc_name = kwargs.pop('mcmc_name', '')
self.mcmc.param_names = self.param_name_plus_group
optimize = kwargs.pop('preoptimize', False)
if optimize:
self.mcmc.optimized_initpar()
self._mcmc_parinit_optimized = True
self.mcmc(*args, **kwargs)
self.mcmc_accepted = self.mcmc.accepted
self.mcmc_nwalker = self.mcmc.Nwalker
self.mcmc_niterate = self.mcmc.Niterate
@property
def skeleton(self):
keys = [
'Ndata', 'index_param_indiv', 'index_bounds_indiv',
'index_param_common', 'index_bounds_common', 'group_indiv',
'group_common', 'fixed', 'detrendvar', 'param_name',
'param_bounds', '_varind', 'mcmc_accepted', 'mcmc_nwalker',
'mcmc_niterate'
]
skeleton = dict(input_data=self.data, param_init=self.param)
for key in keys:
try:
skeleton[key] = self.__getattribute__(key)
except:
skeleton[key] = None
skeleton['mcmc_param_init_optimized'] = self.mcmc.param_init if self._mcmc_parinit_optimized else []
return skeleton
@property
def param_name_plus_group(self):
groupexp = np.full(len(self.param), '', dtype=object)
for i in range(len(self.group_indiv)):
if self.group_indiv[i]:
groupexp[self.index_param_indiv[i]] = '-' + self.group_indiv[i]
for i in range(len(self.group_common)):
if self.group_common[i]:
groupexp[self.index_param_common[i]] = self.group_common[i]
return list(map(''.join, zip(self.param_name, groupexp)))
def saveall(self,
retrieval_skeleton='',
params_mcmc='',
mcmc_params_rawsample='',
results_optimized='',
**kwargs):
import pickle as pkl
if retrieval_skeleton:
pkl.dump(self.skeleton, open(retrieval_skeleton + '.pkl', 'wb'))
if params_mcmc:
if 'header' not in kwargs:
kwargs['header'] = '\t'.join(self.param_name_plus_group)
np.savetxt(params_mcmc, self.params_mcmc, **kwargs)
def load_retrieval_skeleton(self, skeleton='', param_init='init'):
import pickle as pkl
if skeleton:
if type(skeleton) == str and os.path.exists(skeleton + '.pkl'):
skeleton = pkl.load(open(skeleton + '.pkl', 'rb'))
self.data = skeleton.pop('input_data')
parami = skeleton.pop('param_init')
parammo = skeleton.pop('mcmc_param_init_optimized')
paramo = skeleton.pop('param_optimized', [])
if param_init == 'init': self.param = parami
if param_init == 'mcmc_preoptimized': self.param = parammo
if param_init == 'mcmc_final':
self.param = self.median_err_params_mcmc[:, 0]
if param_init == 'optimized_final': self.param = paramo
for key in skeleton:
self.__setattr__(key, skeleton[key])
def load_params_mcmc(self, source):
if source:
if type(source) == str and os.path.exists(source):
source = np.loadtxt(source)
self.params_mcmc = source
def get_flatsamples(self, saveto='', **kwargs):
self.params_mcmc = self.mcmc.get_flatsamples(**kwargs)
def load_backend_mcmc(self, source, mcmc_name=''):
if not hasattr(self, 'mcmc'):
self.mcmc = MCMC([], [])
self.mcmc.mcmc_name = mcmc_name
self.mcmc.param_names = self.param_name_plus_group
self.mcmc.load_backend(source)
self.mcmc.Niterate, self.mcmc.Nwalker, self.mcmc.Ndim = self.mcmc.get_samples(
).shape
@staticmethod
def chooseNflatten_from_sample(samples, burn=0, thin=1, accepted=[]):
return samples[burn::thin, :, :].reshape(
-1, samples.shape[2]) if len(accepted) != 0 else samples[
burn::thin, accepted, :].reshape(-1, samples.shape[2])
def get_best_fit(self, **kwargs):
data = self.data if self.Ndata == 1 else list(
map(list, zip(*self.data)))
param = kwargs.pop('param_init', self.param)
bounds = kwargs.pop('param_bounds', self.param_bounds)
self.mcmc = MCMC(np.array(data),
param,
bounds,
fixpos=self.fixed,
modelfunc=self._model_function,
loglikelihood=self.log_likelihood_gp,
ignorex_loglikelihood=True)
def get_transit_model(self,
param,
t,
named=False,
detrend=True,
denoiseGP=True):
if not named:
if np.ndim(param): param = self.get_named_param(param)
elif np.ndim(param) == 2:
param = list(param)
for i in range(len(param)):
param[i] = self.get_named_param(param[i])
modelop = self._model_function(param, t, named=True, detrend=detrend)
if not np.any(self._contains_gppar(param)) or not denoiseGP:
return modelop
if self.Ndata == 1:
self.gp.kernel = kernel(param['gpa'], param['gptau'])
self.gp.compute(self.data[0], self.data[2])
return self.gp.sample_conditional(self.data[1] - modelop,
t) + modelop
for i in range(self.Ndata):
if self._contains_gppar(param[i]):
self.gp.kernel = kernel(param[i]['gpa'], param[i]['gptau'])
self.gp.compute(self.data[i][0], self.data[i][2])
modelop[i] += self.gp.sample_conditional(self.data[i][1] -
modelop[i],
t,
size=100).mean(0)
return modelop
def get_adjusted_data(self,
param,
named=False,
detrend=True,
denoiseGP=True):
if self.Ndata == 0: return
if not named:
if np.ndim(param): param = self.get_named_param(param)
elif np.ndim(param) == 2:
param = list(param)
for i in range(len(param)):
param[i] = self.get_named_param(param[i])
if detrend and self._calcdetrend:
detfac = get_detrend_factor(param, self.detrendvar)
if self.Ndata == 1:
t, f, e = self.data
if detrend and self._calcdetrend:
f /= detfac
e /= detfac
if denoiseGP and self._calcgp:
self.gp.kernel = kernel(param['gpa'], param['gptau'])
self.gp.compute(t, e)
f -= self.gp.sample_conditional(
f -
self._model_function(param, t, named=True, detrend=False),
t,
size=100).mean(0)
return t, f, e
t, f, e = [], [], []
for i in range(self.Ndata):
ti, fi, ei = self.data[i]
if detrend and self._calcdetrend:
fi /= detfac[i]
ei /= detfac[i]
if denoiseGP and self._calcgp:
self.gp.kernel = kernel(param[i]['gpa'], param[i]['gptau'])
self.gp.compute(ti, ei)
fi -= self.gp.sample_conditional(fi - self._model_function(
param[i], ti, named=True, detrend=False),
t,
size=100).mean(0)
t.append(ti)
f.append(fi)
e.append(ei)
return t, f, e
@property
def median_err_params_mcmc(self):
return get_median_error_from_distribution(self.params_mcmc,
sigma=1,
method='percentile',
saveas='')
def save_median_err_params_mcmc(self,
saveto='',
parnames=[],
display=True):
parfinal = self.median_err_params_mcmc.T
print(parfinal.shape)
if not parnames: parnames = self.mcmc.param_names
if os.path.exists(saveto): os.remove(saveto)
for i, parname in enumerate(parnames):
if display:
print(
parname +
f': {parfinal[i, 0]} +{parfinal[i, 1]} -{parfinal[i, 2]}')
if saveto:
print(
parname +
f': {parfinal[i, 0]} +{parfinal[i, 1]} -{parfinal[i, 2]}',
file=open(saveto, 'a'))
def overplot_model_median_err_params_mcmc(self,
multifig=False,
axis='auto',
figsize=(10, 10)):
import matplotlib.pyplot as plt
params = self.median_err_params_mcmc.T
t, f, e = self.get_adjusted_data(params[:, 0])
if self.Ndata == 1:
fig, ax = plt.subplots(figsize=figsize)
midfluxfit = self.get_transit_model(params[:, 0],
t,
detrend=False,
denoiseGP=False)
ax.errorbar(t, f, e, fmt='.')
ax.plot(t,
midfluxfit,
c='r',
lw=3,
label='Model corr. to median of parameters')
return fig, ax
if not multifig:
if axis == 'auto':
figure, axes = plt.subplots(self.Ndata, 1, figsize=figsize)
elif isinstance(axis, (tuple, list)) and len(axis) == 2:
figure, axes = plt.subplots(axis[0], axis[1], figsize=figsize)
axes = np.ravel(axes)
else:
figure, axes = [], []
for i in range(self.Ndata):
if multifig:
fig, ax = plt.subplots(figsize=figsize)
figure.append(fig)
axes.append(ax)
else:
ax = axes[i]
midfluxfit = self.get_transit_model(params[i][:, 0],
t[i],
detrend=False,
denoiseGP=False)
ax.errorbar(t[i], f[i], e[i], fmt='.')
ax.plot(t[i],
midfluxfit,
c='r',
lw=3,
label='Model corr. to median of parameters')
return figure, axes
