Every photon: a neural localiser tried on a full dSTORM movie, and how it failed

import os

for name in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS"):
    os.environ.setdefault(name, "1")  # before NumPy: one BLAS thread per chain

import logging
import sys
import time

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import nutpie
import pandas as pd
import pymc as pm
import pytensor.tensor as pt
from IPython.display import HTML, Markdown, display
from matplotlib import animation
from matplotlib.colors import LogNorm, TwoSlopeNorm
from matplotlib.lines import Line2D
from matplotlib.patches import Patch, Rectangle
from matplotlib.ticker import NullFormatter
from scipy import ndimage as ndi
from scipy import stats

from pymc_challenges import data
from pymc_challenges.bigrun import PreflightError, Store, banner, estimated, git_hash, measured

# ArviZ leaves a lazy stub under "jax", and Matplotlib's first plot tests `sys.modules.get("jax")`, which
# would import all of JAX for nothing. Nothing here uses JAX, so the stub goes.
sys.modules.pop("jax", None)

RANDOM_SEED = 88
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
logging.getLogger("pymc").setLevel(logging.WARNING)

LEVEL = os.environ.get("PYMC_CHALLENGES_RUN", "").strip().lower() or "replay"
if LEVEL != "replay":
    raise PreflightError(f"PYMC_CHALLENGES_RUN={LEVEL!r}: E88 has one level, replay. The heavy run is not a notebook "
                         "cell: it is tools/run_e88_full.py, and it is not for limited hardware.")

CORES = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count()
PIXEL_NM, CYCLE_MS, FRAMES, SIDE, PATCH = 127.0, 11.742, 15000, 200, 11
RECS = {"r7": "recording 7", "r32": "recording 32"}
GRID = (2000, 3000, 4000, 5000, 6000, 7500, 9000, 10500, 12000, 13500, 15000)
CLASSICAL, NET_A, NET_B, SINGLE, MUTED, LINE_COL = "#0b6e8a", "#b2182b", "#d98200", "#5f6b73", "#9a9a9a", "#111111"
COLOUR = {"chosen": CLASSICAL, "K1": SINGLE, "a": NET_A, "b": NET_B}
NAME = {"chosen": "classical, chosen K", "K1": "classical, K = 1", "a": "network A", "b": "network B"}
DARK = plt.get_cmap("inferno").with_extremes(bad="black")



def markdown_table(frame):
    """A DataFrame as a Markdown table: pandas needs `tabulate` for this, and it is not a dependency here."""
    lines = ["| " + " | ".join(frame.columns) + " |", "|" + "---|" * frame.shape[1]]
    lines += ["| " + " | ".join(str(cell).replace("|", "/") for cell in row) + " |" for row in frame.to_numpy()]
    return Markdown("\n".join(lines))


store = Store("E88")
data.describe("dannhauser2022_e88_summary")
pack = np.load(data.path("dannhauser2022_e88_summary"), allow_pickle=False)
print(f"\nNotice carried in the file: {pack['attribution']}\n")
summary = store.load_json("summary")
certificate = store.load_json("certificate")
paired = store.load_json("paired_report")
halo = store.load_json("halo_report")
replicate = store.load_json("replicate_report")
audit = store.load_json("audit_nuts_report")
overlap_posthoc = store.load_json("p2b_overlap_posthoc")
posthoc = {name: store.load_json(f"posthoc_{name}") for name in ("sim", "real", "halo", "psf", "frames", "stuck", "nuts_cold", "nuts_warm")}
decisions = {"classical": store.load_json("p2b_decision"), "classical_line": store.load_json("p2b_decision_line")}
for net in ("a", "b"):
    for tail in ("", "_line", "_cdrift", "_line_cdrift"):
        decisions[net + tail] = store.load_json(f"decision_{net}{tail}")
started = time.perf_counter()
print(f"{CORES} cores here. Level: {LEVEL}. Written outcome of the stored run: {summary['outcome']}.")

seconds = FRAMES * CYCLE_MS / 1000
print(f"One movie: {FRAMES:,} frames x {CYCLE_MS} ms = {seconds:.1f} s on the camera (TIFF tags, measured).")
print(f"45 junctions per channel: {45 * seconds / 3600:.1f} h of scope time and {45 * 1.2098:.0f} GB "
      "(the 45 is the paper's 22 + 23 NMJs, one recording each is our assumption).")
for n in (10500, 12000):
    print(f"  At {n:,} frames: saves {45 * seconds * (1 - n / FRAMES) / 60:.0f} min and "
          f"{45 * 1.2098 * (1 - n / FRAMES):.0f} GB per channel [extrapolated, same assumption]")

frames_shown = pack["show_frames"]
fig = plt.figure(figsize=(13, 6.6), layout="constrained")
grid = fig.add_gridspec(2, 4, width_ratios=[1, 1, 2.05, 2.05])
vmin, vmax = np.percentile(pack["r7_frames"], [1, 99.9])
for row, f in enumerate(frames_shown):
    for col, rec in enumerate(RECS):
        ax = fig.add_subplot(grid[row, col])
        ax.imshow(pack[f"{rec}_frames"][row], cmap="gray", vmin=vmin, vmax=vmax, interpolation="nearest")
        ax.set_title(f"{RECS[rec]}, frame {f:,}", fontsize=9)
        ax.axis("off")
for col, rec in enumerate(RECS):
    ax = fig.add_subplot(grid[:, 2 + col])
    render = np.ma.masked_equal(pack[f"{rec}_render_chosen"].astype(float), 0)
    ax.imshow(render, cmap=DARK, norm=LogNorm(1, 60), extent=(-0.5, 199.5, 199.5, -0.5), interpolation="nearest")
    ax.contour(pack[f"{rec}_mask"], levels=[0.5], colors="#9fe6ff", linewidths=0.8)
    if rec == "r7":
        ax.add_patch(Rectangle((49.5, 24.5), 40, 40, fill=False, ec="#55ff99", lw=1.3))
        ax.text(51, 22, "E58's crop", color="#55ff99", fontsize=9)
    ax.plot([155, 155 + 5000 / PIXEL_NM], [190, 190], color="w", lw=3)
    ax.text(155 + 2500 / PIXEL_NM, 186, "5 µm", color="w", ha="center", fontsize=9)
    ax.set_title(f"{RECS[rec]}: classical localisations and the mask", fontsize=10)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_facecolor("black");

p3 = store.load_json("p3_calibration")
line = store.load_json("p3_line")
print(f"{'':14}{'offset moves':>14}{'slope moves':>14}{'line':>8}")
for rec in RECS:
    row = p3["recordings"][rec]
    print(f"{RECS[rec]:14}{row['offset_dev_max']:>10.0f} ADU{row['slope_rel_max']:>14.0%}{'20 ADU, 10%':>13}")
print(f"Between the two: {p3['between']['offset']:.0f} ADU, {p3['between']['slope_rel']:.0%}. Verdict: NO-GO, "
      "so the written action applies: calibrate per block and per recording.")
print("Post-hoc, one line per recording across blocks: " + "; ".join(
    f"{RECS[r]} slope {line[r]['slope']:.2f}, offset {line[r]['offset']:.0f}, r {line[r]['r']:.4f}" for r in RECS)
    + f". E58: slope {p3['e58']['slope']}, offset {p3['e58']['offset']:.0f}")

fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.2), layout="constrained")
for rec, marker in zip(RECS, ("o", "s")):
    row = p3["recordings"][rec]
    slopes, offsets = np.array(row["slopes"]), np.array(row["offsets"])
    dots = axes[0].scatter(slopes, offsets, c=frames_at, cmap="viridis", marker=marker, s=34, label=RECS[rec], lw=0)
    medians = pack[f"{rec}_median_per_100"].reshape(30, 5).mean(1)
    axes[1].plot(frames_at, (medians - offsets) / slopes, marker=marker, ms=3.5, color=CLASSICAL, lw=1.2,
                 label=f"{RECS[rec]}, per-block fit")
    axes[1].plot(frames_at, (medians - line[rec]["offset"]) / line[rec]["slope"], marker=marker, ms=3.5, color=NET_B,
                 lw=1.2, label=f"{RECS[rec]}, one line")
    axes[0].plot(line[rec]["slope"], line[rec]["offset"], marker=marker, ms=13, mfc="none", mec="k", mew=1.6)
    axes[0].annotate("post-hoc line", (line[rec]["slope"], line[rec]["offset"]), xytext=(8, -14),
                     textcoords="offset points", fontsize=8.5)
axes[0].plot(p3["e58"]["slope"], p3["e58"]["offset"], "*", ms=15, color=NET_A)
axes[0].annotate("E58", (p3["e58"]["slope"], p3["e58"]["offset"]), xytext=(9, 3), textcoords="offset points", color=NET_A)
fig.colorbar(dots, ax=axes[0], label="frame", shrink=0.8)
axes[0].legend(loc="upper left", fontsize=9)
axes[0].set(xlabel="fitted gain (ADU per count)", ylabel="fitted offset (ADU)")
axes[0].set_title("Block fits slide along one ridge as the movie runs", fontsize=11)
axes[1].legend(fontsize=8, loc="upper left", ncols=2)
axes[1].set(xlabel="frame", ylabel="background (counts per px)", ylim=(None, 86))
axes[1].set_title("Taken literally, the block fits make the background rise", fontsize=11);

def n_star(curve, allowance=0.10):
    """Smallest N on the grid where FRC, count and band agree with 15,000, there and at every larger N."""
    points, bands = curve["points"], curve["bands"]
    full, top = points["15000"], bands["15000"]

    def overlap(q, r):
        return q[0] <= r[1] and r[0] <= q[1]

    answer = FRAMES
    for N in reversed(GRID):
        p, q = points[str(N)], bands[str(N)]
        fine = (abs(p["frc_nm"] / full["frc_nm"] - 1) <= allowance and abs(p["count"] / full["count"] - 1) <= allowance
                and overlap(q["frc_q"], top["frc_q"]) and overlap(q["count_q"], top["count_q"]))
        if not fine:
            break
        answer = N
    return answer


def n_both(name, allowance=0.10):
    return max(n_star(decisions[name][rec], allowance) for rec in RECS)


for name in decisions:
    for rec in RECS:
        assert n_star(decisions[name][rec]) == decisions[name][rec]["N_star"], (name, rec)
print("Live rule matches every stored N*, across", len(decisions), "curve sets and both recordings.\n")
rows = []
for N in (2000, 5000, 7500, 9000, 10500, 12000, 15000):
    row = {"frames": N}
    for rec in RECS:
        p = decisions["classical"][rec]["points"][str(N)]
        c = decisions["classical"][rec]["checks"][str(N)]
        row[f"{rec} FRC nm"], row[f"{rec} clusters"], row[f"{rec} all three"] = round(p["frc_nm"], 1), int(p["count"]), c["all"]
    rows.append(row)
print(pd.DataFrame(rows).to_string(index=False))
print("\nN*, recording 7 / recording 32 / both, classical:")
print(f"  per block (written action, gated):  {decisions['classical']['r7']['N_star']:,} / "
      f"{decisions['classical']['r32']['N_star']:,} / {n_both('classical'):,}")
print(f"  one line (post-hoc sensitivity):    {decisions['classical_line']['r7']['N_star']:,} / "
      f"{decisions['classical_line']['r32']['N_star']:,} / {n_both('classical_line'):,}")
print("\nWhat the 10% allowance does to N for both, per block, then line:")
for allowance in (0.05, 0.075, 0.10, 0.125, 0.15, 0.20):
    print(f"  {allowance:>5.1%}: {n_both('classical', allowance):>6,} {n_both('classical_line', allowance):>7,}")

frames_boot = posthoc["frames"]
print("\nPost-hoc, count rule only, classical, 100 block-bootstrap draws (200-frame blocks):")
for key, title in (("per_block", "per block"), ("line", "one line")):
    row = frames_boot["calibrations"][key]
    low, high = row["both_95"]
    print(f"  {title:10} N for both, 95% range {low:,.0f}-{high:,.0f}. Draws at each N: " + ", ".join(f"{int(n):,}: {k}" for n, k in row["both_counts"].items())
          + f". Recording 32 alone: {min(row['r32']['stars']):,}-{max(row['r32']['stars']):,}")
margins = {int(k): v["margin"] for k, v in frames_boot["band_seeds"].items()}
print("Band overlap at 10,500 of recording 7's count, in clusters, stored seed then five fresh families: "
      + ", ".join(f"{margins[k]:+.1f}" for k in (88, 1, 2, 3, 4, 5)) + ". A negative margin fails the test. FRC's band was not repeated.")
row = frames_boot["calibrations"]["per_block"]["r7"]
print(f"Recording 7, first N frames:           {row['points']}")
print(f"Same event count, thinned from all 15,000: {row['thinned']} (one thinning at each N)")
print(f"10,500 frames hold {row['events'][7] / row['events'][-1]:.0%} of recording 7's events")

over = {rec: decisions["classical"][rec]["overlap"] for rec in RECS}
rows = [("Written primary: 4-8 px annulus, scaled", [decisions["classical"]["overlap_primary"][r] for r in RECS], "gated", CLASSICAL),
        ("Detector, direct (merging hides most)", [over[r]["all"]["detector_under_4"] for r in RECS], "lower bound", SINGLE),
        ("Poisson at mask density", [over[r]["poisson_basis"]["share_under_4"] for r in RECS], "ignores clustering", SINGLE),
        ("Neighbour 1.5-4 px, frame 200 away", [overlap_posthoc[r]["gap200"]["px_1p5_to_4"] for r in RECS], "post-hoc", NET_B),
        ("Neighbour under 4 px, frame 200 away", [overlap_posthoc[r]["gap200"]["under_4"] for r in RECS], "post-hoc, upper side", NET_B),
        ("Two bright emitters in a fit", [over[r]["share_two_bright_1_4px"] for r in RECS], "inflated by misfit", MUTED),
        ("Fit chose K of 2 or more", [over[r]["share_K2plus"] for r in RECS], "inflated by misfit", MUTED)]
print(f"{'':42}{'rec 7':>8}{'rec 32':>8}  status")
for label, values, status, _ in rows:
    print(f"{label:42}{values[0]:>8.2%}{values[1]:>8.2%}  {status}")
dense = 1 - np.exp(-0.70 * np.pi * (4 * PIXEL_NM / 1000) ** 2)
print(f"\nE58's dense simulation, same Poisson sum at 0.70 per um^2: {dense:.0%} [extrapolated]")
assert decisions["classical"]["overlap_rare"], "the written verdict"
print("Written verdict: RARE, so the stop before FULL fired.")

fig, ax = plt.subplots(figsize=(11.5, 3.9), layout="constrained")
for i, (label, values, status, colour) in enumerate(rows[::-1]):
    ax.plot(values, [i, i], color=colour, lw=2.6, zorder=2)
    ax.scatter(values[0], i, s=60, color=colour, zorder=3)
    ax.scatter(values[1], i, s=60, facecolor="w", edgecolor=colour, lw=1.8, zorder=3)
ax.set_yticks(range(len(rows)), [r[0] for r in rows[::-1]], fontsize=9.5)
ax.axvline(0.05, color=LINE_COL, lw=1.2)
ax.text(0.052, 0.45, "gate line, 5%", fontsize=9)
ax.axvline(dense, color=NET_A, lw=1.2, ls="--")
ax.text(dense * 0.97, 5.45, "E58's dense case", fontsize=9, color=NET_A, ha="right")
ax.set_xscale("log")
ax.set_xticks([0.003, 0.01, 0.03, 0.1, 0.3, 1], ["0.3%", "1%", "3%", "10%", "30%", "100%"])
ax.set(xlim=(0.003, 1))
ax.set_xlabel("share of in-mask blobs with a neighbour under 4 px (filled: recording 7, open: recording 32)", fontsize=10)
ax.set_title("The written measure says rare, a better one says just over the line, and either is far from dense", fontsize=11);

train, infer = summary["train"], summary["infer"]
p4_gpu, p4_cpu = store.load_json("p4_gpu"), store.load_json("p4_cpu")
inference_s = sum(row["predict_seconds"] for net in infer.values() for row in net.values())
rates = sorted(row["frames_per_s"] for net in infer.values() for row in net.values())
classical_s = [store.load_json(f"classical_{r}_timing")["seconds"] for r in ("r7", "r32", "line_r7", "line_r32")]
ratio = p4_gpu["steps_per_s"] / p4_cpu["steps_per_s"]
print(f"Network A: {train['a']['steps']:,} steps in {train['a']['seconds']:.0f} s, {train['a']['steps_per_s']:.1f} steps/s, "
      f"{train['a']['vram']['peak_gb']:.2f} GB VRAM, {train['a']['weights']:,} weights, deterministic ops")
print(f"Network B: {train['b']['steps']:,} steps in {train['b']['seconds']:.0f} s, {train['b']['steps_per_s']:.1f} steps/s, "
      f"{train['b']['vram']['peak_gb']:.2f} GB VRAM, warm start from A")
print(f"Inference: 8 movies, 120,000 frames, {inference_s:.0f} s of prediction at {rates[0]:.0f}-{rates[-1]:.0f} frames/s, "
      f"{max(r['vram']['peak_gb'] for n in infer.values() for r in n.values()):.2f} GB VRAM")
print(f"Pilot P4, without deterministic ops: GPU {p4_gpu['steps_per_s']:.0f} steps/s on {p4_gpu['vram']['peak_gb']:.2f} GB, "
      f"CPU {p4_cpu['steps_per_s']:.2f} steps/s on 13 CPUs: {ratio:.0f}x. A pilot rate: the planned full CPU comparison was skipped.")
print(f"Classical: 4 runs in {sum(classical_s) / 60:.1f} min on 13 CPUs (one recording-7 run covers 28 of 30 blocks)")

bars = [("Train A, 40,000 steps, deterministic GPU", train["a"]["seconds"], False, NET_A),
        ("Train B, 12,000 steps, deterministic GPU", train["b"]["seconds"], False, NET_B),
        ("Train A at P4's GPU rate, not deterministic", 40000 / p4_gpu["steps_per_s"], True, NET_A),
        ("Train A at P4's CPU rate, 13 CPUs", 40000 / p4_cpu["steps_per_s"], True, NET_A),
        ("Use: both networks on both movies, both calibrations", inference_s, False, MUTED),
        ("Use: classical, recording 32 on 13 CPUs", classical_s[1], False, CLASSICAL),
        ("Camera time for one movie", seconds, False, SINGLE)]
fig, axes = plt.subplots(1, 2, figsize=(13, 3.9), width_ratios=[1.5, 1], layout="constrained")
ax = axes[0]
for i, (label, value, extra, colour) in enumerate(bars[::-1]):
    ax.barh(i, value, color="none" if extra else colour, edgecolor=colour, hatch="///" if extra else None, lw=1.3)
    ax.text(value * 1.08, i, f"{value / 60:.1f} min" if value < 3600 else f"{value / 3600:.1f} h", va="center", fontsize=9)
ax.set_yticks(range(len(bars)), [b[0] for b in bars[::-1]], fontsize=9.3)
ax.set_xscale("log")
ax.set(xlim=(40, 60000))
ax.set_xlabel("seconds (log). Hatched: extrapolated from a measured rate", fontsize=10)
ax.set_title("Training is one bill, use is another", fontsize=11)
ax = axes[1]
ax.plot(train["a"]["log_steps"], train["a"]["log"], color=NET_A, lw=1.5, label="network A")
ax.plot(np.array(train["b"]["log_steps"]) + 40000, train["b"]["log"], color=NET_B, lw=1.5, label="network B, warm start")
ax.set(xlabel="training step", ylabel="loss at the logged step", ylim=(-45, 15))
ax.set_title("A's and B's losses, on different simulators", fontsize=11)
ax.legend(loc="upper right");

p5 = {"5,000 steps (P5, NO-GO)": store.load_json("p5_p5"), "A, 40,000 steps": store.load_json("p5_a"),
      "B, 12,000 more": store.load_json("p5_b"), "small, 1,500 steps": store.load_json("p5_small")}
rows = []
for name, r in p5.items():
    sparse, bright, dense_run = r["sparse_2000"], r["sparse_e58_brightness"], r["dense_e58_brightness"]
    rows.append({"network": name, "recall": sparse["tol2.0"]["recall"], "x CRLB": sparse["rmse_over_crlb"],
                 "median sd px": sparse["median_sd_px"], "recall, E58 brightness": bright["tol2.0"]["recall"],
                 "x CRLB, E58 brightness": bright["rmse_over_crlb"], "recall, dense": dense_run["tol2.0"]["recall"],
                 "x CRLB, dense": dense_run["rmse_over_crlb"]})
print(pd.DataFrame(rows).round(3).to_string(index=False))
print(f"\nLines at 2,000 photons, sparse: recall 0.90, 1.3 x CRLB. CRLB here {p5['A, 40,000 steps']['crlb_px']:.3f} px, "
      f"{p5['A, 40,000 steps']['crlb_nm']:.1f} nm.")


def cells(net, world, field):
    table = [r for r in certificate["simulation"][net] if r["world"] == world]
    photons, densities = sorted({r["photons"] for r in table}), sorted({r["density"] for r in table})
    out = np.full((len(photons), len(densities)), np.nan)
    for r in table:
        out[photons.index(r["photons"]), densities.index(r["density"])] = r[field]
    return out, photons, densities


def near_nominal(net, world):
    table = [r for r in certificate["simulation"][net] if r["world"] == world and r["in_prior"]]
    return sum(0.90 <= r["coverage95"] <= 0.98 for r in table), len(table)


fig, axes = plt.subplots(2, 3, figsize=(13, 6.6), sharex=True, sharey=True, layout="constrained")
views = (("gauss", "coverage95", "95% coverage, simple world"), ("empirical", "coverage95", "95% coverage, empirical world"),
         ("empirical", "precision", "precision, empirical world"))
for row, net in enumerate(("a", "b")):
    for col, (world, field, title) in enumerate(views):
        ax = axes[row, col]
        values, photons, densities = cells(net, world, field)
        norm = TwoSlopeNorm(0.95, 0.40, 1.0)
        mesh = ax.pcolormesh(np.arange(6), np.arange(5), values, cmap="RdBu", norm=norm, edgecolor="w", lw=1.2)
        for i in range(4):
            for j in range(5):
                ax.text(j + 0.5, i + 0.5, f"{values[i, j]:.2f}", ha="center", va="center", fontsize=9,
                        color="w" if abs(values[i, j] - 0.95) > 0.28 else "k")
        ax.add_patch(Rectangle((4, 0), 1, 4, fill=False, hatch="xxx", ec="#555555", lw=0))
        ax.set_title(f"{NAME[net]}: {title}", fontsize=10)
        ax.set_xticks(np.arange(5) + 0.5, [f"{d:g}" for d in densities])
        ax.set_yticks(np.arange(4) + 0.5, [f"{p:,.0f}" for p in photons])
    axes[row, 0].set_ylabel("photons per emitter")
fig.supxlabel("emitters per µm² per frame. Hatched: outside the training prior", fontsize=10.5)
fig.colorbar(mesh, ax=axes, shrink=0.6, label="value (white is 0.95)")
fig.suptitle("Simulation only. Coverage is closest to nominal in the world nearest each network's training", fontsize=11)
for net in ("a", "b"):
    for world in ("gauss", "empirical"):
        print(f"{NAME[net]} in the {world} world: {near_nominal(net, world)[0]} of {near_nominal(net, world)[1]} "
              "cells in the prior have 95% coverage within 0.90-0.98")
emp_a = [r for r in certificate["simulation"]["a"] if r["world"] == "empirical" and r["in_prior"]]
emp_b = [r for r in certificate["simulation"]["b"] if r["world"] == "empirical" and r["in_prior"]]
print(f"Empirical world, in the prior: A precision under 0.9 in {sum(r['precision'] < 0.9 for r in emp_a)} of 16, lowest "
      f"{min(r['precision'] for r in emp_a):.2f}; B in {sum(r['precision'] < 0.9 for r in emp_b)} of 16, with recall under "
      f"0.9 in {sum(r['recall'] < 0.9 for r in emp_b)}")
own = [r for r in certificate["simulation"]["a"] if r["world"] == "gauss" and r["density"] == 0.7 and r["photons"] in (2000, 4000)]
print(f"A in its own world at 0.70 per um^2, 2,000-4,000 photons: {min(r['rmse_over_crlb'] for r in own):.1f}-"
      f"{max(r['rmse_over_crlb'] for r in own):.1f} x CRLB");

psf = store.load_json("psf_bg")
knots, profile = np.array(psf["knots_px"]), np.array(psf["profile"])
sigma = 1.05
gauss = np.exp(-knots**2 / (2 * sigma**2)) / (2 * np.pi * sigma**2)
print(f"{psf['spots']:,} spots of {psf['pool']:,}: {psf['share_beyond_3px']:.1%} of energy beyond 3 px against "
      f"{psf['share_beyond_3px_gauss_105']:.1%} for a {sigma} px Gaussian; peak {psf['peak']:.3f} against "
      f"{psf['peak_gauss_105']:.3f}; sd of log width {psf['log_w_sd']:.3f}, which includes fit noise")
psf_control = posthoc["psf"]
print(f"Post-hoc control on the same {psf_control['spots']:,} spots. Stack: light over N {psf_control['stack_integral']:.2f}, "
      f"{psf_control['stack_share_signed']:.1%} beyond 3 px counted out to 9 px")
for name, title in (("after_100", "same sites, 100 frames away"), ("after_300", "same sites, 300 frames away"), ("random_sites", "random sites, same frames")):
    row = psf_control["controls"][name]
    print(f"  {title:28} leaves {row['integral_control']:.2f} of N, flat {row['pedestal_3_to_8px']:.4f} per px² at 3-8 px, core {row['core_under_1px']:.4f}. "
          f"Stack minus this: light over N {row['integral_difference']:.2f}, {row['share_difference']:.1%} beyond 3 px, peak {row['peak_difference']:.3f}")
width = psf_control["width"]
print(f"Width scale of a spot: median {width['stored_rule_median']:.2f}, log sd {width['stored_rule_log_sd']:.3f} with the stored fit; "
      f"{width['amplitude_free_median']:.2f}, {width['amplitude_free_log_sd']:.3f} with a free amplitude")
centres = (np.arange(len(psf_control["raw"])) + 0.5) * psf_control["knots_px"][1]
after = psf_control["controls"]["after_100"]
clean = np.array(after["difference"]) / after["integral_difference"]
fig, ax = plt.subplots(figsize=(8.6, 3.8), layout="constrained")
ax.fill_between(knots * PIXEL_NM, np.maximum(pack["psf_band"][0], 1e-6), pack["psf_band"][1], color=NET_B, alpha=0.35, lw=0)
ax.plot(knots * PIXEL_NM, np.maximum(profile, 1e-6), color=NET_B, lw=2, label="the stored stack, which B was trained on")
ax.plot(centres * PIXEL_NM, np.where(clean > 0, clean, np.nan), color=CLASSICAL, lw=2, label="stack minus the same sites 100 frames away, post-hoc")
ax.plot(knots * PIXEL_NM, gauss, color=NET_A, lw=2, ls="--", label="what A was trained on: Gaussian, 1.05 px")
ax.axvline(9 * PIXEL_NM, color=MUTED, lw=1, ls=":")
ax.text(9 * PIXEL_NM + 12, 0.002, "wings past here\nare not measured", fontsize=9, color=SINGLE)
ax.set_yscale("log")
ax.set(ylim=(1e-5, 0.3), xlim=(0, 1500), xlabel="distance from the spot centre (nm)", ylabel="share of photons per px²")
ax.set_title("Most of the stack's wings stay when the spot has gone, so they are not PSF", fontsize=11)
ax.legend(loc="upper right", fontsize=9);

verdicts = paired["verdicts"]
rows = []
for key, title in (("r32", "recording 32, judged"), ("r7", "recording 7"), ("r32_line", "recording 32, line calibration"),
                   ("r7_line", "recording 7, line calibration")):
    rows.append({"where": title, "A": verdicts[f"{key}_a"]["z"], "A 95% half": verdicts[f"{key}_a"]["halfwidth95"],
                 "A 1 se": verdicts[f"{key}_a"]["one_se"], "B": verdicts[f"{key}_b"]["z"],
                 "B 95% half": verdicts[f"{key}_b"]["halfwidth95"], "B 1 se": verdicts[f"{key}_b"]["one_se"],
                 "classical": verdicts[f"{key}_a"]["classical_z"], "K = 1": paired["rows"][key]["K1"]["gap3_R1.5"]["matched_chosen"]["z_sd"],
                 "K = 1, own basis": verdicts[f"{key}_a"]["k1_z"],
                 "absolute": f"{verdicts[f'{key}_a']['absolute']}, {verdicts[f'{key}_b']['absolute']}",
                 "relative": verdicts[f"{key}_a"]["relative"]})
print(pd.DataFrame(rows).round(3).to_string(index=False))
sim = paired["rows"]["simulated"]
print("\nSimulated baseline, gap 3, unweighted, whole field:",
      ", ".join(f"{NAME[m]} {sim[m]['gap3_R1.5']['unweighted']['z_sd']:.2f}" for m in ("a", "b", "chosen", "K1")))
print("It tests the method. It says nothing of the microscope.")
truth = posthoc["sim"]
print("Post-hoc, pairs with both ends on a true emitter: " + ", ".join(f"{NAME[m]} {truth[m]['z_both_true']:.2f}" for m in ("a", "b", "chosen", "K1")))
print("Share of rows within 2 px of a true emitter: " + ", ".join(f"{NAME[m]} {truth[m]['precision']:.2f}" for m in ("a", "b", "chosen", "K1"))
      + f". A's gap-3 pairs: {truth['a']['pairs']:,}, of which {truth['a']['pairs_neither']:,} join two false rows, yet the mixture calls {truth['a']['mixture_true_share']:.0%} true")
print(f"Network B at gap 1 on simulation: z {sim['b']['gap1_R1.5']['unweighted']['z_sd']:.2f}, where its 3-frame input shares pixels")

extra = posthoc["real"]
print("\nPost-hoc, the same links read other ways. Gated values above are unchanged.")
rows = []
for rec in ("r32", "r7"):
    common, floor = extra[rec]["common"], extra[rec]["floor"]
    for m in ("chosen", "K1", "a", "b"):
        row = paired["rows"][rec][m]["gap3_R1.5"]
        rows.append({"recording": RECS[rec], "method": NAME[m], "own pairs": row["pairs"], "matched z": row["matched_chosen"]["z_sd"], "matched to K = 1": row["matched_K1"]["z_sd"],
                     "unweighted z": row["unweighted"]["z_sd"], "unw. upper": row["unweighted"]["z_sd_hi"], "common z": common["z"][m], "common scatter nm": common["scatter_nm"][m],
                     "z tight / mid / wide sd": " / ".join(f"{v:.2f}" for v in extra[rec]["sd_tertiles"][m]), "residual scale": floor[m]["c"], "floor nm per frame": floor[m]["floor_nm"] / np.sqrt(2)})
print(pd.DataFrame(rows).round(3).to_string(index=False))
print(f"Common pairs: {extra['r32']['common']['pairs']:,} on recording 32, {extra['r7']['common']['pairs']:,} on recording 7. Scatter by gap 1 / 3 / 10, K = 1, recording 32: "
      + " / ".join(f"{extra['r32']['by_gap_scatter_nm']['K1'][g]:.1f}" for g in ("1", "3", "10")) + " nm")

fig, axes = plt.subplots(1, 3, figsize=(13.2, 4.3), layout="constrained")
ax = axes[0]
place = {"sim": 3, "r7": 2, "r32": 1}
offsets = {"a": 0.2, "b": 0.0, "chosen": -0.2}
for m, shift in offsets.items():
    judged, other = paired["rows"]["r32"][m]["gap3_R1.5"], paired["rows"]["r7"][m]["gap3_R1.5"]
    for y, row in ((place["r32"], judged), (place["r7"], other)):
        q = row["matched_chosen"]
        ax.errorbar(q["z_sd"], y + shift, xerr=[[q["z_sd"] - q["z_sd_lo"]], [q["z_sd_hi"] - q["z_sd"]]], fmt="o", ms=7,
                    color=COLOUR[m], capsize=3)
    ax.scatter(posthoc["sim"][m]["z_both_true"], place["sim"] + shift, s=48, marker="D", color=COLOUR[m])
ax.axvspan(1.0, 1.15, color="#2a7f62", alpha=0.15, lw=0)
ax.axvspan(1.25, 1.75, color=NET_A, alpha=0.07, lw=0)
ax.axvline(1.25, color=LINE_COL, lw=1.1)
ax.text(1.003, 3.42, "pass", fontsize=9, color="#2a7f62")
ax.text(1.262, 3.42, "failure", fontsize=9, color=NET_A)
ax.set_yticks([1, 2, 3], ["recording 32\n(judged)", "recording 7", "simulated,\ntrue pairs"])
ax.set(xlim=(0.98, 1.75), ylim=(0.6, 3.6))
ax.set_xlabel("z sd (1.0 means honest error bars)", fontsize=10)
ax.set_title("Simulation passes, real frames fail", fontsize=11)
ax.legend(handles=[Line2D([], [], marker="o", ls="", color=COLOUR[m], label=NAME[m]) for m in offsets], loc="center right",
          bbox_to_anchor=(1.0, 0.72), fontsize=8.5)

ax = axes[1]
q = pack["z_quantiles"]
grid_z = np.linspace(-4.5, 4.5, 200)
for m in ("a", "b", "chosen"):
    ax.plot(pack[f"r32_z_{m}"], q, color=COLOUR[m], lw=2)
ax.plot(pack["sim_z_a"], q, color=NET_A, lw=1.2, ls=":", label="network A on simulation, false rows in")
ax.plot(grid_z, stats.norm.cdf(grid_z), color=LINE_COL, lw=1.3, label="honest: N(0, 1)")
ax.plot(grid_z, stats.norm.cdf(grid_z, scale=1.25), color=LINE_COL, lw=1.1, ls="--", label="the failure line: sd 1.25")
ax.set(xlim=(-4.5, 4.5), ylim=(0, 1), ylabel="cumulative share")
ax.set_xlabel("z on recording 32, gap 3 (raw, false links included)", fontsize=10)
ax.set_title("Real curves sit outside the failure line", fontsize=11)
ax.legend(fontsize=8.5, loc="upper left")

ax = axes[2]
radii = (0.75, 1.0, 1.5, 2.0, 3.0)
for m in ("a", "b", "chosen"):
    z = [paired["rows"]["r32"][m][f"gap3_R{r}"]["matched_chosen"]["z_sd"] for r in radii]
    ax.plot(radii, z, "o-", color=COLOUR[m], lw=2)
    gap1 = paired["rows"]["r32"][m]["gap1_R1.5"]["matched_chosen"]["z_sd"]
    ax.scatter(1.5, gap1, marker="x", s=70, color=COLOUR[m], zorder=4)
ax.axhline(1.25, color=LINE_COL, lw=1.1)
ax.axvline(1.5, color=MUTED, lw=1, ls=":")
ax.text(1.53, 1.21, "written radius", fontsize=8.5, color=SINGLE)
ax.set(xlabel="link radius (px)", ylabel="z sd, recording 32")
ax.set_title("No radius rescues it. x marks gap 1", fontsize=11);