SURE-tools 2.1.83__py3-none-any.whl → 2.1.85__py3-none-any.whl

SURE/PerturbFlow.py

@@ -74,6 +74,7 @@ class PerturbFlow(nn.Module):
                  use_cuda: bool = True,
                  seed: int = 42,
                  zero_bias: bool|list = True,
+                 enumrate: bool = False,
                  dtype = torch.float32, # type: ignore
                  ):
         super().__init__()
@@ -102,6 +103,7 @@ class PerturbFlow(nn.Module):
         else:
             self.use_bias = [not zero_bias] * self.cell_factor_size
         #self.use_bias = not zero_bias
+        self.enumerate = enumerate
         
         self.codebook_weights = None
 
@@ -204,7 +206,7 @@ class PerturbFlow(nn.Module):
             for i in np.arange(self.cell_factor_size):
                 if self.use_bias[i]:
                     self.cell_factor_effect.append(MLP(
-                        [self.code_size+1] + self.decoder_hidden_layers + [self.latent_dim],
+                        [self.latent_dim+1] + self.decoder_hidden_layers + [self.latent_dim],
                         activation=activate_fct,
                         output_activation=None,
                         post_layer_fct=post_layer_fct,
@@ -215,7 +217,7 @@ class PerturbFlow(nn.Module):
                     )
                 else:
                     self.cell_factor_effect.append(ZeroBiasMLP(
-                        [self.code_size+1] + self.decoder_hidden_layers + [self.latent_dim],
+                        [self.latent_dim+1] + self.decoder_hidden_layers + [self.latent_dim],
                         activation=activate_fct,
                         output_activation=None,
                         post_layer_fct=post_layer_fct,
@@ -429,7 +431,10 @@ class PerturbFlow(nn.Module):
                 zns = pyro.sample('zn', dist.Gumbel(zn_loc, zn_scale).to_event(1))
 
             if self.cell_factor_size>0:
-                zus = self._total_effects(ns, us)
+                if self.enumerate:
+                    zus = self._total_effects(zn_loc, us)
+                else:
+                    zus = self._total_effects(zns, us)
                 zs = zns+zus
             else:
                 zs = zns
@@ -631,7 +636,10 @@ class PerturbFlow(nn.Module):
                 #        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)])
-                zus = self._total_effects(ns, us)
+                if self.enumerate:
+                    zus = self._total_effects(zn_loc, us)
+                else:
+                    zus = self._total_effects(zns, us)
                 zs = zns+zus
             else:
                 zs = zns
@@ -842,8 +850,7 @@ class PerturbFlow(nn.Module):
     
     def _cell_response(self, xs, factor_idx, perturb):
         #zns,_ = self.encoder_zn(xs)    
-        #zns,_ = self._get_basal_embedding(xs)   
-        zns = self._soft_assignments(xs)
+        zns,_ = self._get_basal_embedding(xs)   
         if perturb.ndim==2:
             ms = self.cell_factor_effect[factor_idx]([zns, perturb])
         else:
@@ -877,8 +884,7 @@ class PerturbFlow(nn.Module):
         return Z
     
     def get_metacell_response(self, factor_idx, perturb):
-        #zs = self._get_codebook()
-        zs = self._get_codebook_identity()
+        zs = self._get_codebook()
         ps = convert_to_tensor(perturb, device=self.get_device())         
         ms = self.cell_factor_effect[factor_idx]([zs,ps])
         return tensor_to_numpy(ms)

SURE/utils/custom_mlp.py

@@ -239,7 +239,10 @@ class ZeroBiasMLP(nn.Module):
     def forward(self, x):
         y = self.mlp(x)
         mask = torch.zeros_like(y)
-        mask[x[1][:,0]>0,:] = 1
+        if len(y.shape)==2:
+            mask[x[1][:,0]>0,:] = 1
+        elif len(y.shape)==3:
+            mask[:,x[1][:,0]>0,:] = 1
         return y*mask
         
 

{sure_tools-2.1.83.dist-info → sure_tools-2.1.85.dist-info}/METADATA

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

{sure_tools-2.1.83.dist-info → sure_tools-2.1.85.dist-info}/RECORD

@@ -1,4 +1,4 @@
-SURE/PerturbFlow.py,sha256=5HzS8oB06iSR3JM5AalGfYi-quxbjkTZeTypjih-VBI,54759
+SURE/PerturbFlow.py,sha256=hu4kD6PphVqAG8GlITk44trwVfE6jz72RhQlTYRbi-o,54991
 SURE/SURE.py,sha256=g8EhovBxjfpbVJA0AkmVkQ_ZW_JFc8TtkTCg8FCybV4,47750
 SURE/__init__.py,sha256=NOJI_K-eCqPgStXXvgl3wIEMp6d8saMTDYLJ7Ga9MqE,293
 SURE/assembly/__init__.py,sha256=jxZLURXKPzXe21LhrZ09LgZr33iqdjlQy4oSEj5gR2Q,172
@@ -14,12 +14,12 @@ SURE/flow/plot_quiver.py,sha256=UbmuScUcgbQHeMmjKmgqxjrIjHhiHx0VWct16UMMwuE,8110
 SURE/perturb/__init__.py,sha256=8TP1dSUhXiZzKpFebHZmm8XMMGbUz_OfQ10xu-6uPPY,43
 SURE/perturb/perturb.py,sha256=1iSsCePcwkA2CyM1nCdq_G8gogUNjhMH0BfhhvhpJQk,5037
 SURE/utils/__init__.py,sha256=YF5jB-PAHJQ40OlcZ7BCZbsN2q1JKuPT6EppilRXQqM,680
-SURE/utils/custom_mlp.py,sha256=C0EXLGYsWkUQpEL49AyBFPSzKmasb2hdvtnJfxbF-YU,9282
+SURE/utils/custom_mlp.py,sha256=HuNb7f8-6RFjsvfEu1XOuNpLrHZkGYHgf8TpJfPSNO0,9382
 SURE/utils/queue.py,sha256=E_5PA5EWcBoGAZj8BkKQnkCK0p4C-4-xcTPqdIXaPXU,1892
 SURE/utils/utils.py,sha256=IUHjDDtYaAYllCWsZyIzqQwaLul6fJRvHRH4vIYcR-c,8462
-sure_tools-2.1.83.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
-sure_tools-2.1.83.dist-info/METADATA,sha256=H-q3GA7c-UxJp8C3OfR-f7YpSkhqaSQD3oZ_qcg9OJo,2678
-sure_tools-2.1.83.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sure_tools-2.1.83.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
-sure_tools-2.1.83.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
-sure_tools-2.1.83.dist-info/RECORD,,
+sure_tools-2.1.85.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
+sure_tools-2.1.85.dist-info/METADATA,sha256=I2sMctindGA7kNf_TAdn64Oq6GL_QcGDuKLV6B9zQKY,2678
+sure_tools-2.1.85.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sure_tools-2.1.85.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
+sure_tools-2.1.85.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
+sure_tools-2.1.85.dist-info/RECORD,,