原文:[2201.09637] DrugOOD: Out-of-Distribution (OOD) Dataset Curator and Benchmark for AI-aided Drug Discovery -- A Focus on Affinity Prediction Problems with Noise Annotations (arxiv.org)
项目地址:GitHub - tencent-ailab/DrugOOD: OOD Dataset Curator and Benchmark for AI-aided Drug Discovery
主要工作
人工智能辅助药物发现 (AIDD) 因其有望使新药研发更快、更便宜、更高效而越来越受欢迎。尽管它在 ADMET 预测、虚拟筛选、蛋白质折叠和生成化学等许多领域得到广泛应用,但在现实世界 AIDD 中不可避免的噪声分布外 (OOD) 学习问题方面却鲜有探索。应用程序。
在新药研发过程中,药物分子的体内过程中的吸收 (Absorption),分布 (Distribution),代谢 (Metabolism),排泄(Excretion),以及毒性 (Toxicity),简称ADMET性质,是其成药性的重要指标。
作者推出了 DrugOOD,一个用于人工智能辅助药物发现的系统 OOD 数据集管理器和基准,它附带一个开源 Python 包,可以完全自动化数据管理和 OOD 基准测试流程。作者关注AIDD中最关键的问题之一:药物靶点结合亲和力预测,它既涉及大分子(蛋白质靶点)又涉及小分子(药物化合物)。
由于分子数据通常使用图神经网络 (GNN) 主干建模为不规则图,因此 DrugOOD 也可作为图 OOD 学习问题的有价值的测试平台。
背景
虚拟筛选的目的是在存在大量候选化合物的情况下精确定位一小组对给定目标蛋白具有高结合亲和力的化合物。解决虚拟筛选问题的一个关键任务是开发计算方法来预测给定药物-靶标对的结合亲和力,而亲和力预测方法可以大致分为两类:基于配体的和基于结构的。
在药物开发过程中,涉及许多预测任务。在这里考虑计算机辅助药物发现的两个关键任务:基于配体的亲和力预测(ligand based affinity prediction, LBAP)和基于结构的亲和力预测(structure based affinity prediction, SBAP)。在LBAP中,我们遵循惯例,不涉及任何蛋白质靶点信息,通常用于某一特定蛋白质靶点的活性预测。在 SBAP 中,我们考虑靶标和药物信息来预测配体的结合活性,旨在开发可以泛化不同蛋白质靶标的模型。
DrugOOD 基于 ChEMBL(Mendez 等人,2019)构建所有数据集,这是一个大型、开放访问的药物发现数据库,旨在捕获整个药物研究和开发过程中的药物化学数据和知识。我们使用最新版本的SQLite格式:ChEMBL 29。
在药物研发过程中,在预测小分子的生物活性时,我们可能会遇到与模型训练集截然不同的分子支架、尺寸等。这些差异也可能反映在SBAP任务的目标中。因此,对于 LBAP 任务,我们考虑以下三个领域:测定、支架和分子大小。对于SBAP任务,除了上面提到的三个域之外,我们还考虑了两个额外的目标特异性域:蛋白质和蛋白质家族。用户还可以通过配置文件轻松自定义域并生成相应的数据集。(即构成了不同的 distribution 的 domain 可以用于 OOD 测试)
代码实现
数据集的读取
项目使用了SQLite数据库从chembl_29.db
获取数据。具体的实现可以在drugood/curators/chembl/sql_exe.py
文件中的SqlFunctions
类找到。
例如,get_all_assay_ids
方法从ASSAYS
表中获取所有的ASSAY_ID
:
def get_all_assay_ids(self, ):
return self.__get_one_col_from_table__(table_name='ASSAYS', col_name='ASSAY_ID')
这个方法调用了__get_one_col_from_table__
方法,该方法执行一个SQL查询来从指定的表中获取指定的列:
def __get_one_col_from_table__(self, table_name, col_name):
sql = 'select {col_name} from {table_name}'.format(table_name=table_name, col_name=col_name)
res = self.sql_exe(sql, return_cur=False)
ans = []
for item in res:
ans.append(item[0])
return ans
在README.md
文件中,项目提供了如何配置数据库路径和数据保存目录的说明。在configs/_base_/curators/lbap_defaults.py
或configs/_base_/curators/sbap_defaults.py
文件中指定source_root
为chembl_29.db
的路径,target_root
为数据保存目录。
然后,运行tools/curate.py
脚本并指定配置文件来生成数据集。例如:
python tools/curate.py --cfg configs/curators/lbap_core_ic50_assay.py
以上是从chembl_29.db
获取数据的基本流程。
数据的格式
根据论文内容,数据集中的每个小分子都表示为一个图,其中节点是原子,边是化学键。按照 Pushing the Boundaries of Molecular Representation for Drug Discovery with the Graph Attention Mechanism | Journal of Medicinal Chemistry (acs.org) 中的预处理策略,通过 RDKit 包对分子进行预处理。输入节点特征是39维向量,包括原子符号、杂化、氢等。输入边特征是 10 维向量,包括键类型、共轭、环和键立体化学。
RDKit 是一个开源的化学信息学和机器学习工具包:GitHub - rdkit/rdkit: The official sources for the RDKit library
在 drugood/utils/smile_to_dgl.py
文件中,get_atom_features
函数用于获取原子特性,返回一个39维的向量,包括原子符号、杂化、氢等特性:
def get_atom_features(atom):
feature = np.zeros(39)
...
return feature
同样在 drugood/utils/smile_to_dgl.py
文件中,get_bond_features
函数用于获取键的特性,返回一个10维的向量,包括键类型、共轭、环和键立体化学等特性:
def get_bond_features(bond):
feature = np.zeros(10)
...
return feature
这两个函数都被用在 smile2graph
函数中,该函数将分子的 SMILES 表示转换为图表示,其中节点和边的特性分别由 get_atom_features
和 get_bond_features
函数提供。
def smile2graph(smile):
mol = Chem.MolFromSmiles(smile)
if (mol is None):
return None
src = []
dst = []
atom_feature = []
bond_feature = []
try:
for atom in mol.GetAtoms():
one_atom_feature = get_atom_features(atom)
atom_feature.append(one_atom_feature)
for bond in mol.GetBonds():
i = bond.GetBeginAtomIdx()
j = bond.GetEndAtomIdx()
one_bond_feature = get_bond_features(bond)
src.append(i)
dst.append(j)
bond_feature.append(one_bond_feature)
src.append(j)
dst.append(i)
bond_feature.append(one_bond_feature)
src = torch.tensor(src).long()
dst = torch.tensor(dst).long()
atom_feature = np.array(atom_feature)
bond_feature = np.array(bond_feature)
atom_feature = torch.tensor(atom_feature).float()
bond_feature = torch.tensor(bond_feature).float()
graph_cur_smile = dgl.graph((src, dst), num_nodes=len(mol.GetAtoms()))
graph_cur_smile.ndata['x'] = atom_feature
graph_cur_smile.edata['x'] = bond_feature
return graph_cur_smile
except RuntimeError:
return None
最后返回的graph_cur_smile
是一个 DGL 的Graph对象,它的结构包括节点和边,以及与节点和边相关联的特征。
标签的生成
def classification_label_generating(self, data):
""" Generate classification label for data
Args:
data (List(Dict)):
Returns:
output (List(Dict))
"""
all_values = [case['reg_label'] for case in data]
all_values = sorted(all_values)
median = all_values[int(len(all_values) * 0.5)]
if self.cfg.classification_threshold.lower_bound <= median <= self.cfg.classification_threshold.upper_bound:
thr = median
else:
thr = self.cfg.classification_threshold.fix_value
self.statistics['thr_for_cls'] = thr
positive_samples = 0
negative_samples = 0
for line in data:
if line['reg_label'] >= thr:
positive_samples += 1
cls_label = 1
else:
negative_samples += 1
cls_label = 0
line['cls_label'] = cls_label
self.statistics['positive_samples'] = positive_samples
self.statistics['negative_samples'] = negative_samples
self.statistics['positive_rate'] = positive_samples / (positive_samples + negative_samples)
return data
drugood/curators/curator.py
中GenericCurator
类中有一个名为 classification_label_generating
的方法接收一个数据列表作为输入,每个元素都是一个字典,然后返回处理后的数据。
首先,它从输入数据中提取出所有的 reg_label
值并求中位数。如果中位数在配置文件中设置的分类阈值的上下界之间,那么就使用这个中位数作为阈值。否则,就使用配置文件中设置的固定阈值。
接着,它遍历输入数据,对每一行数据进行处理。如果某行数据的 reg_label
值大于或等于阈值,那么就将这行数据的分类标签设置为1,并增加正样本的数量。否则,就将这行数据的分类标签设置为0,并增加负样本的数量。
最后,它计算正样本的比例,即正样本数量除以总样本数量(正样本数量加负样本数量)。这个比例被保存在统计信息中。
这个方法的返回值是处理后的数据,每一行数据都增加了一个新的键值对,键是 cls_label
,值是分类标签(0或1)。
LBAP/SBAP 问题视为二元分类问题,其中输入是小分子的图数据(区别在于SBAP多考虑了蛋白质信息),标签是二元亲和力分类的真实值(有效或无效)。