Fix plots and simplify live community search

app.py

@@ -8,6 +8,7 @@ from typing import Dict, List, Tuple
 import gradio as gr
 import numpy as np
 import plotly.express as px
+import plotly.graph_objects as go
 import torch
 from gradio_client import utils as gradio_client_utils
 from transformers import AutoModel, AutoTokenizer
@@ -93,6 +94,14 @@ CSS = """
   color: var(--muted);
   font-size: 0.95rem;
 }
+.search-results {
+  max-height: 320px;
+  overflow-y: auto;
+  border: 1px solid var(--line);
+  border-radius: 16px;
+  background: rgba(255, 255, 255, 0.72);
+  padding: 10px 12px 2px 12px;
+}
 """
 
 
@@ -520,18 +529,35 @@ def _plot_umap(vectors: np.ndarray, labels: List[str], title: str):
     )
     coords = reducer.fit_transform(vectors)
     norms = np.linalg.norm(vectors, axis=1)
-
-    fig = px.scatter(
-        x=coords[:, 0],
-        y=coords[:, 1],
-        color=norms,
-        hover_name=labels,
-        labels={"x": "UMAP 1", "y": "UMAP 2", "color": "vector norm"},
-        title=title,
-        color_continuous_scale="Viridis",
+    x_values = [float(value) for value in coords[:, 0]]
+    y_values = [float(value) for value in coords[:, 1]]
+    color_values = [float(value) for value in norms]
+
+    fig = go.Figure(
+        data=[
+            go.Scatter(
+                x=x_values,
+                y=y_values,
+                mode="markers",
+                text=labels,
+                customdata=np.array(color_values).reshape(-1, 1),
+                hovertemplate="<b>%{text}</b><br>UMAP 1=%{x:.3f}<br>UMAP 2=%{y:.3f}<br>norm=%{customdata[0]:.3f}<extra></extra>",
+                marker={
+                    "size": 10,
+                    "color": color_values,
+                    "colorscale": "Viridis",
+                    "line": {"width": 0.6, "color": "#1d2a1f"},
+                    "opacity": 0.92,
+                    "showscale": True,
+                    "colorbar": {"title": "vector norm"},
+                },
+            )
+        ]
     )
-    fig.update_traces(marker={"size": 10, "line": {"width": 0.6, "color": "#1d2a1f"}, "opacity": 0.9})
     fig.update_layout(
+        title=title,
+        xaxis_title="UMAP 1",
+        yaxis_title="UMAP 2",
         paper_bgcolor="rgba(255,255,255,0)",
         plot_bgcolor="rgba(255,255,255,0.75)",
         margin={"l": 10, "r": 10, "t": 60, "b": 10},
@@ -540,13 +566,18 @@ def _plot_umap(vectors: np.ndarray, labels: List[str], title: str):
 
 
 def _plot_logits(logits: np.ndarray):
-    fig = px.histogram(
-        x=logits,
-        nbins=min(50, max(12, len(logits) // 4)),
-        title="Logit Distribution Over Input DNA Embeddings",
-        color_discrete_sequence=["#d8832f"],
+    hist_values = [float(value) for value in logits]
+    fig = go.Figure(
+        data=[
+            go.Histogram(
+                x=hist_values,
+                nbinsx=min(50, max(12, len(hist_values) // 4)),
+                marker={"color": "#d8832f"},
+            )
+        ]
     )
     fig.update_layout(
+        title="Logit Distribution Over Input DNA Embeddings",
         xaxis_title="Logit",
         yaxis_title="Count",
         paper_bgcolor="rgba(255,255,255,0)",
@@ -759,25 +790,21 @@ with gr.Blocks(title="Microbiome Explorer", css=CSS, theme=gr.themes.Soft()) as
         with gr.Tab("Build A Community"):
             with gr.Column(elem_classes=["soft-card"]):
                 gr.Markdown(
-                    "Search `otus.97.allinfo` by OTU ID, taxon label, or taxonomy string. Add matching OTUs to a custom community, then score the assembled set."
+                    "Search the OTU index by OTU ID, taxon label, or taxonomy string. Matching OTUs appear directly below as you type, so you can add them without opening another widget."
                 )
                 with gr.Row():
                     taxa_query = gr.Textbox(
                         label="Search taxa",
                         placeholder="Try Nitrospira, Lysobacter, Gammaproteobacteria, 97_8697 ...",
-                        scale=5,
+                        scale=6,
                     )
-                    taxa_search_btn = gr.Button("Refresh", variant="secondary", scale=1)
 
                 community_search_status = gr.Markdown(elem_classes=["section-note"])
-                taxa_matches = gr.Dropdown(
+                taxa_matches = gr.CheckboxGroup(
                     label="Matching OTUs",
                     choices=[],
                     value=[],
-                    multiselect=True,
-                    allow_custom_value=False,
-                    interactive=True,
-                    max_choices=MAX_GENES,
+                    elem_classes=["search-results"],
                 )
 
                 with gr.Row():
@@ -822,11 +849,6 @@ with gr.Blocks(title="Microbiome Explorer", css=CSS, theme=gr.themes.Soft()) as
         inputs=[microbeatlas_in],
         outputs=[run_summary, input_umap_plot, final_umap_plot, logits_plot, top_table, member_table],
     )
-    taxa_search_btn.click(
-        fn=search_taxa,
-        inputs=[taxa_query],
-        outputs=[taxa_matches, community_search_status],
-    )
     taxa_query.change(
         fn=search_taxa,
         inputs=[taxa_query],