dataloaders/clonidine_data.py

import torch
import sys
sys.argv = ['']
from sklearn.preprocessing import StandardScaler
import pytorch_lightning as pl
from torch.utils.data import DataLoader
from lightning.pytorch.utilities.combined_loader import CombinedLoader
import pandas as pd
import numpy as np
from functools import partial
from scipy.spatial import cKDTree
from sklearn.cluster import KMeans
from torch.utils.data import TensorDataset

#from train.parsers_tahoe import parse_args
#args = parse_args()

class DrugResponseDataModule(pl.LightningDataModule):
    def __init__(self, args):
        super().__init__()
        self.save_hyperparameters()

        self.batch_size = args.batch_size
        self.max_dim = args.dim
        self.whiten = args.whiten
        self.split_ratios = args.split_ratios
        
        # Path to your combined data
        self.data_path = "/raid/st512/branchsbm/data/pca_and_leiden_labels.csv"
        self.num_timesteps = 2
        self.args = args
        self._prepare_data()

    def _prepare_data(self):
        df = pd.read_csv(self.data_path, comment='#')
        df = df.iloc[:, 1:]
        df = df.replace('', np.nan)
        pc_cols = df.columns[:50]
        for col in pc_cols:
            df[col] = pd.to_numeric(df[col], errors='coerce')
        leiden_dmso_col = 'leiden_DMSO_TF_0.0uM'
        leiden_clonidine_col = 'leiden_Clonidine (hydrochloride)_5.0uM'

        dmso_mask = df[leiden_dmso_col].notna()        # Has leiden value in DMSO column
        clonidine_mask = df[leiden_clonidine_col].notna()  # Has leiden value in Clonidine column
        
        dmso_data = df[dmso_mask].copy()
        clonidine_data = df[clonidine_mask].copy()
        
        top_clonidine_clusters = ['0.0', '4.0']
        
        x1_1_data = clonidine_data[clonidine_data[leiden_clonidine_col].astype(str) == top_clonidine_clusters[0]]
        x1_2_data = clonidine_data[clonidine_data[leiden_clonidine_col].astype(str) == top_clonidine_clusters[1]]
        x1_1_coords = x1_1_data[pc_cols].values 
        x1_2_coords = x1_2_data[pc_cols].values 
        x1_1_coords = x1_1_coords.astype(float)
        x1_2_coords = x1_2_coords.astype(float)

        target_size = min(len(x1_1_coords), len(x1_2_coords))
        
        # Sample endpoint clusters to target size
        np.random.seed(42)
        if len(x1_1_coords) > target_size:
            idx1 = np.random.choice(len(x1_1_coords), target_size, replace=False)
            x1_1_coords = x1_1_coords[idx1]
        
        if len(x1_2_coords) > target_size:
            idx2 = np.random.choice(len(x1_2_coords), target_size, replace=False)
            x1_2_coords = x1_2_coords[idx2]
        
        dmso_cluster_counts = dmso_data[leiden_dmso_col].value_counts()
        
        # DMSO
        largest_dmso_cluster = dmso_cluster_counts.index[0]
        dmso_cluster_data = dmso_data[dmso_data[leiden_dmso_col] == largest_dmso_cluster]
        
        dmso_coords = dmso_cluster_data[pc_cols].values
        
        # Random sampling from largest DMSO cluster to match target size
        np.random.seed(42)
        if len(dmso_coords) >= target_size:
            idx0 = np.random.choice(len(dmso_coords), target_size, replace=False)
            x0_coords = dmso_coords[idx0]
        else:
            # If largest cluster is smaller than target, use all of it and pad with other DMSO cells
            remaining_needed = target_size - len(dmso_coords)
            other_dmso_data = dmso_data[dmso_data[leiden_dmso_col] != largest_dmso_cluster]
            other_dmso_coords = other_dmso_data[pc_cols].values
            
            if len(other_dmso_coords) >= remaining_needed:
                idx_other = np.random.choice(len(other_dmso_coords), remaining_needed, replace=False)
                x0_coords = np.vstack([dmso_coords, other_dmso_coords[idx_other]])
            else:
                # Use all available DMSO cells and reduce target size
                all_dmso_coords = dmso_data[pc_cols].values
                target_size = min(target_size, len(all_dmso_coords))
                idx0 = np.random.choice(len(all_dmso_coords), target_size, replace=False)
                x0_coords = all_dmso_coords[idx0]
        
        # Also resample endpoint clusters to match final target size
        if len(x1_1_coords) > target_size:
            idx1 = np.random.choice(len(x1_1_coords), target_size, replace=False)
            x1_1_coords = x1_1_coords[idx1]
        
        if len(x1_2_coords) > target_size:
            idx2 = np.random.choice(len(x1_2_coords), target_size, replace=False)
            x1_2_coords = x1_2_coords[idx2]
        
        self.n_samples = target_size
        
        x0 = torch.tensor(x0_coords, dtype=torch.float32)
        x1_1 = torch.tensor(x1_1_coords, dtype=torch.float32)
        x1_2 = torch.tensor(x1_2_coords, dtype=torch.float32)
        
        self.coords_t0 = x0
        self.coords_t1 = torch.cat([x1_1, x1_2], dim=0)
        self.time_labels = np.concatenate([
            np.zeros(len(self.coords_t0)),    # t=0
            np.ones(len(self.coords_t1)),     # t=1
        ])
        
        split_index = int(target_size * self.split_ratios[0])
        
        if target_size - split_index < self.batch_size:
            split_index = target_size - self.batch_size
        
        train_x0 = x0[:split_index]
        val_x0 = x0[split_index:]
        train_x1_1 = x1_1[:split_index]
        val_x1_1 = x1_1[split_index:]
        train_x1_2 = x1_2[:split_index]
        val_x1_2 = x1_2[split_index:]
        
        self.val_x0 = val_x0
        
        train_x0_weights = torch.full((train_x0.shape[0], 1), fill_value=1.0)
        train_x1_1_weights = torch.full((train_x1_1.shape[0], 1), fill_value=0.5)
        train_x1_2_weights = torch.full((train_x1_2.shape[0], 1), fill_value=0.5)
        
        val_x0_weights = torch.full((val_x0.shape[0], 1), fill_value=1.0)
        val_x1_1_weights = torch.full((val_x1_1.shape[0], 1), fill_value=0.5)
        val_x1_2_weights = torch.full((val_x1_2.shape[0], 1), fill_value=0.5)
        
        self.train_dataloaders = {
            "x0": DataLoader(TensorDataset(train_x0, train_x0_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
            "x1_1": DataLoader(TensorDataset(train_x1_1, train_x1_1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
            "x1_2": DataLoader(TensorDataset(train_x1_2, train_x1_2_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
        }
        
        self.val_dataloaders = {
            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.batch_size, shuffle=False, drop_last=True),
            "x1_1": DataLoader(TensorDataset(val_x1_1, val_x1_1_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
            "x1_2": DataLoader(TensorDataset(val_x1_2, val_x1_2_weights), batch_size=self.batch_size, shuffle=True, drop_last=True),
        }
        
        all_coords = df[pc_cols].dropna().values.astype(float)
        self.dataset = torch.tensor(all_coords, dtype=torch.float32)
        self.tree = cKDTree(all_coords)

        self.test_dataloaders = {
            "x0": DataLoader(TensorDataset(val_x0, val_x0_weights), batch_size=self.val_x0.shape[0], shuffle=False, drop_last=False),
            "dataset": DataLoader(TensorDataset(self.dataset), batch_size=self.dataset.shape[0], shuffle=False, drop_last=False),
        }
        
        # Metric samples
        km_all = KMeans(n_clusters=3, random_state=0).fit(self.dataset.numpy())
        cluster_labels = km_all.labels_
        
        cluster_0_mask = cluster_labels == 0
        cluster_1_mask = cluster_labels == 1
        cluster_2_mask = cluster_labels == 2
        
        samples = self.dataset.cpu().numpy()
            
        cluster_0_data = samples[cluster_0_mask]
        cluster_1_data = samples[cluster_1_mask]
        cluster_2_data = samples[cluster_2_mask]
        
        self.metric_samples_dataloaders = [
            DataLoader(
                torch.tensor(cluster_2_data, dtype=torch.float32),
                batch_size=cluster_2_data.shape[0],
                shuffle=False,
                drop_last=False,
            ),
            DataLoader(
                torch.tensor(cluster_0_data, dtype=torch.float32),
                batch_size=cluster_0_data.shape[0],
                shuffle=False,
                drop_last=False,
            ),
            
            DataLoader(
                torch.tensor(cluster_1_data, dtype=torch.float32),
                batch_size=cluster_1_data.shape[0],
                shuffle=False,
                drop_last=False,
            ),
        ]
        
    def train_dataloader(self):
        combined_loaders = {
            "train_samples": CombinedLoader(self.train_dataloaders, mode="min_size"),
            "metric_samples": CombinedLoader(
                self.metric_samples_dataloaders, mode="min_size"
            ),
        }
        return CombinedLoader(combined_loaders, mode="max_size_cycle")

    def val_dataloader(self):
        combined_loaders = {
            "val_samples": CombinedLoader(self.val_dataloaders, mode="min_size"),
            "metric_samples": CombinedLoader(
                self.metric_samples_dataloaders, mode="min_size"
            ),
        }

        return CombinedLoader(combined_loaders, mode="max_size_cycle")

    
    
    def test_dataloader(self):
        combined_loaders = {
            "test_samples": CombinedLoader(self.test_dataloaders, mode="min_size"),
            "metric_samples": CombinedLoader(
                self.metric_samples_dataloaders, mode="min_size"
            ),
        }

        return CombinedLoader(combined_loaders, mode="max_size_cycle")

    def get_manifold_proj(self, points):
        """Adapted for 2D cell data - uses local neighborhood averaging instead of plane fitting"""
        return partial(self.local_smoothing_op, tree=self.tree, dataset=self.dataset)

    @staticmethod
    def local_smoothing_op(x, tree, dataset, k=10, temp=1e-3):
        """
        Apply local smoothing based on k-nearest neighbors in the full dataset
        This replaces the plane projection for 2D manifold regularization
        """
        points_np = x.detach().cpu().numpy()
        _, idx = tree.query(points_np, k=k)
        nearest_pts = dataset[idx]  # Shape: (batch_size, k, 2)
        
        # Compute weighted average of neighbors
        dists = (x.unsqueeze(1) - nearest_pts).pow(2).sum(-1, keepdim=True)
        weights = torch.exp(-dists / temp)
        weights = weights / weights.sum(dim=1, keepdim=True)
        
        # Weighted average of neighbors
        smoothed = (weights * nearest_pts).sum(dim=1)
        
        # Blend original point with smoothed version
        alpha = 0.3  # How much smoothing to apply
        return (1 - alpha) * x + alpha * smoothed
    
    def get_timepoint_data(self):
        """Return data organized by timepoints for visualization"""
        return {
            't0': self.coords_t0,
            't1': self.coords_t1, 
            'time_labels': self.time_labels
        }
    
def get_datamodule():
    datamodule = DrugResponseDataModule(args)
    datamodule.setup(stage="fit")   
    return datamodule