feat: clean track names

bigwig_export.py

@@ -0,0 +1,167 @@
+"""
+BigWig export functionality for NTv3 tracks.
+"""
+
+import os
+import uuid
+import tempfile
+import zipfile
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+try:
+    import pyBigWig
+except ImportError:
+    pyBigWig = None
+
+if TYPE_CHECKING:
+    from ntv3_tracks_pipeline import NTv3TracksOutput
+
+
+def _softmax_last(x: np.ndarray) -> np.ndarray:
+    """Compute softmax over the last dimension."""
+    x = x - x.max(axis=-1, keepdims=True)
+    ex = np.exp(x)
+    return ex / ex.sum(axis=-1, keepdims=True)
+
+
+def create_bigwig_zip(
+    out: "NTv3TracksOutput",
+    bigwig_selected: list[str],
+    bed_elements: list[str],
+) -> str:
+    """
+    Create BigWig files for selected tracks and save them in a zip file.
+    
+    Parameters
+    ----------
+    out : NTv3TracksOutput
+        The prediction output from the pipeline.
+    bigwig_selected : list[str]
+        List of BigWig track IDs to export.
+    bed_elements : list[str]
+        List of BED element names to export.
+    
+    Returns
+    -------
+    str
+        Path to the created zip file containing BigWig files.
+    
+    Raises
+    ------
+    ImportError
+        If pyBigWig is not installed.
+    ValueError
+        If no predictions are available or no tracks are selected.
+    """
+    if pyBigWig is None:
+        raise ImportError("pyBigWig is required for BigWig export. Install with: pip install pyBigWig")
+    
+    if out is None:
+        raise ValueError("No predictions available. Please run a prediction first.")
+    
+    bw_names = out.bigwig_track_names or []
+    bw_logits = out.bigwig_tracks_logits
+    bed_names = out.bed_element_names or []
+    bed_logits = out.bed_tracks_logits
+    
+    if bw_logits is None or not bw_names:
+        raise ValueError("No BigWig tracks available in model output.")
+    
+    # Get genomic coordinates
+    chrom = out.chrom
+    if chrom is None:
+        raise ValueError("Chromosome information not available. Use genomic coordinates for BigWig export.")
+    
+    start = out.start
+    end = out.end
+    window_len = out.window_len or (end - start)
+    
+    # Calculate prediction region (center 37.5%)
+    pred_start = out.pred_start or (start + int(window_len * 0.3125))
+    pred_end = out.pred_end or (pred_start + int(window_len * 0.375))
+    
+    # Create temporary directory for BigWig files
+    tmpdir = tempfile.gettempdir()
+    output_dir = os.path.join(tmpdir, f"bigwig_outputs_{uuid.uuid4().hex}")
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Prepare track data list
+    track_data_list = []
+    
+    # Add BigWig tracks
+    for track_id in bigwig_selected:
+        if track_id in bw_names:
+            idx = bw_names.index(track_id)
+            track_data_list.append(("bigwig", track_id, idx, None))
+    
+    # Add BED elements (as probabilities)
+    if bed_logits is not None and bed_elements:
+        probs = _softmax_last(bed_logits)
+        for elem_name in bed_elements:
+            if elem_name in bed_names:
+                eidx = bed_names.index(elem_name)
+                # Store as bed element with probability data
+                track_data_list.append(("bed", elem_name, eidx, probs[:, eidx, 1]))
+    
+    if not track_data_list:
+        raise ValueError("No tracks selected for export.")
+    
+    # Create BigWig files
+    created_files = []
+    for track_type, track_id, track_idx, bed_probs in track_data_list:
+        if track_type == "bigwig":
+            track_data = bw_logits[:, track_idx].astype(np.float32)
+            display_name = track_id
+        else:  # bed
+            if bed_probs is None:
+                continue
+            track_data = bed_probs.astype(np.float32)
+            display_name = track_id
+        
+        # Clean filename
+        clean_name = display_name.replace(" ", "_").replace("/", "_").replace("-", "_")
+        bw_filename = os.path.join(output_dir, f"{clean_name}.bw")
+        
+        # Create BigWig file
+        bw = pyBigWig.open(bw_filename, "w")
+        
+        # Add header - use end of genomic window as chromosome size
+        bw.addHeader([(chrom, end)])
+        
+        # Add entries
+        num_positions = len(track_data)
+        starts = np.arange(pred_start, pred_start + num_positions, dtype=np.int64)
+        ends = starts + 1
+        values = track_data.tolist()
+        
+        bw.addEntries(
+            chroms=[chrom] * len(starts),
+            starts=starts.tolist(),
+            ends=ends.tolist(),
+            values=values
+        )
+        
+        bw.close()
+        created_files.append(bw_filename)
+    
+    # Create zip file
+    zip_path = os.path.join(tmpdir, f"ntv3_tracks_{uuid.uuid4().hex}.zip")
+    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        for bw_file in created_files:
+            zipf.write(bw_file, os.path.basename(bw_file))
+    
+    # Clean up individual BigWig files
+    for bw_file in created_files:
+        try:
+            os.remove(bw_file)
+        except:
+            pass
+    try:
+        os.rmdir(output_dir)
+    except:
+        pass
+    
+    return zip_path
+