SURE-tools 2.2.19__py3-none-any.whl → 2.3.2__py3-none-any.whl

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)
@@ -706,13 +710,13 @@ class DensityFlow(nn.Module):
                 #    zus = self.cell_factor_effect[i](us[:,i].reshape(-1,1))
                 #else:
                 #    zus = self.cell_factor_effect[i]([zns,us[:,i].reshape(-1,1)])
-                zus = self._cell_response(zns, i, us)
+                zus = self._cell_response(zns, i, us[:,i].reshape(-1,1))
             else:
                 #if self.turn_off_cell_specific:
                 #    zus = zus + self.cell_factor_effect[i](us[:,i].reshape(-1,1))
                 #else:
                 #    zus = zus + self.cell_factor_effect[i]([zns,us[:,i].reshape(-1,1)])
-                zus = zus + self._cell_response(zns, i, us)
+                zus = zus + self._cell_response(zns, i, us[:,i].reshape(-1,1))
         return zus
                         
     def _get_codebook_identity(self):
@@ -854,12 +858,12 @@ class DensityFlow(nn.Module):
             us_i = us[:,pert_idx].reshape(-1,1)
             
             # factor effect of xs 
-            dzs0 = self.get_cell_response(zs, factor_idx=pert_idx, perturb=us_i)
+            dzs0 = self.get_cell_response(zs, perturb_idx=pert_idx, perturb_us=us_i)
             
             # perturbation effect 
             ps = np.ones_like(us_i)
             if np.sum(np.abs(ps-us_i))>=1:
-                dzs = self.get_cell_response(zs, factor_idx=pert_idx, perturb=ps)          
+                dzs = self.get_cell_response(zs, perturb_idx=pert_idx, perturb_us=ps)          
                 zs = zs + dzs0 + dzs
             else:
                 zs = zs + dzs0
@@ -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)