SURE-tools 2.4.25__tar.gz → 2.4.42__tar.gz

{sure_tools-2.4.25 → sure_tools-2.4.42}/PKG-INFO

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

{sure_tools-2.4.25 → sure_tools-2.4.42}/SURE/DensityFlow.py

@@ -57,16 +57,16 @@ def set_random_seed(seed):
 class DensityFlow(nn.Module):
     def __init__(self,
                  input_size: int,
-                 codebook_size: int = 200,
+                 codebook_size: int = 30,
                  cell_factor_size: int = 0,
                  turn_off_cell_specific: bool = False,
                  supervised_mode: bool = False,
-                 z_dim: int = 10,
-                 z_dist: Literal['normal','studentt','laplacian','cauchy','gumbel'] = 'gumbel',
-                 loss_func: Literal['negbinomial','poisson','multinomial','bernoulli'] = 'multinomial',
-                 inverse_dispersion: float = 10.0,
+                 z_dim: int = 50,
+                 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,
-                 hidden_layers: list = [500],
+                 hidden_layers: list = [1024],
                  hidden_layer_activation: Literal['relu','softplus','leakyrelu','linear'] = 'relu',
                  nn_dropout: float = 0.1,
                  post_layer_fct: list = ['layernorm'],
@@ -81,10 +81,11 @@ class DensityFlow(nn.Module):
 
         self.input_size = input_size
         self.cell_factor_size = cell_factor_size
-        self.inverse_dispersion = inverse_dispersion
+        self.dispersion = dispersion
         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}")
@@ -258,7 +261,7 @@ class DensityFlow(nn.Module):
                             )
                         )
             
-        self.decoder_concentrate = MLP(
+        self.decoder_log_mu = MLP(
                     [self.latent_dim] + self.decoder_hidden_layers + [self.input_size],
                     activation=activate_fct,
                     output_activation=None,
@@ -348,8 +351,8 @@ class DensityFlow(nn.Module):
         self.options = dict(dtype=xs.dtype, device=xs.device)
         
         if self.loss_func=='negbinomial':
-            total_count = pyro.param("inverse_dispersion", self.inverse_dispersion *
-                                     xs.new_ones(self.input_size), constraint=constraints.positive)
+            dispersion = pyro.param("dispersion", self.dispersion *
+                                            xs.new_ones(self.input_size), constraint=constraints.positive)
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -383,28 +386,32 @@ class DensityFlow(nn.Module):
                 zns = pyro.sample('zn', dist.Gumbel(zn_loc, zn_scale).to_event(1))
 
             zs = zns
-            concentrate = self.decoder_concentrate(zs)
+            log_mu = self.decoder_log_mu(zs)
             if self.loss_func in ['bernoulli']:
-                log_theta = concentrate
+                log_theta = log_mu
+            elif self.loss_func == 'negbinomial':
+                mu = log_mu.exp()
             else:
-                rate = concentrate.exp()
+                rate = log_mu.exp()
                 theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
                 if self.loss_func == 'poisson':
                     rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
+                logits = (mu.log()-dispersion.log()).clamp(min=-10, max=10)
                 if self.use_zeroinflate:
-                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=total_count, probs=theta),gate_logits=gate_logits).to_event(1), obs=xs)
+                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=dispersion, 
+                                                                                         logits=logits),gate_logits=gate_logits).to_event(1), obs=xs)
                 else:
-                    pyro.sample('x', dist.NegativeBinomial(total_count=total_count, probs=theta).to_event(1), obs=xs)
+                    pyro.sample('x', dist.NegativeBinomial(total_count=dispersion, 
+                                                           logits=logits).to_event(1), obs=xs)
             elif self.loss_func == 'poisson':
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
                 else:
                     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)
-                pyro.sample('x', dist.Multinomial(total_count=int(1e8), logits=concentrate), obs=xs)
+                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)
@@ -428,8 +435,8 @@ class DensityFlow(nn.Module):
         self.options = dict(dtype=xs.dtype, device=xs.device)
         
         if self.loss_func=='negbinomial':
-            total_count = pyro.param("inverse_dispersion", self.inverse_dispersion *
-                                     xs.new_ones(self.input_size), constraint=constraints.positive)
+            dispersion = pyro.param("dispersion", self.dispersion *
+                                            xs.new_ones(self.input_size), constraint=constraints.positive)
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -468,28 +475,30 @@ class DensityFlow(nn.Module):
             else:
                 zs = zns
 
-            concentrate = self.decoder_concentrate(zs)
+            log_mu = self.decoder_log_mu(zs)
             if self.loss_func in ['bernoulli']:
-                log_theta = concentrate
+                log_theta = log_mu
+            elif self.loss_func == 'negbinomial':
+                mu = log_mu.exp()
             else:
-                rate = concentrate.exp()
+                rate = log_mu.exp()
                 theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
                 if self.loss_func == 'poisson':
                     rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
+                logits = (mu.log()-dispersion.log()).clamp(min=-10, max=10)
                 if self.use_zeroinflate:
-                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=total_count, probs=theta),gate_logits=gate_logits).to_event(1), obs=xs)
+                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=dispersion, logits=logits),gate_logits=gate_logits).to_event(1), obs=xs)
                 else:
-                    pyro.sample('x', dist.NegativeBinomial(total_count=total_count, probs=theta).to_event(1), obs=xs)
+                    pyro.sample('x', dist.NegativeBinomial(total_count=dispersion, logits=logits).to_event(1), obs=xs)
             elif self.loss_func == 'poisson':
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
                 else:
                     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)
-                pyro.sample('x', dist.Multinomial(total_count=int(1e8), logits=concentrate), obs=xs)
+                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)
@@ -513,8 +522,8 @@ class DensityFlow(nn.Module):
         self.options = dict(dtype=xs.dtype, device=xs.device)
         
         if self.loss_func=='negbinomial':
-            total_count = pyro.param("inverse_dispersion", self.inverse_dispersion *
-                                     xs.new_ones(self.input_size), constraint=constraints.positive)
+            dispersion = pyro.param("dispersion", self.dispersion *
+                                            xs.new_ones(self.input_size), constraint=constraints.positive)
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -565,28 +574,31 @@ class DensityFlow(nn.Module):
 
             zs = zns
 
-            concentrate = self.decoder_concentrate(zs)
+            log_mu = self.decoder_log_mu(zs)
             if self.loss_func in ['bernoulli']:
-                log_theta = concentrate
+                log_theta = log_mu
+            elif self.loss_func in ['negbinomial']:
+                mu = log_mu.exp()
             else:
-                rate = concentrate.exp()
+                rate = log_mu.exp()
                 theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
                 if self.loss_func == 'poisson':
                     rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
+                logits = (mu.log()-dispersion.log()).clamp(min=-10, max=10)
                 if self.use_zeroinflate:
-                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=total_count, probs=theta),gate_logits=gate_logits).to_event(1), obs=xs)
+                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=dispersion, 
+                                                                                         logits=logits),gate_logits=gate_logits).to_event(1), obs=xs)
                 else:
-                    pyro.sample('x', dist.NegativeBinomial(total_count=total_count, probs=theta).to_event(1), obs=xs)
+                    pyro.sample('x', dist.NegativeBinomial(total_count=dispersion, logits=logits).to_event(1), obs=xs)
             elif self.loss_func == 'poisson':
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
                 else:
                     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)
-                pyro.sample('x', dist.Multinomial(total_count=int(1e8), logits=concentrate), obs=xs)
+                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)
@@ -610,8 +622,8 @@ class DensityFlow(nn.Module):
         self.options = dict(dtype=xs.dtype, device=xs.device)
         
         if self.loss_func=='negbinomial':
-            total_count = pyro.param("inverse_dispersion", self.inverse_dispersion *
-                                     xs.new_ones(self.input_size), constraint=constraints.positive)
+            dispersion = pyro.param("dispersion", self.dispersion *
+                                            xs.new_ones(self.input_size), constraint=constraints.positive)
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -672,28 +684,31 @@ class DensityFlow(nn.Module):
             else:
                 zs = zns
 
-            concentrate = self.decoder_concentrate(zs)
+            log_mu = self.decoder_log_mu(zs)
             if self.loss_func in ['bernoulli']:
-                log_theta = concentrate
+                log_theta = log_mu
+            elif self.loss_func in ['negbinomial']:
+                mu = log_mu.exp()
             else:
-                rate = concentrate.exp()
+                rate = log_mu.exp()
                 theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
                 if self.loss_func == 'poisson':
                     rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
+                logits = (mu.log()-dispersion.log()).clamp(min=-10, max=10)
                 if self.use_zeroinflate:
-                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=total_count, probs=theta),gate_logits=gate_logits).to_event(1), obs=xs)
+                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=dispersion, 
+                                                                                         logits=logits),gate_logits=gate_logits).to_event(1), obs=xs)
                 else:
-                    pyro.sample('x', dist.NegativeBinomial(total_count=total_count, probs=theta).to_event(1), obs=xs)
+                    pyro.sample('x', dist.NegativeBinomial(total_count=dispersion, logits=logits).to_event(1), obs=xs)
             elif self.loss_func == 'poisson':
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
                 else:
                     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)
-                pyro.sample('x', dist.Multinomial(total_count=int(1e8), logits=concentrate), obs=xs)
+                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)
@@ -717,13 +732,13 @@ class DensityFlow(nn.Module):
                 #    zus = self.cell_factor_effect[i](us[:,i].reshape(-1,1))
                 #else:
                 #    zus = self.cell_factor_effect[i]([zns,us[:,i].reshape(-1,1)])
-                zus = self._cell_response(zns, i, us[:,i].reshape(-1,1))
+                zus = self._cell_shift(zns, i, us[:,i].reshape(-1,1))
             else:
                 #if self.turn_off_cell_specific:
                 #    zus = zus + self.cell_factor_effect[i](us[:,i].reshape(-1,1))
                 #else:
                 #    zus = zus + self.cell_factor_effect[i]([zns,us[:,i].reshape(-1,1)])
-                zus = zus + self._cell_response(zns, i, us[:,i].reshape(-1,1))
+                zus = zus + self._cell_shift(zns, i, us[:,i].reshape(-1,1))
         return zus
                         
     def _get_codebook_identity(self):
@@ -744,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)
@@ -865,12 +902,12 @@ class DensityFlow(nn.Module):
             us_i = us[:,pert_idx].reshape(-1,1)
             
             # factor effect of xs 
-            dzs0 = self.get_cell_response(zs, perturb_idx=pert_idx, perturb_us=us_i)
+            dzs0 = self.get_cell_shift(zs, perturb_idx=pert_idx, perturb_us=us_i)
             
             # perturbation effect 
             ps = np.ones_like(us_i)
             if np.sum(np.abs(ps-us_i))>=1:
-                dzs = self.get_cell_response(zs, perturb_idx=pert_idx, perturb_us=ps)          
+                dzs = self.get_cell_shift(zs, perturb_idx=pert_idx, perturb_us=ps)          
                 zs = zs + dzs0 + dzs
             else:
                 zs = zs + dzs0
@@ -884,10 +921,11 @@ class DensityFlow(nn.Module):
             library_sizes = library_sizes.reshape(-1,1)
             
         counts = self.get_counts(zs, library_sizes=library_sizes)
+        log_mu = self.get_log_mu(zs)
         
-        return counts, zs
+        return counts, log_mu
     
-    def _cell_response(self, zs, perturb_idx, perturb):
+    def _cell_shift(self, zs, perturb_idx, perturb):
         #zns,_ = self.encoder_zn(xs)    
         #zns,_ = self._get_basal_embedding(xs)   
         zns = zs
@@ -904,7 +942,7 @@ class DensityFlow(nn.Module):
             
         return ms 
 
-    def get_cell_response(self, 
+    def get_cell_shift(self, 
                              zs, 
                              perturb_idx,
                              perturb_us,
@@ -922,46 +960,43 @@ class DensityFlow(nn.Module):
         Z = []
         with tqdm(total=len(dataloader), desc='', unit='batch') as pbar:
             for Z_batch, P_batch, _ in dataloader:
-                zns = self._cell_response(Z_batch, perturb_idx, P_batch)
+                zns = self._cell_shift(Z_batch, perturb_idx, P_batch)
                 Z.append(tensor_to_numpy(zns))
                 pbar.update(1)
 
         Z = np.concatenate(Z)
         return Z
     
-    def _get_expression_response(self, delta_zs):
-        return self.decoder_concentrate(delta_zs)
+    def _log_mu(self, zs):
+        return self.decoder_log_mu(zs)
     
-    def get_expression_response(self, 
-                             delta_zs, 
-                             batch_size: int = 1024):
+    def get_log_mu(self, zs, batch_size: int = 1024):
         """
         Return cells' changes in the feature space induced by specific perturbation of a factor
 
         """
-        delta_zs = convert_to_tensor(delta_zs, device=self.get_device())
-        dataset = CustomDataset(delta_zs)
+        zs = convert_to_tensor(zs, device=self.get_device())
+        dataset = CustomDataset(zs)
         dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
 
         R = []
         with tqdm(total=len(dataloader), desc='', unit='batch') as pbar:
-            for delta_Z_batch, _ in dataloader:
-                r = self._get_expression_response(delta_Z_batch)
+            for Z_batch, _ in dataloader:
+                r = self._log_mu(Z_batch)
                 R.append(tensor_to_numpy(r))
                 pbar.update(1)
 
         R = np.concatenate(R)
         return R
     
-    def _count(self, concentrate, library_size=None):
+    def _count(self, log_mu, library_size=None):
         if self.loss_func == 'bernoulli':
-            #counts = self.sigmoid(concentrate)
-            counts = dist.Bernoulli(logits=concentrate).to_event(1).mean
+            counts = dist.Bernoulli(logits=log_mu).to_event(1).mean
         elif self.loss_func == 'multinomial':
-            theta = dist.Multinomial(total_count=int(1e8), logits=concentrate).mean
+            theta = dist.Multinomial(total_count=int(1e8), logits=log_mu).mean
             counts = theta * library_size
         else:
-            rate = concentrate.exp()
+            rate = log_mu.exp()
             theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
             counts = theta * library_size
         return counts
@@ -983,8 +1018,8 @@ class DensityFlow(nn.Module):
         E = []
         with tqdm(total=len(dataloader), desc='', unit='batch') as pbar:
             for Z_batch, L_batch, _ in dataloader:
-                concentrate = self._get_expression_response(Z_batch)
-                counts = self._count(concentrate, L_batch)
+                log_mu = self._log_mu(Z_batch)
+                counts = self._count(log_mu, L_batch)
                 E.append(tensor_to_numpy(counts))
                 pbar.update(1)
         
@@ -996,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):
@@ -1157,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 = (
@@ -1357,7 +1439,7 @@ def main():
     df = DensityFlow(
         input_size=input_size,
         cell_factor_size=cell_factor_size,
-        inverse_dispersion=args.inverse_dispersion,
+        dispersion=args.dispersion,
         z_dim=args.z_dim,
         hidden_layers=args.hidden_layers,
         hidden_layer_activation=args.hidden_layer_activation,

sure_tools-2.4.25/SURE/DensityFlow2.py → sure_tools-2.4.42/SURE/DensityFlowLinear.py

@@ -54,7 +54,7 @@ def set_random_seed(seed):
     # Set seed for Pyro
     pyro.set_rng_seed(seed)
 
-class DensityFlow2(nn.Module):
+class DensityFlowLinear(nn.Module):
     def __init__(self,
                  input_size: int,
                  codebook_size: int = 200,
@@ -63,8 +63,8 @@ class DensityFlow2(nn.Module):
                  supervised_mode: bool = False,
                  z_dim: int = 10,
                  z_dist: Literal['normal','studentt','laplacian','cauchy','gumbel'] = 'gumbel',
-                 loss_func: Literal['negbinomial','poisson','multinomial','bernoulli'] = 'multinomial',
-                 inverse_dispersion: float = 10.0,
+                 loss_func: Literal['negbinomial','poisson','multinomial','bernoulli'] = 'negbinomial',
+                 dispersion: float = 8.0,
                  use_zeroinflate: bool = False,
                  hidden_layers: list = [500],
                  hidden_layer_activation: Literal['relu','softplus','leakyrelu','linear'] = 'relu',
@@ -81,7 +81,7 @@ class DensityFlow2(nn.Module):
 
         self.input_size = input_size
         self.cell_factor_size = cell_factor_size
-        self.inverse_dispersion = inverse_dispersion
+        self.dispersion = dispersion
         self.latent_dim = z_dim
         self.hidden_layers = hidden_layers
         self.decoder_hidden_layers = hidden_layers[::-1]
@@ -107,10 +107,11 @@ class DensityFlow2(nn.Module):
         
         self.codebook_weights = None
 
+        self.seed = seed
         set_random_seed(seed)
         self.setup_networks()
         
-        print(f"🧬 DensityFlow2 Initialized:")
+        print(f"🧬 DensityFlowLinear Initialized:")
         print(f"   - Latent Dimension: {self.latent_dim}")
         print(f"   - Gene Dimension: {self.input_size}")
         print(f"   - Hidden Dimensions: {self.hidden_layers}")
@@ -208,17 +209,6 @@ class DensityFlow2(nn.Module):
             use_cuda=self.use_cuda,
         )
         
-        if self.loss_func == 'negbinomial':
-            self.encoder_inverse_dispersion = MLP(
-                [self.latent_dim] + hidden_sizes + [[self.input_size, self.input_size]],
-                activation=activate_fct,
-                output_activation=[Exp, Exp],
-                post_layer_fct=post_layer_fct,
-                post_act_fct=post_act_fct,
-                allow_broadcast=self.allow_broadcast,
-                use_cuda=self.use_cuda,
-            )
-        
         if self.cell_factor_size>0:
             self.cell_factor_effect = nn.ModuleList()
             for i in np.arange(self.cell_factor_size):
@@ -269,7 +259,7 @@ class DensityFlow2(nn.Module):
                             )
                         )
             
-        self.decoder_concentrate = MLP(
+        self.decoder_log_mu = MLP(
                     [self.latent_dim] + self.decoder_hidden_layers + [self.input_size],
                     activation=activate_fct,
                     output_activation=None,
@@ -352,15 +342,15 @@ class DensityFlow2(nn.Module):
         return xs
 
     def model1(self, xs):
-        pyro.module('DensityFlow2', self)
+        pyro.module('DensityFlowLinear', self)
 
         eps = torch.finfo(xs.dtype).eps
         batch_size = xs.size(0)
         self.options = dict(dtype=xs.dtype, device=xs.device)
         
         if self.loss_func=='negbinomial':
-            total_count = pyro.param("inverse_dispersion", self.inverse_dispersion *
-                                     xs.new_ones(self.input_size), constraint=constraints.positive)
+            dispersion = pyro.param("dispersion", self.dispersion *
+                                            xs.new_ones(self.input_size), constraint=constraints.positive)
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -394,28 +384,31 @@ class DensityFlow2(nn.Module):
                 zns = pyro.sample('zn', dist.Gumbel(zn_loc, zn_scale).to_event(1))
 
             zs = zns
-            concentrate = self.decoder_concentrate(zs)
+            log_mu = self.decoder_log_mu(zs)
             if self.loss_func in ['bernoulli']:
-                log_theta = concentrate
+                log_theta = log_mu
+            elif self.loss_func in ['negbinomial']:
+                mu = log_mu.exp()
             else:
-                rate = concentrate.exp()
+                rate = log_mu.exp()
                 theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
                 if self.loss_func == 'poisson':
                     rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
+                logits = (mu.log()-dispersion.log()).clamp(min=-15, max=15)
                 if self.use_zeroinflate:
-                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=total_count, probs=theta),gate_logits=gate_logits).to_event(1), obs=xs)
+                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=dispersion, 
+                                                                                         logits=logits),gate_logits=gate_logits).to_event(1), obs=xs)
                 else:
-                    pyro.sample('x', dist.NegativeBinomial(total_count=total_count, probs=theta).to_event(1), obs=xs)
+                    pyro.sample('x', dist.NegativeBinomial(total_count=dispersion, logits=logits).to_event(1), obs=xs)
             elif self.loss_func == 'poisson':
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
                 else:
                     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)
-                pyro.sample('x', dist.Multinomial(total_count=int(1e8), logits=concentrate), obs=xs)
+                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)
@@ -432,17 +425,15 @@ class DensityFlow2(nn.Module):
             ns = pyro.sample('n', dist.OneHotCategorical(logits=alpha))
 
     def model2(self, xs, us=None):
-        pyro.module('DensityFlow2', self)
+        pyro.module('DensityFlowLinear', self)
 
         eps = torch.finfo(xs.dtype).eps
         batch_size = xs.size(0)
         self.options = dict(dtype=xs.dtype, device=xs.device)
         
         if self.loss_func=='negbinomial':
-            #total_count = pyro.param("inverse_dispersion", self.inverse_dispersion *
-            #                         xs.new_ones(self.input_size), constraint=constraints.positive)
-            with pyro.plate("genes", self.input_size):
-                inverse_dispersion = pyro.sample("inverse_dispersion", dist.LogNormal(self.inverse_dispersion, 0.5).to_event(1))
+            dispersion = pyro.param("dispersion", self.dispersion *
+                                            xs.new_ones(self.input_size), constraint=constraints.positive)
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -482,28 +473,28 @@ class DensityFlow2(nn.Module):
                 zs = zns'''
 
             zs = zns
-            concentrate = self.decoder_concentrate(zs)
+            log_mu = self.decoder_log_mu(zs)
             for i in np.arange(self.cell_factor_size):
                 zus = self._cell_shift(zs, i, us[:,i].reshape(-1,1))
-                concentrate += self.decoder_concentrate(zus)
+                log_mu += self.decoder_log_mu(zus)
                 
             if self.loss_func in ['bernoulli']:
-                log_theta = concentrate
+                log_theta = log_mu
             elif self.loss_func in ['negbinomial']:
-                mu = concentrate.exp()
+                mu = log_mu.exp()
             else:
-                rate = concentrate.exp()
+                rate = log_mu.exp()
                 theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
                 if self.loss_func == 'poisson':
                     rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
-                logits = (mu.log()-inverse_dispersion.log()).clamp(min=-10, max=10)
+                logits = (mu.log()-dispersion.log()).clamp(min=-15, max=15)
                 if self.use_zeroinflate:
-                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=inverse_dispersion, 
+                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=dispersion, 
                                                                                          logits=logits),gate_logits=gate_logits).to_event(1), obs=xs)
                 else:
-                    pyro.sample('x', dist.NegativeBinomial(total_count=inverse_dispersion, 
+                    pyro.sample('x', dist.NegativeBinomial(total_count=dispersion, 
                                                            logits=logits).to_event(1), obs=xs)
             elif self.loss_func == 'poisson':
                 if self.use_zeroinflate:
@@ -511,8 +502,7 @@ class DensityFlow2(nn.Module):
                 else:
                     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)
-                pyro.sample('x', dist.Multinomial(total_count=int(1e8), logits=concentrate), obs=xs)
+                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)
@@ -528,21 +518,16 @@ class DensityFlow2(nn.Module):
             alpha = self.encoder_n(zns)
             ns = pyro.sample('n', dist.OneHotCategorical(logits=alpha))
             
-            if self.loss_func == 'negbinomial':
-                id_loc,id_scale = self.encoder_inverse_dispersion(zns)
-                with pyro.plate("genes", self.input_size):
-                    pyro.sample("inverse_dispersion", dist.LogNormal(id_loc, id_scale).to_event(1))
-
     def model3(self, xs, ys, embeds=None):
-        pyro.module('DensityFlow2', self)
+        pyro.module('DensityFlowLinear', self)
 
         eps = torch.finfo(xs.dtype).eps
         batch_size = xs.size(0)
         self.options = dict(dtype=xs.dtype, device=xs.device)
         
         if self.loss_func=='negbinomial':
-            total_count = pyro.param("inverse_dispersion", self.inverse_dispersion *
-                                     xs.new_ones(self.input_size), constraint=constraints.positive)
+            dispersion = pyro.param("dispersion", self.dispersion *
+                                            xs.new_ones(self.input_size), constraint=constraints.positive)
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -593,28 +578,31 @@ class DensityFlow2(nn.Module):
 
             zs = zns
 
-            concentrate = self.decoder_concentrate(zs)
+            log_mu = self.decoder_log_mu(zs)
             if self.loss_func in ['bernoulli']:
-                log_theta = concentrate
+                log_theta = log_mu
+            elif self.loss_func in ['negbinomial']:
+                mu = log_mu.exp()
             else:
-                rate = concentrate.exp()
+                rate = log_mu.exp()
                 theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
                 if self.loss_func == 'poisson':
                     rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
+                logits = (mu.log()-dispersion.log()).clamp(min=-15, max=15)
                 if self.use_zeroinflate:
-                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=total_count, probs=theta),gate_logits=gate_logits).to_event(1), obs=xs)
+                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=dispersion, 
+                                                                                         logits=logits),gate_logits=gate_logits).to_event(1), obs=xs)
                 else:
-                    pyro.sample('x', dist.NegativeBinomial(total_count=total_count, probs=theta).to_event(1), obs=xs)
+                    pyro.sample('x', dist.NegativeBinomial(total_count=dispersion, logits=logits).to_event(1), obs=xs)
             elif self.loss_func == 'poisson':
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
                 else:
                     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)
-                pyro.sample('x', dist.Multinomial(total_count=int(1e8), logits=concentrate), obs=xs)
+                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)
@@ -631,15 +619,15 @@ class DensityFlow2(nn.Module):
                 zns = embeds
 
     def model4(self, xs, us, ys, embeds=None):
-        pyro.module('DensityFlow2', self)
+        pyro.module('DensityFlowLinear', self)
 
         eps = torch.finfo(xs.dtype).eps
         batch_size = xs.size(0)
         self.options = dict(dtype=xs.dtype, device=xs.device)
         
         if self.loss_func=='negbinomial':
-            total_count = pyro.param("inverse_dispersion", self.inverse_dispersion *
-                                     xs.new_ones(self.input_size), constraint=constraints.positive)
+            dispersion = pyro.param("dispersion", self.dispersion *
+                                            xs.new_ones(self.input_size), constraint=constraints.positive)
             
         if self.use_zeroinflate:
             gate_logits = pyro.param("dropout_rate", xs.new_zeros(self.input_size))
@@ -695,32 +683,34 @@ class DensityFlow2(nn.Module):
                 zs = zns'''
 
             zs = zns
-            concentrate = self.decoder_concentrate(zs)
+            log_mu = self.decoder_log_mu(zs)
             for i in np.arange(self.cell_factor_size):
                 zus = self._cell_shift(zs, i, us[:,i].reshape(-1,1))
-                concentrate += self.decoder_concentrate(zus)
+                log_mu += self.decoder_log_mu(zus)
             
             if self.loss_func in ['bernoulli']:
-                log_theta = concentrate
+                log_theta = log_mu
+            elif self.loss_func in ['negbinomial']:
+                mu = log_mu.exp()
             else:
-                rate = concentrate.exp()
+                rate = log_mu.exp()
                 theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
                 if self.loss_func == 'poisson':
                     rate = theta * torch.sum(xs, dim=1, keepdim=True)
 
             if self.loss_func == 'negbinomial':
+                logits = (mu.log()-dispersion.log()).clamp(min=-15, max=15)
                 if self.use_zeroinflate:
-                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=total_count, probs=theta),gate_logits=gate_logits).to_event(1), obs=xs)
+                    pyro.sample('x', dist.ZeroInflatedDistribution(dist.NegativeBinomial(total_count=dispersion, logits=logits),gate_logits=gate_logits).to_event(1), obs=xs)
                 else:
-                    pyro.sample('x', dist.NegativeBinomial(total_count=total_count, probs=theta).to_event(1), obs=xs)
+                    pyro.sample('x', dist.NegativeBinomial(total_count=dispersion, logits=logits).to_event(1), obs=xs)
             elif self.loss_func == 'poisson':
                 if self.use_zeroinflate:
                     pyro.sample('x', dist.ZeroInflatedDistribution(dist.Poisson(rate=rate),gate_logits=gate_logits).to_event(1), obs=xs.round())
                 else:
                     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)
-                pyro.sample('x', dist.Multinomial(total_count=int(1e8), logits=concentrate), obs=xs)
+                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)
@@ -887,6 +877,8 @@ class DensityFlow2(nn.Module):
         
         # basal embedding
         zs = self.get_basal_embedding(xs)
+        log_mu = self.get_log_mu(zs)
+        
         for pert in perturbs_predict:
             pert_idx = int(np.where(perturbs_reference==pert)[0])
             us_i = us[:,pert_idx].reshape(-1,1)
@@ -898,10 +890,12 @@ class DensityFlow2(nn.Module):
             ps = np.ones_like(us_i)
             if np.sum(np.abs(ps-us_i))>=1:
                 dzs = self.get_cell_shift(zs, perturb_idx=pert_idx, perturb_us=ps)          
-                zs = zs + dzs0 + dzs
+                delta = dzs0 + dzs
             else:
-                zs = zs + dzs0
-        
+                delta = dzs0
+
+            log_mu = log_mu + self.get_log_mu(delta)
+
         if library_sizes is None:
             library_sizes = np.sum(xs, axis=1, keepdims=True)
         elif type(library_sizes) == list:
@@ -910,9 +904,9 @@ class DensityFlow2(nn.Module):
         elif len(library_sizes.shape)==1:
             library_sizes = library_sizes.reshape(-1,1)
             
-        counts = self.get_counts(zs, library_sizes=library_sizes)
+        counts = self.get_counts(log_mu, library_sizes=library_sizes)
         
-        return counts, zs
+        return counts, log_mu
     
     def _cell_shift(self, zs, perturb_idx, perturb):
         #zns,_ = self.encoder_zn(xs)    
@@ -956,48 +950,46 @@ class DensityFlow2(nn.Module):
         Z = np.concatenate(Z)
         return Z
     
-    def _get_theta(self, delta_zs):
-        return self.decoder_concentrate(delta_zs)
+    def _log_mu(self, zs):
+        return self.decoder_log_mu(zs)
     
-    def get_theta(self, 
-                             delta_zs, 
-                             batch_size: int = 1024):
+    def get_log_mu(self, zs, batch_size: int = 1024):
         """
         Return cells' changes in the feature space induced by specific perturbation of a factor
 
         """
-        delta_zs = convert_to_tensor(delta_zs, device=self.get_device())
-        dataset = CustomDataset(delta_zs)
+        zs = convert_to_tensor(zs, device=self.get_device())
+        dataset = CustomDataset(zs)
         dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
 
         R = []
         with tqdm(total=len(dataloader), desc='', unit='batch') as pbar:
-            for delta_Z_batch, _ in dataloader:
-                r = self._get_theta(delta_Z_batch)
+            for Z_batch, _ in dataloader:
+                r = self._log_mu(Z_batch)
                 R.append(tensor_to_numpy(r))
                 pbar.update(1)
 
         R = np.concatenate(R)
         return R
     
-    def _count(self, concentrate, library_size=None):
+    def _count(self, log_mu, library_size=None):
         if self.loss_func == 'bernoulli':
             #counts = self.sigmoid(concentrate)
-            counts = dist.Bernoulli(logits=concentrate).to_event(1).mean
+            counts = dist.Bernoulli(logits=log_mu).to_event(1).mean
         elif self.loss_func == 'multinomial':
-            theta = dist.Multinomial(total_count=int(1e8), logits=concentrate).mean
+            theta = dist.Multinomial(total_count=int(1e8), logits=log_mu).mean
             counts = theta * library_size
         else:
-            rate = concentrate.exp()
+            rate = log_mu.exp()
             theta = dist.DirichletMultinomial(total_count=1, concentration=rate).mean
             counts = theta * library_size
         return counts
     
-    def get_counts(self, concentrate, 
+    def get_counts(self, log_mu, 
                    library_sizes,
                    batch_size: int = 1024):
 
-        concentrate = convert_to_tensor(concentrate, device=self.get_device())
+        log_mu = convert_to_tensor(log_mu, device=self.get_device())
         
         if type(library_sizes) == list:
             library_sizes = np.array(library_sizes).reshape(-1,1)
@@ -1005,13 +997,13 @@ class DensityFlow2(nn.Module):
             library_sizes = library_sizes.reshape(-1,1)
         ls = convert_to_tensor(library_sizes, device=self.get_device())
             
-        dataset = CustomDataset2(concentrate,ls)
+        dataset = CustomDataset2(log_mu,ls)
         dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
         
         E = []
         with tqdm(total=len(dataloader), desc='', unit='batch') as pbar:
-            for C_batch, L_batch, _ in dataloader:
-                counts = self._count(C_batch, L_batch)
+            for Mu_batch, L_batch, _ in dataloader:
+                counts = self._count(Mu_batch, L_batch)
                 E.append(tensor_to_numpy(counts))
                 pbar.update(1)
         
@@ -1045,7 +1037,7 @@ class DensityFlow2(nn.Module):
             threshold: int = 0,
             use_jax: bool = True):
         """
-        Train the DensityFlow2 model.
+        Train the DensityFlowLinear model.
 
         Parameters
         ----------
@@ -1071,7 +1063,7 @@ class DensityFlow2(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 DensityFlow2 in the shell command.
+            the Python script or Jupyter notebook. It is OK if it is used when runing DensityFlowLinear in the shell command.
         """
         xs = self.preprocess(xs, threshold=threshold)
         xs = convert_to_tensor(xs, dtype=self.dtype, device=self.get_device())
@@ -1189,12 +1181,12 @@ class DensityFlow2(nn.Module):
 
         
 EXAMPLE_RUN = (
-    "example run: DensityFlow2 --help"
+    "example run: DensityFlowLinear --help"
 )
 
 def parse_args():
     parser = argparse.ArgumentParser(
-        description="DensityFlow2\n{}".format(EXAMPLE_RUN))
+        description="DensityFlowLinear\n{}".format(EXAMPLE_RUN))
 
     parser.add_argument(
         "--cuda", action="store_true", help="use GPU(s) to speed up training"
@@ -1381,10 +1373,10 @@ def main():
     cell_factor_size = 0 if us is None else us.shape[1]
 
     ###########################################
-    df = DensityFlow2(
+    df = DensityFlowLinear(
         input_size=input_size,
         cell_factor_size=cell_factor_size,
-        inverse_dispersion=args.inverse_dispersion,
+        dispersion=args.dispersion,
         z_dim=args.z_dim,
         hidden_layers=args.hidden_layers,
         hidden_layer_activation=args.hidden_layer_activation,
@@ -1412,9 +1404,9 @@ def main():
 
     if args.save_model is not None:
         if args.save_model.endswith('gz'):
-            DensityFlow2.save_model(df, args.save_model, compression=True)
+            DensityFlowLinear.save_model(df, args.save_model, compression=True)
         else:
-            DensityFlow2.save_model(df, args.save_model)
+            DensityFlowLinear.save_model(df, args.save_model)
     
 
 

{sure_tools-2.4.25 → sure_tools-2.4.42}/SURE/__init__.py

@@ -1,6 +1,6 @@
 from .SURE import SURE
 from .DensityFlow import DensityFlow
-from .DensityFlow2 import DensityFlow2
+from .DensityFlowLinear import DensityFlowLinear
 from .PerturbE import PerturbE
 from .TranscriptomeDecoder import TranscriptomeDecoder
 from .SimpleTranscriptomeDecoder import SimpleTranscriptomeDecoder
@@ -12,7 +12,7 @@ from . import utils
 from . import codebook
 from . import SURE
 from . import DensityFlow
-from . import DensityFlow2
+from . import DensityFlowLinear
 from . import atac
 from . import flow 
 from . import perturb
@@ -23,6 +23,6 @@ from . import EfficientTranscriptomeDecoder
 from . import VirtualCellDecoder
 from . import PerturbationAwareDecoder
 
-__all__ = ['SURE', 'DensityFlow', 'DensityFlow2', 'PerturbE', 'TranscriptomeDecoder', 'SimpleTranscriptomeDecoder',
+__all__ = ['SURE', 'DensityFlow', 'DensityFlowLinear', 'PerturbE', 'TranscriptomeDecoder', 'SimpleTranscriptomeDecoder',
            'EfficientTranscriptomeDecoder', 'VirtualCellDecoder', 'PerturbationAwareDecoder', 
            'flow', 'perturb', 'atac', 'utils', 'codebook']

{sure_tools-2.4.25 → sure_tools-2.4.42}/SURE_tools.egg-info/PKG-INFO

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

{sure_tools-2.4.25 → sure_tools-2.4.42}/SURE_tools.egg-info/SOURCES.txt

@@ -2,7 +2,7 @@ LICENSE
 README.md
 setup.py
 SURE/DensityFlow.py
-SURE/DensityFlow2.py
+SURE/DensityFlowLinear.py
 SURE/EfficientTranscriptomeDecoder.py
 SURE/PerturbE.py
 SURE/PerturbationAwareDecoder.py

{sure_tools-2.4.25 → sure_tools-2.4.42}/setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
 
 setup(
     name='SURE-tools',
-    version='2.4.25',
+    version='2.4.42',
     description='Succinct Representation of Single Cells',
     long_description=long_description,
     long_description_content_type="text/markdown",