Created
July 11, 2019 15:30
-
-
Save dfm/b1633319f9314dfdb2245d2f392cd073 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import pymc3 as pm | |
import theano.tensor as tt | |
Ks=np.array([8.0,2.15]) | |
with pm.Model() as model: | |
# Common "shared" orbital parameters | |
logK = pm.Uniform("logK",lower=np.log(0.5), upper=np.log(100.0), testval=np.log(Ks), shape=2) | |
# Mass | |
# logmpl = pm.Normal("logmpl", mu=np.log(float(msini1.value)*M_star[0],float(msini2.value)*M_star[0]), sd=(5.0,5.0),shape=2) | |
# mpl = pm.Deterministic("mpl", tt.exp(logmpl)) | |
BoundedNormal = pm.Bound(pm.Normal, lower=1.0, upper=2000.0) | |
P = BoundedNormal("P", mu=np.array(periods), sd=np.array(period_errs), shape=2, | |
testval=np.array(periods)) | |
#ecc = pm.Uniform("ecc", lower=0, upper=1, testval=0.1,shape=2) | |
ecc = pm.Beta("ecc", alpha=0.867, beta=3.03, shape=2, testval=np.array([0.01, 0.15])) | |
omega = xo.distributions.Angle("omega",shape=2) | |
t0 = pm.Normal("t0", mu=np.array(t0s), sd=np.array(t0_errs), shape=2) | |
# Track some parameters | |
K = pm.Deterministic("K", pm.math.exp(logK)) | |
# We need to include a linear trend too | |
trend = pm.Normal("trend", mu=0.0, sd=100.0) | |
# Set up the orbit | |
orbit = xo.orbits.KeplerianOrbit(period=P, t0=t0, ecc=ecc, omega=omega) | |
# Loop over datasets | |
logss = [] | |
sys_vels = [] | |
gp_params = [] | |
for t in utel: | |
xt = x[tel== t] | |
yt = y[tel == t] | |
yerrt = yerr[tel == t] | |
# Jitter for this instrument | |
logs = pm.Uniform("logs_{0}".format(t), lower=-15, upper=15, testval=0.0) #<--- CHANGE? | |
logss.append(logs) | |
# Systemic velocity for this instrument | |
sys_vel = pm.Uniform("sys_vel_{0}".format(t),lower=-100,upper=100,testval=0.0) | |
sys_vels.append(sys_vel) | |
# Compute the model RV | |
vrad = orbit.get_radial_velocity(xt, K=K) | |
#pm.Deterministic("vrad", vrad) | |
# And the background model | |
bkg_rv = sys_vel + trend * (xt - x_ref) | |
#bkg = pm.Deterministic("bkg_rv",bkg_rv) | |
# Sum over planets and add the background to get the full model | |
#rv_model = pm.Deterministic("rv_model", tt.sum(vrad, axis=-1) + bkg_rv) | |
# Noise model (error bar + jitter) | |
err = pm.math.sqrt(yerrt**2 + pm.math.exp(2*logs)) | |
mod = tt.sum(vrad, axis=-1) + bkg_rv | |
#GP | |
logS1 = pm.Normal("logS1_{0}".format(t), mu=0.0, sd=15.0, testval=np.log(np.var(yt))) | |
logw1 = pm.Normal("logw1_{0}".format(t), mu=0.0, sd=15.0, testval=np.log(3.0)) | |
logS2 = pm.Normal("logS2_{0}".format(t), mu=0.0, sd=15.0, testval=np.log(np.var(yt))) | |
logw2 = pm.Normal("logw2_{0}".format(t), mu=0.0, sd=15.0, testval=np.log(3.0)) | |
logQ = pm.Normal("logQ_{0}".format(t), mu=0.0, sd=15.0, testval=0) | |
gp_params += [logS1, logw1, logS2, logw2, logQ] | |
# Set up the kernel an GP | |
kernel = xo.gp.terms.SHOTerm(log_S0=logS1, log_w0=logw1, Q=1.0/np.sqrt(2)) | |
kernel += xo.gp.terms.SHOTerm(log_S0=logS2, log_w0=logw2, log_Q=logQ) | |
gp = xo.gp.GP(kernel, xt, err) | |
# Add a custom "potential" (log probability function) with the GP likelihood | |
pm.Potential("rv_obs_{0}".format(t), gp.log_likelihood(yt-mod)) | |
pm.Deterministic("gp_pred_{0}".format(t), gp.predict()) | |
# Do 'pfs2' explicitly so can try to connect pfs1 and pfs2 sys_vel terms | |
# xt = x[tel== 'pfs2'] | |
# yt = y[tel == 'pfs2'] | |
# yerrt = yerr[tel == 'pfs2'] | |
# logs = pm.Uniform("logs_{0}".format('pfs2'), lower=-5.0, upper=5.0, testval=0.0) | |
#logs = pm.Normal("logs_{0}".format('pfs2'), mu=np.log(10.0), sd=10.0) | |
# logss.append(logs) | |
# BoundedNormal = pm.Bound(pm.Normal, lower=sys_vel-2.0, upper=sys_vel+2.0) | |
# sys_vel = BoundedNormal("sys_vel_{0}".format('pfs2'),mu=sys_vel, sd=2.0) | |
# sys_vels.append(sys_vel) | |
# vrad = orbit.get_radial_velocity(xt, K=K) | |
# bkg_rv = sys_vel + trend * (xt - x_ref) | |
# err = pm.math.sqrt(yerrt**2 + pm.math.exp(2*logs)) | |
# pm.Normal("obs_{0}".format('pfs2'), mu=tt.sum(vrad, axis=-1) + bkg_rv, sd=err, observed=yt) | |
pm.Normal("pfs", mu=0, sd=3.0, observed=model.sys_vel_pfs1 - model.sys_vel_pfs2) | |
pm.Deterministic("rv_model", orbit.get_radial_velocity(x, K=K)) | |
# These are for plotting | |
rv_model_pred = [] | |
for i in range(2): | |
x_pred = P[i] * phase_pred + t0[i] | |
rv_model_pred.append(orbit.get_radial_velocity(x_pred, K=K)[:, i]) | |
pm.Deterministic("rv_model_pred", tt.stack(rv_model_pred)) | |
pm.Deterministic("gp_pred", gp.predict()) | |
print("\nOptimizing GP params...") | |
map_soln = xo.optimize(vars=sys_vels + logss) | |
map_soln = xo.optimize(map_soln, vars=[t0]) | |
map_soln = xo.optimize(map_soln, vars=gp_params) | |
map_soln = xo.optimize(map_soln, vars=[t0, P, logK, ecc, omega]) | |
map_soln = xo.optimize(map_soln) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment