SURE-tools 2.4.32__py3-none-any.whl → 2.4.43__py3-none-any.whl

SURE/DensityFlow.py

@@ -57,12 +57,12 @@ 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,
                  z_dim: int = 50,
-                 z_dist: Literal['normal','studentt','laplacian','cauchy','gumbel'] = 'gumbel',
+                 z_dist: Literal['normal','studentt','laplacian','cauchy','gumbel'] = 'studentt',
                  loss_func: Literal['negbinomial','poisson','multinomial','bernoulli'] = 'negbinomial',
                  dispersion: float = 8.0,
                  use_zeroinflate: bool = False,
@@ -85,6 +85,7 @@ class DensityFlow(nn.Module):
         self.latent_dim = z_dim
         self.hidden_layers = hidden_layers
         self.decoder_hidden_layers = hidden_layers[::-1]
+        self.config_enum = config_enum
         self.allow_broadcast = config_enum == 'parallel'
         self.use_cuda = use_cuda
         self.loss_func = loss_func
@@ -107,10 +108,12 @@ class DensityFlow(nn.Module):
         
         self.codebook_weights = None
 
+        self.seed = seed
         set_random_seed(seed)
         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}")
@@ -395,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)
@@ -484,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:
@@ -583,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)
@@ -693,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)
@@ -756,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)
@@ -1167,8 +1192,55 @@ 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
+        
+'''    def save(self, path):
+        """Save model checkpoint"""
+        torch.save({
+            'model_state_dict': self.state_dict(),
+            'model_config': {
+                'input_size': self.input_size,
+                'codebook_size': self.code_size,
+                'cell_factor_size': self.cell_factor_size,
+                'turn_off_cell_specific':self.turn_off_cell_specific,
+                'supervised_mode':self.supervised_mode,
+                'z_dim': self.latent_dim,
+                'z_dist': self.latent_dist,
+                'loss_func': self.loss_func,
+                'dispersion': self.dispersion,
+                'use_zeroinflate': self.use_zeroinflate,
+                'hidden_layers':self.hidden_layers,
+                'hidden_layer_activation':self.hidden_layer_activation,
+                'nn_dropout':self.nn_dropout,
+                'post_layer_fct':self.post_layer_fct,
+                'post_act_fct':self.post_act_fct,
+                'config_enum':self.config_enum,
+                'use_cuda':self.use_cuda,
+                'seed':self.seed,
+                'zero_bias':self.use_bias,
+                'dtype':self.dtype,
+            }
+        }, path)
+    
+    @classmethod
+    def load_model(cls, model_path: str):
+        """Load pre-trained model"""
+        checkpoint = torch.load(model_path)
+        model = DensityFlow(**checkpoint.get('model_config'))
+        
+        checkpoint = torch.load(model_path, map_location=model.get_device())
+        model.load_state_dict(checkpoint['model_state_dict'])
+        
+        return model'''
 
         
 EXAMPLE_RUN = (

SURE/DensityFlowLinear.py

@@ -107,6 +107,7 @@ class DensityFlowLinear(nn.Module):
         
         self.codebook_weights = None
 
+        self.seed = seed
         set_random_seed(seed)
         self.setup_networks()
         

{sure_tools-2.4.32.dist-info → sure_tools-2.4.43.dist-info}/METADATA

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

{sure_tools-2.4.32.dist-info → sure_tools-2.4.43.dist-info}/RECORD

@@ -1,6 +1,6 @@
-SURE/DensityFlow.py,sha256=M8yDMCzH-LSVjFB34pUaN2TNUqo0jnZMmaFcLLjtYBU,57218
+SURE/DensityFlow.py,sha256=fqqI8sHnfXuTK9O1il-dL7F4W7gbUMjGHD8uRwpESlc,60218
 SURE/DensityFlow2.py,sha256=BBRCoA4NpU4EjghToOvowo17UtwYokTN75KxWYHTX1E,58404
-SURE/DensityFlowLinear.py,sha256=RfjIwXxO0gWQni7LHRy4DSDUkdz3HpDjlZbxJse65Ts,57515
+SURE/DensityFlowLinear.py,sha256=bYiPHJ6mza4sOXUjlFq7wButu3rNLYZuqWUTtIO06F4,57540
 SURE/EfficientTranscriptomeDecoder.py,sha256=O_x-4edKBU5OJJbOOS-59u3TQElZqhAtOVJMPlpw8m0,21667
 SURE/PerturbE.py,sha256=DxEp-qef--x8-GMZdPfBf8ts8UDDc34h2P5AnpqZ-YM,52265
 SURE/PerturbationAwareDecoder.py,sha256=duhvBvZjOpAk7c2YTfmA2qKbrgVvwT7IW1pxaukq_iU,30231
@@ -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.32.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
-sure_tools-2.4.32.dist-info/METADATA,sha256=9o2WcDEY2JIf1AhclsLePndrjYVexLj9QUHtZjzGc38,2678
-sure_tools-2.4.32.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sure_tools-2.4.32.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
-sure_tools-2.4.32.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
-sure_tools-2.4.32.dist-info/RECORD,,
+sure_tools-2.4.43.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
+sure_tools-2.4.43.dist-info/METADATA,sha256=q0DTzGgBqj5Hi8n2YNmJymHD25dZSUdRVlMrfiy-5Hw,2678
+sure_tools-2.4.43.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sure_tools-2.4.43.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
+sure_tools-2.4.43.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
+sure_tools-2.4.43.dist-info/RECORD,,