app.py

import gradio as gr
import torch
import joblib
import numpy as np
from itertools import product
import torch.nn as nn
import matplotlib.pyplot as plt
import io
from PIL import Image

##############################################################################
# MODEL DEFINITION
##############################################################################

class VirusClassifier(nn.Module):
    def __init__(self, input_shape: int):
        super(VirusClassifier, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(input_shape, 64),
            nn.GELU(),
            nn.BatchNorm1d(64),
            nn.Dropout(0.3),
            nn.Linear(64, 32),
            nn.GELU(),
            nn.BatchNorm1d(32),
            nn.Dropout(0.3),
            nn.Linear(32, 32),
            nn.GELU(),
            nn.Linear(32, 2)
        )

    def forward(self, x):
        return self.network(x)

##############################################################################
# UTILITIES
##############################################################################

def parse_fasta(text):
    """
    Parses FASTA formatted text into a list of (header, sequence).
    """
    sequences = []
    current_header = None
    current_sequence = []
    
    for line in text.strip().split('\n'):
        line = line.strip()
        if not line:
            continue
        if line.startswith('>'):
            if current_header:
                sequences.append((current_header, ''.join(current_sequence)))
            current_header = line[1:]
            current_sequence = []
        else:
            current_sequence.append(line.upper())
    if current_header:
        sequences.append((current_header, ''.join(current_sequence)))
    return sequences

def sequence_to_kmer_vector(sequence: str, k: int = 4) -> np.ndarray:
    """
    Convert a sequence to a k-mer frequency vector of size len(ACGT^k).
    """
    kmers = [''.join(p) for p in product("ACGT", repeat=k)]
    kmer_dict = {km: i for i, km in enumerate(kmers)}
    vec = np.zeros(len(kmers), dtype=np.float32)
    
    for i in range(len(sequence) - k + 1):
        kmer = sequence[i:i+k]
        if kmer in kmer_dict:
            vec[kmer_dict[kmer]] += 1

    total_kmers = len(sequence) - k + 1
    if total_kmers > 0:
        vec = vec / total_kmers

    return vec

def ablation_importance(model, x_tensor):
    """
    Calculates a simple ablation-based importance measure for each feature:
    1. Compute baseline human probability p_base.
    2. For each feature i, set x[i] = 0, re-run inference, compute new p, and 
       measure delta = p_base - p. 
    3. Return array of deltas (positive means that removing that feature 
       *decreases* the probability => that feature was pushing it higher).
    """
    model.eval()
    with torch.no_grad():
        # Baseline probability
        output = model(x_tensor)
        probs = torch.softmax(output, dim=1)
        p_base = probs[0, 1].item()

    # Store the delta importances
    importances = np.zeros(x_tensor.shape[1], dtype=np.float32)

    # For efficiency, we do ablation one feature at a time
    for i in range(x_tensor.shape[1]):
        x_copy = x_tensor.clone()
        x_copy[0, i] = 0.0  # Ablate this feature
        with torch.no_grad():
            output_ablation = model(x_copy)
            probs_ablation = torch.softmax(output_ablation, dim=1)
            p_ablation = probs_ablation[0, 1].item()
        # Delta
        importances[i] = p_base - p_ablation
    
    return importances, p_base

##############################################################################
# PLOTTING
##############################################################################

def create_step_and_frequency_plot(important_kmers, human_prob, title):
    """
    Creates a combined step plot (showing how each k-mer modifies the probability)
    and a frequency vs. sigma bar chart.
    """
    fig = plt.figure(figsize=(15, 10))
    
    # Create grid for subplots
    gs = plt.GridSpec(2, 1, height_ratios=[1.5, 1], hspace=0.3)
    
    # 1. Probability Step Plot
    ax1 = plt.subplot(gs[0])
    current_prob = 0.5
    steps = [('Start', current_prob, 0)]
    
    for kmer_info in important_kmers:
        change = kmer_info['impact']  # positive => pushes up, negative => pushes down
        current_prob += change
        steps.append((kmer_info['kmer'], current_prob, change))
    
    x = range(len(steps))
    y = [step[1] for step in steps]
    
    # Plot steps
    ax1.step(x, y, 'b-', where='post', label='Probability', linewidth=2)
    ax1.plot(x, y, 'b.', markersize=10)
    
    # Add reference line
    ax1.axhline(y=0.5, color='r', linestyle='--', label='Neutral (0.5)')
    
    # Customize plot
    ax1.grid(True, linestyle='--', alpha=0.7)
    ax1.set_ylim(0, 1)
    ax1.set_ylabel('Human Probability')
    ax1.set_title(f'K-mer Contributions to Prediction (final prob: {human_prob:.3f})')
    
    # Add labels for each point
    for i, (kmer, prob, change) in enumerate(steps):
        # Add k-mer label
        ax1.annotate(kmer, 
                     (i, prob),
                     xytext=(0, 10 if i % 2 == 0 else -20),
                     textcoords='offset points',
                     ha='center',
                     rotation=45)
        
        # Add change value
        if i > 0:
            change_text = f'{change:+.3f}'
            color = 'green' if change > 0 else 'red'
            ax1.annotate(change_text,
                         (i, prob),
                         xytext=(0, -20 if i % 2 == 0 else 10),
                         textcoords='offset points',
                         ha='center',
                         color=color)
    
    ax1.legend()
    
    # 2. K-mer Frequency and Sigma Plot
    ax2 = plt.subplot(gs[1])
    
    # Prepare data
    kmers = [k['kmer'] for k in important_kmers]
    frequencies = [k['occurrence'] for k in important_kmers]
    sigmas = [k['sigma'] for k in important_kmers]
    
    # Color the bars: if impact>0 => green, else red
    colors = ['g' if k['impact'] > 0 else 'r' for k in important_kmers]
    
    # Create bar plot for frequencies
    x = np.arange(len(kmers))
    width = 0.35
    
    ax2.bar(x - width/2, frequencies, width, label='Frequency (%)', color=colors, alpha=0.6)
    
    # Twin axis for sigma
    ax2_twin = ax2.twinx()
    # To highlight positive or negative sigma, pick color accordingly
    sigma_colors = []
    for s, c in zip(sigmas, colors):
        if s >= 0:
            sigma_colors.append('blue')  # above average
        else:
            sigma_colors.append('gray')  # below average
            
    ax2_twin.bar(x + width/2, sigmas, width, label='σ from Mean', color=sigma_colors, alpha=0.3)
    
    # Customize plot
    ax2.set_xticks(x)
    ax2.set_xticklabels(kmers, rotation=45)
    ax2.set_ylabel('Frequency (%)')
    ax2_twin.set_ylabel('Standard Deviations (σ) from Mean')
    ax2.set_title('K-mer Frequencies and Statistical Significance')
    
    # Add legends
    lines1, labels1 = ax2.get_legend_handles_labels()
    lines2, labels2 = ax2_twin.get_legend_handles_labels()
    ax2.legend(lines1 + lines2, labels1 + labels2, loc='upper right')
    
    plt.tight_layout()
    return fig

def create_shap_like_bar_plot(impact_values, kmer_list, top_k):
    """
    Creates a horizontal bar plot showing the top_k features by absolute impact.
    impact_values: array of float (length=256).
    kmer_list: list of all k=4 kmers in order.
    top_k: integer, how many top features to display.
    """
    # Sort by absolute impact
    indices_sorted = np.argsort(np.abs(impact_values))[::-1]
    top_indices = indices_sorted[:top_k]
    
    top_impacts = impact_values[top_indices]
    top_kmers = [kmer_list[i] for i in top_indices]
    
    fig = plt.figure(figsize=(8, 6))
    plt.barh(range(len(top_impacts)), top_impacts, color=['green' if i > 0 else 'red' for i in top_impacts])
    plt.yticks(range(len(top_impacts)), top_kmers)
    plt.xlabel("Impact on Human Probability (Ablation)")
    plt.title(f"Top {top_k} K-mers by Absolute Impact")
    plt.gca().invert_yaxis()  # Highest at top
    plt.tight_layout()
    return fig

def create_global_bar_plot(impact_values, kmer_list):
    """
    Creates a bar plot for ALL features (256) to see the global distribution.
    """
    fig = plt.figure(figsize=(12, 6))
    indices_sorted = np.argsort(np.abs(impact_values))[::-1]
    sorted_impacts = impact_values[indices_sorted]
    sorted_kmers = [kmer_list[i] for i in indices_sorted]
    
    plt.bar(range(len(sorted_impacts)), sorted_impacts, 
            color=['green' if i > 0 else 'red' for i in sorted_impacts])
    plt.title("Global Impact of All 256 K-mers (Ablation Method)")
    plt.xlabel("K-mer (sorted by |impact|)")
    plt.ylabel("Impact on Human Probability")
    # Optionally, we can skip labeling all 256 on x-axis. 
    # But we can show only the top/bottom or none for clarity.
    plt.tight_layout()
    return fig

##############################################################################
# MAIN PREDICTION FUNCTION
##############################################################################

def predict(file_obj, top_kmers=10, advanced_plots=False, fasta_text=""):
    """
    Main prediction function called by Gradio.
    - file_obj: optional uploaded FASTA file
    - top_kmers: number of top k-mers to display in the main SHAP-like plot
    - advanced_plots: bool, whether to return global bar plots
    - fasta_text: optional direct-pasted FASTA text
    """
    # Priority: If user pasted text, use that; otherwise use uploaded file.
    if fasta_text.strip():
        text = fasta_text.strip()
    else:
        if file_obj is None:
            return "No FASTA input provided", None, None, None
        try:
            if isinstance(file_obj, str):
                text = file_obj
            else:
                text = file_obj.decode('utf-8')
        except Exception as e:
            return f"Error reading file: {str(e)}", None, None, None

    # Parse FASTA
    sequences = parse_fasta(text)
    if len(sequences) == 0:
        return "No valid FASTA sequences found", None, None, None
    header, seq = sequences[0]

    # Load model + scaler
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = VirusClassifier(256).to(device)
    try:
        state_dict = torch.load('model.pt', map_location=device)
        model.load_state_dict(state_dict)
        scaler = joblib.load('scaler.pkl')
    except Exception as e:
        return f"Error loading model or scaler: {str(e)}", None, None, None

    # Prepare the vector
    raw_freq_vector = sequence_to_kmer_vector(seq, k=4)
    scaled_vector = scaler.transform(raw_freq_vector.reshape(1, -1))
    X_tensor = torch.FloatTensor(scaled_vector).to(device)

    # Compute ablation-based importances
    importances, p_base = ablation_importance(model, X_tensor)
    # p_base is baseline human probability

    # We also want frequency in % and sigma from mean
    # If your scaler is e.g. StandardScaler, then "scaled_vector[0][i]" is
    # how many std devs from the mean that feature is. 
    # We'll gather info in a list of dicts for each k-mer.
    kmers_4 = [''.join(p) for p in product("ACGT", repeat=4)]
    kmer_dict = {km: i for i, km in enumerate(kmers_4)}

    # We'll sort by absolute impact to get the top 10 by default.
    abs_sorted_idx = np.argsort(np.abs(importances))[::-1]
    # But for the final step/frequency plot we only show top_kmers
    top_indices = abs_sorted_idx[:top_kmers]

    # Build a list of the top k-mers
    important_kmers = []
    for idx in top_indices:
        # "impact" is how much that feature changed the probability
        impact = importances[idx]
        # raw frequency => raw_freq_vector[idx] * 100 for %
        freq_pct = float(raw_freq_vector[idx] * 100.0)
        # sigma => scaled_vector[0][idx]
        sigma_val = float(scaled_vector[0][idx])

        important_kmers.append({
            'kmer': kmers_4[idx],
            'impact': impact,
            'occurrence': freq_pct,
            'sigma': sigma_val
        })
    
    # For text output
    # We decide final class based on model's direct output
    with torch.no_grad():
        output = model(X_tensor)
        probs = torch.softmax(output, dim=1)
    pred_class = 1 if probs[0,1] > probs[0,0] else 0
    pred_label = 'human' if pred_class == 1 else 'non-human'
    human_prob = probs[0,1].item()
    nonhuman_prob = probs[0,0].item()
    confidence = max(human_prob, nonhuman_prob)

    results_text = (f"Sequence: {header}\n"
                    f"Prediction: {pred_label}\n"
                    f"Confidence: {confidence:.4f}\n"
                    f"Human probability: {human_prob:.4f}\n"
                    f"Non-human probability: {nonhuman_prob:.4f}\n"
                    f"Most influential k-mers (by ablation impact):\n")
    
    for kmer_info in important_kmers:
        # sign => if impact>0 => removing it lowers p(human), so it was pushing p(human) up
        direction = "UP (toward human)" if kmer_info['impact'] > 0 else "DOWN (toward non-human)"
        results_text += (
            f"  {kmer_info['kmer']}: {direction}, "
            f"Impact={kmer_info['impact']:.4f}, "
            f"Occ={kmer_info['occurrence']:.2f}% of seq, "
            f"{abs(kmer_info['sigma']):.2f}σ "
            + ("above" if kmer_info['sigma']>0 else "below")
            + " mean\n"
        )

    # PLOT 1: A SHAP-like bar plot for the top K features
    shap_fig = create_shap_like_bar_plot(importances, kmers_4, top_kmers)

    # PLOT 2: Step + frequency plot for the top K features
    step_fig = create_step_and_frequency_plot(important_kmers, human_prob, header)

    # PLOT 3 (optional advanced): global bar plot of all 256 features
    global_fig = None
    if advanced_plots:
        global_fig = create_global_bar_plot(importances, kmers_4)

    # Convert figures to PIL Images
    def fig_to_image(fig):
        buf = io.BytesIO()
        fig.savefig(buf, format='png', bbox_inches='tight', dpi=200)
        buf.seek(0)
        im = Image.open(buf)
        plt.close(fig)
        return im

    shap_img = fig_to_image(shap_fig)
    step_img = fig_to_image(step_fig)
    if global_fig is not None:
        global_img = fig_to_image(global_fig)
    else:
        global_img = None

    return results_text, shap_img, step_img, global_img

##############################################################################
# GRADIO INTERFACE
##############################################################################

title_text = "Virus Host Classifier"
description_text = """
Upload or paste a FASTA sequence to predict if it's likely **human** or **non-human** origin.
- **k=4** k-mers are used as features.
- We display ablation-based feature importance for interpretability.
- Advanced plots can be toggled to see the global distribution of all 256 k-mer impacts.
"""

iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.File(label="Upload FASTA file", type="binary", optional=True),
        gr.Slider(label="Number of top k-mers to show", minimum=1, maximum=50, value=10, step=1),
        gr.Checkbox(label="Show advanced (global) plots?", value=False),
        gr.Textbox(label="Or paste FASTA text here", lines=5, placeholder=">header\nACGTACGT...")
    ],
    outputs=[
        gr.Textbox(label="Results", lines=10),
        gr.Image(label="SHAP-like Top-k K-mer Bar Plot"),
        gr.Image(label="Step & Frequency Plot (Top-k)"),
        gr.Image(label="Global 256-K-mer Plot (advanced)", optional=True)
    ],
    title=title_text,
    description=description_text
)

if __name__ == "__main__":
    iface.launch(share=True)