GetConformerRMSMatrixBatch#

nvmolkit.conformerRmsd.GetConformerRMSMatrixBatch(
mols: list[Mol],
prealigned: bool = False,
stream: Stream | None = None,
*,
output_format: str = 'condensed',
) → list[AsyncGpuResult] | list[Tensor]#

Compute pairwise RMSD matrices for a batch of molecules on GPU.

All molecules are processed in a single kernel launch, so their conformer pairs execute concurrently. This improves GPU saturation over repeated single-molecule calls, particularly for molecules with few conformers.

Parameters:

  • mols – List of RDKit molecules, each with zero or more conformers. Strip hydrogens first (Chem.RemoveHs) for heavy-atom RMSD. Molecules with fewer than 2 conformers return an empty result.

  • prealigned – If True, skip Kabsch alignment and compute RMSD on raw coordinates. If False (default), optimally align each pair.

  • stream – CUDA stream to use. If None, uses the current stream.

  • output_format – "condensed" returns one AsyncGpuResult per molecule, each wrapping a 1-D RDKit-style lower-triangle vector. "square" returns one symmetric CUDA torch.Tensor of shape (N, N) per molecule.

Returns:

With output_format="condensed", a list of AsyncGpuResult, one per input molecule, in the same order as mols. Each element wraps a 1-D tensor of shape (N*(N-1)/2,) for a molecule with N conformers, or shape (0,) for molecules with fewer than 2 conformers. The RMSD for conformer pair (i, j) with i > j is at index i*(i-1)//2 + j.

With output_format="square", a list of symmetric CUDA torch.Tensor objects with shapes (N, N).

Raises:

  • ValueError – If any element of mols is None.

  • TypeError – If stream is not a torch.cuda.Stream or None.

Example

>>> from rdkit import Chem
>>> from rdkit.Chem import rdDistGeom
>>> from nvmolkit.conformerRmsd import GetConformerRMSMatrixBatch
>>>
>>> mols = [Chem.RemoveHs(Chem.AddHs(Chem.MolFromSmiles(s)))
...         for s in ["CCCCCC", "c1ccccc1"]]
>>> for mol in mols:
...     rdDistGeom.EmbedMultipleConfs(mol, numConfs=20)
>>>
>>> results = GetConformerRMSMatrixBatch(mols)
>>> # results[0] is AsyncGpuResult for mols[0], etc.