def underdamped(self):
Q = self.Q
f = tt.sqrt(tt.maximum(4.0 * Q ** 2 - 1.0, self.eps))
a = self.S0 * self.w0 * Q
c = 0.5 * self.w0 / Q
empty = tt.zeros(0, dtype=self.dtype)
return (
empty,
empty,
tt.stack([a]),
tt.stack([a / f]),
tt.stack([c]),
tt.stack([c * f]),
)
def overdamped(self):
Q = self.Q
f = tt.sqrt(tt.maximum(1.0 - 4.0 * Q ** 2, self.eps))
empty = tt.zeros(0, dtype=self.dtype)
return (
0.5
* self.S0
* self.w0
* Q
* tt.stack([1.0 + 1.0 / f, 1.0 - 1.0 / f]),
0.5 * self.w0 / Q * tt.stack([1.0 - f, 1.0 + f]),
empty,
empty,
empty,
empty,
)
def get_cl(u1, u2):
u1 = as_tensor_variable(u1)
u2 = as_tensor_variable(u2)
c0 = 1 - u1 - 1.5 * u2
c1 = u1 + 2 * u2
c2 = -0.25 * u2
norm = np.pi * (c0 + c1 / 1.5)
return tt.stack([c0, c1, c2]) / norm
def forward(self, x):
p = x[0]
b = x[1]
pl = self.min_radius
r11 = (b / (1 + pl) - 1) * (1 - self.Ar) + 1
r21 = (p - pl) / self.dr
arg = p - b - self.dr + 1
q1 = (arg / self.dr)**2
q2 = (pl - b + 1) / arg
r12 = q1 * self.Ar
r22 = q2
y = tt.switch(b <= 1, tt.stack((r11, r21), axis=0),
tt.stack((r12, r22), axis=0))
return tt.log(y) - tt.log(1 - y)
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 backward(self, y):
y = tt.nnet.sigmoid(y)
r1 = y[0]
r2 = y[1]
pl, pu = self.min_radius, self.max_radius
b1 = (1 + pl) * (1 + (r1 - 1) / (1 - self.Ar))
p1 = pl + r2 * self.dr
q1 = r1 / self.Ar
q2 = r2
b2 = (1 + pl) + tt.sqrt(q1) * q2 * self.dr
p2 = pu - self.dr * tt.sqrt(q1) * (1 - q2)
pb = tt.switch(
r1 > self.Ar,
tt.stack((p1, b1), axis=0),
tt.stack((p2, b2), axis=0),
)
return pb
def get_light_curve(
self,
orbit=None,
r=None,
t=None,
texp=None,
oversample=7,
order=0,
use_in_transit=None,
light_delay=False,
):
if self.num_planets is None:
try:
vec = orbit.period.tag.test_value
except AttributeError:
raise ValueError(
"Can't compute num_planets, please provide a value")
num_planets = len(np.atleast_1d(vec))
else:
num_planets = int(self.num_planets)
if num_planets <= 1:
func = _wrapper(
self.base_light_curve,
orbit=orbit,
r=r,
texp=texp,
oversample=oversample,
order=order,
use_in_transit=use_in_transit,
light_delay=light_delay,
)
mn = orbit.t0
mx = orbit.t0 + orbit.period
return interp(
0,
tt.mod(t - orbit.t0, orbit.period) + orbit.t0,
mn,
mx,
(mx - mn) / (self.num_phase + 1),
func,
)[:, None]
ys = []
for n in range(num_planets):
func = _wrapper(
self.base_light_curve,
orbit=orbit,
r=r,
texp=texp,
oversample=oversample,
order=order,
use_in_transit=use_in_transit,
light_delay=light_delay,
)
mn = orbit.t0[n]
mx = orbit.t0[n] + orbit.period[n]
ys.append(
interp(
n,
tt.mod(t - orbit.t0[n], orbit.period[n]) + orbit.t0[n],
mn,
mx,
(mx - mn) / (self.num_phase + 1),
func,
))
return tt.stack(ys, axis=-1)
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_model_dt(self, t):
vals = []
for i in range(len(self.ttvs)):
inds = tt.extra_ops.searchsorted(self._bin_edges[i], t)
vals.append(self._bin_values[i][inds])
return tt.stack(vals, -1)
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 backward(self, y):
q = tt.nnet.sigmoid(y)
sqrtq1 = tt.sqrt(q[0])
twoq2 = 2 * q[1]
u = tt.stack([sqrtq1 * twoq2, sqrtq1 * (1 - twoq2)])
return u
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)),
)