alnn_rewrite/analysis/evaluate_models.py

# pyright: basic

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd

from analysis.analysis_modules import (
    run_calibration,
    run_longitudinal,
    run_performance,
    run_physician,
)
from analysis.data_access import load_backend_evaluation, load_clinical_table
from analysis.holdout_evaluation import ensure_backend_netcdf
from analysis.noise_analysis import run_noise_analysis
from analysis.runtime import backend_dir, init_runtime_paths, load_config, write_json


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Run modular evaluation analyses for ensemble and bayesian models. "
            "All outputs are written to alnn_rewrite/analysis_output."
        )
    )
    parser.add_argument(
        "--backend",
        nargs="+",
        choices=["ensemble", "bayesian"],
        default=["ensemble", "bayesian"],
        help="Backends to evaluate.",
    )
    parser.add_argument(
        "--run-name",
        default=None,
        help="Optional run directory name under analysis_output.",
    )
    parser.add_argument(
        "--threshold-start",
        type=float,
        default=0.5,
        help="Threshold sweep start.",
    )
    parser.add_argument(
        "--threshold-stop",
        type=float,
        default=0.95,
        help="Threshold sweep stop (inclusive).",
    )
    parser.add_argument(
        "--threshold-step",
        type=float,
        default=0.05,
        help="Threshold sweep step.",
    )
    parser.add_argument(
        "--decision-threshold",
        type=float,
        default=0.5,
        help="Decision threshold used in noise analysis.",
    )
    parser.add_argument(
        "--positive-class-index",
        type=int,
        default=0,
        help="Positive class index in both predictions.img_class and labels.label.",
    )
    parser.add_argument(
        "--calibration-bins",
        type=int,
        default=10,
        help="Number of reliability bins used for ECE/MCE.",
    )
    parser.add_argument(
        "--skip-noise",
        action="store_true",
        help="Skip Gaussian noise sensitivity analysis.",
    )
    parser.add_argument(
        "--noise-sigmas",
        nargs="+",
        type=float,
        default=[
            0.0,
            0.01,
            0.03,
            0.05,
            0.1,
            0.2,
            0.3,
            0.4,
            0.5,
            0.6,
            0.75,
            1.0,
        ],
        help="Gaussian noise sigmas for sensitivity analysis.",
    )
    parser.add_argument(
        "--bayesian-mc-passes",
        type=int,
        default=20,
        help="MC forward passes for bayesian noise analysis.",
    )
    return parser.parse_args()


def _threshold_array(start: float, stop: float, step: float) -> np.ndarray:
    if step <= 0:
        raise ValueError("threshold-step must be > 0")
    if stop < start:
        raise ValueError("threshold-stop must be >= threshold-start")

    # Include stop when it lands on a step boundary.
    n = int(round((stop - start) / step))
    return np.array([start + i * step for i in range(n + 1)], dtype=float)


def _run_backend(
    config: dict[str, Any],
    root_dir: Path,
    backend: str,
    clinical_df: pd.DataFrame,
    args: argparse.Namespace,
    out_dir: Path,
) -> dict[str, Any]:
    netcdf_path = ensure_backend_netcdf(
        config=config,
        root_dir=root_dir,
        backend=backend,
        bayesian_mc_passes=int(args.bayesian_mc_passes),
    )

    evaluation = load_backend_evaluation(
        config=config,
        backend=backend,
        class_index=int(args.positive_class_index),
    )

    thresholds = _threshold_array(
        start=float(args.threshold_start),
        stop=float(args.threshold_stop),
        step=float(args.threshold_step),
    )

    summary: dict[str, Any] = {
        "backend": backend,
        "netcdf": str(netcdf_path),
        "source_file": str(evaluation.source_file),
        "uncertainty_metric": evaluation.uncertainty_metric,
    }

    summary["performance"] = run_performance(
        evaluation=evaluation,
        output_dir=out_dir,
        thresholds=thresholds,
    )
    summary["calibration"] = run_calibration(
        evaluation=evaluation,
        output_dir=out_dir,
        bins=int(args.calibration_bins),
    )
    summary["physician"] = run_physician(
        evaluation=evaluation,
        clinical_df=clinical_df,
        output_dir=out_dir,
    )
    summary["longitudinal"] = run_longitudinal(
        evaluation=evaluation,
        clinical_df=clinical_df,
        output_dir=out_dir,
    )

    if args.skip_noise:
        summary["noise"] = {"skipped": True, "reason": "--skip-noise supplied"}
    else:
        try:
            summary["noise"] = run_noise_analysis(
                config=config,
                root_dir=Path(__file__).resolve().parents[1],
                backend=backend,
                output_dir=out_dir,
                class_index=int(args.positive_class_index),
                noise_sigmas=[float(x) for x in args.noise_sigmas],
                threshold=float(args.decision_threshold),
                calibration_bins=int(args.calibration_bins),
                bayesian_mc_passes=int(args.bayesian_mc_passes),
            )
        except Exception as exc:
            summary["noise"] = {
                "skipped": True,
                "reason": f"Noise analysis failed: {exc}",
            }

    write_json(out_dir / "backend_summary.json", summary)
    return summary


def main() -> None:
    args = _parse_args()

    analysis_dir = Path(__file__).resolve().parent
    paths = init_runtime_paths(analysis_dir=analysis_dir, run_name=args.run_name)
    config = load_config(paths.root_dir)
    clinical_df = load_clinical_table(config=config, root_dir=paths.root_dir)

    manifest: dict[str, Any] = {
        "run_dir": str(paths.run_dir),
        "output_root": str(paths.output_root),
        "positive_class_index": int(args.positive_class_index),
        "threshold_sweep": {
            "start": float(args.threshold_start),
            "stop": float(args.threshold_stop),
            "step": float(args.threshold_step),
        },
        "calibration_bins": int(args.calibration_bins),
        "noise_sigmas": [float(x) for x in args.noise_sigmas],
        "backends": {},
    }

    for backend in args.backend:
        out_dir = backend_dir(paths, backend)
        manifest["backends"][backend] = _run_backend(
            config=config,
            root_dir=paths.root_dir,
            backend=backend,
            clinical_df=clinical_df,
            args=args,
            out_dir=out_dir,
        )

    write_json(paths.run_dir / "run_manifest.json", manifest)
    print(f"Analysis complete. Results saved to {paths.run_dir}")


if __name__ == "__main__":
    main()