nvmolkit-usage

nvMolKit 用法

用途

以 GPU 加速的批处理方式实现常见的 RDKit 操作。API 尽可能与 RDKit 保持一致，但面向批处理：它们接受 rdkit.Chem.Mol 列表（或指纹列表），并在一个或多个 GPU 上并行处理。nvMolKit 在构建时链接 RDKit；输入和输出都是实际的 RDKit Mol 对象。

本指南介绍已安装的 Python API。不涉及从源代码构建 nvMolKit。

nvMolKit 的适用场景

以下情况适合使用 nvMolKit：

• 工作负载是一大批分子一起处理（通常为数千个或更多）。
• 关注的指标是吞吐量／整批处理的总耗时，而非单个分子的延迟。
• 对整批分子重复执行相同操作（如为分子库生成指纹、嵌入／最小化多个构象、批量计算两两相似度），以便让 GPU 保持高利用率。

要求

• 具备计算能力 7.0（V100）或更高的 NVIDIA GPU
• 与 CUDA 12.6+ 兼容的 CUDA 驱动。
• 可正常运行且支持 CUDA 的 torch 安装（nvMolKit 通过 torch 的 CUDA 数组接口返回 GPU 张量）。

协助安装时，应先让用户选择主机驱动支持的 PyTorch CUDA 后端，再安装 nvMolKit。nvMolKit 的 PyPI wheel 使用 CUDA Toolkit 12.9 构建，并依赖 CUDA 12 运行时软件包，但即使如此，pip/uv 仍可能选择 CUDA 13 PyTorch wheel，除非安装命令另行指定。

• Conda：优先使用 conda-forge 的 pytorch-gpu；将 cuda-version 固定为 12.6 或其他受驱动支持的 CUDA 版本。
• pip：引导用户前往 PyTorch 安装选择器或旧版本页面，先安装使用 CUDA 12.x 后端的 torch，再安装 nvMolKit。
• uv：使用显式后端安装 nvMolKit，例如 uv pip install --torch-backend=cu128 nvmolkit。

输入

• 必需：选择一项操作，并提供来自用户代码或分子数据集的分子或指纹。使用 RDKit 解析 SMILES，并拒绝解析失败的结果（None）。
• 分子操作使用 RDKit Mol 对象。使用 ETKDG 时需添加氢原子；最小化和构象比较需要已有构象。
• 指纹相似度计算接受打包后的 AsyncGpuResult、torch 张量或 NumPy 数组：每行对应一个分子，数据类型为 int32 或 uint32。
• 可选：根据用户要求的工作流设置构象数量、指纹参数、截断值、输出模式和硬件选项；否则使用文档所述的 API 默认值。

限制

• 必须使用 CUDA；不支持 CPU 回退。需要在 CPU 上执行时，请直接使用 RDKit。
• 对于单分子任务或无法批处理的操作，通常更适合使用原版 RDKit。
• ETKDG 不支持自定义边界矩阵、自定义 CPCI、坐标映射或分离片段嵌入。
• 子结构搜索不支持基于手性的匹配、增强立体化学或其他高级 RDKit SubstructMatchParameters 选项。

说明

1. 编写 nvMolKit 代码前，先运行下面的冒烟测试。
2. 从入口点表中选择一项 API，并遵循其输入要求。
3. 按照下文说明处理其结果；读取主机端数据前，应先同步异步 GPU 结果。

编写实际代码前验证安装

import nvmolkit
import torch
from rdkit import Chem
from nvmolkit.fingerprints import MorganFingerprintGenerator

print("nvmolkit:", nvmolkit.__version__)
print("cuda available:", torch.cuda.is_available())
print("device count:", torch.cuda.device_count())

mols = [Chem.MolFromSmiles(smi) for smi in ["CCO", "c1ccccc1", "CC(=O)O"]]
fpgen = MorganFingerprintGenerator(radius=2, fpSize=1024)
result = fpgen.GetFingerprints(mols)
torch.cuda.synchronize()
fps = result.torch()
print("fps shape:", tuple(fps.shape), "dtype:", fps.dtype)
# Expected: shape (3, 32), dtype torch.int32  (1024 bits packed into 32 int32s per row)

如果测试失败，请引导用户参阅安装指南，不要自行猜测。

入口点

任务	模块	主要入口点
Morgan 指纹	nvmolkit.fingerprints	MorganFingerprintGenerator(radius, fpSize).GetFingerprints(mols)
批量 Tanimoto / 余弦相似度	nvmolkit.similarity	crossTanimotoSimilarity(...)、crossCosineSimilarity(...)，以及适用于结果过大、无法放入 GPU 内存的 *MemoryConstrained 变体
ETKDG 构象嵌入	nvmolkit.embedMolecules	EmbedMolecules(molecules, params, confsPerMolecule, ...)
MMFF94 优化（单次执行）	nvmolkit.mmffOptimization	MMFFOptimizeMoleculesConfs(molecules, ..., minimizerKind=..., fireOptions=...)
UFF 优化（单次执行）	nvmolkit.uffOptimization	UFFOptimizeMoleculesConfs(molecules, ..., minimizerKind=..., fireOptions=...)
使用自定义选项和约束的力场	nvmolkit.batchedForcefield	MMFFBatchedForcefield(mols, properties=..., nonBondedThreshold=..., ignoreInterfragInteractions=..., hardwareOptions=...)、UFFBatchedForcefield(mols, vdwThreshold=..., ...)。逐分子视图 ff[i] 提供 add_distance_constraint、add_position_constraint、add_angle_constraint、add_torsion_constraint。方法：.compute_energy()、.compute_gradients()、.minimize(maxIters, forceTol, minimizerKind=..., fireOptions=...)
成对构象 RMSD	nvmolkit.conformerRmsd	GetConformerRMSMatrix(mol)、GetConformerRMSMatrixBatch(mols)
扭转指纹偏差（TFD）	nvmolkit.tfd	GetTFDMatrix(mol)、GetTFDMatrices(mols)
Butina 聚类	nvmolkit.clustering	butina(distance_matrix, cutoff)（预计算矩阵）、fused_butina(fingerprints, cutoff)（内存高效，按需计算）；两者均支持显式 RDKit 和设备输出模式
子结构搜索	nvmolkit.substructure	hasSubstructMatch、countSubstructMatches、getSubstructMatches
最大公共子结构	nvmolkit.mcs	findMCS(mols, ...)，适用于所有配对、显式配对或两个一一配对的分子列表
硬件调优（批大小、GPU ID）	nvmolkit.types	传递给 ETKDG / MMFF / UFF 的 HardwareOptions(...)
可选自动调优	nvmolkit.autotune	tune_embed_molecules、tune_mmff_optimize、tune_uff_optimize、tune_batched_forcefield、tune_substructure、tune_mcs。需要 optuna 包

结果类型和执行模型

根据操作的输出类型，有两种返回形式可承载驻留在 GPU 上的输出。

AsyncGpuResult

用于返回单个扁平张量的操作（指纹、相似度矩阵、RMSD/TFD 向量、Butina 输入）。主要行为：

• 异步执行。调用返回时，内核可能尚未完成。
• result.torch() 返回 GPU 上的零拷贝 torch.Tensor。调用方负责在主机上读取值之前进行同步。
• result.numpy() 会进行同步并返回 CPU numpy 数组。
• 提供 __cuda_array_interface__，因此可直接传入其他 nvMolKit 函数（例如指纹 → 相似度），无需往返主机。

CUDA 流控制

部分返回 AsyncGpuResult 的 API 接受可选参数 stream: torch.cuda.Stream | None = None，调用方可借此将 nvMolKit 工作提交到非默认流，并与自己的内核重叠执行。省略该参数时，调用会使用当前 torch 流。

接受 stream 参数的 API：

• MorganFingerprintGenerator.GetFingerprints
• crossTanimotoSimilarity、crossCosineSimilarity 及其 *MemoryConstrained 变体
• butina、fused_butina
• GetConformerRMSMatrix、GetConformerRMSMatrixBatch

其他 API（ETKDG、MMFF/UFF 优化、TFD、子结构搜索、MCS）对调用方而言均为同步执行——无需配置流。

典型模式：

import torch
from rdkit import Chem
from nvmolkit.fingerprints import MorganFingerprintGenerator
from nvmolkit.similarity import crossTanimotoSimilarity

stream = torch.cuda.Stream()
fpgen = MorganFingerprintGenerator(radius=2, fpSize=1024)
mols = [Chem.MolFromSmiles(smi) for smi in ["CCO", "c1ccccc1", "CC(=O)O"]]

with torch.cuda.stream(stream):
    fps = fpgen.GetFingerprints(mols, stream=stream)
    sim = crossTanimotoSimilarity(fps, stream=stream)
stream.synchronize()
print(sim.torch())

Device3DResult

在使用 output=CoordinateOutput.DEVICE 调用时，用于 ETKDG 嵌入以及 MMFF/UFF 优化（单次调用和 BatchedForcefield）。这是将构象写回 Mol 对象的 GPU 驻留等效方式。字段：

• values：形状为 (total_atoms, 3) 的 float64 AsyncGpuResult。按 CSR 风格布局拼接的构象坐标。
• atom_starts、mol_indices、conf_indices：描述布局的 int32 AsyncGpuResult 缓冲区（values[atom_starts[i]:atom_starts[i+1]] 是构象 i 的原子）。
• energies、converged：仅在 MMFF/UFF 最小化时填充的 AsyncGpuResult 缓冲区（普通 ETKDG 不会填充）。
• gpu_id：缓冲区所在的设备。每个 API 的 targetGpu 参数用于选择设备；targetGpu=-1 使用默认的整合设备。
• .per_molecule() 返回按分子组织的 list[list[torch.Tensor]]，其中包含各构象视图；.dense(pad_value=nan) 会生成填充后的 (n_mols, max_confs, max_atoms, 3) 张量。

默认模式（CoordinateOutput.RDKIT_CONFORMERS）仍会将优化后的坐标写回每个 Mol，并返回 Python 能量和收敛标志列表。需要串联后续 GPU 工作（例如 ETKDG → MMFF → 相似度评分）且不进行主机往返时，请使用 CoordinateOutput.DEVICE。

MCSBatchResult

findMCS 是同步函数，返回由 CPU NumPy 数组支持的 MCSBatchResult。 结果始终为扁平结构：result[k]（或 result.get_result(k)） 会生成配对位置 k 处的结果，而通常不是分子 k 的结果。使用 result.pairs[k] 确认该配对。在 all_pairs 模式下，这些是针对 mols 生成的配对；在 pairs 模式下，它们会严格保留提供的配对序列；在 paired_lists 模式下，第 k 项会比较 mols[k] 和 mols_b[k]，而 result.pairs[k] 使用合并表中的索引 (k, len(mols) + k)。每个 MCSResult 都包含 pair、num_atoms、num_bonds、canceled、 atom_mapping 和 bond_mapping；每个映射的两列分别索引该结果配对中的第一个和第二个分子。

配置

有关 ETKDG、力场、子结构或 MCS 调优，请阅读 高级配置参考。 其中列出了配置字段、默认值、GPU 选择和自动调优 API。

示例

Morgan 指纹 + 批量 Tanimoto 相似度

import torch
from rdkit import Chem
from nvmolkit.fingerprints import MorganFingerprintGenerator
from nvmolkit.similarity import crossTanimotoSimilarity

smiles = ["CCO", "CCN", "c1ccccc1", "CC(=O)O", "CCOCC"]
mols = [Chem.MolFromSmiles(smi) for smi in smiles]

fpgen = MorganFingerprintGenerator(radius=2, fpSize=1024)
fps = fpgen.GetFingerprints(mols)

sim = crossTanimotoSimilarity(fps)
torch.cuda.synchronize()
print(sim.torch())

输入为 list[Mol]。GetFingerprints 的输出是一个 AsyncGpuResult，包装了形状为 (n_mols, fpSize / 32) 的 int32 打包位张量。可直接将其传入 crossTanimotoSimilarity，生成 (n, n) 相似度矩阵；传入两组指纹则可生成 (n, m) 交叉矩阵。对于大到无法在 GPU 上完整存放的集合，请使用 crossTanimotoSimilarityMemoryConstrained（分块计算，并在 CPU 上返回 numpy）。

ETKDG 构象嵌入

from rdkit.Chem import AddHs, MolFromSmiles
from rdkit.Chem.rdDistGeom import ETKDGv3
from nvmolkit.embedMolecules import EmbedMolecules

mols = [AddHs(MolFromSmiles(smi)) for smi in ["C1CCCCC1", "C1CCCCC2CCCCC12", "COO"]]
params = ETKDGv3()

EmbedMolecules(mols, params, confsPerMolecule=10, maxIterations=-1)

for mol in mols:
    print(mol.GetNumConformers())

输入为已清理且添加了氢原子的 list[Mol]（AddHs）。构象会原地添加；不支持的嵌入选项请参阅限制。

批量构象的 MMFF94 最小化

from rdkit.Chem import AddHs, MolFromSmiles
from rdkit.Chem.rdDistGeom import ETKDGv3
from nvmolkit.embedMolecules import EmbedMolecules
from nvmolkit.mmffOptimization import MMFFOptimizeMoleculesConfs

mols = [AddHs(MolFromSmiles(smi)) for smi in ["CCO", "CCN", "c1ccccc1"]]
params = ETKDGv3()
EmbedMolecules(mols, params, confsPerMolecule=5)

energies = MMFFOptimizeMoleculesConfs(mols, maxIters=500)
for mol, mol_energies in zip(mols, energies):
    print(mol.GetNumConformers(), mol_energies)

输入为已填充构象的 list[Mol]（通常由 ETKDG、RDKit 的 EmbedMultipleConfs 或之前的 nvMolKit 调用生成）。坐标会原地更新；返回值为 list[list[float]]，其中包含优化后的能量，顺序与输入分子及其构象索引对应。UFF 的形状完全相同：将其替换为 from nvmolkit.uffOptimization import UFFOptimizeMoleculesConfs。

默认最小化器为 BFGS。若要使用 FIRE，请传入 minimizerKind="FIRE"； 也可以通过 fireOptions= 传入 nvmolkit.types.FireOptions 实例进行自定义。单次调用的 MMFF 和 UFF 函数，以及两种批量力场的 .minimize() 方法，都接受相同的选择参数。

如果任何输入分子为 None 或缺少 MMFF/UFF 原子类型，调用会引发 ValueError。异常的 args[1] 是一个包含 "none" 和 "no_params" 键的字典，其中列出了有问题的索引——这对筛除有噪声的输入集合很有用。

构象 RMSD 和 Butina 聚类

import torch
from rdkit import Chem
from rdkit.Chem.rdDistGeom import EmbedMultipleConfs
from nvmolkit.clustering import OutputMode, butina
from nvmolkit.conformerRmsd import GetConformerRMSMatrixBatch

mols = [Chem.AddHs(Chem.MolFromSmiles(smi)) for smi in ["CCCCCC", "c1ccccc1"]]
for mol in mols:
    EmbedMultipleConfs(mol, numConfs=10)

# Remove hydrogens after embedding for heavy-atom RMSD.
heavy_mols = [Chem.RemoveHs(mol) for mol in mols]

# Default RMSD output is RDKit-compatible condensed lower-triangle form.
condensed = GetConformerRMSMatrixBatch(heavy_mols)

# Butina expects a square distance matrix, so request square GPU tensors.
square = GetConformerRMSMatrixBatch(heavy_mols, output_format="square")
results = [
    butina(distance_matrix, cutoff=0.5, output=OutputMode.DEVICE)
    for distance_matrix in square
]

torch.cuda.synchronize()
for result in results:
    print(result.cluster_ids.torch().cpu().tolist())

默认情况下，两个 Butina 函数都会返回驻留在 GPU 上的结果：

• 默认值 output=OutputMode.DEVICE 会返回聚类 ID、质心和大小。
• output=OutputMode.RDKIT 会在主机端返回 RDKit 聚类元组。每个聚类的第一个元素是其质心。

设备输出字段是 AsyncGpuResult 对象。使用 .torch() 可访问其 CUDA 张量，而无需复制到主机端；使用 .numpy() 则会同步并将字段复制到主机端。

GetConformerRMSMatrix(mol) 和 GetConformerRMSMatrixBatch(mols) 默认使用 output_format="condensed"，返回封装了 RDKit 风格扁平向量的 AsyncGpuResult 对象，向量长度为 N * (N - 1) // 2。在将结果传入 butina() 或其他需要 N x N 距离矩阵的 API 时，请使用 output_format="square"。两种格式的数据都驻留在 GPU 上；对于压缩结果，请调用 .numpy()；对于方形张量，请先同步再将其移至 CPU。

原子-原子路径相似度和有向球排斥聚类

原子-原子路径（AAP）相似度与有向球排斥（DISE）聚类提供设备输出和 RDKit 风格的输出模式：

from rdkit import Chem
from nvmolkit.clustering import OutputMode, fused_dise

molecules = [Chem.MolFromSmiles(smiles) for smiles in ["CCCC", "CCCO", "CCOC"]]
device_result = fused_dise(molecules, 1 - 0.217, metric="aap")
rdkit_clusters = fused_dise(molecules, 1 - 0.217, metric="aap", output=OutputMode.RDKIT)

device_result 包含 cluster_ids、centroids 和 cluster_sizes 字段；聚类 ID 从 0 开始且连续。OutputMode.RDKIT 描述的是质心置前的 RDKit 聚类表示形式，而不是用 RDKit 实现 AAP+DISE 算法。目前的 DISE 控制循环在返回任一模式前都会进行同步；DEVICE 表示稳定的结果结构及其驻留位置，而不表示整个调用是异步执行的。

最大公共子结构搜索

from rdkit import Chem
from nvmolkit.mcs import findMCS

mols = [Chem.MolFromSmiles(smi) for smi in ["CCO", "CCN", "c1ccccc1O"]]
result = findMCS(mols, mode="pairs", pairs=[(0, 1), (0, 2)])

for pair_idx, pair in enumerate(result.pairs):
    item = result[pair_idx]
    print(pair, item.num_atoms, item.num_bonds, item.atom_mapping)

默认的 mode="all_pairs" 会生成包含对角线的上三角矩阵。mode="pairs" 会原样保留显式的配对列表，包括重复项和反向配对。mode="paired_lists" 会将 mols 与大小相同的 mols_b 逐项配对。超时时间按配对计算；请检查 item.canceled，因为超时结果可能包含已找到的最佳部分 MCS。

匹配选项包括原子比较、键比较、价态/形式电荷匹配，以及仅对环内原子/键进行匹配。不支持的 RDKit fMCS 选项会触发报错，而不是静默改变语义。对于重复执行的典型显式配对工作负载，nvmolkit.autotune.tune_mcs 会返回一个 TuneResult。其 best_config 是经过调优的 MCSConfig，可传入 findMCS(..., config=result.best_config)。

自定义力场选项和约束（BatchedForcefield）

如需了解按分子设置的力场参数、几何约束，或分别调用能量和梯度计算，请阅读高级力场指南。

故障排除

现象	可能原因	解决方法
torch.cuda.is_available() 为 false	缺少 GPU 访问权限、驱动不兼容，或 torch CUDA 构建	检查 GPU 和驱动，然后按照上文的安装指南选择兼容的 torch 构建。
RuntimeError: invalid device ordinal	请求的 gpuIds 条目不可见	使用小于 torch.cuda.device_count() 的设备索引，或使用 API 文档中说明的 GPU 默认值。
某个 RDKit 选项被拒绝	该选项不受此 nvMolKit API 支持	仅当受支持的选项能够保留所需行为时才使用；否则，请使用 RDKit 执行该操作。

深入了解

• 完整功能列表、API 参考和指南：https://nvidia-bionemo.github.io/nvMolKit/
• 各版本变更内容：https://nvidia-bionemo.github.io/nvMolKit/changelog.html
• 示例（Jupyter 笔记本）：GitHub 仓库中的 examples/ 目录