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

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)
         
@@ -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,