SATURN：利用蛋白质大模型嵌入进行跨物种整合

最新推荐文章于 2024-11-08 13:43:44 发布

木与长清

最新推荐文章于 2024-11-08 13:43:44 发布

阅读量1.1k

点赞数 30

分类专栏：单细胞空间整合工具学习文章标签： scikit-learn python

本文链接：https://blog.csdn.net/XMQ_MOLLY/article/details/140663310

版权

单细胞空间整合工具学习专栏收录该内容

2 篇文章 0 订阅

订阅专栏

伴随着众多大模型新方法的涌现，跨物种整合领域出现了许多利用蛋白质序列信息进行嵌入的整合方法，有一些具备相当的准确性。

这里记录一下用 Rosen, Y., Brbić, M., Roohani, Y. et al. Toward universal cell embeddings: integrating single-cell RNA-seq datasets across species with SATURN. Nat Methods (2024). https://doi.org/10.1038/s41592-024-02191-z 方法文章中数据进行测试的全流程。

原始代码和数据链接见https://github.com/snap-stanford/SATURN

1.环境准备

https://github.com/snap-stanford/SATURN/blob/main/requirements.txt

建议在配置GPU的环境下训练模型，否则会相当耗时

2.数据准备

物种单细胞数据

SATURN/Vignettes/frog_zebrafish_embryogenesis/dataloader.ipynb at main · snap-stanford/SATURN · GitHub

这里准备的是物种带注释的计数矩阵，并且都需要转成h5ad的格式

物种蛋白质嵌入文件

还需要准备每个物种的蛋白质序列，可以从NCBI-DATASET中下载

提取之前可以进行一定的预处理和数据清洗，见clean_fasta.py(SATURN/protein_embeddings at main · snap-stanford/SATURN · GitHub)

现有的蛋白质大模型很多，这里选择ESM2，提取蛋白质嵌入信息的torch文件详见extract.py

(esm/scripts/extract.py at main · facebookresearch/esm · GitHub)

具体的一些参数解释如下

#Extraction of the final-layer embedding for a FASTA file from the ESM-2 model
python scripts/extract.py esm2_t33_650M_UR50D \
       examples/data/some_proteins.fasta \
       examples/data/some_proteins_emb_esm2 \
       --repr_layers 0 32 33 \   #(default: final only) selects which layers to include embeddings from
       --include mean per_tok
#per_tok includes the full sequence, with an embedding per amino acid (seq_len x hidden_dim).
#mean includes the embeddings averaged over the full sequence, per layer.

我们提取训练的最后一层即可（即第33层）

extract一步会输出一个包含很多蛋白质嵌入的文件夹，最后我们将所有pt整合成一个文件SATURN/protein_embeddings/convert_protein_embeddings_to_gene_embeddings.py at main · snap-stanford/SATURN · GitHub

3.训练模型

准备运行表格

正式训练SATURN前，我们需要准备一个包含所有输入文件的csv表格 frog_zebrafish_run.csv，将自己的单细胞数据路径和蛋白质嵌入pt文件路径放入如下数据框即可

	path	species	embedding_path
0	data/frog.h5ad	frog	data/frog_gene_symbol_to_embedding_ESM2.pt
1	data/zebrafish.h5ad	zebrafish	data/zebrafish_gene_symbol_to_embedding_ESM2.pt

准备匹配打分表

为了在训练过程中对匹配打分，我们还需要准备物种细胞类型注释的匹配表格frog_zebrafish_cell_type_map.csv ，缺失的地方标记为NAN

Unnamed: 0	frog_cell_type	zebrafish_cell_type
0	0	Blastula	NaN
1	1	Germline	Germline
2	2	Neuroectoderm	Neuroectoderm
3	3	Non-neural ectoderm	Non-neural ectoderm
4	4	Involuting marginal zone	Involuting marginal zone
5	5	Spemann organizer	NaN
6	6	Endoderm	Endoderm
7	7	Epidermal progenitor	Epidermal progenitor
8	8	Ionocyte	Ionocyte
9	9	Goblet cell	NaN

训练过程

https://github.com/snap-stanford/SATURN/blob/main/train-saturn.py

Argument	Value	Explanation
in_data	`data/frog_zebrafish_run.csv`	The csv we created containing paths.
in_label_col	`cell_type`	Use the `cell_type` column labels for metric learning. NOTE: SATURN is weakly supervised, it does not share cell type labels across species, so you don't need to match these values across AnnDatas.
ref_label_col	`cell_type`	Extra cell type argument, will be added to our output but won't effect results since we didn't add `--use_ref_labels
num_macrogenes	`2000`	By default, we use 2000 macrogenes.
hv_genes	`8000`	By default, we use the 8000 most highly variable genes.
centroids_init_path	`saturn_results/fz_centroids.pkl`	Since this is the first time we are runinng this command, we will have to initialize our macrogenes using KMeans, which is an expensive operation. We save that initialization to this location so that if we pass this path to this argument in future runs, we can skip that step.
score_adata		By adding this flag, we will score our adatas after pretraining and while fine tuning with metric learning.
ct_map_path	`data/frog_zebrafish_cell_type_map.csv`	The path to our cell type mapping file, needed since we are scoring while training.
work_dir	`./`	SATURN outputs to a folder called `saturn_results`.

参数的解释如上，其中num_macrogenes是提取的宏基因数量，hv_genes是选取的高可变基因数量，添加score_adata可以对预训练以及后续的训练进行打分对比

python code/train-saturn.py --in_data=data/frog_zebrafish_run.csv \
                              --in_label_col=cell_type --ref_label_col=cell_type \
                              --num_macrogenes=2000     --hv_genes=8000          \
                              --centroids_init_path=saturn_results/fz_centroids.pkl \
                              --score_adata --ct_map_path=data/frog_zebrafish_cell_type_map.csv \
                              --work_dir=. \
                              --device_num=1 \

训练过程中的打分

训练结束如下

4.预训练结果

主要的输出结果有两个，一个是物种单细胞整合的h5ad文件test256_data_frog_zebrafish_org_saturn_seed_0.h5ad，

UMAP可视化后

另一个是代表基因匹配宏基因的权重文件 test256_data_frog_zebrafish_org_saturn_seed_0_genes_to_macrogenes.pkl

可以用于进行差异基因表达的查看、对比其他的基因家族等

替换随机数嵌入

为了验证ESM2蛋白质模型对于SATURN结果的影响，我这里将数据准备中的pt嵌入文件中的字典值替换成随机数如下

import torch
import numpy as np

original_file = './data/frog_embedding.torch'
data = torch.load(original_file)

random_dict = {}
for key, value in data.items():
    random_dict[key] = np.random.random(value.shape)

for key, value in data.items():
    value.data = torch.from_numpy(random_dict[key])
    value.data = value.data.float()

output_folder = './'
output_file = output_folder + '/random_frog_embedding.torch'
torch.save(data, output_file)

其他训练模型的步骤一样，输出结果并可视化检查

图形UMAP明显错乱，证明我们的蛋白质嵌入是有作用的

5.多重训练

训练过程

要复制青蛙和斑马鱼胚胎发生的 SATURN 结果，需要用不同的种子运行 SATURN 30 次

下游分析可以选取打分最高的种子结果

这里同样需要准备相应的运行表格 frog_zebrafish_run_multi.csv，https://github.com/snap-stanford/SATURN/blob/main/saturn_multiple_seeds.py

python code/saturn_multiple_seeds.py \
                --run=data/frog_zebrafish_run_multi.csv --embedding_model=ESM2 \
                --gpus 0 1 2 3 \
                --seeds=30 --pe_sim_penalty=0.2

这里gpus指的是当前可以用的GPU设备编号，可以通过下面的代码查询

import torch

if torch.cuda.is_available():
    device_count = torch.cuda.device_count()
    for i in range(device_count):
        device_name = torch.cuda.get_device_name(i)
        print(f"GPU Device {i}: {device_name}")
else:
    print("No GPUs available")

会根据提供的设备分配任务运行

训练种子打分

最后需要对30次训练结果进行打分并保存到一个表格里，需要用到训练模型步骤里的物种细胞类型匹配表 frog_zebrafish_cell_type_map.csv

python ./code/score_adata.py --adata=./data/fz_multi_seeds.csv --scores=1 \
                                 --multiple_files --species1=zebrafish --species2=frog --label=labels2 \
                                 --ct_map=./data/frog_zebrafish_cell_type_map.csv

运行过程如下