Update app.py

app.py

@@ -555,22 +555,209 @@ def analyze_sequence_comparison(file1, file2, fasta1="", fasta2=""):
     except Exception as e:
         error_msg = f"Error during sequence comparison: {str(e)}"
         return error_msg, None, None
+
+###############################################################################
+# 11. GENE FEATURE ANALYSIS
+###############################################################################
+
+def parse_gene_features(text):
+    """Parse gene features from text file in FASTA-like format"""
+    genes = []
+    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:
+                genes.append({
+                    'header': current_header,
+                    'sequence': ''.join(current_sequence),
+                    'metadata': parse_gene_metadata(current_header)
+                })
+            current_header = line[1:]
+            current_sequence = []
+        else:
+            current_sequence.append(line.upper())
+            
+    if current_header:
+        genes.append({
+            'header': current_header,
+            'sequence': ''.join(current_sequence),
+            'metadata': parse_gene_metadata(current_header)
+        })
+    
+    return genes
+
+def parse_gene_metadata(header):
+    """Extract metadata from gene header"""
+    metadata = {}
+    parts = header.split()
+    
+    for part in parts:
+        if '[' in part and ']' in part:
+            key_value = part[1:-1].split('=', 1)
+            if len(key_value) == 2:
+                metadata[key_value[0]] = key_value[1]
+                
+    return metadata
+
+def analyze_gene_features(sequence_file, features_file, fasta_text="", features_text=""):
+    """Analyze SHAP values for each gene feature"""
+    # First analyze whole sequence
+    sequence_results = analyze_sequence(sequence_file, top_kmers=10, fasta_text=fasta_text)
+    if isinstance(sequence_results[0], str) and "Error" in sequence_results[0]:
+        return f"Error in sequence analysis: {sequence_results[0]}", None, None
+        
+    # Get SHAP values
+    shap_means = sequence_results[3]["shap_means"]
+    
+    # Parse gene features
+    if features_text.strip():
+        genes = parse_gene_features(features_text)
+    else:
+        try:
+            with open(features_file, 'r') as f:
+                genes = parse_gene_features(f.read())
+        except Exception as e:
+            return f"Error reading features file: {str(e)}", None, None
+            
+    # Analyze each gene
+    gene_results = []
+    for gene in genes:
+        try:
+            location = gene['metadata'].get('location', '')
+            if not location:
+                continue
+                
+            # Parse location (assuming format like "21729..22861")
+            start, end = map(int, location.split('..'))
+            
+            # Get SHAP values for this region
+            gene_shap = shap_means[start:end]
+            avg_shap = float(np.mean(gene_shap))
+            
+            gene_results.append({
+                'gene_name': gene['metadata'].get('gene', 'Unknown'),
+                'location': location,
+                'avg_shap': avg_shap,
+                'start': start,
+                'end': end,
+                'locus_tag': gene['metadata'].get('locus_tag', ''),
+                'classification': 'Human' if avg_shap > 0 else 'Non-human',
+                'confidence': abs(avg_shap)
+            })
+            
+        except Exception as e:
+            print(f"Error processing gene {gene['metadata'].get('gene', 'Unknown')}: {str(e)}")
+            continue
+            
+    # Create CSV output
+    csv_output = "gene_name,location,avg_shap,classification,confidence,locus_tag\n"
+    for result in gene_results:
+        csv_output += f"{result['gene_name']},{result['location']},{result['avg_shap']:.4f},"
+        csv_output += f"{result['classification']},{result['confidence']:.4f},{result['locus_tag']}\n"
+        
+    # Create genome diagram
+    diagram_img = create_genome_diagram(gene_results, len(shap_means))
+    
+    return gene_results, csv_output, diagram_img
+
+def create_genome_diagram(gene_results, genome_length):
+    """Create genome diagram using BioPython"""
+    from Bio.Graphics import GenomeDiagram
+    from Bio.SeqFeature import SeqFeature, FeatureLocation
+    from reportlab.lib import colors
+    from io import BytesIO
+    from PIL import Image
+    
+    # Create diagram
+    gd_diagram = GenomeDiagram.Diagram("Genome SHAP Analysis")
+    gd_track = gd_diagram.new_track(1, name="Genes")
+    gd_feature_set = gd_track.new_set()
+    
+    # Add features
+    for gene in gene_results:
+        # Create feature
+        feature = SeqFeature(
+            FeatureLocation(gene['start'], gene['end']),
+            type="gene"
+        )
+        
+        # Calculate color based on SHAP value
+        if gene['avg_shap'] > 0:
+            intensity = min(1.0, abs(gene['avg_shap']) * 2)
+            color = colors.Color(1-intensity, 1-intensity, 1)  # Red
+        else:
+            intensity = min(1.0, abs(gene['avg_shap']) * 2)
+            color = colors.Color(1-intensity, 1-intensity, 1)  # Blue
+            
+        # Add to diagram
+        gd_feature_set.add_feature(
+            feature,
+            color=color,
+            label=True,
+            name=f"{gene['gene_name']}\n(SHAP: {gene['avg_shap']:.3f})"
+        )
+        
+    # Draw diagram
+    gd_diagram.draw(
+        format="linear",
+        orientation="landscape",
+        pagesize=(15, 5),
+        start=0,
+        end=genome_length,
+        fragments=1
+    )
+    
+    # Save to BytesIO and convert to PIL Image
+    buffer = BytesIO()
+    gd_diagram.write(buffer, "PNG")
+    buffer.seek(0)
+    return Image.open(buffer)
+
+###############################################################################
+# 12. DOWNLOAD FUNCTIONS
+###############################################################################
+
+def prepare_csv_download(data, filename="analysis_results.csv"):
+    """Prepare CSV data for download"""
+    if isinstance(data, str):
+        return data.encode(), filename
+    elif isinstance(data, (list, dict)):
+        import csv
+        from io import StringIO
+        
+        output = StringIO()
+        writer = csv.DictWriter(output, fieldnames=data[0].keys())
+        writer.writeheader()
+        writer.writerows(data)
+        return output.getvalue().encode(), filename
+    else:
+        raise ValueError("Unsupported data type for CSV download")
         
 ###############################################################################
-# 10. BUILD GRADIO INTERFACE
+# 13. BUILD GRADIO INTERFACE
 ###############################################################################
 
 css = """
 .gradio-container {
     font-family: 'IBM Plex Sans', sans-serif;
 }
+.download-button {
+    margin-top: 10px;
+}
 """
 
 with gr.Blocks(css=css) as iface:
     gr.Markdown("""
     # Virus Host Classifier
     **Step 1**: Predict overall viral sequence origin (human vs non-human) and identify extreme regions.  
-    **Step 2**: Explore subregions to see local SHAP signals, distribution, GC content, etc.
+    **Step 2**: Explore subregions to see local SHAP signals, distribution, GC content, etc.  
+    **Step 3**: Analyze gene features and their contributions.  
+    **Step 4**: Compare sequences and analyze differences.
     
     **Color Scale**: Negative SHAP = Blue, Zero = White, Positive = Red.
     """)
@@ -587,19 +774,20 @@ with gr.Blocks(css=css) as iface:
                 results_box = gr.Textbox(label="Classification Results", lines=12, interactive=False)
                 kmer_img = gr.Image(label="Top k-mer SHAP")
                 genome_img = gr.Image(label="Genome-wide SHAP Heatmap (Blue=neg, White=0, Red=pos)")
+                download_results = gr.File(label="Download Results", visible=False, elem_classes="download-button")
         seq_state = gr.State()
         header_state = gr.State()
         analyze_btn.click(
             analyze_sequence,
             inputs=[file_input, top_k, text_input, win_size],
-            outputs=[results_box, kmer_img, genome_img, seq_state, header_state]
+            outputs=[results_box, kmer_img, genome_img, seq_state, header_state, download_results]
         )
 
     with gr.Tab("2) Subregion Exploration"):
         gr.Markdown("""
         **Subregion Analysis**  
         Select start/end positions to view local SHAP signals, distribution, GC content, etc.
-        The heatmap also uses the same Blue-White-Red scale.
+        The heatmap uses the same Blue-White-Red scale.
         """)
         with gr.Row():
             region_start = gr.Number(label="Region Start", value=0)
@@ -609,13 +797,47 @@ with gr.Blocks(css=css) as iface:
         with gr.Row():
             subregion_img = gr.Image(label="Subregion SHAP Heatmap (B-W-R)")
             subregion_hist_img = gr.Image(label="SHAP Distribution (Histogram)")
+        download_subregion = gr.File(label="Download Subregion Analysis", visible=False, elem_classes="download-button")
         region_btn.click(
             analyze_subregion,
             inputs=[seq_state, header_state, region_start, region_end],
-            outputs=[subregion_info, subregion_img, subregion_hist_img]
+            outputs=[subregion_info, subregion_img, subregion_hist_img, download_subregion]
+        )
+
+    with gr.Tab("3) Gene Features Analysis"):
+        gr.Markdown("""
+        **Analyze Gene Features**  
+        Upload a FASTA file and corresponding gene features file to analyze SHAP values per gene.
+        Gene features should be in the format:
+        ```
+        >gene_name [gene=X] [locus_tag=Y] [location=start..end]
+        SEQUENCE
+        ```
+        The genome viewer will show genes color-coded by their contribution:
+        - Red: Genes pushing toward human origin
+        - Blue: Genes pushing toward non-human origin
+        - Color intensity indicates strength of signal
+        """)
+        with gr.Row():
+            with gr.Column(scale=1):
+                gene_fasta_file = gr.File(label="Upload FASTA file", file_types=[".fasta", ".fa", ".txt"], type="filepath")
+                gene_fasta_text = gr.Textbox(label="Or paste FASTA sequence", placeholder=">sequence_name\nACGTACGT...", lines=5)
+            with gr.Column(scale=1):
+                features_file = gr.File(label="Upload gene features file", file_types=[".txt"], type="filepath")
+                features_text = gr.Textbox(label="Or paste gene features", placeholder=">gene_1 [gene=U12]...\nACGT...", lines=5)
+        
+        analyze_genes_btn = gr.Button("Analyze Gene Features", variant="primary")
+        gene_results = gr.Textbox(label="Gene Analysis Results", lines=12, interactive=False)
+        gene_diagram = gr.Image(label="Genome Diagram with Gene Features")
+        download_gene_results = gr.File(label="Download Gene Analysis", visible=False, elem_classes="download-button")
+        
+        analyze_genes_btn.click(
+            analyze_gene_features,
+            inputs=[gene_fasta_file, features_file, gene_fasta_text, features_text],
+            outputs=[gene_results, download_gene_results, gene_diagram]
         )
     
-    with gr.Tab("3) Comparative Analysis"):
+    with gr.Tab("4) Comparative Analysis"):
         gr.Markdown("""
         **Compare Two Sequences**  
         Upload or paste two FASTA sequences to compare their SHAP patterns.
@@ -638,29 +860,33 @@ with gr.Blocks(css=css) as iface:
         with gr.Row():
             diff_heatmap = gr.Image(label="SHAP Difference Heatmap")
             diff_hist = gr.Image(label="Distribution of SHAP Differences")
+        download_comparison = gr.File(label="Download Comparison Results", visible=False, elem_classes="download-button")
         compare_btn.click(
             analyze_sequence_comparison,
             inputs=[file_input1, file_input2, text_input1, text_input2],
-            outputs=[comparison_text, diff_heatmap, diff_hist]
+            outputs=[comparison_text, diff_heatmap, diff_hist, download_comparison]
         )
     
     gr.Markdown("""
     ### Interface Features
-    - **Overall Classification** (human vs non-human) using k-mer frequencies.
-    - **SHAP Analysis** to see which k-mers push classification toward or away from human.
+    - **Overall Classification** (human vs non-human) using k-mer frequencies
+    - **SHAP Analysis** shows which k-mers push classification toward or away from human
     - **White-Centered SHAP Gradient**: 
-      - Negative (blue), 0 (white), Positive (red), with symmetrical color range around 0. 
-    - **Identify Subregions** with the strongest push for human or non-human.
-    - **Subregion Exploration**: 
-      - Local SHAP heatmap & histogram 
-      - GC content 
-      - Fraction of positions pushing human vs. non-human 
-      - Simple logic-based classification
+      - Negative (blue), 0 (white), Positive (red)
+      - Symmetrical color range around 0
+    - **Identify Subregions** with strongest push for human or non-human
+    - **Gene Feature Analysis**:
+      - Analyze individual genes' contributions
+      - Interactive genome viewer
+      - Gene-level statistics and classification
     - **Sequence Comparison**:
       - Compare two sequences to identify regions of difference
-      - Normalized comparison to handle different sequence lengths
+      - Normalized comparison to handle different lengths
       - Statistical summary of differences
+    - **Data Export**:
+      - Download results as CSV files
+      - Save analysis outputs for further processing
     """)
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()