SURE-tools 2.1.78__tar.gz → 2.1.80__tar.gz

{sure_tools-2.1.78 → sure_tools-2.1.80}/PKG-INFO

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

{sure_tools-2.1.78 → sure_tools-2.1.80}/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))
@@ -812,6 +807,36 @@ class PerturbFlow(nn.Module):
         A = np.concatenate(A)
         return A
     
+    def predict(self, xs, us, perturbs_predict:list, perturbs_reference:list, library_sizes=None):
+        perturbs_reference = np.array(perturbs_reference)
+        
+        # basal embedding
+        zs = self.get_basal_embedding(xs)
+        for pert in perturbs_predict:
+            pert_idx = np.where(perturbs_reference==pert)[0]
+            us_i = us[:,pert_idx].reshape(-1,1)
+            
+            # factor effect of xs 
+            dzs0 = self.get_cell_response(xs, factor_idx=pert_idx, perturb=us_i)
+            
+            # perturbation effect 
+            ps = np.ones_like(us_i)
+            dzs = self.get_cell_response(xs, factor_idx=pert_idx, perturb=ps)
+            
+            zs = zs + dzs0 + dzs
+        
+        if library_sizes is None:
+            library_sizes = np.sum(xs, axis=1, keepdims=True)
+        elif type(library_sizes) == list:
+            library_sizes = np.array(library_sizes)
+            library_sizes = library_sizes.reshape(-1,1)
+        elif len(library_sizes.shape)==1:
+            library_sizes = library_sizes.reshape(-1,1)
+            
+        counts = self.get_counts(zs, library_sizes=library_sizes)
+        
+        return counts, zs
+    
     def _cell_response(self, xs, factor_idx, perturb):
         #zns,_ = self.encoder_zn(xs)    
         zns,_ = self._get_basal_embedding(xs)   
@@ -1065,9 +1090,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_tools-2.1.78 → sure_tools-2.1.80}/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.78 → sure_tools-2.1.80}/SURE_tools.egg-info/PKG-INFO

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

{sure_tools-2.1.78 → sure_tools-2.1.80}/setup.py

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