perform_shap_analysis.py

import numpy as np
import json
from typing import Dict, List, Tuple, Optional, Union
from pathlib import Path

class SHAPAnalysis:
    """Core SHAP analysis methods without plotting functionality."""
    
    def __init__(self, analyzer):
        self.analyzer = analyzer
        self.amino_acids = analyzer.amino_acids
        self.save_dir = analyzer.save_dir
        
    def _setup_directories(self, analysis_type: str = None):
        """Setup organized directory structure for SHAP results."""
        base_dir = self.save_dir.parent.parent  # Go up from explanations/{plastic_type}
        
        # Get analysis type from save_dir if not provided
        if analysis_type is None:
            analysis_type = self.save_dir.name  # e.g., 'promiscuous', 'pp_pet_selective'
        
        self.raw_data_dir = base_dir / "raw_data" / analysis_type
        self.reports_dir = base_dir / "reports" / analysis_type
        self.visualizations_dir = base_dir / "visualizations" / analysis_type
        
        # Create directories
        self.raw_data_dir.mkdir(exist_ok=True, parents=True)
        self.reports_dir.mkdir(exist_ok=True, parents=True)
        self.visualizations_dir.mkdir(exist_ok=True, parents=True)
        
        return {
            'raw_data': self.raw_data_dir,
            'reports': self.reports_dir, 
            'visualizations': self.visualizations_dir
        }
        
    def analyze_position_importance(self, shap_results: Dict) -> Dict:
        """Analyze position-level importance with statistical analysis."""
        print("Analyzing position-level importance...")
        
        if shap_results['analysis_type'] == 'promiscuous':
            analysis = self._analyze_promiscuous_positions(shap_results)
        elif shap_results['analysis_type'] == 'pp_pet_selective':
            analysis = self._analyze_pp_pet_positions(shap_results)
        else:
            analysis = self._analyze_single_plastic_positions(shap_results)
        
        return analysis
    
    def _analyze_promiscuous_positions(self, shap_results: Dict) -> Dict:
        """Analyze position importance for promiscuous designs."""
        sequences = shap_results['sequences']
        avg_shap = shap_results['average_shap_values']  # (n_sequences, 12, 18)
        individual_shap = shap_results['individual_shap_values']
        
        # Flip sign for intuitive interpretation (negative SHAP = better binding = positive importance)
        avg_shap_flipped = -avg_shap
        
        # Calculate position importance (sum of SHAP values across amino acids)
        avg_position_importance = np.sum(avg_shap_flipped, axis=2)  # (n_sequences, 12)
        
        # Calculate for individual plastics
        individual_position_importance = {}
        for plastic, data in individual_shap.items():
            plastic_shap = data['shap_values']
            # Flip sign for intuitive interpretation (negative SHAP = better binding = positive importance)
            plastic_shap_flipped = -plastic_shap
            individual_position_importance[plastic] = np.sum(plastic_shap_flipped, axis=2)
        
        # Calculate statistics for promiscuous model
        promiscuous_stats = {}
        for pos in range(12):
            prom_pos = avg_position_importance[:, pos]
            promiscuous_stats[f'position_{pos+1}'] = {
                'mean': np.mean(prom_pos),
                'std': np.std(prom_pos),
                'median': np.median(prom_pos),
                'min': np.min(prom_pos),
                'max': np.max(prom_pos)
            }
        
        # Calculate statistics for individual plastic models
        individual_stats = {}
        for plastic, importance_matrix in individual_position_importance.items():
            individual_stats[plastic] = {}
            for pos in range(12):
                plastic_pos = importance_matrix[:, pos]
                individual_stats[plastic][f'position_{pos+1}'] = {
                    'mean': np.mean(plastic_pos),
                    'std': np.std(plastic_pos),
                    'median': np.median(plastic_pos),
                    'min': np.min(plastic_pos),
                    'max': np.max(plastic_pos)
                }
        
        return {
            'analysis_type': 'promiscuous',
            'sequences': sequences,
            'promiscuous_position_importance': avg_position_importance,
            'individual_position_importance': individual_position_importance,
            'promiscuous_position_statistics': promiscuous_stats,
            'individual_position_statistics': individual_stats
        }
    
    def _analyze_pp_pet_positions(self, shap_results: Dict) -> Dict:
        """Analyze position importance for PP-PET selective designs."""
        sequences = shap_results['sequences']
        diff_shap = shap_results['difference_shap_values']
        pp_shap = shap_results['pp_shap_values']['shap_values']
        pet_shap = shap_results['pet_shap_values']['shap_values']
        
        # Flip sign for individual models (negative SHAP = better binding = positive importance)
        pp_shap_flipped = -pp_shap
        pet_shap_flipped = -pet_shap
        # Flip sign for selectivity too (negative diff = PP better = should be positive)
        diff_shap_flipped = -diff_shap
        
        # Calculate position importance
        diff_position_importance = np.sum(diff_shap_flipped, axis=2)
        pp_position_importance = np.sum(pp_shap_flipped, axis=2)
        pet_position_importance = np.sum(pet_shap_flipped, axis=2)
        
        # Calculate statistics for selectivity model (PP-PET difference)
        selectivity_stats = {}
        for pos in range(12):
            diff_pos = diff_position_importance[:, pos]
            selectivity_stats[f'position_{pos+1}'] = {
                'mean': np.mean(diff_pos),
                'std': np.std(diff_pos),
                'median': np.median(diff_pos),
                'min': np.min(diff_pos),
                'max': np.max(diff_pos)
            }
        
        # Calculate statistics for PP model
        pp_stats = {}
        for pos in range(12):
            pp_pos = pp_position_importance[:, pos]
            pp_stats[f'position_{pos+1}'] = {
                'mean': np.mean(pp_pos),
                'std': np.std(pp_pos),
                'median': np.median(pp_pos),
                'min': np.min(pp_pos),
                'max': np.max(pp_pos)
            }
        
        # Calculate statistics for PET model
        pet_stats = {}
        for pos in range(12):
            pet_pos = pet_position_importance[:, pos]
            pet_stats[f'position_{pos+1}'] = {
                'mean': np.mean(pet_pos),
                'std': np.std(pet_pos),
                'median': np.median(pet_pos),
                'min': np.min(pet_pos),
                'max': np.max(pet_pos)
            }
        
        return {
            'analysis_type': 'pp_pet_selective',
            'sequences': sequences,
            'selectivity_position_importance': diff_position_importance,
            'pp_position_importance': pp_position_importance,
            'pet_position_importance': pet_position_importance,
            'selectivity_position_statistics': selectivity_stats,
            'pp_position_statistics': pp_stats,
            'pet_position_statistics': pet_stats
        }
    
    def _analyze_single_plastic_positions(self, shap_results: Dict) -> Dict:
        """Analyze position importance for single plastic designs."""
        sequences = shap_results['sequences']
        shap_values = shap_results['shap_values']
        plastic_type = shap_results['plastic_type']
        
        # Flip sign for intuitive interpretation (negative SHAP = better binding = positive importance)
        shap_values_flipped = -shap_values
        
        position_importance = np.sum(shap_values_flipped, axis=2)
        
        position_stats = {}
        for pos in range(12):
            pos_values = position_importance[:, pos]
            position_stats[f'position_{pos+1}'] = {
                'mean': np.mean(pos_values),
                'std': np.std(pos_values),
                'median': np.median(pos_values)
            }
        
        return {
            'analysis_type': 'single_plastic',
            'plastic_type': plastic_type,
            'sequences': sequences,
            'position_importance': position_importance,
            'position_statistics': position_stats
        }
    
    def analyze_amino_acid_contributions(self, shap_results: Dict) -> Dict:
        """Analyze amino acid-level contributions with statistical analysis."""
        print("Analyzing amino acid contributions...")
        
        if shap_results['analysis_type'] == 'promiscuous':
            analysis = self._analyze_promiscuous_amino_acids(shap_results)
        elif shap_results['analysis_type'] == 'pp_pet_selective':
            analysis = self._analyze_pp_pet_amino_acids(shap_results)
        else:
            analysis = self._analyze_single_plastic_amino_acids(shap_results)
        
        return analysis
    
    def _analyze_promiscuous_amino_acids(self, shap_results: Dict) -> Dict:
        """Analyze amino acid contributions for promiscuous designs."""
        sequences = shap_results['sequences']
        avg_shap = shap_results['average_shap_values']
        individual_shap = shap_results['individual_shap_values']
        
        # Flip sign for intuitive interpretation (negative SHAP = better binding = positive importance)
        avg_shap_flipped = -avg_shap
        
        # Global amino acid importance (sum across positions)
        avg_aa_importance = np.sum(avg_shap_flipped, axis=1)  # (n_sequences, 18)
        
        individual_aa_importance = {}
        for plastic, data in individual_shap.items():
            # Flip sign for intuitive interpretation (negative SHAP = better binding = positive importance)
            plastic_shap_flipped = -data['shap_values']
            individual_aa_importance[plastic] = np.sum(plastic_shap_flipped, axis=1)
        
        # Calculate statistics for promiscuous model
        promiscuous_aa_stats = {}
        for aa_idx, aa in enumerate(self.amino_acids):
            prom_aa = avg_aa_importance[:, aa_idx]
            promiscuous_aa_stats[aa] = {
                'mean': np.mean(prom_aa),
                'std': np.std(prom_aa),
                'median': np.median(prom_aa),
                'min': np.min(prom_aa),
                'max': np.max(prom_aa)
            }
        
        # Calculate statistics for individual plastic models
        individual_aa_stats = {}
        for plastic, aa_importance_matrix in individual_aa_importance.items():
            individual_aa_stats[plastic] = {}
            for aa_idx, aa in enumerate(self.amino_acids):
                plastic_aa = aa_importance_matrix[:, aa_idx]
                individual_aa_stats[plastic][aa] = {
                    'mean': np.mean(plastic_aa),
                    'std': np.std(plastic_aa),
                    'median': np.median(plastic_aa),
                    'min': np.min(plastic_aa),
                    'max': np.max(plastic_aa)
                }
        
        return {
            'analysis_type': 'promiscuous',
            'sequences': sequences,
            'promiscuous_aa_importance': avg_aa_importance,
            'individual_aa_importance': individual_aa_importance,
            'promiscuous_aa_statistics': promiscuous_aa_stats,
            'individual_aa_statistics': individual_aa_stats
        }
    
    def _analyze_pp_pet_amino_acids(self, shap_results: Dict) -> Dict:
        """Analyze amino acid contributions for PP-PET selective designs."""
        sequences = shap_results['sequences']
        diff_shap = shap_results['difference_shap_values']
        pp_shap = shap_results['pp_shap_values']['shap_values']
        pet_shap = shap_results['pet_shap_values']['shap_values']
        
        # Flip sign for individual models (negative SHAP = better binding = positive importance)
        pp_shap_flipped = -pp_shap
        pet_shap_flipped = -pet_shap
        # Flip sign for selectivity too (negative diff = PP better = should be positive)
        diff_shap_flipped = -diff_shap
        
        # Global amino acid importance
        diff_aa_importance = np.sum(diff_shap_flipped, axis=1)
        pp_aa_importance = np.sum(pp_shap_flipped, axis=1)
        pet_aa_importance = np.sum(pet_shap_flipped, axis=1)
        
        # Calculate statistics for selectivity model (PP-PET difference)
        selectivity_aa_stats = {}
        for aa_idx, aa in enumerate(self.amino_acids):
            diff_aa = diff_aa_importance[:, aa_idx]
            selectivity_aa_stats[aa] = {
                'mean': np.mean(diff_aa),
                'std': np.std(diff_aa),
                'median': np.median(diff_aa),
                'min': np.min(diff_aa),
                'max': np.max(diff_aa)
            }
        
        # Calculate statistics for PP model
        pp_aa_stats = {}
        for aa_idx, aa in enumerate(self.amino_acids):
            pp_aa = pp_aa_importance[:, aa_idx]
            pp_aa_stats[aa] = {
                'mean': np.mean(pp_aa),
                'std': np.std(pp_aa),
                'median': np.median(pp_aa),
                'min': np.min(pp_aa),
                'max': np.max(pp_aa)
            }
        
        # Calculate statistics for PET model
        pet_aa_stats = {}
        for aa_idx, aa in enumerate(self.amino_acids):
            pet_aa = pet_aa_importance[:, aa_idx]
            pet_aa_stats[aa] = {
                'mean': np.mean(pet_aa),
                'std': np.std(pet_aa),
                'median': np.median(pet_aa),
                'min': np.min(pet_aa),
                'max': np.max(pet_aa)
            }
        
        return {
            'analysis_type': 'pp_pet_selective',
            'sequences': sequences,
            'selectivity_aa_importance': diff_aa_importance,
            'pp_aa_importance': pp_aa_importance,
            'pet_aa_importance': pet_aa_importance,
            'selectivity_aa_statistics': selectivity_aa_stats,
            'pp_aa_statistics': pp_aa_stats,
            'pet_aa_statistics': pet_aa_stats
        }
    
    def _analyze_single_plastic_amino_acids(self, shap_results: Dict) -> Dict:
        """Analyze amino acid contributions for single plastic designs."""
        sequences = shap_results['sequences']
        shap_values = shap_results['shap_values']
        plastic_type = shap_results['plastic_type']
        
        # Flip sign for intuitive interpretation (negative SHAP = better binding = positive importance)
        shap_values_flipped = -shap_values
        
        # Global amino acid importance
        aa_importance = np.sum(shap_values_flipped, axis=1)
        
        # Statistics
        aa_stats = {}
        for aa_idx, aa in enumerate(self.amino_acids):
            aa_values = aa_importance[:, aa_idx]
            aa_stats[aa] = {
                'mean': np.mean(aa_values),
                'std': np.std(aa_values),
                'median': np.median(aa_values)
            }
        
        
        return {
            'analysis_type': 'single_plastic',
            'plastic_type': plastic_type,
            'sequences': sequences,
            'aa_importance': aa_importance,
            'amino_acid_statistics': aa_stats,
        }
    
    def _save_analysis_results(self, position_analysis: Dict, aa_analysis: Dict, shap_results: Dict):
        """Save analysis results: statistics in reports, everything else in raw_data."""
        def convert_numpy(obj):
            if isinstance(obj, np.ndarray):
                return obj.tolist()
            elif isinstance(obj, np.integer):
                return int(obj)
            elif isinstance(obj, np.floating):
                return float(obj)
            elif isinstance(obj, np.bool_):
                return bool(obj)
            elif isinstance(obj, dict):
                return {key: convert_numpy(value) for key, value in obj.items()}
            elif isinstance(obj, list):
                return [convert_numpy(item) for item in obj]
            else:
                return obj
        
        # Setup directories
        analysis_type = shap_results['analysis_type']
        dirs = self._setup_directories(analysis_type)
        
        # === SAVE RAW DATA ===
        # Save arrays, sequences, and non-statistical data
        raw_data = {}
        
        # Extract arrays from position analysis
        for key, value in position_analysis.items():
            if isinstance(value, np.ndarray):
                raw_data[f"position_{key}"] = value
            elif isinstance(value, dict):
                for subkey, subvalue in value.items():
                    if isinstance(subvalue, np.ndarray):
                        raw_data[f"position_{key}_{subkey}"] = subvalue
        
        # Extract arrays from amino acid analysis  
        for key, value in aa_analysis.items():
            if isinstance(value, np.ndarray):
                raw_data[f"aa_{key}"] = value
            elif isinstance(value, dict):
                for subkey, subvalue in value.items():
                    if isinstance(subvalue, np.ndarray):
                        raw_data[f"aa_{key}_{subkey}"] = subvalue
        
        # Add sequences and metadata to raw data
        raw_data['sequences'] = shap_results['sequences']
        raw_data['analysis_type'] = analysis_type
        
        # Save arrays and metadata
        arrays_path = dirs['raw_data'] / "analysis_arrays.npz"
        np.savez_compressed(arrays_path, **raw_data)
        print(f"Saved analysis arrays and metadata to {arrays_path}")
        
        # === SAVE REPORTS (STATISTICS ONLY) ===
        # Only save statistics - no arrays, sequences, or detailed data
        report_data = {
            'position_statistics': convert_numpy(position_analysis.get('position_statistics', {})),
            'amino_acid_statistics': convert_numpy(aa_analysis.get('amino_acid_statistics', {}))
        }
        
        report_path = dirs['reports'] / "analysis_report.json"
        with open(report_path, 'w') as f:
            json.dump(report_data, f, indent=2)
        print(f"Saved statistics to {report_path}")
    
    def generate_comprehensive_report(self, shap_results: Dict) -> Dict:
        """Generate comprehensive analysis report ready for plotting."""
        print("Generating comprehensive SHAP analysis report...")
        
        position_analysis = self.analyze_position_importance(shap_results)
        aa_analysis = self.analyze_amino_acid_contributions(shap_results)
        
        # Save results to files
        self._save_analysis_results(position_analysis, aa_analysis, shap_results)
        
        # Return only what's needed for plotting
        report = {
            'position_analysis': position_analysis,
            'amino_acid_analysis': aa_analysis
        }
        
        print("Comprehensive analysis complete!")
        return report