spaceexpress 0.1.5__py3-none-any.whl

SpaceExpress/__init__.py

@@ -0,0 +1,4 @@
+from .utils import shortest_path, plot_embedding, jaccard_similarity, choose_k, plot_DSE
+from .preprocessing import select_hvg_after_outlier
+from .spaceexpress import train_SpaceExpress, train_SpaceExpress_multi
+from .spaceexpress_dse import SpaceExpress_DSE, summary_DSE

SpaceExpress/preprocessing.py

@@ -0,0 +1,109 @@
+"""Preprocessing helpers for sparse spatial-expression inputs."""
+
+from __future__ import annotations
+
+import anndata as ad
+import numpy as np
+import scanpy as sc
+import scipy.sparse as sp
+
+
+def select_hvg_after_outlier(
+    adata_list,
+    n_top_genes=200,
+    z_threshold=4.0,
+):
+    """Remove pooled high-expression outliers before batch-aware HVG selection.
+
+    The input AnnData objects must already be normalized and log-transformed.
+    For every common gene, values greater than or equal to pooled
+    ``mean + z_threshold * sample_sd`` are set to zero. Seurat HVGs are then
+    selected from the cleaned matrices with each AnnData treated as a batch.
+
+    Returns
+    -------
+    cleaned_hvg : list[AnnData]
+        Copies restricted to the same sorted HVGs and containing cleaned values.
+    hvg : list[str]
+        Sorted selected gene names.
+    diagnostics : dict
+        Counts and nonzero-prevalence summaries for reproducibility.
+    """
+    if len(adata_list) < 2:
+        raise ValueError("At least two AnnData objects are required.")
+    if n_top_genes < 1:
+        raise ValueError("n_top_genes must be positive.")
+    if z_threshold <= 0:
+        raise ValueError("z_threshold must be positive.")
+
+    common = sorted(set.intersection(*(set(item.var_names) for item in adata_list)))
+    if len(common) < n_top_genes:
+        raise ValueError(
+            f"Need at least {n_top_genes} common genes, found {len(common)}."
+        )
+    cleaned = [item[:, common].copy() for item in adata_list]
+    matrices = [
+        item.X.tocsr().astype(np.float64, copy=True)
+        if sp.issparse(item.X)
+        else sp.csr_matrix(np.asarray(item.X, dtype=np.float64))
+        for item in cleaned
+    ]
+
+    pooled = sp.vstack(matrices, format="csr")
+    n_obs = pooled.shape[0]
+    means = np.asarray(pooled.sum(axis=0)).ravel() / n_obs
+    sums_of_squares = np.asarray(pooled.power(2).sum(axis=0)).ravel()
+    variances = (sums_of_squares - n_obs * means**2) / max(n_obs - 1, 1)
+    thresholds = means + z_threshold * np.sqrt(np.maximum(variances, 0.0))
+
+    removed_entries = []
+    for item, matrix in zip(cleaned, matrices):
+        remove = matrix.data >= thresholds[matrix.indices]
+        removed_entries.append(int(remove.sum()))
+        matrix.data[remove] = 0.0
+        matrix.eliminate_zeros()
+        item.X = matrix.astype(np.float32)
+
+    combined = ad.concat(
+        cleaned,
+        label="_hvg_batch",
+        keys=[f"sample{i}" for i in range(len(cleaned))],
+        index_unique="-",
+    )
+    sc.pp.highly_variable_genes(
+        combined,
+        n_top_genes=n_top_genes,
+        flavor="seurat",
+        batch_key="_hvg_batch",
+    )
+    hvg = sorted(combined.var_names[combined.var["highly_variable"]].tolist())
+    if len(hvg) != n_top_genes:
+        raise RuntimeError(f"Expected {n_top_genes} HVGs, found {len(hvg)}.")
+
+    result = [item[:, hvg].copy() for item in cleaned]
+    prevalence = []
+    for item in result:
+        counts = np.asarray((item.X > 0).sum(axis=0)).ravel()
+        prevalence.append(
+            {
+                "min": int(counts.min()),
+                "q25": float(np.quantile(counts, 0.25)),
+                "median": float(np.median(counts)),
+                "q75": float(np.quantile(counts, 0.75)),
+                "max": int(counts.max()),
+            }
+        )
+    diagnostics = {
+        "method": "pooled_mean_plus_sample_sd",
+        "z_threshold": float(z_threshold),
+        "common_genes_before_hvg": len(common),
+        "removed_expression_entries": removed_entries,
+        "selected_hvg_nonzero_spots": prevalence,
+    }
+    for item in result:
+        item.uns["spaceexpress_preprocessing"] = {
+            "mean_sd_outliers_removed": True,
+            "z_threshold": float(z_threshold),
+            "stage": "before_hvg",
+        }
+    return result, hvg, diagnostics

SpaceExpress/spaceexpress.py

@@ -0,0 +1,300 @@
+import pickle, random
+import numpy as np
+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
+from sklearn.neighbors import kneighbors_graph
+from tqdm import tqdm
+import copy
+import scipy
+import scanpy as sc
+
+def get_index (data):
+    n_cells = data.shape[0]
+    mask = torch.triu(torch.ones((n_cells, n_cells)), diagonal=1)        
+    index = torch.nonzero(mask, as_tuple=False)
+    return index
+
+def loaded_paths(file_path):
+    loaded_paths = {}
+
+    with open(file_path, 'rb') as handle:
+        while True:
+            try:
+                data = pickle.load(handle)
+                loaded_paths.update(data)
+            except EOFError:
+                break  
+    shortest_path = [loaded_paths[i][0] for i in range(len(loaded_paths))]
+    shortest_path = np.array(shortest_path)
+    return shortest_path
+
+class SpaceExpress(nn.Module):
+    def __init__(self, input_dim, hidden_dim, latent_dim):
+        super(SpaceExpress, self).__init__()
+
+        self.encoder = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.BatchNorm1d(hidden_dim), 
+            nn.ReLU(),
+            nn.Linear(hidden_dim, latent_dim),
+        )
+
+    def forward(self, x):
+        encoded = self.encoder(x)
+        return encoded
+
+class KKLoss:
+    def __init__(self, device, shortest_path):
+        self.device = device
+        self.shortest_path = torch.tensor(shortest_path, dtype=torch.float32).to(device)
+
+    def __call__(self, p, idx):
+        d = self.shortest_path
+        mask = d != 0
+        d_inv = torch.where(mask, 1 / d**2, torch.zeros_like(d))
+        p = p.unsqueeze(1) - p.unsqueeze(0)
+        p = torch.sum(torch.abs(p),dim=-1)
+        diff = p - d
+        out = d_inv * mask * diff.pow(2)
+        out = out[idx[:, 0], idx[:, 1]]
+        out = torch.sum(out)   
+        return out
+
+def get_avg_neighbor(pos, count, k):
+    A = kneighbors_graph(pos, k, mode='connectivity', include_self=False)
+    neighbor = A.dot(count) / k
+    neighbor_dense = neighbor.toarray() if hasattr(neighbor, "toarray") else neighbor
+    count_dense = count.toarray() if hasattr(count, "toarray") else count
+    out = np.concatenate((count_dense, neighbor_dense), axis=1)
+    return out
+
+def train_SpaceExpress(adata, shortest_file_path, device = None, epochs = 10000, lr = 0.01, hid_dim = 32, emb_dim = 8, 
+                       patience = 100, random_seed = 42, batch_size = 256, num_hvg = 1000, save_model = False):
+    """
+    Train SpaceExpress model.
+    
+    Parameters:
+    adata (AnnData): Anndata object
+    shortest_path_PATH (str): Path to the shortest path file
+    device (str): Device to use
+    lr (float): Learning rate
+    hid_dim (int): Hidden dimension
+    emb_dim (int): Embedding dimension
+    patience (int): Patience for early stopping
+    random_seed (int): Random seed
+    batch_size (int): Batch size
+    
+    Returns:
+    AnnData: Anndata object with SpaceExpress embedding
+    """    
+
+    # Set random seed
+    torch.manual_seed(random_seed)
+    torch.cuda.manual_seed(random_seed)
+    torch.cuda.manual_seed_all(random_seed) 
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    np.random.seed(random_seed)
+    random.seed(random_seed)
+
+    # Set device
+    if device == 'mps':
+        device = torch.device("mps")
+    else:
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    
+    # Preprocessing
+    percentile_95 = [np.percentile(adata.X[:,i], 95) for i in range(adata.shape[1])] 
+    adata.X = np.array([np.clip(adata.X[:,i], a_min=None, a_max=percentile_95[i]) for i in range(adata.shape[1])]).T 
+    
+    is_count_data = np.all(np.equal(np.mod(adata.X, 1), 0))
+    if adata.shape[1] > num_hvg:
+        if is_count_data == True:
+            sc.pp.highly_variable_genes(adata, n_top_genes=num_hvg, flavor='seurat_v3', subset = True) 
+        else:
+            sc.pp.highly_variable_genes(adata, n_top_genes=num_hvg, subset = True) 
+    sc.pp.scale(adata) 
+
+    # Get data
+    data = adata.X 
+    pos = adata.obsm['spatial']
+    shortest_path = loaded_paths(shortest_file_path)      
+    
+    data = torch.tensor(data, dtype=torch.float).to(device) 
+    model = SpaceExpress(data.shape[1], hid_dim, emb_dim).to(device)
+    optimizer = optim.Adam(model.parameters(), lr=lr)
+    index = get_index(data)
+    kkloss = KKLoss(device, shortest_path)
+
+    best_train_loss = float('inf')
+    patience_counter = 0
+    
+    # Training loop
+    print('Start training...')
+    with tqdm(total=epochs) as pbar:
+        for epoch in tqdm(range(epochs)):
+            total_loss = 0
+            
+            idx = index[torch.randperm(index.size(0))[:batch_size**2]]
+            optimizer.zero_grad()
+            embeddings = model(data.to(device))
+            loss = kkloss(embeddings, idx)        
+            loss.backward()
+            optimizer.step()
+            total_loss += loss.item()
+        
+            pbar.set_description(f"Epoch {epoch+1}/{epochs}")
+            pbar.set_postfix(loss=f"{total_loss:.4f}")
+            pbar.update(1)  
+                
+            if total_loss < best_train_loss:
+                best_train_loss = total_loss
+                patience_counter = 0
+                best_model = copy.deepcopy(model)
+            else:
+                patience_counter += 1
+
+            if patience_counter >= patience:
+                break
+
+    # Switch to evaluation mode
+    best_model.eval()
+    emb = best_model(data)
+    emb = emb.cpu().detach().numpy()
+    
+    if save_model == True:
+        return emb, best_model
+    
+    return emb
+
+def train_SpaceExpress_multi(adata_list_input, shortest_file_path_list, device = None, epochs = 10000, lr = 0.01, hid_dim = 32, 
+                             emb_dim = 4, patience = 100, random_seed = 42, batch_size = 256, num_hvg = 1000, save_model = False):
+    """
+    Train SpaceExpress model.
+    
+    Parameters:
+    adata_list (list): List of AnnData objects
+    shortest_path_PATH (str): Path to the shortest path file
+    device (str): Device to use
+    lr (float): Learning rate
+    hid_dim (int): Hidden dimension
+    emb_dim (int): Embedding dimension
+    patience (int): Patience for early stopping
+    random_seed (int): Random seed
+    batch_size (int): Batch size
+    
+    Returns:
+    adata_list (list): List of AnnData objects with SpaceExpress embeddings
+    """    
+    
+    # Set random seed
+    torch.manual_seed(random_seed)
+    torch.cuda.manual_seed(random_seed)
+    torch.cuda.manual_seed_all(random_seed) 
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    np.random.seed(random_seed)
+    random.seed(random_seed)
+    
+    adata_list = [i.copy() for i in adata_list_input]
+    
+    # Set device
+    if device == 'mps':
+        device = torch.device("mps")
+    else:
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    
+    # Preprocessing
+    for i, adata in enumerate(adata_list):
+        # Check if adata.X is sparse and convert if necessary
+        if isinstance(adata.X, (scipy.sparse.csr_matrix, scipy.sparse.csc_matrix)):
+            adata.X = adata.X.toarray()    
+
+        percentile_95 = [np.percentile(adata.X[:,i], 95) for i in range(adata.shape[1])] 
+        adata.X = np.array([np.clip(adata.X[:,i], a_min=None, a_max=percentile_95[i]) for i in range(adata.shape[1])]).T 
+        is_count_data = np.all(np.equal(np.mod(adata.X, 1), 0))
+        if adata.shape[1] > num_hvg:
+            if is_count_data == True:
+                sc.pp.highly_variable_genes(adata, n_top_genes=num_hvg, flavor='seurat_v3', subset = True) 
+            else:
+                sc.pp.highly_variable_genes(adata, n_top_genes=num_hvg, subset = True) 
+        sc.pp.scale(adata)
+        adata_list[i] = adata
+    
+    intersecting_genes = set.intersection(*(set(adata.var_names) for adata in adata_list))
+    adata_list = [adata[:, list(intersecting_genes)] for adata in adata_list]
+    print(f'Size of the input data: {[adata_list[i].shape for i in range(len(adata_list))]}')
+    
+    # Get data
+    num_data = len(adata_list)
+    data_list = [adata.X for adata in adata_list]
+    pos_list = [adata.obsm['spatial'] for adata in adata_list]
+    shortest_path_list = [loaded_paths(i) for i in shortest_file_path_list]
+    
+    num_gene = data_list[0].shape[1]
+    print('number of genes:', num_gene)
+    for i in data_list: 
+        assert num_gene == i.shape[1], "The number of genes should be the same across all datasets"
+    
+    data_list = [torch.tensor(data, dtype=torch.float).to(device) for data in data_list]    
+    model = SpaceExpress(num_gene, hid_dim, emb_dim).to(device)
+    optimizer = optim.Adam(model.parameters(), lr=lr)
+      
+    index_list = [get_index(data) for data in data_list]
+    kkloss_list = [KKLoss(device, shortest_path) for shortest_path in shortest_path_list]
+
+    best_train_loss = float('inf')
+    patience_counter = 0
+    
+    # Training loop
+    print('Start training...')
+    with tqdm(total=epochs) as pbar:
+        for epoch in tqdm(range(epochs)):
+            total_loss = 0
+            reg_loss = 0
+            
+            loss = 0
+            for i in range(num_data):
+                data = data_list[i]
+                index = index_list[i]
+                kkloss = kkloss_list[i]
+                
+                optimizer.zero_grad()
+                embeddings = model(data.to(device))
+                idx = index[torch.randperm(index.size(0))[:batch_size**2]]
+                loss += kkloss(embeddings, idx)      
+                
+            loss.backward()
+            optimizer.step()
+            total_loss += loss.item()        
+
+            pbar.set_description(f"Epoch {epoch+1}/{epochs}")
+            pbar.set_postfix(loss=f"{total_loss:.4f}")
+            pbar.update(1)  # Update the progress bar once per epoch
+            
+            if total_loss < best_train_loss:
+                best_train_loss = total_loss
+                patience_counter = 0
+                best_model = copy.deepcopy(model)
+            else:
+                patience_counter += 1
+
+            if patience_counter >= patience:
+                break
+
+    # Switch to evaluation mode
+    best_model.eval()
+    out = []
+    for i in range(num_data):
+        data = data_list[i]
+        emb = best_model(data)
+        emb = emb.cpu().detach().numpy()
+        out.append(emb)
+        
+    if save_model == True:
+        return out, best_model
+    
+    return out