SURE-tools 2.4.17__py3-none-any.whl → 2.4.22__py3-none-any.whl

SURE/DensityFlow.py

@@ -109,6 +109,13 @@ class DensityFlow(nn.Module):
 
         set_random_seed(seed)
         self.setup_networks()
+        
+        print(f"🧬 DensityFlow Initialized:")
+        print(f"   - Latent Dimension: {self.latent_dim}")
+        print(f"   - Gene Dimension: {self.input_size}")
+        print(f"   - Hidden Dimensions: {self.hidden_layers}")
+        print(f"   - Device: {self.get_device()}")
+        print(f"   - Parameters: {sum(p.numel() for p in self.parameters()):,}")
 
     def setup_networks(self):
         latent_dim = self.latent_dim

SURE/PerturbationAwareDecoder.py

@@ -0,0 +1,723 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+from typing import Dict, List, Optional, Tuple, Union
+import math
+import warnings
+warnings.filterwarnings('ignore')
+
+class PerturbationAwareDecoder:
+    """
+    Advanced transcriptome decoder with perturbation awareness
+    Supports single hidden layer configurations
+    """
+    
+    def __init__(self, 
+                 latent_dim: int = 100,
+                 num_known_perturbations: int = 50,
+                 gene_dim: int = 60000,
+                 hidden_dims: List[int] = [512],  # Now supports single element list
+                 perturbation_embedding_dim: int = 128,
+                 biological_prior_dim: int = 256,
+                 dropout_rate: float = 0.1,
+                 device: str = None):
+        """
+        Multi-modal decoder supporting single hidden layer configurations
+        
+        Args:
+            latent_dim: Latent variable dimension
+            num_known_perturbations: Number of known perturbation types
+            gene_dim: Number of genes
+            hidden_dims: Hidden layer dimensions (can be single element list)
+            perturbation_embedding_dim: Embedding dimension for perturbations
+            biological_prior_dim: Dimension for biological prior knowledge
+            dropout_rate: Dropout rate
+            device: Computation device
+        """
+        self.latent_dim = latent_dim
+        self.num_known_perturbations = num_known_perturbations
+        self.gene_dim = gene_dim
+        self.hidden_dims = hidden_dims
+        self.perturbation_embedding_dim = perturbation_embedding_dim
+        self.biological_prior_dim = biological_prior_dim
+        self.dropout_rate = dropout_rate
+        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        # Validate hidden_dims
+        self._validate_hidden_dims()
+        
+        # Initialize multi-modal model
+        self.model = self._build_single_layer_compatible_model()
+        self.model.to(self.device)
+        
+        # Training state
+        self.is_trained = False
+        self.training_history = None
+        self.best_val_loss = float('inf')
+        self.known_perturbation_names = []
+        self.perturbation_prototypes = None
+        
+        print(f"🧬 PerturbationAwareDecoder Initialized:")
+        print(f"   - Latent Dimension: {latent_dim}")
+        print(f"   - Known Perturbations: {num_known_perturbations}")
+        print(f"   - Gene Dimension: {gene_dim}")
+        print(f"   - Hidden Dimensions: {hidden_dims}")
+        print(f"   - Parameters: {sum(p.numel() for p in self.model.parameters()):,}")
+    
+    def _validate_hidden_dims(self):
+        """Validate hidden_dims parameter"""
+        assert len(self.hidden_dims) >= 1, "hidden_dims must have at least one element"
+        assert all(dim > 0 for dim in self.hidden_dims), "All hidden dimensions must be positive"
+        
+        if len(self.hidden_dims) == 1:
+            print("🔧 Single hidden layer configuration detected")
+        else:
+            print(f"🔧 Multi-layer configuration: {len(self.hidden_dims)} hidden layers")
+    
+    class SingleLayerCompatiblePerturbationEncoder(nn.Module):
+        """Perturbation encoder compatible with single hidden layer"""
+        
+        def __init__(self, num_perturbations: int, embedding_dim: int, hidden_dim: int):
+            super().__init__()
+            self.num_perturbations = num_perturbations
+            
+            # Embedding for perturbation types
+            self.perturbation_embedding = nn.Embedding(num_perturbations, embedding_dim)
+            
+            # Single projection layer (compatible with single hidden layer)
+            self.projection = nn.Sequential(
+                nn.Linear(embedding_dim, hidden_dim),
+                nn.ReLU(),
+                nn.Dropout(0.1)
+            )
+            
+        def forward(self, one_hot_perturbations):
+            batch_size = one_hot_perturbations.shape[0]
+            
+            # Convert one-hot to indices
+            perturbation_indices = torch.argmax(one_hot_perturbations, dim=1)
+            
+            # Get perturbation embeddings
+            perturbation_embeds = self.perturbation_embedding(perturbation_indices)
+            
+            # Single projection
+            hidden_repr = self.projection(perturbation_embeds)
+            
+            return hidden_repr
+    
+    class SingleLayerCrossModalFusion(nn.Module):
+        """Cross-modal fusion compatible with single hidden layer"""
+        
+        def __init__(self, latent_dim: int, perturbation_dim: int, fusion_dim: int):
+            super().__init__()
+            self.latent_projection = nn.Linear(latent_dim, fusion_dim)
+            self.perturbation_projection = nn.Linear(perturbation_dim, fusion_dim)
+            
+            # Simplified fusion for single layer
+            self.fusion_gate = nn.Sequential(
+                nn.Linear(fusion_dim * 2, fusion_dim),
+                nn.Sigmoid()
+            )
+            
+            self.norm = nn.LayerNorm(fusion_dim)
+            self.dropout = nn.Dropout(0.1)
+        
+        def forward(self, latent, perturbation_encoded):
+            # Project both modalities
+            latent_proj = self.latent_projection(latent)
+            perturbation_proj = self.perturbation_projection(perturbation_encoded)
+            
+            # Simple gated fusion (no complex attention for single layer)
+            concatenated = torch.cat([latent_proj, perturbation_proj], dim=-1)
+            fusion_gate = self.fusion_gate(concatenated)
+            
+            # Gated fusion
+            fused = fusion_gate * latent_proj + (1 - fusion_gate) * perturbation_proj
+            fused = self.norm(fused)
+            fused = self.dropout(fused)
+            
+            return fused
+    
+    class SingleLayerPerturbationResponseNetwork(nn.Module):
+        """Response network compatible with single hidden layer"""
+        
+        def __init__(self, fusion_dim: int, gene_dim: int, hidden_dims: List[int]):
+            super().__init__()
+            
+            # Handle single layer case
+            if len(hidden_dims) == 1:
+                # Single hidden layer
+                self.base_network = nn.Sequential(
+                    nn.Linear(fusion_dim, hidden_dims[0]),
+                    nn.BatchNorm1d(hidden_dims[0]),
+                    nn.ReLU(),
+                    nn.Dropout(0.1)
+                )
+                final_input_dim = hidden_dims[0]
+            else:
+                # Multiple hidden layers
+                layers = []
+                input_dim = fusion_dim
+                
+                for hidden_dim in hidden_dims:
+                    layers.extend([
+                        nn.Linear(input_dim, hidden_dim),
+                        nn.BatchNorm1d(hidden_dim),
+                        nn.ReLU(),
+                        nn.Dropout(0.1)
+                    ])
+                    input_dim = hidden_dim
+                
+                self.base_network = nn.Sequential(*layers)
+                final_input_dim = hidden_dims[-1]
+            
+            # Final projection
+            self.final_projection = nn.Linear(final_input_dim, gene_dim)
+            
+            # Perturbation-aware scaling (simplified for single layer)
+            self.scale = nn.Linear(fusion_dim, 1)
+            self.bias = nn.Linear(fusion_dim, 1)
+        
+        def forward(self, fused_representation):
+            base_output = self.base_network(fused_representation)
+            expression = self.final_projection(base_output)
+            
+            # Perturbation-aware scaling
+            scale = torch.sigmoid(self.scale(fused_representation)) * 2
+            bias = self.bias(fused_representation)
+            
+            return F.softplus(expression * scale + bias)
+    
+    class SingleLayerNovelPerturbationPredictor(nn.Module):
+        """Novel perturbation predictor compatible with single hidden layer"""
+        
+        def __init__(self, num_known_perturbations: int, gene_dim: int, hidden_dim: int):
+            super().__init__()
+            self.num_known_perturbations = num_known_perturbations
+            self.gene_dim = gene_dim
+            
+            # Learnable perturbation prototypes
+            self.perturbation_prototypes = nn.Parameter(
+                torch.randn(num_known_perturbations, gene_dim) * 0.1
+            )
+            
+            # Simplified response generator for single layer
+            if hidden_dim > 0:
+                self.response_generator = nn.Sequential(
+                    nn.Linear(num_known_perturbations, hidden_dim),
+                    nn.ReLU(),
+                    nn.Linear(hidden_dim, gene_dim)
+                )
+            else:
+                # Direct projection if no hidden layer
+                self.response_generator = nn.Linear(num_known_perturbations, gene_dim)
+            
+            # Simplified attention
+            self.similarity_attention = nn.Sequential(
+                nn.Linear(num_known_perturbations, num_known_perturbations),
+                nn.Softmax(dim=-1)
+            )
+        
+        def forward(self, similarity_matrix, latent_features=None):
+            batch_size = similarity_matrix.shape[0]
+            
+            # Method 1: Attention-weighted combination of known responses
+            attention_weights = self.similarity_attention(similarity_matrix)
+            weighted_response = torch.matmul(attention_weights, self.perturbation_prototypes)
+            
+            # Method 2: Direct generation from similarity
+            generated_response = self.response_generator(similarity_matrix)
+            
+            # Simple combination
+            combination_weights = torch.sigmoid(similarity_matrix.mean(dim=1, keepdim=True))
+            final_response = (combination_weights * weighted_response + 
+                            (1 - combination_weights) * generated_response)
+            
+            return final_response
+    
+    class SingleLayerCompatibleDecoder(nn.Module):
+        """Main decoder compatible with single hidden layer configurations"""
+        
+        def __init__(self, latent_dim: int, num_known_perturbations: int, gene_dim: int, 
+                    hidden_dims: List[int], perturbation_embedding_dim: int, 
+                    biological_prior_dim: int, dropout_rate: float):
+            super().__init__()
+            
+            self.num_known_perturbations = num_known_perturbations
+            self.latent_dim = latent_dim
+            self.gene_dim = gene_dim
+            
+            # Handle single layer case for hidden_dims
+            if len(hidden_dims) == 1:
+                # Use the single dimension for all components
+                main_hidden_dim = hidden_dims[0]
+                response_hidden_dims = []  # No additional hidden layers for response
+            else:
+                # Multiple layers: first for fusion, rest for response
+                main_hidden_dim = hidden_dims[0]
+                response_hidden_dims = hidden_dims[1:]
+            
+            # Perturbation encoder
+            self.perturbation_encoder = PerturbationAwareDecoder.SingleLayerCompatiblePerturbationEncoder(
+                num_known_perturbations, perturbation_embedding_dim, main_hidden_dim
+            )
+            
+            # Cross-modal fusion
+            self.cross_modal_fusion = PerturbationAwareDecoder.SingleLayerCrossModalFusion(
+                latent_dim, main_hidden_dim, main_hidden_dim
+            )
+            
+            # Response network
+            self.response_network = PerturbationAwareDecoder.SingleLayerPerturbationResponseNetwork(
+                main_hidden_dim, gene_dim, response_hidden_dims
+            )
+            
+            # Novel perturbation predictor
+            self.novel_predictor = PerturbationAwareDecoder.SingleLayerNovelPerturbationPredictor(
+                num_known_perturbations, gene_dim, main_hidden_dim
+            )
+            
+        def forward(self, latent, perturbation_matrix, mode='one_hot'):
+            if mode == 'one_hot':
+                # Known perturbation pathway
+                perturbation_encoded = self.perturbation_encoder(perturbation_matrix)
+                fused = self.cross_modal_fusion(latent, perturbation_encoded)
+                expression = self.response_network(fused)
+                
+            elif mode == 'similarity':
+                # Novel perturbation pathway
+                expression = self.novel_predictor(perturbation_matrix, latent)
+                
+            else:
+                raise ValueError(f"Unknown mode: {mode}. Use 'one_hot' or 'similarity'")
+            
+            return expression
+        
+        def get_perturbation_prototypes(self):
+            """Get learned perturbation response prototypes"""
+            return self.novel_predictor.perturbation_prototypes.detach()
+    
+    def _build_single_layer_compatible_model(self):
+        """Build model compatible with single hidden layer"""
+        return self.SingleLayerCompatibleDecoder(
+            self.latent_dim, self.num_known_perturbations, self.gene_dim,
+            self.hidden_dims, self.perturbation_embedding_dim,
+            self.biological_prior_dim, self.dropout_rate
+        )
+    
+    def train(self,
+              train_latent: np.ndarray,
+              train_perturbations: np.ndarray,
+              train_expression: np.ndarray,
+              val_latent: np.ndarray = None,
+              val_perturbations: np.ndarray = None,
+              val_expression: np.ndarray = None,
+              batch_size: int = 32,
+              num_epochs: int = 200,
+              learning_rate: float = 1e-4,
+              checkpoint_path: str = 'single_layer_decoder.pth') -> Dict:
+        """
+        Train the decoder with single layer compatibility
+        """
+        print("🧬 Starting Training with Single Layer Support...")
+        
+        # Validate one-hot encoding
+        self._validate_one_hot_perturbations(train_perturbations)
+        
+        # Data preparation
+        train_dataset = self._create_dataset(train_latent, train_perturbations, train_expression)
+        
+        if val_latent is not None and val_perturbations is not None and val_expression is not None:
+            val_dataset = self._create_dataset(val_latent, val_perturbations, val_expression)
+            train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+            val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
+        else:
+            train_size = int(0.9 * len(train_dataset))
+            val_size = len(train_dataset) - train_size
+            train_subset, val_subset = torch.utils.data.random_split(train_dataset, [train_size, val_size])
+            train_loader = DataLoader(train_subset, batch_size=batch_size, shuffle=True)
+            val_loader = DataLoader(val_subset, batch_size=batch_size, shuffle=False)
+        
+        print(f"📊 Training samples: {len(train_loader.dataset)}")
+        print(f"📊 Validation samples: {len(val_loader.dataset)}")
+        print(f"🔧 Hidden layers: {len(self.hidden_dims)}")
+        print(f"🔧 Hidden dimensions: {self.hidden_dims}")
+        
+        # Optimizer
+        optimizer = optim.AdamW(
+            self.model.parameters(),
+            lr=learning_rate,
+            weight_decay=1e-5
+        )
+        
+        # Scheduler
+        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)
+        
+        # Loss function
+        def loss_fn(pred, target):
+            mse_loss = F.mse_loss(pred, target)
+            poisson_loss = (pred - target * torch.log(pred + 1e-8)).mean()
+            correlation = self._pearson_correlation(pred, target)
+            correlation_loss = 1 - correlation
+            return mse_loss + 0.3 * poisson_loss + 0.1 * correlation_loss
+        
+        # Training history
+        history = {
+            'train_loss': [], 'val_loss': [],
+            'train_mse': [], 'val_mse': [],
+            'train_correlation': [], 'val_correlation': [],
+            'learning_rates': []
+        }
+        
+        best_val_loss = float('inf')
+        patience = 20
+        patience_counter = 0
+        
+        print("\n🔬 Starting training...")
+        for epoch in range(1, num_epochs + 1):
+            # Training
+            train_metrics = self._train_epoch(train_loader, optimizer, loss_fn)
+            
+            # Validation
+            val_metrics = self._validate_epoch(val_loader, loss_fn)
+            
+            # Update scheduler
+            scheduler.step()
+            current_lr = optimizer.param_groups[0]['lr']
+            
+            # Record history
+            history['train_loss'].append(train_metrics['loss'])
+            history['val_loss'].append(val_metrics['loss'])
+            history['train_mse'].append(train_metrics['mse'])
+            history['val_mse'].append(val_metrics['mse'])
+            history['train_correlation'].append(train_metrics['correlation'])
+            history['val_correlation'].append(val_metrics['correlation'])
+            history['learning_rates'].append(current_lr)
+            
+            # Print progress
+            if epoch % 10 == 0 or epoch == 1:
+                print(f"🧪 Epoch {epoch:3d}/{num_epochs} | "
+                      f"Train: {train_metrics['loss']:.4f} | "
+                      f"Val: {val_metrics['loss']:.4f} | "
+                      f"Corr: {val_metrics['correlation']:.4f} | "
+                      f"LR: {current_lr:.2e}")
+            
+            # Early stopping
+            if val_metrics['loss'] < best_val_loss:
+                best_val_loss = val_metrics['loss']
+                patience_counter = 0
+                self._save_checkpoint(epoch, optimizer, scheduler, best_val_loss, history, checkpoint_path)
+            else:
+                patience_counter += 1
+                if patience_counter >= patience:
+                    print(f"🛑 Early stopping at epoch {epoch}")
+                    break
+        
+        self.is_trained = True
+        self.training_history = history
+        self.best_val_loss = best_val_loss
+        self.perturbation_prototypes = self.model.get_perturbation_prototypes().cpu().numpy()
+        
+        print(f"\n🎉 Training completed! Best val loss: {best_val_loss:.4f}")
+        return history
+    
+    def _validate_one_hot_perturbations(self, perturbations):
+        """Validate that perturbations are proper one-hot encodings"""
+        assert perturbations.shape[1] == self.num_known_perturbations, \
+            f"Perturbation dimension {perturbations.shape[1]} doesn't match expected {self.num_known_perturbations}"
+        
+        row_sums = perturbations.sum(axis=1)
+        valid_rows = np.all((row_sums == 0) | (row_sums == 1))
+        assert valid_rows, "Perturbations should be one-hot encoded (sum to 0 or 1 per row)"
+        
+        print("✅ One-hot perturbations validated")
+    
+    def _create_dataset(self, latent_data, perturbations, expression_data):
+        """Create dataset with one-hot perturbations"""
+        class OneHotDataset(Dataset):
+            def __init__(self, latent, perturbations, expression):
+                self.latent = torch.FloatTensor(latent)
+                self.perturbations = torch.FloatTensor(perturbations)
+                self.expression = torch.FloatTensor(expression)
+            
+            def __len__(self):
+                return len(self.latent)
+            
+            def __getitem__(self, idx):
+                return self.latent[idx], self.perturbations[idx], self.expression[idx]
+        
+        return OneHotDataset(latent_data, perturbations, expression_data)
+    
+    def predict(self, 
+                latent_data: np.ndarray, 
+                perturbations: np.ndarray, 
+                batch_size: int = 32) -> np.ndarray:
+        """
+        Predict expression for known perturbations
+        """
+        if not self.is_trained:
+            warnings.warn("⚠️ Model not trained. Predictions may be inaccurate.")
+        
+        self._validate_one_hot_perturbations(perturbations)
+        
+        self.model.eval()
+        
+        if isinstance(latent_data, np.ndarray):
+            latent_data = torch.FloatTensor(latent_data)
+        if isinstance(perturbations, np.ndarray):
+            perturbations = torch.FloatTensor(perturbations)
+        
+        predictions = []
+        with torch.no_grad():
+            for i in range(0, len(latent_data), batch_size):
+                batch_latent = latent_data[i:i+batch_size].to(self.device)
+                batch_perturbations = perturbations[i:i+batch_size].to(self.device)
+                
+                batch_pred = self.model(batch_latent, batch_perturbations, mode='one_hot')
+                predictions.append(batch_pred.cpu())
+        
+        return torch.cat(predictions).numpy()
+    
+    def predict_novel_perturbation(self,
+                                 latent_data: np.ndarray,
+                                 similarity_matrix: np.ndarray,
+                                 batch_size: int = 32) -> np.ndarray:
+        """
+        Predict response to novel perturbations
+        """
+        if not self.is_trained:
+            warnings.warn("⚠️ Model not trained. Novel perturbation prediction may be inaccurate.")
+        
+        assert similarity_matrix.shape[1] == self.num_known_perturbations, \
+            f"Similarity matrix columns {similarity_matrix.shape[1]} must match known perturbations {self.num_known_perturbations}"
+        
+        self.model.eval()
+        
+        if isinstance(latent_data, np.ndarray):
+            latent_data = torch.FloatTensor(latent_data)
+        if isinstance(similarity_matrix, np.ndarray):
+            similarity_matrix = torch.FloatTensor(similarity_matrix)
+        
+        predictions = []
+        with torch.no_grad():
+            for i in range(0, len(latent_data), batch_size):
+                batch_latent = latent_data[i:i+batch_size].to(self.device)
+                batch_similarity = similarity_matrix[i:i+batch_size].to(self.device)
+                
+                batch_pred = self.model(batch_latent, batch_similarity, mode='similarity')
+                predictions.append(batch_pred.cpu())
+        
+        return torch.cat(predictions).numpy()
+    
+    def get_known_perturbation_prototypes(self) -> np.ndarray:
+        """Get learned perturbation response prototypes"""
+        if not self.is_trained:
+            warnings.warn("⚠️ Model not trained. Prototypes may be uninformative.")
+        
+        if self.perturbation_prototypes is None:
+            self.model.eval()
+            with torch.no_grad():
+                self.perturbation_prototypes = self.model.get_perturbation_prototypes().cpu().numpy()
+        
+        return self.perturbation_prototypes
+    
+    def _pearson_correlation(self, pred, target):
+        """Calculate Pearson correlation coefficient"""
+        pred_centered = pred - pred.mean(dim=1, keepdim=True)
+        target_centered = target - target.mean(dim=1, keepdim=True)
+        
+        numerator = (pred_centered * target_centered).sum(dim=1)
+        denominator = torch.sqrt(torch.sum(pred_centered ** 2, dim=1)) * torch.sqrt(torch.sum(target_centered ** 2, dim=1))
+        
+        return (numerator / (denominator + 1e-8)).mean()
+    
+    def _train_epoch(self, train_loader, optimizer, loss_fn):
+        """Train one epoch"""
+        self.model.train()
+        total_loss = 0
+        total_mse = 0
+        total_correlation = 0
+        
+        for latent, perturbations, target in train_loader:
+            latent = latent.to(self.device)
+            perturbations = perturbations.to(self.device)
+            target = target.to(self.device)
+            
+            optimizer.zero_grad()
+            pred = self.model(latent, perturbations, mode='one_hot')
+            
+            loss = loss_fn(pred, target)
+            loss.backward()
+            
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+            optimizer.step()
+            
+            mse_loss = F.mse_loss(pred, target).item()
+            correlation = self._pearson_correlation(pred, target).item()
+            
+            total_loss += loss.item()
+            total_mse += mse_loss
+            total_correlation += correlation
+        
+        num_batches = len(train_loader)
+        return {
+            'loss': total_loss / num_batches,
+            'mse': total_mse / num_batches,
+            'correlation': total_correlation / num_batches
+        }
+    
+    def _validate_epoch(self, val_loader, loss_fn):
+        """Validate one epoch"""
+        self.model.eval()
+        total_loss = 0
+        total_mse = 0
+        total_correlation = 0
+        
+        with torch.no_grad():
+            for latent, perturbations, target in val_loader:
+                latent = latent.to(self.device)
+                perturbations = perturbations.to(self.device)
+                target = target.to(self.device)
+                
+                pred = self.model(latent, perturbations, mode='one_hot')
+                loss = loss_fn(pred, target)
+                mse_loss = F.mse_loss(pred, target).item()
+                correlation = self._pearson_correlation(pred, target).item()
+                
+                total_loss += loss.item()
+                total_mse += mse_loss
+                total_correlation += correlation
+        
+        num_batches = len(val_loader)
+        return {
+            'loss': total_loss / num_batches,
+            'mse': total_mse / num_batches,
+            'correlation': total_correlation / num_batches
+        }
+    
+    def _save_checkpoint(self, epoch, optimizer, scheduler, best_loss, history, path):
+        """Save model checkpoint"""
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': self.model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'scheduler_state_dict': scheduler.state_dict(),
+            'best_val_loss': best_loss,
+            'training_history': history,
+            'perturbation_prototypes': self.perturbation_prototypes,
+            'model_config': {
+                'latent_dim': self.latent_dim,
+                'num_known_perturbations': self.num_known_perturbations,
+                'gene_dim': self.gene_dim,
+                'hidden_dims': self.hidden_dims
+            }
+        }, path)
+    
+    def load_model(self, model_path: str):
+        """Load pre-trained model"""
+        checkpoint = torch.load(model_path, map_location=self.device)
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        self.perturbation_prototypes = checkpoint.get('perturbation_prototypes')
+        self.is_trained = True
+        self.training_history = checkpoint.get('training_history')
+        self.best_val_loss = checkpoint.get('best_val_loss', float('inf'))
+        print(f"✅ Model loaded! Best val loss: {self.best_val_loss:.4f}")
+
+'''# Example usage with single hidden layer
+def example_single_layer_usage():
+    """Example demonstration with single hidden layer"""
+    
+    # Initialize decoder with single hidden layer
+    decoder = PerturbationAwareDecoder(
+        latent_dim=100,
+        num_known_perturbations=10,
+        gene_dim=2000,  # Reduced for example
+        hidden_dims=[512],  # Single hidden layer
+        perturbation_embedding_dim=128
+    )
+    
+    # Generate example data
+    n_samples = 1000
+    n_perturbations = 10
+    
+    # Latent variables
+    latent_data = np.random.randn(n_samples, 100).astype(np.float32)
+    
+    # One-hot encoded perturbations
+    perturbations = np.zeros((n_samples, n_perturbations))
+    for i in range(n_samples):
+        if i % 10 != 0:  # 90% perturbed, 10% control
+            perturb_id = np.random.randint(0, n_perturbations)
+            perturbations[i, perturb_id] = 1.0
+    
+    # Expression data
+    base_weights = np.random.randn(100, 2000) * 0.1
+    perturbation_effects = np.random.randn(n_perturbations, 2000) * 0.5
+    
+    expression_data = np.tanh(latent_data.dot(base_weights))
+    for i in range(n_samples):
+        if perturbations[i].sum() > 0:
+            perturb_id = np.argmax(perturbations[i])
+            expression_data[i] += perturbation_effects[perturb_id]
+    
+    expression_data = np.maximum(expression_data, 0)
+    
+    print(f"📊 Data shapes: Latent {latent_data.shape}, Perturbations {perturbations.shape}, Expression {expression_data.shape}")
+    print(f"🔧 Hidden layers: {len(decoder.hidden_dims)}")
+    print(f"🔧 Hidden dimensions: {decoder.hidden_dims}")
+    
+    # Train with single hidden layer
+    history = decoder.train(
+        train_latent=latent_data,
+        train_perturbations=perturbations,
+        train_expression=expression_data,
+        batch_size=32,
+        num_epochs=50
+    )
+    
+    # Test predictions
+    test_latent = np.random.randn(10, 100).astype(np.float32)
+    test_perturbations = np.zeros((10, n_perturbations))
+    for i in range(10):
+        test_perturbations[i, i % n_perturbations] = 1.0
+    
+    predictions = decoder.predict(test_latent, test_perturbations)
+    print(f"🔮 Known perturbation prediction shape: {predictions.shape}")
+    
+    # Test novel perturbation prediction
+    test_latent_novel = np.random.randn(5, 100).astype(np.float32)
+    similarity_matrix = np.random.rand(5, n_perturbations)
+    similarity_matrix = similarity_matrix / similarity_matrix.sum(axis=1, keepdims=True)
+    
+    novel_predictions = decoder.predict_novel_perturbation(test_latent_novel, similarity_matrix)
+    print(f"🔮 Novel perturbation prediction shape: {novel_predictions.shape}")
+    
+    return decoder
+
+# Example usage with multiple hidden layers
+def example_multi_layer_usage():
+    """Example demonstration with multiple hidden layers"""
+    
+    # Initialize decoder with multiple hidden layers
+    decoder = PerturbationAwareDecoder(
+        latent_dim=100,
+        num_known_perturbations=10,
+        gene_dim=2000,
+        hidden_dims=[256, 512, 1024],  # Multiple hidden layers
+        perturbation_embedding_dim=128
+    )
+    
+    print(f"🔧 Hidden layers: {len(decoder.hidden_dims)}")
+    print(f"🔧 Hidden dimensions: {decoder.hidden_dims}")
+    
+    return decoder
+
+if __name__ == "__main__":
+    print("=== Single Hidden Layer Example ===")
+    decoder_single = example_single_layer_usage()
+    
+    print("\n=== Multiple Hidden Layers Example ===")
+    decoder_multi = example_multi_layer_usage()'''

SURE/VirtualCellDecoder.py

@@ -56,7 +56,6 @@ class VirtualCellDecoder:
         print(f"   - Biological Prior Dimension: {biological_prior_dim}")
         print(f"   - Device: {self.device}")
         print(f"   - Parameters: {sum(p.numel() for p in self.model.parameters()):,}")
-        print(f"   - Estimated GPU Memory: ~8-10GB (optimized for 20GB)")
     
     class BiologicalPriorNetwork(nn.Module):
         """Biological prior network based on gene regulatory knowledge"""

SURE/__init__.py

@@ -5,6 +5,7 @@ from .TranscriptomeDecoder import TranscriptomeDecoder
 from .SimpleTranscriptomeDecoder import SimpleTranscriptomeDecoder
 from .EfficientTranscriptomeDecoder import EfficientTranscriptomeDecoder 
 from .VirtualCellDecoder import VirtualCellDecoder
+from .PerturbationAwareDecoder import PerturbationAwareDecoder
 
 from . import utils 
 from . import codebook
@@ -18,6 +19,8 @@ from . import TranscriptomeDecoder
 from . import SimpleTranscriptomeDecoder
 from . import EfficientTranscriptomeDecoder
 from . import VirtualCellDecoder
+from . import PerturbationAwareDecoder
 
 __all__ = ['SURE', 'DensityFlow', 'PerturbE', 'TranscriptomeDecoder', 'SimpleTranscriptomeDecoder',
-           'EfficientTranscriptomeDecoder', 'VirtualCellDecoder', 'flow', 'perturb', 'atac', 'utils', 'codebook']
+           'EfficientTranscriptomeDecoder', 'VirtualCellDecoder', 'PerturbationAwareDecoder', 
+           'flow', 'perturb', 'atac', 'utils', 'codebook']

{sure_tools-2.4.17.dist-info → sure_tools-2.4.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: SURE-tools
-Version: 2.4.17
+Version: 2.4.22
 Summary: Succinct Representation of Single Cells
 Home-page: https://github.com/ZengFLab/SURE
 Author: Feng Zeng

{sure_tools-2.4.17.dist-info → sure_tools-2.4.22.dist-info}/RECORD

@@ -1,11 +1,12 @@
-SURE/DensityFlow.py,sha256=YvaE9aPbAC2U7WhTye5i2AMtcw0BI_qS3gv9SP4aE0k,56676
+SURE/DensityFlow.py,sha256=mp2hVxqwg-wVlDEc2Q2flrtzowDVyMzjC7od3rZTm40,57044
 SURE/EfficientTranscriptomeDecoder.py,sha256=O_x-4edKBU5OJJbOOS-59u3TQElZqhAtOVJMPlpw8m0,21667
 SURE/PerturbE.py,sha256=DxEp-qef--x8-GMZdPfBf8ts8UDDc34h2P5AnpqZ-YM,52265
+SURE/PerturbationAwareDecoder.py,sha256=hhdGC9kIOlfFaWrO4hnaF91PEMSEMqVSovzDS-QWytk,30653
 SURE/SURE.py,sha256=MXs7iuvcj-lU4dJ_MwKegpL2Rqk2HB4eFfAgHRA3RtA,47744
 SURE/SimpleTranscriptomeDecoder.py,sha256=mLgYYipfRrmuXlpoaxLPfJS009OVwyshdL3nXTJygIE,22285
 SURE/TranscriptomeDecoder.py,sha256=n2tVB8hNVLwSQ1G1Jpd6WzMl2Iw63eK0_Ujk9d48SJY,20982
-SURE/VirtualCellDecoder.py,sha256=X37TbiddREugON10Y9TJ2M-UzrRKe2o185d2-6FtW_M,27542
-SURE/__init__.py,sha256=ZHatQR813LVibrHOLkemTXFt3eSpyluWGxAi8doQTD4,879
+SURE/VirtualCellDecoder.py,sha256=z1Z7GRTYmTE3DaSKZueofv138R0J7kGFfnh0a_Lee38,27468
+SURE/__init__.py,sha256=7P8BRLq4sPXwQX5qEeafMPxAnRzxFbK2ZO_QTmQOlhk,1021
 SURE/assembly/__init__.py,sha256=jxZLURXKPzXe21LhrZ09LgZr33iqdjlQy4oSEj5gR2Q,172
 SURE/assembly/assembly.py,sha256=6IMdelPOiRO4mUb4dC7gVCoF1Uvfw86-Map8P_jnUag,21477
 SURE/assembly/atlas.py,sha256=ALjmVWutm_tOHTcT1aqOxmuCEQw-XzrtDoMCV_8oXLk,21794
@@ -22,9 +23,9 @@ SURE/utils/__init__.py,sha256=YF5jB-PAHJQ40OlcZ7BCZbsN2q1JKuPT6EppilRXQqM,680
 SURE/utils/custom_mlp.py,sha256=Rn_PQouxPMSda-KKBYrwVVv3GFFuUmCLxp8cV5LszZo,10580
 SURE/utils/queue.py,sha256=E_5PA5EWcBoGAZj8BkKQnkCK0p4C-4-xcTPqdIXaPXU,1892
 SURE/utils/utils.py,sha256=IUHjDDtYaAYllCWsZyIzqQwaLul6fJRvHRH4vIYcR-c,8462
-sure_tools-2.4.17.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
-sure_tools-2.4.17.dist-info/METADATA,sha256=Qv8OlzL7iQmfG1pGSdSQG-uo92pVRFHJTvMASWxdwwo,2678
-sure_tools-2.4.17.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sure_tools-2.4.17.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
-sure_tools-2.4.17.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
-sure_tools-2.4.17.dist-info/RECORD,,
+sure_tools-2.4.22.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
+sure_tools-2.4.22.dist-info/METADATA,sha256=W67v_Fs0SOx1tBPZN4-7U1Y6loSUjS697gL_9Q6piC4,2678
+sure_tools-2.4.22.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sure_tools-2.4.22.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
+sure_tools-2.4.22.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
+sure_tools-2.4.22.dist-info/RECORD,,