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_celerite_matrices(self, x, diag, **kwargs):
dt = self.delta
ar, cr, a, b, c, d = self.term.coefficients
# Real part
cd = cr * dt
delta_diag = 2 * tt.sum(ar * (cd - tt.sinh(cd)) / cd ** 2)
# Complex part
cd = c * dt
dd = d * dt
c2 = c ** 2
d2 = d ** 2
c2pd2 = c2 + d2
C1 = a * (c2 - d2) + 2 * b * c * d
C2 = b * (c2 - d2) - 2 * a * c * d
norm = (dt * c2pd2) ** 2
sinh = tt.sinh(cd)
cosh = tt.cosh(cd)
delta_diag += 2 * tt.sum(
(
C2 * cosh * tt.sin(dd)
- C1 * sinh * tt.cos(dd)
+ (a * c + b * d) * dt * c2pd2
)
/ norm
)
new_diag = diag + delta_diag
return super().get_celerite_matrices(x, new_diag)
def get_value(self, tau):
ar, cr, ac, bc, cc, dc = self.coefficients
tau = tt.abs_(tau)
tau = tt.shape_padright(tau)
K = tt.sum(ar * tt.exp(-cr * tau), axis=-1)
factor = tt.exp(-cc * tau)
K += tt.sum(ac * factor * tt.cos(dc * tau), axis=-1)
K += tt.sum(bc * factor * tt.sin(dc * tau), axis=-1)
return K
def get_psd(self, omega):
ar, cr, ac, bc, cc, dc = self.coefficients
omega = tt.shape_padright(omega)
w2 = omega ** 2
w02 = cc ** 2 + dc ** 2
power = tt.sum(ar * cr / (cr ** 2 + w2), axis=-1)
power += tt.sum(
((ac * cc + bc * dc) * w02 + (ac * cc - bc * dc) * w2)
/ (w2 * w2 + 2.0 * (cc ** 2 - dc ** 2) * w2 + w02 * w02),
axis=-1,
)
return np.sqrt(2.0 / np.pi) * power
def _do_compute(self, quiet):
if quiet:
self._d, self._W, _ = ops.factor_quiet(self._t, self._c, self._a,
self._U, self._V)
self._log_det = tt.switch(tt.any(self._d <= 0.0), -np.inf,
tt.sum(tt.log(self._d)))
else:
self._d, self._W, _ = ops.factor(self._t, self._c, self._a,
self._U, self._V)
self._log_det = tt.sum(tt.log(self._d))
self._norm = -0.5 * (self._log_det + self._size * np.log(2 * np.pi))
def grad(self, inputs, gradients):
b, r = inputs
s, dsdb, dsdr = self(b, r)
bs = gradients[0]
for g in gradients[1:]:
if not isinstance(g.type, aesara.gradient.DisconnectedType):
raise ValueError(
"Backpropagation is only supported for the solution vector"
)
if isinstance(bs.type, aesara.gradient.DisconnectedType):
return [
aesara.gradient.DisconnectedType()(),
aesara.gradient.DisconnectedType()(),
]
return [tt.sum(bs * dsdb, axis=-1), tt.sum(bs * dsdr, axis=-1)]
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 test_velocity():
t_tensor = tt.dvector()
t = np.linspace(0, 100, 1000)
m_planet = 0.1
m_star = 1.3
orbit = KeplerianOrbit(
m_star=m_star,
r_star=1.0,
t0=0.5,
period=100.0,
ecc=0.1,
omega=0.5,
Omega=1.0,
incl=0.25 * np.pi,
m_planet=m_planet,
)
star_pos = orbit.get_star_position(t_tensor)
star_vel = theano.function([], orbit.get_star_velocity(t))()
star_vel_expect = np.empty_like(star_vel)
for i in range(3):
g = theano.grad(tt.sum(star_pos[i]), t_tensor)
star_vel_expect[i] = theano.function([t_tensor], g)(t)
assert np.allclose(star_vel, star_vel_expect)
planet_pos = orbit.get_planet_position(t_tensor)
planet_vel = theano.function([], orbit.get_planet_velocity(t))()
planet_vel_expect = np.empty_like(planet_vel)
for i in range(3):
g = theano.grad(tt.sum(planet_pos[i]), t_tensor)
planet_vel_expect[i] = theano.function([t_tensor], g)(t)
assert np.allclose(planet_vel, planet_vel_expect)
pos = orbit.get_relative_position(t_tensor)
vel = np.array(theano.function([], orbit.get_relative_velocity(t))())
vel_expect = np.empty_like(vel)
for i in range(3):
g = theano.grad(tt.sum(pos[i]), t_tensor)
vel_expect[i] = theano.function([t_tensor], g)(t)
assert np.allclose(vel, vel_expect)
def get_value(self, tau0):
dt = self.delta
ar, cr, a, b, c, d = self.term.coefficients
# Format the lags correctly
tau0 = tt.abs_(tau0)
tau = tau0[..., None]
# Precompute some factors
dpt = dt + tau
dmt = dt - tau
# Real parts:
# tau > Delta
crd = cr * dt
cosh = tt.cosh(crd)
norm = 2 * ar / crd ** 2
K_large = tt.sum(norm * (cosh - 1) * tt.exp(-cr * tau), axis=-1)
# tau < Delta
crdmt = cr * dmt
K_small = K_large + tt.sum(norm * (crdmt - tt.sinh(crdmt)), axis=-1)
# Complex part
cd = c * dt
dd = d * dt
c2 = c ** 2
d2 = d ** 2
c2pd2 = c2 + d2
C1 = a * (c2 - d2) + 2 * b * c * d
C2 = b * (c2 - d2) - 2 * a * c * d
norm = 1.0 / (dt * c2pd2) ** 2
k0 = tt.exp(-c * tau)
cdt = tt.cos(d * tau)
sdt = tt.sin(d * tau)
# For tau > Delta
cos_term = 2 * (tt.cosh(cd) * tt.cos(dd) - 1)
sin_term = 2 * (tt.sinh(cd) * tt.sin(dd))
factor = k0 * norm
K_large += tt.sum(
(C1 * cos_term - C2 * sin_term) * factor * cdt, axis=-1
)
K_large += tt.sum(
(C2 * cos_term + C1 * sin_term) * factor * sdt, axis=-1
)
# tau < Delta
edmt = tt.exp(-c * dmt)
edpt = tt.exp(-c * dpt)
cos_term = (
edmt * tt.cos(d * dmt) + edpt * tt.cos(d * dpt) - 2 * k0 * cdt
)
sin_term = (
edmt * tt.sin(d * dmt) + edpt * tt.sin(d * dpt) - 2 * k0 * sdt
)
K_small += tt.sum(2 * (a * c + b * d) * c2pd2 * dmt * norm, axis=-1)
K_small += tt.sum((C1 * cos_term + C2 * sin_term) * norm, axis=-1)
return tt.switch(tt.le(tau0, dt), K_small, K_large)
def test_approx_transit_depth():
u = np.array([0.3, 0.2])
lc = LimbDarkLightCurve(u[0], u[1])
for b, delta in [
(np.float64(0.5), np.float64(0.01)),
(np.array([0.1, 0.9]), np.array([0.1, 0.5])),
(np.array([0.1, 0.9, 0.3]), np.array([0.1, 0.5, 0.0234])),
]:
dv = tt.as_tensor_variable(delta)
ror, jac = lc.get_ror_from_approx_transit_depth(dv, b, jac=True)
_check_quad(u, b, delta, ror.eval())
assert np.allclose(theano.grad(tt.sum(ror), dv).eval(), jac.eval())
def get_light_curve(
self,
orbit=None,
r=None,
t=None,
texp=None,
oversample=7,
order=0,
use_in_transit=None,
light_delay=False,
):
"""Get the light curve for an orbit at a set of times
Args:
orbit: An object with a ``get_relative_position`` method that
takes a tensor of times and returns a list of Cartesian
coordinates of a set of bodies relative to the central source.
This method should return three tensors (one for each
coordinate dimension) and each tensor should have the shape
``append(t.shape, r.shape)`` or ``append(t.shape, oversample,
r.shape)`` when ``texp`` is given. The first two coordinate
dimensions are treated as being in the plane of the sky and the
third coordinate is the line of sight with positive values
pointing *away* from the observer. For an example, take a look
at :class:`orbits.KeplerianOrbit`.
r (tensor): The radius of the transiting body in the same units as
``r_star``. This should have a shape that is consistent with
the coordinates returned by ``orbit``. In general, this means
that it should probably be a scalar or a vector with one entry
for each body in ``orbit``. Note that this is a different
quantity than the planet-to-star radius ratio; do not confuse
the two!
t (tensor): The times where the light curve should be evaluated.
texp (Optional[tensor]): The exposure time of each observation.
This can be a scalar or a tensor with the same shape as ``t``.
If ``texp`` is provided, ``t`` is assumed to indicate the
timestamp at the *middle* of an exposure of length ``texp``.
oversample (Optional[int]): The number of function evaluations to
use when numerically integrating the exposure time.
order (Optional[int]): The order of the numerical integration
scheme. This must be one of the following: ``0`` for a
centered Riemann sum (equivalent to the "resampling" procedure
suggested by Kipping 2010), ``1`` for the trapezoid rule, or
``2`` for Simpson's rule.
use_in_transit (Optional[bool]): If ``True``, the model will only
be evaluated for the data points expected to be in transit
as computed using the ``in_transit`` method on ``orbit``.
"""
if orbit is None:
raise ValueError("missing required argument 'orbit'")
if r is None:
raise ValueError("missing required argument 'r'")
if t is None:
raise ValueError("missing required argument 't'")
use_in_transit = (not light_delay
if use_in_transit is None else use_in_transit)
r = as_tensor_variable(r)
r = tt.reshape(r, (r.size, ))
t = as_tensor_variable(t)
# If use_in_transit, we should only evaluate the model at times where
# at least one planet is transiting
if use_in_transit:
transit_model = tt.shape_padleft(tt.zeros_like(r),
t.ndim) + tt.shape_padright(
tt.zeros_like(t), r.ndim)
inds = orbit.in_transit(t, r=r, texp=texp, light_delay=light_delay)
t = t[inds]
# Handle exposure time integration
if texp is None:
tgrid = t
rgrid = tt.shape_padleft(r, tgrid.ndim) + tt.shape_padright(
tt.zeros_like(tgrid), r.ndim)
else:
texp = as_tensor_variable(texp)
oversample = int(oversample)
oversample += 1 - oversample % 2
stencil = np.ones(oversample)
# Construct the exposure time integration stencil
if order == 0:
dt = np.linspace(-0.5, 0.5, 2 * oversample + 1)[1:-1:2]
elif order == 1:
dt = np.linspace(-0.5, 0.5, oversample)
stencil[1:-1] = 2
elif order == 2:
dt = np.linspace(-0.5, 0.5, oversample)
stencil[1:-1:2] = 4
stencil[2:-1:2] = 2
else:
raise ValueError("order must be <= 2")
stencil /= np.sum(stencil)
if texp.ndim == 0:
dt = texp * dt
else:
if use_in_transit:
dt = tt.shape_padright(texp[inds]) * dt
else:
dt = tt.shape_padright(texp) * dt
tgrid = tt.shape_padright(t) + dt
# Madness to get the shapes to work out...
rgrid = tt.shape_padleft(r, tgrid.ndim) + tt.shape_padright(
tt.zeros_like(tgrid), 1)
# Evalute the coordinates of the transiting body in the plane of the
# sky
coords = orbit.get_relative_position(tgrid, light_delay=light_delay)
b = tt.sqrt(coords[0]**2 + coords[1]**2)
b = tt.reshape(b, rgrid.shape)
los = tt.reshape(coords[2], rgrid.shape)
lc = self._compute_light_curve(b / orbit.r_star, rgrid / orbit.r_star,
los / orbit.r_star)
if texp is not None:
stencil = tt.shape_padright(tt.shape_padleft(stencil, t.ndim), 1)
lc = tt.sum(stencil * lc, axis=t.ndim)
if use_in_transit:
transit_model = tt.set_subtensor(transit_model[inds], lc)
return transit_model
else:
return lc
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 forward(self, x):
usum = tt.sum(x, axis=0)
q = tt.stack([usum**2, 0.5 * x[0] / usum])
return tt.log(q) - tt.log(1 - q)
def _do_norm(self, y):
alpha = ops.solve_lower(self._t, self._c, self._U, self._W,
y[:, None])[0][:, 0]
return tt.sum(alpha**2 / self._d)
def duration_to_eccentricity(func, duration, ror,
**kwargs): # pragma: no cover
num_planets = kwargs.pop("num_planets", 1)
orbit_type = kwargs.pop("orbit_type", KeplerianOrbit)
name = kwargs.get("name", "dur_ecc")
inputs = _get_consistent_inputs(
kwargs.get("a", None),
kwargs.get("period", None),
kwargs.get("rho_star", None),
kwargs.get("r_star", None),
kwargs.get("m_star", None),
kwargs.get("rho_star_units", None),
kwargs.get("m_planet", 0.0),
kwargs.get("m_planet_units", None),
)
a, period, rho_star, r_star, m_star, m_planet = inputs
b = kwargs.get("b", 0.0)
s = tt.sin(kwargs["omega"])
umax_inv = tt.switch(tt.lt(s, 0), tt.sqrt(1 - s**2), 1.0)
const = (period * tt.shape_padright(r_star) * tt.sqrt((1 + ror)**2 - b**2))
const /= np.pi * a
u = duration / const
e1 = -s * u**2 / ((s * u)**2 + 1)
e2 = tt.sqrt((s**2 - 1) * u**2 + 1) / ((s * u)**2 + 1)
models = []
logjacs = []
logprobs = []
for args in product(*(zip("np", (-1, 1)) for _ in range(num_planets))):
labels, signs = zip(*args)
# Compute the eccentricity branch
ecc = tt.stack([e1[i] + signs[i] * e2[i] for i in range(num_planets)])
# Work out the Jacobian
valid_ecc = tt.and_(tt.lt(ecc, 1.0), tt.ge(ecc, 0.0))
logjac = tt.switch(
tt.all(valid_ecc),
tt.sum(0.5 * tt.log(1 - ecc**2) + 2 * tt.log(s * ecc + 1) -
tt.log(tt.abs_(s + ecc)) - tt.log(const)),
-np.inf,
)
ecc = tt.switch(valid_ecc, ecc, tt.zeros_like(ecc))
# Create a sub-model to capture this component
with pm.Model(name="dur_ecc_" + "_".join(labels)) as model:
pm.Deterministic("ecc", ecc)
orbit = orbit_type(ecc=ecc, **kwargs)
logprob = tt.sum(func(orbit))
models.append(model)
logjacs.append(logjac)
logprobs.append(logprob)
# Compute the marginalized likelihood
logjacs = tt.stack(logjacs)
logprobs = tt.stack(logprobs)
logprob = tt.switch(
tt.gt(1.0 / u, umax_inv),
tt.sum(pm.logsumexp(logprobs + logjacs)),
-np.inf,
)
pm.Potential(name + "_logp", logprob)
pm.Deterministic(name + "_logjacs", logjacs)
pm.Deterministic(name + "_logprobs", logprobs)
# Loop over models and compute the marginalized values for all the
# parameters and deterministics
norm = tt.sum(pm.logsumexp(logjacs))
logw = tt.switch(
tt.gt(1.0 / u, umax_inv),
logjacs - norm,
-np.inf + tt.zeros_like(logjacs),
)
pm.Deterministic(name + "_logw", logw)
for k in models[0].named_vars.keys():
name = k[len(models[0].name) + 1:]
pm.Deterministic(
name,
sum(
tt.exp(logw[i]) * model.named_vars[model.name + "_" + name]
for i, model in enumerate(models)),
)