基于 AI 的个性化癌症治疗方案设计：从多组学数据到可解释模型的全流程实践

# 1. 安装依赖
pip install gdc-client==1.6.1 pandas==2.1.4 anndata==0.10.3
# 2. 下载 TCGA-COAD/READ RNA-seq + 临床
gdc-client download -m gdc_manifest_coad.txt -d tcga_coad
gdc-client download -m gdc_manifest_read.txt -d tcga_read
# 3. 合并并质控
python scripts/merge_tcga.py --min_tpm 1 --min_samples 100

import anndata as ad
import pandas as pd

rna = ad.read_h5ad('tcga_crc_rna.h5ad')
cna = ad.read_h5ad('tcga_crc_cna.h5ad')
mut = pd.read_csv('tcga_crc_mut.csv', index_col=0)

# 对齐样本
common = rna.obs.index.intersection(cna.obs.index).intersection(mut.index)
X = pd.concat([
    pd.DataFrame(rna[common].X, index=common),
    pd.DataFrame(cna[common].X, index=common),
    mut.loc[common]
], axis=1)
X.to_pickle('tcga_crc_multiomics.pkl')

gdsc = pd.read_csv('GDSC2_fitted_dose_response.csv')
gdsc['label'] = (gdsc['IC50'] <= 1).astype(int)
gdsc[['COSMIC_ID', 'DRUG_NAME', 'label']].to_csv('gdsc_drug_labels.csv')

import torch
import torch.nn as nn
from torch.nn import TransformerEncoder, TransformerEncoderLayer

class OmicsEncoder(nn.Module):
    def __init__(self, input_dim, d_model=512, nhead=8):
        super().__init__()
        self.linear = nn.Linear(input_dim, d_model)
        encoder_layer = TransformerEncoderLayer(d_model, nhead, batch_first=True)
        self.transformer = TransformerEncoder(encoder_layer, num_layers=4)
        
    def forward(self, x):
        x = self.linear(x).unsqueeze(1)  # (B,1,d)
        return self.transformer(x).squeeze(1)

class DrugEncoder(nn.Module):
    def __init__(self, vocab_size=2048, d_model=512):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        encoder_layer = TransformerEncoderLayer(d_model, 8, batch_first=True)
        self.transformer = TransformerEncoder(encoder_layer, num_layers=4)
        
    def forward(self, z):
        z = self.embed(z)  # (B,L,d)
        return self.transformer(z).mean(1)  # 全局池化

class PatientDrugModel(pl.LightningModule):
    def __init__(self, omics_dim, vocab_size):
        super().__init__()
        self.omics_enc = OmicsEncoder(omics_dim)
        self.drug_enc = DrugEncoder(vocab_size)
        self.attn = nn.MultiheadAttention(512, 8, batch_first=True)
        self.fc = nn.Linear(512, 1)
        
    def forward(self, x_omics, x_drug):
        h_p = self.omics_enc(x_omics)  # (B,d)
        h_d = self.drug_enc(x_drug)    # (B,d)
        h_p = h_p.unsqueeze(1)
        h_d = h_d.unsqueeze(1)
        attn_out, weights = self.attn(h_p, h_d, h_d)
        out = torch.sigmoid(self.fc(attn_out.squeeze(1)))
        return out, weights

class RankLoss(nn.Module):
    def forward(self, pos, neg):
        return torch.relu(neg - pos + 0.1).mean()

def training_step(self, batch, _):
    x_omics, x_drug, y = batch
    y_hat, _ = self(x_omics, x_drug)
    bce = nn.functional.binary_cross_entropy(y_hat, y.float())
    
    # 构造正负对
    pos_mask = y == 1
    if pos_mask.sum() > 0 and (~pos_mask).sum() > 0:
        rank = RankLoss()(y_hat[pos_mask], y_hat[~pos_mask])
    else:
        rank = 0
    loss = bce + 0.1*rank
    self.log('train_loss', loss)
    return loss

import itertools, pickle
drugs = list(pd.read_csv('approved_drugs.csv')['smiles'])
combos = list(itertools.combinations(drugs, 2)) + \
         list(itertools.combinations(drugs, 3))
pickle.dump(combos, open('candidate_combos.pkl','wb'))  # ~1.2 M

model = PatientDrugModel.load_from_checkpoint('best.ckpt')
omics = torch.tensor(X.loc['TCGA-3L-AA1B']).float().unsqueeze(0).cuda()

def batch_predict(combos, batch_size=2048):
    results = []
    for i in range(0, len(combos), batch_size):
        batch = combos[i:i+batch_size]
        z = torch.tensor([smiles_to_fp(c) for c in batch]).cuda()
        with torch.no_grad():
            scores, _ = model(omics.repeat(z.size(0),1), z)
        results.extend(scores.cpu().numpy())
    return results

topk = np.argsort(scores)[-5:]

import shap
explainer = shap.DeepExplainer(model.omics_enc, background)
shap_values = explainer.shap_values(x_omics)
top_genes = X.columns[np.argsort(-shap_values[0])[:10]]

from gseapy import enrichr
enr = enrichr(gene_list=top_genes.tolist(), 
              gene_sets=['KEGG_2021_Human'])
enr.results.to_csv('pathway_enrichment.csv')

# 构建知识图谱
import networkx as nx
G = nx.Graph()
for gene in top_genes:
    for drug in topk_drugs:
        if gene in drugbank_targets[drug]:
            G.add_edge(gene, drug, relation='target')
nx.write_gml(G, 'explanation_graph.gml')