feat: add features to handle better different species

app.py

@@ -12,7 +12,12 @@ matplotlib.use('Agg')
 
 import matplotlib.pyplot as plt
 
-from ntv3_tracks_pipeline import load_ntv3_tracks_pipeline, BED_ELEMENT_COLORS
+from ntv3_tracks_pipeline import (
+    load_ntv3_tracks_pipeline,
+    BED_ELEMENT_COLORS,
+    ASSEMBLY_TO_SPECIES,
+    SPECIES_WITH_COORDINATE_SUPPORT,
+)
 from bigwig_export import create_bigwig_zip, _softmax_last
 
 
@@ -116,6 +121,28 @@ def _get_bigwig_names(species: str) -> list[str]:
     return _BIGWIG_CACHE[species]
 
 
+def _has_bigwigs(species: str) -> bool:
+    """Check if a species has BigWig tracks available in the current model."""
+    try:
+        tracks = _get_bigwig_names(species)
+        return len(tracks) > 0
+    except (ValueError, AttributeError):
+        # Species not in config or pipeline not loaded
+        return False
+
+
+def _get_species_with_bigwigs() -> set[str]:
+    """Get set of species that have BigWig tracks available in the current model."""
+    if pipe is None:
+        return set()
+    
+    species_with_bigwigs = set()
+    for species in ASSEMBLY_TO_SPECIES.values():
+        if _has_bigwigs(species):
+            species_with_bigwigs.add(species)
+    return species_with_bigwigs
+
+
 def _rank_search(query: str, names: list[str], limit: int) -> list[str]:
     """
     Return up to `limit` candidate track IDs matching `query` using a fast,
@@ -180,7 +207,11 @@ def update_coords_on_species_change(species: str):
 
 def reset_on_species_change(species: str):
     # Clear results + selected when species changes (avoids mismatched IDs)
-    _get_bigwig_names(species)  # warms cache
+    try:
+        _get_bigwig_names(species)  # warms cache if available
+    except (ValueError, AttributeError):
+        # Species doesn't have bigwigs, that's okay
+        pass
     return (
         gr.update(value=""),          # query textbox
         gr.update(choices=[], value=[]),  # results list
@@ -201,7 +232,18 @@ def predict(
     bigwig_selected: list[str],
     bed_elements: list[str],
 ):
+    # Debug: verify species is being passed
+    if not species:
+        raise gr.Error("Species parameter is missing. Please select a species.")
+    
     if use_coords:
+        # Check if this species supports coordinate-based fetching
+        if species not in SPECIES_WITH_COORDINATE_SUPPORT:
+            raise gr.Error(
+                f"Species '{species}' does not support coordinate-based sequence fetching. "
+                f"Please provide a DNA sequence directly or use one of the supported species: "
+                f"{', '.join(sorted(SPECIES_WITH_COORDINATE_SUPPORT))}"
+            )
         if not chrom:
             raise gr.Error("chrom is required when use_coords=True")
         if start is None or end is None or int(end) <= int(start):
@@ -211,7 +253,11 @@ def predict(
         if not seq or not seq.strip():
             raise gr.Error("seq is required when use_coords=False")
         inputs = {"seq": seq.strip(), "species": species}
-
+    
+    # Verify species is in inputs before calling pipeline
+    if "species" not in inputs:
+        raise gr.Error(f"Internal error: species not found in inputs dict. Inputs: {list(inputs.keys())}")
+    
     out = pipe(inputs)
 
     bw_names = out.bigwig_track_names or []
@@ -219,11 +265,18 @@ def predict(
     bed_names = out.bed_element_names or []
     bed_logits = out.bed_tracks_logits
 
-    if bw is None or not bw_names:
-        raise gr.Error("No BigWig tracks available in model output.")
+    # Check if we have any tracks/elements to plot
+    has_bigwigs = bw is not None and len(bw_names) > 0
+    has_bed = bed_logits is not None and len(bed_names) > 0
+    
+    if not has_bigwigs and not has_bed:
+        raise gr.Error("No BigWig tracks or BED elements available for this species in the current model.")
+    
+    if not has_bigwigs and bigwig_selected:
+        raise gr.Error("No BigWig tracks available for this species, but BigWig tracks were selected. Please deselect BigWig tracks or choose a different species.")
 
     # Defaults if user picked none
-    if not bigwig_selected:
+    if has_bigwigs and not bigwig_selected:
         default_bigwig_tracks = [
             "ENCSR056HPM",  # K562 RNA-seq
             "ENCSR921NMD",  # K562 DNAse
@@ -236,21 +289,31 @@ def predict(
         ]
         # Filter to only include tracks that are available for this species/assembly
         bigwig_selected = [tid for tid in default_bigwig_tracks if tid in bw_names]
+    
     if (not bed_elements) and bed_names:
         default_bed_elements = ["protein_coding_gene", "exon", "intron"]
         # Filter to only include elements that are available
         bed_elements = [elem for elem in default_bed_elements if elem in bed_names]
 
     # Validate (important for API usage)
-    missing_tracks = [t for t in bigwig_selected if t not in bw_names]
-    if missing_tracks:
-        raise gr.Error(f"Unknown BigWig track id(s): {missing_tracks}")
-
-    missing_elems = [e for e in bed_elements if e not in bed_names]
-    if missing_elems:
-        raise gr.Error(f"Unknown BED element(s): {missing_elems}")
-
-    L = bw.shape[0]
+    if has_bigwigs and bigwig_selected:
+        missing_tracks = [t for t in bigwig_selected if t not in bw_names]
+        if missing_tracks:
+            raise gr.Error(f"Unknown BigWig track id(s): {missing_tracks}")
+
+    if bed_elements:
+        missing_elems = [e for e in bed_elements if e not in bed_names]
+        if missing_elems:
+            raise gr.Error(f"Unknown BED element(s): {missing_elems}")
+
+    # Determine sequence length from available data
+    if has_bigwigs:
+        L = bw.shape[0]
+    elif has_bed:
+        L = bed_logits.shape[0]
+    else:
+        raise gr.Error("No data available for plotting.")
+    
     stride = _global_stride(L, PLOT_TARGET_POINTS)
 
     x0 = int(out.pred_start or 0)
@@ -258,10 +321,14 @@ def predict(
     x = np.linspace(x0, x1, num=L, endpoint=False)[::stride]
 
     series: list[tuple[str, np.ndarray]] = []
-    for tid in bigwig_selected:
-        idx = bw_names.index(tid)
-        series.append((tid, bw[:, idx][::stride].astype(float)))
+    
+    # Add BigWig tracks if available and selected
+    if has_bigwigs and bigwig_selected:
+        for tid in bigwig_selected:
+            idx = bw_names.index(tid)
+            series.append((tid, bw[:, idx][::stride].astype(float)))
 
+    # Add BED elements if available and selected
     if bed_logits is not None and bed_elements:
         probs = _softmax_last(bed_logits)
         for ename in bed_elements:
@@ -532,9 +599,24 @@ DEFAULT_COORDS = {
 # Get default coordinates for default species
 _default_coords = DEFAULT_COORDS.get(DEFAULT_SPECIES, DEFAULT_COORDS["human"])
 
+# Format species names for display (replace underscores with spaces, capitalize)
+def _format_species_name(species: str) -> str:
+    """Format species name for display."""
+    return species.replace("_", " ").title()
+
+# Get all available species and format them
+_all_species = sorted(ASSEMBLY_TO_SPECIES.values())
+_all_species_formatted = [_format_species_name(s) for s in _all_species]
+_all_species_list = ", ".join(_all_species_formatted)
+
+# Get species with BigWig tracks
+_species_with_bigwigs = _get_species_with_bigwigs()
+_bigwig_species_formatted = sorted([_format_species_name(s) for s in _species_with_bigwigs])
+_bigwig_species_list = ", ".join(_bigwig_species_formatted) if _bigwig_species_formatted else "None (BED elements only)"
+
 with gr.Blocks(title="NTv3 Tracks Demo") as demo:
     gr.Markdown(
-    """
+    f"""
 <div class="intro-hero">
 
 <div class="intro-title">
@@ -543,9 +625,6 @@ with gr.Blocks(title="NTv3 Tracks Demo") as demo:
     Predict and visualize functional genomics signals directly from DNA using
     <strong>Nucleotide Transformer v3</strong>.
   </p>
-  <p style="margin-top: 8px; font-size: 0.95rem; opacity: 0.85;">
-    <strong>Currently available species:</strong> Human, Mouse, Drosophila melanogaster, Arabidopsis thaliana, Gorilla
-  </p>
 </div>
 
 <div class="intro-grid">
@@ -581,6 +660,12 @@ with gr.Blocks(title="NTv3 Tracks Demo") as demo:
   <span><strong>Tip:</strong> The demo includes default settings that you can use to get started, taking ~ 1 minute to run.</span>
 </div>
 
+<div style="margin-top: 16px; padding: 12px; background: rgba(0,0,0,0.03); border-radius: 12px; font-size: 0.95rem;">
+  <strong>Available species:</strong> {_all_species_list}<br>
+  <br>
+  <strong>Species with functional tracks:</strong> {_bigwig_species_list}
+</div>
+
 </div>
 """,
     elem_id="intro_markdown",
@@ -609,17 +694,26 @@ with gr.Blocks(title="NTv3 Tracks Demo") as demo:
     
     model_status = gr.Markdown("", visible=False)
     
-    gr.Markdown("## Input sequence (Genomic coordinate or DNA sequence)")
+    gr.Markdown("## Input sequence (Genomic coordinate or DNA sequence)\n"
+    "Supported species for coordinate-based sequence fetching: " + ", ".join(sorted(SPECIES_WITH_COORDINATE_SUPPORT)) + "\n"
+    )
+    
+    # Get all available species from the pipeline
+    all_species = sorted(ASSEMBLY_TO_SPECIES.values())
     
     with gr.Row():
         species = gr.Dropdown(
-            ["human", "mouse", "drosophila_melanogaster", "arabidopsis_thaliana", "gorilla_gorilla"],
+            choices=all_species,
             value=DEFAULT_SPECIES,
             label="Species",
         )
-        use_coords = gr.Checkbox(True, label="Use genome coordinates")
+        use_coords = gr.Checkbox(
+            True, 
+            label="Use genome coordinate (only for supported species)",
+            visible=DEFAULT_SPECIES in SPECIES_WITH_COORDINATE_SUPPORT
+        )
 
-    with gr.Row():
+    with gr.Row(visible=True) as coords_row:
         chrom = gr.Textbox(label="Chromosome", value=_default_coords["chrom"])
         start = gr.Number(label="Start", value=_default_coords["start"], precision=0)
         end = gr.Number(label="End", value=_default_coords["end"], precision=0)
@@ -650,6 +744,11 @@ with gr.Blocks(title="NTv3 Tracks Demo") as demo:
     )
 
     gr.Markdown("## Select functional tracks")
+    
+    bigwig_no_tracks_msg = gr.Markdown(
+        "⚠️ No functional genomic tracks available for this species in the current model.",
+        visible=False,
+    )
 
     bigwig_selected = gr.CheckboxGroup(
         choices=_init_bigwig_selected,
@@ -667,7 +766,7 @@ with gr.Blocks(title="NTv3 Tracks Demo") as demo:
         label="Results (click to add to Selected)",
     )
 
-    with gr.Row():
+    with gr.Row(visible=True) as bigwig_buttons_row:
         bigwig_clear_btn = gr.Button("Clear results")
         bigwig_remove_btn = gr.Button("Remove checked from Selected")
 
@@ -743,11 +842,52 @@ with gr.Blocks(title="NTv3 Tracks Demo") as demo:
         outputs=[bigwig_query, bigwig_results, bigwig_selected],
     )
     
-    # Update coordinates when species changes
+    # Update coordinates visibility and values when species changes
+    def update_on_species_change(species: str, use_coords_val: bool):
+        """Update coordinates visibility and values when species changes."""
+        is_supported = species in SPECIES_WITH_COORDINATE_SUPPORT
+        has_bigwigs = _has_bigwigs(species)
+        coords = DEFAULT_COORDS.get(species, DEFAULT_COORDS["human"])
+        # Show coordinates only if species is supported AND use_coords is True
+        show_coords = is_supported and use_coords_val
+        return (
+            gr.update(visible=show_coords, value=coords["chrom"]),
+            gr.update(visible=show_coords, value=coords["start"]),
+            gr.update(visible=show_coords, value=coords["end"]),
+            gr.update(value=is_supported, visible=is_supported),  # Show/hide and enable use_coords only if supported
+            gr.update(visible=show_coords),  # Show/hide the row
+            gr.update(visible=not has_bigwigs),  # Show "no tracks" message if no bigwigs
+            gr.update(visible=has_bigwigs),  # Show bigwig selection if available
+            gr.update(visible=has_bigwigs),  # Show bigwig query if available
+            gr.update(visible=has_bigwigs),  # Show bigwig results if available
+            gr.update(visible=has_bigwigs),  # Show bigwig buttons if available
+        )
+    
+    # Update coordinates visibility when checkbox changes
+    def update_coords_visibility(use_coords_val: bool, species: str):
+        """Update coordinates visibility when checkbox changes."""
+        is_supported = species in SPECIES_WITH_COORDINATE_SUPPORT
+        show_coords = is_supported and use_coords_val
+        return (
+            gr.update(visible=show_coords),
+            gr.update(visible=show_coords),
+            gr.update(visible=show_coords),
+            gr.update(visible=show_coords),  # Show/hide the row
+        )
+    
     species.change(
-        fn=update_coords_on_species_change,
-        inputs=[species],
-        outputs=[chrom, start, end],
+        fn=update_on_species_change,
+        inputs=[species, use_coords],
+        outputs=[
+            chrom, start, end, use_coords, coords_row,
+            bigwig_no_tracks_msg, bigwig_selected, bigwig_query, bigwig_results, bigwig_buttons_row
+        ],
+    )
+    
+    use_coords.change(
+        fn=update_coords_visibility,
+        inputs=[use_coords, species],
+        outputs=[chrom, start, end, coords_row],
     )
 
     btn.click(

ntv3_tracks_pipeline.py

@@ -38,8 +38,6 @@ ASSEMBLY_TO_SPECIES = {
     "Glycine_max_v2.1": "glycine_max",
     "IWGSC": "triticum_aestivum",
     "Gossypium_hirsutum_v2.1": "gossypium_hirsutum",
-    "ASM228892v3": "delphinapterus_leucas",
-    "ASM334442v1": "ursus_americanus",
     "AmpOce1": "amphiprion_ocellaris",
     "Bison_UMD1": "bison_bison_bison",
     "ChiLan1": "chinchilla_lanigera",
@@ -47,7 +45,6 @@ ASSEMBLY_TO_SPECIES = {
     "GRCz11": "danio_rerio",
     "KH": "ciona_intestinalis",
     "Mnem_1": "macaca_nemestrina",
-    "R64": "saccharomyces_cerevisiae",
     "ROS_Cfam_1": "canis_lupus_familiaris",
     "SCA1": "serinus_canaria",
     "TETRAODON8": "tetraodon_nigroviridis",
@@ -56,11 +53,23 @@ ASSEMBLY_TO_SPECIES = {
     "fSalTru1": "salmo_trutta",
     "gorGor4": "gorilla_gorilla",
     "mRatBN7": "rattus_norvegicus",
-    "SL3": "solanum_lycopersicum",
-    "ARS-UCD2.0": "bos_taurus",
 }
 SPECIES_TO_ASSEMBLY = {v: k for k, v in ASSEMBLY_TO_SPECIES.items()}
 
+# ---------------------------------------------------------------------
+# Species that support coordinate-based sequence fetching
+# ---------------------------------------------------------------------
+# List of species that can fetch DNA sequences from genomic coordinates via API.
+# Species not in this list can still be used but require direct DNA sequence input.
+SPECIES_WITH_COORDINATE_SUPPORT = {
+    "human",  # hg38 - UCSC API
+    "mouse",  # mm10 - UCSC API
+    "drosophila_melanogaster",  # dm6 - UCSC API
+    "arabidopsis_thaliana",  # TAIR10 - UCSC hub API
+    "gorilla_gorilla",  # gorGor4 - UCSC API
+    # Add more species as API URLs are configured
+}
+
 # ---------------------------------------------------------------------
 # Assembly -> API URL template mapping
 # ---------------------------------------------------------------------