add more robust error checking

app.py

@@ -16,6 +16,7 @@ import numpy as np
 import scanpy as sc
 import pyarrow as pa
 import pyarrow.parquet as pq
+import matplotlib.pyplot as plt
 from pathlib import Path
 from composer import Trainer, Callback
 from tahoex.model.model import ComposerTX
@@ -45,6 +46,10 @@ with open("./symbol-to-ensembl.json", "r") as f:
 PARQUET_INDEX_COL = "index"
 PARQUET_EMB_COL = "tx1-70m"
 
+# constants for UMAP recoloring
+OBS_NONE_OPTION = "(none)"
+MAX_CATEGORIES = 50
+
 # helper to read AnnData header
 def read_anndata_header(fileobj):
     adata = sc.read_h5ad(fileobj.name, backed="r")
@@ -105,21 +110,130 @@ def ensure_dropdowns(fileobj):
         return (
             gr.Dropdown(choices=["<use .X>"], value="<use .X>"),
             gr.Dropdown(choices=[], value=None),
-            gr.Dropdown(choices=[], value=None),
         )
     try:
-        layers, var_cols, obs_cols = read_anndata_header(fileobj)
+        layers, var_cols, _ = read_anndata_header(fileobj)
         return (
             gr.Dropdown(choices=["<use .X>"] + layers, value="<use .X>"),
             gr.Dropdown(choices=var_cols, value=(var_cols[0] if var_cols else None)),
-            gr.Dropdown(choices=obs_cols, value=None),
         )
     except Exception:
         return (
             gr.Dropdown(choices=["<use .X>"], value="<use .X>"),
             gr.Dropdown(choices=[], value=None),
-            gr.Dropdown(choices=[], value=None),
         )
+    
+# draw an uncolored UMAP
+def draw_uncolored(coords, title_suffix=None):
+    fig = plt.figure(figsize=(5.5, 5.0))
+    ax = fig.add_subplot(111)
+    ax.scatter(coords[:, 0], coords[:, 1], s=3, alpha=0.75)
+    ttl = "UMAP of Tx1-70M embeddings"
+    if title_suffix:
+        ttl += f" ({title_suffix})"
+    ax.set_title(ttl)
+    ax.set_xlabel("UMAP1")
+    ax.set_ylabel("UMAP2")
+    fig.tight_layout()
+    out_png = _unique_output("umap.png")
+    fig.savefig(out_png, dpi=160)
+    plt.close(fig)
+    return out_png
+
+# recolor UMAP given obs column
+def recolor_umap(obs_col, coords, h5ad_path):
+
+    # make sure inputs are valid
+    if coords is None or h5ad_path is None:
+        raise gr.Error("Run embeddings first to compute UMAP.")
+    coords = np.asarray(coords)
+    if coords.ndim != 2 or coords.shape[1] != 2:
+        raise gr.Error(f"UMAP coordinates look wrong, shape = {coords.shape}. Please recompute.")
+
+    # handle no-coloring option
+    if obs_col == OBS_NONE_OPTION:
+        out_png = draw_uncolored(coords)
+        return str(out_png.resolve())
+
+    # read obs column
+    try:
+        adata = sc.read_h5ad(h5ad_path, backed="r")
+        series = adata.obs[obs_col]
+        n = series.shape[0]
+        if n != coords.shape[0]:
+            gr.Warning(f"Length mismatch: obs has {n} rows, UMAP has {coords.shape[0]}. Using minimum length.")
+            m = min(n, coords.shape[0])
+            series = series.iloc[:m]
+            coords = coords[:m]
+    except Exception as e:
+        raise gr.Error(f"Failed to read .obs column '{obs_col}': {e}")
+    
+    # sanitize values
+    s = series.copy()
+    numeric_candidate = pd.to_numeric(s, errors="coerce")
+    n_numeric_valid = int(np.isfinite(numeric_candidate.astype(float)).sum())
+    n_total = int(len(s))
+
+    # plot numerical labels
+    if n_numeric_valid >= max(5, 0.5 * n_total):
+
+        # check for sufficient numeric values or constant values
+        vals = pd.to_numeric(s, errors="coerce").astype(float).values
+        mask = np.isfinite(vals)
+        if mask.sum() < max(10, 0.1 * len(vals)):
+            gr.Warning(f"Too few finite numeric values in '{obs_col}'. Showing uncolored UMAP.")
+            return draw_uncolored(f"{obs_col}: insufficient numeric values")
+        if np.nanmax(vals[mask]) == np.nanmin(vals[mask]):
+            gr.Info(f"'{obs_col}' is constant. Showing uncolored UMAP.")
+            return draw_uncolored(f"{obs_col}: constant")
+
+        # draw with colorbar
+        fig = plt.figure(figsize=(5.5, 5.0))
+        ax = fig.add_subplot(111)
+        scatt = ax.scatter(coords[mask, 0], coords[mask, 1], s=3, alpha=0.85, c=vals[mask])
+        fig.colorbar(scatt, ax=ax, shrink=0.7, label=obs_col)
+        if (~mask).any():
+            ax.scatter(coords[~mask, 0], coords[~mask, 1], s=3, alpha=0.25)
+        ax.set_title(f"UMAP colored by {obs_col}")
+        ax.set_xlabel("UMAP1")
+        ax.set_ylabel("UMAP2")
+        fig.tight_layout()
+        out_png = _unique_output("umap.png")
+        fig.savefig(out_png, dpi=160)
+        plt.close(fig)
+        return str(out_png.resolve())
+    
+    # categorical coloring
+    else:
+
+        # check for too many or too few categories
+        cats = s.astype(str).fillna("NA").values
+        uniq = pd.unique(cats)
+        n_cat = len(uniq)
+        if n_cat > MAX_CATEGORIES:
+            gr.Warning(f"'{obs_col}' has too many categories. Showing uncolored UMAP.")
+            out_png = draw_uncolored(coords, f"{obs_col}: {n_cat} categories")
+            return str(out_png.resolve())
+        if n_cat <= 1:
+            gr.Info(f"'{obs_col}' has a single category. Showing uncolored UMAP.")
+            out_png = draw_uncolored(f"{obs_col}: 1 category")
+            return str(out_png.resolve())
+
+        # draw with legend
+        fig = plt.figure(figsize=(5.5, 5.0))
+        ax = fig.add_subplot(111)
+        for cat in sorted(map(str, uniq)):
+            mask = (cats == cat)
+            ax.scatter(coords[mask, 0], coords[mask, 1], s=3, alpha=0.85, label=cat)
+        ax.legend(markerscale=3, fontsize=8, loc="best", frameon=True, ncol=1)
+        ax.set_title(f"UMAP colored by {obs_col}")
+        ax.set_xlabel("UMAP1")
+        ax.set_ylabel("UMAP2")
+        fig.tight_layout()
+        out_png = _unique_output("umap.png")
+        fig.savefig(out_png, dpi=160)
+        plt.close(fig)
+        return str(out_png.resolve())
 
 # custom callback to report progress to Gradio
 class GradioProgressCallback(Callback):
@@ -337,47 +451,6 @@ def _embed(adata_bytes, layer_name, feature_col, use_symbols, progress):
     # return embeddings and metadata
     return cell_array, layers, var_cols, obs_cols, adata_persisted
 
-# recolor UMAP given obs column
-def recolor_umap(obs_col, coords, h5ad_path):
-
-    # make sure inputs are valid
-    if not obs_col:
-        raise gr.Error("Pick a .obs column to color by.")
-    if coords is None or h5ad_path is None:
-        raise gr.Error("Compute embeddings before selecting a .obs column.")
-
-    # read obs column
-    adata = sc.read_h5ad(h5ad_path, backed="r")
-    if obs_col not in adata.obs.columns:
-        raise gr.Error(f"Column '{obs_col}' not found in .obs.")
-
-    # construct plot
-    color_series = adata.obs[obs_col]
-    import matplotlib.pyplot as plt
-    fig = plt.figure(figsize=(5.5, 5.0))
-    ax = fig.add_subplot(111)
-    if pd.api.types.is_numeric_dtype(color_series):
-        scatt = ax.scatter(coords[:, 0], coords[:, 1], s=3, alpha=0.85, c=color_series.values)
-        fig.colorbar(scatt, ax=ax, shrink=0.7, label=obs_col)
-        ax.set_title(f"UMAP colored by {obs_col}")
-    else:
-        cs = color_series.astype(str).values
-        for cat in sorted(pd.unique(cs)):
-            mask = (cs == cat)
-            ax.scatter(coords[mask, 0], coords[mask, 1], s=3, alpha=0.85, label=str(cat))
-        ax.legend(markerscale=3, fontsize=8, loc="best", frameon=True)
-        ax.set_title(f"UMAP colored by {obs_col}")
-
-    # finalize and save plot
-    ax.set_xlabel("UMAP1")
-    ax.set_ylabel("UMAP2")
-    fig.tight_layout()
-    out_png = _unique_output("umap.png")
-    fig.savefig(out_png, dpi=160)
-    plt.close(fig)
-    return str(out_png.resolve())
-
-
 # processing pipeline given user inputs
 def run_pipeline(fileobj, layer_choice, var_choice, use_symbols, progress=gr.Progress(track_tqdm=False)):
 
@@ -415,7 +488,6 @@ def run_pipeline(fileobj, layer_choice, var_choice, use_symbols, progress=gr.Pro
 
     # plot UMAP (no coloring by default)
     progress(0.90, desc="plotting UMAP")
-    import matplotlib.pyplot as plt
     fig = plt.figure(figsize=(5.5, 5.0))
     ax = fig.add_subplot(111)
     ax.scatter(coords[:, 0], coords[:, 1], s=3, alpha=0.75)
@@ -428,7 +500,7 @@ def run_pipeline(fileobj, layer_choice, var_choice, use_symbols, progress=gr.Pro
     plt.close(fig)
 
     # enable coloring dropdown
-    update_obs_dd = gr.Dropdown(choices=obs_cols, value=None, interactive=True)
+    update_obs_dd = gr.Dropdown(choices=[OBS_NONE_OPTION] + obs_cols, value=OBS_NONE_OPTION, interactive=True)
 
     # save other outputs and return paths
     progress(0.95, desc="saving outputs")