Update app.py

app.py

@@ -2,14 +2,13 @@ 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
-matplotlib.use("Agg")  # If no display environment
 import matplotlib.pyplot as plt
 import io
 from PIL import Image
-import shap
 
 ###############################################################################
 # Model Definition
@@ -33,35 +32,30 @@ class VirusClassifier(nn.Module):
 
     def forward(self, x):
         return self.network(x)
-
-
-###############################################################################
-# Torch Model Wrapper for SHAP
-###############################################################################
-class TorchModelWrapper:
-    """
-    A simple callable that converts incoming NumPy arrays to torch Tensors,
-    does a forward pass, and returns NumPy arrays. This is needed for shap.
-    """
-    def __init__(self, model: nn.Module, device='cpu'):
-        self.model = model
-        self.device = device
-
-    def __call__(self, x_np: np.ndarray):
-        x_torch = torch.from_numpy(x_np).float().to(self.device)
-        with torch.no_grad():
-            out = self.model(x_torch).cpu().numpy()  # shape=(batch,2)
-        return out
-
+    
+    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):
-    """
-    Parse FASTA text, returning a list of (header, sequence).
-    We'll only use the *first* sequence in practice.
-    """
+    """Parses text input in FASTA format into a list of (header, sequence)."""
     sequences = []
     current_header = None
     current_sequence = []
@@ -82,9 +76,7 @@ def parse_fasta(text):
     return sequences
 
 def sequence_to_kmer_vector(sequence: str, k: int = 4) -> np.ndarray:
-    """
-    Convert a single sequence to a 4^k dimension k-mer frequency vector.
-    """
+    """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)
@@ -96,223 +88,403 @@ def sequence_to_kmer_vector(sequence: str, k: int = 4) -> np.ndarray:
 
     total_kmers = len(sequence) - k + 1
     if total_kmers > 0:
-        vec /= total_kmers
+        vec = vec / total_kmers  # normalize frequencies
 
     return vec
 
 ###############################################################################
-# Visualization
+# Additional Plots
 ###############################################################################
-def create_waterfall_plot(shap_values, base_value, data, max_display=15):
+def create_probability_bar_plot(prob_human, prob_nonhuman):
     """
-    Create a SHAP waterfall plot for a single sample's single class
-    (shap_values: shape=(num_features,))
-    (base_value: scalar)
-    (data: original input data for that sample, shape=(num_features,))
+    Simple bar plot comparing human vs. non-human probabilities.
     """
-    # Build a shap Explanation object
-    expl = shap.Explanation(
-        values=shap_values,
-        base_values=base_value,
-        data=data,
-        feature_names=[f"feat_{i}" for i in range(len(shap_values))]
-    )
-
-    fig = plt.figure(figsize=(6, 4), dpi=75)
-    shap.plots.waterfall(expl, max_display=max_display, show=False)
-    buf = io.BytesIO()
-    plt.savefig(buf, format='png', bbox_inches='tight')
-    buf.seek(0)
-    im = Image.open(buf)
-    plt.close(fig)
-    return im
+    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_freq_sigma_plot(shap_values, raw_freq, scaled_vec, kmer_list, title="Top-10 k-mers"):
+def create_frequency_sigma_plot(important_kmers, title):
     """
-    Show top-10 k-mers by absolute shap value with frequency (%) & z-score.
-    shap_values: (256,)
-    raw_freq: (256,) unscaled frequency
-    scaled_vec: (256,) scaled frequency (z-scores)
-    kmer_list: list of length=256
+    Creates a bar plot of the top k-mers (by importance) showing 
+    frequency (%) and σ from mean.
     """
-    abs_vals = np.abs(shap_values)
-    top_indices = np.argsort(abs_vals)[-10:][::-1]  # top 10
-    top_data = []
+    # 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]
     
-    for idx in top_indices:
-        idx = int(idx)  # ensure it's an integer
-        top_data.append({
-            "kmer": kmer_list[idx],
-            "shap": shap_values[idx],
-            "abs_shap": abs_vals[idx],
-            "freq": raw_freq[idx] * 100.0,
-            "sigma": scaled_vec[idx]
-        })
-
-    # Sort again by abs_shap desc
-    top_data.sort(key=lambda x: x["abs_shap"], reverse=True)
-
-    kmers = [d["kmer"] for d in top_data]
-    freqs = [d["freq"] for d in top_data]
-    sigmas = [d["sigma"] for d in top_data]
-    colors = ["green" if d["shap"] >= 0 else "red" for d in top_data]
-
     x = np.arange(len(kmers))
     width = 0.4
 
-    fig, ax = plt.subplots(figsize=(6, 4), dpi=75)
+    fig, ax_bar = plt.subplots(figsize=(10, 5))
 
-    ax.bar(x - width/2, freqs, width, color=colors, alpha=0.7, label="Frequency (%)")
-    ax.set_ylabel("Frequency (%)")
-    if freqs:
-        ax.set_ylim(0, max(freqs)*1.25)
+    # 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")
 
-    ax2 = ax.twinx()
-    ax2.bar(x + width/2, sigmas, width, color="gray", alpha=0.5, label="Z-score")
-    ax2.set_ylabel("Standard Deviations (σ)")
+    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(title)
-
-    lines1, labels1 = ax.get_legend_handles_labels()
-    lines2, labels2 = ax2.get_legend_handles_labels()
-    ax.legend(lines1 + lines2, labels1 + labels2, loc='upper 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()
-    buf = io.BytesIO()
-    fig.savefig(buf, format='png', bbox_inches='tight')
-    buf.seek(0)
-    im = Image.open(buf)
-    plt.close(fig)
-    return im
+    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
 
 ###############################################################################
-# Main Gradio Logic
+# Prediction Function
 ###############################################################################
-def classify_and_explain(file_obj):
+def predict(file_obj):
     """
-    - Reads the first sequence from the uploaded FASTA
-    - Loads model & scaler
-    - Produces a single prediction + shap explanation
-    - Returns Markdown + Waterfall image + freq/sigma image
+    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
     """
-    # 1. Read text from user input
+    # 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:
-        return "Error reading file!", None, None
+            text = file_obj.decode('utf-8')
+    except Exception as e:
+        return (
+            f"Error reading file: {str(e)}",
+            None,
+            None,
+            None,
+            None
+        )
 
-    # 2. Parse
+    # 1. Parse FASTA
     sequences = parse_fasta(text)
-    if not sequences:
-        return "No valid FASTA sequences found!", None, None
-    header, seq = sequences[0]  # only the first sequence
+    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
 
-    # 3. Convert to k-mer
+    # 2. Prepare model, scaler, and input
     k = 4
-    raw_freq_vec = sequence_to_kmer_vector(seq, k=k)  # shape=(256,)
-
-    # 4. Load model & scaler
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = VirusClassifier(4**k).to(device)
     try:
-        state_dict = torch.load("model.pt", map_location=device, weights_only=True)
+        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)
-        model.eval()
         scaler = joblib.load("scaler.pkl")
-    except Exception as e:
-        return f"Error loading model/scaler: {str(e)}", None, None
-
-    # 5. Scale data
-    scaled_data = scaler.transform(raw_freq_vec.reshape(1, -1))  # shape=(1, 256)
-    X_tensor = torch.FloatTensor(scaled_data).to(device)
-
-    # 6. Predict
-    with torch.no_grad():
-        out = model(X_tensor)  # shape=(1,2)
-        probs = torch.softmax(out, dim=1).cpu().numpy()[0]  # shape=(2,)
-    pred_class = np.argmax(probs)
-    pred_label = "human" if pred_class == 1 else "non-human"
-    confidence = float(np.max(probs))
-    human_prob = float(probs[1])
-    nonhuman_prob = float(probs[0])
-
-    # 7. SHAP Explanation (single sample)
-    # We'll wrap the model for shap
-    wrapped_model = TorchModelWrapper(model, device)
-    background_data = scaled_data  # 1 sample as background (a bit silly, but simpler)
-    explainer = shap.Explainer(wrapped_model, background_data)
-    shap_values = explainer(scaled_data)  # shape => (1, 2, 256) for 2-class output
-
-    # We'll pick class=1's shap values
-    sample_shap_vals = shap_values.values[0, 1, :]  # shape=(256,)
-    base_value       = shap_values.base_values[0, 1]
-    sample_data      = shap_values.data[0]          # shape=(256,)
-
-    # 8. Create the two plots
-    wf_img = create_waterfall_plot(
-        shap_values=sample_shap_vals, 
-        base_value=base_value, 
-        data=sample_data, 
-        max_display=15
-    )
-
-    # freq-sigma plot
-    kmers = [''.join(p) for p in product("ACGT", repeat=k)]
-    freq_img = create_freq_sigma_plot(
-        shap_values=sample_shap_vals,
-        raw_freq=raw_freq_vec,
-        scaled_vec=scaled_data[0],
-        kmer_list=kmers,
-        title=f"{header[:25]}... (Top-10 K-mers)"
-    )
-
-    # 9. Markdown result
-    result_md = f"""
-# Classification Result
-
-**Header**: {header}
-
-**Predicted Label**: {pred_label}  
-**Confidence**: {confidence:.4f}  
-
-**Human Probability**: {human_prob:.4f}  
-**Non-human Probability**: {nonhuman_prob:.4f}
+        model.eval()
 
-Above are the SHAP-based analyses (class=1).
-    """
+        scaled_vector = scaler.transform(raw_freq_vector.reshape(1, -1))
+        X_tensor = torch.FloatTensor(scaled_vector).to(device)
 
-    return result_md, wf_img, freq_img
+        # 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="Single-Sequence Virus Host Classifier") as demo:
-    gr.Markdown("## Upload a FASTA file containing **one** (or more) sequences. We only use the **first**.")
-
-    file_input = gr.File(label="Upload FASTA", type="binary")
-    run_btn = gr.Button("Classify & Explain")
+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("Results"):
+        with gr.Tab("Prediction Results"):
             md_out = gr.Markdown()
-        with gr.Tab("SHAP Waterfall"):
-            wf_out = gr.Image(label="Waterfall Plot (class=1)")
-        with gr.Tab("Top-10 K-mer Plot"):
-            freq_out = gr.Image(label="Frequency & Sigma (class=1)")
-
-    run_btn.click(
-        fn=classify_and_explain,
-        inputs=[file_input],
-        outputs=[md_out, wf_out, freq_out]
+        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]
     )
 
-# No share=True -> avoid HF Spaces warning
 if __name__ == "__main__":
-    demo.launch()
+    # 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)