def get_celerite_matrices(self, x, diag, **kwargs):
x = tt.as_tensor_variable(x)
diag = tt.as_tensor_variable(diag)
ar, cr, ac, bc, cc, dc = self.coefficients
a = diag + tt.sum(ar) + tt.sum(ac)
arg = dc[None, :] * x[:, None]
cos = tt.cos(arg)
sin = tt.sin(arg)
z = tt.zeros_like(x)
U = tt.concatenate(
(
ar[None, :] + z[:, None],
ac[None, :] * cos + bc[None, :] * sin,
ac[None, :] * sin - bc[None, :] * cos,
),
axis=1,
)
V = tt.concatenate(
(tt.ones_like(ar)[None, :] + z[:, None], cos, sin),
axis=1,
)
c = tt.concatenate((cr, cc, cc))
return c, a, U, V
def _get_position_and_velocity(self, t, parallax=None):
sinf, cosf = self._get_true_anomaly(t)
if self.ecc is None:
r = 1.0
vx, vy, vz = self._rotate_vector(-self.K0 * sinf, self.K0 * cosf)
else:
r = (1.0 - self.ecc**2) / (1 + self.ecc * cosf)
vx, vy, vz = self._rotate_vector(-self.K0 * sinf,
self.K0 * (cosf + self.ecc))
x, y, z = self._rotate_vector(r * cosf, r * sinf)
pos = tt.stack((x, y, z), axis=-1)
pos = tt.concatenate(
(
tt.sum(tt.shape_padright(self.a_star) * pos,
axis=0,
keepdims=True),
tt.shape_padright(self.a_planet) * pos,
),
axis=0,
)
vel = tt.stack((vx, vy, vz), axis=-1)
vel = tt.concatenate(
(
tt.sum(
tt.shape_padright(self.m_planet) * vel,
axis=0,
keepdims=True,
),
-tt.shape_padright(self.m_star) * vel,
),
axis=0,
)
if parallax is not None:
# convert r into arcseconds
pos = pos * (parallax * au_per_R_sun)
vel = vel * (parallax * au_per_R_sun)
return pos, vel
def __init__(self, *args, **kwargs):
ttvs = kwargs.pop("ttvs", None)
transit_times = kwargs.pop("transit_times", None)
transit_inds = kwargs.pop("transit_inds", None)
if ttvs is None and transit_times is None:
raise ValueError("one of 'ttvs' or 'transit_times' must be "
"defined")
if ttvs is not None:
self.ttvs = [as_tensor_variable(ttv, ndim=1) for ttv in ttvs]
if transit_inds is None:
self.transit_inds = [
tt.arange(ttv.shape[0]) for ttv in self.ttvs
]
else:
self.transit_inds = [
tt.cast(as_tensor_variable(inds, ndim=1), "int64")
for inds in transit_inds
]
else:
# If transit times are given, compute the least squares period and
# t0 based on these times.
self.transit_times = []
self.ttvs = []
self.transit_inds = []
period = []
t0 = []
for i, times in enumerate(transit_times):
times = as_tensor_variable(times, ndim=1)
if transit_inds is None:
inds = tt.arange(times.shape[0])
else:
inds = tt.cast(as_tensor_variable(transit_inds[i]),
"int64")
self.transit_inds.append(inds)
# A convoluted version of linear regression; don't ask
N = times.shape[0]
sumx = tt.sum(inds)
sumx2 = tt.sum(inds**2)
sumy = tt.sum(times)
sumxy = tt.sum(inds * times)
denom = N * sumx2 - sumx**2
slope = (N * sumxy - sumx * sumy) / denom
intercept = (sumx2 * sumy - sumx * sumxy) / denom
expect = intercept + inds * slope
period.append(slope)
t0.append(intercept)
self.ttvs.append(times - expect)
self.transit_times.append(times)
kwargs["t0"] = tt.stack(t0)
# We'll have two different periods: one that is the mean difference
# between transit times and one that is a parameter that sets the
# transit shape. If a "period" parameter is not given, these will
# be the same. Users will probably want to put a prior relating the
# two periods if they use separate values.
self.ttv_period = tt.stack(period)
if "period" not in kwargs:
if "delta_log_period" in kwargs:
kwargs["period"] = tt.exp(
tt.log(self.ttv_period) +
kwargs.pop("delta_log_period"))
else:
kwargs["period"] = self.ttv_period
super(TTVOrbit, self).__init__(*args, **kwargs)
if ttvs is not None:
self.ttv_period = self.period
self.transit_times = [
self.t0[i] + self.period[i] * self.transit_inds[i] + ttv
for i, ttv in enumerate(self.ttvs)
]
# Compute the full set of transit times
self.all_transit_times = []
for i, inds in enumerate(self.transit_inds):
expect = self.t0[i] + self.period[i] * tt.arange(inds.max() + 1)
self.all_transit_times.append(
tt.set_subtensor(expect[inds], self.transit_times[i]))
# Set up a histogram for identifying the transit offsets
self._bin_edges = [
tt.concatenate((
[tts[0] - 0.5 * self.ttv_period[i]],
0.5 * (tts[1:] + tts[:-1]),
[tts[-1] + 0.5 * self.ttv_period[i]],
)) for i, tts in enumerate(self.all_transit_times)
]
self._bin_values = [
tt.concatenate(([tts[0]], tts, [tts[-1]]))
for i, tts in enumerate(self.all_transit_times)
]
def get_coefficients(self):
c1 = self.term1.coefficients
c2 = self.term2.coefficients
# First compute real terms
ar = []
cr = []
ar.append(tt.flatten(c1[0][:, None] * c2[0][None, :]))
cr.append(tt.flatten(c1[1][:, None] + c2[1][None, :]))
# Then the complex terms
ac = []
bc = []
cc = []
dc = []
# real * complex
ac.append(tt.flatten(c1[0][:, None] * c2[2][None, :]))
bc.append(tt.flatten(c1[0][:, None] * c2[3][None, :]))
cc.append(tt.flatten(c1[1][:, None] + c2[4][None, :]))
dc.append(tt.flatten(tt.zeros_like(c1[1])[:, None] + c2[5][None, :]))
ac.append(tt.flatten(c2[0][:, None] * c1[2][None, :]))
bc.append(tt.flatten(c2[0][:, None] * c1[3][None, :]))
cc.append(tt.flatten(c2[1][:, None] + c1[4][None, :]))
dc.append(tt.flatten(tt.zeros_like(c2[1])[:, None] + c1[5][None, :]))
# complex * complex
aj, bj, cj, dj = c1[2:]
ak, bk, ck, dk = c2[2:]
ac.append(
tt.flatten(
0.5 * (aj[:, None] * ak[None, :] + bj[:, None] * bk[None, :])
)
)
bc.append(
tt.flatten(
0.5 * (bj[:, None] * ak[None, :] - aj[:, None] * bk[None, :])
)
)
cc.append(tt.flatten(cj[:, None] + ck[None, :]))
dc.append(tt.flatten(dj[:, None] - dk[None, :]))
ac.append(
tt.flatten(
0.5 * (aj[:, None] * ak[None, :] - bj[:, None] * bk[None, :])
)
)
bc.append(
tt.flatten(
0.5 * (bj[:, None] * ak[None, :] + aj[:, None] * bk[None, :])
)
)
cc.append(tt.flatten(cj[:, None] + ck[None, :]))
dc.append(tt.flatten(dj[:, None] + dk[None, :]))
return [
tt.concatenate(vals, axis=0)
if len(vals)
else tt.zeros(0, dtype=self.dtype)
for vals in (ar, cr, ac, bc, cc, dc)
]
def get_coefficients(self):
coeffs = (t.coefficients for t in self.terms)
return tuple(tt.concatenate(a, axis=0) for a in zip(*coeffs))