OSCR

Single-cell-level perturbation-induced and condition-related signal estimation with batch effect removal using NDreamer.

Code ↔ Paper

24 matches between paragraphs of the paper and lines of its authors' code, computed by the harvester (lexical-v1). Click a colored paragraph or line to see its counterpart.

The 24 matches · 5 of them tie a paragraph to a whole file, not to given lines: weak matches, whose lines are not tinted
  1. [1] § Results › Evaluation metrics ↔ R_plot_codes/experimental_benchmark_plot_codes/ablation.R, the whole file · a weak match · score 0.85 · ASW_label, bLISI, mixing metrics, ASW_batch, cLISI, kBET
  2. [2] § Results › Evaluation metrics ↔ R_plot_codes/experimental_benchmark_plot_codes/plot_experimental.R, lines 13–86 · score 0.81 · ASW_label, bLISI, ASW_batch, cLISI, kBET, biological conservation
  3. [3] § Methods › Estimation of the effect modifier embedding › Normalized mutual information loss ↔ ndreamer/DL_loss_func.py, lines 386–461 · score 0.75 · hot encoded, joint probability, assignment vector, matrix, mutual, loss
  4. [4] § Results › Evaluation metrics ↔ R_plot_codes/case_control_code/radar.R, lines 14–122 · score 0.75 · EM bio conservation, denoised batch mixing, EM batch, signal preservation, Radar, distortion
  5. [5] § Results › Evaluation metrics › Simulation dataset ↔ R_plot_codes/experimental_benchmark_plot_codes/overall_plot_experimental.R, lines 37–92 · score 0.74 · scCAPE, scGen, CINEMA OT, biological conservation, Mixscape, NDreamer
  6. [6] § Results › Exploration of batch effects, individual treatment effect, and sensitivity test › Ablation study ↔ ndreamer/model.py, lines 221–306 · score 0.72 · VQ VAE, local neighborhood loss, triplet loss, independent loss, model, NDreamer
  7. [7] § Results › NDreamer demonstrates robustly improved performance across experimental perturbation datasets ↔ R_plot_codes/experimental_benchmark_plot_codes/overall_plot_experimental.R, lines 37–92 · score 0.67 · scCAPE, scGen, CINEMA OT, Mixscape, NDreamer, benchmarking
  8. [8] § Results › NDreamer demonstrates robustly improved performance across experimental perturbation datasets ↔ R_plot_codes/experimental_benchmark_plot_codes/ablation.R, the whole file · a weak match · score 0.66 · bLISI, ASW_batch, mixing metrics, kBET, benchmarking, NDreamer
  9. [9] § Results › NDreamer uncovers differential gene expression patterns in a large Alzheimer's disease cohort ↔ R_plot_codes/AD_ADNC_code/AD_neurons_GO.R, lines 1–10 · score 0.65 · PCDH9 AS2, PPP1R9A AS1, GO, genes, NDreamer
  10. [10] § Methods › Overview of the NDreamer deep learning model ↔ tutorial_experimental_PBMC.ipynb, lines 25–180 · score 0.63 · local neighborhood loss, Gaussian kernel, triplet loss function, NDreamer, modifier space, reconstruct
  11. [11] § Results › Evaluation metrics ↔ R_plot_codes/case_control_code/radar_mouse.R, lines 14–120 · score 0.63 · denoised batch mixing, EM batch, signal preservation, Radar, distortion, mouse
  12. [12] § Results › Exploration of batch effects, individual treatment effect, and sensitivity test › Exploration of batch effects ↔ ndreamer/pipeline.py, lines 292–362 · score 0.63 · CD8T, local neighborhood loss, independent loss, UMAP, NDreamer, PBMC
  13. [13] § Results › NDreamer demonstrates robustly improved performance across experimental perturbation datasets ↔ ndreamer/metrics.py, lines 164–278 · score 0.61 · bLISI, asw_batch, kBET, silhouette, metrics, neighbor
  14. [14] § Methods › Overview of the NDreamer deep learning model ↔ ndreamer/model.py, lines 221–306 · score 0.61 · local neighborhood loss, dependent loss, triplet loss, modifier space, VAE, reconstruct
  15. [15] § Methods › Estimation of the effect modifier embedding › Local neighborhood loss ↔ ndreamer/DL_loss_func.py, lines 464–551 · score 0.59 · Gaussian kernel, local neighborhood loss, nearest neighbors, distance, cluster, space
  16. [16] § Methods › Estimation of the effect modifier embedding › Normalized mutual information loss ↔ ndreamer/DL_loss_func.py, lines 386–461 · score 0.59 · mutual information loss, encourage independence, entropy, assignments, probability, batch
  17. [17] § Results › Evaluation metrics ↔ ndreamer/metrics.py, lines 164–278 · score 0.57 · bLISI, asw_batch, kBET, metrics, cell
  18. [18] § Results › NDreamer uncovers differential gene expression patterns in a large Alzheimer's disease cohort ↔ R_plot_codes/case_control_code/simple_GO.R, the whole file · a weak match · score 0.57 · enriched GO terms, GO enrichment, KIT, TRPC3, genes
  19. [19] § Results › NDreamer demonstrates robustly improved performance across experimental perturbation datasets ↔ R_plot_codes/experimental_benchmark_plot_codes/plot_experimental_have_batch_ndreamer.R, lines 13–96 · score 0.55 · cLISI, conservation metrics, ARI, NMI, ASW, perturbation
  20. [20] § Results › NDreamer demonstrates robustly improved performance across experimental perturbation datasets ↔ R_plot_codes/experimental_benchmark_plot_codes/plot_experimental_have_batch.R, lines 13–96 · score 0.54 · cLISI, conservation metrics, ARI, NMI, ASW, perturbation
  21. [21] § Methods › Estimation of the effect modifier embedding › Encoder with neural discrete representation learning and variational inference ↔ ndreamer/model.py, lines 188–219 · score 0.53 · KL divergence, divergence loss, VAE
  22. [22] § Results › NDreamer uncovers differential gene expression patterns in a large Alzheimer's disease cohort ↔ R_plot_codes/AD_ADNC_code/AD_respiration_GO.R, the whole file · a weak match · score 0.53 · enriched GO terms, GO enrichment, PLCG2, genes
  23. [23] § Results › NDreamer uncovers differential gene expression patterns in a large Alzheimer's disease cohort ↔ R_plot_codes/AD_ADNC_code/AD_respiration_GO.R, the whole file · a weak match · score 0.52 · ATP6, COX2, CYTB, ND3, ND4, respiration
  24. [24] § Methods › Estimation of the effect modifier embedding › Encoder with neural discrete representation learning and variational inference ↔ ndreamer/model.py, lines 123–186 · score 0.51 · VQ VAE, commitment loss, codebook, Encoder, embeddings

Paper

Loaded from Europe PMC by your browser, not stored by OSCR: doi.org · Europe PMC

The paper is loaded when this pane is shown.

The authors' code

Python · 482 lines · 26 KB · GPL-3.0 · 4 matches

  1. import os
  2. import numpy as np
  3. import pandas as pd
  4. import torch
  5. import torch.nn as nn
  6. import torch.nn.functional as F
  7. from torch.cuda import device
  9. from ndreamer.model_DL import NDreamer_generator,Discriminator,set_seed
  10. from ndreamer.DL_loss_func import CrossEntropy, create_triplets_within_groups, IndependenceLoss, kl_divergence_loss, \
  11. IndependenceLoss_label, OrthogonalityRegularization, EntropyPenalty, create_triplets_within_groups_logits, \
  12. IndependenceLoss_between_matrix, compute_mmd, compute_local_neighborhood_loss, create_triplets_within_groups_tensor
  13. from ndreamer.data_preprocess import process_adata,generate_balanced_dataloader,generate_adata_to_dataloader
  14. from ndreamer.DL_loss_func import reconstruction_error
  15. from ndreamer.statistics import *
  17. from sklearn.decomposition import PCA
  18. # Dynamic import of tqdm based on the environment
  19. import sys
  20. if 'ipykernel' in sys.modules:
  21. from tqdm.notebook import tqdm
  22. else:
  23. from tqdm import tqdm
  25. from ndreamer.plot import plot_distribution
  27. class NDreamer_DL(nn.Module):
  28. def __init__(self, adata, condition_key, contorl_name, num_hvg, require_batch=False, batch_key=None,
  29. resolution_low=0.5, resolution_high=7, cluster_method="Leiden", embedding_dim=512, codebooks=None,
  30. codebook_dim=8, encoder_hidden=None, decoder_hidden=None, z_dim=256,
  31. cos_loss_scaler=20, random_seed=123, batch_size=2048, epoches=10, lr=1e-3, triplet_margin=5,
  32. independent_loss_scaler=1000, save_pth="./model/", developer_test_mode=False,
  33. library_size_normalize_adata=False, save_preprocessed_adata_path="./model/preprocessed.h5ad",
  34. KL_scaler=5e-3, reconstruct_scaler=1, triplet_scaler=5, num_triplets_per_label=15,
  35. tau=0.01, commitment_loss_scaler=1, cluster_correlation_scaler=50,
  36. try_identify_perturb_escaped_cell=False, reset_threshold=1 / 1024, reset_interval=30,
  37. try_identify_cb_specific_subtypes=False, local_neighborhood_loss_scaler=1,
  38. local_neighbor_sigma=1, n_neighbors=20, local_neighbor_across_cluster_scaler=20,
  39. have_negative_data=False):
  40. super(NDreamer_DL, self).__init__()
  41. input_dim=num_hvg
  43. if codebooks is None:
  44. codebooks = [1024 for i in range(32)]
  45. if encoder_hidden is None:
  46. encoder_hidden = [2048, 1024]
  47. if decoder_hidden is None:
  48. decoder_hidden = [512, 1024]
  50. self.developer_test_mode=developer_test_mode
  51. self.require_batch=require_batch
  52. self.batch_key = batch_key
  53. self.condition_key = condition_key
  54. self.triplet_margin=triplet_margin
  55. if not os.path.exists(save_pth):
  56. os.mkdir(save_pth)
  57. self.save_pth=save_pth
  58. self.random_seed = random_seed
  59. self.batch_size = batch_size
  60. self.epoches = epoches
  61. self.lr = lr
  62. self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
  63. print("Using device:",self.device)
  64. self.kl_scaler = KL_scaler
  65. self.reconstruct_scaler = reconstruct_scaler
  66. self.triplet_scaler = triplet_scaler
  67. self.num_triplets_per_label=num_triplets_per_label
  68. self.independent_loss_scaler = independent_loss_scaler
  69. self.codebooks=codebooks
  70. self.local_neighborhood_loss_scaler=local_neighborhood_loss_scaler
  71. print(local_neighborhood_loss_scaler)
  72. self.local_neighbor_sigma=local_neighbor_sigma
  73. self.try_identify_cb_specific_subtypes=try_identify_cb_specific_subtypes
  74. self.try_identify_perturb_escaped_cell=try_identify_perturb_escaped_cell
  75. self.n_neighbors=n_neighbors
  76. self.local_neighbor_across_cluster_scaler=local_neighbor_across_cluster_scaler
  78. #data preprocessing
  79. if not developer_test_mode:
  80. print("Start data preprocessing")
  81. adata, batch_dict, condition_dict = process_adata(adata=adata, condition_key=condition_key,
  82. input_dim=input_dim, control_name=contorl_name,
  83. require_batch=require_batch, batch_key=batch_key,
  84. resolution_low=resolution_low,
  85. resolution_high=resolution_high,
  86. cluster_method=cluster_method,
  87. library_size_normalize_adata=library_size_normalize_adata)
  90. self.adata = adata
  91. self.batch_dict = batch_dict
  92. self.condition_dict = condition_dict
  93. print("Data preprocessing done")
  94. if save_preprocessed_adata_path is not None:
  95. adata.write(save_preprocessed_adata_path)
  96. else:
  97. self.adata=adata
  99. print("Remaining number of cells:",self.adata.shape[0])
  101. calcluated_epoches=15*self.adata.shape[0]//(self.batch_size*len(np.unique(self.adata.obs["group"])))+1
  102. if calcluated_epoches>self.epoches:
  103. if len(np.unique(self.adata.obs["group"]))<=10:
  104. calcluated_epoches = calcluated_epoches + reset_interval - calcluated_epoches % reset_interval - 2
  105. self.epoches=calcluated_epoches
  106. print("Too few epoches (steps, if rigorously speaking). Changing epoch to", self.epoches, "to adjust for number of cells")
  108. num_batches = []#np.unique(adata.obs["batch"]).shape[0]
  109. if not self.developer_test_mode and self.require_batch:
  110. if isinstance(batch_key, str):
  111. num_batches = [max(self.batch_dict[batch_key].values()) + 1]
  112. elif isinstance(batch_key, list):
  113. num_batches=[]
  114. for batch_keyi in batch_key:
  115. num_batches.append(max(self.batch_dict[batch_keyi].values()) + 1)
  116. else:
  117. num_batches=[]
  118. self.num_batches=num_batches
  119. num_treatments = np.unique(adata.obs["condition"]).shape[0]
  120. if not developer_test_mode:
  121. num_treatments=max(num_treatments,max(self.condition_dict.values())+1)
  123. # define the model
  124. self.VQ_VAE = NDreamer_generator(input_dim=input_dim, num_treatments=num_treatments,z_dim=z_dim,
  125. num_batches=num_batches,embedding_dim=embedding_dim,tau=tau,
  126. codebooks=codebooks,codebook_dim=codebook_dim,encoder_hidden=encoder_hidden,
  127. decoder_hidden=decoder_hidden,commitment_loss_scaler=commitment_loss_scaler,
  128. reset_threshold=reset_threshold,reset_interval=reset_interval,
  129. try_identify_cb_specific_subtypes=try_identify_cb_specific_subtypes,
  130. try_identify_perturb_escaped_cell=try_identify_perturb_escaped_cell,
  131. have_negative_data=have_negative_data)
  133. self.require_batch = require_batch
  134. print("Require batch:",self.require_batch)
  136. self.independence_loss=IndependenceLoss(scaler=independent_loss_scaler)
  137. self.independence_loss_label=IndependenceLoss_label(scaler=cluster_correlation_scaler/num_treatments)
  138. self.independence_loss_between_codebook=IndependenceLoss_between_matrix(scaler=independent_loss_scaler/100)
  139. self.cross_entropy = CrossEntropy()
  140. self.cos_loss_scaler = cos_loss_scaler
  141. #self.MINE=MINE(x_dim=input_dim,z_dim=z_dim,hidden_dim=1024)
  143. # init the model
  144. self.VQ_VAE.to(self.device)
  145. self.independence_loss.to(self.device)
  146. self.independence_loss_label.to(self.device)
  147. #self.cross_entropy.to(self.device)
  148. #self.orthogonality_regularization.to(self.device)
  149. #self.entropy_penalty.to(self.device)
  150. self.independence_loss_between_codebook.to(self.device)
  151. #self.MINE.to(self.device)
  153. self.logits=None
  154. self.df_latent=None
  156. def train_model(self):
  157. optimizer_G = torch.optim.AdamW(self.VQ_VAE.parameters(), lr=self.lr)
  159. progress_bar = tqdm(total=self.epoches, desc="Overall Progress", leave=True, miniters=1, mininterval=0)
  160. for epoch in range(self.epoches):
  161. set_seed(self.random_seed+epoch)
  162. data_loader = generate_balanced_dataloader(self.adata, batch_size=self.batch_size)
  163. self.VQ_VAE.train()
  164. all_losses = 0
  165. T_loss = 0
  166. V_loss = 0
  167. I_loss = 0
  168. K_loss = 0
  169. N_loss = 0
  170. Commitment_loss=0
  171. Dependent_loss=0
  172. for i, (exp, condition, batch, labels_low, labels_high) in enumerate(data_loader):
  173. #print(exp.shape, condition.shape, batch.shape, labels_low.shape, labels_high.shape)
  174. # convert to cuda
  175. exp = exp.to(self.device)
  176. condition = condition.to(self.device)
  177. batch = batch.to(self.device)
  178. labels_low = labels_low.to(self.device)
  179. labels_high = labels_high.to(self.device)
  180. # run VQ-VAE
  181. z, variance, reconstructed, logits, commitment_loss,escape_judger_choice, cb_specifc_embedding=self.VQ_VAE(exp=exp, treatment=condition, batch=batch)
  182. commitment_loss=commitment_loss/len(self.codebooks)
  184. if self.VQ_VAE.encoder.codebooks[0].just_reset_codebook:
  185. print("Finish resetting codebook embeddings, current step (epoch):",epoch)
  186. continue
  188. # calculate the reconstruction loss
  189. reconst_loss_mse = reconstruction_error(exp, reconstructed)*self.reconstruct_scaler
  190. reconst_loss_cos = (1 - torch.sum(F.normalize(reconstructed, p=2) * F.normalize(exp, p=2), 1)).mean()
  191. reconst_loss_cos = self.cos_loss_scaler * reconst_loss_cos
  192. reconst_loss = reconst_loss_mse + reconst_loss_cos
  193. reconst_loss = torch.clamp(reconst_loss, max=1e5)
  195. KL_loss=kl_divergence_loss(z_mean=z,z_var=variance)*self.kl_scaler
  196. KL_loss = torch.clamp(KL_loss, max=1e5)
  198. # independence loss for VAE
  199. independent_loss=0
  200. independent_loss=independent_loss+self.independence_loss(condition, logits)
  201. for i in range(len(self.num_batches)):
  202. independent_loss = independent_loss + self.independence_loss(batch[:,i], logits)
  203. independent_loss = torch.clamp(independent_loss, max=1e5)
  205. # convert matrix batch to unique int number vector
  206. base = batch.max() + 1 # Base larger than the largest number in the matrix
  207. powers = base ** torch.arange(batch.size(1), device=exp.device).flip(0).float() # Positional encoding powers
  208. batch_vector = torch.matmul(batch.float(), powers).long() # Encoded as unique integers
  210. dependent_loss = 0
  211. dependent_loss=dependent_loss-self.independence_loss_label(assignments=labels_high,
  212. probabilities=logits,
  213. condition=condition,
  214. batch=batch_vector)
  216. dependent_loss=dependent_loss-self.independence_loss_label(assignments=labels_low,
  217. probabilities=logits,
  218. condition=condition,
  219. batch=batch_vector)
  221. # triplet loss
  222. triplet_loss = create_triplets_within_groups_tensor(z, labels_low, labels_high, batch_vector,
  223. condition, margin=self.triplet_margin,
  224. num_triplets=self.num_triplets_per_label)
  225. triplet_loss=triplet_loss*self.triplet_scaler
  227. neighbor_loss = compute_local_neighborhood_loss(latent=z, inputs=exp,
  228. condition=condition, batch=batch_vector,
  229. sigma=self.local_neighbor_sigma,
  230. n_neighbors=self.n_neighbors,
  231. label_low=labels_low,
  232. label_high=labels_high,
  233. local_neighbor_across_cluster_scaler=self.local_neighbor_across_cluster_scaler
  234. ) * self.local_neighborhood_loss_scaler
  236. all_loss = reconst_loss + independent_loss +KL_loss + triplet_loss + neighbor_loss + commitment_loss + dependent_loss# + embedding_loss +entropy_loss
  238. all_loss = torch.clamp(all_loss, max=1e5)
  240. all_loss.backward()
  241. optimizer_G.step()
  242. optimizer_G.zero_grad()
  243. all_losses += all_loss.item()/len(data_loader)
  244. N_loss += neighbor_loss.item()/len(data_loader)
  245. T_loss += triplet_loss.item()/len(data_loader)
  246. V_loss += reconst_loss.item()/len(data_loader)
  247. I_loss +=independent_loss.item()/len(data_loader)
  248. K_loss +=KL_loss.item()/len(data_loader)
  249. Dependent_loss+=dependent_loss.item()/len(data_loader)
  250. #Entropy_loss+=entropy_loss.item()
  251. #Embedding_loss += 0#embedding_loss.item()
  252. Commitment_loss+=commitment_loss.item()/len(data_loader)
  254. print(f"Epoch: {epoch + 1}/{self.epoches} | "
  255. f"All Loss: {all_losses:.4f} | "
  256. f"Neighborhood Loss: {N_loss:.4f} | "
  257. f"Triplet Loss: {T_loss:.4f} | "
  258. f"Reconstruction Loss: {V_loss:.4f} | "
  259. f"Independent Loss: {I_loss:.4f} | "
  260. f"KL Loss: {K_loss:.4f} | "
  261. f"Commitment Loss: {Commitment_loss:.4f} | "
  262. f"Dependent Loss: {Dependent_loss:.4f}")
  263. progress_bar.update(1) # Increment the progress bar by one for each batch processed
  264. progress_bar.set_postfix(epoch=f"{epoch + 1}/{self.epoches}", all_loss=all_losses, neigh_loss=N_loss, triplet_loss=T_loss, reconst_loss=V_loss, independent_loss=I_loss, KL_loss=K_loss, Commitment_loss=Commitment_loss, Dependent_loss=Dependent_loss)
  266. progress_bar.close()
  267. torch.save(self.state_dict(), os.path.join(self.save_pth, 'ndreamer' + '.pth'))
  269. def get_modifier_space(self,dim=50, save_path_adata=None, save_path_latent_df=None,
  270. print_codebook_statisitics=False, kBET_test=False, save_others=False):
  271. self.VQ_VAE.eval()
  272. data_loader = generate_adata_to_dataloader(self.adata)
  274. device = 'cuda' if torch.cuda.is_available() else 'cpu'
  275. all_z = []
  276. all_indices = []
  277. logits=[]
  278. cb_specific_subtypes=[]
  279. escape_probs=[]
  280. with torch.no_grad():
  281. for i, (x, indices, condition, batch) in enumerate(data_loader):
  282. x = x.to(device)
  283. condition = condition.to(device)
  284. batch = batch.to(device)
  286. z, variance, reconstructed, logit, commitment_loss, escape_judger_choice, cb_specifc_embedding = self.VQ_VAE(
  287. exp=x, treatment=condition, batch=batch)
  288. all_z.append(z.cpu().detach())
  289. all_indices.extend(indices.tolist())
  291. if self.try_identify_perturb_escaped_cell:
  292. escape_probs.append(escape_judger_choice.detach().cpu())
  294. if self.try_identify_cb_specific_subtypes:
  295. cb_specific_subtypes.append(cb_specifc_embedding.detach().cpu())
  297. if len(logits) == 0:
  298. logits = [i.detach().cpu() for i in logit]
  299. else:
  300. logits = [torch.concatenate([logits[i], logit[i].detach().cpu()], dim=0) for i in
  301. range(len(logits))]
  303. all_z_combined = torch.cat(all_z, dim=0)
  304. all_indices_tensor = torch.tensor(all_indices)
  305. all_z_reordered = all_z_combined[all_indices_tensor.argsort()]
  306. all_z_np = all_z_reordered.numpy()
  308. # Create anndata object with reordered embeddings
  309. self.adata.obsm['X_effect_modifier_space'] = all_z_np
  310. pca = PCA(n_components=dim)
  311. # Fit and transform the data
  312. X_pca = pca.fit_transform(self.adata.obsm['X_effect_modifier_space'])
  314. # Store the PCA-reduced data back into adata.obsm
  315. self.adata.obsm['X_effect_modifier_space_PCA'] = X_pca
  316. if save_path_adata is not None:
  317. if save_path_adata.find(".h5ad")<0:
  318. save_path_adata = save_path_adata+".h5ad"
  319. self.adata.write(save_path_adata)
  320. else:
  321. self.adata.write(os.path.join(self.save_pth, 'adata' + '.h5ad'))
  322. print("Effect modifier space saved.")
  324. if self.try_identify_cb_specific_subtypes:
  325. print(cb_specific_subtypes)
  326. cb_specific_subtypes = torch.concat(cb_specific_subtypes, dim=0)
  327. plot_distribution(cb_specific_subtypes.reshape(-1).numpy())
  329. self.adata.obs["original_order"]=np.array(range(self.adata.shape[0]))
  330. if not save_others:
  331. return
  333. torch.save(logits, os.path.join(self.save_pth, "logits.pth"))
  334. self.logits = logits
  336. if self.try_identify_perturb_escaped_cell:
  337. escape_probs = torch.concat(escape_probs, dim=0)
  338. torch.save(escape_probs, os.path.join(self.save_pth, "escape_probs.pth"))
  340. if self.try_identify_cb_specific_subtypes:
  341. cb_specific_subtypes = torch.concat(cb_specific_subtypes, dim=0)
  342. torch.save(cb_specific_subtypes, os.path.join(self.save_pth, "cb_specific_subtypes.pth"))
  344. df_latent=pd.DataFrame(data=self.adata.obsm['X_effect_modifier_space_PCA'],columns=['latent_'+str(i) for i in range(X_pca.shape[1])])
  345. df_latent["batch"]=np.array(self.adata.obs["batch"])
  346. df_latent["condition"]=np.array(self.adata.obs["condition"])
  347. df_latent["group"]=np.array(self.adata.obs["group"])
  348. if save_path_latent_df is not None:
  349. if save_path_latent_df.find(".csv")<0:
  350. save_path_latent_df = save_path_latent_df+".csv"
  351. df_latent.to_csv(save_path_latent_df)
  352. else:
  353. df_latent.to_csv(os.path.join(self.save_pth, 'latent.csv'))
  354. self.df_latent=df_latent
  356. if print_codebook_statisitics:
  357. print(torch.max(self.logits[0],dim=-1))
  358. print("Statistic values of the codebook selection")
  359. print("mean:\n",[torch.mean(logits[i], dim=0) for i in range(len(logits))])
  360. print("variance:\n",[torch.std(logits[i], dim=0) for i in range(len(logits))])
  362. if kBET_test:
  363. run_kbet(df_latent,do_PCA=False)
  365. def run_kBET_test(self):
  366. run_kbet(self.df_latent, do_PCA=False)
  369. if __name__ == '__main__':
  370. test_run_mode=False
  371. train_model=True
  372. if test_run_mode:
  373. import scanpy as sc
  374. import anndata as ad
  376. '''exp = torch.abs(torch.randn((2048, 2000)))
  377. condition = torch.randint(low=0, high=3, size=(2048,), dtype=torch.long)
  378. batch1 = torch.randint(low=0, high=4, size=(2048,), dtype=torch.long) # np.zeros(2048)
  379. batch2 = torch.randint(low=0, high=2, size=(2048,), dtype=torch.long)
  381. adata = ad.AnnData(X=exp.numpy())
  382. adata.obs["condition"] = condition.numpy()
  383. adata.obs["group"] = condition.numpy()
  384. adata.obs["batch1"] = batch1
  385. adata.obs["batch"] = batch1
  386. adata.obs["batch2"] = batch2
  387. adata.obs["leiden1"] = torch.randint(low=0, high=3, size=(2048,), dtype=torch.long).numpy()
  388. adata.obs["leiden1"] = adata.obs["leiden1"].astype("category")
  389. adata.obs["leiden2"] = torch.randint(low=0, high=9, size=(2048,), dtype=torch.long).numpy()
  390. adata.obs["leiden2"] = adata.obs["leiden2"].astype("category")
  392. model = NDreamer_DL(adata=adata, condition_key="condition", contorl_name=0, num_hvg=2000,
  393. developer_test_mode=False, require_batch=True, batch_key="batch1",#["batch1","batch2"],
  394. batch_size=512)
  395. model.train_model()'''
  397. adata = sc.read("../ECCITE_preprocessed.h5ad")
  398. adata.obsm["batch"]=np.expand_dims(adata.obs["batch"].copy(),axis=-1)
  399. print(adata)
  400. print(np.unique(adata.obs["condition"]))
  401. model = NDreamer_DL(adata, condition_key='perturbation', contorl_name='NT', num_hvg=2000, require_batch=True,
  402. batch_key='replicate',
  403. resolution_low=0.5, resolution_high=7, cluster_method="Leiden", embedding_dim=512,
  404. codebooks=[1024 for i in range(32)],
  405. codebook_dim=8, encoder_hidden=[1024, 512], decoder_hidden=[512, 1024], z_dim=256,
  406. cos_loss_scaler=20, random_seed=123, batch_size=1024, epoches=1, lr=1e-3,
  407. triplet_margin=5, independent_loss_scaler=1000, save_pth="./model/",
  408. developer_test_mode=True,
  409. library_size_normalize_adata=False,
  410. save_preprocessed_adata_path=None,
  411. KL_scaler=5e-3, reconstruct_scaler=1, triplet_scaler=5, num_triplets_per_label=15,
  412. tau=0.01, commitment_loss_scaler=1, cluster_correlation_scaler=50, reset_threshold=1 / 1024,
  413. reset_interval=30, try_identify_cb_specific_subtypes=False,
  414. local_neighborhood_loss_scaler=1, local_neighbor_sigma=1,
  415. try_identify_perturb_escaped_cell=False, n_neighbors=20,
  416. local_neighbor_across_cluster_scaler=20)
  417. model.train_model()
  418. print()
  419. model.get_modifier_space(save_path_adata="../ECCITE_results.h5ad")
  420. adata1 = model.adata.copy()
  421. sc.pp.neighbors(adata1, use_rep='X_effect_modifier_space_PCA', n_neighbors=15)
  422. sc.tl.umap(adata1)
  423. sc.pl.umap(adata1, color=["MULTI_ID","HTO_classification",
  424. "gene_target","perturbation",
  425. "replicate","Phase"], frameon=False, ncols=1)
  426. else:
  427. import os
  429. os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
  430. import scanpy as sc
  432. #adata=sc.read_h5ad("../data/PBMC.h5ad")
  433. #adata = sc.read("../PBMC_preprocessed.h5ad")
  434. adata = sc.read("../virus_preprocessed.h5ad")
  435. #adata=sc.read_h5ad("../PBMC_imbalance_preprocessed.h5ad")
  436. #adata=sc.read_h5ad("../common1.h5ad")
  437. if "condition" not in adata.obsm.keys():
  438. adata.obsm["batch"] = np.ones((adata.shape[0],1))#np.expand_dims(adata.obs["condition"].copy(), axis=1)
  439. print(adata)
  440. print(np.unique(adata.obs["condition"]))
  441. model = NDreamer_DL(adata, condition_key="condition", contorl_name="control", num_hvg=min(adata.shape[1],3608),
  442. require_batch=False,
  443. batch_key=None,
  444. resolution_low=0.5, resolution_high=7, cluster_method="Leiden", embedding_dim=512,
  445. codebooks=[1024 for i in range(32)],
  446. codebook_dim=8, encoder_hidden=[2048, 1024], decoder_hidden=[512, 1024], z_dim=256,
  447. cos_loss_scaler=20, random_seed=123, batch_size=1024, epoches=100, lr=1e-3,
  448. triplet_margin=5,independent_loss_scaler=1000, save_pth="./model/",
  449. developer_test_mode=True,
  450. library_size_normalize_adata=False,
  451. save_preprocessed_adata_path=None,
  452. KL_scaler=5e-3, reconstruct_scaler=1, triplet_scaler=5, num_triplets_per_label=15,
  453. tau=0.01, commitment_loss_scaler=1, cluster_correlation_scaler=50,reset_threshold=1/1024,
  454. reset_interval=30,try_identify_cb_specific_subtypes=False,
  455. local_neighborhood_loss_scaler=1,local_neighbor_sigma=1,
  456. try_identify_perturb_escaped_cell=False,n_neighbors=20,
  457. local_neighbor_across_cluster_scaler=20, have_negative_data=True)
  458. '''model = NDreamer_DL(adata, condition_key="condition", contorl_name="control", num_hvg=min(adata.shape[1], 3608),
  459. require_batch=False,
  460. batch_key=None,have_negative_data=True,developer_test_mode=True)'''
  461. if train_model:
  462. model.train_model()
  463. print()
  464. model.load_state_dict(torch.load(os.path.join(model.save_pth, 'ndreamer' + '.pth')))
  465. model.get_modifier_space(save_path_adata="../PBMC_results.h5ad")
  467. adata1 = model.adata.copy()
  468. sc.pp.neighbors(adata1, use_rep='X_effect_modifier_space_PCA',n_neighbors=15)
  469. sc.tl.umap(adata1)
  470. try:
  471. sc.pl.umap(adata1, color=['condition', 'cell_type'], frameon=False, ncols=1)
  472. except:
  473. sc.pl.umap(adata1, color=['condition', 'cell_type1021'], frameon=False, ncols=1)
  475. sc.pp.neighbors(adata1, use_rep='X_effect_modifier_space',n_neighbors=15)
  476. sc.tl.umap(adata1)
  477. try:
  478. sc.pl.umap(adata1, color=['condition', 'cell_type'], frameon=False, ncols=1)
  479. except:
  480. sc.pl.umap(adata1, color=['condition', 'cell_type1021'], frameon=False, ncols=1)
  482. #model.run_kBET_test()

model.py at commit dfff8e9, under GPL-3.0 · at the source

Overview

Authors: Xiao Xiao1, Hongyu Zhao1,2,3, Zuoheng Wang1,4
ORCID iDs: Hongyu Zhao, Zuoheng Wang
  1. Department of Biostatistics, Yale University School of Public Health, 300 George Street, New Haven, CT 06511, United States
  2. Department of Genetics, Yale University School of Medicine, 333 Cedar Street, New Haven, CT 06520, United States
  3. Department of Statistics and Data Science, Yale University, 219 Prospect Street, New Haven, CT 06511, United States
  4. Department of Biomedical Informatics & Data Science, Yale University School of Medicine, 101 College Street, New Haven, CT 06510, United States
Institutions: Yale University (United States)
Journal: Briefings in bioinformatics, volume 27, issue 5, article bbag485
Dates: received 17 February 2026; accepted 6 August 2026; published online 14 September 2026; in print September 2026
Type: Research article · Language: English
License: CC BY-NC
Identifiers: DOI 10.1093/bib/bbag485 · PMID 42734923 · PMCID PMC13573602 · OpenAlex W7213096958
Open access: gold, a free copy (OpenAlex)
Status: code verified
Categories: genetics / omics (modality), human (organism), Alzheimer's / dementia (population)
Methods: Connectivity, Preprocessing, Machine learning
Keywords: scRNA-seq, batch effect, neural discrete representation learning, counterfactual causal matching, perturbation-induced signal, disease-associated signal
MeSH: Alzheimer Disease*, Single-Cell Analysis*, Algorithms, Animals, Gene Expression Profiling, Humans, Single-Cell Gene Expression Analysis (* major topic)
Journal subjects: Problem Solving Protocol
Topic: Single-cell and spatial transcriptomics (Molecular Biology, Biochemistry, Genetics and Molecular Biology), according to OpenAlex
Funding: NIH (R01LM014087)
Citations: not cited yet (Europe PMC); 50 references in the paper

Abstract

Advances in sequencing technologies and the growing volume of single-cell data have created unprecedented opportunities for uncovering gene expression patterns causally induced by experimental perturbations or statistically associated, but not necessarily causal, with disease conditions. However, current analytical methods inadequately account for batch effects and data sparsity or fail to capture the inherent non-linearity in single-cell data, leading to biased estimation. To address these limitations, we developed NDreamer that combines neural discrete representation learning and matching to remove batch effects and estimate perturbation-induced or condition-associated signals at single-cell resolution. NDreamer outperformed existing methods by using mutual information loss on discrete latent variables to disentangle cells’ intrinsic features from conditions or batch effects, while preserving both global and local variance within batches and conditions via triplet and local neighborhood loss. We applied NDreamer to multiple datasets across platforms, organs, and species and validated and benchmarked its performance in removing batch effects and estimating perturbation-induced or condition-associated signals. In particular, we applied NDreamer to an Alzheimer’s disease cohort, revealing biologically relevant gene expression patterns that distinguish dementia patients from controls.

Reproduced under the paper's license (CC BY-NC), from the paper cited above.

Repositories

Its files are read in the Code ↔ Paper reader above, with 24 matches between paragraphs and lines of code.

lugia-xiao/NDreamer

License: GPL-3.0
State: the link answers, verified on 26 September 2026
Evidence: files inventoried
Commit: dfff8e9fb8c0a9bb246468d06d42aee5e666b1bf, 20 March 2025
Languages: Python (13), Jupyter (2)
Size: 20 files, 15 scripts
Software Heritage: not archived
Found in: “Code availability”
Holds: README, license file, environment (setup.py), 2 notebooks
Not found: CITATION.cff, tests, continuous integration, documentation
Tools: NumPy (12 files), PyTorch (11 files), Scanpy (9 files), SciPy (6 files), pandas (5 files), scikit-learn (5 files), anndata (4 files), Matplotlib (4 files), rpy2 (3 files), seaborn (2 files), statsmodels (1 file), UMAP (1 file)
Availability: 1 check, the latest on 26 September 2026: the link answers
  • 26 September 2026: the link answers
17 files
At the source: github.com/lugia-xiao/NDreamer

lugia-xiao/NDreamer_reproducible

License: GPL-3.0
State: the link answers, verified on 26 September 2026
Evidence: files inventoried
Commit: 118684aa70fe23fcb60985c54121984a27ede9d2, 16 May 2026
Languages: Jupyter (136), Python (65), R (13)
Size: 471 files, 214 scripts
Software Heritage: not archived
Found in: “Code availability”
Holds: README, license file, environment (case_control_benchmark/scDisInFact_code/scDisInFact-main/setup.py), tests, 136 notebooks
Not found: CITATION.cff, continuous integration, documentation
Tools: NumPy (15 files), Scanpy (15 files), Matplotlib (12 files), rpy2 (11 files), ggplot2 (9 files), PyTorch (9 files), scikit-learn (9 files), pandas (7 files), anndata (6 files), seaborn (6 files), tidyverse (6 files), clusterProfiler (5 files), SciPy (3 files), reshape2 (1 file)
Availability: 1 check, the latest on 26 September 2026: the link answers
  • 26 September 2026: the link answers
29 files
At the source: github.com/lugia-xiao/NDreamer_reproducible

Code availability

The NDreamer package is available at https://github.com/lugia-xiao/NDreamer with detailed tutorials. All code used to reproduce the results shown in the article are available at https://github.com/lugia-xiao/NDreamer_reproducible.

Reproduced under the paper's license (CC BY-NC), from the paper cited above.

Tracing map

Proposed by the machine: these links were found in the paper and verified at the source, without human review. The map will receive a Zenodo DOI once one of the paper's authors has validated it with their ORCID.

What the map holds:

  • 2 repositories of the authors' code, each at its verified commit, with its license and how the link was found in the paper;
  • 43 scripts, each with its path and the digest of its content;
  • 24 matches between paragraphs of the paper and lines of the code (method lexical-v1);
  • neither the text of the paper nor the code itself.

Its JSON (tracing-map.json) is deposited on Zenodo with its DOI once the map is validated.

Data

Datasets cited

Data availability

The PBMC dataset is available from the GEO under accession number GSE96583 and can be downloaded from the scGen tutorial https://drive.google.com/uc?id=1r87vhoLLq6PXAYdmyyd89zG90eJOFYLk. The rhinovirus infection dataset and another PBMC dataset are available at https://datadryad.org/stash/dataset/doi:10.5061/dryad.4xgxd25g1 (https://doi.org/10.5061/dryad.4xgxd25g1). The ECCITE-seq dataset is available from the GEO under accession number GSE153056 and can be downloaded from pertpy [2] using the code ‘pertpy.data.papalexi_2021()’. The ASD dataset is available from the GEO under accession number GSE157977. The T1D dataset is available from the GEO under accession number GSE148073 and can be downloaded from https://cellxgene.cziscience.com/collections/51544e44-293b-4c2b-8c26-560678423380. The kidney dataset is available from the GEO under accession number GSE211785, and we use the version ‘GSE211785_Susztak_SC_SN_ATAC_merged_PreSCVI_final.h5ad.gz’. The mouse radiation therapy dataset is available from the GEO under accession number GSE280883. The SEA-AD MTG snRNA-seq dataset is available at https://cellxgene.cziscience.com/collections/1ca90a2d-2943-483d-b678-b809bf464c30. The SEA-AD dataset is available at https://cellxgene.cziscience.com/collections/1ca90a2d-2943-483d-b678-b809bf464c30.

Reproduced under the paper's license (CC BY-NC), from the paper cited above.

Versions

The history of this record: each version stored by the harvester or made by a correction of its authors or of the maintainers of its code, and what changed in its facts. The texts of the paper (its abstract, its availability statements) are not part of it; versions that changed only those are not listed.

Version 1, 27 September 2026: the first record

Recorded: type, language, journal, volume, issue, pages, dates, 3 authors, 6 keywords, 7 MeSH terms, 1 funder, 50 references.

Cite

This paper

Xiao, X., Zhao, H., & Wang, Z. (2026). Single-cell-level perturbation-induced and condition-related signal estimation with batch effect removal using NDreamer. Briefings in bioinformatics, 27(5), bbag485. https://doi.org/10.1093/bib/bbag485

BibTeX

@article{xiao2026single,
author = {Xiao, Xiao and Zhao, Hongyu and Wang, Zuoheng},
title = {{Single-cell-level perturbation-induced and condition-related signal estimation with batch effect removal using NDreamer}},
journal = {Briefings in bioinformatics},
year = {2026},
month = sep,
volume = {27},
number = {5},
pages = {bbag485},
publisher = {Oxford University Press},
issn = {1467-5463},
doi = {10.1093/bib/bbag485},
url = {https://doi.org/10.1093/bib/bbag485},
pmid = {42734923},
pmcid = {PMC13573602}
}

RIS

TY - JOUR
AU - Xiao, Xiao
AU - Zhao, Hongyu
AU - Wang, Zuoheng
TI - Single-cell-level perturbation-induced and condition-related signal estimation with batch effect removal using NDreamer
T2 - Briefings in bioinformatics
J2 - Brief Bioinform
PY - 2026
DA - 2026/09/01
VL - 27
IS - 5
SP - bbag485
SN - 1467-5463
PB - Oxford University Press
DO - 10.1093/bib/bbag485
UR - https://doi.org/10.1093/bib/bbag485
LA - en
ER -

CSL-JSON

{
"id": "10.1093/bib/bbag485",
"type": "article-journal",
"title": "Single-cell-level perturbation-induced and condition-related signal estimation with batch effect removal using NDreamer",
"container-title": "Briefings in bioinformatics",
"author": [
{
"family": "Xiao",
"given": "Xiao"
},
{
"family": "Zhao",
"given": "Hongyu"
},
{
"family": "Wang",
"given": "Zuoheng"
}
],
"container-title-short": "Brief Bioinform",
"volume": "27",
"issue": "5",
"page": "bbag485",
"DOI": "10.1093/bib/bbag485",
"PMID": "42734923",
"PMCID": "PMC13573602",
"ISSN": "1467-5463",
"publisher": "Oxford University Press",
"URL": "https://doi.org/10.1093/bib/bbag485",
"language": "en",
"issued": {
"date-parts": [
[
2026,
9,
1
]
]
}
}

The tracing map gets a citation of its own once an author has validated it and it has a DOI.

Similar papers

The papers with a page that share the most with this one: the tools found in their code, their categories, datasets, cited references and authors, the rarest counting most.

[1] doi:10.1093/bib/bbag175 [code]
Uncovering causal relationships in single-cell omic studies with causarray.
Journal: Briefings in bioinformatics
In common: anndata, clusterProfiler, Scanpy, 9 other tools, Alzheimer's / dementia, 4 references
[2] doi:10.1016/j.cpblue.2026.100007 [code]
An integrated single-cell and spatial proteotranscriptomics atlas of fibroblast-driven immunoregulation within the human adult oral cavity.
Journal: Cell press blue
In common: rpy2, UMAP, anndata, 12 other tools, 1 reference
[3] doi:10.1038/s41586-026-10490-y [code]
Lineage and organ signals sequentially build organ intrinsic nervous systems.
Journal: Nature
In common: UMAP, anndata, Scanpy, 10 other tools, author Hongyu Zhao
[4] doi:10.1016/j.isci.2026.116055 [code]
Mapping the transcriptional diversity of calcium signaling in the mouse and human brain.
Journal: iScience
In common: rpy2, UMAP, anndata, 11 other tools, genetics / omics, 1 reference
[5] doi:10.1016/j.xcrm.2026.102766 [code]
A longitudinal single-cell and spatial multiomic atlas of pediatric high-grade glioma.
Journal: Cell reports. Medicine
In common: rpy2, UMAP, anndata, 11 other tools, genetics / omics
[6] doi:10.1038/s41514-026-00391-9 [code]
Region-specific transcriptional signatures of brain aging in the absence of neuropathology at the single-cell level.
Journal: npj aging
In common: anndata, Scanpy, reshape2, 10 other tools, genetics / omics, 2 references
[7] doi:10.1038/s41586-026-10629-x [code]
Whole-genome duplication shaped cell-type evolution in the vertebrate brain.
Journal: Nature
In common: UMAP, anndata, clusterProfiler, 11 other tools, genetics / omics
[8] doi:10.1371/journal.pcbi.1014327 [code]
Supervised deep learning with gene functional annotation for cell classification.
Journal: PLoS computational biology
In common: anndata, Scanpy, reshape2, 10 other tools, Alzheimer's / dementia, genetics / omics, 1 reference
[9] doi:10.1016/j.celrep.2026.117073 [code]
Single-cell epigenomics uncovers heterochromatin instability and transcription factor dysfunction during mouse brain aging.
Journal: Cell reports
In common: rpy2, anndata, clusterProfiler, 9 other tools, genetics / omics, 1 reference
[10] doi:10.21203/rs.3.rs-9676637/v1 [code]
A Comprehensive Benchmarking of Spatial Deconvolution and Domain Detection Methods across Diverse Tissues and Spatial Transcriptomic Technologies
Journal: Research Square (preprint)
In common: rpy2, UMAP, anndata, 10 other tools, genetics / omics

Contribute

The authors of this paper can claim it, correct its record and validate its tracing map, and the maintainers of its code (its owner, or a public member of its organization) correct what it says of their repository; anyone signed in can ask for its removal. Every request goes to OSCR's own machine, which answers it; your account page follows them.

Sign in with ORCID to claim this paper as one of its authors, correct its record or validate its tracing map: when the paper's metadata lists your ORCID iD, you are recognized at once. Maintainers of its code: sign in with GitHub, then claim the repository on your account page.

Request its removal

To ask OSCR to remove this record, the copies of its authors' scripts or its tracing map, use the removal request page: signed in, you say who you are, what to remove and why, then review and confirm the request. Published rules decide every request (how).

Discussion, reproductions, activity

Discussion: questions and error reports about this paper and its code, from signed-in readers and its authors. It opens with sign-in.

Reproductions: reports from readers who ran the authors' code: what they reproduced, with which environment, commit and data. It opens with sign-in.

Activity: what happens around this paper: new versions of its record, its map's validation, discussions and reproductions. It opens with sign-in.