SURE-tools 2.0.10__tar.gz → 2.1.1__tar.gz

{sure_tools-2.0.10 → sure_tools-2.1.1}/PKG-INFO

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

sure_tools-2.0.10/SURE/SURE.py → sure_tools-2.1.1/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
@@ -97,7 +97,6 @@ class SURE(nn.Module):
                  input_size: int,
                  codebook_size: int = 200,
                  cell_factor_size: int = 0,
-                 cell_factor_names: list = None,
                  supervised_mode: bool = False,
                  z_dim: int = 10,
                  z_dist: Literal['normal','studentt','laplacian','cauchy','gumbel'] = 'normal',
@@ -135,7 +134,6 @@ class SURE(nn.Module):
         self.post_layer_fct = post_layer_fct
         self.post_act_fct = post_act_fct
         self.hidden_layer_activation = hidden_layer_activation
-        self.cell_factor_names = cell_factor_names
         
         self.codebook_weights = None
 
@@ -330,7 +328,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 +405,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 +493,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 +585,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,11 +811,8 @@ class SURE(nn.Module):
         A = np.concatenate(A)
         return A
     
-    def _cell_move(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])
-        
+    def _cell_response(self, xs, factor_idx, perturb):
+        zns,_ = self.encoder_zn(xs)       
         if perturb.ndim==2:
             ms = self.cell_factor_effect[factor_idx]([zns, perturb])
         else:
@@ -825,7 +820,7 @@ class SURE(nn.Module):
             
         return ms 
 
-    def get_cell_move(self, 
+    def get_cell_response(self, 
                              xs, 
                              factor_idx,
                              perturb,
@@ -843,27 +838,23 @@ 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_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())
-        
-        if type(factor_idx) == str:
-            factor_idx = int(np.where(self.cell_factor_names==factor_idx)[0])
-            
+        ps = convert_to_tensor(perturb, device=self.get_device())         
         ms = self.cell_factor_effect[factor_idx]([zs,ps])
         return tensor_to_numpy(ms)
     
-    def _get_expression_responses(self, delta_zs):
+    def _get_expression_response(self, delta_zs):
         return self.decoder_concentrate(delta_zs)
     
-    def get_expression_responses(self, 
+    def get_expression_response(self, 
                              delta_zs, 
                              batch_size: int = 1024):
         """
@@ -911,7 +902,7 @@ class SURE(nn.Module):
             threshold: int = 0,
             use_jax: bool = False):
         """
-        Train the SURE model.
+        Train the PerturbFlow model.
 
         Parameters
         ----------
@@ -937,7 +928,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())
@@ -1055,12 +1046,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"
@@ -1278,7 +1269,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,
@@ -1301,7 +1292,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,
@@ -1313,9 +1304,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)