SURE-tools 2.2.28__tar.gz → 2.3.2__tar.gz

{sure_tools-2.2.28 → sure_tools-2.3.2}/PKG-INFO

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

{sure_tools-2.2.28 → sure_tools-2.3.2}/SURE/DensityFlow.py

@@ -396,7 +396,8 @@ class DensityFlow(nn.Module):
                 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=int(1e8), probs=theta), obs=xs)
+                pyro.sample('x', dist.Multinomial(total_count=int(1e8), logits=concentrate), 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)
@@ -480,7 +481,8 @@ class DensityFlow(nn.Module):
                 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=int(1e8), probs=theta), obs=xs)
+                pyro.sample('x', dist.Multinomial(total_count=int(1e8), logits=concentrate), 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)
@@ -576,7 +578,8 @@ class DensityFlow(nn.Module):
                 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=int(1e8), probs=theta), obs=xs)
+                pyro.sample('x', dist.Multinomial(total_count=int(1e8), logits=concentrate), 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)
@@ -682,7 +685,8 @@ class DensityFlow(nn.Module):
                 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=int(1e8), probs=theta), obs=xs)
+                pyro.sample('x', dist.Multinomial(total_count=int(1e8), logits=concentrate), 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)
@@ -946,6 +950,9 @@ class DensityFlow(nn.Module):
         if self.loss_func == 'bernoulli':
             #counts = self.sigmoid(concentrate)
             counts = dist.Bernoulli(logits=concentrate).to_event(1).mean
+        elif self.loss_func == 'multinomial':
+            theta = dist.Multinomial(total_count=int(1e8), logits=concentrate).mean
+            counts = theta * library_size
         else:
             rate = concentrate.exp()
             theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
@@ -1340,7 +1347,7 @@ def main():
     cell_factor_size = 0 if us is None else us.shape[1]
 
     ###########################################
-    DensityFlow = DensityFlow(
+    df = DensityFlow(
         input_size=input_size,
         cell_factor_size=cell_factor_size,
         inverse_dispersion=args.inverse_dispersion,
@@ -1359,7 +1366,7 @@ def main():
         dtype=dtype,
     )
 
-    DensityFlow.fit(xs, us=us, 
+    df.fit(xs, us=us, 
              num_epochs=args.num_epochs,
              learning_rate=args.learning_rate,
              batch_size=args.batch_size,
@@ -1371,9 +1378,9 @@ def main():
 
     if args.save_model is not None:
         if args.save_model.endswith('gz'):
-            DensityFlow.save_model(DensityFlow, args.save_model, compression=True)
+            DensityFlow.save_model(df, args.save_model, compression=True)
         else:
-            DensityFlow.save_model(DensityFlow, args.save_model)
+            DensityFlow.save_model(df, args.save_model)