探索「老藥新用」最短路徑：亞馬遜AI Lab開源大規模藥物重定位知識圖譜DRKG

Drug Repurposing Knowledge Graph

• https://github.com/gnn4dr/DRKG/blob/master/embedding_analysis/Train_embeddings.ipynb

• https://github.com/gnn4dr/DRKG/blob/master/drug_repurpose/COVID-19_drug_repurposing.ipynb

• 下載 DRKG 知識圖譜，DRKG 知識圖譜已開源：https://dgl-data.s3-us-west-2.amazonaws.com/dataset/DRKG/drkg.tar.gz

import syssys.path.insert(1, '../utils')from utils import download_and_extractdownload_and_extract()drkg_file = '../data/drkg/drkg.tsv'

• DRKG 知識圖譜包含一個 tsv 格式檔案 drkg.tsv，其中包含了知識圖譜的所有三元組，在訓練之前，我們將資料集隨機按照 0.9:0.05:0.05 的比例劃分成訓練集、驗證集和測試集。

import pandas as pdimport numpy as npdf = pd.read_csv(drkg_file, sep="\t")triples = df.values.tolist()seed = np.arange(num_triples)np.random.shuffle(seed)train_cnt = int(num_triples * 0.9)valid_cnt = int(num_triples * 0.05)train_set = seed[:train_cnt]train_set = train_set.tolist()valid_set = seed[train_cnt:train_cnt+valid_cnt].tolist()test_set = seed[train_cnt+valid_cnt:].tolist()with open("train/drkg_train.tsv", 'w+') as f:  for idx in train_set:     f.writelines("{}\t{}\t{}\n".format(triples[idx][0], triples[idx][1], triples[idx][2]))with open("train/drkg_valid.tsv", 'w+') as f:  for idx in valid_set:    f.writelines("{}\t{}\t{}\n".format(triples[idx][0], triples[idx][1], triples[idx][2]))    with open("train/drkg_test.tsv", 'w+') as f:   for idx in test_set:     f.writelines("{}\t{}\t{}\n".format(triples[idx][0], triples[idx][1], triples[idx][2]))

• 隨後直接呼叫 DGL-KE 軟體包的命令列進行 DRKG 知識圖譜的低緯嵌入向量表示訓練，案例中我們選用 TransE_l2 知識圖譜嵌入演算法，並使用了 AWS p3.16xlarge 例項進行多 GPU 並行進行訓練（使用其他知識圖譜嵌入演算法以及其他機型可以參考 https://aws-dglke.readthedocs.io/en/latest/index.html 中的說明）。

!DGLBACKEND=pytorch dglke_train --dataset DRKG --data_path ./train --data_files drkg_train.tsv drkg_valid.tsv drkg_test.tsv --format 'raw_udd_hrt' --model_name TransE_l2 --batch_size 2048 \--neg_sample_size 256 --hidden_dim 400 --gamma 12.0 --lr 0.1 --max_step 100000 --log_interval 1000 --batch_size_eval 16 -adv --regularization_coef 1.00E-07 --test --num_thread 1 --gpu 0 1 2 3 4 5 6 7 --num_proc 8 --neg_sample_size_eval 10000 --async_update

• 訓練完成後我們將得到兩個檔案: 1) DRKG_TransE_l2_entity.npy, DRKG 中實體的低維向量表示和 2）DRKG_TransE_l2_relation.npy，DRKG 中關係的低維向量表示。後續我們可以使用訓練好的實體和關係的低維向量表示進行藥物預測。

node_emb = np.load('./ckpts/TransE_l2_DRKG_0/DRKG_TransE_l2_entity.npy')relation_emb = np.load('./ckpts/TransE_l2_DRKG_0/DRKG_TransE_l2_relation.npy')

• 設定目標病毒實體、藥物實體和治療關係。

# 目標新冠病毒相關實體COV_disease_list = ['Disease::SARS-CoV2 E','Disease::SARS-CoV2 M', ...]# 藥物疾病治療相關關係treatment = ['Hetionet::CtD::Compound:Disease','GNBR::T::Compound:Disease']# 獲取來自 Drugbank 的分子量 (molecule weight) 大於 250 的 FDA 獲准藥物實體（已在 infer_drug.tsv 中提供drug_list = []with open("./infer_drug.tsv", newline='', encoding='utf-8') as csvfile:   reader = csv.DictReader(csvfile, delimiter='\t', fieldnames=['drug','ids'])     for row_val in reader:       drug_list.append(row_val['drug'])

• 獲取預訓練 DRKG 知識圖譜的嵌入表示

# 讀取預訓練 embeddingentity_emb = np.load('../data/drkg/embed/DRKG_TransE_l2_entity.npy')rel_emb = np.load('../data/drkg/embed/DRKG_TransE_l2_relation.npy')drug_ids = th.tensor(drug_ids).long()disease_ids = th.tensor(disease_ids).long()treatment_rid = th.tensor(treatment_rid)drug_emb = th.tensor(entity_emb[drug_ids])treatment_embs = [th.tensor(rel_emb[rid]) for rid in treatment_rid]

• 所有可能的(藥物，治療，病毒）三元組組合在 TrainsE_l2 演算法下的分數(score)，計算公式如下：

import torch.nn.functional as fngamma=12.0def transE_l2(head, rel, tail):  score = head + rel - tail  return gamma - th.norm(score, p=2, dim=-1)  scores_per_disease = [] dids = [] # 針對兩種治療關係分別計算(藥物，治療，病毒）三元組的分數，並最終合併 for rid in range(len(treatment_embs)):    treatment_emb=treatment_embs[rid]    for disease_id in disease_ids:       disease_emb = entity_emb[disease_id]        score = fn.logsigmoid(transE_l2(drug_emb, treatment_emb, disease_emb))         scores_per_disease.append(score)         dids.append(drug_ids)scores = th.cat(scores_per_disease)

• 對分數進行排序

idx = th.flip(th.argsort(scores), dims=[0])scores = scores[idx].numpy()dids = dids[idx].numpy()

• 獲取最終 topk 的藥物推薦

topk=100_, unique_indices = np.unique(dids, return_index=True)topk_indices = np.sort(unique_indices)[:topk]# top100 的藥物 IDproposed_dids = dids[topk_indices]# top100 的分數proposed_scores = scores[topk_indices]

• 最終得到的藥物中，目前已經處於臨床實驗的有 6 例，具體結果如下（排名，藥物名稱，相關分數）

[0]    Ribavirin          -0.21416784822940826[4]    Dexamethasone      -0.9984006881713867[8]    Colchicine         -1.080674648284912[16]   Methylprednisolone -1.1618402004241943[49]   Oseltamivir        -1.3885014057159424[87]   Deferoxamine       -1.513066053390503