SURE-tools 2.4.34__py3-none-any.whl → 2.4.42__py3-none-any.whl

SURE/DensityFlow.py

@@ -57,7 +57,7 @@ def set_random_seed(seed):
 class DensityFlow(nn.Module):
     def __init__(self,
                  input_size: int,
-                 codebook_size: int = 100,
+                 codebook_size: int = 30,
                  cell_factor_size: int = 0,
                  turn_off_cell_specific: bool = False,
                  supervised_mode: bool = False,
@@ -113,6 +113,7 @@ class DensityFlow(nn.Module):
         self.setup_networks()
         
         print(f"🧬 DensityFlow Initialized:")
+        print(f"   - Codebook size: {self.code_size}")
         print(f"   - Latent Dimension: {self.latent_dim}")
         print(f"   - Gene Dimension: {self.input_size}")
         print(f"   - Hidden Dimensions: {self.hidden_layers}")
@@ -397,7 +398,7 @@ class DensityFlow(nn.Module):
                     rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
-                logits = (mu.log()-dispersion.log()).clamp(min=-15, max=15)
+                logits = (mu.log()-dispersion.log()).clamp(min=-10, max=10)
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=dispersion, 
                                                                                          logits=logits),gate_logits=gate_logits).to_event(1), obs=xs)
@@ -486,7 +487,7 @@ class DensityFlow(nn.Module):
                     rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
-                logits = (mu.log()-dispersion.log()).clamp(min=-15, max=15)
+                logits = (mu.log()-dispersion.log()).clamp(min=-10, max=10)
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=dispersion, logits=logits),gate_logits=gate_logits).to_event(1), obs=xs)
                 else:
@@ -585,7 +586,7 @@ class DensityFlow(nn.Module):
                     rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
-                logits = (mu.log()-dispersion.log()).clamp(min=-15, max=15)
+                logits = (mu.log()-dispersion.log()).clamp(min=-10, max=10)
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=dispersion, 
                                                                                          logits=logits),gate_logits=gate_logits).to_event(1), obs=xs)
@@ -695,7 +696,7 @@ class DensityFlow(nn.Module):
                     rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
-                logits = (mu.log()-dispersion.log()).clamp(min=-15, max=15)
+                logits = (mu.log()-dispersion.log()).clamp(min=-10, max=10)
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=dispersion, 
                                                                                          logits=logits),gate_logits=gate_logits).to_event(1), obs=xs)
@@ -758,6 +759,28 @@ class DensityFlow(nn.Module):
         cb = self._get_codebook()
         cb = tensor_to_numpy(cb)
         return cb
+    
+    def _get_complete_embedding(self, xs, us):
+        basal,_ = self._get_basal_embedding(xs)
+        dzs = self._total_effects(basal, us)
+        return basal + dzs 
+    
+    def get_complete_embedding(self, xs, us, batch_size:int=1024):
+        xs = self.preprocess(xs)
+        xs = convert_to_tensor(xs, device=self.get_device())
+        us = convert_to_tensor(us, device=self.get_device())
+        dataset = CustomDataset2(xs, us)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+        Z = []
+        with tqdm(total=len(dataloader), desc='', unit='batch') as pbar:
+            for X_batch, U_batch, _ in dataloader:
+                zns = self._get_complete_embedding(X_batch, U_batch)
+                Z.append(tensor_to_numpy(zns))
+                pbar.update(1)
+
+        Z = np.concatenate(Z)
+        return Z
 
     def _get_basal_embedding(self, xs):           
         loc, scale = self.encoder_zn(xs)
@@ -1008,7 +1031,7 @@ class DensityFlow(nn.Module):
             ad = sc.AnnData(xs)
             binarize(ad, threshold=threshold)
             xs = ad.X.copy()
-        else:
+        elif self.loss_func == 'poisson':
             xs = np.round(xs)
             
         if sparse.issparse(xs):
@@ -1142,7 +1165,7 @@ class DensityFlow(nn.Module):
                 pbar.set_postfix({'loss': str_loss})
                 pbar.update(1)
                 
-    '''@classmethod
+    @classmethod
     def save_model(cls, model, file_path, compression=False):
         """Save the model to the specified file path."""
         file_path = os.path.abspath(file_path)
@@ -1169,10 +1192,18 @@ class DensityFlow(nn.Module):
         else:
             with open(file_path, 'rb') as pickle_file:
                 model = pickle.load(pickle_file)
+                
+        print(f"🧬 DensityFlow Initialized:")
+        print(f"   - Codebook size: {model.code_size}")
+        print(f"   - Latent Dimension: {model.latent_dim}")
+        print(f"   - Gene Dimension: {model.input_size}")
+        print(f"   - Hidden Dimensions: {model.hidden_layers}")
+        print(f"   - Device: {model.get_device()}")
+        print(f"   - Parameters: {sum(p.numel() for p in model.parameters()):,}")
         
-        return model'''
+        return model
         
-    def save_model(self, path):
+'''    def save(self, path):
         """Save model checkpoint"""
         torch.save({
             'model_state_dict': self.state_dict(),
@@ -1209,7 +1240,7 @@ class DensityFlow(nn.Module):
         checkpoint = torch.load(model_path, map_location=model.get_device())
         model.load_state_dict(checkpoint['model_state_dict'])
         
-        return model
+        return model'''
 
         
 EXAMPLE_RUN = (

{sure_tools-2.4.34.dist-info → sure_tools-2.4.42.dist-info}/METADATA

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

{sure_tools-2.4.34.dist-info → sure_tools-2.4.42.dist-info}/RECORD

@@ -1,4 +1,4 @@
-SURE/DensityFlow.py,sha256=KsngVc_ciwNAxXQZD4JFqJ-sDFLU_3Ra6Nyjcs6NvWo,58878
+SURE/DensityFlow.py,sha256=hYuPv1X9lsbKBIAGNRfzfCRYE1szCZPVJkOnIlK-Oc4,60246
 SURE/DensityFlow2.py,sha256=BBRCoA4NpU4EjghToOvowo17UtwYokTN75KxWYHTX1E,58404
 SURE/DensityFlowLinear.py,sha256=bYiPHJ6mza4sOXUjlFq7wButu3rNLYZuqWUTtIO06F4,57540
 SURE/EfficientTranscriptomeDecoder.py,sha256=O_x-4edKBU5OJJbOOS-59u3TQElZqhAtOVJMPlpw8m0,21667
@@ -25,9 +25,9 @@ SURE/utils/__init__.py,sha256=YF5jB-PAHJQ40OlcZ7BCZbsN2q1JKuPT6EppilRXQqM,680
 SURE/utils/custom_mlp.py,sha256=Rn_PQouxPMSda-KKBYrwVVv3GFFuUmCLxp8cV5LszZo,10580
 SURE/utils/queue.py,sha256=E_5PA5EWcBoGAZj8BkKQnkCK0p4C-4-xcTPqdIXaPXU,1892
 SURE/utils/utils.py,sha256=IUHjDDtYaAYllCWsZyIzqQwaLul6fJRvHRH4vIYcR-c,8462
-sure_tools-2.4.34.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
-sure_tools-2.4.34.dist-info/METADATA,sha256=dRu5kHcrc-T8DPNWOoz5_aNQM_1tc0sJ3BR40ZHGkec,2678
-sure_tools-2.4.34.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sure_tools-2.4.34.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
-sure_tools-2.4.34.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
-sure_tools-2.4.34.dist-info/RECORD,,
+sure_tools-2.4.42.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
+sure_tools-2.4.42.dist-info/METADATA,sha256=g9G4ftWszZ4prtVpgsS3GVIhY-nvh1g8zhOunpDGiqg,2678
+sure_tools-2.4.42.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sure_tools-2.4.42.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
+sure_tools-2.4.42.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
+sure_tools-2.4.42.dist-info/RECORD,,