geney 1.4.40__py3-none-any.whl → 1.4.41__py3-none-any.whl

geney/__init__.py

@@ -4,12 +4,17 @@ from .engines import (
     sai_predict_probs,
     run_spliceai_seq,
     run_splicing_engine,
+    predict_splicing,
+    adjoin_splicing_outcomes,
 )
 from .transcripts import TranscriptLibrary
-from .splicing_table import adjoin_splicing_outcomes
 from .splice_graph import SpliceSimulator
-from .pipelines import oncosplice_pipeline_single_transcript
-from .samples import *
+from .pipelines import (
+    oncosplice_pipeline,
+    oncosplice_top_isoform,
+    max_splicing_delta,
+    oncosplice_pipeline_single_transcript,  # backwards compat
+)
 
 __all__ = [
     "Mutation",
@@ -18,8 +23,16 @@ __all__ = [
     "sai_predict_probs",
     "run_spliceai_seq",
     "run_splicing_engine",
-    "TranscriptLibrary",
+    "predict_splicing",
     "adjoin_splicing_outcomes",
+    "TranscriptLibrary",
     "SpliceSimulator",
+    "oncosplice_pipeline",
+    "oncosplice_top_isoform",
+    "max_splicing_delta",
     "oncosplice_pipeline_single_transcript",
-]
+]
+
+
+mut_id = 'KRAS:12:25227343:G:T'
+epistasis_id = 'KRAS:12:25227343:G:T|KRAS:12:25227344:A:T'

geney/engines.py

@@ -2,137 +2,102 @@
 from __future__ import annotations
 
 from typing import Dict, List, Tuple, Optional, Union
-
 import numpy as np
 
-# These are your existing helpers; keep them in separate modules if you want.
-# from ._spliceai_utils import one_hot_encode, sai_models  # type: ignore
-# from ._pangolin_utils import pangolin_predict_probs, pang_models  # type: ignore
-import torch
-from pkg_resources import resource_filename
-from pangolin.model import *
-import numpy as np
-import sys
+# Lazy-loaded model containers (loaded automatically on first use)
+_pang_models = None
+_sai_models = None
+_pang_device = None
+_sai_device = None
+
 
-pang_model_nums = [0, 1, 2, 3, 4, 5, 6, 7]
-pang_models = []
+def _get_torch_device():
+    """Get the best available device for PyTorch."""
+    import sys
+    import torch
 
-def get_best_device():
-    """Get the best available device for computation."""
     if sys.platform == 'darwin' and torch.backends.mps.is_available():
         try:
-            # Test MPS availability
             torch.tensor([1.0], device="mps")
             return torch.device("mps")
         except RuntimeError:
-            print("Warning: MPS not available, falling back to CPU")
             return torch.device("cpu")
     elif torch.cuda.is_available():
         return torch.device("cuda")
-    else:
-        return torch.device("cpu")
+    return torch.device("cpu")
+
+
+def _get_tensorflow_device():
+    """Get the best available TensorFlow device."""
+    import sys
+    import tensorflow as tf
+
+    try:
+        if tf.config.list_physical_devices('GPU'):
+            return '/GPU:0'
+        elif sys.platform == 'darwin' and tf.config.list_physical_devices('MPS'):
+            return '/device:GPU:0'
+    except Exception:
+        pass
+    return '/CPU:0'
+
+
+def _load_pangolin_models():
+    """Lazy load Pangolin models."""
+    global _pang_models, _pang_device
+
+    if _pang_models is not None:
+        return _pang_models
+
+    import torch
+    from pkg_resources import resource_filename
+    from pangolin.model import Pangolin, L, W, AR
+
+    _pang_device = _get_torch_device()
+    print(f"Pangolin loading to {_pang_device}...")
 
-device = get_best_device()
-print(f"Pangolin loaded to {device}.")
+    _pang_models = []
+    pang_model_nums = [0, 1, 2, 3, 4, 5, 6, 7]
 
-# Initialize models with improved error handling
-try:
     for i in pang_model_nums:
         for j in range(1, 6):
             try:
-                model = Pangolin(L, W, AR).to(device)
-                
-                # Load weights with proper device mapping
+                model = Pangolin(L, W, AR).to(_pang_device)
                 model_path = resource_filename("pangolin", f"models/final.{j}.{i}.3")
-                weights = torch.load(model_path, weights_only=True, map_location=device)
-                
+                weights = torch.load(model_path, weights_only=True, map_location=_pang_device)
                 model.load_state_dict(weights)
                 model.eval()
-                pang_models.append(model)
-                
+                _pang_models.append(model)
             except Exception as e:
                 print(f"Warning: Failed to load Pangolin model {j}.{i}: {e}")
-                continue
-                
-except Exception as e:
-    print(f"Error initializing Pangolin models: {e}")
-    pang_models = []
 
+    print(f"Pangolin loaded ({len(_pang_models)} models).")
+    return _pang_models
 
-def pang_one_hot_encode(seq: str) -> np.ndarray:
-    """One-hot encode DNA sequence for Pangolin model.
-    
-    Args:
-        seq: DNA sequence string
-        
-    Returns:
-        One-hot encoded array of shape (len(seq), 4)
-        
-    Raises:
-        ValueError: If sequence contains invalid characters
-    """
-    if not isinstance(seq, str):
-        raise TypeError(f"Expected string, got {type(seq).__name__}")
-    
-    IN_MAP = np.asarray([[0, 0, 0, 0],  # N or unknown
-                         [1, 0, 0, 0],  # A
-                         [0, 1, 0, 0],  # C
-                         [0, 0, 1, 0],  # G
-                         [0, 0, 0, 1]]) # T
-    
-    # Validate sequence
-    valid_chars = set('ACGTN')
-    if not all(c.upper() in valid_chars for c in seq):
-        raise ValueError("Sequence contains invalid characters (only A, C, G, T, N allowed)")
-    
-    # Convert to numeric representation
-    seq = seq.upper().replace('A', '1').replace('C', '2')
-    seq = seq.replace('G', '3').replace('T', '4').replace('N', '0')
-    
-    try:
-        seq_array = np.asarray(list(map(int, list(seq))))
-        return IN_MAP[seq_array.astype('int8')]
-    except (ValueError, IndexError) as e:
-        raise ValueError(f"Failed to encode sequence: {e}") from e
 
+def _load_spliceai_models():
+    """Lazy load SpliceAI models."""
+    global _sai_models, _sai_device
 
+    if _sai_models is not None:
+        return _sai_models
 
+    import os
+    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
 
-import os
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+    import sys
+    import tensorflow as tf
+    from keras.models import load_model
+    from importlib import resources
 
-import absl.logging
-absl.logging.set_verbosity(absl.logging.ERROR)
+    import absl.logging
+    absl.logging.set_verbosity(absl.logging.ERROR)
 
-import os
-import sys
-import tensorflow as tf
-import numpy as np
-from keras.models import load_model
-from importlib import resources
+    _sai_device = _get_tensorflow_device()
+    print(f"SpliceAI loading to {_sai_device}...")
 
+    _sai_models = []
 
-# Force device selection with error handling
-def get_best_tensorflow_device():
-    """Get the best available TensorFlow device."""
-    try:
-        # Try GPU first
-        if tf.config.list_physical_devices('GPU'):
-            return '/GPU:0'
-        # Try MPS on macOS
-        elif sys.platform == 'darwin' and tf.config.list_physical_devices('MPS'):
-            return '/device:GPU:0'
-        else:
-            return '/CPU:0'
-    except Exception as e:
-        print(f"Warning: Device selection failed, using CPU: {e}")
-        return '/CPU:0'
-
-device = get_best_tensorflow_device()
-
-# Model loading paths with error handling
-def load_spliceai_models():
-    """Load SpliceAI models with proper error handling."""
     try:
         if sys.platform == 'darwin':
             model_filenames = [f"models/spliceai{i}.h5" for i in range(1, 6)]
@@ -140,168 +105,312 @@ def load_spliceai_models():
         else:
             model_paths = [f"/tamir2/nicolaslynn/tools/SpliceAI/spliceai/models/spliceai{i}.h5"
                            for i in range(1, 6)]
-        
-        # Load models onto correct device
-        models = []
-        with tf.device(device):
+
+        with tf.device(_sai_device):
             for i, model_path in enumerate(model_paths):
                 try:
                     model = load_model(str(model_path))
-                    models.append(model)
+                    _sai_models.append(model)
                 except Exception as e:
                     print(f"Warning: Failed to load SpliceAI model {i+1}: {e}")
-                    continue
-        
-        if not models:
-            raise RuntimeError("No SpliceAI models could be loaded")
-            
-        return models
-        
     except Exception as e:
         print(f"Error loading SpliceAI models: {e}")
-        return []
 
-sai_models = load_spliceai_models()
+    print(f"SpliceAI loaded ({len(_sai_models)} models).")
+    return _sai_models
+
+
+def pang_one_hot_encode(seq: str) -> np.ndarray:
+    """One-hot encode DNA sequence for Pangolin model."""
+    if not isinstance(seq, str):
+        raise TypeError(f"Expected string, got {type(seq).__name__}")
+
+    IN_MAP = np.asarray([[0, 0, 0, 0],  # N
+                         [1, 0, 0, 0],  # A
+                         [0, 1, 0, 0],  # C
+                         [0, 0, 1, 0],  # G
+                         [0, 0, 0, 1]]) # T
+
+    valid_chars = set('ACGTN')
+    if not all(c.upper() in valid_chars for c in seq):
+        raise ValueError("Sequence contains invalid characters")
+
+    seq = seq.upper().replace('A', '1').replace('C', '2')
+    seq = seq.replace('G', '3').replace('T', '4').replace('N', '0')
 
+    seq_array = np.asarray(list(map(int, list(seq))))
+    return IN_MAP[seq_array.astype('int8')]
 
-print(f"SpliceAI loaded to {device}.")
 
 def one_hot_encode(seq: str) -> np.ndarray:
-    """One-hot encode DNA sequence for SpliceAI model.
-    
-    Args:
-        seq: DNA sequence string
-        
-    Returns:
-        One-hot encoded array of shape (len(seq), 4)
-        
-    Raises:
-        ValueError: If sequence contains invalid characters
-    """
+    """One-hot encode DNA sequence for SpliceAI model."""
     if not isinstance(seq, str):
         raise TypeError(f"Expected string, got {type(seq).__name__}")
-    
-    # Validate sequence
+
     valid_chars = set('ACGTN')
     if not all(c.upper() in valid_chars for c in seq):
-        raise ValueError("Sequence contains invalid characters (only A, C, G, T, N allowed)")
-    
-    encoding_map = np.asarray([[0, 0, 0, 0],  # N or unknown
+        raise ValueError("Sequence contains invalid characters")
+
+    encoding_map = np.asarray([[0, 0, 0, 0],  # N
                                [1, 0, 0, 0],  # A
                                [0, 1, 0, 0],  # C
                                [0, 0, 1, 0],  # G
                                [0, 0, 0, 1]]) # T
 
-    # Convert to numeric representation
     seq = seq.upper().replace('A', '\x01').replace('C', '\x02')
     seq = seq.replace('G', '\x03').replace('T', '\x04').replace('N', '\x00')
 
-    try:
-        return encoding_map[np.frombuffer(seq.encode('latin1'), np.int8) % 5]
-    except Exception as e:
-        raise ValueError(f"Failed to encode sequence: {e}") from e
+    return encoding_map[np.frombuffer(seq.encode('latin1'), np.int8) % 5]
 
 
-def sai_predict_probs(seq: str, models: list) -> tuple[np.ndarray, np.ndarray]:
-    """
-    Predict donor and acceptor probabilities for each nt in seq using SpliceAI.
-    Returns (acceptor_probs, donor_probs) as np.ndarray of shape (L,).
+def pangolin_predict_probs(seq: str, models: list = None) -> Tuple[List[float], List[float]]:
+    """Predict splice site probabilities using Pangolin.
+
+    Pangolin outputs shape (1, 12, seq_len) where:
+    - 12 channels = 4 tissues × 3 prediction types
+    - For each tissue: [site_usage, acceptor_gain, donor_gain] or similar
+
+    We aggregate by taking max across tissues.
     """
+    import torch
+
+    if models is None:
+        models = _load_pangolin_models()
+
     if not models:
-        raise ValueError("No SpliceAI models loaded")
+        raise ValueError("No Pangolin models loaded")
 
-    if not isinstance(seq, str):
-        raise TypeError(f"Expected string, got {type(seq).__name__}")
+    x = pang_one_hot_encode(seq)
+    x = torch.tensor(x.T[None, :, :], dtype=torch.float32, device=_pang_device)
+
+    preds = []
+    with torch.no_grad():
+        for model in models:
+            pred = model(x)
+            preds.append(pred.cpu().numpy())
+
+    y = np.mean(preds, axis=0)  # Shape: (1, 12, seq_len)
+
+    # Pangolin has 12 channels organized as:
+    # Indices 0,3,6,9: site usage scores for 4 tissues
+    # Indices 1,4,7,10: acceptor gain scores for 4 tissues
+    # Indices 2,5,8,11: donor gain scores for 4 tissues
+    # Take max across the 4 tissues for each type
+
+    # Acceptor: max of channels 1, 4, 7, 10
+    acceptor_channels = y[0, [1, 4, 7, 10], :]  # (4, seq_len)
+    acceptor_probs = np.max(acceptor_channels, axis=0).tolist()
+
+    # Donor: max of channels 2, 5, 8, 11
+    donor_channels = y[0, [2, 5, 8, 11], :]  # (4, seq_len)
+    donor_probs = np.max(donor_channels, axis=0).tolist()
+
+    return donor_probs, acceptor_probs
+
+
+def sai_predict_probs(seq: str, models: list = None) -> Tuple[np.ndarray, np.ndarray]:
+    """Predict donor and acceptor probabilities using SpliceAI."""
+    if models is None:
+        models = _load_spliceai_models()
+
+    if not models:
+        raise ValueError("No SpliceAI models loaded")
 
     if len(seq) < 1000:
         raise ValueError(f"Sequence too short: {len(seq)} (expected >= 1000)")
 
-    try:
-        x = one_hot_encode(seq)[None, :]
-        preds = []
-        for i, model in enumerate(models):
-            try:
-                pred = model.predict(x, verbose=0)
-                preds.append(pred)
-            except Exception as e:
-                print(f"Warning: SpliceAI model {i+1} failed: {e}")
-        if not preds:
-            raise RuntimeError("All SpliceAI model predictions failed")
+    x = one_hot_encode(seq)[None, :].astype(np.float32)
 
-        y = np.mean(preds, axis=0)          # (1, L, 3)
-        y = y[0, :, 1:].T                   # (2, L) -> [acceptor, donor]
-        return y[0, :], y[1, :]
+    # Use direct model call instead of .predict() to avoid Jupyter kernel issues
+    preds = []
+    for model in models:
+        pred = model(x, training=False)
+        if hasattr(pred, 'numpy'):
+            pred = pred.numpy()
+        preds.append(pred)
 
-    except Exception as e:
-        raise RuntimeError(f"SpliceAI prediction failed: {e}") from e
+    y = np.mean(preds, axis=0)
+    y = y[0, :, 1:].T
+    return y[0, :], y[1, :]
 
 
 def run_spliceai_seq(
     seq: str,
     indices: Union[List[int], np.ndarray],
     threshold: float = 0.0,
-) -> tuple[Dict[int, float], Dict[int, float]]:
-    """
-    Run SpliceAI on seq and return donor / acceptor sites above threshold.
-    Returns (donor_indices, acceptor_indices) as dict[pos -> prob]
-    """
-    if not isinstance(seq, str):
-        raise TypeError(f"Expected string sequence, got {type(seq).__name__}")
-
-    if not isinstance(indices, (list, np.ndarray)):
-        raise TypeError(f"Expected list or array for indices, got {type(indices).__name__}")
-
+) -> Tuple[Dict[int, float], Dict[int, float]]:
+    """Run SpliceAI on seq and return donor/acceptor sites above threshold."""
     if len(indices) != len(seq):
         raise ValueError(f"indices length ({len(indices)}) must match sequence length ({len(seq)})")
 
-    if not isinstance(threshold, (int, float)):
-        raise TypeError(f"Threshold must be numeric, got {type(threshold).__name__}")
+    acc_probs, don_probs = sai_predict_probs(seq)
+    acceptor = {pos: p for pos, p in zip(indices, acc_probs) if p >= threshold}
+    donor = {pos: p for pos, p in zip(indices, don_probs) if p >= threshold}
+    return donor, acceptor
 
-    try:
-        acc_probs, don_probs = sai_predict_probs(seq, models=sai_models)
-        acceptor = {pos: p for pos, p in zip(indices, acc_probs) if p >= threshold}
-        donor = {pos: p for pos, p in zip(indices, don_probs) if p >= threshold}
-        return donor, acceptor
-    except Exception as e:
-        raise RuntimeError(f"SpliceAI sequence analysis failed: {e}") from e
+
+def _generate_random_sequence(length: int) -> str:
+    """Generate a random DNA sequence of given length."""
+    import random
+    return ''.join(random.choices('ACGT', k=length))
 
 
 def run_splicing_engine(
     seq: Optional[str] = None,
     engine: str = "spliceai",
 ) -> Tuple[List[float], List[float]]:
+    """Run specified splicing engine to predict splice site probabilities."""
+    if seq is None:
+        seq = _generate_random_sequence(15_001)
+
+    if not isinstance(seq, str) or not seq:
+        raise ValueError("Sequence must be a non-empty string")
+
+    valid_chars = set("ACGTN")
+    if not all(c.upper() in valid_chars for c in seq):
+        raise ValueError("Sequence contains invalid nucleotides")
+
+    match engine:
+        case "spliceai":
+            acc, don = sai_predict_probs(seq)
+            return don.tolist(), acc.tolist()
+        case "pangolin":
+            return pangolin_predict_probs(seq)
+        case _:
+            raise ValueError(f"Engine '{engine}' not implemented. Available: 'spliceai', 'pangolin'")
+
+
+# ------------------------------------------------------------------------------
+# Higher-level prediction utilities (formerly in splicing_table.py)
+# ------------------------------------------------------------------------------
+
+def predict_splicing(s, position: int, engine: str = 'spliceai', context: int = 7500):
     """
-    Run specified splicing engine to predict splice site probabilities.
+    Predict splicing probabilities at a given position using the specified engine.
+
+    Args:
+        s: Sequence object with .seq, .index, .clone(), .rev attributes
+        position: The genomic position to predict splicing probabilities for.
+        engine: The prediction engine to use. Supported: 'spliceai', 'pangolin'.
+        context: The length of the target central region (default: 7500).
 
     Returns:
-        (donor_probs, acceptor_probs) as lists
+        pd.DataFrame with position index and columns: donor_prob, acceptor_prob, nucleotides
     """
-    from .utils import generate_random_sequence  # type: ignore
+    import pandas as pd
 
-    if seq is None:
-        seq = generate_random_sequence(15_001)
+    if position < s.index.min() or position > s.index.max():
+        raise ValueError(f"Position {position} is outside sequence bounds [{s.index.min()}, {s.index.max()}]")
 
-    if not isinstance(seq, str):
-        raise TypeError(f"Sequence must be string, got {type(seq).__name__}")
-    if not seq:
-        raise ValueError("Sequence cannot be empty")
+    target = s.clone(position - context, position + context)
+
+    if len(target.seq) == 0:
+        raise ValueError(f"No sequence data found around position {position} with context {context}")
+
+    seq, indices = target.seq, target.index
+
+    if len(indices) == 0:
+        raise ValueError(f"No indices found in sequence around position {position}")
+
+    rel_pos = np.abs(indices - position).argmin()
+    left_missing, right_missing = max(0, context - rel_pos), max(0, context - (len(seq) - rel_pos))
+
+    if left_missing > 0 or right_missing > 0:
+        step = -1 if s.rev else 1
+
+        if left_missing > 0:
+            left_pad = np.arange(indices[0] - step * left_missing, indices[0], step)
+        else:
+            left_pad = np.array([], dtype=indices.dtype)
+
+        if right_missing > 0:
+            right_pad = np.arange(indices[-1] + step, indices[-1] + step * (right_missing + 1), step)
+        else:
+            right_pad = np.array([], dtype=indices.dtype)
+
+        seq = 'N' * left_missing + seq + 'N' * right_missing
+        indices = np.concatenate([left_pad, indices, right_pad])
+
+    donor_probs, acceptor_probs = run_splicing_engine(seq=seq, engine=engine)
+
+    seq = seq[5000:-5000]
+    indices = indices[5000:-5000]
+    expected_len = len(seq)
+
+    if len(donor_probs) != expected_len:
+        if len(donor_probs) > expected_len:
+            offset = (len(donor_probs) - expected_len) // 2
+            donor_probs = donor_probs[offset:offset + expected_len]
+            acceptor_probs = acceptor_probs[offset:offset + expected_len]
+        else:
+            pad_len = expected_len - len(donor_probs)
+            donor_probs = donor_probs + [0.0] * pad_len
+            acceptor_probs = acceptor_probs + [0.0] * pad_len
+
+    df = pd.DataFrame({
+        'position': indices,
+        'donor_prob': donor_probs,
+        'acceptor_prob': acceptor_probs,
+        'nucleotides': list(seq)
+    }).set_index('position').round(3)
+
+    df.attrs['name'] = s.name
+    return df
 
-    valid_chars = set("ACGTN")
-    if not all(c.upper() in valid_chars for c in seq):
-        raise ValueError("Sequence contains invalid nucleotides (only A, C, G, T, N allowed)")
+
+def adjoin_splicing_outcomes(
+    splicing_predictions: Dict[str, 'pd.DataFrame'],
+    transcript: Optional[object] = None,
+) -> 'pd.DataFrame':
+    """
+    Combine splicing predictions for multiple mutations into a multi-index DataFrame.
+
+    Args:
+        splicing_predictions: {label -> DF with 'donor_prob','acceptor_prob','nucleotides'}
+        transcript: optional transcript (must have .acceptors, .donors, .rev)
+    """
+    import pandas as pd
+
+    if not splicing_predictions:
+        raise ValueError("splicing_predictions cannot be empty")
+
+    dfs = []
+    for label, df in splicing_predictions.items():
+        if not isinstance(df, pd.DataFrame):
+            raise TypeError(f"Expected DataFrame for '{label}', got {type(df).__name__}")
+
+        required_cols = ["donor_prob", "acceptor_prob", "nucleotides"]
+        missing = [c for c in required_cols if c not in df.columns]
+        if missing:
+            raise ValueError(f"DataFrame for '{label}' missing required columns: {missing}")
+
+        var_df = df.rename(
+            columns={
+                "donor_prob": ("donors", f"{label}_prob"),
+                "acceptor_prob": ("acceptors", f"{label}_prob"),
+                "nucleotides": ("nts", f"{label}"),
+            }
+        )
+        dfs.append(var_df)
 
     try:
-        match engine:
-            case "spliceai":
-                acc, don = sai_predict_probs(seq, models=sai_models)
-                donor_probs, acceptor_probs = don.tolist(), acc.tolist()
-            case "spliceai-pytorch":
-                raise ValueError("spliceai-pytorch engine has been removed. Use 'spliceai' instead.")
-            case "pangolin":
-                donor_probs, acceptor_probs = pangolin_predict_probs(seq, models=pang_models)
-            case _:
-                raise ValueError(f"Engine '{engine}' not implemented. Available: 'spliceai', 'pangolin'")
-    except ImportError as e:
-        raise ImportError(f"Failed to import engine '{engine}': {e}") from e
-
-    return donor_probs, acceptor_probs
+        full_df = pd.concat(dfs, axis=1)
+    except Exception as e:
+        raise ValueError(f"Failed to concatenate DataFrames: {e}") from e
+
+    if not isinstance(full_df.columns, pd.MultiIndex):
+        full_df.columns = pd.MultiIndex.from_tuples(full_df.columns)
+
+    if transcript is not None:
+        full_df[("acceptors", "annotated")] = full_df.apply(
+            lambda row: row.name in transcript.acceptors, axis=1
+        )
+        full_df[("donors", "annotated")] = full_df.apply(
+            lambda row: row.name in transcript.donors, axis=1
+        )
+        full_df.sort_index(axis=1, level=0, inplace=True)
+        full_df.sort_index(ascending=not transcript.rev, inplace=True)
+    else:
+        full_df.sort_index(axis=1, level=0, inplace=True)
+
+    return full_df