SURE-tools 2.1.42__py3-none-any.whl → 2.1.44__py3-none-any.whl

SURE/PerturbFlow.py

@@ -73,7 +73,7 @@ class PerturbFlow(nn.Module):
                  config_enum: str = 'parallel',
                  use_cuda: bool = False,
                  seed: int = 42,
-                 zero_bias: bool = True,
+                 zero_bias: bool|list = True,
                  dtype = torch.float32, # type: ignore
                  ):
         super().__init__()
@@ -97,7 +97,11 @@ class PerturbFlow(nn.Module):
         self.post_layer_fct = post_layer_fct
         self.post_act_fct = post_act_fct
         self.hidden_layer_activation = hidden_layer_activation
-        self.use_bias = not zero_bias
+        if type(zero_bias) == list:
+            self.use_bias = [not x for x in zero_bias]
+        else:
+            self.use_bias = [not zero_bias] * self.cell_factor_size
+        #self.use_bias = not zero_bias
         
         self.codebook_weights = None
 
@@ -185,10 +189,29 @@ class PerturbFlow(nn.Module):
                 use_cuda=self.use_cuda,
             )
 
-        self.encoder_zn = MLP(
-            [self.input_size] + hidden_sizes + [[latent_dim, latent_dim]],
+        #self.encoder_zn = MLP(
+        #    [self.input_size] + hidden_sizes + [[latent_dim, latent_dim]],
+        #    activation=activate_fct,
+        #    output_activation=[None, Exp],
+        #    post_layer_fct=post_layer_fct,
+        #    post_act_fct=post_act_fct,
+        #    allow_broadcast=self.allow_broadcast,
+        #    use_cuda=self.use_cuda,
+        #)
+        
+        self.encoder_zn_loc = MLP(
+            [self.input_size] + hidden_sizes + [self.input_size * latent_dim],
+            activation=activate_fct,
+            output_activation=None,
+            post_layer_fct=post_layer_fct,
+            post_act_fct=post_act_fct,
+            allow_broadcast=self.allow_broadcast,
+            use_cuda=self.use_cuda,
+        )
+        self.encoder_zn_scale = MLP(
+            [self.input_size] + hidden_sizes + [latent_dim],
             activation=activate_fct,
-            output_activation=[None, Exp],
+            output_activation=Exp,
             post_layer_fct=post_layer_fct,
             post_act_fct=post_act_fct,
             allow_broadcast=self.allow_broadcast,
@@ -198,7 +221,7 @@ class PerturbFlow(nn.Module):
         if self.cell_factor_size>0:
             self.cell_factor_effect = nn.ModuleList()
             for i in np.arange(self.cell_factor_size):
-                if self.use_bias:
+                if self.use_bias[i]:
                     self.cell_factor_effect.append(MLP(
                         [self.latent_dim+1] + self.decoder_hidden_layers + [self.latent_dim],
                         activation=activate_fct,
@@ -374,7 +397,11 @@ class PerturbFlow(nn.Module):
 
     def guide1(self, xs):
         with pyro.plate('data'):
-            zn_loc, zn_scale = self.encoder_zn(xs)
+            batch_size = xs.shape[0]
+            #zn_loc, zn_scale = self.encoder_zn(xs)
+            zn_loc, zn_scale = self.encoder_zn_loc(xs), self.encoder_zn_scale(xs)
+            zn_loc_3d = zn_loc.view(batch_size, self.input_size, self.latent_dim)
+            zn_loc = zn_loc_3d.sum(dim=1)
             zns = pyro.sample('zn', dist.Normal(zn_loc, zn_scale).to_event(1))
 
             alpha = self.encoder_n(zns)
@@ -462,7 +489,11 @@ class PerturbFlow(nn.Module):
 
     def guide2(self, xs, us=None):
         with pyro.plate('data'):
-            zn_loc, zn_scale = self.encoder_zn(xs)
+            batch_size = xs.shape[0]
+            #zn_loc, zn_scale = self.encoder_zn(xs)
+            zn_loc, zn_scale = self.encoder_zn_loc(xs), self.encoder_zn_scale(xs)
+            zn_loc_3d = zn_loc.view(batch_size, self.input_size, self.latent_dim)
+            zn_loc = zn_loc_3d.sum(dim=1)
             zns = pyro.sample('zn', dist.Normal(zn_loc, zn_scale).to_event(1))
 
             alpha = self.encoder_n(zns)
@@ -557,7 +588,11 @@ class PerturbFlow(nn.Module):
     def guide3(self, xs, ys, embeds=None):
         with pyro.plate('data'):
             if embeds is None:
-                zn_loc, zn_scale = self.encoder_zn(xs)
+                batch_size = xs.shape[0]
+                #zn_loc, zn_scale = self.encoder_zn(xs)
+                zn_loc, zn_scale = self.encoder_zn_loc(xs), self.encoder_zn_scale(xs)
+                zn_loc_3d = zn_loc.view(batch_size, self.input_size, self.latent_dim)
+                zn_loc = zn_loc_3d.sum(dim=1)
                 zns = pyro.sample('zn', dist.Normal(zn_loc, zn_scale).to_event(1))
 
     def model4(self, xs, us, ys, embeds=None):
@@ -659,7 +694,11 @@ class PerturbFlow(nn.Module):
     def guide4(self, xs, us, ys, embeds=None):
         with pyro.plate('data'):
             if embeds is None:
-                zn_loc, zn_scale = self.encoder_zn(xs)
+                batch_size = xs.shape[0]
+                #zn_loc, zn_scale = self.encoder_zn(xs)
+                zn_loc, zn_scale = self.encoder_zn_loc(xs), self.encoder_zn_scale(xs)
+                zn_loc_3d = zn_loc.view(batch_size, self.input_size, self.latent_dim)
+                zn_loc = zn_loc_3d.sum(dim=1)
                 zns = pyro.sample('zn', dist.Normal(zn_loc, zn_scale).to_event(1))
 
     def _total_effects(self, zns, us):
@@ -688,7 +727,11 @@ class PerturbFlow(nn.Module):
         return cb
 
     def _get_basal_embedding(self, xs):           
-        zns, _ = self.encoder_zn(xs)
+        #zns, _ = self.encoder_zn(xs)
+        batch_size = xs.shape[0]
+        zns = self.encoder_zn_loc(xs)
+        zns_3d = zns.view(batch_size, self.input_size, self.latent_dim)
+        zns = zns_3d.sum(dim=1)
         return zns 
     
     def get_basal_embedding(self, 
@@ -727,7 +770,8 @@ class PerturbFlow(nn.Module):
         if self.supervised_mode:
             alpha = self.encoder_n(xs)
         else:
-            zns,_ = self.encoder_zn(xs)
+            #zns,_ = self.encoder_zn(xs)
+            zns = self._get_basal_embedding(xs)
             alpha = self.encoder_n(zns)
         return alpha
     
@@ -797,7 +841,8 @@ class PerturbFlow(nn.Module):
         return A
     
     def _cell_response(self, xs, factor_idx, perturb):
-        zns,_ = self.encoder_zn(xs)       
+        #zns,_ = self.encoder_zn(xs)    
+        zns = self._get_basal_embedding(xs)   
         if perturb.ndim==2:
             ms = self.cell_factor_effect[factor_idx]([zns, perturb])
         else:

{sure_tools-2.1.42.dist-info → sure_tools-2.1.44.dist-info}/METADATA

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

{sure_tools-2.1.42.dist-info → sure_tools-2.1.44.dist-info}/RECORD

@@ -1,4 +1,4 @@
-SURE/PerturbFlow.py,sha256=4oz02P9kTRNJ897ITR1oOxKYGVw2sgvFKazz2FC-6f0,52118
+SURE/PerturbFlow.py,sha256=EbbEQBf9HQ6XeUw3g4DKFbe2IVCJRGmCwjQcHwCnQxU,54309
 SURE/SURE.py,sha256=ko15a9BhvUqHviogZ0YCdTQjM-2zqkO9OvHZSpnGbg0,47458
 SURE/__init__.py,sha256=NOJI_K-eCqPgStXXvgl3wIEMp6d8saMTDYLJ7Ga9MqE,293
 SURE/assembly/__init__.py,sha256=jxZLURXKPzXe21LhrZ09LgZr33iqdjlQy4oSEj5gR2Q,172
@@ -17,9 +17,9 @@ SURE/utils/__init__.py,sha256=QJUOfrXzdWSmoM0P3LH8oKEHttzCWqpDy2UF0F0dtN4,673
 SURE/utils/custom_mlp.py,sha256=rHnx9jEef02zfCUdbYVCmbuHcDdIBmRgt__wpdpZvYg,8104
 SURE/utils/queue.py,sha256=E_5PA5EWcBoGAZj8BkKQnkCK0p4C-4-xcTPqdIXaPXU,1892
 SURE/utils/utils.py,sha256=IUHjDDtYaAYllCWsZyIzqQwaLul6fJRvHRH4vIYcR-c,8462
-sure_tools-2.1.42.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
-sure_tools-2.1.42.dist-info/METADATA,sha256=hwV5P2BtOm4qQ-v0seI9yavlU5zNiHr0_jKovfBl-0E,2651
-sure_tools-2.1.42.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sure_tools-2.1.42.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
-sure_tools-2.1.42.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
-sure_tools-2.1.42.dist-info/RECORD,,
+sure_tools-2.1.44.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
+sure_tools-2.1.44.dist-info/METADATA,sha256=GOqZOvGRfh--3bd3Bgpkiv7E05gfRoDNflkE3c4Mt6s,2651
+sure_tools-2.1.44.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sure_tools-2.1.44.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
+sure_tools-2.1.44.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
+sure_tools-2.1.44.dist-info/RECORD,,