embed.py

'''
Code for embedding omics data with transformer or (V)AE

'''

import pytorch_lightning as pl 
import torch
import torch.nn as nn
import torch.nn.functional as F
import pandas as pd
import numpy as np
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm
from augment_omics.model.nn import init_model, AE, VAE
from typing import List



def mask_matrix(
    matrix: torch.Tensor,
    mask_ratio: float = 0.15,
    impute_value: float = 0.0,
):
    """
    For each row, sample a random mask ratio in [0, max_mask_ratio],
    then mask that fraction of non-NaN entries in the row.

    Args:
        matrix: 2D torch.Tensor (rows x cols) with possible NaNs
        max_mask_ratio: upper bound for per-row mask ratio (0–1)
        impute_value: value to fill for masked and NaN entries

    Returns:
        masked_matrix: matrix with masked & NaN cells replaced by impute_value
        masked_bool: boolean tensor where True = masked
        nan_bool: boolean tensor where True = originally NaN
        row_ratios: numpy array of per-row sampled mask ratios
    """
    if matrix.dim() != 2:
        raise ValueError("matrix must be 2D")

    matrix = matrix.clone()
    device = matrix.device
    rows, cols = matrix.shape

    non_nan = ~torch.isnan(matrix)
    nan_bool = ~non_nan
    masked_bool = torch.zeros_like(non_nan, dtype=torch.bool)

    # sample random ratios per row from [0, max_mask_ratio]
    row_ratios = np.random.uniform(0, mask_ratio, size=rows)

    for i in range(rows):
        valid_cols = torch.nonzero(non_nan[i], as_tuple=True)[0].cpu().numpy()
        n = len(valid_cols)
        if n == 0:
            continue

        k = int(row_ratios[i] * n)
        if k <= 0:
            continue

        chosen = np.random.choice(valid_cols, size=k, replace=False)
        masked_bool[i, torch.tensor(chosen, device=device)] = True

    # replace NaNs and masked values with impute_value
    matrix[nan_bool] = impute_value
    matrix[masked_bool] = impute_value

    return matrix, masked_bool, nan_bool



class BaseEmbed(pl.LightningModule):
    def __init__(self):
        super().__init__() 
        
    def training_step(self,):
        raise NotImplementedError
    
    def validation_step(self,):
        raise NotImplementedError

    def embed(self, data: pd.DataFrame, batch_size: int):
        raise NotImplementedError
        
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr,
                                     weight_decay=self.hparams.weight_decay) ## add weight decay
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer,
            mode='min',
            factor=0.5,
            patience=5,)
        
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "monitor": self.hparams.monitor,
                "frequency": 1,
                "interval": "epoch",
            },
        }
        
    def mse_loss(self, y_pred, y_true):
        return F.mse_loss(y_pred, y_true)
        
    
    def bert_init_params(self, module): 
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=0.02)
        elif isinstance(module, nn.BatchNorm1d):
            torch.nn.init.normal_(module.weight.data, mean=1, std=0.02)
            torch.nn.init.constant_(module.bias.data, 0)
            
class OmicsTransformer(BaseEmbed):
    '''
        Transformer encoder for proteomics and metabolomics
    
    '''
    def __init__(self,input_dim,hidden_dims,num_head,num_layers,dropout_prob,lr,mask_ratio,
                 weight_decay,monitor):
        super().__init__()
        
        self.hparams.input_dim = input_dim
        self.hparams.hidden_dims = hidden_dims
        self.hparams.num_head = num_head
        self.hparams.num_layers = num_layers
        self.hparams.dropout_prob = dropout_prob
        self.hparams.lr = lr
        self.hparams.mask_ratio = mask_ratio
        self.hparams.weight_decay = weight_decay
        self.hparams.monitor = monitor
        self.save_hyperparameters()
        
        self.expression_proj = nn.Linear(1,hidden_dims)
        self.omics_emb = nn.Embedding(input_dim,hidden_dims) # embed per feature
        
        transformer_layer = nn.TransformerEncoderLayer(
            batch_first=True,
            d_model=hidden_dims,
            nhead=num_head,
            dim_feedforward=hidden_dims * 2,
            dropout=dropout_prob  
        ) 
        self.transformer = nn.TransformerEncoder(transformer_layer, num_layers=num_layers)
        
        self.linear = nn.Sequential(
            nn.Linear(hidden_dims, hidden_dims // 2),
            nn.ReLU(),
            nn.Linear(hidden_dims // 2, hidden_dims //4),
            nn.ReLU(),
            nn.Linear(hidden_dims // 4, 1)
        )
        
        self.apply(self.bert_init_params)
        
    def training_step(self,batch,batch_index):
        
        batch = batch[0]
        
        batch_masked, mask_indices, nan_indices = mask_matrix(batch, mask_ratio=self.hparams.mask_ratio)
        
        x = batch_masked.unsqueeze(-1)
        x = self.expression_proj(x)
        x = x + self.omics_emb.weight
        x = self.transformer(x)
        x = self.linear(x)
        x = x.squeeze(-1) 
        
        loss = self.mse_loss(x, 
                             batch
                             )
        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, 
             sync_dist=True, batch_size=len(batch[0]))
        
        return loss
    
    def validation_step(self,batch,batch_index):
        
        batch = batch[0]
        
        batch_masked, mask_indices, nan_indices = mask_matrix(batch, mask_ratio=self.hparams.mask_ratio)
        
        x = batch_masked.unsqueeze(-1)
        x = self.expression_proj(x)
        x = x + self.omics_emb.weight
        x = self.transformer(x)
        x = self.linear(x)
        x = x.squeeze(-1) 
        
        loss = self.mse_loss(x,
                             batch
                             )
        self.log('val_loss', loss, on_step=True, on_epoch=True, prog_bar=True, 
             sync_dist=True, batch_size=len(batch[0]))
        return x
    
    @torch.no_grad()
    def inference(self, data: pd.DataFrame,batch_size:int=128):
        
        if isinstance(data, pd.DataFrame):
            dataset = TensorDataset(torch.tensor(data.values).float())
        else:
            dataset = TensorDataset(torch.tensor(data).float())
            
        
        with torch.no_grad():
            self.eval()
            
            preds = []
            
            data_loader = DataLoader(dataset,
                                    batch_size=batch_size,
                                    shuffle=False,
                                    num_workers=0,
                                    persistent_workers=False,
                                    )
            
            for batch in tqdm(data_loader, desc="Embedding"):
                batch = batch[0]
                batch = batch.to(self.device)
                batch = batch.unsqueeze(-1)
                batch = self.expression_proj(batch)
                batch = batch + self.omics_emb.weight
                batch = self.transformer(batch)
                embedding = batch.mean(dim=1)
                preds.append(embedding.cpu().numpy())
                
            preds = np.concatenate(preds, axis=0)
            
            if isinstance(data, pd.DataFrame):
                preds = pd.DataFrame(preds,index=data.index,)
            
        return preds
    
class Autoencoder(pl.LightningModule):
    """
    Autoencoder (AE) or Variational Autoencoder (VAE) for tabular data embedding.  

    This module is built on PyTorch Lightning and supports both deterministic autoencoders (AE) 
    and variational autoencoders (VAE). It can be used to train models that learn to reconstruct 
    partially masked input data and impute missing values at inference time.  
    """
    
    def __init__(self,
                 model_type: str,
                 input_dim: int,
                 hidden_dims: List[int],
                 activation: str,
                 use_batch_norm: bool,
                 dropout: float,
                 weight_decay: float,
                 lr: float,   
                 init_method: str,
                 optim_monitor: str = 'val_loss',
                 ):
        super().__init__()
        
        self.hparams.activation = activation
        self.hparams.use_batch_norm = use_batch_norm
        self.hparams.dropout = dropout
        self.hparams.lr = lr
        self.hparams.weight_decay = weight_decay
        self.hparams.init_method = init_method
        
        self.hparams.optim_monitor = optim_monitor
        self.model_type = model_type
        
        self.hparams.module_name = self.__class__.__name__
        
        self.hparams_dict = {
            "module_name": self.__class__.__name__,
            "input_dim": input_dim,
            "hidden_dims": hidden_dims,
            "activation": activation,
            "use_batch_norm": use_batch_norm,
            "dropout": dropout,
            "weight_decay": weight_decay,
            "lr": lr,
            "init_method": init_method,
            "model_type": model_type,
            "optim_monitor": optim_monitor,
        }
        
        self.save_hyperparameters(self.hparams_dict)
        ## initialize model
        if model_type == 'AE':
            self.model = AE(input_dim=input_dim, hidden_dims=hidden_dims, 
                            activation=activation, use_batch_norm=use_batch_norm, dropout=dropout)
        elif model_type == 'VAE':
            self.model = VAE(input_dim=input_dim, hidden_dims=hidden_dims, 
                             activation=activation, use_batch_norm=use_batch_norm, dropout=dropout)
        else:
            raise ValueError(f"Unsupported model_type: {model_type}. Choose 'AE' or 'VAE'.")
        
        ## initialize weights
        init_model(self.model, init_method)
        
        ## loss function
        self.mse_loss = nn.MSELoss(reduction='sum')

    def _vae_loss(self, pred, target, mu, logvar):
        """Compute VAE loss = reconstruction loss + KL divergence loss
        KL divergence loss is computed with free bits.
        """
        # Reconstruction loss
        recon_loss = self.mse_loss(pred, target)
        
        # KL divergence loss
        kld_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        
        recon_loss = recon_loss / len(target)
        kld_loss = kld_loss / len(target)
        
        return recon_loss, kld_loss
    
    def _mse_loss(self, pred, target):
        # batch-wise mse loss
        loss = self.mse_loss(pred, target)

        return loss / len(target)


    def training_step(self, batch, batch_idx,):
        batch = batch[0]        
        pred = self.model(batch)
        
        ## loss
        if self.model_type == 'AE':
            
            loss = self._mse_loss(pred, batch)
        
        elif self.model_type == 'VAE':
            reconstructed, mu, logvar = pred
            recon_loss, kld_loss = self._vae_loss(reconstructed, batch, 
                                  mu, logvar)
            self.log('train_recon_loss', recon_loss, on_step=True, on_epoch=True, prog_bar=False, 
                 sync_dist=True, batch_size=len(batch))
            self.log('train_kld_loss', kld_loss, on_step=True, on_epoch=True, prog_bar=False, 
                 sync_dist=True, batch_size=len(batch))
            loss = recon_loss + kld_loss
        
        else:
            raise ValueError(f"Unsupported model_type: {self.model_type}. Choose 'AE' or 'VAE'.")
        
        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, 
             sync_dist=True, batch_size=len(batch))
        
        return loss
    
    def validation_step(self, batch, batch_idx,):
        batch = batch[0]        
        pred = self.model(batch)
        
        ## loss
        if self.model_type == 'AE':
            loss = self._mse_loss(pred, batch)
        
        elif self.model_type == 'VAE':
            reconstructed, mu, logvar = pred

            recon_loss, kld_loss = self._vae_loss(reconstructed,
                                  batch, mu, logvar)
            loss = recon_loss + kld_loss
            self.log('val_recon_loss', recon_loss, on_step=False, on_epoch=True, prog_bar=False, 
                 sync_dist=True, batch_size=len(batch))
            self.log('val_kld_loss', kld_loss, on_step=False, on_epoch=True, prog_bar=False, 
                 sync_dist=True, batch_size=len(batch))

        else:
            raise ValueError(f"Unsupported model_type: {self.model_type}. Choose 'AE' or 'VAE'.")
        
        self.log('val_loss', loss, on_step=False, on_epoch=True, prog_bar=True, 
             sync_dist=True, batch_size=len(batch))

        return loss

    
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr,
                                     weight_decay=self.hparams.weight_decay) ## add weight decay
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer,
            mode='min',
            factor=0.5,
            patience=5,
        )
        
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "monitor": self.hparams.optim_monitor,
                "frequency": 1,
                "interval": "epoch",
            },
        }
        
    @torch.no_grad()
    def inference(self, data: pd.DataFrame,batch_size: int = 10) -> pd.DataFrame:
        """Impute missing values in the input DataFrame.
        
        Args:
            data (pd.DataFrame): Input data with missing values (NaNs).
            batch_size (int): Batch size for DataLoader.
        
        Returns:
            pd.DataFrame: Imputed DataFrame with missing values filled.
        """
        if isinstance(data, pd.DataFrame):
            dataset = TensorDataset(torch.tensor(data.values).float())
        else:
            dataset = TensorDataset(torch.tensor(data).float())

        with torch.no_grad():
            self.eval()
            
            preds = []
            
            data_loader = DataLoader(dataset,
                                    batch_size=batch_size,
                                    shuffle=False,
                                    num_workers=0,
                                    persistent_workers=False)

            for batch in tqdm(data_loader, desc="Embedding"):
                batch = batch[0]
                batch = batch.to(self.device)
                
                if self.model_type == "AE":
                    latent = self.model.encode(batch)
                elif self.model_type == "VAE":
                    mu, logvar = self.model.encode(batch)
                    latent = self.model.reparameterize(mu, logvar)
                else:
                    raise ValueError(f"Unsupported model_type: {self.model_type}. Choose 'AE' or 'VAE'.")
                
                preds.append(latent.cpu().numpy())
                
            
            preds = np.concatenate(preds, axis=0)
            if isinstance(data, pd.DataFrame):
                preds = pd.DataFrame(preds, index=data.index,)
            return preds

class TransformerAE(BaseEmbed):
    '''
    Transformer-based Autoencoder for Tabular Data (Proteomics/Metabolomics).
    
    Structure:
      [Input: Batch x n_features] 
           ↓ (Feature Embedding + Value Projection)
      [Seq: Batch x n_features x hidden_dim]
           ↓ (Transformer Encoder)
      [Latent Token Aggregation] -> [Latent z: Batch x latent_dim]
           ↓ (Project back to sequence)
      [Seq: Batch x n_features x hidden_dim]
           ↓ (Transformer Decoder)
      [Output: Batch x n_features]
    '''
    def __init__(self, 
                 input_dim: int,          # Number of features (proteins)
                 hidden_dims: int,        # d_model for Transformer
                 latent_dim: int,         # Bottleneck dimension (z)
                 num_head: int,
                 num_encoder_layers: int,
                 num_decoder_layers: int,
                 dropout_prob: float,
                 lr: float,
                 weight_decay: float,
                 monitor: str = 'val_loss',
                 mask_ratio: float = 0.0): # Set 0 for pure AE, >0 for Denoising AE
        
        super().__init__()
        self.save_hyperparameters()
        
        # --- Shared Components ---
        # Feature embeddings (learnable vector for each protein ID)
        self.feature_embeddings = nn.Embedding(input_dim, hidden_dims)
        # Projects the float scalar value of a protein to a vector
        self.value_proj = nn.Linear(1, hidden_dims)
        
        # --- Encoder ---
        encoder_layer = nn.TransformerEncoderLayer(
            batch_first=True,
            d_model=hidden_dims,
            nhead=num_head,
            dim_feedforward=hidden_dims * 2,
            dropout=dropout_prob,
            activation='gelu'
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        
        # Compression to latent space
        # We use a "CLS"-like token approach or simple flattening. 
        # Here, to keep it flexible, we project the flattened sequence or mean pool.
        # Let's use a learnable Readout token (CLS) appended to input, or Mean Pooling.
        # Approach: Mean Pooling over feature dimension -> Linear -> Latent
        self.to_latent = nn.Sequential(
            nn.Linear(hidden_dims, hidden_dims // 2),
            nn.ReLU(),
            nn.Linear(hidden_dims // 2, latent_dim)
        )

        # --- Decoder ---
        # Project latent z back to the sequence size (Batch x n_features x hidden_dim)
        self.from_latent = nn.Sequential(
            nn.Linear(latent_dim, hidden_dims // 2),
            nn.ReLU(),
            nn.Linear(hidden_dims // 2, hidden_dims),
        )
        
        decoder_layer = nn.TransformerEncoderLayer( # We use EncoderLayer architecture for the decoder block as well (Self-attention)
            batch_first=True,
            d_model=hidden_dims,
            nhead=num_head,
            dim_feedforward=hidden_dims * 2,
            dropout=dropout_prob,
            activation='gelu'
        )
        self.decoder = nn.TransformerEncoder(decoder_layer, num_layers=num_decoder_layers)
        
        # Final prediction head: Hidden Dim -> Scalar Value
        self.readout_head = nn.Linear(hidden_dims, 1)

        self.apply(self.bert_init_params)

    def forward(self, x):
        # x shape: [Batch, n_features]
        batch_size, n_features = x.shape
        device = x.device
        
        # --- Input Embedding ---
        # 1. Project values: [B, F] -> [B, F, 1] -> [B, F, H]
        val_emb = self.value_proj(x.unsqueeze(-1))
        
        # 2. Add feature identity embeddings
        # feature_indices: [0, 1, ..., F-1]
        feature_idx = torch.arange(n_features, device=device).expand(batch_size, n_features)
        feat_emb = self.feature_embeddings(feature_idx)
        
        # Combined embedding
        x_emb = val_emb + feat_emb # [B, F, H]
        
        # --- Encoder ---
        enc_out = self.encoder(x_emb) # [B, F, H]
        
        # Aggregate to Latent Vector (Mean Pooling)
        # You could also use a [CLS] token strategy here
        pooled = enc_out.mean(dim=1) # [B, H]
        z = self.to_latent(pooled)   # [B, latent_dim]
        
        # --- Decoder ---
        # 1. Expand Latent back to Sequence
        # We need to broadcast z back to every feature position to condition generation
        z_proj = self.from_latent(z) # [B, H]
        z_expanded = z_proj.unsqueeze(1).expand(-1, n_features, -1) # [B, F, H]
        
        # 2. Add Feature embeddings again so decoder knows which position is which protein
        # Note: In a pure AE, we usually feed the latent + position info.
        dec_input = z_expanded + feat_emb 
        
        # 3. Apply Transformer Decoder Blocks
        dec_out = self.decoder(dec_input) # [B, F, H]
        
        # 4. Readout to scalar
        recon = self.readout_head(dec_out).squeeze(-1) # [B, F]
        
        return recon, z

    def training_step(self, batch, batch_idx):
        x = batch[0] # [Batch, Features]
        
        # Optional: Masking for Denoising Autoencoder behavior
        if self.hparams.mask_ratio > 0:
            x_corrupted, _, _ = mask_matrix(x, mask_ratio=self.hparams.mask_ratio)
        else:
            x_corrupted = x

        recon, z = self.forward(x_corrupted)
        
        # Reconstruction Loss
        loss = self.mse_loss(recon, x) # Compare against original unmasked x
        
        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x = batch[0]
        recon, z = self.forward(x)
        
        loss = self.mse_loss(recon, x)
        self.log('val_loss', loss, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
        return loss

    @torch.no_grad()
    def inference(self, data: pd.DataFrame, batch_size: int = 128, return_recon=False):
        '''
        Returns latent embeddings by default. 
        Set return_recon=True to get reconstructed data (imputation).
        '''
        if isinstance(data, pd.DataFrame):
            dataset = TensorDataset(torch.tensor(data.values).float())
        else:
            dataset = TensorDataset(torch.tensor(data).float())
        
        self.eval()
        results = []
        
        data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
        
        for batch in tqdm(data_loader, desc="Inference"):
            batch = batch[0].to(self.device)
            recon, z = self.forward(batch)
            
            if return_recon:
                results.append(recon.cpu().numpy())
            else:
                results.append(z.cpu().numpy())
                
        results = np.concatenate(results, axis=0)
        
        if isinstance(data, pd.DataFrame):
            results = pd.DataFrame(results, index=data.index)
            
        return results