Utilisation de nvMolKit

Objectif

Implémentations GPU-accélérées et par batch d'opérations RDKit courantes. Les APIs reflètent RDKit autant que possible mais sont orientées batch : elles prennent des listes de rdkit.Chem.Mol (ou des listes d'empreintes) et les traitent en parallèle sur un ou plusieurs GPUs. nvMolKit se lie à RDKit au moment de la compilation ; les entrées et sorties sont de vrais objets RDKit Mol.

Cette skill couvre l'API Python installée. La compilation de nvMolKit depuis les sources dépasse le cadre de ce document.

Où nvMolKit excelle

Utilisez nvMolKit quand :

• La charge de travail est un grand batch de molécules traitées ensemble (généralement des milliers ou plus).
• La métrique est débit / temps total écoulé sur le batch, pas la latence par molécule.
• La même opération est répétée à l'identique sur le batch (empreintage d'une bibliothèque, intégration/minimisation de nombreux conformères, similarité par paire en bulk), pour que le GPU reste saturé.

Prérequis

• Un GPU NVIDIA avec capacité de calcul 7.0 (V100) ou supérieure
• Un driver CUDA compatible avec CUDA 12.6+.
• Une installation torch fonctionnelle avec support CUDA (nvMolKit retourne des tenseurs GPU via l'interface de tableau CUDA de torch).

Pour l'aide à l'installation, faites en sorte que l'utilisateur choisisse un backend PyTorch CUDA que le driver hôte supporte avant d'installer nvMolKit. Les wheels PyPI de nvMolKit sont compilées avec CUDA Toolkit 12.9 et dépendent des paquets runtime CUDA 12, mais pip/uv peuvent quand même sélectionner une wheel PyTorch CUDA 13 à moins que la commande d'installation ne le précise.

• Conda : préférez pytorch-gpu sur conda-forge ; épinglez cuda-version=12.6 ou une autre version CUDA supportée par le driver.
• pip : renvoyez l'utilisateur au sélecteur d'installation PyTorch ou à la page des versions antérieures pour installer torch avec un backend CUDA 12.x avant d'installer nvMolKit.
• uv : installez nvMolKit avec un backend explicite, par exemple uv pip install --torch-backend=cu128 nvmolkit.

Entrées

• Obligatoire : choisissez une opération et fournissez des molécules ou des empreintes depuis le code ou le dataset moléculaire de l'utilisateur. Analysez les SMILES avec RDKit et rejetez les analyses échouées (None).
• Les opérations moléculaires utilisent des objets RDKit Mol. Ajoutez des hydrogènes pour ETKDG ; la minimisation et les comparaisons de conformères nécessitent des conformères existants.
• La similarité d'empreinte accepte des résultats AsyncGpuResult empaquetés, des tenseurs torch, ou des arrays NumPy : une molécule par ligne, avec des mots int32 ou uint32.
• Optionnel : prenez les comptes de conformères, les paramètres d'empreinte, les seuils, les modes de sortie et les options matérielles selon le flux de travail demandé par l'utilisateur ; sinon utilisez les défauts de l'API documentée.

Limitations

• CUDA est obligatoire ; il n'y a pas de fallback CPU. Utilisez RDKit directement quand l'exécution CPU est nécessaire.
• RDKit pur est généralement préférable pour le travail sur une seule molécule ou les opérations qui ne peuvent pas être batchées.
• ETKDG ne supporte pas les matrices de limites personnalisées, les CPCI personnalisés, les cartes de coordonnées, ou l'intégration de fragments séparés.
• La recherche de sous-structure ne supporte pas les appariements conscients de la chiralité, la stéréochimie améliorée, ou d'autres options avancées de RDKit SubstructMatchParameters.

Instructions

1. Exécutez le test de fumée ci-dessous avant d'écrire du code nvMolKit.
2. Choisissez une API dans le tableau des points d'entrée et appliquez ses exigences d'entrée.
3. Gérez son résultat comme décrit ci-dessous ; synchronisez les résultats GPU asynchrones avant la lecture hôte.

Vérifiez l'installation avant d'écrire du vrai code

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)
# Attendu : shape (3, 32), dtype torch.int32  (1024 bits empaquetés dans 32 int32s par ligne)

Si cela échoue, renvoyez l'utilisateur au guide d'installation plutôt que de deviner.

Points d'entrée

Types de résultats et modèle d'exécution

Deux formes de retour portent la sortie résidente GPU, selon ce que l'opération produit.

AsyncGpuResult

Utilisée par les opérations qui retournent un tenseur plat unique (empreintes, matrices de similarité, vecteurs RMSD/TFD, entrées Butina). Comportements clés :

• Asynchrone. Le kernel peut ne pas avoir complété quand l'appel revient.
• result.torch() retourne un torch.Tensor sur le GPU sans copie. L'appelant est responsable de la synchronisation avant de lire les valeurs sur l'hôte.
• result.numpy() synchronise et retourne un array numpy CPU.
• Expose __cuda_array_interface__, donc il peut être passé directement dans d'autres fonctions nvMolKit (ex. empreintes → similarité) sans aller-retour hôte.

Contrôle du flux CUDA

Un sous-ensemble des APIs retournant AsyncGpuResult accepte un argument optionnel stream: torch.cuda.Stream | None = None pour que les appelants soumettent le travail nvMolKit à un flux non-défaut et le chevauchent avec leurs propres kernels. Quand omis, l'appel utilise le flux torch courant.

APIs qui prennent un argument stream :

• MorganFingerprintGenerator.GetFingerprints
• crossTanimotoSimilarity, crossCosineSimilarity, et leurs variants *MemoryConstrained
• butina, fused_butina
• GetConformerRMSMatrix, GetConformerRMSMatrixBatch

Les autres APIs (ETKDG, optimisation MMFF/UFF, TFD, recherche de sous-structure, MCS) sont synchrones à l'appelant — pas besoin de tuyautage de flux.

Motif typique :

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

Utilisée par l'intégration ETKDG et l'optimisation MMFF/UFF (one-shot et BatchedForcefield) quand appelée avec output=CoordinateOutput.DEVICE. L'équivalent résident GPU d'écrire des conformères dans les objets Mol. Champs :

• values : AsyncGpuResult de shape (total_atoms, 3) float64. Coordonnées de conformère concaténées en disposition CSR.
• atom_starts, mol_indices, conf_indices : buffers AsyncGpuResult int32 décrivant la disposition (values[atom_starts[i]:atom_starts[i+1]] sont les atomes du conformère i).
• energies, converged : buffers AsyncGpuResult remplis seulement pour la minimisation MMFF/UFF (pas pour le plain ETKDG).
• gpu_id : device sur lequel les buffers résident. L'argument targetGpu sur chaque API choisit celui-ci ; targetGpu=-1 utilise le device de consolidation par défaut.
• .per_molecule() retourne des vues imbriquées list[list[torch.Tensor]] par conformère ; .dense(pad_value=nan) matérialise un tenseur rembourré (n_mols, max_confs, max_atoms, 3).

Le mode par défaut (CoordinateOutput.RDKIT_CONFORMERS) écrit toujours les coordonnées optimisées dans chaque Mol et retourne des listes Python d'énergies/drapeaux de convergence. Utilisez CoordinateOutput.DEVICE quand vous chaînez du travail GPU en aval (ex. ETKDG → MMFF → scoring de similarité) sans aller-retours hôte.

MCSBatchResult

findMCS est synchrone et retourne un MCSBatchResult soutenu par des arrays NumPy CPU. Les résultats sont toujours plats : result[k] (ou result.get_result(k)) matérialise le résultat à la position de paire k, pas généralement le résultat pour la molécule k. Utilisez result.pairs[k] pour identifier cette paire. En mode all_pairs ce sont les paires générées sur mols ; en mode pairs elles préservent exactement la séquence de paires fournie ; en mode paired_lists l'élément k compare mols[k] avec mols_b[k], tandis que result.pairs[k] utilise les indices de table combinée (k, len(mols) + k). Chaque MCSResult a pair, num_atoms, num_bonds, canceled, atom_mapping, et bond_mapping ; les deux colonnes de chaque mapping indexent la première et la deuxième molécule de cette paire de résultat, respectivement.

Configuration

Pour le tuning ETKDG, champ de force, sous-structure, ou MCS, lisez la référence avancée de configuration. Elle liste les champs de configuration, les défauts, la sélection GPU, et les APIs d'autotuning.

Exemples

Empreintes Morgan + similarité Tanimoto en bulk

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())

Les entrées sont list[Mol]. La sortie de GetFingerprints est un AsyncGpuResult enveloppant un tenseur int32 (n_mols, fpSize / 32) de bits empaquetés. Passez-le directement dans crossTanimotoSimilarity pour une matrice de similarité (n, n) ; passez deux ensembles d'empreintes pour une matrice croisée (n, m). Pour les ensembles trop volumineux pour être matérialisés sur le GPU, utilisez crossTanimotoSimilarityMemoryConstrained (calcul en chunks, retourne numpy sur CPU).

Intégration de conformères 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())

Les entrées sont des list[Mol] sanitisées avec hydrogènes ajoutés (AddHs). Les conformères sont ajoutés en place ; voir Limitations pour les options d'intégration non supportées.

Minimisation MMFF94 d'un batch de conformères

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)

Les entrées sont des list[Mol] avec conformères déjà remplis (généralement par ETKDG, la EmbedMultipleConfs de RDKit, ou un appel nvMolKit antérieur). Les coordonnées sont mises à jour en place ; le retour est list[list[float]] des énergies optimisées alignées avec l'ordre des molécules d'entrée et l'index du conformère. UFF est identique en forme : remplacez par from nvmolkit.uffOptimization import UFFOptimizeMoleculesConfs.

BFGS est le minimiseur par défaut. Pour utiliser FIRE, passez minimizerKind="FIRE" ; optionnellement le personnalisez avec une instance nvmolkit.types.FireOptions via fireOptions=. Les fonctions MMFF et UFF one-shot et les deux méthodes .minimize() de batched-forcefield acceptent le même sélecteur.

Si une molécule d'entrée est None ou manque de types d'atomes MMFF/UFF, l'appel lève ValueError. Le args[1] de l'exception est un dict avec clés "none" et "no_params" listant les indices offensants - utile pour filtrer un ensemble d'entrée bruyant.

RMSD de conformères et clustering Butina

import torch
from rdkit import Chem
from rdkit.Chem.rdDistGeom import EmbedMultipleConfs
from nvmolkit.clustering import ButinaOutputMode, 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)

# Retirez les hydrogènes après intégration pour un RMSD sur les atomes lourds.
heavy_mols = [Chem.RemoveHs(mol) for mol in mols]

# La sortie RMSD par défaut est sous forme condensée du triangle inférieur compatible RDKit.
condensed = GetConformerRMSMatrixBatch(heavy_mols)

# Butina attend une matrice distance carrée, donc demandez des tenseurs GPU carrés.
square = GetConformerRMSMatrixBatch(heavy_mols, output_format="square")
results = [
    butina(distance_matrix, cutoff=0.5, output=ButinaOutputMode.DEVICE)
    for distance_matrix in square
]

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

Les deux fonctions Butina retournent des résultats résidents GPU par défaut :

• Par défaut, output=ButinaOutputMode.DEVICE, retourne les IDs de cluster, centroïdes et tailles.
• output=ButinaOutputMode.RDKIT retourne les tuples de cluster RDKit sur l'hôte. Le premier élément de chaque cluster est son centroïde.

Les champs de sortie device sont des objets AsyncGpuResult. Utilisez .torch() pour accéder à leurs tenseurs CUDA sans copie hôte ou .numpy() pour synchroniser et copier un champ vers l'hôte.

GetConformerRMSMatrix(mol) et GetConformerRMSMatrixBatch(mols) défaillent à output_format="condensed", retournant des objets AsyncGpuResult qui enveloppent des vecteurs plats RDKit-style de longueur N * (N - 1) // 2. Utilisez output_format="square" quand vous chaînez dans butina() ou n'importe quelle autre API qui attend une matrice distance N x N. Les deux formes vivent sur le GPU ; appelez .numpy() sur les résultats condensés ou synchronisez avant de déplacer des tenseurs carrés vers le CPU.

Similarité de chemin Atome-Atome et clustering exclusion sphère dirigée

La similarité de chemin Atome-Atome (AAP) avec clustering exclusion sphère dirigée (DISE) fournit les modes de sortie device et RDKit :

from rdkit import Chem
from nvmolkit.clustering import DISEOutputMode, aap_dise

molecules = [Chem.MolFromSmiles(smiles) for smiles in ["CCCC", "CCCO", "CCOC"]]
device_result = aap_dise(molecules)
rdkit_clusters = aap_dise(molecules, output=DISEOutputMode.RDKIT)

device_result a les champs cluster_ids, centroids, et cluster_sizes ; les IDs de cluster sont zéro-indexés et contigus. DISEOutputMode.RDKIT décrit la représentation de cluster RDKit centroïde-d'abord, pas une implémentation RDKit de l'algorithme AAP+DISE. La boucle de contrôle DISE courante synchronise avant de retourner l'un ou l'autre mode ; DEVICE décrit le schéma stable et où le résultat réside, pas l'exécution asynchrone de l'appel global.

Recherche de sous-structure commune maximale

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)

Le défaut mode="all_pairs" génère le triangle supérieur incluant la diagonale. mode="pairs" préserve une liste de paires explicite exactement, incluant les doublons et les paires inversées. mode="paired_lists" ferme mols avec un mols_b de même taille. Les timeouts sont par paire ; inspectez item.canceled car un résultat timeouté peut contenir le meilleur MCS partiel trouvé.

Les options d'appariement incluent atom_compare, bond_compare, correspondance de valence/charge formelle, et correspondance anneau-seul atome/liaison. Les options fMCS RDKit non supportées lèvent plutôt que de changer silencieusement la sémantique. Pour les charges de travail répétées de paires explicites représentatives, nvmolkit.autotune.tune_mcs retourne un TuneResult. Son best_config est le MCSConfig tunné à passer à findMCS(..., config=result.best_config).

Options de champ de force personnalisées + contraintes (BatchedForcefield)

Pour les paramètres de champ de force par molécule, les contraintes géométriques, ou les appels séparés d'énergie et gradient, lisez la recette avancée de champ de force.

Dépannage

Aller plus loin

• Liste complète de fonctionnalités, référence API et guides : https://nvidia-bionemo.github.io/nvMolKit/
• Ce qui a changé dans chaque release : https://nvidia-bionemo.github.io/nvMolKit/changelog.html
• Exemples travaillés (notebooks Jupyter) : le répertoire examples/ dans le dépôt GitHub