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 get_aor_from_transit_duration(duration, period, b, ror=None):
"""Get the semimajor axis implied by a circular orbit and duration
Args:
duration: The transit duration
period: The orbital period
b: The impact parameter of the transit
ror: The radius ratio of the planet to the star
Returns:
The semimajor axis in units of the stellar radius and the Jacobian
``d a / d duration``
"""
if ror is None:
ror = as_tensor_variable(0.0)
b2 = b**2
opk2 = (1 + ror)**2
phi = np.pi * duration / period
sinp = tt.sin(phi)
cosp = tt.cos(phi)
num = tt.sqrt(opk2 - b2 * cosp**2)
aor = num / sinp
grad = np.pi * cosp * (b2 - opk2) / (num * period * sinp**2)
return aor, grad
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_coefficients(self):
ar, cr, a, b, c, d = self.term.coefficients
# Real componenets
crd = cr * self.delta
coeffs = [2 * ar * (tt.cosh(crd) - 1) / crd ** 2, cr]
# Imaginary coefficients
cd = c * self.delta
dd = d * self.delta
c2 = c ** 2
d2 = d ** 2
factor = 2.0 / (self.delta * (c2 + d2)) ** 2
cos_term = tt.cosh(cd) * tt.cos(dd) - 1
sin_term = tt.sinh(cd) * tt.sin(dd)
C1 = a * (c2 - d2) + 2 * b * c * d
C2 = b * (c2 - d2) - 2 * a * c * d
coeffs += [
factor * (C1 * cos_term - C2 * sin_term),
factor * (C2 * cos_term + C1 * sin_term),
c,
d,
]
return coeffs
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 __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 _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 compute_gradient(self, t):
t = tt.as_tensor_variable(t).astype("float64")
dcos = tt.cos(self.omega * t)
dsin = tt.sin(self.omega * t)
df = 2 * np.pi * t * (self.sin * dcos - self.cos * dsin)
return (df, dcos, dsin)
def get_value(self, t):
t = tt.as_tensor_variable(t).astype("float64")
return (
self.cos * tt.cos(self.omega * t)
+ self.sin * tt.sin(self.omega * t)
)
def compute_gradient(self, t):
t = tt.as_tensor_variable(t).astype("float64")
damp = tt.sin(self.omega * t + self.phase)
dphi = self.amp * tt.cos(self.omega * t + self.phase)
df = 2 * np.pi * t * dphi
return (df, damp, dphi)
def get_value(self, t):
t = tt.as_tensor_variable(t).astype("float64")
return self.amp * tt.sin(self.omega * t + self.phase)
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_psd(self, omega):
psd0 = self.term.get_psd(omega)
arg = 0.5 * self.delta * omega
sinc = tt.switch(tt.neq(arg, 0), tt.sin(arg) / arg, tt.ones_like(arg))
return psd0 * sinc ** 2
