gene_knockout_analysis_attention_layer.py

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import torch
from scipy.spatial.distance import cosine
from scipy.stats import spearmanr
from tqdm import tqdm

class AttentionLayerKnockoutAnalyzer:
    """
    Analyzes the impact of gene knockouts on attention layer outputs.
    
    This measures how removing genes affects the intermediate representations
    learned by the attention mechanism, complementing the classification-based
    knockout analysis.
    """
    
    def __init__(self, model, device='cpu'):
        """
        Initialize the attention layer knockout analyzer.
        
        Args:
            model: Trained PyTorchTransformerModel
            device: Device to run computations on
        """
        self.model = model
        self.device = device
        self.model.eval()
    
    def extract_attention_outputs(self, sequences):
        """
        Extract attention layer outputs for sequences.
        
        Args:
            sequences: Input sequences tensor (num_genes, timesteps, features)
            
        Returns:
            Attention layer outputs (num_genes, timesteps, embed_dim)
        """
        with torch.no_grad():
            if not isinstance(sequences, torch.Tensor):
                sequences = torch.tensor(sequences, dtype=torch.float32)
            sequences = sequences.to(self.device)
            
            # Extract attention outputs
            attention_outputs, _ = self.model.multi_head_attention(
                sequences, 
                return_attention_weights=True
            )
            
        return attention_outputs
    
    def get_baseline_attention_outputs(self, sequences):
        """
        Get baseline attention outputs without any knockouts.
        
        Args:
            sequences: Input sequences tensor
            
        Returns:
            Dictionary with baseline attention outputs and derived metrics
        """
        attention_outputs = self.extract_attention_outputs(sequences)
        
        # Convert to numpy for easier manipulation
        if isinstance(attention_outputs, torch.Tensor):
            outputs_np = attention_outputs.cpu().numpy()
        else:
            outputs_np = attention_outputs
        
        # Compute baseline statistics
        baseline = {
            'outputs': outputs_np,
            'mean_activation': np.mean(np.abs(outputs_np), axis=(1, 2)),  # Per gene
            'temporal_profile': np.mean(np.abs(outputs_np), axis=2),  # (genes, time)
            'feature_activation': np.mean(np.abs(outputs_np), axis=1),  # (genes, features)
            'global_mean': np.mean(np.abs(outputs_np)),
            'global_std': np.std(outputs_np)
        }
        
        return baseline
    
    def knockout_single_gene_attention(self, sequences, gene_idx, knockout_value=0.0):
        """
        Create knockout version and extract attention outputs.
        
        Args:
            sequences: Input sequences
            gene_idx: Index of gene to knockout
            knockout_value: Value to set the gene to
            
        Returns:
            Attention outputs with gene knocked out
        """
        # Create knockout sequences
        if isinstance(sequences, torch.Tensor):
            sequences_ko = sequences.clone()
        else:
            sequences_ko = sequences.copy()
        
        sequences_ko[gene_idx, :, :] = knockout_value
        
        # Extract attention outputs
        attention_outputs_ko = self.extract_attention_outputs(sequences_ko)
        
        return attention_outputs_ko
    
    def compute_representation_distance(self, baseline_outputs, knockout_outputs):
        """
        Compute distance between baseline and knockout representations.
        
        Args:
            baseline_outputs: Baseline attention outputs
            knockout_outputs: Knockout attention outputs
            
        Returns:
            Dictionary with distance metrics
        """
        if isinstance(baseline_outputs, torch.Tensor):
            baseline_np = baseline_outputs.cpu().numpy()
        else:
            baseline_np = baseline_outputs
            
        if isinstance(knockout_outputs, torch.Tensor):
            knockout_np = knockout_outputs.cpu().numpy()
        else:
            knockout_np = knockout_outputs
        
        # Flatten for distance computation
        baseline_flat = baseline_np.flatten()
        knockout_flat = knockout_np.flatten()
        
        # Compute various distance metrics
        l2_distance = np.linalg.norm(baseline_flat - knockout_flat)
        l1_distance = np.sum(np.abs(baseline_flat - knockout_flat))
        
        # Normalize by number of elements
        normalized_l2 = l2_distance / np.sqrt(len(baseline_flat))
        normalized_l1 = l1_distance / len(baseline_flat)
        
        # Cosine similarity
        cos_sim = 1 - cosine(baseline_flat, knockout_flat)
        
        # Per-gene activation changes
        baseline_gene_activation = np.mean(np.abs(baseline_np), axis=(1, 2))
        knockout_gene_activation = np.mean(np.abs(knockout_np), axis=(1, 2))
        activation_changes = np.abs(baseline_gene_activation - knockout_gene_activation)
        
        distances = {
            'l2_distance': l2_distance,
            'l1_distance': l1_distance,
            'normalized_l2': normalized_l2,
            'normalized_l1': normalized_l1,
            'cosine_similarity': cos_sim,
            'mean_activation_change': np.mean(activation_changes),
            'max_activation_change': np.max(activation_changes),
            'activation_changes_per_gene': activation_changes,
            'num_genes_affected': np.sum(activation_changes > 1e-6)  # Genes with non-zero change
        }
        
        return distances
    
    def analyze_single_gene_knockout_attention(self, sequences, gene_idx, 
                                               baseline_outputs):
        """
        Analyze impact of single gene knockout on attention outputs.
        
        Args:
            sequences: Input sequences
            gene_idx: Index of gene to knockout
            baseline_outputs: Baseline attention outputs and metrics
            
        Returns:
            Dictionary with knockout impact on attention layer
        """
        # Get knockout attention outputs
        knockout_outputs = self.knockout_single_gene_attention(sequences, gene_idx)
        
        # Compute distances
        distances = self.compute_representation_distance(
            baseline_outputs['outputs'], 
            knockout_outputs
        )
        
        # Compute change in gene's own activation
        if isinstance(knockout_outputs, torch.Tensor):
            knockout_np = knockout_outputs.cpu().numpy()
        else:
            knockout_np = knockout_outputs
        
        baseline_np = baseline_outputs['outputs']
        
        # Self-impact: change in the knocked-out gene's representation
        self_activation_baseline = np.mean(np.abs(baseline_np[gene_idx]))
        self_activation_knockout = np.mean(np.abs(knockout_np[gene_idx]))
        self_impact = np.abs(self_activation_baseline - self_activation_knockout)
        
        # Cross-impact: changes in other genes' representations
        other_genes_mask = np.ones(baseline_np.shape[0], dtype=bool)
        other_genes_mask[gene_idx] = False
        
        baseline_other = baseline_np[other_genes_mask]
        knockout_other = knockout_np[other_genes_mask]
        cross_impact = np.mean(np.abs(baseline_other - knockout_other))
        
        impact = {
            'gene_idx': gene_idx,
            'normalized_l2_distance': distances['normalized_l2'],
            'normalized_l1_distance': distances['normalized_l1'],
            'cosine_similarity': distances['cosine_similarity'],
            'mean_activation_change': distances['mean_activation_change'],
            'max_activation_change': distances['max_activation_change'],
            'num_genes_affected': distances['num_genes_affected'],
            'self_impact': self_impact,
            'cross_impact': cross_impact,
            'self_cross_ratio': self_impact / (cross_impact + 1e-10)
        }
        
        return impact
    
    def analyze_all_genes_knockout_attention(self, sequences, gene_names, verbose=True):
        """
        Analyze attention layer knockout impact for all genes.
        
        Args:
            sequences: Input sequences
            gene_names: List of gene names
            verbose: Whether to show progress
            
        Returns:
            DataFrame with attention knockout results for all genes
        """
        print("Computing baseline attention outputs...")
        baseline = self.get_baseline_attention_outputs(sequences)
        
        print("Baseline attention statistics:")
        print(f"  Mean activation: {baseline['global_mean']:.4f}")
        print(f"  Std activation: {baseline['global_std']:.4f}")
        
        print(f"\nAnalyzing attention knockout impact for {len(gene_names)} genes...")
        
        results = []
        iterator = tqdm(range(len(gene_names))) if verbose else range(len(gene_names))
        
        for gene_idx in iterator:
            impact = self.analyze_single_gene_knockout_attention(
                sequences, gene_idx, baseline
            )
            impact['gene_name'] = gene_names[gene_idx]
            impact['baseline_activation'] = baseline['mean_activation'][gene_idx]
            results.append(impact)
        
        # Create dataframe and sort by impact
        results_df = pd.DataFrame(results)
        results_df = results_df.sort_values('normalized_l2_distance', ascending=False)
        results_df['rank'] = range(1, len(results_df) + 1)
        
        return results_df, baseline
    
    def identify_affected_genes(self, sequences, gene_idx, baseline_outputs, 
                               threshold=0.01):
        """
        Identify which other genes' representations are affected by knocking out a gene.
        
        Args:
            sequences: Input sequences
            gene_idx: Index of gene to knockout
            baseline_outputs: Baseline attention outputs
            threshold: Threshold for significant change
            
        Returns:
            DataFrame with affected genes and their change magnitudes
        """
        # Get knockout outputs
        knockout_outputs = self.knockout_single_gene_attention(sequences, gene_idx)
        
        if isinstance(knockout_outputs, torch.Tensor):
            knockout_np = knockout_outputs.cpu().numpy()
        else:
            knockout_np = knockout_outputs
        
        baseline_np = baseline_outputs['outputs']
        
        # Compute per-gene activation changes
        baseline_activations = np.mean(np.abs(baseline_np), axis=(1, 2))
        knockout_activations = np.mean(np.abs(knockout_np), axis=(1, 2))
        changes = np.abs(baseline_activations - knockout_activations)
        
        # Find genes above threshold
        affected_indices = np.where(changes > threshold)[0]
        
        affected_data = []
        for idx in affected_indices:
            affected_data.append({
                'affected_gene_idx': idx,
                'activation_change': changes[idx],
                'baseline_activation': baseline_activations[idx],
                'knockout_activation': knockout_activations[idx],
                'percent_change': (changes[idx] / (baseline_activations[idx] + 1e-10)) * 100
            })
        
        affected_df = pd.DataFrame(affected_data)
        if len(affected_df) > 0:
            affected_df = affected_df.sort_values('activation_change', ascending=False)
        
        return affected_df
    
    def visualize_attention_knockout_results(self, results_df, save_path=None,
                                            patient_id=None, top_n=20):
        """
        Visualize attention layer knockout analysis results.
        
        Args:
            results_df: DataFrame with attention knockout results
            save_path: Path to save figure
            patient_id: Patient identifier
            top_n: Number of top genes to show
        """
        fig, axes = plt.subplots(2, 2, figsize=(16, 12))
        
        # 1. Top genes by representation distance
        top_genes = results_df.head(top_n)
        axes[0, 0].barh(range(len(top_genes)), top_genes['normalized_l2_distance'],
                       color='darkviolet', alpha=0.7)
        axes[0, 0].set_yticks(range(len(top_genes)))
        axes[0, 0].set_yticklabels(top_genes['gene_name'], fontsize=9)
        axes[0, 0].invert_yaxis()
        axes[0, 0].set_xlabel('Normalized L2 Distance', fontsize=11)
        axes[0, 0].set_title(f'Top {top_n} Genes by Attention Representation Impact', 
                            fontsize=12)
        axes[0, 0].grid(axis='x', alpha=0.3)
        
        # 2. Self-impact vs Cross-impact
        axes[0, 1].scatter(results_df['self_impact'], results_df['cross_impact'],
                          alpha=0.5, s=30, c='teal')
        axes[0, 1].plot([0, results_df['self_impact'].max()],
                       [0, results_df['self_impact'].max()],
                       'r--', alpha=0.5, linewidth=2, label='Equal Impact Line')
        axes[0, 1].set_xlabel('Self-Impact (Own Representation Change)', fontsize=11)
        axes[0, 1].set_ylabel('Cross-Impact (Other Genes Change)', fontsize=11)
        axes[0, 1].set_title('Self-Impact vs Cross-Impact', fontsize=12)
        axes[0, 1].legend()
        axes[0, 1].grid(alpha=0.3)
        
        # Label high cross-impact genes
        high_cross = results_df.nlargest(5, 'cross_impact')
        for _, row in high_cross.iterrows():
            axes[0, 1].annotate(row['gene_name'],
                              (row['self_impact'], row['cross_impact']),
                              fontsize=8, alpha=0.7)
        
        # 3. Number of genes affected
        axes[1, 0].hist(results_df['num_genes_affected'], bins=30,
                       color='coral', edgecolor='black', alpha=0.7)
        axes[1, 0].axvline(results_df['num_genes_affected'].mean(),
                          color='red', linestyle='--', linewidth=2,
                          label=f'Mean: {results_df["num_genes_affected"].mean():.1f}')
        axes[1, 0].set_xlabel('Number of Other Genes Affected', fontsize=11)
        axes[1, 0].set_ylabel('Frequency', fontsize=11)
        axes[1, 0].set_title('Distribution of Gene Influence', fontsize=12)
        axes[1, 0].legend()
        axes[1, 0].grid(axis='y', alpha=0.3)
        
        # 4. Cosine similarity distribution
        axes[1, 1].hist(results_df['cosine_similarity'], bins=30,
                       color='mediumseagreen', edgecolor='black', alpha=0.7)
        axes[1, 1].axvline(results_df['cosine_similarity'].mean(),
                          color='red', linestyle='--', linewidth=2,
                          label=f'Mean: {results_df["cosine_similarity"].mean():.3f}')
        axes[1, 1].set_xlabel('Cosine Similarity (Baseline vs Knockout)', fontsize=11)
        axes[1, 1].set_ylabel('Frequency', fontsize=11)
        axes[1, 1].set_title('Representation Similarity After Knockout', fontsize=12)
        axes[1, 1].legend()
        axes[1, 1].grid(axis='y', alpha=0.3)
        
        if patient_id:
            fig.suptitle(f'Attention Layer Knockout Analysis - Patient {patient_id}',
                        fontsize=14, y=0.995)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Saved attention knockout visualization to {save_path}")
        
        plt.close()
    
    def compare_attention_vs_classification_knockout(self, attention_knockout_df,
                                                     classification_knockout_df,
                                                     save_path=None, patient_id=None):
        """
        Compare attention layer knockout impact with classification knockout impact.
        
        Args:
            attention_knockout_df: Attention layer knockout results
            classification_knockout_df: Classification knockout results
            save_path: Path to save figure
            patient_id: Patient identifier
        """
        # Merge dataframes
        merged = pd.merge(
            attention_knockout_df[['gene_name', 'normalized_l2_distance', 'rank']],
            classification_knockout_df[['gene_name', 'accuracy_drop', 'rank']],
            on='gene_name',
            how='inner',
            suffixes=('_attention', '_classification')
        )
        
        # Calculate correlations
        spearman_corr, spearman_p = spearmanr(
            merged['rank_attention'], 
            merged['rank_classification']
        )
        
        from scipy.stats import pearsonr
        pearson_corr, pearson_p = pearsonr(
            merged['normalized_l2_distance'],
            merged['accuracy_drop']
        )
        
        fig, axes = plt.subplots(1, 2, figsize=(14, 5))
        
        # 1. Scatter: Attention impact vs Classification impact
        axes[0].scatter(merged['normalized_l2_distance'], merged['accuracy_drop'],
                       alpha=0.5, s=40, c='darkorange')
        axes[0].set_xlabel('Attention Layer Impact (L2 Distance)', fontsize=11)
        axes[0].set_ylabel('Classification Impact (Accuracy Drop)', fontsize=11)
        axes[0].set_title('Representation Impact vs Classification Impact\n' +
                         f'Pearson r = {pearson_corr:.3f} (p = {pearson_p:.3e})',
                         fontsize=12)
        axes[0].grid(alpha=0.3)
        
        # Trendline
        z = np.polyfit(merged['normalized_l2_distance'], merged['accuracy_drop'], 1)
        p = np.poly1d(z)
        x_line = np.linspace(merged['normalized_l2_distance'].min(),
                            merged['normalized_l2_distance'].max(), 100)
        axes[0].plot(x_line, p(x_line), "r--", alpha=0.8, linewidth=2)
        
        # Label extreme points
        # High attention impact, low classification impact
        high_att_low_class = merged.nlargest(3, 'normalized_l2_distance').nsmallest(3, 'accuracy_drop')
        for _, row in high_att_low_class.iterrows():
            axes[0].annotate(row['gene_name'],
                           (row['normalized_l2_distance'], row['accuracy_drop']),
                           fontsize=8, alpha=0.7, color='blue')
        
        # High classification impact, low attention impact
        high_class_low_att = merged.nlargest(3, 'accuracy_drop').nsmallest(3, 'normalized_l2_distance')
        for _, row in high_class_low_att.iterrows():
            axes[0].annotate(row['gene_name'],
                           (row['normalized_l2_distance'], row['accuracy_drop']),
                           fontsize=8, alpha=0.7, color='red')
        
        # 2. Rank comparison
        axes[1].scatter(merged['rank_attention'], merged['rank_classification'],
                       alpha=0.5, s=40, c='mediumvioletred')
        axes[1].plot([merged['rank_attention'].min(), merged['rank_attention'].max()],
                    [merged['rank_attention'].min(), merged['rank_attention'].max()],
                    'k--', alpha=0.5, linewidth=2, label='Perfect Agreement')
        axes[1].set_xlabel('Attention Layer Knockout Rank', fontsize=11)
        axes[1].set_ylabel('Classification Knockout Rank', fontsize=11)
        axes[1].set_title('Ranking Comparison\n' +
                         f'Spearman ρ = {spearman_corr:.3f} (p = {spearman_p:.3e})',
                         fontsize=12)
        axes[1].legend()
        axes[1].grid(alpha=0.3)
        axes[1].invert_xaxis()
        axes[1].invert_yaxis()
        
        if patient_id:
            fig.suptitle(f'Attention vs Classification Knockout Comparison - Patient {patient_id}',
                        fontsize=14, y=1.00)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Saved comparison visualization to {save_path}")
        
        plt.close()
        
        return {
            'spearman_correlation': spearman_corr,
            'spearman_pvalue': spearman_p,
            'pearson_correlation': pearson_corr,
            'pearson_pvalue': pearson_p,
            'merged_data': merged
        }
    
    def analyze_gene_network_effects(self, sequences, gene_names, top_genes_df,
                                     baseline_outputs, top_n=10, save_path=None,
                                     patient_id=None):
        """
        Analyze network effects: which genes affect which other genes.
        
        Args:
            sequences: Input sequences
            gene_names: List of gene names
            top_genes_df: DataFrame with top genes (from attention knockout)
            baseline_outputs: Baseline attention outputs
            top_n: Number of top genes to analyze
            save_path: Path to save results
            patient_id: Patient identifier
        """
        print(f"\nAnalyzing network effects for top {top_n} genes...")
        
        # Create interaction matrix
        top_genes = top_genes_df.head(top_n)
        interaction_matrix = np.zeros((top_n, len(gene_names)))
        
        for i, (idx, row) in enumerate(top_genes.iterrows()):
            gene_idx = gene_names.index(row['gene_name'])
            
            # Get affected genes
            affected_df = self.identify_affected_genes(
                sequences, gene_idx, baseline_outputs, threshold=0.001
            )
            
            # Fill interaction matrix
            for _, affected_row in affected_df.iterrows():
                affected_idx = int(affected_row['affected_gene_idx'])
                interaction_matrix[i, affected_idx] = affected_row['activation_change']
        
        # Visualize interaction matrix (focus on top genes affecting other top genes)
        if save_path:
            fig, ax = plt.subplots(figsize=(12, 10))
            
            # Extract submatrix for top genes affecting top genes
            top_gene_indices = [gene_names.index(name) for name in top_genes['gene_name']]
            submatrix = interaction_matrix[:, top_gene_indices]
            
            sns.heatmap(submatrix, cmap='YlOrRd', annot=True, fmt='.3f',
                       xticklabels=top_genes['gene_name'],
                       yticklabels=top_genes['gene_name'],
                       cbar_kws={'label': 'Activation Change'},
                       ax=ax)
            
            ax.set_xlabel('Affected Gene', fontsize=11)
            ax.set_ylabel('Knocked Out Gene', fontsize=11)
            
            if patient_id:
                ax.set_title(f'Gene Interaction Network (Top {top_n}) - Patient {patient_id}',
                           fontsize=12)
            else:
                ax.set_title(f'Gene Interaction Network (Top {top_n})', fontsize=12)
            
            plt.tight_layout()
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Saved gene network visualization to {save_path}")
            plt.close()
        
        return interaction_matrix
    
    def generate_attention_knockout_report(self, sequences, gene_names,
                                          output_dir, patient_id,
                                          classification_knockout_df=None):
        """
        Generate comprehensive attention layer knockout analysis report.
        
        Args:
            sequences: Input sequences
            gene_names: List of gene names
            output_dir: Directory to save outputs
            patient_id: Patient identifier
            classification_knockout_df: Optional classification knockout results for comparison
        """
        print(f"\nGenerating attention layer knockout analysis for patient {patient_id}...")
        
        # Create subdirectory
        attn_knockout_dir = os.path.join(output_dir, 
                                         f"{patient_id}_attention_knockout_analysis")
        os.makedirs(attn_knockout_dir, exist_ok=True)
        
        # 1. Analyze all genes
        print("Step 1/4: Analyzing attention knockout for all genes...")
        results_df, baseline = self.analyze_all_genes_knockout_attention(
            sequences, gene_names, verbose=True
        )
        
        results_df.to_csv(
            os.path.join(attn_knockout_dir, f"{patient_id}_attention_knockout_results.csv"),
            index=False
        )
        
        # 2. Visualize results
        print("Step 2/4: Creating visualizations...")
        self.visualize_attention_knockout_results(
            results_df,
            save_path=os.path.join(attn_knockout_dir, 
                                  f"{patient_id}_attention_knockout_impact.png"),
            patient_id=patient_id,
            top_n=20
        )
        
        # 3. Analyze gene network effects
        print("Step 3/4: Analyzing gene interaction network...")
        interaction_matrix = self.analyze_gene_network_effects(
            sequences, gene_names, results_df, baseline, top_n=15,
            save_path=os.path.join(attn_knockout_dir,
                                  f"{patient_id}_gene_interaction_network.png"),
            patient_id=patient_id
        )
        
        # 4. Compare with classification knockout if available
        comparison_stats = None
        if classification_knockout_df is not None:
            print("Step 4/4: Comparing with classification knockout results...")
            comparison_stats = self.compare_attention_vs_classification_knockout(
                results_df, classification_knockout_df,
                save_path=os.path.join(attn_knockout_dir,
                                      f"{patient_id}_attention_vs_classification_knockout.png"),
                patient_id=patient_id
            )
            
            # Save comparison data
            comparison_stats['merged_data'].to_csv(
                os.path.join(attn_knockout_dir,
                           f"{patient_id}_knockout_comparison.csv"),
                index=False
            )
        else:
            print("Step 4/4: Skipped (no classification knockout data provided)")
        
        # 5. Generate summary report
        self._generate_attention_knockout_summary(
            results_df, baseline, comparison_stats,
            attn_knockout_dir, patient_id
        )
        
        print(f"✓ Attention knockout analysis complete! Results saved to {attn_knockout_dir}")
        
        return results_df, comparison_stats
    
    def _generate_attention_knockout_summary(self, results_df, baseline,
                                            comparison_stats, output_dir, patient_id):
        """Generate text summary of attention knockout analysis."""
        summary_path = os.path.join(output_dir, 
                                   f"{patient_id}_attention_knockout_summary.txt")
        
        with open(summary_path, 'w') as f:
            f.write("="*80 + "\n")
            f.write(f"ATTENTION LAYER KNOCKOUT ANALYSIS - Patient {patient_id}\n")
            f.write("="*80 + "\n\n")
            
            f.write("BASELINE ATTENTION STATISTICS\n")
            f.write("-"*80 + "\n")
            f.write(f"Mean activation: {baseline['global_mean']:.4f}\n")
            f.write(f"Std activation:  {baseline['global_std']:.4f}\n\n")
            
            f.write("TOP 10 GENES BY REPRESENTATION IMPACT\n")
            f.write("-"*80 + "\n")
            for idx, row in results_df.head(10).iterrows():
                f.write(f"{row['rank']:2d}. {row['gene_name']:20s} | "
                       f"L2 Dist: {row['normalized_l2_distance']:.4f} | "
                       f"Affected: {row['num_genes_affected']:.0f} genes | "
                       f"Cross-Impact: {row['cross_impact']:.4f}\n")
            
            f.write("\n\nREPRESENTATION IMPACT STATISTICS\n")
            f.write("-"*80 + "\n")
            f.write(f"Mean L2 distance:          {results_df['normalized_l2_distance'].mean():.4f}\n")
            f.write(f"Std L2 distance:           {results_df['normalized_l2_distance'].std():.4f}\n")
            f.write(f"Max L2 distance:           {results_df['normalized_l2_distance'].max():.4f}\n")
            f.write(f"Mean cosine similarity:    {results_df['cosine_similarity'].mean():.4f}\n")
            f.write(f"Mean genes affected:       {results_df['num_genes_affected'].mean():.1f}\n")
            f.write(f"Max genes affected:        {results_df['num_genes_affected'].max():.0f}\n")
            
            if comparison_stats is not None:
                f.write("\n\nCOMPARISON WITH CLASSIFICATION KNOCKOUT\n")
                f.write("-"*80 + "\n")
                f.write(f"Spearman correlation: {comparison_stats['spearman_correlation']:.4f} "
                       f"(p = {comparison_stats['spearman_pvalue']:.3e})\n")
                f.write(f"Pearson correlation:  {comparison_stats['pearson_correlation']:.4f} "
                       f"(p = {comparison_stats['pearson_pvalue']:.3e})\n\n")
                
                f.write("INTERPRETATION:\n")
                if comparison_stats['spearman_correlation'] > 0.7:
                    f.write("  ✓ Strong agreement: Representation changes correlate with "
                           "classification impact\n")
                elif comparison_stats['spearman_correlation'] > 0.4:
                    f.write("  ~ Moderate agreement: Some representation changes don't affect "
                           "classification\n")
                else:
                    f.write("  ! Weak agreement: Representation and classification impacts diverge\n")
                f.write("\n" + "="*80 + "\n")

        print(f"Saved attention knockout summary to {summary_path}")

    def analyze_attention_layer_knockouts(model, sequences, gene_names, output_dir, 
                                          patient_id, classification_knockout_df=None, 
                                          device='cpu'):
        """
        Args:
        model: Trained PyTorchTransformerModel
        sequences: Input sequences
        gene_names: List of gene names
        output_dir: Output directory
        patient_id: Patient identifier
        classification_knockout_df: Optional classification knockout results
        device: Device to run on
        
        Returns:
            Analyzer instance and results
        """
        analyzer = AttentionLayerKnockoutAnalyzer(model, device=device)
        results_df, comparison_stats = analyzer.generate_attention_knockout_report(
            sequences, gene_names, output_dir, 
            patient_id, classification_knockout_df
        )

        return analyzer, results_df, comparison_stats