Planets around other stars, and a visitor from another star

import logging
import time
import warnings

import arviz as az
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numba as nb
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import plotly.io as pio
import pymc as pm
import pytensor.tensor as pt
from matplotlib.colors import LogNorm
from pytensor import wrap_jax
from scipy import integrate, stats

from pymc_challenges import data

jax.config.update("jax_enable_x64", True)
pio.renderers.default = "plotly_mimetype+notebook_connected"
RANDOM_SEED = 78
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
logging.getLogger("pymc").setLevel(logging.WARNING)
warnings.filterwarnings("ignore", category=RuntimeWarning, module="arviz")
warnings.filterwarnings("ignore", message=".*object mode.*")
warnings.filterwarnings("ignore", category=integrate.IntegrationWarning)
pd.set_option("display.width", 170)
pd.set_option("display.max_columns", 20)
T_START = time.time()

@nb.njit(cache=True)
def _biweight_loc(x, c=5.0):
    m = np.median(x)
    for _ in range(5):
        mad = np.median(np.abs(x - m))
        if mad == 0:
            return m
        num = 0.0
        den = 0.0
        for i in range(len(x)):
            u = (x[i] - m) / (c * mad)
            if abs(u) < 1:
                w = (1 - u * u) ** 2
                num += w * (x[i] - m)
                den += w
        if den == 0:
            return m
        m = m + num / den
    return m


@nb.njit(parallel=True, cache=True)
def biweight_trend(t, y, use, window):
    """Sliding biweight location of y within +-window/2 of each point, from points with use=True."""
    n = len(t)
    out = np.empty(n)
    lo = np.searchsorted(t, t - window / 2)
    hi = np.searchsorted(t, t + window / 2)
    for i in nb.prange(n):
        seg = y[lo[i]:hi[i]][use[lo[i]:hi[i]]]
        out[i] = _biweight_loc(seg) if len(seg) > 3 else np.nan
    return out


trend = biweight_trend(t_all, y_all, np.ones(len(t_all), bool), 1.0)
resid = y_all - trend
SIG = 1.4826 * np.median(np.abs(resid))
keep = resid < 5 * SIG            # clip upward outliers only: a transit is a DOWNWARD excursion
t_all, y_all, resid, trend = t_all[keep], y_all[keep], resid[keep], trend[keep]
print(f"robust sd of the detrended flux: {SIG:.1f} ppm; {np.sum(~keep)} upward outliers clipped")

fig, axes = plt.subplots(2, 1, figsize=(11, 4.5), sharex=True)
zoom = (t_all > 585) & (t_all < 600)
axes[0].plot(t_all[zoom], y_all[zoom], ".", ms=3, color="0.5")
axes[0].plot(t_all[zoom], trend[zoom], color="C1", label="1-day biweight trend")
axes[0].legend()
axes[0].set(ylabel="flux (ppm)")
axes[1].plot(t_all[zoom], resid[zoom], ".", ms=3, color="C0")
axes[1].set(xlabel="time (BKJD, days)", ylabel="detrended (ppm)");

@nb.njit(parallel=True, cache=True)
def bls(t, y, w, periods, dur_lo, dur_hi, durations, dbin):
    """Box least squares. For each period: the best dip over epochs and allowed durations.
    Returns dchi2, duration, mid-time and depth of the best box."""
    n = len(t)
    wsum = w.sum()
    ybar = (w * y).sum() / wsum
    tref = t[0]
    tt = t - tref
    npd = len(periods)
    power = np.zeros(npd)
    bdur = np.zeros(npd)
    bt0 = np.zeros(npd)
    bdep = np.zeros(npd)
    for ip in nb.prange(npd):
        P = periods[ip]
        nbins = int(np.ceil(P / dbin))
        binw = P / nbins
        invP = 1.0 / P
        sw = np.zeros(nbins)
        sy = np.zeros(nbins)
        for i in range(n):
            ph = tt[i] * invP
            ph -= np.floor(ph)
            b = min(int(ph * nbins), nbins - 1)
            sw[b] += w[i]
            sy[b] += w[i] * (y[i] - ybar)
        best = 0.0
        for idur in range(len(durations)):
            if durations[idur] < dur_lo[ip] or durations[idur] > dur_hi[ip]:
                continue
            k = int(np.round(durations[idur] / binw))
            if k < 1 or k >= nbins // 2:
                continue
            W = 0.0
            S = 0.0
            for j in range(k):
                W += sw[j]
                S += sy[j]
            for j0 in range(nbins):
                if W > 0 and S < 0:
                    p = S * S / W * wsum / (wsum - W)
                    if p > best:
                        best = p
                        bdur[ip] = k * binw
                        bt0[ip] = tref + (j0 + k / 2) * binw
                        bdep[ip] = -S / W
                jin = (j0 + k) % nbins
                W += sw[jin] - sw[j0]
                S += sy[jin] - sy[j0]
        power[ip] = best
    return power, bdur, bt0, bdep


G_CGS = 6.674e-8
CADENCE = CAD_DAYS = 29.4244 / 1440           # Kepler long cadence, days


def d0_days(P, rho=1.41):
    """Central transit duration (days) of a small planet around a star of mean density rho (g/cm3)."""
    aR = (G_CGS * rho * (P * 86400) ** 2 / (3 * np.pi)) ** (1 / 3)
    return P / (np.pi * aR)


def period_grid(pmin, pmax, T, oversample=2.0):
    ps = [pmin]
    while ps[-1] < pmax:
        ps.append(ps[-1] * (1 + 0.4 * d0_days(ps[-1]) / (oversample * T)))
    return np.array(ps)


DURS = np.array([1, 1.25, 1.5, 2, 2.5, 3, 3.5, 4, 5, 6, 7, 8.5, 10, 12, 14]) / 24


def run_bls(t, y, periods, sig=SIG):
    d0 = d0_days(periods)
    pw, dur, t0, dep = bls(t, y, np.full(len(t), sig**-2), periods, 0.4 * d0, 1.6 * d0, DURS, 1 / 72)
    return pd.DataFrame({"P": periods, "snr": np.sqrt(pw), "dur": dur, "t0": t0, "depth": dep})


def fold_phase(t, P, t0):
    """Time from the nearest transit centre (days)."""
    return (t - t0 + 0.5 * P) % P - 0.5 * P


PERIODS = period_grid(0.5, 100, t_all.max() - t_all.min())
print(f"{len(PERIODS):,} trial periods")
tic = time.time()
search = []
t_s, r_s = t_all.copy(), resid.copy()
for it in range(3):
    res = run_bls(t_s, r_s, PERIODS)
    best = res.loc[res.snr.idxmax()]
    sde = (best.snr - res.snr.mean()) / res.snr.std()
    search.append((res, best, t_s.copy(), r_s.copy()))
    print(f"search {it + 1}: P = {best.P:.5f} d, duration {best.dur * 24:.1f} h, depth {best.depth:.0f} ppm, "
          f"SNR {best.snr:.1f}, SDE {sde:.1f}")
    # mask this signal (1.5 durations either side of each transit) and detrend again
    m = np.abs(fold_phase(t_s, best.P, best.t0)) < 1.5 * best.dur
    y_s = y_all[np.isin(t_all, t_s)][~m]
    t_s = t_s[~m]
    r_s = y_s - biweight_trend(t_s, y_s, np.ones(len(t_s), bool), 1.0)
print(f"three searches: {time.time() - tic:.0f} s")

fig, axes = plt.subplots(3, 2, figsize=(12, 8.5), gridspec_kw={"width_ratios": [2.2, 1]})
for i, (res, best, t_s, r_s) in enumerate(search):
    ax = axes[i, 0]
    ax.plot(res.P, res.snr, lw=0.4, color="C0")
    ax.plot(best.P, best.snr, "v", color="C1", ms=7)
    ax.set(xscale="log", ylabel="BLS SNR", xlabel="trial period (days)" if i == 2 else None,
           title=f"search {i + 1}: best P = {best.P:.4f} d, SNR {best.snr:.1f}")
    ph = fold_phase(t_s, best.P, best.t0) * 24
    w = np.abs(ph) < 5 * best.dur * 24
    ax = axes[i, 1]
    ax.plot(ph[w], r_s[w], ",", color="0.6")
    edges = np.linspace(-5 * best.dur * 24, 5 * best.dur * 24, 41)
    k = np.digitize(ph[w], edges)
    yb = pd.Series(r_s[w]).groupby(k).mean()
    xb = pd.Series(ph[w]).groupby(k).mean()
    ax.plot(xb, yb, "o", ms=3, color="C1")
    ax.set(ylim=(-600, 250), xlabel="hours from mid-transit" if i == 2 else None, ylabel="ppm")
axes[0, 0].annotate("2P", xy=(2 * search[0][1].P, 118), xytext=(3.5, 150), fontsize=8,
                    arrowprops=dict(arrowstyle="->", lw=0.6))
axes[0, 0].annotate("10c", xy=(45.3, 123), xytext=(20, 150), fontsize=8, arrowprops=dict(arrowstyle="->", lw=0.6))
res3 = search[2][0]
peaks = res3.sort_values("snr", ascending=False)
top = []
for row in peaks.itertuples():            # five highest peaks at clearly different periods
    if all(abs(row.P / p_ - 1) > 0.01 for p_ in top):
        top.append(row.P)
    if len(top) == 5:
        break
print("search 3, five highest peaks:", ", ".join(f"P = {p_:.4f} d (SNR {res3.snr[res3.P == p_].iloc[0]:.1f})" for p_ in top))
single = search[0][1].depth / SIG * np.sqrt(search[0][1].dur / CAD_DAYS)
print(f"SNR of a single transit of the 0.84-day signal: about {single:.1f}");

GL_X, GL_W = np.polynomial.legendre.leggauss(16)
TH = (GL_X + 1) * np.pi / 2           # theta nodes on (0, pi)
TW = GL_W * np.pi / 2


def blocked_np(z, k, u1, u2):
    """Fraction of the flux of a quadratically limb-darkened star blocked by an opaque disk."""
    z = np.maximum(np.asarray(z, float), 1e-9)
    r0, r1 = np.abs(z - k), np.minimum(z + k, 1.0)
    half = np.maximum(r1 - r0, 0) / 2
    r = ((r0 + r1) / 2)[..., None] - half[..., None] * np.cos(TH)
    dr = half[..., None] * np.sin(TH)
    alpha = np.arccos(np.clip((r**2 + z[..., None] ** 2 - k**2) / (2 * r * z[..., None] + 1e-300), -1, 1))
    mu = np.sqrt(np.maximum(1 - r**2, 0))
    part = ((1 - u1 * (1 - mu) - u2 * (1 - mu) ** 2) * 2 * alpha * r * dr * TW).sum(-1)
    R = np.maximum(k - z, 0) / 2          # disk of radius k - z fully covered when z < k
    rr = R[..., None] * (1 - np.cos(TH))
    mu2 = np.sqrt(np.maximum(1 - rr**2, 0))
    full = ((1 - u1 * (1 - mu2) - u2 * (1 - mu2) ** 2) * 2 * np.pi * rr * R[..., None] * np.sin(TH) * TW).sum(-1)
    return (part + full) / (np.pi * (1 - u1 / 3 - u2 / 6))


def uniform_overlap(z, k):
    """Exact overlap area / pi of two disks (radii 1 and k): the uniform-source transit."""
    z = np.asarray(z, float)
    out = np.where(z <= 1 - k, k**2, 0.0)
    part = (z > abs(1 - k)) & (z < 1 + k)
    zz = z[part]
    k0 = np.arccos((k**2 + zz**2 - 1) / (2 * k * zz))
    k1 = np.arccos((1 - k**2 + zz**2) / (2 * zz))
    out[part] = (k**2 * k0 + k1 - 0.5 * np.sqrt(4 * zz**2 - (1 + zz**2 - k**2) ** 2)) / np.pi
    return out


def blocked_bruteforce(z, k, u1, u2):
    """Reference: 2-D adaptive integration of I over the planet's disk."""
    def f(s, phi):
        x, y = z + s * np.cos(phi), s * np.sin(phi)
        r2 = x * x + y * y
        if r2 >= 1:
            return 0.0
        mu = np.sqrt(1 - r2)
        return (1 - u1 * (1 - mu) - u2 * (1 - mu) ** 2) * s
    val, _ = integrate.dblquad(f, 0, 2 * np.pi, 0, k, epsabs=1e-13, epsrel=1e-10)
    return val / (np.pi * (1 - u1 / 3 - u2 / 6))


zs = np.array([0.0, 0.006, 0.02, 0.4, 0.9, 0.975, 0.99, 1.0, 1.015])
print("uniform source vs exact overlap, max |error| / k^2:",
      f"{np.abs(blocked_np(zs, 0.02, 0, 0) - uniform_overlap(zs, 0.02)).max() / 0.02**2:.1e}")
errs = [abs(blocked_np(np.array([z]), 0.02, 0.45, 0.2)[0] - blocked_bruteforce(z, 0.02, 0.45, 0.2)) / 0.02**2
        for z in [0.0, 0.01, 0.5, 0.98, 1.0, 1.01]]
print(f"limb-darkened vs 2-D brute force (k = 0.02), max |error| / k^2: {max(errs):.1e}")
small = blocked_np(np.array([0.0]), 1e-3, 0.45, 0.2)[0] / (1e-6 / (1 - 0.45 / 3 - 0.2 / 6))
print(f"small-planet limit at disk centre (should be 1): {small:.6f}")

def z_of_t(t, P, t0, aR, b):
    ph = 2 * np.pi * (t - t0) / P
    return np.where(np.cos(ph) > 0, np.sqrt(aR**2 * np.sin(ph) ** 2 + b**2 * np.cos(ph) ** 2), 10.0)


def offsets(n):
    """Sub-exposure times (days) that average a model over one long-cadence exposure."""
    return (np.arange(n) - (n - 1) / 2) / n * CADENCE


NSUB = {"b": 11, "c": 5}
aR_b_guess = (G_CGS * 1.068 * (Pb * 86400) ** 2 / (3 * np.pi)) ** (1 / 3)
fig, axes = plt.subplots(1, 2, figsize=(12, 3.8))
tt = np.linspace(-0.1, 0.1, 2001)
for b_ in [0.0, 0.5, 0.8, 0.95, 1.0]:
    aR = 20.0
    axes[0].plot(tt * 24, -1e6 * blocked_np(z_of_t(tt, 5.0, 0, aR, b_), 0.05, 0.45, 0.2), label=f"b = {b_}")
axes[0].set(xlabel="hours from mid-transit", ylabel="ppm", title="k = 0.05, P = 5 d: shape against b")
axes[0].legend(fontsize=8)
inst = -1e6 * blocked_np(z_of_t(tt, Pb, 0, aR_b_guess, 0.33), 0.0125, 0.45, 0.2)
ax = axes[1]
ax.plot(tt * 24, inst, color="0.5", label="instantaneous")
dev = {}
for ns in [41, 5, 11]:
    sm = -1e6 * blocked_np(z_of_t(tt[:, None] + offsets(ns), Pb, 0, aR_b_guess, 0.33), 0.0125, 0.45, 0.2).mean(1)
    dev[ns] = sm
    ax.plot(tt * 24, sm, label=f"29.4-min average ({ns} sub-steps)", lw=1.5 if ns == 41 else 1)
for ns in [5, 11]:
    d = (dev[ns] - dev[41])[np.abs(dev[41]) > 0]
    print(f"{ns:2d} sub-steps vs 41: largest difference {np.abs(d).max():.1f} ppm, rms over the transit {np.sqrt(np.mean(d**2)):.1f} ppm")
ax.set(xlabel="hours from mid-transit", ylabel="ppm", title="Kepler-10b-like: the exposure smooths the shape")
ax.legend(fontsize=8);

S_GRID = np.r_[np.linspace(0, 1, 201)[:-1], np.linspace(1, 2, 65)]


def s_to_z_np(s, k):
    return np.where(s < 1, s * (1 - k), 1 - k + (s - 1) * 2 * k)


def blocked_table(k, u1, u2):
    """The blocked fraction on S_GRID, as a PyTensor graph (same quadrature as blocked_np)."""
    s = pt.as_tensor(S_GRID)
    z = pt.maximum(pt.switch(s < 1, s * (1 - k), 1 - k + (s - 1) * 2 * k), 1e-9)
    r0, r1 = pt.abs(z - k), pt.minimum(z + k, 1.0)
    half = pt.maximum(r1 - r0, 0) / 2
    r = ((r0 + r1) / 2)[:, None] - half[:, None] * np.cos(TH)
    dr = half[:, None] * np.sin(TH)
    c = pt.clip((r**2 + z[:, None] ** 2 - k**2) / (2 * r * z[:, None]), -1 + 1e-12, 1 - 1e-12)
    mu = pt.sqrt(pt.maximum(1 - r**2, 1e-12))
    part = ((1 - u1 * (1 - mu) - u2 * (1 - mu) ** 2) * 2 * pt.arccos(c) * r * dr * TW).sum(-1)
    R = pt.maximum(k - z, 0) / 2
    rr = R[:, None] * (1 - np.cos(TH))
    mu2 = pt.sqrt(pt.maximum(1 - rr**2, 1e-12))
    full = ((1 - u1 * (1 - mu2) - u2 * (1 - mu2) ** 2) * 2 * np.pi * rr * R[:, None] * np.sin(TH) * TW).sum(-1)
    return (part + full) / (np.pi * (1 - u1 / 3 - u2 / 6))


def interp_blocked(z, k, tab):
    """Linear interpolation in s on the piecewise-uniform grid (200 steps on [0,1), 64 on [1,2])."""
    s = pt.clip(pt.switch(z < 1 - k, z / (1 - k), 1 + (z - (1 - k)) / (2 * k)), 0, 2)
    x = pt.switch(s < 1, s * 200, 200 + (s - 1) * 64)
    j = pt.clip(pt.floor(x), 0, len(S_GRID) - 2).astype("int64")
    w = x - j
    return tab[j] * (1 - w) + tab[j + 1] * w


# the same interpolation in NumPy, against direct quadrature
zz = np.linspace(0, 1.03, 3000)
tab_np = blocked_np(np.maximum(s_to_z_np(S_GRID, 0.0125), 1e-9), 0.0125, 0.45, 0.2)
sv = np.clip(np.where(zz < 1 - 0.0125, zz / (1 - 0.0125), 1 + (zz - (1 - 0.0125)) / 0.025), 0, 2)
xv = np.where(sv < 1, sv * 200, 200 + (sv - 1) * 64)
jv = np.clip(np.floor(xv), 0, len(S_GRID) - 2).astype(int)
approx = tab_np[jv] * (1 - (xv - jv)) + tab_np[jv + 1] * (xv - jv)
print(f"interpolation error, k = 0.0125: max {1e6 * np.abs(approx - blocked_np(zz, 0.0125, 0.45, 0.2)).max():.2f} ppm")

def make_windows(t, y, centres, halfwidths):
    """Points within the (merged) windows; per-window orthonormal polynomial bases."""
    iv = sorted((c - h, c + h) for c, h in zip(centres, halfwidths))
    merged = []
    for a, b in iv:
        if merged and a <= merged[-1][1]:
            merged[-1][1] = max(merged[-1][1], b)
        else:
            merged.append([a, b])
    merged = np.array(merged)
    j = np.searchsorted(merged[:, 0], t, side="right") - 1
    inw = (j >= 0) & (t <= merged[np.maximum(j, 0), 1])
    tw, yw = t[inw], y[inw]
    _, win = np.unique(j[inw], return_inverse=True)
    counts = np.bincount(win)
    ends = np.cumsum(counts)
    starts = ends - counts
    Q = np.zeros((len(tw), 3))
    npar = 0
    for a, b in zip(starts, ends):
        tt_ = tw[a:b] - tw[a:b].mean()
        deg = min(2 if np.ptp(tw[a:b]) > 0.5 else 1, b - a - 1)
        Q[a:b, :deg + 1] = np.linalg.qr(np.vander(tt_, deg + 1, increasing=True))[0]
        npar += deg + 1
    return dict(t=tw, y=yw, Q=Q, starts=starts, ends=ends, npar=npar)


def centres_of(P, t0, lo, hi, offset=0.0):
    n = np.arange(np.floor((lo - t0) / P) - 1, np.ceil((hi - t0) / P) + 2)
    return t0 + (n + offset) * P


def marginal_loglike(W, model_flux, log_sigma):
    r = pt.as_tensor(W["y"]) - model_flux
    cs = pt.concatenate([pt.zeros((1, 3)), pt.cumsum(r[:, None] * W["Q"], axis=0)], axis=0)
    proj = cs[W["ends"]] - cs[W["starts"]]
    chi2 = (pt.sum(r**2) - pt.sum(proj**2)) / pt.exp(2 * log_sigma)
    return -0.5 * chi2 - (len(W["t"]) - W["npar"]) * log_sigma


# raw flux (outliers already clipped), reference epochs near the middle of the baseline
tmid = 0.5 * (t_all.min() + t_all.max())
EPH = {}
for name, (P, t0, dur) in {"b": (Pb, t0b, durb), "c": (Pc, t0c, durc)}.items():
    EPH[name] = (P, t0 + np.round((tmid - t0) / P) * P, dur)
lo, hi = t_all.min(), t_all.max()
cen_b = centres_of(EPH["b"][0], EPH["b"][1], lo, hi)
cen_bo = centres_of(EPH["b"][0], EPH["b"][1], lo, hi, 0.5)
cen_c = centres_of(EPH["c"][0], EPH["c"][1], lo, hi)
W_c = make_windows(t_all, y_all, cen_c, np.full(len(cen_c), 0.6))
W_bc = make_windows(t_all, y_all, np.r_[cen_b, cen_bo, cen_c],
                    np.r_[np.full(len(cen_b) + len(cen_bo), 0.16), np.full(len(cen_c), 0.6)])
for nm, W in [("10c only", W_c), ("10b + occultation + 10c", W_bc)]:
    print(f"{nm}: {len(W['t']):,} points in {len(W['starts'])} windows, {W['npar']} baseline coefficients")

def add_planet(name, W, rho, u1, u2, k_bounds=(0.003, 0.1), odd_even=False, occultation=False):
    """Adds one planet's parameters and returns its (negative) contribution to the flux in ppm."""
    P0, T00, dur = EPH[name]
    P = pm.Deterministic(f"P_{name}", P0 + (2e-5 if name == "b" else 2e-3) * pm.Normal(f"dP_{name}", 0, 1))
    T0 = pm.Deterministic(f"T0_{name}", T00 + 0.01 * pm.Normal(f"dT_{name}", 0, 1))
    k = pm.Deterministic(f"k_{name}", pt.exp(pm.Uniform(f"log_k_{name}", *np.log(k_bounds))))
    b = pm.Deterministic(f"b_{name}", pm.Uniform(f"ub_{name}", 0.0, 1.0) * (1 + k))
    if rho is None:   # parameterised by the transit duration D instead of the stellar density
        D = pt.exp(pm.Uniform(f"log_D_{name}", np.log(1 / 24), np.log(1.0)))
        aR = P / (np.pi * D) * pt.sqrt((1 + k) ** 2 - b**2)
        pm.Deterministic("rho", 3 * np.pi * aR**3 / (G_CGS * (P * 86400) ** 2))
    else:
        aR = (G_CGS * rho * (P * 86400) ** 2 / (3 * np.pi)) ** (1 / 3)
    flux = pt.zeros(len(W["t"]))
    phase = fold_phase(W["t"], P0, T00)
    near = np.flatnonzero(np.abs(phase) < 0.75 * dur + CADENCE)   # only these can be in transit
    ph = 2 * np.pi * (W["t"][near][:, None] + offsets(NSUB[name]) - T0) / P
    z = pt.sqrt(aR**2 * pt.sin(ph) ** 2 + b**2 * pt.cos(ph) ** 2)
    blk = interp_blocked(z, k, blocked_table(k, u1, u2))
    if odd_even:   # fractional depth difference between odd and even transits
        eps = pm.Normal(f"eps_{name}", 0.0, 0.5)
        odd = (np.round((W["t"][near] - T00) / P0) % 2 == 1)[:, None]
        blk = blk * pt.switch(odd, 1 + eps / 2, 1 - eps / 2)
    flux = pt.inc_subtensor(flux[near], -1e6 * blk.mean(1))
    if occultation:   # planet behind the star: its own light is hidden (uniform disk), depth in ppm
        occ = pm.Normal(f"occ_{name}", 0.0, 50.0)
        near_o = np.flatnonzero(np.abs(np.abs(phase) - 0.5 * P0) < 1.2 * dur)
        ph = 2 * np.pi * (W["t"][near_o] - T0) / P
        z = pt.sqrt(aR**2 * pt.sin(ph) ** 2 + b**2 * pt.cos(ph) ** 2)
        hidden = interp_blocked(z, k, blocked_table(k, 0.0, 0.0) / k**2)
        flux = pt.inc_subtensor(flux[near_o], -occ * hidden)
    return flux


def transit_model(W, planets, rho_prior, **kw):
    with pm.Model() as m:
        if rho_prior == "vague":
            rho = pm.Deterministic("rho", pt.exp(pm.Uniform("log_rho", np.log(0.1), np.log(10.0))))
        elif rho_prior == "duration":
            rho = None
        else:
            rho = pm.Normal("rho", *rho_prior)
        q1 = pm.Uniform("q1", 0.0, 1.0)
        q2 = pm.Uniform("q2", 0.0, 1.0)
        u1, u2 = 2 * pt.sqrt(q1) * q2, pt.sqrt(q1) * (1 - 2 * q2)
        log_sigma = pm.Normal("log_sigma", np.log(60.0), 1.0)
        flux = sum(add_planet(p, W, rho, u1, u2, **kw.get(p, {})) for p in planets)
        pm.Potential("likelihood", marginal_loglike(W, flux, log_sigma))
    return m


def summarise(idata, names):
    s = az.summary(idata, var_names=names, round_to=6)
    ss = idata.sample_stats
    print(f"divergences {int(ss['diverging'].sum())}, mean tree depth {float(ss['depth'].mean()):.1f}, "
          f"max r_hat {s['r_hat'].max():.3f}, min bulk ESS {s['ess_bulk'].min():.0f}")
    return s.iloc[:, [0, 1, 2, 3, 4, 6]]