export_block_catalog_distributions.py

#!/usr/bin/env python3
"""Export labeled block catalogs and per-sample metric distributions."""

from __future__ import annotations

import argparse
import csv
import math
from pathlib import Path
from typing import Dict, List, Sequence

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


METRICS = [
    ("omit_length_um", "Omit Length (um)"),
    ("correct_length_um", "Recall Length (um)"),
    ("split_length_um", "Split Length (um)"),
    ("merge_count", "Merge Count"),
    ("self_merge_count", "Self-Merge Count"),
    ("avg_intensity_mean", "Average Intensity"),
]

ALL_COMBINED_LABEL = "all-carveouts-combined"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source-csv", required=True, type=Path)
    parser.add_argument("--output-dir", required=True, type=Path)
    parser.add_argument("--voxel-size-xy", type=float, default=0.748)
    return parser.parse_args()


def read_csv_rows(path: Path) -> List[dict]:
    with path.open("r", newline="") as f:
        return list(csv.DictReader(f))


def write_csv(path: Path, rows: Sequence[dict], fieldnames: Sequence[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def safe_float(value: str) -> float:
    return float(value) if value not in ("", None) else 0.0


def safe_int(value: str) -> int:
    return int(value) if value not in ("", None) else 0


def add_derived_metrics(row: dict, voxel_size_xy: float) -> dict:
    out = row.copy()
    counts = {
        label: safe_int(row[f"{label}_point_count"])
        for label in ("omit", "correct", "split")
    }
    means = {
        label: safe_float(row[f"{label}_intensity_mean"])
        for label in ("omit", "correct", "split")
    }
    stds = {
        label: safe_float(row[f"{label}_intensity_std"])
        for label in ("omit", "correct", "split")
    }
    total_points = sum(counts.values())
    if total_points > 0:
        mean = sum(counts[label] * means[label] for label in counts) / total_points
        second = sum(counts[label] * (stds[label] ** 2 + means[label] ** 2) for label in counts) / total_points
        out["avg_intensity_mean"] = mean
        out["avg_intensity_std"] = math.sqrt(max(second - mean * mean, 0.0))
    else:
        out["avg_intensity_mean"] = 0.0
        out["avg_intensity_std"] = 0.0
    out["total_length_um"] = sum(safe_float(row[f"{label}_length_um"]) for label in ("omit", "correct", "split"))
    out["merge_total"] = safe_int(row["merge_count"]) + safe_int(row["self_merge_count"])
    out["cube_center_x_phys_um"] = safe_float(row["cube_center_x_vox"]) * voxel_size_xy
    out["cube_center_y_phys_um"] = safe_float(row["cube_center_y_vox"]) * voxel_size_xy
    out["cube_center_z_phys_um"] = safe_float(row["cube_center_z_vox"])
    out["center_volume_coordinate_vox"] = (
        f"[{int(round(safe_float(row['cube_center_x_vox'])))}"
        f",{int(round(safe_float(row['cube_center_y_vox'])))}"
        f",{int(round(safe_float(row['cube_center_z_vox'])))}]"
    )
    return out


def build_catalog(rows: Sequence[dict]) -> List[dict]:
    grouped: Dict[str, List[dict]] = {}
    for row in rows:
        grouped.setdefault(row["carveout"], []).append(row)

    catalog_rows: List[dict] = []
    for carveout, group in sorted(grouped.items()):
        group_sorted = sorted(
            group,
            key=lambda r: (
                safe_int(r["cube_z"]),
                safe_int(r["cube_y"]),
                safe_int(r["cube_x"]),
            ),
        )
        width = max(3, len(str(len(group_sorted) - 1)))
        for idx, row in enumerate(group_sorted):
            out = row.copy()
            out["block_label"] = f"{idx:0{width}d}"
            catalog_rows.append(out)
    return catalog_rows


def plot_distributions(rows: Sequence[dict], output_dir: Path) -> None:
    grouped: Dict[str, List[dict]] = {}
    for row in rows:
        grouped.setdefault(row["carveout"], []).append(row)
    grouped[ALL_COMBINED_LABEL] = list(rows)

    dist_dir = output_dir / "distributions"
    dist_dir.mkdir(parents=True, exist_ok=True)
    summary_rows = []

    for carveout, group in sorted(grouped.items()):
        fig, axes = plt.subplots(2, 3, figsize=(15, 9), facecolor="white")
        axes = axes.ravel()
        for ax, (metric, title) in zip(axes, METRICS):
            vals = np.asarray([safe_float(row[metric]) for row in group], dtype=float)
            ax.hist(vals, bins=30, color="#4f81bd", edgecolor="white", linewidth=0.6)
            ax.axvline(float(np.mean(vals)), color="#c0504d", linewidth=1.3, label="mean")
            ax.axvline(float(np.median(vals)), color="#9bbb59", linewidth=1.3, label="median")
            ax.set_title(title)
            ax.set_xlabel(metric)
            ax.set_ylabel("block count")
            ax.legend(loc="upper right", fontsize=8)
            summary_rows.append(
                {
                    "carveout": carveout,
                    "metric": metric,
                    "count": len(vals),
                    "mean": float(np.mean(vals)),
                    "median": float(np.median(vals)),
                    "std": float(np.std(vals)),
                    "min": float(np.min(vals)),
                    "max": float(np.max(vals)),
                }
            )
        title = (
            "all cells / blocks combined metric distributions"
            if carveout == ALL_COMBINED_LABEL
            else f"{carveout} block metric distributions"
        )
        fig.suptitle(title, fontsize=14)
        plt.tight_layout()
        plt.savefig(dist_dir / f"{carveout}_block_metric_distributions.png", dpi=220, bbox_inches="tight")
        plt.close(fig)

    write_csv(
        output_dir / "distribution_summary.csv",
        summary_rows,
        ["carveout", "metric", "count", "mean", "median", "std", "min", "max"],
    )


def main() -> None:
    args = parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)

    source_rows = [add_derived_metrics(row, args.voxel_size_xy) for row in read_csv_rows(args.source_csv)]
    catalog_rows = build_catalog(source_rows)
    fieldnames = [
        "carveout",
        "block_label",
        "center_volume_coordinate_vox",
        "cube_z",
        "cube_y",
        "cube_x",
        "cube_origin_z_vox",
        "cube_origin_y_vox",
        "cube_origin_x_vox",
        "cube_center_z_phys_um",
        "cube_center_y_phys_um",
        "cube_center_x_phys_um",
        "omit_length_um",
        "correct_length_um",
        "split_length_um",
        "total_length_um",
        "merge_count",
        "self_merge_count",
        "merge_total",
        "omit_point_count",
        "correct_point_count",
        "split_point_count",
        "omit_intensity_mean",
        "omit_intensity_std",
        "correct_intensity_mean",
        "correct_intensity_std",
        "split_intensity_mean",
        "split_intensity_std",
        "avg_intensity_mean",
        "avg_intensity_std",
    ]
    write_csv(args.output_dir / "block_catalog.csv", catalog_rows, fieldnames)
    plot_distributions(catalog_rows, args.output_dir)


if __name__ == "__main__":
    main()