app.py

"""
TDA Explorer - Interactive Topological Data Analysis
by Dr. Milan Amrutkumar Joshi
"""

import gradio as gr
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.patches import Circle as MplCircle
from scipy.spatial.distance import pdist, squareform
from ripser import ripser
import traceback
import warnings

warnings.filterwarnings("ignore")

# ── Color Palette ─────────────────────────────────────────────
PALETTE = {
    "point": "#E63946",
    "edge": "#457B9D",
    "circle": "#A8DADC",
    "h0": "#E63946",
    "h1": "#457B9D",
    "h2": "#2A9D8F",
    "bg": "#F1FAEE",
    "dark": "#1D3557",
}

# ── Dataset Generators ────────────────────────────────────────

def make_circle(n=150, noise=0.05):
    t = np.linspace(0, 2 * np.pi, n, endpoint=False)
    pts = np.column_stack([np.cos(t), np.sin(t)])
    return pts + np.random.normal(0, noise, pts.shape)


def make_figure_eight(n=150, noise=0.05):
    half = n // 2
    t1 = np.linspace(0, 2 * np.pi, half, endpoint=False)
    t2 = np.linspace(0, 2 * np.pi, n - half, endpoint=False)
    left = np.column_stack([np.cos(t1) - 1, np.sin(t1)])
    right = np.column_stack([np.cos(t2) + 1, np.sin(t2)])
    pts = np.vstack([left, right])
    return pts + np.random.normal(0, noise, pts.shape)


def make_two_circles(n=150, noise=0.05):
    half = n // 2
    t1 = np.linspace(0, 2 * np.pi, half, endpoint=False)
    t2 = np.linspace(0, 2 * np.pi, n - half, endpoint=False)
    outer = np.column_stack([np.cos(t1), np.sin(t1)])
    inner = np.column_stack([0.4 * np.cos(t2), 0.4 * np.sin(t2)])
    pts = np.vstack([outer, inner])
    return pts + np.random.normal(0, noise, pts.shape)


def make_clusters(n=150, noise=0.05):
    k = 3
    per = n // k
    centers = np.array([[0, 0], [2.5, 0], [1.25, 2.2]])
    pts = []
    for i in range(k):
        count = per if i < k - 1 else n - per * (k - 1)
        pts.append(centers[i] + np.random.normal(0, 0.15 + noise, (count, 2)))
    return np.vstack(pts)


def make_moons(n=150, noise=0.08):
    half = n // 2
    t1 = np.linspace(0, np.pi, half)
    t2 = np.linspace(0, np.pi, n - half)
    upper = np.column_stack([np.cos(t1), np.sin(t1)])
    lower = np.column_stack([1 - np.cos(t2), 1 - np.sin(t2) - 0.5])
    pts = np.vstack([upper, lower])
    return pts + np.random.normal(0, noise, pts.shape)


def make_sphere(n=200, noise=0.05):
    phi = np.random.uniform(0, 2 * np.pi, n)
    cos_theta = np.random.uniform(-1, 1, n)
    theta = np.arccos(cos_theta)
    x = np.sin(theta) * np.cos(phi)
    y = np.sin(theta) * np.sin(phi)
    z = np.cos(theta)
    pts = np.column_stack([x, y, z])
    return pts + np.random.normal(0, noise, pts.shape)


def make_torus(n=300, noise=0.05):
    R, r = 2.0, 0.7
    theta = np.random.uniform(0, 2 * np.pi, n)
    phi = np.random.uniform(0, 2 * np.pi, n)
    x = (R + r * np.cos(theta)) * np.cos(phi)
    y = (R + r * np.cos(theta)) * np.sin(phi)
    z = r * np.sin(theta)
    pts = np.column_stack([x, y, z])
    return pts + np.random.normal(0, noise, pts.shape)


def make_random(n=150, noise=0.05):
    return np.random.uniform(-2, 2, (n, 2))


DATASETS = {
    "Circle (H1=1)": make_circle,
    "Figure-8 (H1=2)": make_figure_eight,
    "Concentric Circles": make_two_circles,
    "3 Clusters (H0=3)": make_clusters,
    "Moons": make_moons,
    "Sphere 3D (H2=1)": make_sphere,
    "Torus 3D (H1=2, H2=1)": make_torus,
    "Random Noise": make_random,
}

# ── Shared state ──────────────────────────────────────────────
_state = {"pts": None, "result": None, "auto_eps": 0.5}


# ── Visualization Functions ───────────────────────────────────

def plot_point_cloud(pts):
    fig, ax = plt.subplots(figsize=(6, 6), dpi=100)
    fig.patch.set_facecolor("#FAFAFA")
    ax.set_facecolor("#FAFAFA")
    dim = pts.shape[1]
    if dim >= 3:
        ax.remove()
        ax = fig.add_subplot(111, projection="3d")
        ax.scatter(pts[:, 0], pts[:, 1], pts[:, 2],
                   c=PALETTE["point"], s=12, alpha=0.7, edgecolors="none")
        ax.set_xlabel("X", fontsize=9)
        ax.set_ylabel("Y", fontsize=9)
        ax.set_zlabel("Z", fontsize=9)
        ax.xaxis.pane.fill = False
        ax.yaxis.pane.fill = False
        ax.zaxis.pane.fill = False
    else:
        ax.scatter(pts[:, 0], pts[:, 1],
                   c=PALETTE["point"], s=18, alpha=0.7, edgecolors="none")
        ax.set_aspect("equal")
    ax.set_title(f"Point Cloud  ({len(pts)} pts, {dim}D)", fontsize=12,
                 fontweight="bold", color=PALETTE["dark"])
    if dim < 3:
        for spine in ax.spines.values():
            spine.set_visible(False)
    ax.tick_params(labelsize=8)
    plt.tight_layout()
    return fig


def plot_persistence(result):
    fig, ax = plt.subplots(figsize=(6, 6), dpi=100)
    fig.patch.set_facecolor("#FAFAFA")
    ax.set_facecolor("#FAFAFA")
    dgms = result["dgms"]
    colors = [PALETTE["h0"], PALETTE["h1"], PALETTE["h2"]]
    labels = ["H0 (components)", "H1 (loops)", "H2 (voids)"]
    all_finite = []
    for dgm in dgms:
        finite = dgm[dgm[:, 1] < np.inf]
        if len(finite):
            all_finite.append(finite)
    mx = max(np.vstack(all_finite).max(), 0.1) * 1.15 if all_finite else 1.0
    ax.plot([0, mx], [0, mx], "--", color="#CCC", linewidth=1, zorder=0)
    for dim_i, dgm in enumerate(dgms):
        finite = dgm[dgm[:, 1] < np.inf]
        inf_pts = dgm[dgm[:, 1] == np.inf]
        c = colors[dim_i % len(colors)]
        lbl = labels[dim_i] if dim_i < len(labels) else f"H{dim_i}"
        if len(finite):
            ax.scatter(finite[:, 0], finite[:, 1], c=c, s=30, alpha=0.75,
                       edgecolors="white", linewidths=0.5, label=lbl, zorder=2)
        if len(inf_pts):
            ax.scatter(inf_pts[:, 0], [mx * 0.95] * len(inf_pts), c=c, s=60,
                       marker="^", alpha=0.9, edgecolors="white", linewidths=0.5,
                       zorder=3)
    ax.set_xlabel("Birth", fontsize=11)
    ax.set_ylabel("Death", fontsize=11)
    ax.set_title("Persistence Diagram", fontsize=13, fontweight="bold",
                 color=PALETTE["dark"])
    ax.legend(fontsize=9, framealpha=0.9)
    ax.set_xlim(-0.02 * mx, mx)
    ax.set_ylim(-0.02 * mx, mx)
    for spine in ax.spines.values():
        spine.set_color("#DDD")
    ax.tick_params(labelsize=8)
    plt.tight_layout()
    return fig


def plot_barcode(result):
    fig, ax = plt.subplots(figsize=(8, 5), dpi=100)
    fig.patch.set_facecolor("#FAFAFA")
    ax.set_facecolor("#FAFAFA")
    dgms = result["dgms"]
    colors = [PALETTE["h0"], PALETTE["h1"], PALETTE["h2"]]
    labels_list = ["H0", "H1", "H2"]
    y = 0
    yticks, ytick_labels = [], []
    max_death = 0
    for dgm in dgms:
        finite = dgm[dgm[:, 1] < np.inf]
        if len(finite):
            max_death = max(max_death, finite[:, 1].max())
    cap = max_death * 1.2 if max_death > 0 else 1.0
    for dim_i, dgm in enumerate(dgms):
        c = colors[dim_i % len(colors)]
        lbl = labels_list[dim_i] if dim_i < len(labels_list) else f"H{dim_i}"
        sorted_dgm = dgm[np.argsort(dgm[:, 1] - dgm[:, 0])[::-1]]
        for birth, death in sorted_dgm:
            d = death if death < np.inf else cap
            ax.plot([birth, d], [y, y], color=c, linewidth=2.5, solid_capstyle="round")
            if death == np.inf:
                ax.plot(d, y, ">", color=c, markersize=5)
            y += 1
        if len(dgm):
            yticks.append(y - len(dgm) / 2)
            ytick_labels.append(lbl)
    ax.set_yticks(yticks)
    ax.set_yticklabels(ytick_labels, fontsize=10, fontweight="bold")
    ax.set_xlabel("Filtration Value", fontsize=11)
    ax.set_title("Persistence Barcode", fontsize=13, fontweight="bold",
                 color=PALETTE["dark"])
    ax.invert_yaxis()
    for spine in ["top", "right"]:
        ax.spines[spine].set_visible(False)
    for spine in ["bottom", "left"]:
        ax.spines[spine].set_color("#DDD")
    ax.tick_params(labelsize=8)
    plt.tight_layout()
    return fig


def plot_filtration(pts, epsilon):
    if pts.shape[1] > 2:
        pts_2d = pts[:, :2]
        title_suffix = " (projected to 2D)"
    else:
        pts_2d = pts
        title_suffix = ""

    fig, ax = plt.subplots(figsize=(6, 6), dpi=100)
    fig.patch.set_facecolor("#FAFAFA")
    ax.set_facecolor("#FAFAFA")

    dist_mat = squareform(pdist(pts_2d))

    for pt in pts_2d:
        circle = MplCircle(pt, epsilon / 2, alpha=0.06,
                           color=PALETTE["circle"], linewidth=0)
        ax.add_patch(circle)

    edges = []
    for i in range(len(pts_2d)):
        for j in range(i + 1, len(pts_2d)):
            if dist_mat[i, j] <= epsilon:
                edges.append([pts_2d[i], pts_2d[j]])

    if epsilon > 0 and len(pts_2d) <= 200:
        for i in range(len(pts_2d)):
            for j in range(i + 1, len(pts_2d)):
                if dist_mat[i, j] > epsilon:
                    continue
                for k in range(j + 1, len(pts_2d)):
                    if dist_mat[i, k] <= epsilon and dist_mat[j, k] <= epsilon:
                        tri = plt.Polygon(
                            [pts_2d[i], pts_2d[j], pts_2d[k]],
                            alpha=0.08, color=PALETTE["h1"], linewidth=0
                        )
                        ax.add_patch(tri)

    if edges:
        lc = LineCollection(edges, colors=PALETTE["edge"], alpha=0.35,
                            linewidths=0.8, zorder=2)
        ax.add_collection(lc)

    ax.scatter(pts_2d[:, 0], pts_2d[:, 1], c=PALETTE["point"], s=20,
               zorder=5, edgecolors="white", linewidths=0.3, alpha=0.9)

    n_edges = len(edges)
    ax.set_title(
        f"Vietoris-Rips Complex  \u03b5={epsilon:.3f}  |  {n_edges} edges{title_suffix}",
        fontsize=11, fontweight="bold", color=PALETTE["dark"]
    )
    ax.set_aspect("equal")
    ax.autoscale()
    margin = epsilon * 0.6 + 0.1
    xlim = ax.get_xlim()
    ylim = ax.get_ylim()
    ax.set_xlim(xlim[0] - margin, xlim[1] + margin)
    ax.set_ylim(ylim[0] - margin, ylim[1] + margin)
    for spine in ax.spines.values():
        spine.set_visible(False)
    ax.tick_params(labelsize=7, colors="#999")
    plt.tight_layout()
    return fig


def betti_summary(result, eps):
    dgms = result["dgms"]
    lines = []
    for dim_i, dgm in enumerate(dgms):
        alive = sum(1 for b, d in dgm if b <= eps and (d > eps or d == np.inf))
        name = ["Components (H0)", "Loops (H1)", "Voids (H2)"][dim_i] if dim_i < 3 else f"H{dim_i}"
        lines.append(f"**\u03b2{dim_i} = {alive}**  \u2014  {name}")
    total_features = sum(len(dgm) for dgm in dgms)
    lines.append(f"\n---\nTotal features: **{total_features}** | \u03b5 = **{eps:.3f}**")
    return "\n\n".join(lines)


# ── Callbacks ─────────────────────────────────────────────────

def generate_data(dataset_name, n_points, noise):
    try:
        n_points = int(n_points)
        gen = DATASETS.get(dataset_name, make_circle)
        pts = gen(n=n_points, noise=noise)
        _state["pts"] = pts

        max_dim = 2 if pts.shape[1] >= 3 else 1
        result = ripser(pts, maxdim=max_dim, thresh=3.0)
        _state["result"] = result

        dist_mat = squareform(pdist(pts))
        np.fill_diagonal(dist_mat, np.inf)
        nn_dists = dist_mat.min(axis=1)
        auto_eps = float(np.median(nn_dists) * 2)
        _state["auto_eps"] = auto_eps

        fig_cloud = plot_point_cloud(pts)
        fig_diag = plot_persistence(result)
        fig_barcode = plot_barcode(result)
        fig_filt = plot_filtration(pts, auto_eps)
        summary = betti_summary(result, auto_eps)

        return fig_cloud, fig_diag, fig_barcode, fig_filt, summary
    except Exception as e:
        traceback.print_exc()
        empty = plt.figure(figsize=(4, 4))
        return empty, empty, empty, empty, f"Error: {e}"


def update_filtration(epsilon):
    try:
        pts = _state.get("pts")
        result = _state.get("result")
        if pts is None or result is None:
            return None, "*Click Generate & Analyze first*"
        fig = plot_filtration(pts, epsilon)
        summary = betti_summary(result, epsilon)
        return fig, summary
    except Exception as e:
        traceback.print_exc()
        return None, f"Error: {e}"


# ── CSS ───────────────────────────────────────────────────────

CUSTOM_CSS = """
.gradio-container {
    max-width: 1200px !important;
    font-family: 'Inter', 'Segoe UI', sans-serif;
}
h1 { color: #1D3557; }
.gr-button-primary {
    background: linear-gradient(135deg, #E63946, #457B9D) !important;
    border: none !important;
}
footer { display: none !important; }
"""

# ── Gradio App ────────────────────────────────────────────────

with gr.Blocks(css=CUSTOM_CSS, title="TDA Explorer", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Topological Data Analysis Explorer
        **Explore persistent homology, Betti numbers & simplicial complexes interactively.**

        Built by [Dr. Milan Joshi](https://huggingface.co/mlnjsh) | Research: Persistent Homology, TDA for ML
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Generate Point Cloud")
            dataset_dd = gr.Dropdown(
                choices=list(DATASETS.keys()),
                value="Circle (H1=1)",
                label="Dataset",
            )
            n_slider = gr.Slider(50, 500, value=150, step=10, label="Points")
            noise_slider = gr.Slider(0.0, 0.3, value=0.05, step=0.01, label="Noise")
            gen_btn = gr.Button("Generate & Analyze", variant="primary", size="lg")

            gr.Markdown("---")
            gr.Markdown("### Filtration Control")
            eps_slider = gr.Slider(
                0.0, 4.0, value=0.5, step=0.01,
                label="Epsilon - connectivity radius",
            )
            betti_md = gr.Markdown("*Click Generate & Analyze to start*")

        with gr.Column(scale=3):
            with gr.Row():
                cloud_plot = gr.Plot(label="Point Cloud")
                filt_plot = gr.Plot(label="Filtration Complex")
            with gr.Row():
                diag_plot = gr.Plot(label="Persistence Diagram")
                barcode_plot = gr.Plot(label="Persistence Barcode")

    with gr.Accordion("What is TDA? (click to learn)", open=False):
        gr.Markdown(
            """
            ### Topological Data Analysis in a Nutshell

            TDA studies the **shape** of data using tools from algebraic topology.

            | Concept | What it captures | Example |
            |---|---|---|
            | **H0** (connected components) | Clusters / groups | 3 clusters -> B0 = 3 |
            | **H1** (loops / tunnels) | Circular holes | Ring of points -> B1 = 1 |
            | **H2** (voids / cavities) | Enclosed volumes | Sphere surface -> B2 = 1 |

            #### How it works
            1. Start with a **point cloud** (your data).
            2. Grow balls of radius e around each point.
            3. When balls overlap, connect points to build a **simplicial complex**.
            4. Track which features (components, loops, voids) **appear** and **disappear** as e increases.
            5. Plot births vs deaths: the **persistence diagram**. Features far from the diagonal are real; near diagonal = noise.

            #### Betti Numbers
            At filtration value e, the **Betti numbers** (B0, B1, B2, ...) count each type of feature currently alive.

            #### Why it matters for ML
            - **Feature engineering**: Persistent homology gives robust shape descriptors.
            - **Outlier detection**: Noise features die quickly; real structure persists.
            - **Clustering**: H0 persistence reveals natural cluster counts.

            *This demo uses [Ripser](https://ripser.scikit-tda.org/) for fast Vietoris-Rips persistent homology computation.*
            """
        )

    # ── Wiring (no slider update, no demo.load) ──────────────
    gen_btn.click(
        fn=generate_data,
        inputs=[dataset_dd, n_slider, noise_slider],
        outputs=[cloud_plot, diag_plot, barcode_plot, filt_plot, betti_md],
    )

    eps_slider.change(
        fn=update_filtration,
        inputs=[eps_slider],
        outputs=[filt_plot, betti_md],
    )


if __name__ == "__main__":
    demo.launch()