#!/usr/bin/env python3
"""
Generate a standalone Plotly HTML page with an interactive t-SNE scatter plot.
Click any point to see the 4-channel waveform of that event.
Usage:
python generate_tsne_viewer.py # use saved model_weights.pt
python generate_tsne_viewer.py --train # retrain from scratch
python generate_tsne_viewer.py --output viewer.html
"""
import argparse
import base64
import json
import os
import sys
import warnings
warnings.filterwarnings('ignore')
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, random_split
import torchmetrics
import lightning as L
from sklearn.manifold import TSNE
# ---------------------------------------------------------------------------
# Reproduce the exact architecture from the notebook
# ---------------------------------------------------------------------------
seed = 17052026
L.seed_everything(seed, verbose=False)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
class ScintDataset(Dataset):
def __init__(self, x, y):
self.x = torch.from_numpy(x).to(torch.float)
self.y = torch.from_numpy(y).to(torch.float)
def __len__(self):
return len(self.y)
def __getitem__(self, idx):
return self.x[idx], self.y[idx]
class NNModel(nn.Module):
def __init__(self, input_dim, hidden_dim_1, hidden_dim_2, output_dim):
super().__init__()
self.linear_1 = nn.Linear(input_dim, hidden_dim_1)
self.linear_2 = nn.Linear(hidden_dim_1, hidden_dim_2)
self.linear_3 = nn.Linear(hidden_dim_2, output_dim)
self.act = nn.ReLU()
def forward(self, x):
x = self.act(self.linear_1(x))
x = self.act(self.linear_2(x))
return self.linear_3(x)
def getlatent(self, x):
x = self.act(self.linear_1(x))
return self.act(self.linear_2(x))
class Classifier(L.LightningModule):
def __init__(self, model, model_params, data_config):
super().__init__()
self.model = model(**model_params)
self.data_config = data_config
self.accuracy = torchmetrics.Accuracy(task="binary", num_classes=2)
def forward(self, x):
return self.model(x)
def configure_optimizers(self):
return torch.optim.AdamW(self.parameters(), self.data_config["learning_rate"])
def _calculate_loss(self, batch, mode):
X, y = batch
logits = self(X)
target = y.float()
loss = F.binary_cross_entropy_with_logits(logits, target)
acc = self.accuracy(logits, target)
return loss, acc
def _step(self, batch, mode):
loss, acc = self._calculate_loss(batch, mode)
self.log(f"{mode}_loss", loss, batch_size=self.data_config["batch_size"],
prog_bar=True, on_epoch=True, on_step=(mode == "train"), sync_dist=True)
self.log(f"{mode}_acc", acc, batch_size=self.data_config["batch_size"],
prog_bar=True, on_epoch=True, on_step=(mode == "train"), sync_dist=True)
return loss
def training_step(self, batch):
return self._step(batch, mode="train")
def validation_step(self, batch):
return self._step(batch, mode="val")
def test_step(self, batch):
return self._step(batch, mode="test")
# ---------------------------------------------------------------------------
# Data pipeline
# ---------------------------------------------------------------------------
def load_data(path="data.npz"):
data = np.load(path)
Li6 = data["Li6"]
Po = data["Po"]
Li6_labels = np.zeros((len(Li6), 2))
Li6_labels[:, 0] = 1
Po_labels = np.zeros((len(Po), 2))
Po_labels[:, 1] = 1
y = np.concatenate([Li6_labels, Po_labels])
x = np.concatenate([Li6.sum(axis=1), Po.sum(axis=1)])
x_orig = np.concatenate([Li6, Po])
x = x / x.sum(axis=1, dtype=np.float32)[:, np.newaxis]
return x, x_orig, y
def split_data(x, x_orig, y, train_frac=0.8):
rng = np.random.default_rng(seed)
indices = rng.permutation(len(x))
k = int(len(x) * train_frac)
return {
"xtrain": x[indices[:k]],
"ytrain": y[indices[:k]],
"xtest": x[indices[k:]],
"ytest": y[indices[k:]],
"xtest_orig": x_orig[indices[k:]],
}
# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
parser = argparse.ArgumentParser(description="Generate interactive t-SNE + waveform viewer (Plotly HTML)")
parser.add_argument("--train", action="store_true", help="Retrain model (default: load model_weights.pt)")
parser.add_argument("--output", default="tsne_viewer.html", help="Output HTML file")
parser.add_argument("--data", default="data.npz", help="Path to data.npz")
parser.add_argument("--n-samples", type=int, default=1000, help="Number of t-SNE samples")
parser.add_argument("--downsample", type=int, default=8,
help="Downsample factor for waveforms (4000 / factor = timepoints)")
args = parser.parse_args()
print("[1/6] Loading data …")
x, x_orig, y = load_data(args.data)
split = split_data(x, x_orig, y)
print("[2/6] Setting up model …")
cl = Classifier(
model=NNModel,
model_params={"input_dim": 4000, "hidden_dim_1": 250, "hidden_dim_2": 50, "output_dim": 2},
data_config={"batch_size": 32, "learning_rate": 1e-3},
)
if args.train:
print("[3/6] Training …")
dset = ScintDataset(split["xtrain"], split["ytrain"])
train_size = int(np.around(0.8 * len(dset)))
val_size = len(dset) - train_size
dset_tr, dset_val = torch.utils.data.random_split(dset, [train_size, val_size])
t_dl = DataLoader(dset_tr, batch_size=32, shuffle=True, num_workers=0)
v_dl = DataLoader(dset_val, batch_size=32, shuffle=False, num_workers=0)
trainer = L.Trainer(max_epochs=10, accelerator="auto", devices=1, strategy="auto")
trainer.fit(model=cl, train_dataloaders=t_dl, val_dataloaders=v_dl)
torch.save(cl.model.state_dict(), "model_weights.pt")
print(" Weights saved to model_weights.pt")
else:
if not os.path.exists("model_weights.pt"):
print("ERROR: model_weights.pt not found. Run with --train first.", file=sys.stderr)
sys.exit(1)
cl.model.load_state_dict(torch.load("model_weights.pt", weights_only=True))
print(" Weights loaded from model_weights.pt")
print("[4/6] Extracting latent features and probabilities …")
cl.eval()
rng = np.random.default_rng(seed)
n = args.n_samples
sample_idx = rng.choice(len(split["xtest"]), n, replace=False)
x_sampled = torch.from_numpy(split["xtest"][sample_idx]).to(torch.float)
with torch.no_grad():
latents = cl.model.getlatent(x_sampled).numpy()
logits = cl.model(x_sampled).numpy()
probs = F.softmax(torch.from_numpy(logits), dim=1).numpy()
probs_li6 = probs[:, 0]
labels = (split["ytest"][sample_idx][:, 0] == 0).astype(int) # 1 = Li6, 0 = Po
sample_labels = np.where(labels == 1, "Li6", "Po")
waveforms_orig = split["xtest_orig"][sample_idx]
print(f" Latents: {latents.shape}, Waveforms: {waveforms_orig.shape}")
print("[5/6] Running t-SNE …")
tsne = TSNE(n_components=2, perplexity=30, random_state=seed,
max_iter=1000, learning_rate="auto", init="pca")
tsne_results = tsne.fit_transform(latents)
print("[6/6] Generating HTML …")
html = build_html(
tsne_results=tsne_results,
probs_li6=probs_li6,
labels=labels,
sample_labels=sample_labels,
sample_idx=sample_idx,
waveforms_orig=waveforms_orig,
downsample=args.downsample,
)
with open(args.output, "w") as f:
f.write(html)
size_mb = os.path.getsize(args.output) / 1e6
print(f" Wrote {args.output} ({size_mb:.1f} MB)")
# ---------------------------------------------------------------------------
# HTML generation
# ---------------------------------------------------------------------------
def build_html(tsne_results, probs_li6, labels, sample_labels, sample_idx,
waveforms_orig, downsample):
n = len(tsne_results)
t = waveforms_orig.shape[2] // downsample
wf_down = waveforms_orig[:, :, ::downsample].astype(np.float32)
wf_bytes = wf_down.tobytes()
wf_b64 = base64.b64encode(wf_bytes).decode("ascii")
meta = {
"tsne_x": tsne_results[:, 0].round(4).tolist(),
"tsne_y": tsne_results[:, 1].round(4).tolist(),
"probs": probs_li6.round(4).tolist(),
"labels": labels.tolist(),
"indices": sample_idx.tolist(),
"n": n, "ch": 4, "t": t,
"downsample": downsample,
}
meta_json = json.dumps(meta)
# Use a plain string template (not f-string) to avoid JS/Python brace conflicts
html = """
t-SNE visualisation of NN latent space and waveform signals -
t-SNE of Last Hidden Layer
Colored by P(Li6) — click any point to see its 4-channel waveform
Click a point
Sample index—
True label—
P(Li6)—
Waveform (4 channels)
Click a t-SNE point to view its waveform
PhyNuBe4 notebook — Kaciel Béraud, LPC Caen
interface vibe-coded by A. Vacheret
"""
html = html.replace("'PLACEHOLDER_WF'", "'" + wf_b64 + "'")
html = html.replace("PLACEHOLDER_META", meta_json)
return html
if __name__ == "__main__":
main()