SURE-tools 4.0.1__py3-none-any.whl → 4.0.3__py3-none-any.whl

SURE/SURE.py

@@ -79,7 +79,7 @@ class SURE(nn.Module):
         Number of features (e.g., genes, peaks, proteins, etc.) per cell.
     codebook_size
         Number of metacells.
-    covariate_size
+    covariate_sizes
         Number of cell-level factors. 
     transforms
         Number of neural spline flows
@@ -111,7 +111,7 @@ class SURE(nn.Module):
                  codebook_size: int,
                  context_sizes: list = [0],
                  perturb_size: int = 0,
-                 covariate_size: int = 0,
+                 covariate_sizes: int = 0,
                  method: Literal['flow','vae'] = 'vae',
                  transforms: int = 1,
                  z_dim: int = 30,
@@ -137,7 +137,7 @@ class SURE(nn.Module):
                                  codebook_size=codebook_size,
                                  context_sizes=context_sizes,
                                  perturb_size=perturb_size,
-                                 covariate_size=covariate_size,
+                                 covariate_sizes=covariate_sizes,
                                  transforms=transforms,
                                  z_dim=z_dim,
                                  z_dist=z_dist,
@@ -159,7 +159,7 @@ class SURE(nn.Module):
                                   codebook_size=codebook_size,
                                   context_sizes=context_sizes,
                                   perturb_size=perturb_size,
-                                  covariate_size=covariate_size,
+                                  covariate_sizes=covariate_sizes,
                                   z_dim=z_dim,
                                   z_dist=z_dist,
                                   loss_func=loss_func,
@@ -245,7 +245,7 @@ class SURE(nn.Module):
     def fit(self, xs:np.array, 
             css:list = None, 
             ps:np.array = None,
-            fs:np.array = None,
+            fss:list = None,
             num_epochs: int = 100, 
             learning_rate: float = 0.0001, 
             use_mask: bool = False,
@@ -291,7 +291,7 @@ class SURE(nn.Module):
             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.
         """
-        self.engine.fit(xs=xs, css=css, ps=ps, fs=fs, num_epochs=num_epochs, learning_rate=learning_rate, use_mask=use_mask, mask_ratio=mask_ratio, batch_size=batch_size, algo=algo,
+        self.engine.fit(xs=xs, css=css, ps=ps, fss=fss, num_epochs=num_epochs, learning_rate=learning_rate, use_mask=use_mask, mask_ratio=mask_ratio, batch_size=batch_size, algo=algo,
                         beta_1=beta_1, weight_decay=weight_decay, decay_rate=decay_rate, config_enum=config_enum, threshold=threshold,
                         use_jax=use_jax, show_progress=show_progress, patience=patience, min_delta=min_delta, restore_best_weights=restore_best_weights,
                         monitor=monitor)

SURE/SURE_nsf.py

@@ -89,7 +89,7 @@ class SURENF(nn.Module):
         Number of features (e.g., genes, peaks, proteins, etc.) per cell.
     codebook_size
         Number of metacells.
-    covariate_size
+    covariate_sizes
         Number of cell-level factors. 
     transforms
         Number of neural spline flows
@@ -121,7 +121,7 @@ class SURENF(nn.Module):
                  codebook_size: int,
                  context_sizes: list = [0],
                  perturb_size: int = 0,
-                 covariate_size: int = 0,
+                 covariate_sizes: int = 0,
                  transforms: int = 1,
                  z_dim: int = 50,
                  z_dist: Literal['normal','studentt','laplacian','cauchy','gumbel'] = 'studentt',
@@ -144,7 +144,7 @@ class SURENF(nn.Module):
         self.input_dim = input_dim
         self.context_sizes = context_sizes
         self.perturb_size = perturb_size
-        self.covariate_size = covariate_size
+        self.covariate_sizes = covariate_sizes
         self.dispersion = dispersion
         self.latent_dim = z_dim
         self.latent_dist = z_dist
@@ -272,16 +272,19 @@ class SURENF(nn.Module):
                 allow_broadcast=self.allow_broadcast,
                 use_cuda=self.use_cuda,
             )
-        if self.covariate_size>0:
-            self.covariate_effect = ZeroBiasMLP2(
-                [self.covariate_size] + self.decoder_hidden_layers + [self.latent_dim],
-                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 np.sum(self.covariate_sizes)>0:
+            self.covariate_effects = nn.ModuleList()
+            for covariate_size in self.covariate_sizes:
+                self.covariate_effects.append(ZeroBiasMLP2(
+                        [covariate_size] + self.decoder_hidden_layers + [self.latent_dim],
+                        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,
+                    )
+                )
             
         self.decoder_log_mu = MLP(
                 [self.latent_dim+self.latent_dim+self.latent_dim] + self.decoder_hidden_layers + [self.input_dim],
@@ -421,8 +424,13 @@ class SURENF(nn.Module):
                 zps = self.perturb_effect([ps, zs+zcs])
             else:
                 zps = torch.zeros_like(zs)
-            if (self.covariate_size>0) and (fs is not None):
-                zfs = self.covariate_effect(fs)
+            if (self.covariate_sizes>0) and (fs is not None):
+                zfs = torch.zeros_like(zs)
+                shift = 0
+                for i, covariate_size in enumerate(self.covariate_sizes):
+                    fs_i = fs[:,shift:(shift+covariate_size)]
+                    zfs += self.covariate_effects[i](fs_i)
+                    shift += covariate_size
             else:
                 zfs = torch.zeros_like(zs)
 
@@ -891,7 +899,7 @@ class SURENF(nn.Module):
     def fit(self, xs:np.array, 
             css:list = None, 
             ps: np.array = None,
-            fs:np.array = None,
+            fss:list = None,
             num_epochs: int = 100, 
             learning_rate: float = 0.0001, 
             use_mask: bool = False,
@@ -957,7 +965,8 @@ class SURENF(nn.Module):
             cs = convert_to_tensor(cs, dtype=self.dtype, device='cpu')
         if ps is not None:
             ps = convert_to_tensor(ps, dtype=self.dtype, device='cpu')
-        if fs is not None:
+        if fss is not None:
+            fs = np.hstack(fss)
             fs = convert_to_tensor(fs, dtype=self.dtype, device='cpu')
 
         dataset = CustomDataset(xs)
@@ -1019,7 +1028,7 @@ class SURENF(nn.Module):
                             batch_p = None
                         else:
                             batch_p = ps[idx].to(self.get_device())
-                        if fs is None:
+                        if fss is None:
                             batch_f = None
                         else:
                             batch_f = fs[idx].to(self.get_device())

SURE/SURE_vae.py

@@ -121,7 +121,7 @@ class SUREVAE(nn.Module):
                  codebook_size: int,
                  context_sizes: list = [0],
                  perturb_size: int = 0,
-                 covariate_size: int = 0,
+                 covariate_sizes: int = 0,
                  transforms: int = 1,
                  z_dim: int = 50,
                  z_dist: Literal['normal','studentt','laplacian','cauchy','gumbel'] = 'studentt',
@@ -144,7 +144,7 @@ class SUREVAE(nn.Module):
         self.input_dim = input_dim
         self.context_sizes = context_sizes
         self.perturb_size = perturb_size
-        self.covariate_size = covariate_size
+        self.covariate_sizes = covariate_sizes
         self.dispersion = dispersion
         self.latent_dim = z_dim
         self.latent_dist = z_dist
@@ -273,7 +273,7 @@ class SUREVAE(nn.Module):
         
         if self.perturb_size>0:
             self.perturb_effect = ZeroBiasMLP3(
-                [self.perturb_size+self.latent_dim] + self.decoder_hidden_layers + [self.latent_dim],
+                [self.perturb_size+self.latent_dim+self.latent_dim] + self.decoder_hidden_layers + [self.latent_dim],
                 activation=activate_fct,
                 output_activation=None,
                 post_layer_fct=post_layer_fct,
@@ -282,16 +282,19 @@ class SUREVAE(nn.Module):
                 use_cuda=self.use_cuda,
             )
             
-        if self.covariate_size>0:
-            self.covariate_effect = ZeroBiasMLP2(
-                [self.covariate_size] + self.decoder_hidden_layers + [self.latent_dim],
-                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 np.sum(self.covariate_sizes)>0:
+            self.covariate_effects = nn.ModuleList()
+            for covariate_size in self.covariate_sizes:
+                self.covariate_effects.append(ZeroBiasMLP2(
+                        [covariate_size] + self.decoder_hidden_layers + [self.latent_dim],
+                        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,
+                    )
+                )
             
         self.decoder_log_mu = MLP(
                 [self.latent_dim] + self.decoder_hidden_layers + [self.input_dim],
@@ -428,11 +431,16 @@ class SUREVAE(nn.Module):
             else:
                 zcs = torch.zeros_like(zs)
             if (self.perturb_size>0) and (ps is not None):
-                zps = self.perturb_effect([ps, zs+zcs])
+                zps = self.perturb_effect([ps, zs, zcs])
             else:
                 zps = torch.zeros_like(zs)
-            if (self.covariate_size>0) and (fs is not None):
-                zfs = self.covariate_effect(fs)
+            if (np.sum(self.covariate_sizes)>0) and (fs is not None):
+                zfs = torch.zeros_like(zs)
+                shift = 0
+                for i, covariate_size in enumerate(self.covariate_sizes):
+                    fs_i = fs[:,shift:(shift+covariate_size)]
+                    zfs += self.covariate_effects[i](fs_i)
+                    shift += covariate_size
             else:
                 zfs = torch.zeros_like(zs)
 
@@ -640,7 +648,7 @@ class SUREVAE(nn.Module):
                 C_batch = zcs[idx].to(self.get_device())
                 P_batch = ps[idx].to(self.get_device())
                 
-                dzs = self.perturb_effect([P_batch,Z_batch+C_batch])
+                dzs = self.perturb_effect([P_batch,Z_batch,C_batch])
                 
                 A.append(tensor_to_numpy(dzs))
                 pbar.update(1)
@@ -693,7 +701,7 @@ class SUREVAE(nn.Module):
                     
                 if ps is not None:
                     P_batch = ps[idx].to(self.get_device())
-                    zps = self.perturb_effect([P_batch,z_basal+zcs])
+                    zps = self.perturb_effect([P_batch,z_basal,zcs])
                 else:
                     zps = torch.zeros_like(z_basal)
                     
@@ -825,7 +833,7 @@ class SUREVAE(nn.Module):
     def fit(self, xs: np.array, 
             css: list = None, 
             ps: np.array = None,
-            fs: np.array = None,
+            fss: list = None,
             num_epochs: int = 100, 
             learning_rate: float = 0.0001, 
             use_mask: bool = False,
@@ -891,7 +899,8 @@ class SUREVAE(nn.Module):
             cs = convert_to_tensor(cs, dtype=self.dtype, device='cpu')
         if ps is not None:
             ps = convert_to_tensor(ps, dtype=self.dtype, device='cpu')
-        if fs is not None:
+        if fss is not None:
+            fs = np.hstack(fss)
             fs = convert_to_tensor(fs, dtype=self.dtype, device='cpu')
 
         dataset = CustomDataset(xs)
@@ -953,7 +962,7 @@ class SUREVAE(nn.Module):
                             batch_p = None
                         else:
                             batch_p = ps[idx].to(self.get_device())
-                        if fs is None:
+                        if fss is None:
                             batch_f = None
                         else:
                             batch_f = fs[idx].to(self.get_device())

{sure_tools-4.0.1.dist-info → sure_tools-4.0.3.dist-info}/METADATA

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

{sure_tools-4.0.1.dist-info → sure_tools-4.0.3.dist-info}/RECORD

@@ -1,7 +1,7 @@
-SURE/SURE.py,sha256=8E39np6zhLbT1cp1xYOg5xLwzFHBupKIm1ydLxNKtqM,16654
+SURE/SURE.py,sha256=YcJ6Lxt4ziQXZ5Yk02Ss-fJRl_M1Z0QghLY_yeHt5Pc,16659
 SURE/SUREMO.py,sha256=hN0G0ZEBNQdmj0gGlBIy1wjKGKqAMUemDjQeICzvNUY,47644
-SURE/SURE_nsf.py,sha256=VR6YgEiIfu7mRH0XLbovfjd2X3WFuom2J4_AgHrq2dM,49040
-SURE/SURE_vae.py,sha256=b3xluu49kmuDlcpZnTAJc92SvQI8MPDN9ntBfnKTJOQ,45850
+SURE/SURE_nsf.py,sha256=uKpau3nDdlX58lLzPzecZULKobyOYB7xKtcYo6vwbak,49541
+SURE/SURE_vae.py,sha256=X4aNb6Wv9Xn-w1Xc-wUuio6YsxX0dIwdJqsWG44hCa0,46375
 SURE/SURE_vanilla.py,sha256=I1RHHCpzk8ml1vMdH_gITOzAFXrHYUA7IAwbVSoxmBo,27327
 SURE/__init__.py,sha256=eJN0vlGblWir1JHaoiQqbQHzS_C1PNYcA_ls7UviTqc,444
 SURE/atac/__init__.py,sha256=3smP8IKHfwNCd1G_sZH3pKHXuLkLpFuLtjUTUSy7_As,34
@@ -15,8 +15,8 @@ SURE/utils/custom_mlp.py,sha256=XvviNUYeaZ5D1evqETyWLxgtL56mmFQR_pxsYpKw0yY,1170
 SURE/utils/label.py,sha256=joKO1mSkjZXeLvSeC7GluQk4-_qgGgPqlwWixdcbKMQ,4648
 SURE/utils/queue.py,sha256=E_5PA5EWcBoGAZj8BkKQnkCK0p4C-4-xcTPqdIXaPXU,1892
 SURE/utils/utils.py,sha256=IUHjDDtYaAYllCWsZyIzqQwaLul6fJRvHRH4vIYcR-c,8462
-sure_tools-4.0.1.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
-sure_tools-4.0.1.dist-info/METADATA,sha256=m4l4x8KnzdyfZquikfwrSqZXegaeRl6X9plr-fcuGrc,1661
-sure_tools-4.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sure_tools-4.0.1.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
-sure_tools-4.0.1.dist-info/RECORD,,
+sure_tools-4.0.3.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
+sure_tools-4.0.3.dist-info/METADATA,sha256=DlTyOxPOevVRibU7zjiIorXK-Te5I5uMI3zec3H94zw,1661
+sure_tools-4.0.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sure_tools-4.0.3.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
+sure_tools-4.0.3.dist-info/RECORD,,