Order SNR legends and dropdowns

app.py

@@ -321,6 +321,10 @@ def plot_tsne(
     sampled_df["x"] = projections[:, 0]
     sampled_df["y"] = projections[:, 1]
 
+    category_orders = {}
+    if color_column == "snr":
+        category_orders["snr"] = [snr for snr in SNR_ORDER if snr in sampled_df["snr"].unique()]
+
     fig = px.scatter(
         sampled_df,
         x="x",
@@ -329,6 +333,7 @@ def plot_tsne(
         hover_data=["tech", "snr", "mod", "mob"],
         title=f"t-SNE of {representation} ({len(sampled_df)} samples)",
         template="plotly_white",
+        category_orders=category_orders,
     )
     height = 680 if color_label == "SNR" else 640
     fig.update_layout(legend_title_text=color_label, width=640, height=height)
@@ -700,7 +705,7 @@ has_moe_column = df["moe_embedding"].apply(lambda x: x is not None)
 joint_eval_df = df[has_moe_column & df["joint_label_id"].notna()]
 
 tech_choices = sorted(df["tech"].unique())
-snr_choices = sorted(df["snr"].unique())
+snr_choices = _sort_snrs(df["snr"].unique())
 mod_choices = sorted(df["mod"].unique())
 mob_choices = sorted(df["mob"].unique())
 
@@ -713,9 +718,10 @@ COLOR_OPTIONS: Dict[str, str] = {
     "Modulation": "mod",
     "Mobility": "mob",
 }
+SNR_ORDER = ["SNR-5dB", "SNR0dB", "SNR5dB", "SNR10dB", "SNR15dB", "SNR20dB", "SNR25dB"]
 TECH_EXPERT_ORDER = ["LTE", "WiFi", "5G"]
 TECH_TO_EXPERT_IDX = {name: idx for idx, name in enumerate(TECH_EXPERT_ORDER)}
-DEFAULT_TSNE_SAMPLES_PER_SNR = 200
+DEFAULT_TSNE_SAMPLES_PER_SNR = 500
 
 default_tech = tech_choices[:1] if tech_choices else []
 initial_spec_mod_choices = TECH_TO_MODS.get(default_tech[0], mod_choices) if default_tech else mod_choices
@@ -727,6 +733,11 @@ def update_modulation_choices(selected_tech: Optional[str]):
         choices = TECH_TO_MODS.get(selected_tech, mod_choices)
     return gr.Dropdown.update(choices=choices, value=None)
 
+
+def _sort_snrs(labels: List[str] | np.ndarray) -> List[str]:
+    ordering = {snr: idx for idx, snr in enumerate(SNR_ORDER)}
+    return sorted(labels, key=lambda x: ordering.get(x, len(ordering)))
+
 with gr.Blocks(title="LWM-Spectro Lab") as demo:
     gr.Markdown("# 🔬 LWM-Spectro Interactive Demo")
     gr.Markdown(