Skip to content

Add UFF forcefield C++ implementation#123

Merged
scal444 merged 4 commits intoNVIDIA-Digital-Bio:mainfrom
scal444:pr5_uff_cpp
Apr 2, 2026
Merged

Add UFF forcefield C++ implementation#123
scal444 merged 4 commits intoNVIDIA-Digital-Bio:mainfrom
scal444:pr5_uff_cpp

Conversation

@scal444
Copy link
Copy Markdown
Collaborator

@scal444 scal444 commented Mar 31, 2026

Adds UFF bond stretch, angle bend, torsion, inversion, and van der Waals terms with batched GPU evaluation. Includes RDKit UFF parameter flattening, batched forcefield wrapper, and C++ tests.

scal444 added 2 commits March 31, 2026 11:35
Cherry-picked from pr5_uff (e7477a9). Adds UFF bond stretch, angle bend,
torsion, inversion, and van der Waals terms with batched GPU evaluation.
Includes RDKit UFF parameter flattening, batched forcefield wrapper, and
C++ tests.

Made-with: Cursor
@scal444 scal444 requested a review from evasnow1992 March 31, 2026 15:49
@greptile-apps
Copy link
Copy Markdown
Contributor

greptile-apps bot commented Mar 31, 2026

Greptile Summary

This PR adds a complete GPU-accelerated UFF forcefield implementation covering bond stretch, angle bend, torsion, inversion, and van der Waals terms, along with a batched forcefield wrapper and RDKit parameter flattener. The energy and most gradient kernels are mathematically sound and well-tested against RDKit reference values.

Key findings:

  • Inversion gradient sign error (uff_kernels_device.cuh): uffInversionGrad constructs the JI–JK plane normal via crossProduct(-rJIx, ..., rJKx, ...) (negated rJI), while uffCalculateCosY (used in the energy) uses positive rJI. After normalisation the gradient's cosY equals -sinW rather than +sinW. The dE_dW formula does not account for this sign flip, producing K·sinW·(C1 − 4C2·cosW) instead of the correct K·(C1 − 4C2·sinW). The error is proportional to the out-of-plane displacement, so it is small for typical near-planar sp2 geometries (which is why the current gradient tests likely pass within the 2×10⁻⁴ tolerance), but can be a factor-of-2 or larger error for significantly non-planar centres.
  • Three pre-existing issues acknowledged in comments: zero-distance bond gradient fallback, VdW gradient/energy inconsistency at dist ≤ 0, and the UFF→MMFF header coupling — all tracked with TODOs.
  • Energy kernels, torsion and angle gradient derivations, VdW LJ gradient formula, and Chebyshev polynomial implementations are all correct.
  • The block-per-molecule shared-memory gradient kernel correctly caps at 256 atoms and falls back to global memory for larger systems.

Confidence Score: 4/5

  • Safe to merge with caution — one P1 inversion gradient bug will produce incorrect gradients for non-planar sp2 centres, affecting minimisation quality on such molecules.
  • The bulk of the implementation (energy terms, torsion/angle/VdW gradients, batching infrastructure, build system) is correct and well-tested. However, the inversion gradient sign error is a definite present defect: the plane normal direction used in the gradient computation is inverted relative to the energy function, yielding a systematically wrong derivative for non-planar inversion centres. The current tests are unlikely to catch this because they compare at near-planar equilibrium geometries where the inversion contribution is inherently small. This warrants a 4/5 rather than 5/5.
  • src/forcefields/uff_kernels_device.cuh — specifically uffInversionGrad around the crossProduct call and the dE_dW formula.

Important Files Changed

Filename Overview
src/forcefields/uff_kernels_device.cuh Core device-side kernel implementations for all UFF energy/gradient terms; contains a P1 inversion gradient bug where the plane normal used in uffInversionGrad is negated relative to uffCalculateCosY, causing incorrect gradient magnitudes for non-planar sp2 centres.
rdkit_extensions/uff_flattened_builder.cpp Flattens RDKit UFF parameters into SoA host buffers; logic closely mirrors RDKit's UFF builder. The unusual neighborSlot increment pattern in addInversions (skipping slot 1 for the central atom) is correct but hard to follow.
src/forcefields/uff.cu Host-side batch management: stream assignment, host-to-device transfers, addMoleculeToBatch, and kernel dispatchers for energy and gradient. Implementation is consistent and handles all term types including constraint terms borrowed from MMFF.
src/forcefields/uff_kernels.cu Kernel launchers and two block-per-molecule kernels (combined energy + gradient). The blockSizePerMol = 128 shared-memory gradient path correctly limits to 256 atoms; fallback to global memory is correct.
tests/test_uff.cu Validates GPU UFF against RDKit reference for energy, gradient, batched, and minimizer paths. Tests cover edge cases (single-atom SMILES, small molecules) and a 25-molecule validation set. Gradient tolerance of 2e-4 may be too loose to catch the inversion gradient bug on near-planar test molecules.

Reviews (2): Last reviewed commit: "Apply copyright script" | Re-trigger Greptile

@scal444
Copy link
Copy Markdown
Collaborator Author

scal444 commented Mar 31, 2026

Partial implementation for #114

has not yet been optimized, we'll go through the kernels but lower priority.

Copy link
Copy Markdown
Collaborator

@evasnow1992 evasnow1992 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changes look good to me. Understand that further optimizations may be applied later. Also a note that python APIs haven't been added yet and document files like docs/index.rst may need updates to include UFF alongside MMFF.

@scal444 scal444 merged commit a6d6195 into NVIDIA-Digital-Bio:main Apr 2, 2026
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants