使用RDKit在单张图中叠加高亮小分子药效团特征

问题描述 #

现有代码可正常检测并高亮小分子的药效团特征，但每个特征单独显示在一张图片中。需求是将所有药效团特征叠加显示在同一张图片内，并用不同颜色区分不同类型的特征。

修改后代码 #

from rdkit import Chem
from rdkit.Chem import ChemicalFeatures
from rdkit import rdBase
from rdkit.RDPaths import RDDataDir
from rdkit.Chem import Draw
import os

print(rdBase.rdkitVersion)

# 构建特征工厂
fdefFile = os.path.join(RDDataDir,'BaseFeatures.fdef')
featFact = ChemicalFeatures.BuildFeatureFactory(fdefFile)

# 加载分子
mols = [Chem.MolFromSmiles('C1=NC(=C2C(=N1)N(C=N2)C3C(C(C(O3)CO)O)O)N')]
# 获取分子的药效团特征
featslists = [featFact.GetFeaturesForMol(mol) for mol in mols]

def drawp4core(mol, feats):
    # 自定义不同药效团类型的高亮颜色（RGB格式，取值0-1）
    color_map = {
        'Donor': (0, 1, 0),        # 绿色表示氢键供体
        'Acceptor': (1, 0, 0),      # 红色表示氢键受体
        'Aromatic': (0, 0, 1),      # 蓝色表示芳香环
        'Hydrophobe': (1, 1, 0),    # 黄色表示疏水基团
        'PosIonizable': (1, 0, 1),  # 紫色表示正离子化位点
        'NegIonizable': (0, 1, 1)   # 青色表示负离子化位点
    }
    
    highlight_atom_groups = []
    highlight_color_list = []
    legend_items = []
    
    for feat in feats:
        atom_ids = feat.GetAtomIds()
        feat_type = feat.GetType()
        # 收集要高亮的原子组
        highlight_atom_groups.append(atom_ids)
        # 匹配对应颜色，无匹配则用灰色
        highlight_color_list.append(color_map.get(feat_type, (0.5, 0.5, 0.5)))
        # 生成图例项
        legend_items.append(f"{feat_type}: 原子ID {atom_ids}")
    
    # 合并图例内容
    combined_legend = "\n".join(legend_items)
    # 绘制单图叠加所有高亮特征
    return Draw.MolToImage(
        mol,
        highlightAtomLists=highlight_atom_groups,
        highlightColors=highlight_color_list,
        legend=combined_legend,
        size=(600, 600)  # 可自定义图片尺寸
    )

# 生成图片
im = drawp4core(mols[0], featslists[0])
print(type(im))
im.show()

关键改动说明 #

• 取消多分子副本生成逻辑，改为针对单个分子处理所有药效团特征
• 新增color_map字典，为不同类型的药效团分配专属颜色，直观区分特征类型
• 使用highlightAtomLists参数传递多组原子ID，highlightColors参数对应每组原子的高亮颜色，实现同图叠加效果
• 汇总所有特征的类型和原子ID到图例，方便查看每个高亮区域对应的药效团信息

内容的提问来源于stack exchange，提问作者Zayyan Masud