Seeing below the diffraction limit: Bayesian localisation, detection and counting in STORM microscopy

SQ2, SQ2PI = np.sqrt(2.0), np.sqrt(2 * np.pi)


def pixel_gauss(u, sig, n):
    """Pixel-integrated 1-D Gaussian at pixels 0..n-1 for centres u (...,) and widths sig (...,),
    with its derivatives with respect to u and sig. Returns three (..., n) arrays."""
    i = np.arange(n)
    lo = (i - 0.5 - u[..., None]) / sig[..., None]
    hi = (i + 0.5 - u[..., None]) / sig[..., None]
    E = 0.5 * (erf(hi / SQ2) - erf(lo / SQ2))
    glo, ghi = np.exp(-0.5 * lo**2) / SQ2PI, np.exp(-0.5 * hi**2) / SQ2PI
    return E, (glo - ghi) / sig[..., None], (lo * glo - hi * ghi) / sig[..., None]


def model_and_jacobian(th, K, P, jacobian=True):
    """Expected photons (B, P*P) for parameter rows th = [log b, log s, (x, y, log N) x K] and,
    optionally, the Jacobian (B, P*P, 2 + 3K)."""
    B = th.shape[0]
    b, s = np.exp(th[:, 0]), np.exp(th[:, 1])
    mu = np.repeat(b[:, None, None], P, 1).repeat(P, 2)
    J = np.zeros((B, P, P, 2 + 3 * K)) if jacobian else None
    if jacobian:
        J[..., 0] = b[:, None, None]
    for k in range(K):
        x, y, N = th[:, 2 + 3 * k], th[:, 3 + 3 * k], np.exp(th[:, 4 + 3 * k])
        Ex, dEx, sEx = pixel_gauss(x, s, P)
        Ey, dEy, sEy = pixel_gauss(y, s, P)
        img = Ey[:, :, None] * Ex[:, None, :]
        mu = mu + N[:, None, None] * img
        if jacobian:
            Nk = N[:, None, None]
            J[..., 2 + 3 * k] = Nk * Ey[:, :, None] * dEx[:, None, :]
            J[..., 3 + 3 * k] = Nk * dEy[:, :, None] * Ex[:, None, :]
            J[..., 4 + 3 * k] = Nk * img
            J[..., 1] += Nk * s[:, None, None] * (sEy[:, :, None] * Ex[:, None, :] + Ey[:, :, None] * sEx[:, None, :])
    mu = mu.reshape(B, -1)
    return (mu, J.reshape(B, P * P, -1)) if jacobian else mu


def simulate_emccd(photons, rng, gain=GAIN, offset=OFFSET, read_sd=5.0):
    """Expected photons per pixel -> ADU: Poisson photons, Gamma gain register, read noise, offset."""
    n = rng.poisson(photons)
    charge = np.where(n > 0, rng.gamma(np.maximum(n, 1), gain), 0.0)
    return np.round(charge + offset + rng.normal(0.0, read_sd, n.shape))


def to_counts(adu, offset=OFFSET, scale=SCALE):
    """ADU -> effective counts (Poisson-like, mean = photons / 2)."""
    return np.maximum(np.round((adu - offset) / scale), 0.0)


# a check of the "half the photons" claim on simulated pixels
lam = np.array([5.0, 50.0, 500.0])
zz = to_counts(simulate_emccd(np.repeat(lam[:, None], 200_000, 1), np.random.default_rng(1)))
for l_, row in zip(lam, zz):
    print(f"{l_:5.0f} photons: effective counts mean {row.mean():7.2f}, variance {row.var():7.2f} "
          f"(Poisson with half the photons: {l_ / 2:.1f})")

FOUR = np.array([[[0, 0, 0]] * 3, [[0, 1, 0], [1, 1, 1], [0, 1, 0]], [[0, 0, 0]] * 3])


def detect(z, bg, thr_sd=5.0):
    """Connected blobs of the smoothed, background-subtracted movie above thr_sd noise sds
    (4-connected within a frame, never across frames)."""
    S = ndi.gaussian_filter(z - bg, sigma=(0, 1.0, 1.0))
    thr = thr_sd * 1.4826 * np.median(np.abs(S - np.median(S)))
    lab, n = ndi.label(S > thr, structure=FOUR)
    return S, lab, n, thr


def blob_patches(z, lab, n, P=11):
    """A P x P patch around each blob, a mask that excludes pixels of *other* blobs (dilated by
    2 px), the patch origin (frame, row, col) and the blob centre in patch coordinates."""
    T, H, W = z.shape
    Z, Wm = np.zeros((n, P, P)), np.zeros((n, P, P))
    org, cen = np.zeros((n, 3), int), np.zeros((n, 2))
    st = ndi.generate_binary_structure(2, 1)
    for i, sl in enumerate(ndi.find_objects(lab)):
        t = sl[0].start
        cy, cx = (sl[1].start + sl[1].stop - 1) / 2, (sl[2].start + sl[2].stop - 1) / 2
        r0 = int(np.clip(round(cy) - P // 2, 0, H - P))
        c0 = int(np.clip(round(cx) - P // 2, 0, W - P))
        Z[i] = z[t, r0:r0 + P, c0:c0 + P]
        L = lab[t, r0:r0 + P, c0:c0 + P]
        Wm[i] = ~ndi.binary_dilation((L > 0) & (L != i + 1), st, iterations=2)
        org[i], cen[i] = (t, r0, c0), (cx - c0, cy - r0)
    return Z.reshape(n, -1), Wm.reshape(n, -1), org, cen


t0 = time.time()
S_real, lab_real, n_real, thr_real = detect(z_real, bg_real)
Z_real, W_real, org_real, cen_real = blob_patches(z_real, lab_real, n_real)
size = np.bincount(lab_real.ravel())[1:]
amp = np.array([S_real[sl].max() for sl in ndi.find_objects(lab_real)])
alone = (W_real.min(1) == 1) & (size > 12) & (size < 40)       # no other blob in its patch
pick = np.nonzero(alone & (amp > np.quantile(amp[alone], 0.45)) & (amp < np.quantile(amp[alone], 0.55)))[0][0]
print(f"{n_real} blobs in {T_REAL} frames ({time.time() - t0:.1f} s); fitting blob {pick} "
      f"(frame {org_real[pick, 0]})")

P = 11
PIX = np.arange(P)


def pixel_gauss_pt(u, s):
    return 0.5 * (pt.erf((PIX + 0.5 - u[..., None]) / (SQ2 * s)) - pt.erf((PIX - 0.5 - u[..., None]) / (SQ2 * s)))


def spot_model(z_patch, K, centre, logb_mu, N_mu=1500.0, pos_sd=2.0):
    """K emitters in one P x P patch: priors identical to the batched fitter below."""
    with pm.Model() as m:
        logb = pm.Normal("logb", logb_mu, 0.5)
        logs = pm.Normal("logs", 0.0, 0.2)
        x = pm.Normal("x", centre[0], pos_sd, shape=K)
        y = pm.Normal("y", centre[1], pos_sd, shape=K)
        logN = pm.Normal("logN", np.log(N_mu), 1.0, shape=K)
        s = pt.exp(logs)
        img = pt.sum(pt.exp(logN)[:, None, None] * pixel_gauss_pt(y, s)[:, :, None]
                     * pixel_gauss_pt(x, s)[:, None, :], axis=0)
        mu = pm.Deterministic("mu", pt.exp(logb) + img)
        pm.Poisson("z", mu.ravel() / 2, observed=z_patch)
    return m


LOGB_REAL = np.log(2 * np.median(bg_real))
z_spot = Z_real[pick]
with spot_model(z_spot, 1, cen_real[pick], LOGB_REAL) as m_spot:
    idata_spot = pm.sample(random_seed=RANDOM_SEED, progressbar=False, var_names=["logb", "logs", "x", "y", "logN"])
post = az.extract(idata_spot, var_names=["x", "y", "logN", "logb", "logs"])
print(f"sampler: nutpie, {idata_spot.posterior.attrs.get('tuning_steps', '?')} tuning steps; divergences: "
      f"{int(idata_spot.sample_stats['diverging'].sum())}")
az.summary(idata_spot, var_names=["x", "y", "logN", "logb", "logs"], round_to=3)

class Prior:
    """Independent normal priors: log b, log s, positions (x, y) ~ N(centre, pos_sd), log N."""

    def __init__(self, logb, logs=(0.0, 0.2), logN=(np.log(1500.0), 1.0), centre=None, pos_sd=2.0):
        self.logb, self.logs, self.logN, self.centre, self.pos_sd = logb, logs, logN, centre, pos_sd

    def mean_prec(self, K):
        c = np.zeros((1, 2)) if self.centre is None else np.atleast_2d(self.centre)
        B = c.shape[0]
        m = np.concatenate([np.full((B, 1), self.logb[0]), np.full((B, 1), self.logs[0])]
                           + [np.c_[c, np.full(B, self.logN[0])]] * K, axis=1)
        p = np.array([self.logb[1] ** -2, self.logs[1] ** -2]
                     + [self.pos_sd ** -2, self.pos_sd ** -2, self.logN[1] ** -2] * K)
        return m, p

    def log_norm(self, K):
        """Normalising constant of the Gaussian prior (needed for evidences)."""
        return (-0.5 * np.log(2 * np.pi * self.logb[1] ** 2) - 0.5 * np.log(2 * np.pi * self.logs[1] ** 2)
                + K * (-0.5 * np.log(2 * np.pi * self.logN[1] ** 2) - np.log(2 * np.pi * self.pos_sd ** 2)))


def info(Je, wt):
    """Fisher information J^T diag(wt) J for a batch, (B, D, D), without a (B, P*P, D, D) temporary."""
    return np.matmul(np.swapaxes(Je * wt[..., None], 1, 2), Je)


def log_lik(th, z, w, K, P, excess=2.0, lsq=False):
    m = model_and_jacobian(th, K, P, jacobian=False) / excess
    if lsq:
        return -0.5 * np.sum(w * (z - m) ** 2, axis=1)
    return np.sum(w * (z * np.log(m) - m - gammaln(z + 1)), axis=1)


def fit_batch(th0, z, K, P, prior=None, w=None, excess=2.0, lsq=False, iters=40):
    """Batched Levenberg-Marquardt (Fisher scoring) for B patches at once.
    Returns the optimum (B, D), its log posterior (B,), and the Fisher + prior precision (B, D, D)."""
    th = th0.copy()
    B, D = th.shape
    w = np.ones_like(z) if w is None else w
    pm_, pp = prior.mean_prec(K) if prior is not None else (np.zeros((1, D)), np.zeros(D))
    obj = lambda t: log_lik(t, z, w, K, P, excess, lsq) - 0.5 * np.sum(pp * (t - pm_) ** 2, axis=1)
    lp, lam = obj(th), np.full(B, 1e-2)

    def grad_info(t):
        mu, J = model_and_jacobian(t, K, P)
        m, Je = mu / excess, J / excess
        if lsq:
            return np.einsum("bpd,bp->bd", Je, w * (z - m)), info(Je, w)
        return np.einsum("bpd,bp->bd", Je, w * (z / m - 1.0)), info(Je, w / m)

    for _ in range(iters):
        g, I = grad_info(th)
        g = g - pp * (th - pm_)
        I = I + np.diag(pp)
        damp = lam[:, None, None] * (np.eye(D) * np.diagonal(I, axis1=1, axis2=2)[:, :, None] + 1e-9 * np.eye(D))
        new = th + np.clip(np.linalg.solve(I + damp, g[..., None])[..., 0], -1.0, 1.0)
        for k in range(K):                                    # keep emitters inside the patch
            new[:, 2 + 3 * k:4 + 3 * k] = np.clip(new[:, 2 + 3 * k:4 + 3 * k], -0.5, P - 0.5)
        lpn = obj(new)
        ok = np.isfinite(lpn) & (lpn >= lp)
        th[ok], lp[ok] = new[ok], lpn[ok]
        lam = np.clip(np.where(ok, lam * 0.3, lam * 10.0), 1e-6, 1e6)
    _, I = grad_info(th) if not lsq else (None, None)
    if lsq:                                                   # report the Poisson information anyway
        mu, J = model_and_jacobian(th, K, P)
        I = info(J / excess, w / (mu / excess))
    return th, lp, I + np.diag(pp)


def start_one(z, P, centre):
    """Crude starting values for one emitter per patch."""
    zz = z.reshape(len(z), -1)
    b0 = np.maximum(np.median(zz, 1), 0.5)
    N0 = np.maximum(2 * (zz.sum(1) - b0 * P * P), 100.0)
    return np.c_[np.log(2 * b0), np.zeros(len(z)), centre, np.log(N0)]


th0 = start_one(z_spot[None], P, cen_real[pick][None])
th_mle, _, I_mle = fit_batch(th0, z_spot[None], 1, P)
th_lsq, _, _ = fit_batch(th0, z_spot[None], 1, P, lsq=True)
prior_spot = Prior(logb=(LOGB_REAL, 0.5), centre=cen_real[pick][None])
th_map, _, I_map = fit_batch(th0, z_spot[None], 1, P, prior_spot)
sd_lap = np.sqrt(np.diag(np.linalg.inv(I_map[0])))
print(f"NUTS posterior : x = {np.mean(xs) * PX:6.1f} +/- {np.std(xs) * PX:4.1f} nm, y = {np.mean(ys) * PX:6.1f} +/- {np.std(ys) * PX:4.1f} nm")
print(f"Laplace (MAP)  : x = {th_map[0, 2] * PX:6.1f} +/- {sd_lap[2] * PX:4.1f} nm, y = {th_map[0, 3] * PX:6.1f} +/- {sd_lap[3] * PX:4.1f} nm")
print(f"MLE            : x = {th_mle[0, 2] * PX:6.1f}, y = {th_mle[0, 3] * PX:6.1f} nm")
print(f"least squares  : x = {th_lsq[0, 2] * PX:6.1f}, y = {th_lsq[0, 3] * PX:6.1f} nm")

th_pm = np.array([[np.mean(post["logb"]), np.mean(post["logs"]), np.mean(xs), np.mean(ys), np.mean(post["logN"])]])
m_pm = model_and_jacobian(th_pm, 1, P, jacobian=False)[0] / 2
resid = (z_spot - m_pm) / np.sqrt(m_pm)

fig, axs = plt.subplots(1, 4, figsize=(15, 3.8))
axs[0].imshow(2 * z_spot.reshape(P, P), cmap="magma")
axs[0].set(title="data (photons)", xticks=[], yticks=[])
axs[1].imshow(2 * m_pm.reshape(P, P), cmap="magma")
axs[1].set(title="posterior-mean model", xticks=[], yticks=[])
im = axs[2].imshow(resid.reshape(P, P), cmap="RdBu_r", vmin=-3, vmax=3)
axs[2].set(title=f"standardised residual\n(sd {resid.std():.2f})", xticks=[], yticks=[])
fig.colorbar(im, ax=axs[2], shrink=0.8)
ax = axs[3]
ax.scatter((xs - np.mean(xs)) * PX, (ys - np.mean(ys)) * PX, s=2, alpha=0.2, color=BLUE, label="NUTS draws")
for thv, c, lab_ in [(th_mle, INK, "MLE"), (th_lsq, ORANGE, "least squares"), (th_map, AQUA, "MAP")]:
    ax.scatter((thv[0, 2] - np.mean(xs)) * PX, (thv[0, 3] - np.mean(ys)) * PX, s=60, marker="x", color=c, label=lab_)
tt_ = np.linspace(0, 2 * np.pi, 100)
ax.plot(2 * sd_lap[2] * PX * np.cos(tt_), 2 * sd_lap[3] * PX * np.sin(tt_), color=AQUA, lw=1.5, label="Laplace 2 sd")
ax.set(xlabel="x - posterior mean (nm)", ylabel="y - posterior mean (nm)", aspect="equal", title="where is the molecule?")
ax.legend(fontsize=7, loc="upper right");

def safe_sd(I, j):
    """Laplace sd of parameter j; inf where the information matrix is numerically singular."""
    v = np.linalg.pinv(I)[:, j, j]
    return np.sqrt(np.where(v > 0, v, np.inf))


def mortensen(N, b, s=1.0, a=1.0, F2=2.0):
    sa2 = s**2 + a**2 / 12
    tau = 2 * np.pi * sa2 * b / (N * a**2)
    t = np.linspace(1e-9, 1, 4001)
    integ = np.array([np.trapezoid(np.log(t) / (1 + t / ta), t) for ta in np.atleast_1d(tau)])
    return (np.atleast_1d(np.sqrt(F2 * sa2 / N / (1 + integ))),
            np.atleast_1d(np.sqrt(F2 * sa2 / N * (16 / 9 + 4 * tau))))


SIG_TRUE = 1.05                                              # PSF sd in pixels (133 nm)
N_GRID = np.array([200, 400, 800, 1500, 3000, 6000, 12000])
prec_rows = []
rng4 = np.random.default_rng(4)
t0 = time.time()
for bgp in [10.0, 110.0]:
    for N in N_GRID:
        B = 1000
        x, y = rng4.uniform(4.5, 5.5, B), rng4.uniform(4.5, 5.5, B)
        tt = np.c_[np.full(B, np.log(bgp)), np.full(B, np.log(SIG_TRUE)), x, y, np.full(B, np.log(N))]
        adu = simulate_emccd(model_and_jacobian(tt, 1, P, jacobian=False), rng4)
        z = to_counts(adu)
        z_naive = np.maximum(np.round((adu - OFFSET) / GAIN), 0.0)   # "photons", ignoring excess noise
        th0 = start_one(z, P, np.full((B, 2), 5.0))
        th_m, _, I_m = fit_batch(th0, z, 1, P)
        th_l, _, _ = fit_batch(th0, z, 1, P, lsq=True)
        th_n, _, I_n = fit_batch(start_one(z_naive / 2, P, np.full((B, 2), 5.0)), z_naive, 1, P, excess=1.0)
        _, Jt = model_and_jacobian(tt, 1, P)
        crlb = np.sqrt(np.mean(np.linalg.inv(info(Jt / 2, 2 / model_and_jacobian(tt, 1, P, False)))[:, 2, 2]))
        sd_m, sd_n = safe_sd(I_m, 2), safe_sd(I_n, 2)
        mo_mle, mo_lsq = mortensen(N, bgp, SIG_TRUE)
        prec_rows.append(dict(bg=bgp, N=N, mle=np.sqrt(np.mean((th_m[:, 2] - x) ** 2)),
                              lsq=np.sqrt(np.mean((th_l[:, 2] - x) ** 2)), crlb=crlb,
                              mort_mle=mo_mle[0], mort_lsq=mo_lsq[0],
                              cover=np.mean(np.abs(th_m[:, 2] - x) < 1.645 * sd_m),
                              cover_naive=np.mean(np.abs(th_n[:, 2] - x) < 1.645 * sd_n)))
print(f"{len(prec_rows) * 1000 * 3} fits in {time.time() - t0:.1f} s")
print(" bg      N   MLE  LSQ  CRLB  Mortensen(MLE, LSQ)   90% coverage: EMCCD model / ignoring excess noise")
for r in prec_rows:
    print(f"{r['bg']:4.0f} {r['N']:6d}  {r['mle'] * PX:5.1f} {r['lsq'] * PX:5.1f} {r['crlb'] * PX:5.1f}   "
          f"{r['mort_mle'] * PX:5.1f} {r['mort_lsq'] * PX:5.1f}            {r['cover']:.2f} / {r['cover_naive']:.2f}")

def count_emitters(z, P, prior, Kmax=3, w=None, iters=40):
    """MAP fits and Laplace log evidences for K = 0..Kmax emitters in every patch.
    Returns a list over K of (theta (B, 2+3K), log evidence (B,), precision (B, D, D))."""
    B = z.shape[0]
    w = np.ones_like(z) if w is None else w
    out = []
    th = np.c_[np.log(np.maximum(2 * np.median(z, 1), 1.0)), np.full(B, prior.logs[0])]
    for K in range(Kmax + 1):
        if K == 0:
            starts = [th]
        else:
            prev = out[-1][0]
            resid = ((z - model_and_jacobian(prev, K - 1, P, False) / 2) * w).reshape(B, P, P)
            rs = ndi.uniform_filter(resid, size=(1, 3, 3)).reshape(B, -1)
            yy, xx = np.divmod(rs.argmax(1), P)
            new = np.c_[prev, xx, yy, np.log(np.maximum(36 * rs.max(1), 100.0))]
            starts = [new]
            if K >= 2:                                        # or split the brightest emitter
                split = new.copy()
                j = 2 + 3 * prev[:, 4::3].argmax(1)
                rows = np.arange(B)
                split[rows, j] -= 0.5
                split[:, -3] = prev[rows, j] + 0.5
                split[:, -2] = prev[rows, j + 1]
                split[rows, j + 2] -= np.log(2)
                split[:, -1] = prev[rows, j + 2] - np.log(2)
                starts.append(split)
        best = None
        for s0 in starts:
            fitK = fit_batch(s0, z, K, P, prior, w=w, iters=iters)
            if best is None:
                best = [a.copy() for a in fitK]
            else:
                better = fitK[1] > best[1]
                for a, b_ in zip(best, fitK):
                    a[better] = b_[better]
        thK, lp, I = best
        logZ = lp + prior.log_norm(K) + 0.5 * (2 + 3 * K) * np.log(2 * np.pi) \
            - 0.5 * np.linalg.slogdet(I)[1] + gammaln(K + 1)
        # the log-likelihood in lp omits nothing: log_lik includes gammaln(z + 1)
        out.append((thK, logZ, I))
    return out


def two_emitter_patches(sep, B, rng, N_med=1500.0, bgp=110.0):
    """B patches with two emitters `sep` pixels apart around the centre, random orientation."""
    c = 5 + rng.uniform(-0.5, 0.5, (B, 2))
    a = rng.uniform(0, np.pi, B)
    d = 0.5 * sep * np.c_[np.cos(a), np.sin(a)]
    N1, N2 = rng.lognormal(np.log(N_med), 0.4, (2, B))
    tt = np.c_[np.full(B, np.log(bgp)), np.full(B, np.log(SIG_TRUE)), c - d, np.log(N1), c + d, np.log(N2)]
    return to_counts(simulate_emccd(model_and_jacobian(tt, 2, P, False), rng)), tt


def last_finite(obj):
    """Last finite entry of a (possibly ragged, nested) per-chain SMC statistic."""
    flat = lambda o: [v for x in o for v in (flat(x) if isinstance(x, (list, tuple, np.ndarray)) else [x])]
    vals = np.array(flat(obj if isinstance(obj, (list, tuple, np.ndarray)) else [obj]), dtype=float)
    return vals[np.isfinite(vals)][-1]


rng5 = np.random.default_rng(5)
z_pair, tt_pair = two_emitter_patches(1.5, 1, rng5)
prior_sim = Prior(logb=(np.log(110.0), 0.5), centre=np.array([[5.0, 5.0]]))
lap = count_emitters(z_pair, P, prior_sim)
smc_logZ = {}
t0 = time.time()
for K in [1, 2, 3]:
    with spot_model(z_pair[0], K, (5.0, 5.0), np.log(110.0)):
        idata_smc = pm.sample_smc(draws=2000, chains=4, cores=1, random_seed=RANDOM_SEED, progressbar=False,
                                  compute_convergence_checks=False)
    lml = idata_smc.sample_stats["log_marginal_likelihood"].values
    smc_logZ[K] = [last_finite(c) for c in lml]
    if K == 2:
        smc2 = az.extract(idata_smc, var_names=["x", "y"])
    del idata_smc
print(f"SMC: {time.time() - t0:.0f} s")
print(" K   Laplace log evidence   SMC log evidence (4 chains)")
for K in [1, 2, 3]:
    print(f" {K}   {lap[K][1][0]:10.1f}            " + ", ".join(f"{v:.1f}" for v in smc_logZ[K]))

H_SIM = 40
rng6 = np.random.default_rng(6)
pores = []
while len(pores) < 48:
    c = rng6.uniform(4, H_SIM - 4, 2)
    if all(np.hypot(*(c - q)) > 300 / PX for q in pores):
        pores.append(c)
pores = np.array(pores)
R_NPC = 55 / PX
mol = []
for i, c in enumerate(pores):
    rot = rng6.uniform(0, 2 * np.pi / 8)
    for k in range(8):
        for j in range(4):
            a = rot + 2 * np.pi * k / 8 + rng6.normal(0, 0.05)
            mol.append((i, c[0] + R_NPC * np.cos(a) + rng6.normal(0, 3 / PX), c[1] + R_NPC * np.sin(a) + rng6.normal(0, 3 / PX)))
mol = np.array(mol)
labelled = rng6.random(len(mol)) < 0.6
mols = mol[labelled]
M_TRUE = np.bincount(mols[:, 0].astype(int), minlength=len(pores))
MEAN_EXTRA = 4.0
blinks = 1 + rng6.geometric(1 / (1 + MEAN_EXTRA), len(mols)) - 1
ev_mol = np.repeat(np.arange(len(mols)), blinks)
ev_N = rng6.lognormal(np.log(1500), 0.6, len(ev_mol))
print(f"{len(pores)} pores, {len(mols)} labelled molecules (true copies per pore {M_TRUE.min()}-{M_TRUE.max()}), "
      f"{len(ev_mol)} blinks")


def make_movie(T, rng):
    ev_t = rng.integers(0, T, len(ev_mol))
    mu = np.full((T, H_SIM, H_SIM), 110.0)
    s = np.full(len(ev_mol), SIG_TRUE)
    Ex, Ey = pixel_gauss(mols[ev_mol, 1], s, H_SIM)[0], pixel_gauss(mols[ev_mol, 2], s, H_SIM)[0]
    np.add.at(mu, ev_t, ev_N[:, None, None] * Ey[:, :, None] * Ex[:, None, :])
    return to_counts(simulate_emccd(mu, rng)), ev_t


def match(t_true, xy_true, t_est, xy_est, tol=1.0):
    """Recall, precision, RMSE (nm) and the matched true-event index of every estimate (-1: false)."""
    owner = np.full(len(t_est), -1)
    err = []
    for t in np.intersect1d(t_true, t_est):
        a, b_ = np.nonzero(t_true == t)[0], np.nonzero(t_est == t)[0]
        D = np.hypot(xy_true[a, None, 0] - xy_est[None, b_, 0], xy_true[a, None, 1] - xy_est[None, b_, 1])
        r, c = linear_sum_assignment(D)
        ok = D[r, c] < tol
        owner[b_[c[ok]]] = a[r[ok]]
        err += list(D[r, c][ok])
    tp = (owner >= 0).sum()
    return tp / len(t_true), tp / max(len(t_est), 1), np.sqrt(np.mean(np.square(err))) * PX, owner


def threshold_pipeline(z, S, lab, thr, P7=7):
    T, H, W = z.shape
    t_, r_, c_ = np.nonzero((S == ndi.maximum_filter(S, size=(1, 3, 3))) & (S > thr) & (lab > 0))
    r0 = np.clip(r_ - P7 // 2, 0, H - P7)
    c0 = np.clip(c_ - P7 // 2, 0, W - P7)
    Zs = np.stack([z[t, a:a + P7, b:b + P7] for t, a, b in zip(t_, r0, c0)]).reshape(len(t_), -1)
    th, _, I = fit_batch(start_one(Zs, P7, np.c_[c_ - c0, r_ - r0]), Zs, 1, P7)
    sd = np.sqrt(0.5 * safe_sd(I, 2) ** 2 + 0.5 * safe_sd(I, 3) ** 2)
    return t_, np.c_[th[:, 2] + c0, th[:, 3] + r0], sd, np.exp(th[:, 4])


def bayes_pipeline(z, lab, n, logb, Kmax=4):
    Z, Wm, org, cen = blob_patches(z, lab, n)
    res = count_emitters(Z, P, Prior(logb=(logb, 0.5), centre=cen), Kmax=Kmax, w=Wm)
    logZ = np.array([r[1] for r in res]).T
    K = logZ.argmax(1)
    t_, xy, sd, Nph, Kof, blob = [], [], [], [], [], []
    for k_ in range(1, Kmax + 1):
        s = K == k_
        th, I = res[k_][0][s], res[k_][2][s]
        for j in range(k_):
            t_.append(org[s, 0])
            xy.append(np.c_[th[:, 2 + 3 * j] + org[s, 2], th[:, 3 + 3 * j] + org[s, 1]])
            sd.append(np.sqrt(0.5 * safe_sd(I, 2 + 3 * j) ** 2 + 0.5 * safe_sd(I, 3 + 3 * j) ** 2))
            Nph.append(np.exp(th[:, 4 + 3 * j]))
            Kof.append(np.full(s.sum(), k_))
            blob.append(np.nonzero(s)[0])
    cat = lambda v: np.concatenate(v)
    return cat(t_), np.vstack(xy), cat(sd), cat(Nph), K, cat(Kof), cat(blob)


movies = {}
t0 = time.time()
for name, T in [("sparse", 2000), ("dense", 250)]:
    z, ev_t = make_movie(T, rng6)
    bg = np.median(z, axis=0)
    S, lab, n, thr = detect(z, bg)
    xy_true = mols[ev_mol, 1:3]
    thr_res = threshold_pipeline(z, S, lab, thr)
    bay_res = bayes_pipeline(z, lab, n, np.log(2 * np.median(bg)))
    m_thr = match(ev_t, xy_true, thr_res[0], thr_res[1])
    m_bay = match(ev_t, xy_true, bay_res[0], bay_res[1])
    movies[name] = dict(T=T, ev_t=ev_t, thr=thr_res, bay=bay_res, m_thr=m_thr, m_bay=m_bay, n_blobs=n,
                        frame=z[int(np.argmax([np.sum(ev_t == t) for t in range(T)]))] if name == "dense" else z[0],
                        frame_t=int(np.argmax([np.sum(ev_t == t) for t in range(T)])) if name == "dense" else 0)
    print(f"{name:6s} ({T} frames, {len(ev_mol) / T:.1f} blinks/frame): {n} blobs; Bayesian K counts "
          f"{np.bincount(bay_res[4], minlength=5)}")
    for lab_, mm, res_ in [("thresholding", m_thr, thr_res), ("Bayesian", m_bay, bay_res)]:
        ok = mm[3] >= 0
        err = res_[1][ok] - xy_true[mm[3][ok]]
        cov90 = np.mean(np.abs(err) < 1.645 * res_[2][ok, None])
        print(f"   {lab_:12s}: {len(res_[0]):5d} localisations, recall {mm[0]:.3f}, precision {mm[1]:.3f}; "
              f"matched: RMSE {mm[2]:.1f} nm, median error {np.median(np.hypot(*err.T)) * PX:.1f} nm, "
              f"90% intervals (per axis) cover {cov90:.2f}")
    del z, S, lab
print(f"pipelines: {time.time() - t0:.0f} s")