preprocess.py

# -*- coding: utf-8 -*-
import torch
import numpy as np
import h5py
import pandas as pd
import scipy.sparse as sp
from pathlib import Path
from sklearn.neighbors import NearestNeighbors

try:
    import scanpy as sc
    import anndata
except ImportError:
    sc = None
    anndata = None

from utils import prepare_dual_view_data


######################################## Graph Augmentation ########################################

class GraphAugmentation:
    """Graph augmentation module for dual-view generation"""

    def __init__(self, n_neighbors=15, edge_drop_rate=0.1, ppr_alpha=0.2, noise_std=0.1):
        self.n_neighbors = n_neighbors
        self.edge_drop_rate = edge_drop_rate
        self.ppr_alpha = ppr_alpha
        self.noise_std = noise_std

    def build_knn_graph(self, X, device=None):
        """Build KNN graph using cosine similarity"""
        if isinstance(X, np.ndarray):
            X = torch.from_numpy(X)
        if device is None:
            device = X.device
        X = X.to(device)
        n_samples = X.size(0)

        X_norm = torch.nn.functional.normalize(X, p=2, dim=1)
        sim_matrix = torch.mm(X_norm, X_norm.t())

        k = min(self.n_neighbors + 1, n_samples)
        topk_values, topk_indices = torch.topk(sim_matrix, k=k, dim=1)

        neighbor_indices = topk_indices[:, 1:]
        neighbor_values = topk_values[:, 1:]

        adj = torch.zeros((n_samples, n_samples), device=device, dtype=torch.float32)
        row_indices = torch.arange(n_samples, device=device).unsqueeze(1).expand(-1, k-1)
        adj[row_indices, neighbor_indices] = neighbor_values
        adj[neighbor_indices, row_indices] = neighbor_values

        return adj

    def normalize_adjacency(self, adj):
        """Symmetric normalization: D^(-1/2) @ A @ D^(-1/2)"""
        n = adj.size(0)
        device = adj.device
        adj = adj + torch.eye(n, device=device)
        degrees = adj.sum(dim=1)
        degrees[degrees == 0] = 1
        deg_inv_sqrt = torch.pow(degrees, -0.5)
        deg_inv_sqrt[torch.isinf(deg_inv_sqrt)] = 0.
        deg_inv_sqrt_mat = torch.diag(deg_inv_sqrt)
        return deg_inv_sqrt_mat @ adj @ deg_inv_sqrt_mat

    def edge_perturbed_graph(self, adj, similarity):
        """Drop edges with lowest similarity"""
        edge_mask = adj > 0
        edge_similarities = similarity[edge_mask]
        n_edges = edge_mask.sum().item()
        n_edges_to_drop = int(n_edges * self.edge_drop_rate)

        if n_edges_to_drop > 0:
            threshold = torch.kthvalue(edge_similarities, n_edges_to_drop)[0]
            adj_perturbed = adj.clone()
            drop_mask = (similarity <= threshold) & edge_mask
            adj_perturbed[drop_mask] = 0
        else:
            adj_perturbed = adj.clone()

        return adj_perturbed

    def pagerank_diffusion(self, adj):
        """PageRank diffusion: PPR = α * (I - (1-α) * A_norm)^(-1)"""
        n = adj.size(0)
        device = adj.device

        adj_self = adj + torch.eye(n, device=device)
        d = adj_self.sum(dim=1)
        d_inv_sqrt = torch.pow(d, -0.5)
        d_inv_sqrt[torch.isinf(d_inv_sqrt)] = 0.0
        norm_adj = d_inv_sqrt.view(-1, 1) * adj_self * d_inv_sqrt.view(1, -1)

        identity = torch.eye(n, device=device)
        ppr_matrix = self.ppr_alpha * torch.linalg.inv(
            identity - (1 - self.ppr_alpha) * norm_adj
        )
        ppr_matrix = (ppr_matrix + ppr_matrix.T) / 2

        threshold = ppr_matrix.mean() + 0.5 * ppr_matrix.std()
        ppr_matrix[ppr_matrix < threshold] = 0

        return ppr_matrix

    def __call__(self, X_features, X_raw):
        """Generate dual-view graphs"""
        device = X_features.device

        adj_base = self.build_knn_graph(X_features, device=device)

        X_norm = X_features / (X_features.norm(dim=1, keepdim=True) + 1e-8)
        similarity = X_norm @ X_norm.T

        adj_r = self.edge_perturbed_graph(adj_base, similarity)
        adj_r = self.normalize_adjacency(adj_r)

        adj_d = self.pagerank_diffusion(adj_base)

        noise = torch.normal(mean=1.0, std=self.noise_std, size=X_features.shape,
                            device=device, dtype=X_features.dtype)
        X_aug = torch.clamp(X_features * noise, min=0.0)

        return {'X_aug': X_aug, 'adj_r': adj_r, 'adj_d': adj_d, 'X_raw': X_raw}


######################################## Dataset Loader ########################################

class DatasetLoader:
    """Dataset loader for scRNA-seq data"""

    def __init__(self, dataset_name, data_dir='datasets', k=15,
                 edge_drop_rate=0.1, ppr_alpha=0.15, noise_std=0.1):
        self.dataset_name = dataset_name
        self.data_dir = data_dir
        self.k = k
        self.edge_drop_rate = edge_drop_rate
        self.ppr_alpha = ppr_alpha
        self.noise_std = noise_std
        self.data = {}

    def preprocess_scanpy(self, X, y, n_top_genes=2000):
        """Preprocess with Scanpy"""
        if torch.is_tensor(X):
            X = X.cpu().numpy()
        if torch.is_tensor(y):
            y = y.cpu().numpy()

        adata = anndata.AnnData(X=X)
        adata.obs['label'] = y
        adata.layers["counts"] = adata.X.copy()

        sc.pp.filter_genes(adata, min_cells=3)
        sc.pp.filter_cells(adata, min_counts=1)
        sc.pp.normalize_total(adata, target_sum=1e4)
        sc.pp.log1p(adata)

        if n_top_genes and n_top_genes < adata.n_vars:
            sc.pp.highly_variable_genes(adata, n_top_genes=n_top_genes,
                                        subset=True, flavor='seurat')

        X_features = adata.X
        if sp.issparse(X_features):
            X_features = X_features.toarray()

        X_raw = adata.layers["counts"]
        if sp.issparse(X_raw):
            X_raw = X_raw.toarray()

        y = adata.obs['label'].values

        return (torch.from_numpy(X_features).float(),
                torch.from_numpy(X_raw).float(),
                torch.from_numpy(y).long())

    def _load_h5_dataset(self, data_path, device):
        """Load data from .h5 file"""
        def read_clean(data):
            if hasattr(data, 'dtype') and data.dtype.type is np.bytes_:
                if data.size == 1:
                    return data.item().decode('utf-8')
                return np.array([item.decode('utf-8') for item in data.flat])
            return data.item() if data.size == 1 else data

        def dict_from_group(group):
            d = {}
            for key in group:
                if isinstance(group[key], h5py.Group):
                    d[key] = dict_from_group(group[key])
                else:
                    d[key] = read_clean(group[key][:])
            return d

        with h5py.File(data_path, "r") as f:
            obs = None
            if "obs" in f:
                obs = pd.DataFrame(dict_from_group(f["obs"]),
                                  index=read_clean(f["obs_names"][:]))

            if "exprs" in f:
                exprs_handle = f["exprs"]
                if isinstance(exprs_handle, h5py.Group):
                    mat = sp.csr_matrix(
                        (exprs_handle["data"][:],
                         exprs_handle["indices"][:],
                         exprs_handle["indptr"][:]),
                        shape=exprs_handle["shape"][:]
                    ).toarray().astype(np.float32)
                else:
                    mat = exprs_handle[:].astype(np.float32)
            elif "X" in f:
                mat = f["X"][:].astype(np.float32)
            else:
                raise ValueError("No expression data found")

            cell_name = None
            if obs is not None and "cell_type1" in obs.columns:
                cell_name = np.array(obs["cell_type1"])
            elif obs is not None and "celltype" in obs.columns:
                cell_name = np.array(obs["celltype"])
            elif "Y" in f:
                y = f["Y"][:]
            else:
                y = None

        if cell_name is not None:
            _, y = np.unique(cell_name, return_inverse=True)
        elif y is not None:
            if y.min() == 1:
                y = y - 1
        else:
            y = np.zeros(mat.shape[0], dtype=np.int64)

        X = torch.from_numpy(mat).float().to(device)
        y = torch.from_numpy(y).long().to(device)

        metadata = {
            'dataset_name': self.dataset_name,
            'n_cells': X.shape[0],
            'n_genes': X.shape[1],
            'n_clusters': len(torch.unique(y))
        }
        return X, y, metadata

    def load(self, device=None, n_top_genes=2000):
        """Load and preprocess dataset"""
        if device is None:
            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        print("Loading data...")
        data_path = Path(self.data_dir) / f"{self.dataset_name}.h5"
        X_raw, y, metadata = self._load_h5_dataset(data_path, device)

        print("Preprocessing with Scanpy...")
        X_features, X_raw, y = self.preprocess_scanpy(X_raw, y, n_top_genes)
        X_features = X_features.to(device)
        X_raw = X_raw.to(device)
        y = y.to(device)

        metadata['n_genes'] = X_features.shape[1]
        metadata['n_cells'] = X_features.shape[0]

        print("Building dual-view graphs...")
        aug = GraphAugmentation(
            n_neighbors=self.k,
            edge_drop_rate=self.edge_drop_rate,
            ppr_alpha=self.ppr_alpha,
            noise_std=self.noise_std
        )
        aug_data = aug(X_features=X_features, X_raw=X_raw)

        print("Computing graph embeddings...")
        X_aug_np = aug_data['X_aug'].cpu().numpy()
        dual_view_data = prepare_dual_view_data(
            X_aug_np, aug_data['adj_r'], aug_data['adj_d'], k=self.k
        )

        self.data = {
            'X': X_raw, 'X2': X_raw, 'X_features': X_features,
            'X_aug': aug_data['X_aug'], 'y': y,
            'A': aug_data['adj_r'], 'A1': aug_data['adj_r'],
            'A2': aug_data['adj_d'], 'metadata': metadata,
            **dual_view_data
        }
        return self.data

    def get_data(self):
        return self.data