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

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'] = 'negbinomial',
+                 loss_func: Literal['negbinomial','poisson','multinomial','bernoulli','gamma-poisson'] = 'negbinomial',
                  inverse_dispersion: float = 10.0,
                  use_zeroinflate: bool = False,
                  hidden_layers: list = [300],
@@ -225,15 +225,36 @@ class PerturbFlow(nn.Module):
                         )
                     )
             
-        self.decoder_concentrate = MLP(
-                [self.latent_dim] + self.decoder_hidden_layers + [self.input_size],
-                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,
-            )
+        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.latent_dim] + self.decoder_hidden_layers + [self.input_size],
+                    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,
+                )
 
         if self.latent_dist == 'studentt':
             self.codebook = MLP(
@@ -317,12 +338,14 @@ 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))
             
         acs_scale = pyro.param("codebook_scale", xs.new_ones(self.latent_dim), constraint=constraints.positive)
-
+        
         I = torch.eye(self.code_size)
         if self.latent_dist=='studentt':
             acs_dof,acs_loc = self.codebook(I)
@@ -350,13 +373,17 @@ class PerturbFlow(nn.Module):
                 zns = pyro.sample('zn', dist.Gumbel(zn_loc, zn_scale).to_event(1))
 
             zs = zns
-            concentrate = self.decoder_concentrate(zs)
-            if self.loss_func == 'bernoulli':
-                log_theta = concentrate
+            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))
             else:
-                rate = concentrate.exp()
-                if self.loss_func != 'poisson':
-                    theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                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
 
             if self.loss_func == 'negbinomial':
                 if self.use_zeroinflate:
@@ -368,8 +395,13 @@ 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=int(1e8), probs=theta), obs=xs)
+                pyro.sample('x', dist.Multinomial(total_count=total_count, 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)
@@ -384,6 +416,10 @@ 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)
@@ -395,6 +431,8 @@ 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))
@@ -439,13 +477,17 @@ class PerturbFlow(nn.Module):
             else:
                 zs = zns
 
-            concentrate = self.decoder_concentrate(zs)
-            if self.loss_func == 'bernoulli':
-                log_theta = concentrate
+            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))
             else:
-                rate = concentrate.exp()
-                if self.loss_func != 'poisson':
-                    theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                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
 
             if self.loss_func == 'negbinomial':
                 if self.use_zeroinflate:
@@ -457,8 +499,13 @@ 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=int(1e8), probs=theta), obs=xs)
+                pyro.sample('x', dist.Multinomial(total_count=total_count, 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)
@@ -473,6 +520,10 @@ 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)
@@ -484,6 +535,8 @@ 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))
@@ -534,13 +587,17 @@ class PerturbFlow(nn.Module):
 
             zs = zns
 
-            concentrate = self.decoder_concentrate(zs)
-            if self.loss_func == 'bernoulli':
-                log_theta = concentrate
+            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))
             else:
-                rate = concentrate.exp()
-                if self.loss_func != 'poisson':
-                    theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                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
 
             if self.loss_func == 'negbinomial':
                 if self.use_zeroinflate:
@@ -552,8 +609,13 @@ 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=int(1e8), probs=theta), obs=xs)
+                pyro.sample('x', dist.Multinomial(total_count=total_count, 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)
@@ -566,6 +628,12 @@ class PerturbFlow(nn.Module):
                 #zn_loc, zn_scale = self.encoder_zn(xs)
                 zn_loc, zn_scale = self._get_basal_embedding(xs)
                 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)
@@ -577,6 +645,8 @@ 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))
@@ -637,13 +707,17 @@ class PerturbFlow(nn.Module):
             else:
                 zs = zns
 
-            concentrate = self.decoder_concentrate(zs)
-            if self.loss_func == 'bernoulli':
-                log_theta = concentrate
+            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))
             else:
-                rate = concentrate.exp()
-                if self.loss_func != 'poisson':
-                    theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+                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
 
             if self.loss_func == 'negbinomial':
                 if self.use_zeroinflate:
@@ -655,8 +729,13 @@ 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=int(1e8), probs=theta), obs=xs)
+                pyro.sample('x', dist.Multinomial(total_count=total_count, 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)
@@ -669,6 +748,12 @@ class PerturbFlow(nn.Module):
                 #zn_loc, zn_scale = self.encoder_zn(xs)
                 zn_loc, zn_scale = self._get_basal_embedding(xs)
                 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
@@ -847,7 +932,12 @@ class PerturbFlow(nn.Module):
         return tensor_to_numpy(ms)
     
     def _get_expression_response(self, delta_zs):
-        return self.decoder_concentrate(delta_zs)
+        if self.loss_func == 'gamma-poisson':
+            alpha,beta = self.encoder_concentrate(delta_zs)
+            xs = dist.Gamma(alpha,beta).to_event(1).mean
+        else:
+            xs = self.decoder_concentrate(delta_zs)
+        return xs
     
     def get_expression_response(self, 
                              delta_zs, 
@@ -880,10 +970,12 @@ class PerturbFlow(nn.Module):
             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)
+            counts = dist.NegativeBinomial(total_count=total_count, probs=theta).to_event(1).mean
         elif self.loss_func == 'poisson':
             rate = concentrate.exp()
-            counts = dist.Poisson(rate=rate).to_event(1)
+            counts = dist.Poisson(rate=rate).to_event(1).mean
+        elif self.loss_func == 'gamma-poisson':
+            counts = dist.Poisson(rate=concentrate).to_event(1).mean
         return counts
     
     def _count_sample(self,concentrate):

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

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

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

@@ -1,4 +1,4 @@
-SURE/PerturbFlow.py,sha256=yt0kOW4buZKpJQ3Jn_8Zd2uEKUq29DZYkSzcgEP55EA,53211
+SURE/PerturbFlow.py,sha256=hOVEsBrMAs7T5yi3LW7KV6hwPuwyjZtKG2wyMF6R08E,58614
 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.52.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
-sure_tools-2.1.52.dist-info/METADATA,sha256=ARO36IQ9aKV9Sp4F9AkLR6zzXSsPnMGguljP7XU95Mk,2678
-sure_tools-2.1.52.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sure_tools-2.1.52.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
-sure_tools-2.1.52.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
-sure_tools-2.1.52.dist-info/RECORD,,
+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,,