app.py

import gradio as gr
import torch
import joblib
import numpy as np
import shap
import random
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)
    
    def get_feature_importance(self, x):
        """
        Calculate gradient-based feature importance, specifically for the 
        'human' class (index=1) by computing gradient of that probability wrt x.
        """
        x.requires_grad_(True)
        output = self.network(x)
        probs = torch.softmax(output, dim=1)
        
        # Probability of 'human' class (index=1)
        human_prob = probs[..., 1]
        if x.grad is not None:
            x.grad.zero_()
        human_prob.backward()
        importance = x.grad  # shape: (batch_size, n_features)
        
        return importance, float(human_prob)

###############################################################################
# Utility Functions
###############################################################################
def parse_fasta(text):
    """Parses text input in FASTA format into a list of (header, sequence)."""
    sequences = []
    current_header = None
    current_sequence = []
    
    for line in text.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 single nucleotide sequence to a k-mer frequency vector."""
    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  # normalize frequencies

    return vec

###############################################################################
# Additional Plots
###############################################################################
def create_probability_bar_plot(prob_human, prob_nonhuman):
    """
    Simple bar plot comparing human vs. non-human probabilities.
    """
    labels = ["Non-human", "Human"]
    probs = [prob_nonhuman, prob_human]
    colors = ["red", "green"]

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.bar(labels, probs, color=colors, alpha=0.7)
    ax.set_ylim(0, 1)
    for i, v in enumerate(probs):
        ax.text(i, v+0.02, f"{v:.3f}", ha='center', color='black', fontsize=11)

    ax.set_title("Predicted Probabilities")
    ax.set_ylabel("Probability")
    plt.tight_layout()
    return fig

def create_frequency_sigma_plot(important_kmers, title):
    """
    Creates a bar plot of the top k-mers (by importance) showing 
    frequency (%) and σ from mean.
    """
    # Sort by absolute impact
    sorted_kmers = sorted(important_kmers, key=lambda x: x['impact'], reverse=True)
    kmers = [k["kmer"] for k in sorted_kmers]
    frequencies = [k["occurrence"] for k in sorted_kmers]  # in %
    sigmas = [k["sigma"] for k in sorted_kmers]
    directions = [k["direction"] for k in sorted_kmers]
    
    x = np.arange(len(kmers))
    width = 0.4

    fig, ax_bar = plt.subplots(figsize=(10, 5))

    # Bar for frequency
    bars_freq = ax_bar.bar(
        x - width/2, frequencies, width, alpha=0.7,
        color=["green" if d=="human" else "red" for d in directions],
        label="Frequency (%)"
    )
    ax_bar.set_ylabel("Frequency (%)")
    ax_bar.set_ylim(0, max(frequencies) * 1.2 if len(frequencies) > 0 else 1)

    # Twin axis for σ
    ax_bar_twin = ax_bar.twinx()
    bars_sigma = ax_bar_twin.bar(
        x + width/2, sigmas, width, alpha=0.5, color="gray", label="σ from Mean"
    )
    ax_bar_twin.set_ylabel("Standard Deviations (σ)")

    ax_bar.set_title(f"Frequency & σ from Mean for Top k-mers — {title}")
    ax_bar.set_xticks(x)
    ax_bar.set_xticklabels(kmers, rotation=45, ha='right')

    # Combined legend
    lines1, labels1 = ax_bar.get_legend_handles_labels()
    lines2, labels2 = ax_bar_twin.get_legend_handles_labels()
    ax_bar.legend(lines1 + lines2, labels1 + labels2, loc="upper right")

    plt.tight_layout()
    return fig

def create_importance_bar_plot(important_kmers, title):
    """
    Create a simple bar chart showing the absolute gradient magnitude 
    for the top k-mers, sorted descending.
    """
    sorted_kmers = sorted(important_kmers, key=lambda x: x['impact'], reverse=True)
    kmers = [k['kmer'] for k in sorted_kmers]
    impacts = [k['impact'] for k in sorted_kmers]
    directions = [k["direction"] for k in sorted_kmers]

    x = np.arange(len(kmers))

    fig, ax = plt.subplots(figsize=(10, 5))
    bar_colors = ["green" if d=="human" else "red" for d in directions]

    ax.bar(x, impacts, color=bar_colors, alpha=0.7, edgecolor='black')
    ax.set_xticks(x)
    ax.set_xticklabels(kmers, rotation=45, ha='right')
    ax.set_title(f"Absolute Feature Importance (Top k-mers) — {title}")
    ax.set_ylabel("Gradient Magnitude")
    ax.grid(axis="y", alpha=0.3)

    plt.tight_layout()
    return fig

###############################################################################
# SHAP Beeswarm
###############################################################################
def create_shap_beeswarm_plot(
    model, 
    input_vector: np.ndarray, 
    background_data: np.ndarray, 
    feature_names: list
):
    """
    Creates a SHAP beeswarm plot using KernelExplainer for the given model and data.
    
    Parameters
    ----------
    model : nn.Module
        Trained PyTorch model (binary classifier).
    input_vector : np.ndarray
        The 1-sample input (or multiple samples) we want SHAP values for.
    background_data : np.ndarray
        Background samples for KernelExplainer. Should have shape (N, #features).
    feature_names : list
        Names for each feature (k-mers).
    
    Returns
    -------
    fig : matplotlib Figure
        Beeswarm plot figure.
    """

    # We'll define a prediction function that shap can call
    # The model outputs logits for shape [N, 2]
    # We want the raw outputs for each class. SHAP will handle the link function if needed.
    def predict_fn(data):
        """
        data: shape (N, #features)
        returns: shape (N, 2) for 2-class logits
        """
        with torch.no_grad():
            x = torch.FloatTensor(data)
            logits = model(x)
            return logits.detach().cpu().numpy()

    # Create KernelExplainer
    explainer = shap.KernelExplainer(
        model=predict_fn,
        data=background_data
    )

    # Compute SHAP values
    # For a 2-class model, shap_values is a list of length 2 => [class0 array, class1 array]
    # Each array is shape (N, #features).
    shap_values = explainer.shap_values(input_vector)

    # We’ll produce a beeswarm for the 'human' class (class index=1).
    # If we have only 1 sample, the beeswarm won't be too interesting, but let's do it anyway.
    class_idx = 1  # 'human'
    
    # If we only have one sample, place it in an array for shap summary plotting:
    # We can do shap_values[class_idx].shape => (1, #features) for a single sample
    # Beeswarm typically expects multiple samples. We'll plot anyway.
    shap.plots.beeswarm(
        shap_values[class_idx],
        feature_names=feature_names,
        show=False
    )

    fig = plt.gcf()
    fig.set_size_inches(8, 6)
    plt.title("SHAP Beeswarm Plot (Class: Human)")

    plt.tight_layout()
    return fig

###############################################################################
# Prediction Function
###############################################################################
def predict(file_obj):
    """
    Main function for Gradio:
      1. Reads the uploaded FASTA file or text.
      2. Loads the model and scaler.
      3. Generates predictions, probabilities, and top k-mers.
      4. Creates multiple outputs:
         - Text summary (Markdown)
         - Probability Bar Plot
         - SHAP Beeswarm Plot
         - Frequency & σ Plot
         - Absolute Feature Importance Bar Plot
    """
    # 0. Basic file read
    if file_obj is None:
        return (
            "Please upload a FASTA file.",
            None,
            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,
            None
        )

    # 1. Parse FASTA
    sequences = parse_fasta(text)
    if len(sequences) == 0:
        return (
            "No valid FASTA sequences found. Please check your input.",
            None,
            None,
            None,
            None
        )
    header, seq = sequences[0]  # We'll classify only the first sequence

    # 2. Prepare model, scaler, and input
    k = 4
    device = "cuda" if torch.cuda.is_available() else "cpu"
    try:
        raw_freq_vector = sequence_to_kmer_vector(seq, k=k)

        # Load model & scaler
        model = VirusClassifier(input_shape=4**k).to(device)
        state_dict = torch.load("model.pt", map_location=device)
        model.load_state_dict(state_dict)
        scaler = joblib.load("scaler.pkl")
        model.eval()

        scaled_vector = scaler.transform(raw_freq_vector.reshape(1, -1))
        X_tensor = torch.FloatTensor(scaled_vector).to(device)

        # 3. Predict
        with torch.no_grad():
            logits = model(X_tensor)
            probs = torch.softmax(logits, dim=1)
        human_prob = float(probs[0][1])
        non_human_prob = float(probs[0][0])
        pred_label = "human" if human_prob >= non_human_prob else "non-human"
        confidence = float(max(probs[0]))

        # 4. Gradient-based feature importance
        importance, hum_prob_grad = model.get_feature_importance(X_tensor)
        importances = importance[0].cpu().numpy()  # shape: (#features,)
        abs_importances = np.abs(importances)

        # 5. Gather k-mer strings
        kmers_list = [''.join(p) for p in product("ACGT", repeat=k)]
        # top 10 by absolute importance
        top_k = 10
        top_idxs = np.argsort(abs_importances)[-top_k:][::-1]
        important_kmers = []
        for idx in top_idxs:
            direction = "human" if importances[idx] > 0 else "non-human"
            freq_percent = float(raw_freq_vector[idx] * 100.0)
            sigma_val = float(scaled_vector[0][idx])  # scaled / standardized val
            important_kmers.append({
                'kmer': kmers_list[idx],
                'idx': idx,
                'impact': abs_importances[idx],
                'direction': direction,
                'occurrence': freq_percent,
                'sigma': sigma_val
            })

        # 6. Generate text summary
        text_summary = (
            f"**Sequence Header**: {header}\n\n"
            f"**Predicted Label**: {pred_label}\n"
            f"**Confidence**: {confidence:.4f}\n\n"
            f"**Human Probability**: {human_prob:.4f}\n"
            f"**Non-human Probability**: {non_human_prob:.4f}\n\n"
            "### Most Influential k-mers:\n"
        )
        for km in important_kmers:
            direction_text = f"(pushes toward {km['direction']})"
            freq_text = f"{km['occurrence']:.2f}%"
            sigma_text = (
                f"{abs(km['sigma']):.2f}σ "
                + ("above" if km['sigma'] > 0 else "below")
                + " mean"
            )
            text_summary += (
                f"- **{km['kmer']}**: impact={km['impact']:.4f}, {direction_text}, "
                f"occurrence={freq_text}, ({sigma_text})\n"
            )

        # 7. Probability Bar Plot
        fig_prob = create_probability_bar_plot(human_prob, non_human_prob)
        buf_prob = io.BytesIO()
        fig_prob.savefig(buf_prob, format='png', bbox_inches='tight', dpi=120)
        buf_prob.seek(0)
        prob_img = Image.open(buf_prob)
        plt.close(fig_prob)

        # 8. SHAP Beeswarm Plot
        # We need some background data for KernelExplainer. Let's create a small random sample
        # or sample from the scaled_vector itself in a repeated manner. Real usage: choose a valid background set.
        background_size = 5  # keep small for speed
        # We'll pick random sequences from normal(0,1) or from scaled_vector repeated
        background_data = []
        for _ in range(background_size):
            # Option A: random small variations around scaled_vector
            # new_sample = scaled_vector[0] + np.random.normal(0, 0.5, size=scaled_vector.shape[1])
            # Option B: just clone the same scaled vector multiple times
            new_sample = scaled_vector[0]
            background_data.append(new_sample)
        background_data = np.stack(background_data, axis=0)  # shape (5, #features)

        fig_bee = create_shap_beeswarm_plot(
            model=model,
            input_vector=scaled_vector,      # our single sample
            background_data=background_data, # background for KernelExplainer
            feature_names=kmers_list
        )
        buf_bee = io.BytesIO()
        fig_bee.savefig(buf_bee, format='png', bbox_inches='tight', dpi=120)
        buf_bee.seek(0)
        bee_img = Image.open(buf_bee)
        plt.close(fig_bee)

        # 9. Frequency & σ Plot
        fig_freq = create_frequency_sigma_plot(important_kmers, header)
        buf_freq = io.BytesIO()
        fig_freq.savefig(buf_freq, format='png', bbox_inches='tight', dpi=120)
        buf_freq.seek(0)
        freq_img = Image.open(buf_freq)
        plt.close(fig_freq)

        # 10. Absolute Feature Importance Bar Plot
        fig_imp = create_importance_bar_plot(important_kmers, header)
        buf_imp = io.BytesIO()
        fig_imp.savefig(buf_imp, format='png', bbox_inches='tight', dpi=120)
        buf_imp.seek(0)
        imp_img = Image.open(buf_imp)
        plt.close(fig_imp)

        return text_summary, prob_img, bee_img, freq_img, imp_img
    
    except Exception as e:
        return (
            f"Error during prediction or visualization: {str(e)}",
            None,
            None,
            None,
            None
        )


###############################################################################
# Gradio Interface
###############################################################################
with gr.Blocks(title="Advanced Virus Host Classifier with SHAP Beeswarm") as demo:
    gr.Markdown(
        """
        # Advanced Virus Host Classifier (SHAP Beeswarm Edition)

        **Upload a FASTA file** containing a single nucleotide sequence. 
        The model will predict whether this sequence is **human** or **non-human**, 
        provide a confidence score, and highlight the most influential k-mers. 
        We also produce a **SHAP beeswarm** plot for the features.
        
        ---
        **Note**: Beeswarm plots are usually most insightful with multiple samples. 
        Here, we demonstrate usage with a single sample plus a small synthetic background.
        """
    )
    
    with gr.Row():
        file_in = gr.File(label="Upload FASTA", type="binary")
        btn = gr.Button("Run Prediction")

    # We will create multiple tabs for our outputs
    with gr.Tabs():
        with gr.Tab("Prediction Results"):
            md_out = gr.Markdown()
        with gr.Tab("Probability Plot"):
            prob_out = gr.Image()
        with gr.Tab("SHAP Beeswarm Plot"):
            bee_out = gr.Image()
        with gr.Tab("Frequency & σ Plot"):
            freq_out = gr.Image()
        with gr.Tab("Importance Bar Plot"):
            imp_out = gr.Image()

    # Link the button
    btn.click(
        fn=predict,
        inputs=[file_in],
        outputs=[md_out, prob_out, bee_out, freq_out, imp_out]
    )

if __name__ == "__main__":
    # By default, share=False. You can set share=True for external access.
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)