app.py

"""
ParticleNet — Graph Neural Network for Particle Collision Event Classification
==============================================================================
Author  : Your Name  (edit this before pushing to Hugging Face)
Project : AI + Physics Portfolio — Project 7
Dataset : Synthetic CERN-style jet data (self-generated, no download needed)
Model   : 3-layer Graph Convolutional Network (GCN) built with PyTorch Geometric
Demo    : Hugging Face Spaces · Gradio interface

Physics context
---------------
At colliders like the LHC at CERN, protons smash together millions of times per
second, producing sprays of particles called *jets*. Identifying what kind of
particle initiated a jet — a quark, gluon, W boson, top quark, or Higgs boson —
is a fundamental task in particle physics and a perfect graph learning problem:
each jet is naturally a graph where particles are nodes and their proximity in
momentum space defines the edges.

This demo trains a small GCN on synthetic data that mimics real jet substructure
features (pT, eta, phi, charge, particle ID), then lets you generate a random
event and watch the model classify it in real time, with full visual explanation.
"""

import gradio as gr
import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, global_mean_pool
from torch_geometric.data import Data, Batch
import numpy as np
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import os, json, time

# ─────────────────────────────────────────────────────────────
# 0.  CONSTANTS  — physics-inspired class labels
# ─────────────────────────────────────────────────────────────
CLASS_NAMES  = ["Gluon jet", "Light-quark jet", "W boson jet",
                "Top quark jet", "Higgs boson jet"]
CLASS_COLORS = ["#378ADD", "#1D9E75", "#EF9F27", "#D85A30", "#7F77DD"]
NUM_CLASSES  = len(CLASS_NAMES)

# Node feature names — each particle in the jet carries these 5 features
FEATURE_NAMES = ["transverse momentum (pT)",
                 "pseudorapidity (η)",
                 "azimuthal angle (φ)",
                 "electric charge",
                 "particle type ID"]

# ─────────────────────────────────────────────────────────────
# 1.  SYNTHETIC DATA GENERATOR
#     Produces CERN-style jet graphs with realistic feature
#     distributions for each class.  No internet required.
# ─────────────────────────────────────────────────────────────
def generate_jet(label: int, seed: int | None = None) -> Data:
    """
    Generate one synthetic jet graph.

    Each jet has between 8 and 24 constituent particles (nodes).
    Edges connect every particle to its 3 nearest neighbours in
    (η, φ) space — this is how real jet algorithms work.

    Feature distributions are loosely inspired by particle physics:
      - Gluon jets: many soft particles, wide angular spread
      - Quark jets: fewer, harder particles
      - W jets:     two sub-clusters (W → qq decay signature)
      - Top jets:   three sub-clusters (t → bqq)
      - Higgs jets: two sub-clusters with b-quark enrichment
    """
    rng = np.random.default_rng(seed)

    # --- number of particles varies by jet type ---
    n_particles_map = {0: (12, 24), 1: (8, 18), 2: (10, 20),
                       3: (14, 24), 4: (10, 20)}
    lo, hi = n_particles_map[label]
    n = rng.integers(lo, hi + 1)

    # --- pT spectrum: power-law (harder for quarks/bosons) ---
    alpha = {0: 3.5, 1: 2.8, 2: 2.5, 3: 2.2, 4: 2.4}[label]
    pt = rng.pareto(alpha, n) * 10 + 1          # GeV, >1

    # --- angular spread in (eta, phi) ---
    spread = {0: 0.5, 1: 0.3, 2: 0.25, 3: 0.35, 4: 0.28}[label]

    if label in (2, 3, 4):
        # Multi-prong: split particles into sub-clusters
        n_prongs = {2: 2, 3: 3, 4: 2}[label]
        # Place cluster centres
        centres_eta = rng.uniform(-spread, spread, n_prongs)
        centres_phi = rng.uniform(-spread, spread, n_prongs)
        assign = rng.integers(0, n_prongs, n)
        eta = centres_eta[assign] + rng.normal(0, spread / 3, n)
        phi = centres_phi[assign] + rng.normal(0, spread / 3, n)
    else:
        eta = rng.normal(0, spread, n)
        phi = rng.normal(0, spread, n)

    # --- charge: mostly neutral for gluons, mix for others ---
    charge_prob = {0: 0.2, 1: 0.45, 2: 0.5, 3: 0.55, 4: 0.4}[label]
    charge = rng.choice([-1, 0, 1], n,
                        p=[charge_prob / 2, 1 - charge_prob, charge_prob / 2])

    # --- particle type ID: 0=photon,1=neutral hadron,2=charged hadron,3=electron,4=muon ---
    pid_probs = {
        0: [0.15, 0.40, 0.35, 0.07, 0.03],
        1: [0.10, 0.25, 0.50, 0.10, 0.05],
        2: [0.08, 0.20, 0.55, 0.12, 0.05],
        3: [0.05, 0.15, 0.60, 0.12, 0.08],
        4: [0.12, 0.28, 0.48, 0.09, 0.03],
    }
    pid = rng.choice(5, n, p=pid_probs[label])

    # --- build node feature matrix (n × 5) ---
    # Normalise to roughly [-1, 1] so the GCN trains easily
    pt_norm    = np.log1p(pt) / 5.0        # log scale for pT
    eta_norm   = eta / 1.0
    phi_norm   = phi / np.pi
    chg_norm   = charge.astype(float)
    pid_norm   = pid.astype(float) / 4.0   # [0,1]

    x = torch.tensor(
        np.stack([pt_norm, eta_norm, phi_norm, chg_norm, pid_norm], axis=1),
        dtype=torch.float
    )

    # --- k-NN graph in (eta, phi) space, k=3 ---
    coords = np.stack([eta, phi], axis=1)
    from sklearn.neighbors import NearestNeighbors
    k = min(3, n - 1)
    nbrs = NearestNeighbors(n_neighbors=k + 1).fit(coords)
    _, indices = nbrs.kneighbors(coords)
    src, dst = [], []
    for i, neighbours in enumerate(indices):
        for j in neighbours[1:]:          # skip self
            src.append(i); dst.append(j)
            src.append(j); dst.append(i)  # undirected

    edge_index = torch.tensor([src, dst], dtype=torch.long)
    y          = torch.tensor([label], dtype=torch.long)

    # Store raw coords for visualisation
    data = Data(x=x, edge_index=edge_index, y=y)
    data.pt  = torch.tensor(pt,  dtype=torch.float)
    data.eta = torch.tensor(eta, dtype=torch.float)
    data.phi = torch.tensor(phi, dtype=torch.float)
    return data


def generate_dataset(n_per_class: int = 200, seed: int = 42) -> list[Data]:
    """Create a balanced training set with n_per_class jets per category."""
    dataset = []
    for label in range(NUM_CLASSES):
        for i in range(n_per_class):
            dataset.append(generate_jet(label, seed=seed * 1000 + label * 100 + i))
    rng = np.random.default_rng(seed)
    perm = rng.permutation(len(dataset))
    return [dataset[i] for i in perm]


# ─────────────────────────────────────────────────────────────
# 2.  GCN MODEL
#     Three graph convolutional layers followed by global mean
#     pooling and a linear classifier head.
# ─────────────────────────────────────────────────────────────
class ParticleGCN(torch.nn.Module):
    """
    A 3-layer Graph Convolutional Network for jet classification.

    Architecture
    ------------
    Input (5 features per particle node)
        → GCNConv(5 → 64)  + ReLU + Dropout(0.2)
        → GCNConv(64 → 128) + ReLU + Dropout(0.2)
        → GCNConv(128 → 64) + ReLU
        → GlobalMeanPool (aggregate all particles into one jet vector)
        → Linear(64 → 32)  + ReLU
        → Linear(32 → 5)   (one logit per class)

    Why GCNs for physics?
    ---------------------
    Unlike CNNs (which need a grid) or RNNs (which need a sequence),
    GCNs operate on arbitrary graphs — perfect for jets where the
    number of particles varies and their spatial relationships matter.
    The message-passing mechanism lets each particle "talk to" its
    neighbours, building up a representation of local jet substructure
    before the global pool summarises the whole event.
    """
    def __init__(self, in_channels: int = 5, hidden: int = 64,
                 out_channels: int = NUM_CLASSES):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden)
        self.conv2 = GCNConv(hidden, hidden * 2)
        self.conv3 = GCNConv(hidden * 2, hidden)
        self.lin1  = torch.nn.Linear(hidden, 32)
        self.lin2  = torch.nn.Linear(32, out_channels)
        self.drop  = torch.nn.Dropout(p=0.2)

    def forward(self, x, edge_index, batch):
        # Message passing through the jet graph
        x = F.relu(self.conv1(x, edge_index))
        x = self.drop(x)
        x = F.relu(self.conv2(x, edge_index))
        x = self.drop(x)
        x = F.relu(self.conv3(x, edge_index))

        # Pool all particle embeddings into a single jet embedding
        x = global_mean_pool(x, batch)

        # Classification head
        x = F.relu(self.lin1(x))
        x = self.lin2(x)
        return x


# ─────────────────────────────────────────────────────────────
# 3.  TRAINING
#     Runs once at startup; saves model to disk for reuse.
# ─────────────────────────────────────────────────────────────
MODEL_PATH = "particlenet_gcn.pt"
HISTORY_PATH = "training_history.json"

def train_model(n_per_class: int = 300, epochs: int = 60,
                lr: float = 1e-3) -> tuple[ParticleGCN, dict]:
    """Train the GCN on synthetic jet data and return model + history."""
    print("Generating synthetic jet dataset...")
    dataset = generate_dataset(n_per_class=n_per_class)

    # 80/20 train-val split
    split = int(0.8 * len(dataset))
    train_data, val_data = dataset[:split], dataset[split:]

    def make_batch(subset):
        return Batch.from_data_list(subset)

    device = torch.device("cpu")   # CPU is fine for this model size
    model  = ParticleGCN().to(device)
    opt    = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
    sched  = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)

    history = {"train_loss": [], "val_loss": [],
               "train_acc": [], "val_acc": []}

    print(f"Training on {len(train_data)} jets for {epochs} epochs...")
    for epoch in range(epochs):
        # ── train ──
        model.train()
        batch = make_batch(train_data)
        batch = batch.to(device)
        opt.zero_grad()
        out  = model(batch.x, batch.edge_index, batch.batch)
        loss = F.cross_entropy(out, batch.y)
        loss.backward()
        opt.step()
        sched.step()
        train_loss = loss.item()
        train_acc  = (out.argmax(1) == batch.y).float().mean().item()

        # ── validate ──
        model.eval()
        with torch.no_grad():
            vbatch = make_batch(val_data).to(device)
            vout   = model(vbatch.x, vbatch.edge_index, vbatch.batch)
            vloss  = F.cross_entropy(vout, vbatch.y).item()
            vacc   = (vout.argmax(1) == vbatch.y).float().mean().item()

        history["train_loss"].append(round(train_loss, 4))
        history["val_loss"].append(round(vloss, 4))
        history["train_acc"].append(round(train_acc, 4))
        history["val_acc"].append(round(vacc, 4))

        if (epoch + 1) % 10 == 0:
            print(f"  Epoch {epoch+1:3d}/{epochs}  "
                  f"loss={train_loss:.4f}  val_loss={vloss:.4f}  "
                  f"val_acc={vacc:.2%}")

    torch.save(model.state_dict(), MODEL_PATH)
    with open(HISTORY_PATH, "w") as f:
        json.dump(history, f)
    print(f"Model saved to {MODEL_PATH}")
    return model, history


def load_or_train() -> tuple[ParticleGCN, dict]:
    """Load pretrained model if available, otherwise train from scratch."""
    model = ParticleGCN()
    if os.path.exists(MODEL_PATH) and os.path.exists(HISTORY_PATH):
        print("Loading pre-trained model...")
        model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
        with open(HISTORY_PATH) as f:
            history = json.load(f)
    else:
        model, history = train_model()
    model.eval()
    return model, history


# ─────────────────────────────────────────────────────────────
# 4.  VISUALISATION HELPERS
# ─────────────────────────────────────────────────────────────
def plot_jet_graph(data: Data, pred_label: int,
                   true_label: int, probs: np.ndarray) -> go.Figure:
    """
    Build an interactive Plotly figure showing:
      Left  — the jet graph in (η, φ) space, nodes sized by pT
      Right — the class probability bar chart
    """
    eta = data.eta.numpy()
    phi = data.phi.numpy()
    pt  = data.pt.numpy()
    src, dst = data.edge_index.numpy()

    fig = make_subplots(
        rows=1, cols=2,
        column_widths=[0.6, 0.4],
        subplot_titles=[
            f"Jet graph — {len(eta)} particles in (η, φ) space",
            "Class probabilities"
        ]
    )

    # ── edge traces ──
    ex, ey = [], []
    for s, d in zip(src, dst):
        ex += [eta[s], eta[d], None]
        ey += [phi[s], phi[d], None]
    fig.add_trace(
        go.Scatter(x=ex, y=ey, mode="lines",
                   line=dict(color="#B4B2A9", width=0.8),
                   hoverinfo="skip", name="Edges"),
        row=1, col=1
    )

    # ── node traces, coloured by particle type ──
    pid_labels = ["Photon", "Neutral hadron", "Charged hadron",
                  "Electron", "Muon"]
    pid_colors = ["#EF9F27", "#5DCAA5", "#378ADD", "#D85A30", "#7F77DD"]
    pid_arr    = (data.x[:, 4].numpy() * 4).round().astype(int)

    for pid_val in range(5):
        mask = pid_arr == pid_val
        if not mask.any():
            continue
        fig.add_trace(
            go.Scatter(
                x=eta[mask], y=phi[mask],
                mode="markers",
                marker=dict(
                    size=np.clip(np.log1p(pt[mask]) * 6, 4, 22),
                    color=pid_colors[pid_val],
                    line=dict(width=0.8, color="#2C2C2A"),
                    opacity=0.85
                ),
                name=pid_labels[pid_val],
                hovertemplate=(
                    f"<b>{pid_labels[pid_val]}</b><br>"
                    "η = %{x:.3f}<br>φ = %{y:.3f}<br>"
                    "pT ≈ %{customdata:.1f} GeV<extra></extra>"
                ),
                customdata=pt[mask]
            ),
            row=1, col=1
        )

    # ── probability bars ──
    bar_colors = [
        CLASS_COLORS[i] if i == pred_label
        else "#D3D1C7"
        for i in range(NUM_CLASSES)
    ]
    fig.add_trace(
        go.Bar(
            x=probs * 100,
            y=CLASS_NAMES,
            orientation="h",
            marker_color=bar_colors,
            text=[f"{p*100:.1f}%" for p in probs],
            textposition="outside",
            hovertemplate="%{y}: %{x:.2f}%<extra></extra>",
            name="Probabilities"
        ),
        row=1, col=2
    )

    # ── ground truth marker on bar chart ──
    fig.add_vline(x=0, row=1, col=2, line_width=0)  # dummy for spacing

    correct = pred_label == true_label
    result_color = "#1D9E75" if correct else "#D85A30"
    result_text  = "CORRECT" if correct else "WRONG"

    fig.update_layout(
        title=dict(
            text=(
                f"<b>Prediction: {CLASS_NAMES[pred_label]}</b>  "
                f"<span style='color:{result_color}'>[{result_text}]</span>  "
                f"· True label: {CLASS_NAMES[true_label]}"
            ),
            font_size=15, x=0.02
        ),
        showlegend=True,
        legend=dict(x=0.01, y=-0.15, orientation="h",
                    font_size=11, bgcolor="rgba(0,0,0,0)"),
        height=480,
        margin=dict(l=40, r=40, t=60, b=80),
        plot_bgcolor="rgba(0,0,0,0)",
        paper_bgcolor="rgba(0,0,0,0)",
        font=dict(color="#3d3d3a"),
        xaxis=dict(title="Pseudorapidity (η)", gridcolor="#E8E6DF",
                   zeroline=True, zerolinecolor="#B4B2A9"),
        yaxis=dict(title="Azimuthal angle φ (rad)", gridcolor="#E8E6DF"),
        xaxis2=dict(title="Probability (%)", range=[0, 115],
                    gridcolor="#E8E6DF"),
        yaxis2=dict(autorange="reversed"),
        bargap=0.25,
    )
    return fig


def plot_training_history(history: dict) -> go.Figure:
    """Plot training and validation loss + accuracy curves."""
    epochs = list(range(1, len(history["train_loss"]) + 1))

    fig = make_subplots(
        rows=1, cols=2,
        subplot_titles=["Loss (cross-entropy)", "Accuracy"]
    )
    for split, color, dash in [("train", "#378ADD", "solid"),
                                ("val",   "#1D9E75", "dash")]:
        fig.add_trace(go.Scatter(
            x=epochs, y=history[f"{split}_loss"],
            name=f"{split} loss",
            line=dict(color=color, dash=dash, width=2),
            mode="lines"
        ), row=1, col=1)
        fig.add_trace(go.Scatter(
            x=epochs, y=[v * 100 for v in history[f"{split}_acc"]],
            name=f"{split} accuracy",
            line=dict(color=color, dash=dash, width=2),
            mode="lines"
        ), row=1, col=2)

    final_val_acc = history["val_acc"][-1] * 100
    fig.update_layout(
        title=dict(
            text=f"Training curves — final validation accuracy: {final_val_acc:.1f}%",
            font_size=14, x=0.02
        ),
        height=320,
        margin=dict(l=40, r=40, t=50, b=40),
        plot_bgcolor="rgba(0,0,0,0)",
        paper_bgcolor="rgba(0,0,0,0)",
        font=dict(color="#3d3d3a"),
        legend=dict(orientation="h", y=-0.15, font_size=11,
                    bgcolor="rgba(0,0,0,0)"),
        xaxis=dict(title="Epoch", gridcolor="#E8E6DF"),
        yaxis=dict(title="Loss",  gridcolor="#E8E6DF"),
        xaxis2=dict(title="Epoch", gridcolor="#E8E6DF"),
        yaxis2=dict(title="Accuracy (%)", gridcolor="#E8E6DF"),
    )
    return fig


def physics_explanation(label: int) -> str:
    """Return a short physics description of each jet class."""
    explanations = {
        0: (
            "**Gluon jet** — Gluons are the force carriers of the strong nuclear "
            "force (QCD). They produce jets with many soft, wide-angle particles "
            "because gluons radiate more than quarks (higher colour charge). "
            "Typical at the LHC: ~70% of jets in inclusive samples are gluon jets."
        ),
        1: (
            "**Light-quark jet** — Up, down, or strange quarks produce narrower, "
            "harder jets with fewer particles. The jet charge and the ratio of "
            "charged to neutral particles help distinguish these from gluon jets."
        ),
        2: (
            "**W boson jet** — When a highly boosted W decays hadronically "
            "(W → qq̄), both daughter quarks are caught inside a single large-R "
            "jet. This creates a distinctive two-prong substructure visible in "
            "the (η, φ) graph as two clusters of particles."
        ),
        3: (
            "**Top quark jet** — The heaviest known elementary particle (173 GeV). "
            "A boosted top decays as t → bW → bqq̄, producing a *three-prong* "
            "substructure. The b-quark sub-jet leaves a secondary vertex signature "
            "that b-taggers exploit."
        ),
        4: (
            "**Higgs boson jet** — The Higgs decays predominantly to bb̄ at low "
            "mass. In the boosted regime both b-quarks merge into one fat jet. "
            "Like the W but enriched in b-quarks. Identifying these jets is "
            "crucial for measuring Higgs couplings at the LHC."
        ),
    }
    return explanations[label]


# ─────────────────────────────────────────────────────────────
# 5.  GRADIO INTERFACE
# ─────────────────────────────────────────────────────────────
print("Loading / training ParticleNet GCN...")
MODEL, HISTORY = load_or_train()
print("Ready.")

# Pre-generate a pool of events we can index into with a slider
POOL_SIZE = 500
EVENT_POOL = [
    generate_jet(label=i % NUM_CLASSES, seed=9999 + i)
    for i in range(POOL_SIZE)
]
# Shuffle the pool so classes are not simply in order
rng = np.random.default_rng(777)
pool_perm = rng.permutation(POOL_SIZE)
EVENT_POOL = [EVENT_POOL[i] for i in pool_perm]


def classify_event(event_index: int):
    """
    Core inference function called by Gradio.
    Returns (jet_graph_figure, explanation_text, training_figure).
    """
    data = EVENT_POOL[int(event_index)]

    # ── run GCN inference ──
    with torch.no_grad():
        batch = Batch.from_data_list([data])
        logits = MODEL(batch.x, batch.edge_index, batch.batch)
        probs  = F.softmax(logits, dim=1).numpy()[0]
        pred   = int(probs.argmax())
        true   = int(data.y.item())

    jet_fig   = plot_jet_graph(data, pred, true, probs)
    train_fig = plot_training_history(HISTORY)

    n_particles = data.x.shape[0]
    n_edges     = data.edge_index.shape[1] // 2
    confidence  = probs[pred] * 100

    expl = (
        f"### Event #{int(event_index)+1}  ·  {n_particles} particles  ·  "
        f"{n_edges} edges\n\n"
        f"**Model prediction:** {CLASS_NAMES[pred]}  "
        f"(confidence: {confidence:.1f}%)\n\n"
        f"**True label:** {CLASS_NAMES[true]}\n\n"
        f"---\n\n"
        f"{physics_explanation(true)}\n\n"
        f"---\n\n"
        f"**How the GCN works:** Each particle sends a message to its "
        f"{min(3, n_particles-1)} nearest neighbours in (η, φ) space. "
        f"After 3 rounds of message passing, all particle embeddings are "
        f"averaged (global mean pool) to produce a single 64-dimensional "
        f"jet embedding. A 2-layer MLP then maps this to the 5 class logits. "
        f"Node size in the graph is proportional to log(pT) — larger nodes "
        f"carry more transverse momentum."
    )
    return jet_fig, expl, train_fig


# ── custom CSS ──
CSS = """
#title-block { padding: 1.2rem 0 0.4rem; }
#title-block h1 { font-size: 1.6rem; font-weight: 500; margin: 0; }
#title-block p  { font-size: 0.92rem; color: #5F5E5A; margin: 0.3rem 0 0; }
.badge {
    display: inline-block;
    font-size: 11px; font-weight: 500; padding: 2px 9px;
    border-radius: 99px; margin-right: 4px;
    background: #E6F1FB; color: #0C447C;
}
.gr-button-primary { background: #185FA5 !important; }
footer { display: none !important; }
"""

with gr.Blocks(css=CSS, theme=gr.themes.Default(
    primary_hue="blue",
    font=gr.themes.GoogleFont("Inter")
)) as demo:

    # ── header ──
    gr.HTML("""
    <div id="title-block">
      <h1>ParticleNet — GNN Particle Collision Classifier</h1>
      <p>
        <span class="badge">Graph Neural Networks</span>
        <span class="badge">CERN-style Jets</span>
        <span class="badge">PyTorch Geometric</span>
        <span class="badge">HEP Physics</span>
        <span class="badge">AI + Physics Portfolio</span>
      </p>
      <p style="margin-top:0.6rem">
        A 3-layer Graph Convolutional Network classifies particle collision events
        (jets) into 5 categories — the same task performed at the LHC at CERN.
        Each jet is a graph: particles are nodes, kNN edges in (η, φ) momentum space.
      </p>
    </div>
    """)

    gr.Markdown("---")

    with gr.Row():
        with gr.Column(scale=3):
            event_slider = gr.Slider(
                minimum=0, maximum=POOL_SIZE - 1, value=0, step=1,
                label="Event index (scroll through 500 synthetic jets)",
                info="Each position is a different randomly generated collision event"
            )
            classify_btn = gr.Button(
                "Classify this jet  →", variant="primary", size="lg"
            )

        with gr.Column(scale=1):
            gr.Markdown("""
            **Quick guide**

            1. Move the slider to pick an event
            2. Click **Classify this jet**
            3. See the jet graph, GCN prediction, and probability bars
            4. Read the physics explanation below

            Node colour = particle type.
            Node size = transverse momentum (pT).
            """)

    jet_plot   = gr.Plot(label="Jet graph and class probabilities")
    expl_box   = gr.Markdown(label="Physics explanation & model reasoning")
    train_plot = gr.Plot(label="GCN training history")

    classify_btn.click(
        fn=classify_event,
        inputs=[event_slider],
        outputs=[jet_plot, expl_box, train_plot]
    )

    # ── auto-run on slider change for snappy UX ──
    event_slider.release(
        fn=classify_event,
        inputs=[event_slider],
        outputs=[jet_plot, expl_box, train_plot]
    )

    gr.Markdown("---")

    with gr.Accordion("Model architecture & physics background", open=False):
        gr.Markdown("""
        ### Graph Convolutional Network architecture

        | Layer | Type | Input dim | Output dim |
        |-------|------|-----------|------------|
        | 1 | GCNConv + ReLU + Dropout(0.2) | 5 | 64 |
        | 2 | GCNConv + ReLU + Dropout(0.2) | 64 | 128 |
        | 3 | GCNConv + ReLU | 128 | 64 |
        | 4 | GlobalMeanPool | 64 × n_nodes | 64 |
        | 5 | Linear + ReLU | 64 | 32 |
        | 6 | Linear (classifier) | 32 | 5 |

        ### Node features (per particle)
        1. **log(pT)** — transverse momentum on a log scale (GeV)
        2. **η** — pseudorapidity (relates to polar angle)
        3. **φ** — azimuthal angle (radians)
        4. **charge** — electric charge {-1, 0, +1}
        5. **PID** — particle type {photon, neutral hadron, charged hadron, electron, muon}

        ### Why GNNs for particle physics?
        Traditional jet classifiers use image-based CNNs (calorimeter images) or
        dense networks on fixed-length feature vectors. GNNs are more natural because:
        - jets have a **variable number** of particles → no padding needed
        - the **spatial relationship** between particles (proximity in momentum space) matters
        - **permutation invariance** is built in — the order of particles in the list is irrelevant
        - **message passing** lets the model learn multi-particle correlations automatically

        ### Related real-world work
        This demo is inspired by the ParticleNet paper (Qu & Gouskos, 2020) and the
        IAIFI group at MIT (Prof. Jesse Thaler), who apply similar techniques to real
        LHC data from the CMS experiment.

        ### Dataset
        Synthetic data generated with physics-motivated distributions (power-law pT
        spectra, multi-prong angular structure for boosted bosons). For a real project,
        replace the generator with the **JetNet** or **Top Quark Tagging** datasets
        available on Zenodo.
        """)

    gr.HTML("""
    <div style="font-size:12px;color:#888780;padding:1rem 0 0.5rem;border-top:1px solid #E8E6DF;margin-top:1rem">
      ParticleNet · AI + Physics Portfolio Project 7 ·
      Built with PyTorch Geometric, Gradio, and Plotly ·
      Inspired by CERN/LHC jet physics and the NSF IAIFI at MIT
    </div>
    """)

if __name__ == "__main__":
    # classify the first event on startup so the UI is not blank
    demo.launch(show_error=True)