SURE-tools 2.1.56__tar.gz → 2.1.64__tar.gz

{sure_tools-2.1.56 → sure_tools-2.1.64}/PKG-INFO

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

{sure_tools-2.1.56 → sure_tools-2.1.64}/SURE/PerturbFlow.py

@@ -62,7 +62,7 @@ class PerturbFlow(nn.Module):
                  supervised_mode: bool = False,
                  z_dim: int = 10,
                  z_dist: Literal['normal','studentt','laplacian','cauchy','gumbel'] = 'normal',
-                 loss_func: Literal['negbinomial','poisson','multinomial','bernoulli','gamma-poisson'] = 'negbinomial',
+                 loss_func: Literal['negbinomial','poisson','multinomial','bernoulli'] = 'multinomial',
                  inverse_dispersion: float = 10.0,
                  use_zeroinflate: bool = False,
                  hidden_layers: list = [300],
@@ -225,28 +225,7 @@ class PerturbFlow(nn.Module):
                         )
                     )
             
-        if self.loss_func == 'gamma-poisson':
-            self.decoder_concentrate = MLP(
-                    [self.latent_dim] + self.decoder_hidden_layers + [[self.input_size,self.input_size]],
-                    activation=activate_fct,
-                    output_activation=[Exp,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_concentrate = MLP(
-            #        [self.latent_dim] + self.decoder_hidden_layers + [[self.input_size,self.input_size]],
-            #        activation=activate_fct,
-            #        output_activation=[Exp,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_concentrate = self.decoder_concentrate
-        else:
-            self.decoder_concentrate = MLP(
+        self.decoder_concentrate = MLP(
                     [self.latent_dim] + self.decoder_hidden_layers + [self.input_size],
                     activation=activate_fct,
                     output_activation=None,
@@ -338,8 +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)
-        elif self.loss_func == 'multinomial':
-            total_count = pyro.param('total_count', int(1e8), constraint=constraints.positive_integer)
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -373,17 +350,14 @@ class PerturbFlow(nn.Module):
                 zns = pyro.sample('zn', dist.Gumbel(zn_loc, zn_scale).to_event(1))
 
             zs = zns
-            if self.loss_func == 'gamma-poisson':
-                con_alpha,con_beta = self.decoder_concentrate(zs)
-                rate = pyro.sample('cs', dist.Gamma(con_alpha, con_beta).to_event(1))
+            concentrate = self.decoder_concentrate(zs)
+            if self.loss_func in ['bernoulli']:
+                log_theta = concentrate
             else:
-                concentrate = self.decoder_concentrate(zs)
-                if self.loss_func == 'bernoulli':
-                    log_theta = concentrate
-                else:
-                    rate = concentrate.exp()
-                    if self.loss_func != 'poisson':
-                        theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                rate = concentrate.exp()
+                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:
@@ -395,13 +369,8 @@ class PerturbFlow(nn.Module):
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
                 else:
                     pyro.sample('x', dist.Poisson(rate=rate).to_event(1), obs=xs.round())
-            elif self.loss_func == 'gamma-poisson':
-                if self.use_zeroinflate:
-                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
-                else:
-                    pyro.sample('x', dist.Poisson(rate=rate).to_event(1), obs=xs.round())
             elif self.loss_func == 'multinomial':
-                pyro.sample('x', dist.Multinomial(total_count=total_count, probs=theta), obs=xs)
+                pyro.sample('x', dist.Multinomial(total_count=int(1e8), probs=theta), obs=xs)
             elif self.loss_func == 'bernoulli':
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Bernoulli(logits=log_theta),gate_logits=gate_logits).to_event(1), obs=xs)
@@ -416,10 +385,6 @@ 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))
 
     def model2(self, xs, us=None):
         pyro.module('PerturbFlow', self)
@@ -431,8 +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)
-        elif self.loss_func == 'multinomial':
-            total_count = pyro.param('total_count', int(1e8), constraint=constraints.positive_integer)
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -466,28 +429,19 @@ class PerturbFlow(nn.Module):
                 zns = pyro.sample('zn', dist.Gumbel(zn_loc, zn_scale).to_event(1))
 
             if self.cell_factor_size>0:
-                #zus = None
-                #for i in np.arange(self.cell_factor_size):
-                #    if i==0:
-                #        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(zns, us)
                 zs = zns+zus
             else:
                 zs = zns
 
-            if self.loss_func == 'gamma-poisson':
-                con_alpha,con_beta = self.decoder_concentrate(zs)
-                rate = pyro.sample('cs', dist.Gamma(con_alpha, con_beta).to_event(1))
+            concentrate = self.decoder_concentrate(zs)
+            if self.loss_func in ['bernoulli']:
+                log_theta = concentrate
             else:
-                concentrate = self.decoder_concentrate(zs)
-                if self.loss_func == 'bernoulli':
-                    log_theta = concentrate
-                else:
-                    rate = concentrate.exp()
-                    if self.loss_func != 'poisson':
-                        theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                rate = concentrate.exp()
+                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:
@@ -499,13 +453,8 @@ class PerturbFlow(nn.Module):
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
                 else:
                     pyro.sample('x', dist.Poisson(rate=rate).to_event(1), obs=xs.round())
-            elif self.loss_func == 'gamma-poisson':
-                if self.use_zeroinflate:
-                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
-                else:
-                    pyro.sample('x', dist.Poisson(rate=rate).to_event(1), obs=xs.round())
             elif self.loss_func == 'multinomial':
-                pyro.sample('x', dist.Multinomial(total_count=total_count, probs=theta), obs=xs)
+                pyro.sample('x', dist.Multinomial(total_count=int(1e8), probs=theta), obs=xs)
             elif self.loss_func == 'bernoulli':
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Bernoulli(logits=log_theta),gate_logits=gate_logits).to_event(1), obs=xs)
@@ -520,10 +469,6 @@ 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))
 
     def model3(self, xs, ys, embeds=None):
         pyro.module('PerturbFlow', self)
@@ -535,8 +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)
-        elif self.loss_func == 'multinomial':
-            total_count = pyro.param('total_count', int(1e8), constraint=constraints.positive_integer)
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -587,17 +530,14 @@ class PerturbFlow(nn.Module):
 
             zs = zns
 
-            if self.loss_func == 'gamma-poisson':
-                con_alpha,con_beta = self.decoder_concentrate(zs)
-                rate = pyro.sample('cs', dist.Gamma(con_alpha, con_beta).to_event(1))
+            concentrate = self.decoder_concentrate(zs)
+            if self.loss_func in ['bernoulli']:
+                log_theta = concentrate
             else:
-                concentrate = self.decoder_concentrate(zs)
-                if self.loss_func == 'bernoulli':
-                    log_theta = concentrate
-                else:
-                    rate = concentrate.exp()
-                    if self.loss_func != 'poisson':
-                        theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                rate = concentrate.exp()
+                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:
@@ -609,13 +549,8 @@ class PerturbFlow(nn.Module):
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
                 else:
                     pyro.sample('x', dist.Poisson(rate=rate).to_event(1), obs=xs.round())
-            elif self.loss_func == 'gamma-poisson':
-                if self.use_zeroinflate:
-                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
-                else:
-                    pyro.sample('x', dist.Poisson(rate=rate).to_event(1), obs=xs.round())
             elif self.loss_func == 'multinomial':
-                pyro.sample('x', dist.Multinomial(total_count=total_count, probs=theta), obs=xs)
+                pyro.sample('x', dist.Multinomial(total_count=int(1e8), probs=theta), obs=xs)
             elif self.loss_func == 'bernoulli':
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Bernoulli(logits=log_theta),gate_logits=gate_logits).to_event(1), obs=xs)
@@ -630,10 +565,6 @@ class PerturbFlow(nn.Module):
                 zns = pyro.sample('zn', dist.Normal(zn_loc, zn_scale).to_event(1))
             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))
 
     def model4(self, xs, us, ys, embeds=None):
         pyro.module('PerturbFlow', self)
@@ -645,8 +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)
-        elif self.loss_func == 'multinomial':
-            total_count = pyro.param('total_count', int(1e8), constraint=constraints.positive_integer)
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -707,17 +636,14 @@ class PerturbFlow(nn.Module):
             else:
                 zs = zns
 
-            if self.loss_func == 'gamma-poisson':
-                con_alpha,con_beta = self.decoder_concentrate(zs)
-                rate = pyro.sample('cs', dist.Gamma(con_alpha, con_beta).to_event(1))
+            concentrate = self.decoder_concentrate(zs)
+            if self.loss_func in ['bernoulli']:
+                log_theta = concentrate
             else:
-                concentrate = self.decoder_concentrate(zs)
-                if self.loss_func == 'bernoulli':
-                    log_theta = concentrate
-                else:
-                    rate = concentrate.exp()
-                    if self.loss_func != 'poisson':
-                        theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                rate = concentrate.exp()
+                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:
@@ -729,13 +655,8 @@ class PerturbFlow(nn.Module):
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
                 else:
                     pyro.sample('x', dist.Poisson(rate=rate).to_event(1), obs=xs.round())
-            elif self.loss_func == 'gamma-poisson':
-                if self.use_zeroinflate:
-                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
-                else:
-                    pyro.sample('x', dist.Poisson(rate=rate).to_event(1), obs=xs.round())
             elif self.loss_func == 'multinomial':
-                pyro.sample('x', dist.Multinomial(total_count=total_count, probs=theta), obs=xs)
+                pyro.sample('x', dist.Multinomial(total_count=int(1e8), probs=theta), obs=xs)
             elif self.loss_func == 'bernoulli':
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Bernoulli(logits=log_theta),gate_logits=gate_logits).to_event(1), obs=xs)
@@ -750,10 +671,6 @@ class PerturbFlow(nn.Module):
                 zns = pyro.sample('zn', dist.Normal(zn_loc, zn_scale).to_event(1))
             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))
 
     def _total_effects(self, zns, us):
         zus = None
@@ -932,12 +849,7 @@ class PerturbFlow(nn.Module):
         return tensor_to_numpy(ms)
     
     def _get_expression_response(self, delta_zs):
-        if self.loss_func == 'gamma-poisson':
-            alpha,beta = self.decoder_concentrate(delta_zs)
-            xs = dist.Gamma(alpha,beta).to_event(1).mean
-        else:
-            xs = self.decoder_concentrate(delta_zs)
-        return xs
+        return self.decoder_concentrate(delta_zs)
     
     def get_expression_response(self, 
                              delta_zs, 
@@ -966,16 +878,16 @@ class PerturbFlow(nn.Module):
             counts = dist.Bernoulli(logits=concentrate).to_event(1).mean
         elif self.loss_func == 'negbinomial':
             #counts = concentrate.exp()
-            rate = concentrate.exp()
-            theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+            #rate = concentrate.exp()
+            #theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
             
             total_count = pyro.param("inverse_dispersion")
-            counts = dist.NegativeBinomial(total_count=total_count, probs=theta).to_event(1).mean
+            counts = dist.NegativeBinomial(total_count=total_count, logits=concentrate).to_event(1).mean
         elif self.loss_func == 'poisson':
             rate = concentrate.exp()
+            theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+            rate = theta * library_size
             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
@@ -997,16 +909,17 @@ class PerturbFlow(nn.Module):
                         use_sampler: bool = False):
 
         zs = convert_to_tensor(zs, device=self.get_device())
-        ls = zs 
         
-        if self.loss_func == 'multinomial':
-            assert library_sizes!=None, 'Library sizes are required for multinomial!'
+        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
             
         dataset = CustomDataset2(zs,ls)
         dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)

{sure_tools-2.1.56 → sure_tools-2.1.64}/SURE_tools.egg-info/PKG-INFO

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

{sure_tools-2.1.56 → sure_tools-2.1.64}/setup.py

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