PCA / t-SNE / UMAP 2D Projection with Confidence Ellipses

#!/usr/bin/env python3
"""FigureCraft Chart Engine: PCA / UMAP 2D Projection with Confidence Ellipses (c14_pca_umap)."""

import os
import sys
import argparse
from pathlib import Path
from typing import List, Dict, Optional
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as patches

try:
    from engines.styles import apply_style, get_palette
except ImportError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from engines.styles import apply_style, get_palette

CHART_ID = "c14_pca_umap"


def compute_confidence_ellipse(x: np.ndarray, y: np.ndarray, confidence: float = 0.95) -> Optional[patches.Ellipse]:
    """Calculates theoretical 95% confidence ellipse from bivariate sample covariance."""
    if len(x) < 3:
        return None

    cov = np.cov(x, y)
    vals, vecs = np.linalg.eigh(cov)
    order = vals.argsort()[::-1]
    vals = vals[order]
    vecs = vecs[:, order]

    # Chi-square critical value: df=2, 0.95 -> 5.991
    # For general confidence: -2 * log(1 - confidence)
    s = -2.0 * np.log(1.0 - confidence)
    width = 2.0 * np.sqrt(s * np.maximum(vals[0], 1e-12))
    height = 2.0 * np.sqrt(s * np.maximum(vals[1], 1e-12))

    theta = np.degrees(np.arctan2(vecs[1, 0], vecs[0, 0]))
    center = (np.mean(x), np.mean(y))

    return patches.Ellipse(xy=center, width=width, height=height, angle=theta)


def render(
    data_path: Optional[str] = None,
    style: str = "nature",
    output_dir: str = "output",
    formats: Optional[List[str]] = None,
    dpi: int = 600,
) -> Dict[str, str]:
    if formats is None:
        formats = ["svg", "pdf", "png", "tiff"]
    if data_path is None:
        data_path = str(Path(__file__).resolve().parents[1] / "data" / f"{CHART_ID}.csv")

    apply_style(style)
    palette = get_palette(style)
    df = pd.read_csv(data_path)
    df = df.dropna(subset=["dim_1", "dim_2", "cluster"])

    clusters = list(df["cluster"].unique())

    fig, ax = plt.subplots(figsize=(6.5, 5.2))

    for i, clus in enumerate(clusters):
        sub_df = df[df["cluster"] == clus]
        x = sub_df["dim_1"].values
        y = sub_df["dim_2"].values
        color = palette[i % len(palette)]
        clus_label = clus.replace("_", " ")

        # Scatter points
        ax.scatter(
            x,
            y,
            c=color,
            label=clus_label,
            s=34,
            alpha=0.8,
            edgecolors="#FFFFFF",
            linewidth=0.5,
            zorder=3,
        )

        # 95% Confidence Ellipse
        ellipse = compute_confidence_ellipse(x, y, confidence=0.95)
        if ellipse:
            ellipse.set_facecolor(color)
            ellipse.set_alpha(0.12)
            ellipse.set_edgecolor(color)
            ellipse.set_linewidth(1.2)
            ellipse.set_linestyle("-")
            ellipse.set_zorder(2)
            ax.add_patch(ellipse)

        # Cluster Centroid label
        cx, cy = np.mean(x), np.mean(y)
        ax.scatter(cx, cy, color=color, s=70, marker="X", edgecolors="#111827", linewidth=0.8, zorder=5)

    ax.set_xlabel("Principal Coordinate 1 (38.4% variance explained)")
    ax.set_ylabel("Principal Coordinate 2 (24.1% variance explained)")
    ax.set_title("Single-Cell Representation Manifold with 95% Confidence Ellipses")
    ax.legend(loc="upper right", frameon=False, fontsize=7.5)

    os.makedirs(output_dir, exist_ok=True)
    generated_files = {}
    for fmt in formats:
        out_path = os.path.join(output_dir, f"{CHART_ID}_{style}.{fmt}")
        fig.savefig(out_path, format=fmt, dpi=dpi, bbox_inches="tight")
        generated_files[fmt] = str(Path(out_path).resolve())

    plt.close(fig)
    return generated_files


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Render PCA / UMAP Projection with Ellipses")
    parser.add_argument("--data", type=str, default=None, help="Path to input CSV dataset")
    parser.add_argument("--style", type=str, default="nature", choices=["nature", "cell", "lancet", "ieee", "economist"], help="Journal style pack")
    parser.add_argument("--output-dir", type=str, default="output", help="Output directory")
    parser.add_argument("--formats", type=str, default="svg,pdf,png,tiff", help="Comma-separated export formats")
    parser.add_argument("--dpi", type=int, default=600, help="Raster DPI")
    args = parser.parse_args()

    fmt_list = [f.strip() for f in args.formats.split(",") if f.strip()]
    results = render(args.data, style=args.style, output_dir=args.output_dir, formats=fmt_list, dpi=args.dpi)
    for fmt, path in results.items():
        print(f"Generated {fmt.upper()}: {path}")