SURE-tools 2.4.2__py3-none-any.whl → 2.4.5__py3-none-any.whl

SURE/TranscriptomeDecoder.py

@@ -86,14 +86,16 @@ class TranscriptomeDecoder:
             output = torch.matmul(latent_proj, gene_embeds)  # [batch, gene_dim]
             
             return self.activation(output)
-    
-    class ChunkedTransformer(nn.Module):
-        """Process genes in chunks to reduce memory usage"""
-        def __init__(self, gene_dim, hidden_dim, chunk_size=1000, num_layers=4):
+        
+    class ChunkedTransformer(nn.Module):    
+        def __init__(self, gene_dim, hidden_dim=512, chunk_size=2000, num_layers=3):
             super().__init__()
             self.chunk_size = chunk_size
+            self.hidden_dim = hidden_dim
             self.num_chunks = (gene_dim + chunk_size - 1) // chunk_size
-            self.layers = nn.ModuleList([
+        
+            # 共享的Transformer层
+            self.transformer_layers = nn.ModuleList([
                 nn.Sequential(
                     nn.Linear(hidden_dim, hidden_dim),
                     nn.GELU(),
@@ -101,22 +103,40 @@ class TranscriptomeDecoder:
                     nn.Linear(hidden_dim, hidden_dim),
                 ) for _ in range(num_layers)
             ])
-            
+        
+            # 每个chunk独立的投影层
+            self.input_projections = nn.ModuleList([
+                nn.Linear(min(chunk_size, gene_dim - i * chunk_size), hidden_dim)
+                for i in range(self.num_chunks)
+            ])
+            self.output_projections = nn.ModuleList([
+                nn.Linear(hidden_dim, min(chunk_size, gene_dim - i * chunk_size))
+                for i in range(self.num_chunks)
+            ])
+        
         def forward(self, x):
-            # Process in chunks to save memory
-            batch_size = x.shape[0]
+            batch_size, gene_dim = x.shape
             output = torch.zeros_like(x)
-            
+        
             for i in range(self.num_chunks):
                 start_idx = i * self.chunk_size
-                end_idx = min((i + 1) * self.chunk_size, x.shape[1])
-                
-                chunk = x[:, start_idx:end_idx]
-                for layer in self.layers:
-                    chunk = layer(chunk) + chunk  # Residual connection
-                
-                output[:, start_idx:end_idx] = chunk
+                end_idx = min((i + 1) * self.chunk_size, gene_dim)
+                current_chunk_size = end_idx - start_idx
+            
+                chunk = x[:, start_idx:end_idx]  # [batch_size, current_chunk_size]
+            
+                # 投影到hidden_dim
+                chunk_proj = self.input_projections[i](chunk)  # [batch_size, hidden_dim]
             
+                # Transformer处理
+                for layer in self.transformer_layers:
+                    chunk_proj = layer(chunk_proj) + chunk_proj
+            
+                # 投影回原始维度
+                chunk_out = self.output_projections[i](chunk_proj)  # [batch_size, current_chunk_size]
+            
+                output[:, start_idx:end_idx] = chunk_out
+        
             return output
     
     class Decoder(nn.Module):
@@ -166,6 +186,8 @@ class TranscriptomeDecoder:
             self.output_scale = nn.Parameter(torch.ones(1))
             self.output_bias = nn.Parameter(torch.zeros(1))
             
+            self.latent_to_gene = nn.Linear(hidden_dim, gene_dim)
+            
             self._init_weights()
         
         def _init_weights(self):
@@ -185,8 +207,8 @@ class TranscriptomeDecoder:
             gene_features = self.gene_projection(latent)
             
             # 3. Add latent information
-            print(f'{gene_features.shape}; {latent_expanded.shape}')
-            gene_features = gene_features + latent_expanded.unsqueeze(1)
+            latent_gene_injection = self.latent_to_gene(latent_expanded)
+            gene_features = gene_features + latent_gene_injection
             
             # 4. Chunked processing (memory efficient)
             gene_features = self.chunked_processor(gene_features)

{sure_tools-2.4.2.dist-info → sure_tools-2.4.5.dist-info}/METADATA

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

{sure_tools-2.4.2.dist-info → sure_tools-2.4.5.dist-info}/RECORD

@@ -1,7 +1,7 @@
 SURE/DensityFlow.py,sha256=YvaE9aPbAC2U7WhTye5i2AMtcw0BI_qS3gv9SP4aE0k,56676
 SURE/PerturbE.py,sha256=DxEp-qef--x8-GMZdPfBf8ts8UDDc34h2P5AnpqZ-YM,52265
 SURE/SURE.py,sha256=MXs7iuvcj-lU4dJ_MwKegpL2Rqk2HB4eFfAgHRA3RtA,47744
-SURE/TranscriptomeDecoder.py,sha256=z0NtjvpT1nN4H9WC97VJPWuTvqaZYn36MbpaIWjXgGU,20324
+SURE/TranscriptomeDecoder.py,sha256=fjTl2wC-nGTdbQGgFDbTmWYI8RoEg6J4cHPmoUoJJfI,21286
 SURE/__init__.py,sha256=pNSGQ4BMqMXBAPHpFOYNB8_0vFW-RqPy3rr5fvdEEyU,473
 SURE/assembly/__init__.py,sha256=jxZLURXKPzXe21LhrZ09LgZr33iqdjlQy4oSEj5gR2Q,172
 SURE/assembly/assembly.py,sha256=6IMdelPOiRO4mUb4dC7gVCoF1Uvfw86-Map8P_jnUag,21477
@@ -19,9 +19,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.2.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
-sure_tools-2.4.2.dist-info/METADATA,sha256=gorPWZ40-GBJ3Rz6CH-eWZFz0ZKW8LacnDEJbg1cik4,2677
-sure_tools-2.4.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sure_tools-2.4.2.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
-sure_tools-2.4.2.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
-sure_tools-2.4.2.dist-info/RECORD,,
+sure_tools-2.4.5.dist-info/licenses/LICENSE,sha256=TFHKwmrAViXQbSX5W-NDItkWFjm45HWOeUniDrqmnu0,1065
+sure_tools-2.4.5.dist-info/METADATA,sha256=2GjCK_HUQ_Vs6b8AT2PIelOadhiVeOakI8B_OqbRyi0,2677
+sure_tools-2.4.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sure_tools-2.4.5.dist-info/entry_points.txt,sha256=-nJI8rVe_qqrR0HmfAODzj-JNfEqCcSsyVh6okSqyHk,83
+sure_tools-2.4.5.dist-info/top_level.txt,sha256=BtFTebdiJeqra4r6mm-uEtwVRFLZ_IjYsQ7OnalrOvY,5
+sure_tools-2.4.5.dist-info/RECORD,,