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 grad(self, inputs, gradients):
outputs = self(*inputs)
grads = (
tt.zeros_like(outputs[n])
if isinstance(b.type, theano.gradient.DisconnectedType)
else b
for n, b in enumerate(gradients[: len(self.spec["outputs"])])
)
return self.rev_op(*chain(inputs, outputs, grads))
def __init__(self, gp, *args, **kwargs):
if not HAS_PYMC3:
raise ImportError(
"pymc3 is required to use the CeleriteNormal distribution")
super().__init__(*args, **kwargs)
self.gp = gp
self.mean = (
self.median) = self.mode = self.gp.mean_value + tt.zeros_like(
self.gp._t)
def grad(self, inputs, gradients):
M, e = inputs
sinf, cosf = self(M, e)
ecosf = e * cosf
ome2 = 1 - e**2
dfdM = (1 + ecosf)**2 / ome2**1.5
dfde = (2 + ecosf) * sinf / ome2
bM = tt.zeros_like(M)
be = tt.zeros_like(M)
if not isinstance(gradients[0].type, aesara.gradient.DisconnectedType):
bM += gradients[0] * cosf * dfdM
be += gradients[0] * cosf * dfde
if not isinstance(gradients[1].type, aesara.gradient.DisconnectedType):
bM -= gradients[1] * sinf * dfdM
be -= gradients[1] * sinf * dfde
return [bM, be]
def _compute_light_curve(self, b, r, los=None):
"""Compute the light curve for a set of impact parameters and radii
.. note:: The stellar radius is *not* included in this method so the
coordinates should be in units of the star's radius.
Args:
b (tensor): A tensor of impact parameter values.
r (tensor): A tensor of radius ratios with the same shape as ``b``.
los (Optional[tensor]): The coordinates of the body along the
line-of-sight. If ``los > 0`` the body is between the observer
and the source.
"""
b = as_tensor_variable(b)
if los is None:
los = tt.ones_like(b)
lc = quad_limbdark_light_curve(self.c, b, r)
return tt.switch(tt.gt(los, 0), lc, tt.zeros_like(lc))
def get_relative_position(self, t, light_delay=False):
"""The planets' positions relative to the star
Args:
t: The times where the position should be evaluated.
Returns:
The components of the position vector at ``t`` in units of
``R_sun``.
"""
if light_delay:
raise NotImplementedError(
"Light travel time delay is not implemented for simple orbits")
dt = tt.mod(tt.shape_padright(t) - self._ref_time, self.period)
dt -= self._half_period
x = tt.squeeze(self.speed * dt)
y = tt.squeeze(self._b_norm + tt.zeros_like(dt))
m = tt.abs_(dt) < 0.5 * self.duration
z = tt.squeeze(m * 1.0 - (~m) * 1.0)
return x, y, z
def _get_consistent_inputs(a, period, rho_star, r_star, m_star, m_planet):
if a is None and period is None:
raise ValueError("values must be provided for at least one of a "
"and period")
if m_planet is not None:
m_planet = as_tensor_variable(to_unit(m_planet, u.M_sun))
if a is not None:
a = as_tensor_variable(to_unit(a, u.R_sun))
if m_planet is None:
m_planet = tt.zeros_like(a)
if period is not None:
period = as_tensor_variable(to_unit(period, u.day))
if m_planet is None:
m_planet = tt.zeros_like(period)
# Compute the implied density if a and period are given
implied_rho_star = False
if a is not None and period is not None:
if rho_star is not None or m_star is not None:
raise ValueError("if both a and period are given, you can't "
"also define rho_star or m_star")
# Default to a stellar radius of 1 if not provided
if r_star is None:
r_star = as_tensor_variable(1.0)
else:
r_star = as_tensor_variable(to_unit(r_star, u.R_sun))
# Compute the implied mass via Kepler's 3rd law
m_tot = 4 * np.pi * np.pi * a**3 / (G_grav * period**2)
# Compute the implied density
m_star = m_tot - m_planet
vol_star = 4 * np.pi * r_star**3 / 3.0
rho_star = m_star / vol_star
implied_rho_star = True
# Make sure that the right combination of stellar parameters are given
if r_star is None and m_star is None:
r_star = 1.0
if rho_star is None:
m_star = 1.0
if (not implied_rho_star) and sum(
arg is None for arg in (rho_star, r_star, m_star)) != 1:
raise ValueError("values must be provided for exactly two of "
"rho_star, m_star, and r_star")
if rho_star is not None and not implied_rho_star:
if has_unit(rho_star):
rho_star = as_tensor_variable(
to_unit(rho_star, u.M_sun / u.R_sun**3))
else:
rho_star = as_tensor_variable(rho_star) / gcc_per_sun
if r_star is not None:
r_star = as_tensor_variable(to_unit(r_star, u.R_sun))
if m_star is not None:
m_star = as_tensor_variable(to_unit(m_star, u.M_sun))
# Work out the stellar parameters
if rho_star is None:
rho_star = 3 * m_star / (4 * np.pi * r_star**3)
elif r_star is None:
r_star = (3 * m_star / (4 * np.pi * rho_star))**(1 / 3)
elif m_star is None:
m_star = 4 * np.pi * r_star**3 * rho_star / 3.0
# Work out the planet parameters
if a is None:
a = (G_grav * (m_star + m_planet) * period**2 / (4 * np.pi**2))**(1.0 /
3)
elif period is None:
period = (2 * np.pi * a**(3 / 2) / (tt.sqrt(G_grav *
(m_star + m_planet))))
return a, period, rho_star * gcc_per_sun, r_star, m_star, m_planet
def __init__(self,
period=None,
a=None,
t0=None,
t_periastron=None,
incl=None,
b=None,
duration=None,
ecc=None,
omega=None,
sin_omega=None,
cos_omega=None,
Omega=None,
m_planet=0.0,
m_star=None,
r_star=None,
rho_star=None,
ror=None,
model=None,
**kwargs):
add_citations_to_model(self.__citations__, model=model)
if "m_planet_units" in kwargs:
deprecation_warning(
"'m_planet_units' is deprecated; Use `with_unit` instead")
m_planet = with_unit(m_planet, kwargs.pop("m_planet_units"))
if "rho_star_units" in kwargs:
deprecation_warning(
"'rho_star_units' is deprecated; Use `with_unit` instead")
rho_star = with_unit(rho_star, kwargs.pop("rho_star_units"))
self.jacobians = defaultdict(lambda: defaultdict(None))
daordtau = None
if ecc is None and duration is not None:
if r_star is None:
r_star = as_tensor_variable(1.0)
if b is None:
raise ValueError(
"'b' must be provided for a circular orbit with a "
"'duration'")
if ror is None:
warnings.warn(
"When using the 'duration' parameter in KeplerianOrbit, "
"the 'ror' parameter should also be provided.",
UserWarning,
)
aor, daordtau = get_aor_from_transit_duration(duration,
period,
b,
ror=ror)
a = r_star * aor
duration = None
inputs = _get_consistent_inputs(a, period, rho_star, r_star, m_star,
m_planet)
(
self.a,
self.period,
self.rho_star,
self.r_star,
self.m_star,
self.m_planet,
) = inputs
self.m_total = self.m_star + self.m_planet
self.n = 2 * np.pi / self.period
self.a_star = self.a * self.m_planet / self.m_total
self.a_planet = -self.a * self.m_star / self.m_total
# Track the Jacobian between the duration and a
if daordtau is not None:
dadtau = self.r_star * daordtau
self.jacobians["duration"]["a"] = dadtau
self.jacobians["duration"]["a_star"] = (dadtau * self.m_planet /
self.m_total)
self.jacobians["duration"]["a_planet"] = (-dadtau * self.m_star /
self.m_total)
# rho = 3 * pi * (a/R)**3 / (G * P**2)
# -> drho / d(a/R) = 9 * pi * (a/R)**2 / (G * P**2)
self.jacobians["duration"]["rho_star"] = (
9 * np.pi * (self.a / self.r_star)**2 * daordtau *
gcc_per_sun / (G_grav * self.period**2))
self.K0 = self.n * self.a / self.m_total
if Omega is None:
self.Omega = None
else:
self.Omega = as_tensor_variable(Omega)
self.cos_Omega = tt.cos(self.Omega)
self.sin_Omega = tt.sin(self.Omega)
# Eccentricity
if ecc is None:
self.ecc = None
self.M0 = 0.5 * np.pi + tt.zeros_like(self.n)
incl_factor = 1
else:
self.ecc = as_tensor_variable(ecc)
if omega is not None:
if sin_omega is not None and cos_omega is not None:
raise ValueError(
"either 'omega' or 'sin_omega' and 'cos_omega' can be "
"provided")
self.omega = as_tensor_variable(omega)
self.cos_omega = tt.cos(self.omega)
self.sin_omega = tt.sin(self.omega)
elif sin_omega is not None and cos_omega is not None:
self.cos_omega = as_tensor_variable(cos_omega)
self.sin_omega = as_tensor_variable(sin_omega)
self.omega = tt.arctan2(self.sin_omega, self.cos_omega)
else:
raise ValueError("both e and omega must be provided")
opsw = 1 + self.sin_omega
E0 = 2 * tt.arctan2(
tt.sqrt(1 - self.ecc) * self.cos_omega,
tt.sqrt(1 + self.ecc) * opsw,
)
self.M0 = E0 - self.ecc * tt.sin(E0)
ome2 = 1 - self.ecc**2
self.K0 /= tt.sqrt(ome2)
incl_factor = (1 + self.ecc * self.sin_omega) / ome2
# The Jacobian for the transform cos(i) -> b
self.dcosidb = self.jacobians["b"]["cos_incl"] = (incl_factor *
self.r_star / self.a)
if b is not None:
if incl is not None or duration is not None:
raise ValueError(
"only one of 'incl', 'b', and 'duration' can be given")
self.b = as_tensor_variable(b)
self.cos_incl = self.dcosidb * self.b
self.incl = tt.arccos(self.cos_incl)
elif incl is not None:
if duration is not None:
raise ValueError(
"only one of 'incl', 'b', and 'duration' can be given")
self.incl = as_tensor_variable(incl)
self.cos_incl = tt.cos(self.incl)
self.b = self.cos_incl / self.dcosidb
elif duration is not None:
# This assertion should never be hit because of the first
# conditional in this method, but let's keep it here anyways
assert self.ecc is not None
self.duration = as_tensor_variable(to_unit(duration, u.day))
c = tt.sin(np.pi * self.duration * incl_factor / self.period)
c2 = c * c
aor = self.a_planet / self.r_star
esinw = self.ecc * self.sin_omega
self.b = tt.sqrt(
(aor**2 * c2 - 1) / (c2 * esinw**2 + 2 * c2 * esinw + c2 -
self.ecc**4 + 2 * self.ecc**2 - 1))
self.b *= 1 - self.ecc**2
self.cos_incl = self.dcosidb * self.b
self.incl = tt.arccos(self.cos_incl)
else:
zla = tt.zeros_like(self.a)
self.incl = 0.5 * np.pi + zla
self.cos_incl = zla
self.b = zla
if t0 is not None and t_periastron is not None:
raise ValueError("you can't define both t0 and t_periastron")
if t0 is None and t_periastron is None:
t0 = tt.zeros_like(self.period)
if t0 is None:
self.t_periastron = as_tensor_variable(t_periastron)
self.t0 = self.t_periastron + self.M0 / self.n
else:
self.t0 = as_tensor_variable(t0)
self.t_periastron = self.t0 - self.M0 / self.n
self.tref = self.t_periastron - self.t0
self.sin_incl = tt.sin(self.incl)
def in_transit(self, t, r=0.0, texp=None, light_delay=False):
"""Get a list of timestamps that are in transit
Args:
t (vector): A vector of timestamps to be evaluated.
r (Optional): The radii of the planets.
texp (Optional[float]): The exposure time.
Returns:
The indices of the timestamps that are in transit.
"""
if light_delay:
raise NotImplementedError(
"Light travel time delay not yet implemented for `in_transit`")
z = tt.zeros_like(self.a)
r = as_tensor_variable(r) + z
R = self.r_star + z
# Wrap the times into time since transit
hp = 0.5 * self.period
dt = tt.mod(self._warp_times(t) + hp, self.period) - hp
if self.ecc is None:
# Equation 14 from Winn (2010)
k = r / R
arg = tt.square(1 + k) - tt.square(self.b)
factor = R / (self.a * self.sin_incl)
hdur = hp * tt.arcsin(factor * tt.sqrt(arg)) / np.pi
t_start = -hdur
t_end = hdur
flag = z
else:
M_contact = ops.contact_points(
self.a,
self.ecc,
self.cos_omega,
self.sin_omega,
self.cos_incl + z,
self.sin_incl + z,
R + r,
)
flag = M_contact[2]
t_start = (M_contact[0] - self.M0) / self.n
t_start = tt.mod(t_start + hp, self.period) - hp
t_end = (M_contact[1] - self.M0) / self.n
t_end = tt.mod(t_end + hp, self.period) - hp
t_start = tt.switch(tt.gt(t_start, 0.0), t_start - self.period,
t_start)
t_end = tt.switch(tt.lt(t_end, 0.0), t_end + self.period, t_end)
if texp is not None:
t_start -= 0.5 * texp
t_end += 0.5 * texp
mask = tt.any(tt.and_(dt >= t_start, dt <= t_end), axis=-1)
result = ifelse(
tt.all(tt.eq(flag, 0)),
tt.arange(t.shape[0])[mask],
tt.arange(t.shape[0]),
)
return result
def _get_true_anomaly(self, t, _pad=True):
M = (self._warp_times(t, _pad=_pad) - self.tref) * self.n
if self.ecc is None:
return tt.sin(M), tt.cos(M)
sinf, cosf = ops.kepler(M, self.ecc + tt.zeros_like(M))
return sinf, cosf
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 logp(self, value):
return tt.zeros_like(tt.as_tensor_variable(value))
def get_star_position(self, t, light_delay=False):
nothing = tt.zeros_like(as_tensor_variable(t))
return nothing, nothing, nothing
def _zeros_like(self, tensor):
return tt.zeros_like(tensor)
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)),
)
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)
]
