gsMap 1.65__py3-none-any.whl → 1.67__py3-none-any.whl

gsMap/GNN_VAE/adjacency_matrix.py

@@ -1,95 +1,75 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-Created on Tue Jul  4 21:31:27 2023
-
-@author: songliyang
-"""
 import numpy as np
 import pandas as pd
 import scipy.sparse as sp
-import sklearn.neighbors
+from sklearn.neighbors import NearestNeighbors
 import torch
 
-
-def Cal_Spatial_Net(adata, n_neighbors=5, verbose=True):
-    """\
-    Construct the spatial neighbor networks.
-    """
-    #- 
+def cal_spatial_net(adata, n_neighbors=5, verbose=True):
+    """Construct the spatial neighbor network."""
     if verbose:
         print('------Calculating spatial graph...')
-    coor = pd.DataFrame(adata.obsm['spatial'])
-    coor.index = adata.obs.index
-    #- 
-    nbrs = sklearn.neighbors.NearestNeighbors(n_neighbors=n_neighbors).fit(coor)
-    #- 
-    distances, indices = nbrs.kneighbors(coor, return_distance=True)
-    KNN_list = []
-    for it in range(indices.shape[0]):
-        KNN_list.append(pd.DataFrame(zip([it]*indices[it].shape[0], indices[it], distances[it])))
-    #- 
-    KNN_df = pd.concat(KNN_list)
-    KNN_df.columns = ['Cell1', 'Cell2', 'Distance']
-    #- 
-    Spatial_Net = KNN_df.copy()
-    Spatial_Net = Spatial_Net.loc[Spatial_Net['Distance']>0,]
-    id_cell_trans = dict(zip(range(coor.shape[0]), np.array(coor.index), ))
-    Spatial_Net['Cell1'] = Spatial_Net['Cell1'].map(id_cell_trans)
-    Spatial_Net['Cell2'] = Spatial_Net['Cell2'].map(id_cell_trans)
-    #- 
-    return Spatial_Net
-
+    coor = pd.DataFrame(adata.obsm['spatial'], index=adata.obs.index)
+    nbrs = NearestNeighbors(n_neighbors=n_neighbors).fit(coor)
+    distances, indices = nbrs.kneighbors(coor)
+    n_cells, n_neighbors = indices.shape
+    cell_indices = np.arange(n_cells)
+    cell1 = np.repeat(cell_indices, n_neighbors)
+    cell2 = indices.flatten()
+    distance = distances.flatten()
+    knn_df = pd.DataFrame({'Cell1': cell1, 'Cell2': cell2, 'Distance': distance})
+    knn_df = knn_df[knn_df['Distance'] > 0].copy()
+    cell_id_map = dict(zip(cell_indices, coor.index))
+    knn_df['Cell1'] = knn_df['Cell1'].map(cell_id_map)
+    knn_df['Cell2'] = knn_df['Cell2'].map(cell_id_map)
+    return knn_df
 
 def sparse_mx_to_torch_sparse_tensor(sparse_mx):
     """Convert a scipy sparse matrix to a torch sparse tensor."""
     sparse_mx = sparse_mx.tocoo().astype(np.float32)
-    indices = torch.from_numpy(np.vstack((sparse_mx.row, sparse_mx.col)).astype(np.int64))
+    indices = torch.from_numpy(
+        np.vstack((sparse_mx.row, sparse_mx.col)).astype(np.int64)
+    )
     values = torch.from_numpy(sparse_mx.data)
     shape = torch.Size(sparse_mx.shape)
     return torch.sparse.FloatTensor(indices, values, shape)
 
-
 def preprocess_graph(adj):
+    """Symmetrically normalize the adjacency matrix."""
     adj = sp.coo_matrix(adj)
     adj_ = adj + sp.eye(adj.shape[0])
-    rowsum = np.array(adj_.sum(1))
-    degree_mat_inv_sqrt = sp.diags(np.power(rowsum, -0.5).flatten())
+    rowsum = np.array(adj_.sum(1)).flatten()
+    degree_mat_inv_sqrt = sp.diags(np.power(rowsum, -0.5))
     adj_normalized = adj_.dot(degree_mat_inv_sqrt).transpose().dot(degree_mat_inv_sqrt).tocoo()
     return sparse_mx_to_torch_sparse_tensor(adj_normalized)
 
-
-
-def Construct_Adjacency_Matrix(adata,Params, verbose=True):
-    # Construct the neighbor graph 
-    Spatial_Net = Cal_Spatial_Net(adata, n_neighbors=Params.n_neighbors)
-    #- 
+def construct_adjacency_matrix(adata, params, verbose=True):
+    """Construct the adjacency matrix from spatial data."""
+    spatial_net = cal_spatial_net(adata, n_neighbors=params.n_neighbors, verbose=verbose)
     if verbose:
-        print('The graph contains %d edges, %d cells.' %(Spatial_Net.shape[0], adata.n_obs))
-        print('%.2f neighbors per cell on average.' %(Spatial_Net.shape[0]/adata.n_obs))
-    #-  
-    cells = np.array(adata.obs.index)
-    cells_id_tran = dict(zip(cells, range(cells.shape[0])))
-    #- 
-    G_df = Spatial_Net.copy()
-    G_df['Cell1'] = G_df['Cell1'].map(cells_id_tran)
-    G_df['Cell2'] = G_df['Cell2'].map(cells_id_tran)
-    #- 
-    if Params.weighted_adj:
-        distance_normalized = G_df.Distance/(max(G_df.Distance)+1)
-        adj_org = sp.coo_matrix((np.exp(-distance_normalized**2/(2)), (G_df['Cell1'], G_df['Cell2'])), shape=(adata.n_obs, adata.n_obs))
+        num_edges = spatial_net.shape[0]
+        num_cells = adata.n_obs
+        print(f'The graph contains {num_edges} edges, {num_cells} cells.')
+        print(f'{num_edges / num_cells:.2f} neighbors per cell on average.')
+    cell_ids = {cell: idx for idx, cell in enumerate(adata.obs.index)}
+    spatial_net['Cell1'] = spatial_net['Cell1'].map(cell_ids)
+    spatial_net['Cell2'] = spatial_net['Cell2'].map(cell_ids)
+    if params.weighted_adj:
+        distance_normalized = spatial_net['Distance'] / (spatial_net['Distance'].max() + 1)
+        weights = np.exp(-0.5 * distance_normalized ** 2)
+        adj_org = sp.coo_matrix(
+            (weights, (spatial_net['Cell1'], spatial_net['Cell2'])),
+            shape=(adata.n_obs, adata.n_obs)
+        )
     else:
-        adj_org = sp.coo_matrix((np.ones(G_df.shape[0]), (G_df['Cell1'], G_df['Cell2'])), shape=(adata.n_obs, adata.n_obs))
-    #- 
-    adj_m1 = adj_org
-    adj_norm_m1 = preprocess_graph(adj_m1)
-    adj_label_m1 = adj_m1 + sp.eye(adj_m1.shape[0])
-    norm_m1 = adj_m1.shape[0] * adj_m1.shape[0] / float((adj_m1.shape[0] * adj_m1.shape[0] - adj_m1.sum()) * 2)
-    #- 
+        adj_org = sp.coo_matrix(
+            (np.ones(spatial_net.shape[0]), (spatial_net['Cell1'], spatial_net['Cell2'])),
+            shape=(adata.n_obs, adata.n_obs)
+        )
+    adj_norm = preprocess_graph(adj_org)
+    norm_value = adj_org.shape[0] ** 2 / ((adj_org.shape[0] ** 2 - adj_org.sum()) * 2)
     graph_dict = {
         "adj_org": adj_org,
-        "adj_norm": adj_norm_m1,
-        "norm_value": norm_m1
+        "adj_norm": adj_norm,
+        "norm_value": norm_value
     }
-    #- 
     return graph_dict

gsMap/GNN_VAE/model.py

@@ -1,87 +1,89 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-Created on Mon Jul  3 11:42:44 2023
-
-@author: songliyang
-"""
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch_geometric.nn import GATConv
 
-
 def full_block(in_features, out_features, p_drop):
-    return nn.Sequential(nn.Linear(in_features, out_features),
-                         nn.BatchNorm1d(out_features),
-                         nn.ELU(),
-                         nn.Dropout(p=p_drop))
+    return nn.Sequential(
+        nn.Linear(in_features, out_features),
+        nn.BatchNorm1d(out_features),
+        nn.ELU(),
+        nn.Dropout(p=p_drop)
+    )
 
-
-class GNN(nn.Module):
-    def __init__(self, in_features, out_features, dr=0, act=F.relu,heads=1):
+class GATModel(nn.Module):
+    def __init__(self, input_dim, params, num_classes=1):
         super().__init__()
-        self.conv1 = GATConv(in_features, out_features,heads)
-        self.act = act
-        self.dr = dr
-    #-
-    def forward(self, x, edge_index):
-        out = self.conv1(x, edge_index)
-        out = self.act(out)
-        out = F.dropout(out, self.dr, self.training)
-        return out
-
-
-    
-class GNN_VAE_Model(nn.Module):
-    def __init__(self, input_dim,params,num_classes=1):
-        super(GNN_VAE_Model, self).__init__()
         self.var = params.var
         self.num_classes = num_classes
-            
+        self.params = params
+
         # Encoder
-        self.encoder = nn.Sequential()
-        self.encoder.add_module('encoder_L1', full_block(input_dim, params.feat_hidden1, params.p_drop))
-        self.encoder.add_module('encoder_L2', full_block(params.feat_hidden1, params.feat_hidden2, params.p_drop))
-        
-        # GNN (GAT)
-        self.gn1 = GNN(params.feat_hidden2, params.gat_hidden1, params.p_drop, act=F.relu,heads = params.nheads)
-        self.gn2 = GNN(params.gat_hidden1*params.nheads, params.gat_hidden2, params.p_drop, act=lambda x: x)
-        self.gn3 = GNN(params.gat_hidden1*params.nheads, params.gat_hidden2, params.p_drop, act=lambda x: x) 
-        
+        self.encoder = nn.Sequential(
+            full_block(input_dim, params.feat_hidden1, params.p_drop),
+            full_block(params.feat_hidden1, params.feat_hidden2, params.p_drop)
+        )
+
+        # GAT Layers
+        self.gat1 = GATConv(
+            in_channels=params.feat_hidden2,
+            out_channels=params.gat_hidden1,
+            heads=params.nheads,
+            dropout=params.p_drop
+        )
+        self.gat2 = GATConv(
+            in_channels=params.gat_hidden1 * params.nheads,
+            out_channels=params.gat_hidden2,
+            heads=1,
+            concat=False,
+            dropout=params.p_drop
+        )
+        if self.var:
+            self.gat3 = GATConv(
+                in_channels=params.gat_hidden1 * params.nheads,
+                out_channels=params.gat_hidden2,
+                heads=1,
+                concat=False,
+                dropout=params.p_drop
+            )
+
         # Decoder
-        self.decoder = nn.Sequential()
-        self.decoder.add_module('decoder_L1', full_block(params.gat_hidden2, params.feat_hidden2, params.p_drop))
-        self.decoder.add_module('decoder_L2', full_block(params.feat_hidden2, params.feat_hidden1, params.p_drop))
-        self.decoder.add_module('decoder_output', nn.Sequential(nn.Linear(params.feat_hidden1, input_dim)))
-        
-        # Cluster
-        self.cluster = nn.Sequential()
-        self.cluster.add_module('cluster_L1', full_block(params.gat_hidden2, params.feat_hidden2, params.p_drop))
-        self.cluster.add_module('cluster_output', nn.Linear(params.feat_hidden2, self.num_classes))
-           
-    def encode(self, x, adj):
-        feat_x = self.encoder(x)
-        hidden1 = self.gn1(feat_x, adj)
-        mu = self.gn2(hidden1, adj)
+        self.decoder = nn.Sequential(
+            full_block(params.gat_hidden2, params.feat_hidden2, params.p_drop),
+            full_block(params.feat_hidden2, params.feat_hidden1, params.p_drop),
+            nn.Linear(params.feat_hidden1, input_dim)
+        )
+
+        # Clustering Layer
+        self.cluster = nn.Sequential(
+            full_block(params.gat_hidden2, params.feat_hidden2, params.p_drop),
+            nn.Linear(params.feat_hidden2, self.num_classes)
+        )
+
+    def encode(self, x, edge_index):
+        x = self.encoder(x)
+        x = self.gat1(x, edge_index)
+        x = F.relu(x)
+        x = F.dropout(x, p=self.params.p_drop, training=self.training)
+
+        mu = self.gat2(x, edge_index)
         if self.var:
-            logvar = self.gn3(hidden1, adj)
+            logvar = self.gat3(x, edge_index)
             return mu, logvar
         else:
             return mu, None
-     
+
     def reparameterize(self, mu, logvar):
         if self.training and logvar is not None:
-            std = torch.exp(logvar)
+            std = torch.exp(0.5 * logvar)
             eps = torch.randn_like(std)
-            return eps.mul(std).add_(mu)
+            return eps * std + mu
         else:
             return mu
-    
-    def forward(self, x, adj):
-        mu, logvar = self.encode(x, adj)
-        gnn_z = self.reparameterize(mu, logvar)
-        x_reconstructed = self.decoder(gnn_z)
-        pred_label = F.softmax(self.cluster(gnn_z),dim=1)
-        return pred_label, x_reconstructed, gnn_z, mu, logvar
+
+    def forward(self, x, edge_index):
+        mu, logvar = self.encode(x, edge_index)
+        z = self.reparameterize(mu, logvar)
+        x_reconstructed = self.decoder(z)
+        pred_label = F.softmax(self.cluster(z), dim=1)
+        return pred_label, x_reconstructed, z, mu, logvar

gsMap/GNN_VAE/train.py

@@ -1,97 +1,86 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-Created on Tue Jul  4 19:58:58 2023
-
-@author: songliyang
-"""
+import logging
 import time
 
 import torch
+import torch.nn.functional as F
 from progress.bar import Bar
 
-from gsMap.GNN_VAE.model import GNN_VAE_Model
+from gsMap.GNN_VAE.model import GATModel
+
+logger = logging.getLogger(__name__)
 
 
 def reconstruction_loss(decoded, x):
-    loss_fn = torch.nn.MSELoss()
-    loss = loss_fn(decoded, x)
-    return loss
+    """Compute the mean squared error loss."""
+    return F.mse_loss(decoded, x)
 
 
 def label_loss(pred_label, true_label):
-    loss_fn = torch.nn.CrossEntropyLoss()
-    loss = loss_fn(pred_label, true_label)
-    return loss
-
-
-class Model_Train:
-    def __init__(self, node_X, graph_dict, params, label=None):
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        torch.cuda.empty_cache()
+    """Compute the cross-entropy loss."""
+    return F.cross_entropy(pred_label, true_label)
 
+class ModelTrainer:
+    def __init__(self, node_x, graph_dict, params, label=None):
+        """Initialize the ModelTrainer with data and hyperparameters."""
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.params = params
-        self.device = device
         self.epochs = params.epochs
-        self.node_X = torch.FloatTensor(node_X.copy()).to(device)
-        self.adj_norm = graph_dict["adj_norm"].to(device).coalesce()
+        self.node_x = torch.FloatTensor(node_x).to(self.device)
+        self.adj_norm = graph_dict["adj_norm"].to(self.device).coalesce()
         self.label = label
         self.num_classes = 1
-        
-        if not self.label is None:
+
+        if self.label is not None:
             self.label = torch.tensor(self.label).to(self.device)
-            self.num_classes = len(self.label.unique())
-        
-        # Set Model 
-        self.model = GNN_VAE_Model(self.params.feat_cell,self.params,self.num_classes).to(device)
-        self.optimizer = torch.optim.Adam(params = list(self.model.parameters()),
-                                          lr = self.params.gat_lr, weight_decay = self.params.gcn_decay)               
-    
-    # Train    
+            self.num_classes = len(torch.unique(self.label))
+
+        # Set up the model
+        self.model = GATModel(self.params.feat_cell, self.params, self.num_classes).to(self.device)
+        self.optimizer = torch.optim.Adam(
+            self.model.parameters(),
+            lr=self.params.gat_lr,
+            weight_decay=self.params.gcn_decay
+        )
+
     def run_train(self):
+        """Train the model."""
         self.model.train()
         prev_loss = float('inf')
-        
-        bar = Bar('GAT-AE model train:', max = self.epochs)
-        bar.check_tty = False 
+        bar = Bar('GAT-AE model train:', max=self.epochs)
+        bar.check_tty = False
+
+        logger.info('Start training...')
         for epoch in range(self.epochs):
             start_time = time.time()
-            self.model.train()
             self.optimizer.zero_grad()
-            pred_label, de_feat, latent_z, mu, logvar = self.model(self.node_X, self.adj_norm)
-            loss_rec = reconstruction_loss(de_feat, self.node_X)
-            
-            # Check whether annotation was provided
-            if not self.label is None:
+            pred_label, de_feat, latent_z, mu, logvar = self.model(self.node_x, self.adj_norm)
+            loss_rec = reconstruction_loss(de_feat, self.node_x)
+
+            if self.label is not None:
                 loss_pre = label_loss(pred_label, self.label)
-                loss = (self.params.rec_w * loss_rec) + (self.params.label_w * loss_pre)
+                loss = self.params.rec_w * loss_rec + self.params.label_w * loss_pre
             else:
                 loss = loss_rec
-                
+
             loss.backward()
             self.optimizer.step()
-            
-            # Update process
-            end_time = time.time()
-            batch_time = end_time - start_time
-            
-            
-            bar_str = '{} / {} | Left time: {batch_time:.2f} mins| Loss: {loss:.4f}'
-            bar.suffix = bar_str.format(epoch + 1,self.epochs,
-                                        batch_time = batch_time * (self.epochs - epoch) / 60, loss=loss.item())
+
+            batch_time = time.time() - start_time
+            left_time = batch_time * (self.epochs - epoch - 1) / 60  # in minutes
+
+            bar.suffix = f'{epoch + 1} / {self.epochs} | Left time: {left_time:.2f} mins | Loss: {loss.item():.4f}'
             bar.next()
-            
-            # Check convergence
+
             if abs(loss.item() - prev_loss) <= self.params.convergence_threshold and epoch >= 200:
-                print('\nConvergence reached. Training stopped.')
+                logger.info('\nConvergence reached. Training stopped.')
                 break
 
             prev_loss = loss.item()
-            
         bar.finish()
-    #-    
+
     def get_latent(self):
+        """Retrieve the latent representation from the model."""
         self.model.eval()
-        pred, de_fea, latent_z, mu, logvar = self.model(self.node_X, self.adj_norm)
-        latent_z = latent_z.data.cpu().numpy()
-        return latent_z
+        with torch.no_grad():
+            _, _, latent_z, _, _ = self.model(self.node_x, self.adj_norm)
+        return latent_z.cpu().numpy()

gsMap/__init__.py

@@ -2,4 +2,4 @@
 Genetics-informed pathogenic spatial mapping
 '''
 
-__version__ = '1.65'
+__version__ = '1.67'

gsMap/config.py

@@ -55,7 +55,8 @@ def add_find_latent_representations_args(parser):
     add_shared_args(parser)
     parser.add_argument('--input_hdf5_path', required=True, type=str, help='Path to the input HDF5 file.')
     parser.add_argument('--annotation', required=True, type=str, help='Name of the annotation in adata.obs to use.')
-    parser.add_argument('--data_layer', required=True, type=str, help='Data layer for gene expression (e.g., "counts", "log1p").')
+    parser.add_argument('--data_layer', type=str, default='counts', required=True,
+                        help='Data layer for gene expression (e.g., "count", "counts", "log1p").')
     parser.add_argument('--epochs', type=int, default=300, help='Number of training epochs.')
     parser.add_argument('--feat_hidden1', type=int, default=256, help='Neurons in the first hidden layer.')
     parser.add_argument('--feat_hidden2', type=int, default=128, help='Neurons in the second hidden layer.')
@@ -66,7 +67,6 @@ def add_find_latent_representations_args(parser):
     parser.add_argument('--n_neighbors', type=int, default=11, help='Number of neighbors for GAT.')
     parser.add_argument('--n_comps', type=int, default=300, help='Number of principal components for PCA.')
     parser.add_argument('--weighted_adj', action='store_true', help='Use weighted adjacency in GAT.')
-    parser.add_argument('--var', action='store_true', help='Enable variance calculations.')
     parser.add_argument('--convergence_threshold', type=float, default=1e-4, help='Threshold for convergence.')
     parser.add_argument('--hierarchically', action='store_true', help='Enable hierarchical latent representation finding.')
 
@@ -236,8 +236,8 @@ def add_run_all_mode_args(parser):
                         help='Path to the input spatial transcriptomics data (H5AD format).')
     parser.add_argument('--annotation', type=str, required=True,
                         help='Name of the annotation in adata.obs to use.')
-    parser.add_argument('--data_layer', type=str, default='X',
-                        help='Data layer of h5ad for gene expression (e.g., "counts", "log1p", "X").')
+    parser.add_argument('--data_layer', type=str, default='counts', required=True,
+                        help='Data layer for gene expression (e.g., "count", "counts", "log1p").')
 
     # GWAS Data Parameters
     parser.add_argument('--trait_name', type=str, help='Name of the trait for GWAS analysis (required if sumstats_file is provided).')