tools/generate_appE_rows.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Generate Appendix E row-only LaTeX fragments for headline PSNR / SSIM / LPIPS.

Expected input CSV schema, long format:

    method_key,scene,metric,value
    3dgsmcmc,Bonsai,psnr,32.398
    3dgsmcmc,Bonsai,ssim,0.9421
    3dgsmcmc,Bonsai,lpips,0.0584

If multiple rows exist for the same (method_key, scene, metric), they are averaged
first. This supports multi-seed cells by taking the seed mean.

If your actual CSV is wide format or uses different column names, replace only
load_long_csv() below. Keep the aggregation / LaTeX logic unchanged.

Usage:
    python tools/generate_appE_rows.py path/to/per_cell_metrics_long.csv

Outputs:
    tex/appE_psnr_rows.tex
    tex/appE_ssim_rows.tex
    tex/appE_lpips_rows.tex
"""

from __future__ import annotations

import argparse
import csv
import math
import os
from collections import defaultdict
from typing import Dict, Iterable, List, Tuple, Optional


DEFAULT_INPUT_CSV = "outputs/phase1/per_cell_metrics_long.csv"
OUT_DIR = "tex"

MASK_METHOD_FAILURE = False
METHOD_FAILURE_CELLS = {
    ("erankgs", "Lego"),
}

DATASETS = {
    "T&T": [
        "Auditorium", "Ballroom", "Barn", "Caterpillar", "Courtroom", "Lighthouse",
        "Museum", "Palace", "Playground", "Temple", "Train", "Truck",
    ],
    "Mip-360": [
        "Bicycle", "Bonsai", "Counter", "Flowers", "Garden", "Kitchen", "Room",
        "Stump", "Treehill",
    ],
    "NS": [
        "Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship",
    ],
    "DB": [
        "DrJohnson", "Playroom",
    ],
}

DATASET_ORDER = ["T&T", "Mip-360", "NS", "DB"]
ALL_CANONICAL_SCENES = [scene for ds in DATASET_ORDER for scene in DATASETS[ds]]

SCENE_CANONICAL = {s.lower(): s for s in ALL_CANONICAL_SCENES}
SCENE_CANONICAL.update({
    "auditorium": "Auditorium",
    "ballroom": "Ballroom",
    "barn": "Barn",
    "caterpillar": "Caterpillar",
    "courtroom": "Courtroom",
    "lighthouse": "Lighthouse",
    "museum": "Museum",
    "palace": "Palace",
    "playground": "Playground",
    "temple": "Temple",
    "train": "Train",
    "truck": "Truck",
    "bicycle": "Bicycle",
    "bonsai": "Bonsai",
    "counter": "Counter",
    "flowers": "Flowers",
    "garden": "Garden",
    "kitchen": "Kitchen",
    "room": "Room",
    "stump": "Stump",
    "treehill": "Treehill",
    "chair": "Chair",
    "drums": "Drums",
    "ficus": "Ficus",
    "hotdog": "Hotdog",
    "lego": "Lego",
    "materials": "Materials",
    "mic": "Mic",
    "ship": "Ship",
    "drjohnson": "DrJohnson",
    "dr_johnson": "DrJohnson",
    "playroom": "Playroom",
})

CATEGORIES = [
    ("FREE", [
        "3dgsmcmc", "absgs", "atomgs", "conegs", "ges", "ghap", "gslpm",
        "lapisgs", "reactgs", "vanilla_3dgs",
    ]),
    ("GEO_REG", [
        "2dgs", "gaussian_surfel", "gof", "pgsr", "scaffoldgs",
    ]),
    ("PRIM_CTRL", [
        "coadaptgs", "erankgs", "gaussianpro", "minisplatting",
        "opti3dgs", "steepgs",
    ]),
    ("RENDER", [
        "3dgs_dr", "analyticsplatting", "lod_gs", "mipsplatting",
        "pixelgs",
    ]),
    ("COMPRESS", [
        "cdcgs", "hogs", "lightgaussian", "octree_gs", "trimgs",
    ]),
]

DISPLAY_NAMES = {
    "3dgsmcmc": "3DGS-MCMC",
    "absgs": "AbsGS",
    "atomgs": "AtomGS",
    "conegs": "ConeGS",
    "ges": "GES",
    "ghap": "GHAP",
    "gslpm": "GSLPM",
    "lapisgs": "LapisGS",
    "reactgs": "ReactGS",
    "vanilla_3dgs": "3DGS",
    "2dgs": "2DGS",
    "gaussian_surfel": "Gaussian Surfels",
    "gof": "GOF",
    "pgsr": "PGSR",
    "scaffoldgs": "Scaffold-GS",
    "coadaptgs": "CoAdaptGS",
    "erankgs": "eRankGS",
    "gaussianpro": "GaussianPro",
    "minisplatting": "MiniSplatting",
    "opti3dgs": "Opti3DGS",
    "steepgs": "SteepGS",
    "3dgs_dr": "3DGS-DR",
    "analyticsplatting": "AnalyticSplatting",
    "lod_gs": "LoD-GS",
    "mipsplatting": "Mip-Splatting",
    "pixelgs": "PixelGS",
    "cdcgs": "CDCGS",
    "hogs": "HoGS",
    "lightgaussian": "LightGaussian",
    "octree_gs": "Octree-GS",
    "trimgs": "TrimGS",
}

METRIC_FORMATS = {
    "psnr": "{:.2f}",
    "ssim": "{:.4f}",
    "lpips": "{:.4f}",
}

OUTPUT_FILES = {
    "psnr": "appE_psnr_rows.tex",
    "ssim": "appE_ssim_rows.tex",
    "lpips": "appE_lpips_rows.tex",
}


def latex_escape(text: str) -> str:
    replacements = {
        "\\": r"\textbackslash{}",
        "_": r"\_",
        "&": r"\&",
        "%": r"\%",
        "$": r"\$",
        "#": r"\#",
        "{": r"\{",
        "}": r"\}",
    }
    return "".join(replacements.get(ch, ch) for ch in text)


def canonical_scene(scene: str) -> Optional[str]:
    scene = scene.strip()
    return SCENE_CANONICAL.get(scene.lower())


def parse_float(value: str) -> Optional[float]:
    value = value.strip()
    if value == "" or value.lower() in {"nan", "none", "null", "na", "n/a", "---"}:
        return None
    try:
        x = float(value)
    except ValueError:
        return None
    if math.isnan(x) or math.isinf(x):
        return None
    return x


def mean_or_none(values: Iterable[float]) -> Optional[float]:
    vals = list(values)
    if not vals:
        return None
    return sum(vals) / len(vals)


def fmt_metric(metric: str, value: Optional[float]) -> str:
    if value is None:
        return "---"
    return METRIC_FORMATS[metric].format(value)


def load_long_csv(path: str) -> Dict[Tuple[str, str, str], float]:
    """
    Load long-format per-cell metrics and average duplicate rows.

    Returns:
        dict[(method_key, canonical_scene, metric)] = mean_value

    TODO if your actual CSV is wide format:
        Replace this function so it emits the same dictionary, e.g.
        for columns method, scene, psnr, ssim, lpips:
            out[(method, scene, "psnr")] = psnr_value
            out[(method, scene, "ssim")] = ssim_value
            out[(method, scene, "lpips")] = lpips_value
    """
    buckets: Dict[Tuple[str, str, str], List[float]] = defaultdict(list)

    with open(path, "r", newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        if reader.fieldnames is None:
            raise ValueError(f"Input CSV has no header: {path}")

        fields = set(reader.fieldnames)
        method_col = "method_key" if "method_key" in fields else "method"
        scene_col = "scene"
        metric_col = "metric"
        value_col = "value"

        required = {method_col, scene_col, metric_col, value_col}
        missing = required - fields
        if missing:
            raise ValueError(
                f"Input CSV missing columns {sorted(missing)}. "
                f"Expected long schema: method_key,scene,metric,value. "
                f"Actual columns: {reader.fieldnames}"
            )

        for row in reader:
            method = row[method_col].strip()
            scene_raw = row[scene_col].strip()
            metric = row[metric_col].strip().lower()
            value = parse_float(row[value_col])

            if metric not in {"psnr", "ssim", "lpips"}:
                continue
            if value is None:
                continue

            scene = canonical_scene(scene_raw)
            if scene is None:
                continue

            if MASK_METHOD_FAILURE and (method, scene) in METHOD_FAILURE_CELLS:
                continue

            buckets[(method, scene, metric)].append(value)

    return {key: mean_or_none(vals) for key, vals in buckets.items() if vals}


def aggregate_for_method(
    cell_values: Dict[Tuple[str, str, str], float],
    method: str,
    metric: str,
) -> Tuple[Optional[float], Optional[float], Optional[float], Optional[float], Optional[float]]:
    dataset_means: List[Optional[float]] = []

    for ds in DATASET_ORDER:
        vals = []
        for scene in DATASETS[ds]:
            v = cell_values.get((method, scene, metric))
            if v is not None:
                vals.append(v)
        dataset_means.append(mean_or_none(vals))

    overall_vals = []
    for scene in ALL_CANONICAL_SCENES:
        v = cell_values.get((method, scene, metric))
        if v is not None:
            overall_vals.append(v)

    overall = mean_or_none(overall_vals)
    return (*dataset_means, overall)


def render_metric_fragment(
    cell_values: Dict[Tuple[str, str, str], float],
    metric: str,
) -> List[str]:
    lines: List[str] = []
    lines.append("% AUTOGENERATED by tools/generate_appE_rows.py - DO NOT EDIT BY HAND.")

    for cat_idx, (category, methods) in enumerate(CATEGORIES):
        if cat_idx > 0:
            lines.append(r"\midrule")

        category_tex = latex_escape(category)
        lines.append(rf"\multicolumn{{6}}{{@{{}}l}}{{\emph{{{category_tex}}}}} \\")

        for method in methods:
            display = latex_escape(DISPLAY_NAMES[method])
            values = aggregate_for_method(cell_values, method, metric)
            formatted = [fmt_metric(metric, v) for v in values]
            lines.append(
                f"{display} & {formatted[0]} & {formatted[1]} & "
                f"{formatted[2]} & {formatted[3]} & {formatted[4]} \\\\"
            )

    return lines


def write_fragments(cell_values: Dict[Tuple[str, str, str], float]) -> None:
    os.makedirs(OUT_DIR, exist_ok=True)

    for metric, filename in OUTPUT_FILES.items():
        path = os.path.join(OUT_DIR, filename)
        lines = render_metric_fragment(cell_values, metric)
        with open(path, "w", encoding="utf-8", newline="\n") as f:
            f.write("\n".join(lines))
            f.write("\n")
        print(f"Wrote {path} ({len(lines)} lines)")


def sanity_check_methods() -> None:
    flat_methods = [m for _, methods in CATEGORIES for m in methods]
    if len(flat_methods) != 31:
        raise ValueError(f"Expected 31 methods, got {len(flat_methods)}")
    if len(set(flat_methods)) != 31:
        dupes = sorted({m for m in flat_methods if flat_methods.count(m) > 1})
        raise ValueError(f"Duplicate methods in inventory: {dupes}")
    missing_display = [m for m in flat_methods if m not in DISPLAY_NAMES]
    if missing_display:
        raise ValueError(f"Missing display names: {missing_display}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "csv",
        nargs="?",
        default=DEFAULT_INPUT_CSV,
        help="Input long-format CSV: method_key,scene,metric,value",
    )
    args = parser.parse_args()

    sanity_check_methods()
    cell_values = load_long_csv(args.csv)
    write_fragments(cell_values)

    print(
        "Generated 3 fragments: 31 rows + 4 category dividers + 1 header "
        "multicolumn each = 36 lines per file"
    )


if __name__ == "__main__":
    main()