SURE-tools 2.0.9__tar.gz → 2.1.0__tar.gz

{sure_tools-2.0.9 → sure_tools-2.1.0}/PKG-INFO

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

sure_tools-2.0.9/SURE/SURE.py → sure_tools-2.1.0/SURE/PerturbFlow.py

@@ -54,7 +54,7 @@ def set_random_seed(seed):
     # Set seed for Pyro
     pyro.set_rng_seed(seed)
 
-class SURE(nn.Module):
+class PerturbFlow(nn.Module):
     """SUccinct REpresentation of single-omics cells
 
     Parameters
@@ -330,7 +330,7 @@ class SURE(nn.Module):
         return xs
 
     def model1(self, xs):
-        pyro.module('sure', self)
+        pyro.module('PerturbFlow', self)
 
         eps = torch.finfo(xs.dtype).eps
         batch_size = xs.size(0)
@@ -407,7 +407,7 @@ class SURE(nn.Module):
             ns = pyro.sample('n', dist.OneHotCategorical(logits=alpha))
 
     def model2(self, xs, us=None):
-        pyro.module('sure', self)
+        pyro.module('PerturbFlow', self)
 
         eps = torch.finfo(xs.dtype).eps
         batch_size = xs.size(0)
@@ -495,7 +495,7 @@ class SURE(nn.Module):
             ns = pyro.sample('n', dist.OneHotCategorical(logits=alpha))
 
     def model3(self, xs, ys, embeds=None):
-        pyro.module('sure', self)
+        pyro.module('PerturbFlow', self)
 
         eps = torch.finfo(xs.dtype).eps
         batch_size = xs.size(0)
@@ -587,7 +587,7 @@ class SURE(nn.Module):
                 zns = pyro.sample('zn', dist.Normal(zn_loc, zn_scale).to_event(1))
 
     def model4(self, xs, us, ys, embeds=None):
-        pyro.module('sure', self)
+        pyro.module('PerturbFlow', self)
 
         eps = torch.finfo(xs.dtype).eps
         batch_size = xs.size(0)
@@ -813,7 +813,7 @@ class SURE(nn.Module):
         A = np.concatenate(A)
         return A
     
-    def _cell_move(self, xs, factor_idx, perturb):
+    def _cell_state_response(self, xs, factor_idx, perturb):
         zns,_ = self.encoder_zn(xs)
         if type(factor_idx) == str:
             factor_idx = int(np.where(self.cell_factor_names==factor_idx)[0])
@@ -825,7 +825,7 @@ class SURE(nn.Module):
             
         return ms 
 
-    def get_cell_move(self, 
+    def get_cell_state_response(self, 
                              xs, 
                              factor_idx,
                              perturb,
@@ -843,14 +843,14 @@ class SURE(nn.Module):
         Z = []
         with tqdm(total=len(dataloader), desc='', unit='batch') as pbar:
             for X_batch, P_batch, _ in dataloader:
-                zns = self._cell_move(X_batch, factor_idx, P_batch)
+                zns = self._cell_state_response(X_batch, factor_idx, P_batch)
                 Z.append(tensor_to_numpy(zns))
                 pbar.update(1)
 
         Z = np.concatenate(Z)
         return Z
     
-    def get_metacell_move(self, factor_idx, perturb):
+    def get_metacell_response(self, factor_idx, perturb):
         zs = self._get_codebook()
         ps = convert_to_tensor(perturb, device=self.get_device())
         
@@ -860,6 +860,30 @@ class SURE(nn.Module):
         ms = self.cell_factor_effect[factor_idx]([zs,ps])
         return tensor_to_numpy(ms)
     
+    def _get_expression_response(self, delta_zs):
+        return self.decoder_concentrate(delta_zs)
+    
+    def get_expression_response(self, 
+                             delta_zs, 
+                             batch_size: int = 1024):
+        """
+        Return cells' changes in the latent space induced by specific perturbation of a factor
+
+        """
+        delta_zs = convert_to_tensor(delta_zs, device=self.get_device())
+        dataset = CustomDataset(delta_zs)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+        R = []
+        with tqdm(total=len(dataloader), desc='', unit='batch') as pbar:
+            for delta_Z_batch, _ in dataloader:
+                r = self._cell_move(delta_Z_batch)
+                R.append(tensor_to_numpy(r))
+                pbar.update(1)
+
+        R = np.concatenate(R)
+        return R
+    
     def preprocess(self, xs, threshold=0):
         if self.loss_func == 'bernoulli':
             ad = sc.AnnData(xs)
@@ -887,7 +911,7 @@ class SURE(nn.Module):
             threshold: int = 0,
             use_jax: bool = False):
         """
-        Train the SURE model.
+        Train the PerturbFlow model.
 
         Parameters
         ----------
@@ -913,7 +937,7 @@ class SURE(nn.Module):
             Parameter for optimization.
         use_jax
             If toggled on, Jax will be used for speeding up. CAUTION: This will raise errors because of unknown reasons when it is called in
-            the Python script or Jupyter notebook. It is OK if it is used when runing SURE in the shell command.
+            the Python script or Jupyter notebook. It is OK if it is used when runing PerturbFlow in the shell command.
         """
         xs = self.preprocess(xs, threshold=threshold)
         xs = convert_to_tensor(xs, dtype=self.dtype, device=self.get_device())
@@ -1031,12 +1055,12 @@ class SURE(nn.Module):
 
         
 EXAMPLE_RUN = (
-    "example run: SURE --help"
+    "example run: PerturbFlow --help"
 )
 
 def parse_args():
     parser = argparse.ArgumentParser(
-        description="SURE\n{}".format(EXAMPLE_RUN))
+        description="PerturbFlow\n{}".format(EXAMPLE_RUN))
 
     parser.add_argument(
         "--cuda", action="store_true", help="use GPU(s) to speed up training"
@@ -1254,7 +1278,7 @@ def main():
     latent_dist = args.z_dist
 
     ###########################################
-    sure = SURE(
+    perturbflow = PerturbFlow(
         input_size=input_size,
         cell_factor_size=cell_factor_size,
         inverse_dispersion=args.inverse_dispersion,
@@ -1277,7 +1301,7 @@ def main():
         dtype=dtype,
     )
 
-    sure.fit(xs, us=us, 
+    perturbflow.fit(xs, us=us, 
              num_epochs=args.num_epochs,
              learning_rate=args.learning_rate,
              batch_size=args.batch_size,
@@ -1289,9 +1313,9 @@ def main():
 
     if args.save_model is not None:
         if args.save_model.endswith('gz'):
-            SURE.save_model(sure, args.save_model, compression=True)
+            PerturbFlow.save_model(perturbflow, args.save_model, compression=True)
         else:
-            SURE.save_model(sure, args.save_model)
+            PerturbFlow.save_model(perturbflow, args.save_model)