SURE-tools 1.0.10__tar.gz → 2.0.2__tar.gz

{sure_tools-1.0.10 → sure_tools-2.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: SURE-tools
-Version: 1.0.10
+Version: 2.0.2
 Summary: Succinct Representation of Single Cells
 Home-page: https://github.com/ZengFLab/SURE
 Author: Feng Zeng

{sure_tools-1.0.10 → sure_tools-2.0.2}/SURE/SURE.py

@@ -97,6 +97,7 @@ class SURE(nn.Module):
                  input_size: int,
                  codebook_size: int = 200,
                  cell_factor_size: int = 0,
+                 cell_factor_names: list = None,
                  supervised_mode: bool = False,
                  z_dim: int = 10,
                  z_dist: Literal['normal','studentt','laplacian','cauchy','gumbel'] = 'normal',
@@ -134,6 +135,7 @@ class SURE(nn.Module):
         self.post_layer_fct = post_layer_fct
         self.post_act_fct = post_act_fct
         self.hidden_layer_activation = hidden_layer_activation
+        self.cell_factor_names = cell_factor_names
         
         self.codebook_weights = None
 
@@ -232,8 +234,10 @@ class SURE(nn.Module):
         )
 
         if self.cell_factor_size>0:
-            self.cell_factor_effect = MLP(
-                    [self.cell_factor_size] + self.decoder_hidden_layers + [self.latent_dim],
+            self.cell_factor_effect = nn.ModuleList()
+            for i in np.arange(self.cell_factor_size):
+                self.cell_factor_effect.append(MLP(
+                    [self.z_dim+1] + hidden_sizes + [self.z_dim],
                     activation=activate_fct,
                     output_activation=None,
                     post_layer_fct=post_layer_fct,
@@ -241,6 +245,7 @@ class SURE(nn.Module):
                     allow_broadcast=self.allow_broadcast,
                     use_cuda=self.use_cuda,
                 )
+                )
             
         self.decoder_concentrate = MLP(
                 [self.latent_dim] + self.decoder_hidden_layers + [self.input_size],
@@ -444,8 +449,14 @@ class SURE(nn.Module):
                 zns = pyro.sample('zn', dist.Gumbel(zn_loc, zn_scale).to_event(1))
 
             if self.cell_factor_size>0:
-                zus = self.cell_factor_effect(us)
-                zs = zns + zus
+                #zus = self.decoder_undesired([zns,us])
+                zus = None
+                for i in np.arange(self.cell_factor_size):
+                    if i==0:
+                        zus = self.cell_factor_effect[i]([zns,us[:,i].reshape(-1,1)])
+                    else:
+                        zus = zus + self.cell_factor_effect[i]([zns,us[:,i].reshape(-1,1)])
+                zs = zns+zus
             else:
                 zs = zns
 
@@ -634,8 +645,14 @@ class SURE(nn.Module):
                     zns = pyro.sample('zn', dist.Gumbel(zn_loc, zn_scale).to_event(1), obs=embeds)
 
             if self.cell_factor_size>0:
-                zus = self.cell_factor_effect(us)
-                zs = zus + zns
+                #zus = self.decoder_undesired([zns,us])
+                zus = None
+                for i in np.arange(self.cell_factor_size):
+                    if i==0:
+                        zus = self.cell_factor_effect[i]([zns,us[:,i].reshape(-1,1)])
+                    else:
+                        zus = zus + self.cell_factor_effect[i]([zns,us[:,i].reshape(-1,1)])
+                zs = zns+zus
             else:
                 zs = zns
 
@@ -796,6 +813,42 @@ class SURE(nn.Module):
         A = np.concatenate(A)
         return A
     
+    def _cell_move(self, xs, factor_idx, perturb):
+        zns = self.encoder_zn(xs)
+        if type(factor_idx) == str:
+            factor_idx = np.where(self.cell_factor_names==factor_idx)
+        
+        if perturb.ndim==2:
+            ms = self.cell_factor_effect[factor_idx]([zns, perturb])
+        else:
+            ms = self.cell_factor_effect[factor_idx]([zns, perturb.reshape(-1,1)])
+            
+        return ms 
+
+    def get_cell_move(self, 
+                             xs, 
+                             factor_idx,
+                             perturb,
+                             batch_size: int = 1024):
+        """
+        Return cells' changes in the latent space induced by specific perturbation of a factor
+
+        """
+        xs = self.preprocess(xs)
+        xs = convert_to_tensor(xs, device=self.get_device())
+        dataset = CustomDataset(xs)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+        Z = []
+        with tqdm(total=len(dataloader), desc='', unit='batch') as pbar:
+            for X_batch, _ in dataloader:
+                zns = self._cell_move(X_batch, factor_idx, perturb)
+                Z.append(tensor_to_numpy(zns))
+                pbar.update(1)
+
+        Z = np.concatenate(Z)
+        return Z
+    
     def preprocess(self, xs, threshold=0):
         if self.loss_func == 'bernoulli':
             ad = sc.AnnData(xs)

{sure_tools-1.0.10 → sure_tools-2.0.2}/SURE/__init__.py

@@ -6,5 +6,6 @@ from . import codebook
 from . import SURE
 from . import SURE2
 from . import atac
+from . import flow 
 
-__all__ = ['SURE','SURE2', 'atac', 'utils', 'codebook']
+__all__ = ['SURE','SURE2', 'flow', 'atac', 'utils', 'codebook']

sure_tools-2.0.2/SURE/flow/__init__.py

@@ -0,0 +1 @@
+from .quiver import plot_quiver

sure_tools-2.0.2/SURE/flow/quiver.py

@@ -0,0 +1,52 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from sklearn.decomposition import PCA
+import scanpy as sc
+
+def plot_quiver_high_dim(z_points, delta_z, method='umap', figsize=(6,4), dpi=200):
+    """
+    从高维潜在空间选择2个维度进行quiver可视化
+    """
+    if method == 'variance':
+        # 方法1: 选择方差最大的2个维度
+        variances = np.var(z_points, axis=0)
+        dims = np.argsort(variances)[-2:]  # 选择方差最大的两个维度
+        dim_names = [f'z[{d}]' for d in dims]
+        
+    elif method == 'pca':
+        # 方法2: 使用PCA的前两个主成分
+        pca = PCA(n_components=2)
+        z_2d = pca.fit_transform(z_points)
+        delta_z_2d = pca.transform(z_points + delta_z) - z_2d
+        dim_names = ['PC1', 'PC2']
+        
+    elif method == 'manual':
+        # 方法3: 手动选择感兴趣的维度
+        dims = [0, 1]  # 选择前两个维度
+        z_2d = z_points[:, dims]
+        delta_z_2d = delta_z[:, dims]
+        dim_names = [f'z[{d}]' for d in dims]
+        
+    elif method == 'umap':
+        ad = sc.AnnData(np.vstack([z_points, z_points+delta_z]))
+        sc.pp.neighbors(ad)
+        sc.tl.umap(ad)
+        z_2d = ad[:z_points.shape[0]].obsm['X_umap']
+        delta_z_2d = ad[z_points.shape[0]:] - z_2d
+        dim_names = ['UMAP1', 'UMAP2']
+    
+    # 绘制quiver图
+    plt.figure(figsize=figsize, dpi=dpi)
+    plt.quiver(z_2d[:, 0], z_2d[:, 1], 
+               delta_z_2d[:, 0], delta_z_2d[:, 1],
+               angles='xy', scale_units='xy', scale=1,
+               color='blue', alpha=0.6, width=0.005)
+    
+    plt.scatter(z_2d[:, 0], z_2d[:, 1], c='gray', alpha=0.5, s=10)
+    plt.xlabel(dim_names[0])
+    plt.ylabel(dim_names[1])
+    plt.title(f"Latent Space Movement ({method} projection)")
+    plt.grid(alpha=0.3)
+    plt.show()
+    
+    return z_2d, delta_z_2d

{sure_tools-1.0.10 → sure_tools-2.0.2}/SURE_tools.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: SURE-tools
-Version: 1.0.10
+Version: 2.0.2
 Summary: Succinct Representation of Single Cells
 Home-page: https://github.com/ZengFLab/SURE
 Author: Feng Zeng

{sure_tools-1.0.10 → sure_tools-2.0.2}/SURE_tools.egg-info/SOURCES.txt

@@ -2,7 +2,6 @@ LICENSE
 README.md
 setup.py
 SURE/SURE.py
-SURE/SURE2.py
 SURE/__init__.py
 SURE/assembly/__init__.py
 SURE/assembly/assembly.py
@@ -11,6 +10,8 @@ SURE/atac/__init__.py
 SURE/atac/utils.py
 SURE/codebook/__init__.py
 SURE/codebook/codebook.py
+SURE/flow/__init__.py
+SURE/flow/quiver.py
 SURE/utils/__init__.py
 SURE/utils/custom_mlp.py
 SURE/utils/queue.py

{sure_tools-1.0.10 → sure_tools-2.0.2}/setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
 
 setup(
     name='SURE-tools',
-    version='1.0.10',
+    version='2.0.2',
     description='Succinct Representation of Single Cells',
     long_description=long_description,
     long_description_content_type="text/markdown",