In [ ]:
import os
from os import path
from astropy.time import Time
from astropy.io import fits, ascii
import astropy.units as u
from astropy.table import Table
from astropy.constants import G
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
import h5py
from thejoker import JokerSamples, JokerParams, TheJoker
from thejoker.sampler.mcmc import TheJokerMCMCModel
from thejoker.plot import plot_rv_curves
from twoface.config import TWOFACE_CACHE_PATH
from twoface.db import (db_connect, AllStar, AllVisit, AllVisitToAllStar,
StarResult, Status, JokerRun, initialize_db)
from twoface.data import APOGEERVData
from twoface.plot import plot_two_panel, plot_mcmc_diagnostic, plot_phase_fold
from twoface.mass import stellar_radius, get_m2_min
In [ ]:
plot_path = path.abspath('../plots/mcmc-unimodal')
os.makedirs(plot_path, exist_ok=True)
In [ ]:
samples_file = path.join(TWOFACE_CACHE_PATH, 'apogee-jitter.hdf5')
mcmc_samples_file = path.join(TWOFACE_CACHE_PATH, 'apogee-jitter-mcmc.hdf5')
In [ ]:
Session, _ = db_connect(path.join(TWOFACE_CACHE_PATH, 'apogee.sqlite'))
session = Session()
In [ ]:
run = session.query(JokerRun).filter(JokerRun.name == 'apogee-jitter').one()
params = run.get_joker_params()
In [ ]:
apogee_ids = []
with h5py.File(mcmc_samples_file) as f:
for i, apogee_id in enumerate(f.keys()):
if i == 0:
R = np.zeros((len(f.keys()), len(f[apogee_id]['chain-stats/gelman_rubin'])))
R[i] = f[apogee_id]['chain-stats/gelman_rubin']
apogee_ids.append(apogee_id)
apogee_ids = np.array(apogee_ids)
In [ ]:
_, bins, _ = plt.hist(R.max(axis=1), bins=np.linspace(1, 5, 64));
_, bins, _ = plt.hist(np.median(R, axis=1), bins=bins);
plt.axvline(1.1)
In [ ]:
idx = np.mean(R, axis=1) < 1.1
In [ ]:
with h5py.File(mcmc_samples_file) as f:
for key in apogee_ids[idx]:
samples = JokerSamples.from_hdf5(f[key])
data = AllStar.get_apogee_id(session, key).apogeervdata()
samples.t0 = data.t0
fig = plot_two_panel(data, samples, plot_data_orbits_kw=dict(highlight_P_extrema=False))
fig.savefig(path.join(plot_path, '{0}.png'.format(key)), dpi=250)
plt.close(fig)
fig = plot_phase_fold(data, samples[np.random.randint(len(samples))])
fig.savefig(path.join(plot_path, '{0}-phase.png'.format(key)), dpi=250)
plt.close(fig)
In [ ]:
Ps = []
Perrs = []
eccs = []
eccerrs = []
loggs = []
Ms = []
with h5py.File(mcmc_samples_file) as f:
for key in apogee_ids:
Ps.append(np.mean(f[key]['P']))
eccs.append(np.mean(f[key]['e']))
Perrs.append(np.std(f[key]['P']))
eccerrs.append(np.std(f[key]['e']))
star = AllStar.get_apogee_id(session, key)
loggs.append(star.logg)
Ms.append(star.martig_mass)
full_tbl = Table({
'P': Ps, 'P_err': Perrs,
'e': eccs, 'e_err': eccerrs,
'logg': loggs, 'M1': Ms
})
tbl = full_tbl[idx]
In [ ]:
_t1 = tbl[ (tbl['logg'] > 2.)]
fig, ax = plt.subplots(1, 1, figsize=(6, 6))
ax.errorbar(_t1['P'], _t1['e'], xerr=_t1['P_err'], yerr=_t1['e_err'],
marker='o', linestyle='none')
ax.set_xscale('log')
ax.set_xlim(1, 2000)
ax.set_ylim(-0.02, 1.02);
ax.set_xlabel('period, $P$ [day]')
ax.set_ylabel('eccentricity, $e$')
fig.tight_layout()
In [ ]:
xs = []
xerrs = []
ys = []
yerrs = []
with h5py.File(mcmc_samples_file) as f:
for key in apogee_ids[idx][(tbl['logg'] > 2.) & (tbl['M1'] > 0)]:
star = AllStar.get_apogee_id(session, key)
M = star.martig_mass*u.Msun
R = stellar_radius(star, M)
xs.append(np.mean(f[key]['P']) / np.sqrt(R**3 / (G*M)).to(u.day).value)
ys.append(np.mean(f[key]['e']))
xerrs.append(np.std(f[key]['P']) / np.sqrt(R**3 / (G*M)).to(u.day).value)
yerrs.append(np.std(f[key]['e']))
In [ ]:
fig, ax = plt.subplots(1, 1, figsize=(6, 6))
ax.errorbar(xs, ys, xerr=xerrs, yerr=yerrs,
marker='o', linestyle='none')
ax.set_xscale('log')
ax.set_xlim(1, 10000)
ax.set_ylim(-0.02, 1.02);
ax.set_xlabel(r'$P / P_\odot$')
ax.set_ylabel('eccentricity, $e$')
fig.tight_layout()
In [ ]:
fig, ax = plt.subplots(1, 1, figsize=(6, 6))
# ax.errorbar(full_tbl['P'], full_tbl['logg'],
# marker='o', linestyle='none')
ax.errorbar(tbl['P'], tbl['logg'], xerr=tbl['P_err'],
marker='o', linestyle='none')
ax.set_xscale('log')
ax.set_xlim(1, 2000)
ax.set_ylim(4, 0);
ax.set_xlabel('period, $P$ [day]')
# ax.set_ylabel('eccentricity, $e$')
fig.tight_layout()
In [ ]:
In [ ]:
import pickle
with open('../cache/emcee/model.pickle', 'rb') as f:
model = pickle.load(f)
In [ ]:
star = AllStar.get_apogee_id(session, '2M03002442+4423109')
data = star.apogeervdata()
In [ ]:
with h5py.File(samples_file, 'r') as f:
samples = JokerSamples.from_hdf5(f[star.apogee_id])
In [ ]:
_ = plot_data_orbits(data, samples)
In [ ]:
# model = TheJokerMCMCModel(params, data)
params2 = JokerParams(params.P_min, params.P_max, jitter_unit=u.m/u.s, jitter=(8.5, 0.9))
model = TheJokerMCMCModel(params2, data)
sams = samples[0]
pars = model.pack_samples(sams)
print(model.ln_likelihood(pars))
In [ ]:
pars = model.pack_samples(sams)
print(model.ln_posterior(model.to_mcmc_params(pars)))
pars[0] = np.exp(8.)
pars[4] = 1.
print(model.ln_posterior(model.to_mcmc_params(pars)))
In [ ]:
%%time
with schwimmbad.MultiPool() as pool:
# joker = TheJoker(params)
joker = TheJoker(params2)
model, mcmc_samples, sampler = joker.mcmc_sample(data, samples[0],
n_steps=8192, n_burn=0,
return_sampler=True)
In [ ]:
_ = plot_mcmc_diagnostic(sampler.chain)
In [ ]:
derps = model.unpack_samples_mcmc(sampler.chain[:, -1])
derps.t0 = data.t0
In [ ]:
_ = plot_phase_fold(star, derps[0])
In [ ]:
_ = plot_data_orbits(data, derps, xlim_choice='tight', highlight_P_extrema=False)
In [ ]:
# needs mcmc
stars = session.query(AllStar).join(StarResult, JokerRun, Status)\
.filter(JokerRun.name == 'apogee-jitter')\
.filter(Status.id == 2).all()
len(stars)
In [ ]:
star = stars[1]
data = star.apogeervdata()
with h5py.File(samples_file) as f:
samples0 = JokerSamples.from_hdf5(f[star.apogee_id])
_ = plot_data_orbits(data, samples0, xlim_choice='tight')
In [ ]:
%%time
with schwimmbad.MultiPool() as pool:
joker = TheJoker(params, pool=pool)
model, samples, sampler = joker.mcmc_sample(data, samples0, n_steps=32768,
n_walkers=256, n_burn=1024,
return_sampler=True)
In [ ]:
ndim = sampler.chain.shape[-1]
fig, axes = plt.subplots(ndim, 3, figsize=(12, 16))
for k in range(ndim):
for walker in sampler.chain[..., k]:
axes[k, 0].plot(walker, marker='', drawstyle='steps-mid', alpha=0.1)
axes[k, 1].plot(np.median(sampler.chain[..., k], axis=0),
marker='', drawstyle='steps-mid')
# std = np.std(sampler.chain[..., k], axis=0)
std = 1.5 * median_absolute_deviation(sampler.chain[..., k], axis=0)
axes[k, 2].plot(std, marker='', drawstyle='steps-mid')
fig.tight_layout()
In [ ]:
plt.scatter(samples['P'].value, samples['e'].value, alpha=0.5, linewidth=0)
In [ ]:
In [ ]:
_ = plot_data_orbits(data, samples0, xlim_choice='tight')
_ = plot_data_orbits(data, samples, xlim_choice='tight', highlight_P_extrema=False)
In [ ]:
import astropy.units as u
from astropy.stats import median_absolute_deviation
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
import glob
import pickle
from os import path
In [ ]:
names = [r'$\ln P$', r'$\sqrt{K}\,\cos M_0$', r'$\sqrt{K}\,\sin M_0$',
r'$\sqrt{e}\,\cos \omega$', r'$\sqrt{e}\,\sin \omega$',
r'$\ln s^2$', '$v_0$']
In [ ]:
for filename in glob.glob('../scripts/test-mcmc-*.pickle'):
with open(filename, 'rb') as f:
sampler = pickle.load(f)
ndim = sampler.chain.shape[-1]
fig, axes = plt.subplots(ndim, 3, figsize=(12, 16))
for k in range(ndim):
axes[k, 0].set_ylabel(names[k])
axes[k, 0].plot(sampler.chain[..., k].T, marker='', drawstyle='steps-mid',
alpha=0.1, rasterized=True)
axes[k, 1].plot(np.median(sampler.chain[..., k], axis=0),
marker='', drawstyle='steps-mid')
# std = np.std(sampler.chain[..., k], axis=0)
std = 1.5 * median_absolute_deviation(sampler.chain[..., k], axis=0)
axes[k, 2].plot(std, marker='', drawstyle='steps-mid')
axes[0, 0].set_title('walkers')
axes[0, 1].set_title('med(walkers)')
axes[0, 2].set_title('1.5 MAD(walkers)')
fig.tight_layout()
fig.savefig('../scripts/{0}.png'.format(path.splitext(path.basename(filename))[0]), dpi=250)
plt.close('all')
In [ ]:
for filename in glob.glob('../scripts/test-mcmc-*.pickle'):
*_, apogee_id = path.splitext(filename)[0].split('-')
star = session.query(AllStar).filter(AllStar.apogee_id == apogee_id).limit(1).one()
data = star.apogeervdata()
model = TheJokerMCMCModel(joker_params=params, data=data)
with open(filename, 'rb') as f:
sampler = pickle.load(f)
samples = model.unpack_samples_mcmc(sampler.chain[:, -1])
samples.t0 = Time(data._t0_bmjd, format='mjd', scale='tcb')
fig = plot_data_orbits(data, samples, n_orbits=256)
fig.savefig('../scripts/{0}-samples.png'.format(apogee_id), dpi=260)
plt.close('all')
In [ ]:
In [ ]:
In [ ]: