SURE-tools 2.1.54__py3-none-any.whl → 2.1.56__py3-none-any.whl

SURE/PerturbFlow.py

@@ -244,7 +244,7 @@ class PerturbFlow(nn.Module):
             #        allow_broadcast=self.allow_broadcast,
             #        use_cuda=self.use_cuda,
             #    )
-            self.encoder_concentrate = self.decoder_concentrate
+            #self.encoder_concentrate = self.decoder_concentrate
         else:
             self.decoder_concentrate = MLP(
                     [self.latent_dim] + self.decoder_hidden_layers + [self.input_size],
@@ -417,9 +417,9 @@ class PerturbFlow(nn.Module):
             alpha = self.encoder_n(zns)
             ns = pyro.sample('n', dist.OneHotCategorical(logits=alpha))
             
-            if self.loss_func == 'gamma-poisson':
-                con_alpha,con_beta = self.encoder_concentrate(zns)
-                rate = pyro.sample('cs', dist.Gamma(con_alpha, con_beta).to_event(1))
+            #if self.loss_func == 'gamma-poisson':
+            #    con_alpha,con_beta = self.encoder_concentrate(zns)
+            #    rate = pyro.sample('cs', dist.Gamma(con_alpha, con_beta).to_event(1))
 
     def model2(self, xs, us=None):
         pyro.module('PerturbFlow', self)
@@ -521,9 +521,9 @@ class PerturbFlow(nn.Module):
             alpha = self.encoder_n(zns)
             ns = pyro.sample('n', dist.OneHotCategorical(logits=alpha))
             
-            if self.loss_func == 'gamma-poisson':
-                con_alpha,con_beta = self.encoder_concentrate(zns)
-                rate = pyro.sample('cs', dist.Gamma(con_alpha, con_beta).to_event(1))
+            #if self.loss_func == 'gamma-poisson':
+            #    con_alpha,con_beta = self.encoder_concentrate(zns)
+            #    rate = pyro.sample('cs', dist.Gamma(con_alpha, con_beta).to_event(1))
 
     def model3(self, xs, ys, embeds=None):
         pyro.module('PerturbFlow', self)
@@ -631,9 +631,9 @@ class PerturbFlow(nn.Module):
             else:
                 zns = embeds
                 
-            if self.loss_func == 'gamma-poisson':
-                con_alpha,con_beta = self.encoder_concentrate(zns)
-                rate = pyro.sample('cs', dist.Gamma(con_alpha, con_beta).to_event(1))
+            #if self.loss_func == 'gamma-poisson':
+            #    con_alpha,con_beta = self.encoder_concentrate(zns)
+            #    rate = pyro.sample('cs', dist.Gamma(con_alpha, con_beta).to_event(1))
 
     def model4(self, xs, us, ys, embeds=None):
         pyro.module('PerturbFlow', self)
@@ -751,9 +751,9 @@ class PerturbFlow(nn.Module):
             else:
                 zns = embeds
             
-            if self.loss_func == 'gamma-poisson':
-                con_alpha,con_beta = self.encoder_concentrate(zns)
-                rate = pyro.sample('cs', dist.Gamma(con_alpha, con_beta).to_event(1))
+            #if self.loss_func == 'gamma-poisson':
+            #    con_alpha,con_beta = self.encoder_concentrate(zns)
+            #    rate = pyro.sample('cs', dist.Gamma(con_alpha, con_beta).to_event(1))
 
     def _total_effects(self, zns, us):
         zus = None
@@ -933,7 +933,7 @@ class PerturbFlow(nn.Module):
     
     def _get_expression_response(self, delta_zs):
         if self.loss_func == 'gamma-poisson':
-            alpha,beta = self.encoder_concentrate(delta_zs)
+            alpha,beta = self.decoder_concentrate(delta_zs)
             xs = dist.Gamma(alpha,beta).to_event(1).mean
         else:
             xs = self.decoder_concentrate(delta_zs)
@@ -960,7 +960,7 @@ class PerturbFlow(nn.Module):
         R = np.concatenate(R)
         return R
     
-    def _count(self,concentrate):
+    def _count(self,concentrate, library_size=None):
         if self.loss_func == 'bernoulli':
             #counts = self.sigmoid(concentrate)
             counts = dist.Bernoulli(logits=concentrate).to_event(1).mean
@@ -976,6 +976,11 @@ class PerturbFlow(nn.Module):
             counts = dist.Poisson(rate=rate).to_event(1).mean
         elif self.loss_func == 'gamma-poisson':
             counts = dist.Poisson(rate=concentrate).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 = counts * library_size
         return counts
     
     def _count_sample(self,concentrate):
@@ -987,22 +992,35 @@ class PerturbFlow(nn.Module):
             counts = dist.Poisson(rate=counts).to_event(1).sample()
         return counts
     
-    def get_counts(self, zs, 
+    def get_counts(self, zs, library_sizes = None,
                         batch_size: int = 1024, 
                         use_sampler: bool = False):
 
         zs = convert_to_tensor(zs, device=self.get_device())
-        dataset = CustomDataset(zs)
+        ls = zs 
+        
+        if self.loss_func == 'multinomial':
+            assert library_sizes!=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)
+            
+        dataset = CustomDataset2(zs,ls)
         dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
         
         E = []
         with tqdm(total=len(dataloader), desc='', unit='batch') as pbar:
-            for Z_batch, _ in dataloader:
+            for Z_batch, L_batch, _ in dataloader:
+                if self.loss_func != 'multinomial':
+                    L_batch = None
                 concentrate = self._get_expression_response(Z_batch)
                 if use_sampler:
                     counts = self._count_sample(concentrate)
                 else:
-                    counts = self._count(concentrate)
+                    counts = self._count(concentrate, L_batch)
                 E.append(tensor_to_numpy(counts))
                 pbar.update(1)
         

{sure_tools-2.1.54.dist-info → sure_tools-2.1.56.dist-info}/METADATA

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

{sure_tools-2.1.54.dist-info → sure_tools-2.1.56.dist-info}/RECORD

@@ -1,4 +1,4 @@
-SURE/PerturbFlow.py,sha256=hOVEsBrMAs7T5yi3LW7KV6hwPuwyjZtKG2wyMF6R08E,58614
+SURE/PerturbFlow.py,sha256=CvnmX1QVo4UK4rkmFQd0RR9YHrNjL1GCHM1aj-BHVqM,59536
 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=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.54.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
-sure_tools-2.1.54.dist-info/METADATA,sha256=VDqYvGzqSz_HeBiPxGgwwl_i-uunVvd3t0MVIa4n6iI,2678
-sure_tools-2.1.54.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sure_tools-2.1.54.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
-sure_tools-2.1.54.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
-sure_tools-2.1.54.dist-info/RECORD,,
+sure_tools-2.1.56.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
+sure_tools-2.1.56.dist-info/METADATA,sha256=kWBC-87jEjWE-JHxXGcrerFaxT9G5buo8zwZhkDxu9o,2678
+sure_tools-2.1.56.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sure_tools-2.1.56.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
+sure_tools-2.1.56.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
+sure_tools-2.1.56.dist-info/RECORD,,