def gp_fit(t, y, yerr, t_grid, integrated=False, exp_time=60.):
# optimize kernel hyperparameters and return fit + predictions
with pm.Model() as model:
logS0 = pm.Normal("logS0", mu=0.4, sd=5.0, testval=np.log(np.var(y)))
logw0 = pm.Normal("logw0", mu=-3.9, sd=0.1)
logQ = pm.Normal("logQ", mu=3.5, sd=5.0)
# Set up the kernel and GP
kernel = terms.SHOTerm(log_S0=logS0, log_w0=logw0, log_Q=logQ)
if integrated:
kernel_int = terms.IntegratedTerm(kernel, exp_time)
gp = GP(kernel_int, t, yerr**2)
else:
gp = GP(kernel, t, yerr**2)
# Add a custom "potential" (log probability function) with the GP likelihood
pm.Potential("gp", gp.log_likelihood(y))
with model:
map_soln = xo.optimize(start=model.test_point)
mu, var = xo.eval_in_model(gp.predict(t_grid, return_var=True),
map_soln)
sd = np.sqrt(var)
y_pred = xo.eval_in_model(gp.predict(t), map_soln)
return map_soln, mu, sd, y_pred
def test_integrated(kernel, seed=1234):
np.random.seed(seed)
x = np.sort(np.random.uniform(0, 100, 100))
dt = 0.4 * np.min(np.diff(x))
y = np.sin(x)
yerr = np.random.uniform(0.1, 0.5, len(x))
diag = yerr**2
_check_model(kernel, x, diag, y)
kernel = terms.IntegratedTerm(kernel, dt)
_check_model(kernel, x, diag, y)
def test_integrated_diag(seed=1234):
np.random.seed(seed)
x = np.sort(np.random.uniform(0, 100, 100))
dt = 0.4 * np.min(np.diff(x))
yerr = np.random.uniform(0.1, 0.5, len(x))
diag = yerr**2
kernel = terms.SHOTerm(log_S0=0.1, log_Q=1.0, log_w0=0.5)
kernel += terms.RealTerm(log_a=0.1, log_c=0.4)
a = kernel.get_celerite_matrices(x, diag)[0].eval()
k0 = kernel.value(tt.zeros(1)).eval()
assert np.allclose(a, k0 + diag)
kernel = terms.IntegratedTerm(kernel, dt)
a = kernel.get_celerite_matrices(x, diag)[0].eval()
k0 = kernel.value(tt.zeros(1)).eval()
assert np.allclose(a, k0 + diag)
def gp_predict(t,
y,
yerr,
t_grid,
logS0=0.4,
logw0=-3.9,
logQ=3.5,
integrated=False,
exp_time=60.):
# take kernel hyperparameters as fixed inputs, train + predict
with pm.Model() as model:
kernel = terms.SHOTerm(log_S0=logS0, log_w0=logw0, log_Q=logQ)
if integrated:
kernel_int = terms.IntegratedTerm(kernel, exp_time)
gp = GP(kernel_int, t, yerr**2)
else:
gp = GP(kernel, t, yerr**2)
gp.condition(y)
mu, var = xo.eval_in_model(gp.predict(t_grid, return_var=True))
sd = np.sqrt(var)
y_pred = xo.eval_in_model(gp.predict(t))
return y_pred, mu, sd
def multi_gp_predict(t, y, yerr, t_grid, integrated=False, exp_time=60.):
# this code is GARBAGE. but in principle does gp_predict() for a full comb of modes.
a_max = 0.55 # amplitude of central mode in m/s
nu_max = 3.1e-3 # peak frequency in Hz
c_env = 0.331e-3 # envelope width in Hz
delta_nu = 0.00013 # Hz
gamma = 1. / (2 * 24. * 60. * 60.) # s^-1 ; 2-day damping timescale
freq_grid = np.arange(nu_max - 0.001, nu_max + 0.001,
delta_nu) # magic numbers
amp_grid = a_max**2 * np.exp(-(freq_grid - nu_max)**2 /
(2. * c_env**2)) # amplitudes in m/s
driving_amp_grid = np.sqrt(amp_grid * gamma * dt)
log_S0_grid = [
np.log(d**2 / (dt * o)) for o, d in zip(omega_grid, driving_amp_grid)
]
with pm.Model() as model:
kernel = None
for o, lS in zip(omega_grid, log_S0_grid):
if kernel is None:
kernel = terms.SHOTerm(log_S0=lS,
log_w0=np.log(o),
log_Q=np.log(o / gamma))
else:
kernel += terms.SHOTerm(log_S0=lS,
log_w0=np.log(o),
log_Q=np.log(o / gamma))
if integrated:
kernel_int = terms.IntegratedTerm(kernel, exp_time)
gp = GP(kernel_int, t, yerr**2)
else:
gp = GP(kernel, t, yerr**2)
gp.condition(y)
mu, var = xo.eval_in_model(gp.predict(t_grid, return_var=True))
sd = np.sqrt(var)
y_pred = xo.eval_in_model(gp.predict(t))
return y_pred, mu, sd