SURE-tools 1.0.2__py3-none-any.whl → 1.0.4__py3-none-any.whl

SURE/SURE.py

@@ -58,11 +58,11 @@ class SURE(nn.Module):
     ----------
     inpute_size
         Number of features (e.g., genes, peaks, proteins, etc.) per cell.
-    undesired_size
-        Number of undesired factors. It would be used to adjust for undesired variations like batch effect.
     codebook_size
         Number of metacells.
-    latent_dim
+    cell_factor_size
+        Number of cell-level factors. 
+    z_dim
         Dimensionality of latent states and metacells. 
     hidden_layers
         A list give the numbers of neurons for each hidden layer.
@@ -73,10 +73,7 @@ class SURE(nn.Module):
         * ``'negbinomial'`` -  negative binomial distribution (default)
         * ``'poisson'`` - poisson distribution
         * ``'multinomial'`` - multinomial distribution
-    user_dirichlet
-        A boolean option. If toggled on, SURE characterizes single-cell data using a hierarchical model, such as
-        dirichlet-negative binomial.
-    latent_dist
+    z_dist
         The distribution model for latent states. 
         
         One of the following:
@@ -100,7 +97,7 @@ class SURE(nn.Module):
                  supervised_mode: bool = False,
                  z_dim: int = 10,
                  z_dist: Literal['normal','studentt','laplacian','cauchy','gumbel'] = 'normal',
-                 loss_func: Literal['negbinomial','poisson','multinomial'] = 'negbinomial',
+                 loss_func: Literal['negbinomial','poisson','multinomial','bernoulli'] = 'negbinomial',
                  inverse_dispersion: float = 10.0,
                  use_zeroinflate: bool = True,
                  hidden_layers: list = [500],
@@ -368,9 +365,12 @@ class SURE(nn.Module):
 
             zs = zns
             concentrate = self.decoder_concentrate(zs)
-            rate = concentrate.exp()
-            if self.loss_func != 'poisson':
-                theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+            if self.loss_func == 'bernoulli':
+                log_theta = concentrate
+            else:
+                rate = concentrate.exp()
+                if self.loss_func == 'negbinomial':
+                    theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
 
             if self.loss_func == 'negbinomial':
                 if self.use_zeroinflate:
@@ -384,6 +384,11 @@ class SURE(nn.Module):
                     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)
+            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)
+                else:
+                    pyro.sample('x', dist.Bernoulli(logits=log_theta).to_event(1), obs=xs)
 
     def guide1(self, xs):
         with pyro.plate('data'):
@@ -442,9 +447,12 @@ class SURE(nn.Module):
                 zs = zns
 
             concentrate = self.decoder_concentrate(zs)
-            rate = concentrate.exp()
-            if self.loss_func != 'poisson':
-                theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+            if self.loss_func == 'bernoulli':
+                log_theta = concentrate
+            else:
+                rate = concentrate.exp()
+                if self.loss_func == 'negbinomial':
+                    theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
 
             if self.loss_func == 'negbinomial':
                 if self.use_zeroinflate:
@@ -458,6 +466,11 @@ class SURE(nn.Module):
                     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)
+            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)
+                else:
+                    pyro.sample('x', dist.Bernoulli(logits=log_theta).to_event(1), obs=xs)
 
     def guide2(self, xs, us=None):
         with pyro.plate('data'):
@@ -528,9 +541,12 @@ class SURE(nn.Module):
             zs = zns
 
             concentrate = self.decoder_concentrate(zs)
-            rate = concentrate.exp()
-            if self.loss_func != 'poisson':
-                theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+            if self.loss_func == 'bernoulli':
+                log_theta = concentrate
+            else:
+                rate = concentrate.exp()
+                if self.loss_func == 'negbinomial':
+                    theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
 
             if self.loss_func == 'negbinomial':
                 if self.use_zeroinflate:
@@ -544,6 +560,11 @@ class SURE(nn.Module):
                     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)
+            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)
+                else:
+                    pyro.sample('x', dist.Bernoulli(logits=log_theta).to_event(1), obs=xs)
 
     def guide3(self, xs, ys, embeds=None):
         with pyro.plate('data'):
@@ -616,9 +637,12 @@ class SURE(nn.Module):
                 zs = zns
 
             concentrate = self.decoder_concentrate(zs)
-            rate = concentrate.exp()
-            if self.loss_func != 'poisson':
-                theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
+            if self.loss_func == 'bernoulli':
+                log_theta = concentrate
+            else:
+                rate = concentrate.exp()
+                if self.loss_func == 'negbinomial':
+                    theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
 
             if self.loss_func == 'negbinomial':
                 if self.use_zeroinflate:
@@ -632,6 +656,11 @@ class SURE(nn.Module):
                     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)
+            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)
+                else:
+                    pyro.sample('x', dist.Bernoulli(logits=log_theta).to_event(1), obs=xs)
 
     def guide4(self, xs, us, ys, embeds=None):
         with pyro.plate('data'):
@@ -764,7 +793,14 @@ class SURE(nn.Module):
         A = np.concatenate(A)
         return A
     
-    def preprocess(self, xs):
+    def preprocess(self, xs, threshold=0):
+        if self.loss_func == 'bernoulli':
+            ad = sc.AnnData(xs)
+            binarize(ad, threshold=threshold)
+            xs = ad.X.copy()
+        else:
+            xs = np.round(xs)
+            
         if sparse.issparse(xs):
             xs = xs.toarray()
         return xs 
@@ -781,6 +817,7 @@ class SURE(nn.Module):
             weight_decay: float = 0.005, 
             decay_rate: float = 0.9,
             config_enum: str = 'parallel',
+            threshold: int = 0,
             use_jax: bool = False):
         """
         Train the SURE model.
@@ -790,7 +827,7 @@ class SURE(nn.Module):
         xs
             Single-cell experssion matrix. It should be a Numpy array or a Pytorch Tensor. Rows are cells and columns are features.
         us
-            Undesired factor matrix. It should be a Numpy array or a Pytorch Tensor. Rows are cells and columns are undesired factors.
+            cell-level factor matrix. 
         ys
             Desired factor matrix. It should be a Numpy array or a Pytorch Tensor. Rows are cells and columns are desired factors.
         num_epochs
@@ -811,7 +848,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.
         """
-        xs = self.preprocess(xs)
+        xs = self.preprocess(xs, threshold=threshold)
         xs = convert_to_tensor(xs, dtype=self.dtype, device=self.get_device())
         if us is not None:
             us = convert_to_tensor(us, dtype=self.dtype, device=self.get_device())
@@ -964,11 +1001,11 @@ def parse_args():
         help="the data file",
     )
     parser.add_argument(
-        "-undesired",
-        "--undesired-factor-file",
+        "-cf",
+        "--cell-factor-file",
         default=None,
         type=str,
-        help="the file for the record of undesired factors",
+        help="the file for the record of cell-level factors",
     )
     parser.add_argument(
         "-delta",
@@ -1148,18 +1185,18 @@ def main():
 
     xs = dt.fread(file=args.data_file, header=True).to_numpy()
     us = None 
-    if args.undesired_factor_file is not None:
-        us = dt.fread(file=args.undesired_factor_file, header=True).to_numpy()
+    if args.cell_factor_file is not None:
+        us = dt.fread(file=args.cell_factor_file, header=True).to_numpy()
 
     input_size = xs.shape[1]
-    undesired_size = 0 if us is None else us.shape[1]
+    cell_factor_size = 0 if us is None else us.shape[1]
 
     latent_dist = args.z_dist
 
     ###########################################
     sure = SURE(
         input_size=input_size,
-        undesired_size=undesired_size,
+        cell_factor_size=cell_factor_size,
         inverse_dispersion=args.inverse_dispersion,
         latent_dim=args.latent_dim,
         hidden_layers=args.hidden_layers,

{SURE_tools-1.0.2.dist-info → sure_tools-1.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: SURE-tools
-Version: 1.0.2
+Version: 1.0.4
 Summary: Succinct Representation of Single Cells
 Home-page: https://github.com/ZengFLab/SURE
 Author: Feng Zeng
@@ -26,6 +26,16 @@ Requires-Dist: networkx
 Requires-Dist: matplotlib
 Requires-Dist: seaborn
 Requires-Dist: fa2-modified
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 # SURE: SUccinct REpresentation of cells
  SURE introduces a vector quantization-based probabilistic generative model for calling metacells and use them as landmarks that form a coordinate system for cell ID. Analyzing single-cell omics data in a manner analogous to reference genome-based genomic analysis.

{SURE_tools-1.0.2.dist-info → sure_tools-1.0.4.dist-info}/RECORD

@@ -1,4 +1,4 @@
-SURE/SURE.py,sha256=AMU2EZKJVIE_n4W_K4crVvcpvFOtiiplczoGOCPvvPY,46959
+SURE/SURE.py,sha256=YhsWt3ndKpiIngKTjKOt58_jNyaBDF0wsoeIVNg2Di0,48758
 SURE/__init__.py,sha256=SbIRwAVBnNhza9vbsUH4N04atr0q_Abp04pCUTBhNio,127
 SURE/assembly/__init__.py,sha256=jxZLURXKPzXe21LhrZ09LgZr33iqdjlQy4oSEj5gR2Q,172
 SURE/assembly/assembly.py,sha256=6IMdelPOiRO4mUb4dC7gVCoF1Uvfw86-Map8P_jnUag,21477
@@ -9,9 +9,9 @@ SURE/utils/__init__.py,sha256=Htqv4KqVKcRiaaTBsR-6yZ4LSlbhbzutjNKXGD9-uds,660
 SURE/utils/custom_mlp.py,sha256=07TYX1HgxfEjb_3i5MpiZfNhOhx3dKntuwGkrpteWiM,7036
 SURE/utils/queue.py,sha256=E_5PA5EWcBoGAZj8BkKQnkCK0p4C-4-xcTPqdIXaPXU,1892
 SURE/utils/utils.py,sha256=IUHjDDtYaAYllCWsZyIzqQwaLul6fJRvHRH4vIYcR-c,8462
-SURE_tools-1.0.2.dist-info/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
-SURE_tools-1.0.2.dist-info/METADATA,sha256=Fy0Bc3luEPlFIfMRunpTkb18J_tPoqKK4jsfPFnRAlo,2431
-SURE_tools-1.0.2.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-SURE_tools-1.0.2.dist-info/entry_points.txt,sha256=u12payZYgCBy5FCwRHP6AlSQhKCiWSEDwj68r1DVdn8,40
-SURE_tools-1.0.2.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
-SURE_tools-1.0.2.dist-info/RECORD,,
+sure_tools-1.0.4.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
+sure_tools-1.0.4.dist-info/METADATA,sha256=2dPXR-pUr_8fNXewgDEo4oueSuRE0If05GQU94wypEo,2650
+sure_tools-1.0.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sure_tools-1.0.4.dist-info/entry_points.txt,sha256=u12payZYgCBy5FCwRHP6AlSQhKCiWSEDwj68r1DVdn8,40
+sure_tools-1.0.4.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
+sure_tools-1.0.4.dist-info/RECORD,,

{SURE_tools-1.0.2.dist-info → sure_tools-1.0.4.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.1.0)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any