fix: restrict bed elements to trained ones

app.py

@@ -308,6 +308,21 @@ def _get_bigwig_names(species: str) -> list[str]:
     return _BIGWIG_CACHE[species]
 
 
+def _get_bed_element_names(species: str) -> list[str]:
+    """Get BED element names available for a given species (filtered by training data)."""
+    if pipe is None:
+        return []
+    try:
+        return pipe.available_bed_element_names(species)
+    except (ValueError, AttributeError):
+        return []
+
+
+def _format_bed_element_for_display(element_name: str) -> str:
+    """Format BED element name for display: replace underscores with spaces and capitalize."""
+    return element_name.replace("_", " ").title()
+
+
 def _has_bigwigs(species: str) -> bool:
     """Check if a species has BigWig tracks available in the current model."""
     try:
@@ -652,7 +667,7 @@ def predict(
     if bed_logits is not None and bed_elements:
         probs = _softmax_last(bed_logits)
         for ename in bed_elements:
-            display_name = ename.replace("_", " ")
+            display_name = ename.replace("_", " ").lower()
             eidx = bed_names.index(ename)
             series.append((display_name, probs[:, eidx, 1][::stride].astype(float)))
 
@@ -751,7 +766,8 @@ observer.observe(document.body, {
 """
 
 # BED list is small enough to keep as dropdown
-_init_bed = pipe.available_bed_element_names()
+# Filter by default species to show only elements available for training
+_init_bed = pipe.available_bed_element_names(DEFAULT_SPECIES)
 
 # Default BigWig tracks
 DEFAULT_BIGWIG_TRACKS = [
@@ -781,6 +797,12 @@ _init_bigwig_selected = [
 # Filter default BED elements to only those available
 _init_bed_selected = [elem for elem in DEFAULT_BED_ELEMENTS if elem in _init_bed]
 
+# Format BED elements for display: use tuples (display_name, value) for dropdown
+_init_bed_choices = [
+    (_format_bed_element_for_display(elem), elem) for elem in _init_bed
+]
+_init_bed_selected_values = _init_bed_selected  # Keep original values for selection
+
 # Default coordinates per species
 DEFAULT_COORDS = {
     "human": {"chrom": "chr19", "start": 6_700_000, "end": 6_831_072},
@@ -1063,8 +1085,8 @@ with gr.Blocks(title="NTv3 Tracks Demo") as demo:
     gr.Markdown("# Select genome annotation elements")
 
     bed_elements = gr.Dropdown(
-        choices=_init_bed,
-        value=_init_bed_selected if _init_bed_selected else [],
+        choices=_init_bed_choices,
+        value=_init_bed_selected_values if _init_bed_selected_values else [],
         multiselect=True,
         label="Genome annotation elements (search + select)",
         elem_id="bed_elements_dropdown",
@@ -1268,6 +1290,17 @@ with gr.Blocks(title="NTv3 Tracks Demo") as demo:
             except Exception:
                 pass
 
+        # Get BED elements available for this species
+        bed_element_names = _get_bed_element_names(species)
+        # Filter default BED elements to only those available for this species
+        default_bed_selected = [
+            elem for elem in DEFAULT_BED_ELEMENTS if elem in bed_element_names
+        ]
+        # Format BED elements for display: use tuples (display_name, value)
+        bed_element_choices = [
+            (_format_bed_element_for_display(elem), elem) for elem in bed_element_names
+        ]
+
         return (
             gr.update(visible=show_coords, value=coords["chrom"]),
             gr.update(visible=show_coords, value=coords["start"]),
@@ -1291,6 +1324,10 @@ with gr.Blocks(title="NTv3 Tracks Demo") as demo:
             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
+            gr.update(
+                choices=bed_element_choices,
+                value=default_bed_selected,
+            ),  # Update BED elements dropdown with species-specific elements
         )
 
     # Update input type radio visibility and value when species changes
@@ -1339,6 +1376,7 @@ with gr.Blocks(title="NTv3 Tracks Demo") as demo:
             bigwig_query,
             bigwig_results,
             bigwig_buttons_row,
+            bed_elements,
         ],
     )
 

ntv3_tracks_pipeline.py

@@ -110,6 +110,79 @@ BED_ELEMENT_COLORS = {
 }
 
 
+def _filter_bed_elements_by_species(
+    bed_element_names: list[str], species: str
+) -> list[str]:
+    """
+    Filter BED element names based on species-specific training data availability.
+    
+    Rules:
+    - Human: all tracks
+    - Mouse: only polyA_signal
+    - Other species: everything except promoter, enhancer, ctcf, lncrna
+    
+    Parameters
+    ----------
+    bed_element_names : list[str]
+        Full list of BED element names from the model config
+    species : str
+        Species name (e.g., "human", "mouse", "drosophila_melanogaster")
+    
+    Returns
+    -------
+    list[str]
+        Filtered list of BED element names available for this species
+    """
+    if not bed_element_names:
+        return []
+    
+    # Elements to exclude for "other species" (everything except human and mouse)
+    excluded_for_other_species = {
+        "promoter Tissue specific",
+        "promoter Tissue invariant",
+        "enhancer Tissue specific",
+        "enhancer Tissue invariant",
+        "CTCF-bound",
+        "lncRNA",
+    }
+    
+    # Normalize element names (handle both with/without underscores/spaces)
+    normalized_excluded = set()
+    for elem in excluded_for_other_species:
+        normalized_excluded.add(elem)
+        normalized_excluded.add(elem.replace(" ", "_"))
+    
+    if species == "human":
+        # Human: all tracks
+        return list(bed_element_names)
+    else:
+        # Other species: everything except promoter, enhancer, ctcf, lncrna
+        # Normalize element names for comparison (handle spaces, underscores, case)
+        normalized_bed_names = {
+            elem.lower().replace("_", " "): elem
+            for elem in bed_element_names
+        }
+        normalized_excluded_lower = {
+            elem.lower().replace("_", " ")
+            for elem in excluded_for_other_species
+        }
+        
+        # Also check for keywords in element names
+        excluded_keywords = ["promoter", "enhancer", "ctcf", "lnc"]
+        
+        filtered_normalized = [
+            norm_name
+            for norm_name, orig_elem in normalized_bed_names.items()
+            if norm_name not in normalized_excluded_lower
+            and not any(keyword in norm_name for keyword in excluded_keywords)
+        ]
+        
+        # Return original element names (preserving original format)
+        return [
+            normalized_bed_names[norm_name] for norm_name in filtered_normalized
+        ]
+
+
 def _sanitize_dna(seq: str) -> str:
     seq = seq.upper()
     return "".join(ch if ch in ("A", "C", "G", "T", "N") else "N" for ch in seq)
@@ -373,11 +446,26 @@ class NTv3TracksPipeline(Pipeline):
 
         return list(self.config.bigwigs_per_species[species])
 
-    def available_bed_element_names(self) -> list[str]:
+    def available_bed_element_names(self, species: str | None = None) -> list[str]:
         """
-        Return BED element names available in this checkpoint (no forward pass).
+        Return BED element names available in this checkpoint for the given species.
+        Filters elements based on species-specific training data availability.
+        
+        Parameters
+        ----------
+        species : str | None
+            Species name (e.g., "human", "mouse"). If None, returns all elements
+            without filtering (for backward compatibility).
+        
+        Returns
+        -------
+        list[str]
+            Filtered list of BED element names available for this species
         """
-        return list(self.bed_element_names or [])
+        all_elements = list(self.bed_element_names or [])
+        if species is None:
+            return all_elements
+        return _filter_bed_elements_by_species(all_elements, species)
 
     def _sanitize_parameters(self, **kwargs):
         return {}, {}, {}
@@ -523,6 +611,33 @@ class NTv3TracksPipeline(Pipeline):
         if mlm_np.ndim == 3:
             mlm_np = mlm_np[0]
 
+        # Filter BED elements based on species
+        species = meta.get("species")
+        all_bed_element_names = self.bed_element_names or []
+        if species and all_bed_element_names:
+            filtered_bed_element_names = _filter_bed_elements_by_species(
+                all_bed_element_names, species
+            )
+            # Filter bed_tracks_logits to only include elements available for this species
+            if filtered_bed_element_names != all_bed_element_names:
+                # Create mapping from filtered element names to original indices
+                element_indices = [
+                    all_bed_element_names.index(elem)
+                    for elem in filtered_bed_element_names
+                    if elem in all_bed_element_names
+                ]
+                if element_indices:
+                    # bed_np shape is (L, E, C) where E is number of elements
+                    bed_np = bed_np[:, element_indices, :]
+                    # Update filtered list to only include elements that were found
+                    filtered_bed_element_names = [
+                        elem
+                        for elem in filtered_bed_element_names
+                        if elem in all_bed_element_names
+                    ]
+        else:
+            filtered_bed_element_names = all_bed_element_names
+
         return NTv3TracksOutput(
             bigwig_tracks_logits=bigwig_np,
             bed_tracks_logits=bed_np,
@@ -533,7 +648,7 @@ class NTv3TracksPipeline(Pipeline):
             species=meta.get("species"),
             assembly=meta.get("assembly"),
             bigwig_track_names=meta.get("bigwig_track_names"),
-            bed_element_names=self.bed_element_names,
+            bed_element_names=filtered_bed_element_names,
             window_len=meta.get("window_len"),
             pred_start=meta.get("pred_start"),
             pred_end=meta.get("pred_end"),