nvMolKit

Operations that return GPU-resident data (like fingerprinting or similarity) return an AsyncGpuResult object. This allows you to queue multiple GPU operations without waiting for each to complete, enabling efficient pipeline construction.

Accessing Results: To retrieve the data, you must synchronize the GPU and then convert the result to your desired format:

• .torch(): Returns a zero-copy torch Tensor (stays on GPU).
• .numpy(): Returns a numpy array (moves data to CPU; involves synchronization).

Streams: Most operations accept an optional stream parameter (a torch.cuda.Stream) to control execution.

Synchronization: Use torch.cuda.synchronize() to ensure all queued operations are complete before accessing results.

from nvmolkit.fingerprints import MorganFingerprintGenerator
import torch

fpgen = MorganFingerprintGenerator(radius=2, fpSize=1024)
# Returns AsyncGpuResult
result = fpgen.GetFingerprints(mols)

# Accessing results
torch.cuda.synchronize()
fps_torch = result.torch()  # Zero-copy on GPU
fps_numpy = result.numpy()  # Moves to CPU

The AsyncGpuResult type is used by operations that return a single flat tensor (e.g., fingerprints, similarity matrices, RMSD/TFD vectors, Butina inputs).

Key Behaviors:

• Asynchronous: The kernel may not have completed when the call returns. You must synchronize before reading values on the host.
• Zero-copy Torch: result.torch() returns a zero-copy torch.Tensor on the GPU.
• Numpy Conversion: result.numpy() synchronizes and returns a CPU numpy array.
• Interoperability: It exposes __cuda_array_interface__, allowing it to be passed directly into other nvMolKit functions (e.g., fingerprints $\rightarrow$ similarity) without a host round-trip.

CUDA Stream Control: A subset of APIs (fingerprints, similarity, Butina, RMSD) accept an optional stream: torch.cuda.Stream | None = None argument. This allows you to overlap nvMolKit work with your own kernels using non-default streams.

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

Use MMFFBatchedForcefield when you need per-molecule settings or atom-specific constraints. Unlike the top-level functions, this class allows you to pass sequences (lists) for scalar arguments to apply different settings to each molecule in the batch.

Supported per-molecule arguments:

• properties: RDKit MMFFMolProperties objects (or None for defaults).
• nonBondedThreshold: Non-bonded cutoff distance.
• ignoreInterfragInteractions: Boolean to omit non-bonded terms between fragments.

Constraints: Constraints are added via a per-molecule view ff[i]. Constraints applied to molecule i affect all its conformers. Use the relative=True flag on distance, angle, or torsion constraints to use offsets from the current geometry instead of absolute values.

from rdkit.Chem import AllChem
from nvmolkit.batchedForcefield import MMFFBatchedForcefield

# Setup per-molecule properties
props_b = AllChem.MMFFGetMoleculeProperties(mols[1])
props_c = AllChem.MMFFGetMoleculeProperties(mols[2])
props_c.SetMMFFVariant("MMFF94s")
props_c.SetMMFFDielectricConstant(4.0)
props_c.SetMMFFDielectricModel(2)  # 1 = constant, 2 = distance-dependent

# Initialize batched forcefield with per-molecule lists
ff = MMFFBatchedForcefield(
    mols,
    properties=[None, props_b, props_c],
    nonBondedThreshold=[100.0, 25.0, 25.0],
    ignoreInterfragInteractions=[True, False, True],
)

# Add constraints to specific molecules
ff[0].add_distance_constraint(idx1=0, idx2=2, relative=True, min_len=-0.2, max_len=0.5, force_constant=20.0)
ff[0].add_position_constraint(idx=0, max_displ=0.1, force_constant=50.0)
ff[1].add_angle_constraint(idx1=0, idx2=1, idx3=2, relative=False, min_angle_deg=100.0, max_angle_deg=120.0, force_constant=25.0)
ff[2].add_torsion_constraint(idx1=0, idx2=1, idx3=2, idx4=3, relative=False, min_dihedral_deg=-10.0, max_dihedral_deg=10.0, force_constant=8.0)

# Run minimization
energies, converged = ff.minimize(maxIters=200, forceTol=1e-4)

nvMolKit provides GPU-accelerated Tanimoto and cosine similarity computations.

Core Behavior:

• Each similarity call returns an n × m matrix representing the similarity between two batches of fingerprints.
• If only one batch is provided, the function performs an all-to-all similarity computation, returning an n × n matrix.
• The results can be accessed as a PyTorch tensor by calling the .torch() method on the returned AsyncGpuResult object.

Input Requirements:

• Data Type: Inputs must be int32 or uint32.
• Shape: Fingerprints must be packed into 32-bit integers with shape (num_fingerprints, fp_size/32). Each row represents a fingerprint, and each column represents a 32-bit block.
• Consistency: If two batches are provided, they must both have the same fp_size.

# Example of accessing the result as a torch tensor
sims = crossTanimotoSimilarity(fps).torch()

The Device3DResult type is used by ETKDG embedding and MMFF/UFF optimization (both one-shot and BatchedForcefield) when called with output=CoordinateOutput.DEVICE. This is the GPU-resident equivalent of writing conformers back to Mol objects, ideal for chaining downstream GPU work without host round-trips.

Fields:

• values: AsyncGpuResult of shape (total_atoms, 3) float64. Concatenated conformer coordinates in CSR-style layout.
• atom_starts, mol_indices, conf_indices: AsyncGpuResult int32 buffers describing the layout (values[atom_starts[i]:atom_starts[i+1]] is conformer i's atoms).
• energies, converged: AsyncGpuResult buffers (populated only for MMFF/UFF minimization).
• gpu_id: The device the buffers live on.

Methods:

• .per_molecule(): Returns nested list[list[torch.Tensor]] of per-conformer views.
• .dense(pad_value=nan): Materializes a padded (n_mols, max_confs, max_atoms, 3) tensor.