Blame for size_report.py

"""
size_report.py - cell size distribution report for CellposeCellCounter (v6)

After the concentration, report the size of the counted cells the way a
commercial counter does: a Total / Live / Dead table (percentage, concentration,
count, average size) and a live/dead size histogram in micrometres.

Sizes are equivalent diameters (diameter of the circle with the mask's area),
converted to um with the 1 mm x 1 mm counting block, so they are only available
for squares cropped with the 4-point block crop. They are pooled over all
squares; the concentrations use the mean count per square, exactly as the
concentration block above them does.

Pure numpy; matplotlib for the histogram.
"""

import numpy as np

HIST_BIN_UM = 1.0     # histogram bin width, um
HIST_MAX_UM = 60.0    # right edge of the histogram (objects above go in the last bin)


def cell_diameters_um(masks, label_map, um_per_px):
    """{'live': [...], 'dead': [...]} equivalent diameters in um.

    `label_map` is {cell_id: 0 live, 1 dead} from the viability step. Cells the
    viability step did not score are left out, so the size counts always match
    the live/dead counts.
    """
    counts = np.bincount(np.asarray(masks).ravel().astype(np.int64))
    live, dead = [], []
    for cid, lab in (label_map or {}).items():
        cid = int(cid)
        if cid <= 0 or cid >= counts.size or counts[cid] == 0:
            continue
        d = 2.0 * np.sqrt(counts[cid] / np.pi) * float(um_per_px)
        (dead if int(lab) == 1 else live).append(round(float(d), 3))
    return {"live": live, "dead": dead, "um_per_px": float(um_per_px)}


def pool(results):
    """Pool per-square size dicts. Returns (live, dead, n_squares_with_sizes)."""
    live, dead, n = [], [], 0
    for r in results:
        s = (r or {}).get("sizes")
        if not s:
            continue
        live += list(s.get("live", []))
        dead += list(s.get("dead", []))
        n += 1
    return np.asarray(live, float), np.asarray(dead, float), n


def _stats(a):
    if a.size == 0:
        return {"n": 0, "mean": None, "median": None, "sd": None}
    return {"n": int(a.size), "mean": float(a.mean()),
            "median": float(np.median(a)),
            "sd": float(a.std(ddof=1)) if a.size > 1 else 0.0}


def _sci(x):
    """2.05 x 10^6 style."""
    if not x:
        return "0"
    e = int(np.floor(np.log10(abs(x))))
    return f"{x / 10 ** e:.2f} x 10^{e}"


def summarize(results, conc_presets):
    """Size summary over all squares.

    Returns a dict with per-class stats plus `n_squares` (squares that had
    sizes) and `mean_per_square` counts used for the concentrations, or None if
    no square carries sizes.
    """
    live, dead, n = pool(results)
    if n == 0:
        return None
    total = np.concatenate([live, dead])
    out = {"n_squares": n, "live": live, "dead": dead, "total": total,
           "stats": {"Total": _stats(total), "Live": _stats(live),
                     "Dead": _stats(dead)},
           "conc_presets": list(conc_presets)}
    return out


def format_table(summary):
    """Monospace text table in the layout of the commercial counter report."""
    if summary is None:
        return ("Size distribution needs the 4-point block crop (it sets the "
                "um scale). No square had one.")
    st, n = summary["stats"], summary["n_squares"]
    tot = st["Total"]["n"] or 0
    presets = summary["conc_presets"]
    head = f"{'':<7}{'Percent':>8}"
    for name, mult in presets:
        # chamber factor is 10^4 per mL, so the dilution is mult / 10^4
        head += f"{('Conc (' + format(mult / 1e4, 'g') + 'x dil.)'):>19}"
    head += f"{'Count':>7}{'Avg size':>11}{'Median':>9}{'SD':>7}"
    lines = [head, "-" * len(head)]
    for cls in ("Total", "Live", "Dead"):
        s = st[cls]
        pct = "" if cls == "Total" else (f"{100 * s['n'] / tot:.0f}%" if tot else "-")
        row = f"{cls:<7}{pct:>8}"
        for _, mult in presets:
            row += f"{(_sci(s['n'] / n * mult) + ' /mL'):>19}"
        row += f"{s['n']:>7}"
        row += (f"{s['mean']:>9.2f}um{s['median']:>9.2f}{s['sd']:>7.2f}"
                if s["mean"] is not None else f"{'-':>11}{'-':>9}{'-':>7}")
        lines.append(row)
    lines.append("")
    lines.append(f"Count = cells in {n} square{'s' if n != 1 else ''} pooled; "
                 f"concentration = count / {n} x factor.")
    lines.append("Size = equivalent diameter (circle of equal area), um, "
                 "scale from the 1 mm block.")
    lines.append("These match the concentration above when every square "
                 "was block-cropped.")
    return "\n".join(lines)


def histogram_counts(summary, bin_um=HIST_BIN_UM, max_um=HIST_MAX_UM):
    """(edges, live_counts, dead_counts); values above max_um fall in the last bin."""
    edges = np.arange(0.0, max_um + bin_um, bin_um)
    clip = lambda a: np.clip(a, 0, max_um - 1e-6)
    lc, _ = np.histogram(clip(summary["live"]), bins=edges)
    dc, _ = np.histogram(clip(summary["dead"]), bins=edges)
    return edges, lc, dc


def render_histogram(summary, bin_um=HIST_BIN_UM, max_um=HIST_MAX_UM):
    """PIL image: live (green) and dead (red) size histogram, stacked."""
    import io
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from PIL import Image

    fig, ax = plt.subplots(figsize=(6.0, 3.2), dpi=120)
    if summary is not None and summary["total"].size:
        edges, lc, dc = histogram_counts(summary, bin_um, max_um)
        centres = edges[:-1] + bin_um / 2
        ax.bar(centres, lc, width=bin_um * 0.9, color="#2ca02c",
               label=f"Live  n={summary['stats']['Live']['n']}")
        ax.bar(centres, dc, width=bin_um * 0.9, bottom=lc, color="#d62728",
               label=f"Dead  n={summary['stats']['Dead']['n']}")
        for cls, col in (("Live", "#2ca02c"), ("Dead", "#d62728")):
            m = summary["stats"][cls]["mean"]
            if m is not None:
                ax.axvline(m, color=col, lw=1, ls="--")
        txt = "   ".join(f"{c} avg {summary['stats'][c]['mean']:.2f} um"
                         for c in ("Live", "Dead")
                         if summary["stats"][c]["mean"] is not None)
        ax.set_title(txt, fontsize=9)
        ax.legend(fontsize=8, frameon=False)
    else:
        ax.text(0.5, 0.5, "no sized cells", ha="center", va="center",
                transform=ax.transAxes, color="grey")
    ax.set_xlim(0, max_um)
    ax.set_xlabel("Cell size, equivalent diameter (um)")
    ax.set_ylabel("Cell quantity")
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)
    fig.tight_layout()
    buf = io.BytesIO()
    fig.savefig(buf, format="png")
    plt.close(fig)
    buf.seek(0)
    return Image.open(buf).convert("RGB")


def csv_rows(summary, bin_um=HIST_BIN_UM, max_um=HIST_MAX_UM):
    """Histogram as rows for a CSV: bin_start_um, bin_end_um, live, dead, total."""
    if summary is None:
        return []
    edges, lc, dc = histogram_counts(summary, bin_um, max_um)
    return [[f"{edges[i]:g}", f"{edges[i + 1]:g}", int(lc[i]), int(dc[i]),
             int(lc[i] + dc[i])] for i in range(len(lc))]