SURE-tools 2.1.77__py3-none-any.whl → 2.1.79__py3-none-any.whl

SURE/PerturbFlow.py

@@ -104,7 +104,6 @@ class PerturbFlow(nn.Module):
         #self.use_bias = not zero_bias
         
         self.codebook_weights = None
-        self.total_count = None
 
         set_random_seed(seed)
         self.setup_networks()
@@ -318,7 +317,6 @@ class PerturbFlow(nn.Module):
         if self.loss_func=='negbinomial':
             total_count = pyro.param("inverse_dispersion", self.inverse_dispersion *
                                      xs.new_ones(self.input_size), constraint=constraints.positive)
-            self.total_count = total_count
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -398,7 +396,6 @@ class PerturbFlow(nn.Module):
         if self.loss_func=='negbinomial':
             total_count = pyro.param("inverse_dispersion", self.inverse_dispersion *
                                      xs.new_ones(self.input_size), constraint=constraints.positive)
-            self.total_count = total_count
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -483,7 +480,6 @@ class PerturbFlow(nn.Module):
         if self.loss_func=='negbinomial':
             total_count = pyro.param("inverse_dispersion", self.inverse_dispersion *
                                      xs.new_ones(self.input_size), constraint=constraints.positive)
-            self.total_count = total_count
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -580,7 +576,6 @@ class PerturbFlow(nn.Module):
         if self.loss_func=='negbinomial':
             total_count = pyro.param("inverse_dispersion", self.inverse_dispersion *
                                      xs.new_ones(self.input_size), constraint=constraints.positive)
-            self.total_count = total_count
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -881,25 +876,11 @@ class PerturbFlow(nn.Module):
         if self.loss_func == 'bernoulli':
             #counts = self.sigmoid(concentrate)
             counts = dist.Bernoulli(logits=concentrate).to_event(1).mean
-        elif self.loss_func == 'negbinomial':
-            rate = concentrate.exp()
-            theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
-            
-            total_count = self.total_count
-            #total_count = pyro.param("inverse_dispersion")
-            #store = pyro.get_param_store()
-            #total_count = store['inverse_dispersion']
-            counts = dist.NegativeBinomial(total_count=total_count, probs=theta).to_event(1).mean
-        elif self.loss_func == 'poisson':
+        else:
             rate = concentrate.exp()
             theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
             counts = theta * library_size
             #counts = dist.Poisson(rate=rate).to_event(1).mean
-        elif self.loss_func == 'multinomial':
-            rate = concentrate.exp()
-            theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
-            #counts = dist.Multinomial(total_count=int(1e8), probs=theta).mean
-            counts = theta * library_size
         return counts
     
     def _count_sample(self,concentrate):
@@ -911,22 +892,17 @@ class PerturbFlow(nn.Module):
             counts = dist.Poisson(rate=counts).to_event(1).sample()
         return counts
     
-    def get_counts(self, zs, library_sizes = None,
+    def get_counts(self, zs, library_sizes,
                         batch_size: int = 1024, 
                         use_sampler: bool = False):
 
         zs = convert_to_tensor(zs, device=self.get_device())
         
-        if self.loss_func in ['multinomial','poisson']:
-            assert library_sizes is not None, 'Library sizes are required for multinomial!'
-            
-            if type(library_sizes) == list:
-                library_sizes = np.array(library_sizes).view(-1,1)
-            elif len(library_sizes.shape)==1:
-                library_sizes = library_sizes.view(-1,1)
-            ls = convert_to_tensor(library_sizes, device=self.get_device())
-        else:
-            ls = zs
+        if type(library_sizes) == list:
+            library_sizes = np.array(library_sizes).view(-1,1)
+        elif len(library_sizes.shape)==1:
+            library_sizes = library_sizes.view(-1,1)
+        ls = convert_to_tensor(library_sizes, device=self.get_device())
             
         dataset = CustomDataset2(zs,ls)
         dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
@@ -1084,9 +1060,6 @@ class PerturbFlow(nn.Module):
                 pbar.set_postfix({'loss': str_loss})
                 pbar.update(1)
                 
-        if self.loss_func == 'negbinomial':
-            self.total_count = pyro.param('inverse_dispersion')
-
     @classmethod
     def save_model(cls, model, file_path, compression=False):
         """Save the model to the specified file path."""

SURE/SURE.py

@@ -369,12 +369,13 @@ class SURE(nn.Module):
 
             zs = zns
             concentrate = self.decoder_concentrate(zs)
-            if self.loss_func == 'bernoulli':
+            if self.loss_func in ['bernoulli']:
                 log_theta = concentrate
             else:
                 rate = concentrate.exp()
-                if self.loss_func != 'poisson':
-                    theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                if self.loss_func == 'poisson':
+                    rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
                 if self.use_zeroinflate:
@@ -451,12 +452,13 @@ class SURE(nn.Module):
                 zs = zns
 
             concentrate = self.decoder_concentrate(zs)
-            if self.loss_func == 'bernoulli':
+            if self.loss_func in ['bernoulli']:
                 log_theta = concentrate
             else:
                 rate = concentrate.exp()
-                if self.loss_func != 'poisson':
-                    theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                if self.loss_func == 'poisson':
+                    rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
                 if self.use_zeroinflate:
@@ -545,12 +547,13 @@ class SURE(nn.Module):
             zs = zns
 
             concentrate = self.decoder_concentrate(zs)
-            if self.loss_func == 'bernoulli':
+            if self.loss_func in ['bernoulli']:
                 log_theta = concentrate
             else:
                 rate = concentrate.exp()
-                if self.loss_func != 'poisson':
-                    theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                if self.loss_func == 'poisson':
+                    rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
                 if self.use_zeroinflate:
@@ -641,13 +644,14 @@ class SURE(nn.Module):
                 zs = zns
 
             concentrate = self.decoder_concentrate(zs)
-            if self.loss_func == 'bernoulli':
+            if self.loss_func in ['bernoulli']:
                 log_theta = concentrate
             else:
                 rate = concentrate.exp()
-                if self.loss_func != 'poisson':
-                    theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
-
+                theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                if self.loss_func == 'poisson':
+                    rate = theta * torch.sum(xs, dim=1, keepdim=True)
+                    
             if self.loss_func == 'negbinomial':
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=total_count, probs=theta),gate_logits=gate_logits).to_event(1), obs=xs)

{sure_tools-2.1.77.dist-info → sure_tools-2.1.79.dist-info}/METADATA

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

{sure_tools-2.1.77.dist-info → sure_tools-2.1.79.dist-info}/RECORD

@@ -1,5 +1,5 @@
-SURE/PerturbFlow.py,sha256=lRBSHIvJiAeyOyRDKukQsHpGxMU9q_jNEYKYcyyQtIc,54717
-SURE/SURE.py,sha256=ko15a9BhvUqHviogZ0YCdTQjM-2zqkO9OvHZSpnGbg0,47458
+SURE/PerturbFlow.py,sha256=jvN2CD2L7YX6EiqP0Zw--zwegwHDiGkQYL2v_So5YYA,53382
+SURE/SURE.py,sha256=g8EhovBxjfpbVJA0AkmVkQ_ZW_JFc8TtkTCg8FCybV4,47750
 SURE/__init__.py,sha256=NOJI_K-eCqPgStXXvgl3wIEMp6d8saMTDYLJ7Ga9MqE,293
 SURE/assembly/__init__.py,sha256=jxZLURXKPzXe21LhrZ09LgZr33iqdjlQy4oSEj5gR2Q,172
 SURE/assembly/assembly.py,sha256=6IMdelPOiRO4mUb4dC7gVCoF1Uvfw86-Map8P_jnUag,21477
@@ -17,9 +17,9 @@ SURE/utils/__init__.py,sha256=YF5jB-PAHJQ40OlcZ7BCZbsN2q1JKuPT6EppilRXQqM,680
 SURE/utils/custom_mlp.py,sha256=C0EXLGYsWkUQpEL49AyBFPSzKmasb2hdvtnJfxbF-YU,9282
 SURE/utils/queue.py,sha256=E_5PA5EWcBoGAZj8BkKQnkCK0p4C-4-xcTPqdIXaPXU,1892
 SURE/utils/utils.py,sha256=IUHjDDtYaAYllCWsZyIzqQwaLul6fJRvHRH4vIYcR-c,8462
-sure_tools-2.1.77.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
-sure_tools-2.1.77.dist-info/METADATA,sha256=R2y9rOLrZXegdYLXXdqUMiIc0wYi-w1cjFO198SnkZo,2678
-sure_tools-2.1.77.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sure_tools-2.1.77.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
-sure_tools-2.1.77.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
-sure_tools-2.1.77.dist-info/RECORD,,
+sure_tools-2.1.79.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
+sure_tools-2.1.79.dist-info/METADATA,sha256=B8i3DryW1NcFrOvzEROG5YDlZG14RHa_q2cVbNVUwmg,2678
+sure_tools-2.1.79.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sure_tools-2.1.79.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
+sure_tools-2.1.79.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
+sure_tools-2.1.79.dist-info/RECORD,,