Upload 2 files

app.py

@@ -0,0 +1,694 @@
+"""
+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)

requirements.txt

@@ -0,0 +1,27 @@
+# ParticleNet — GNN Particle Collision Classifier
+# Hugging Face Spaces requirements
+# Python 3.10+
+#
+# Install order matters on HF Spaces:
+# torch must come before torch_geometric and its companions.
+
+# ── core deep learning ──
+torch==2.2.2
+torchvision==0.17.2   # pulled in by torch ecosystem, keeps versions aligned
+
+# ── graph neural networks ──
+# PyTorch Geometric and its sparse/scatter backends
+torch-geometric==2.5.3
+torch-scatter==2.1.2
+torch-sparse==0.6.18
+torch-cluster==1.6.3
+
+# ── scientific computing ──
+numpy==1.26.4
+scikit-learn==1.4.2    # used for kNN graph construction in the jet generator
+
+# ── visualisation ──
+plotly==5.22.0         # interactive jet graph + training curves
+
+# ── Gradio UI ──
+gradio==4.36.1