Add all-in-one ETKDG force field kernels

rdkit_extensions/CMakeLists.txt

@@ -44,8 +44,8 @@ target_link_libraries(mmff_optimize INTERFACE ${RDKit_LIBS} CUDA::cudart
 add_library(rdkit_dist_geom_flattened dist_geom_flattened_builder.cpp)
 target_link_libraries(
   rdkit_dist_geom_flattened
-  PRIVATE ${RDKit_LIBS} device_vector nvmolkit_versions
-  PUBLIC dist_geom)
+  PRIVATE ${RDKit_LIBS} device_vector
+  PUBLIC dist_geom nvmolkit_versions)
 
 add_library(rdkit_bounds_matrix bounds_matrix.cpp)
 target_link_libraries(rdkit_bounds_matrix PRIVATE ${RDKit_LIBS})

rdkit_extensions/dist_geom_flattened_builder.cpp

@@ -6,6 +6,7 @@
 #include <RDGeneral/Invariant.h>
 
 #include <boost/dynamic_bitset.hpp>
+#include <cmath>
 #include <map>
 #include <sstream>
 
@@ -492,6 +493,89 @@ nvMolKit::DistGeom::EnergyForceContribsHost constructForceFieldContribs(
   return contribs;
 }
 
+#if RDKIT_AIO_ETKDG_API
+nvMolKit::DistGeom::AllInOneForceContribsHost constructAllInOneForceFieldContribs(
+  const int                                         dim,
+  const ::DistGeom::BoundsMatrix&                   mmat,
+  const ::DistGeom::VECT_CHIRALSET&                 csets,
+  const ::ForceFields::CrystalFF::CrystalFFDetails& etkdgDetails,
+  const bool                                        useExpTorsionAnglePrefs,
+  const bool                                        useBasicKnowledge) {
+  PRECONDITION(etkdgDetails.distMat, "all-in-one ETKDG requires a topological distance matrix");
+  AllInOneForceContribsHost contribs;
+  const unsigned int        numAtoms    = mmat.numRows();
+  const auto&               forceConsts = etkdgDetails.forceConsts;
+
+  // DistGeom::addAngleTerms
+  for (const auto& angle : etkdgDetails.angles) {
+    if (angle[3]) {
+      addAngleConstraintContrib(contribs.linearAngleTerms,
+                                angle[0],
+                                angle[1],
+                                angle[2],
+                                TRIPLE_BOND_MIN_ANGLE,
+                                TRIPLE_BOND_MAX_ANGLE);
+    }
+  }
+
+  // DistGeom::addChiralityTerms
+  if (dim == 4) {
+    addChiralViolationContribs(contribs.distanceGeometry, numAtoms, csets, forceConsts.chiral);
+    addFourthDimContribs(contribs.distanceGeometry, dim, numAtoms, forceConsts.fourthDim);
+  }
+
+  // DistGeom::addDistanceTerms: harmonic 1-2 and 1-3 terms, distance violation terms for every other pair
+  for (unsigned int i = 1; i < numAtoms; ++i) {
+    for (unsigned int j = 0; j < i; ++j) {
+      const double lowerBound          = mmat.getLowerBound(i, j);
+      const double upperBound          = mmat.getUpperBound(i, j);
+      const double topologicalDistance = etkdgDetails.distMat[i * numAtoms + j];
+      if (std::abs(topologicalDistance - 1.0) < 1e-4 || std::abs(topologicalDistance - 2.0) < 1e-4) {
+        addDistanceConstraintContrib(contribs.harmonicDistanceTerms, i, j, lowerBound, upperBound, 10.0);
+        continue;
+      }
+      auto& distTerms = contribs.distanceGeometry.distTerms;
+      distTerms.idx1.push_back(static_cast<int>(i));
+      distTerms.idx2.push_back(static_cast<int>(j));
+      distTerms.lb2.push_back(lowerBound * lowerBound);
+      distTerms.ub2.push_back(upperBound * upperBound);
+      distTerms.weight.push_back(forceConsts.distance * etkdgDetails.boundsMatForceScaling);
+    }
+  }
+
+  // DistGeom::addTorsionTerms
+  if (useBasicKnowledge) {
+    constexpr int improperIndices[3][3] = {
+      {0, 2, 3},
+      {0, 3, 2},
+      {2, 3, 0}
+    };
+    for (const auto& improperAtom : etkdgDetails.improperAtoms) {
+      for (const auto& permutation : improperIndices) {
+        contribs.planarityTerms.idx1.push_back(improperAtom[permutation[0]]);
+        contribs.planarityTerms.idx2.push_back(improperAtom[1]);
+        contribs.planarityTerms.idx3.push_back(improperAtom[permutation[1]]);
+        contribs.planarityTerms.idx4.push_back(improperAtom[permutation[2]]);
+      }
+    }
+  }
+  if (useBasicKnowledge || useExpTorsionAnglePrefs) {
+    boost::dynamic_bitset<>   atomPairs(numAtoms * numAtoms);
+    Energy3DForceContribsHost torsionContribs;
+    addExperimentalTorsionTerms(torsionContribs, etkdgDetails, numAtoms, atomPairs);
+    contribs.experimentalTorsionTerms = std::move(torsionContribs.experimentalTorsionTerms);
+  }
+
+  return contribs;
+}
+
+nvMolKit::DistGeom::AllInOneForceWeights allInOneForceWeights(
+  const ::ForceFields::CrystalFF::CrystalFFDetails& etkdgDetails) {
+  const auto& forceConsts = etkdgDetails.forceConsts;
+  return {forceConsts.chiral, forceConsts.fourthDim, forceConsts.kTermAngle, forceConsts.kTermImproper};
+}
+#endif
+
 nvMolKit::DistGeom::Energy3DForceContribsHost construct3DForceFieldContribs(
   const ::DistGeom::BoundsMatrix&                   mmat,
   const ::ForceFields::CrystalFF::CrystalFFDetails& etkdgDetails,

rdkit_extensions/dist_geom_flattened_builder.h

@@ -8,6 +8,7 @@
 #include <map>
 
 #include "src/forcefields/dist_geom.h"
+#include "versions.h"
 namespace nvMolKit {
 namespace DistGeom {
 
@@ -27,6 +28,24 @@ nvMolKit::DistGeom::Energy3DForceContribsHost construct3DForceFieldContribs(
   int                                               dim,
   bool                                              useBasicKnowledge = true);
 
+#if RDKIT_AIO_ETKDG_API
+//! Builds one molecule's terms of RDKit's all-in-one ETKDG force field, as DistGeom::constructAllInOneForceField
+//! followed by DistGeom::addTorsionTerms. Force constants come from etkdgDetails.forceConsts, the 1-2/1-3 pairs from
+//! etkdgDetails.distMat, and the experimental torsion constants in etkdgDetails must already carry RDKit's all-in-one
+//! scaling.
+nvMolKit::DistGeom::AllInOneForceContribsHost constructAllInOneForceFieldContribs(
+  int                                               dim,
+  const ::DistGeom::BoundsMatrix&                   mmat,
+  const ::DistGeom::VECT_CHIRALSET&                 csets,
+  const ::ForceFields::CrystalFF::CrystalFFDetails& etkdgDetails,
+  bool                                              useExpTorsionAnglePrefs,
+  bool                                              useBasicKnowledge);
+
+//! The batch-wide force constants that go with constructAllInOneForceFieldContribs for the same etkdgDetails.
+nvMolKit::DistGeom::AllInOneForceWeights allInOneForceWeights(
+  const ::ForceFields::CrystalFF::CrystalFFDetails& etkdgDetails);
+#endif
+
 }  // namespace DistGeom
 }  // namespace nvMolKit
 

src/forcefields/CMakeLists.txt

@@ -67,3 +67,9 @@ target_link_libraries(
   etk_batched_forcefield
   PUBLIC batched_forcefield dist_geom
   PRIVATE device_vector)
+
+add_library(aio_etkdg_batched_forcefield aio_etkdg_batched_forcefield.cu)
+target_link_libraries(
+  aio_etkdg_batched_forcefield
+  PUBLIC batched_forcefield dist_geom
+  PRIVATE device_vector)

src/forcefields/aio_etkdg_batched_forcefield.cu

@@ -0,0 +1,133 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+#include "src/forcefields/aio_etkdg_batched_forcefield.h"
+#include "src/forcefields/dist_geom_kernels.h"
+#include "src/utils/device_convert.cuh"
+
+namespace nvMolKit {
+
+AllInOneETKDGBatchedForcefield::AllInOneETKDGBatchedForcefield(const DistGeom::BatchedAllInOneSystemHost& systemHost,
+                                                               const std::vector<int>&               atomStartsHost,
+                                                               const DistGeom::AllInOneForceWeights& weights,
+                                                               BatchedForcefieldMetadata             metadata,
+                                                               const cudaStream_t                    stream,
+                                                               const PrecisionMode                   precision)
+    : BatchedForcefield(ForceFieldType::AIO_ETKDG, 4, atomStartsHost, nullptr, std::move(metadata)),
+      weights_(weights) {
+  atomStartsDevice_.setStream(stream);
+  if (usesSinglePrecision(precision)) {
+    DistGeom::sendAllInOneSystemToDevice(systemHost,
+                                         systemDevice_.emplace<DistGeom::BatchedAllInOneDeviceBuffersSingle>(),
+                                         stream);
+  } else {
+    DistGeom::sendAllInOneSystemToDevice(systemHost,
+                                         systemDevice_.emplace<DistGeom::BatchedAllInOneDeviceBuffers>(),
+                                         stream);
+  }
+  atomStartsDevice_.setFromVector(atomStartsHost);
+  setAtomStartsDevice(atomStartsDevice_.data());
+}
+
+cudaError_t AllInOneETKDGBatchedForcefield::computeEnergy(double*        energyOuts,
+                                                          const double*  positions,
+                                                          const uint8_t* activeSystemMask,
+                                                          cudaStream_t   stream) {
+  if (std::holds_alternative<DistGeom::BatchedAllInOneDeviceBuffersSingle>(systemDevice_)) {
+    singleConversion_.positions.setStream(stream);
+    singleConversion_.energies.setStream(stream);
+    singleConversion_.positions.resize(totalPositions());
+    singleConversion_.energies.resize(numMolecules());
+    auto err = detail::convertDeviceArray(singleConversion_.positions.data(), positions, totalPositions(), stream);
+    if (err == cudaSuccess) {
+      err =
+        computeEnergy(singleConversion_.energies.data(), singleConversion_.positions.data(), activeSystemMask, stream);
+    }
+    return err == cudaSuccess ?
+             detail::convertDeviceArray(energyOuts, singleConversion_.energies.data(), numMolecules(), stream) :
+             err;
+  }
+  const auto& buffers = std::get<DistGeom::BatchedAllInOneDeviceBuffers>(systemDevice_);
+  return DistGeom::launchBlockPerMolEnergyKernelAIO(
+    numMolecules(),
+    DistGeom::toAllInOneContribsDevicePtr(buffers, atomStartsDevice_.data()),
+    positions,
+    energyOuts,
+    weights_,
+    torsionTermsEnabled_,
+    activeSystemMask,
+    stream);
+}
+
+cudaError_t AllInOneETKDGBatchedForcefield::computeGradients(double*        grad,
+                                                             const double*  positions,
+                                                             const uint8_t* activeSystemMask,
+                                                             cudaStream_t   stream) {
+  if (std::holds_alternative<DistGeom::BatchedAllInOneDeviceBuffersSingle>(systemDevice_)) {
+    singleConversion_.positions.setStream(stream);
+    singleConversion_.gradients.setStream(stream);
+    singleConversion_.positions.resize(totalPositions());
+    singleConversion_.gradients.resize(totalPositions());
+    auto err = detail::convertDeviceArray(singleConversion_.positions.data(), positions, totalPositions(), stream);
+    if (err != cudaSuccess)
+      return err;
+    singleConversion_.gradients.zero();
+    err = computeGradients(singleConversion_.gradients.data(),
+                           singleConversion_.positions.data(),
+                           activeSystemMask,
+                           stream);
+    return err == cudaSuccess ?
+             detail::convertDeviceArray(grad, singleConversion_.gradients.data(), totalPositions(), stream) :
+             err;
+  }
+  const auto& buffers = std::get<DistGeom::BatchedAllInOneDeviceBuffers>(systemDevice_);
+  return DistGeom::launchBlockPerMolGradKernelAIO(
+    numMolecules(),
+    DistGeom::toAllInOneContribsDevicePtr(buffers, atomStartsDevice_.data()),
+    positions,
+    grad,
+    weights_,
+    torsionTermsEnabled_,
+    activeSystemMask,
+    stream);
+}
+
+cudaError_t AllInOneETKDGBatchedForcefield::computeEnergy(float*         energyOuts,
+                                                          const float*   positions,
+                                                          const uint8_t* activeSystemMask,
+                                                          cudaStream_t   stream) {
+  if (!std::holds_alternative<DistGeom::BatchedAllInOneDeviceBuffersSingle>(systemDevice_)) {
+    return cudaErrorInvalidValue;
+  }
+  const auto& buffers = std::get<DistGeom::BatchedAllInOneDeviceBuffersSingle>(systemDevice_);
+  return DistGeom::launchBlockPerMolEnergyKernelAIO(
+    numMolecules(),
+    DistGeom::toAllInOneContribsDevicePtr(buffers, atomStartsDevice_.data()),
+    positions,
+    energyOuts,
+    weights_,
+    torsionTermsEnabled_,
+    activeSystemMask,
+    stream);
+}
+
+cudaError_t AllInOneETKDGBatchedForcefield::computeGradients(float*         grad,
+                                                             const float*   positions,
+                                                             const uint8_t* activeSystemMask,
+                                                             cudaStream_t   stream) {
+  if (!std::holds_alternative<DistGeom::BatchedAllInOneDeviceBuffersSingle>(systemDevice_)) {
+    return cudaErrorInvalidValue;
+  }
+  const auto& buffers = std::get<DistGeom::BatchedAllInOneDeviceBuffersSingle>(systemDevice_);
+  return DistGeom::launchBlockPerMolGradKernelAIO(
+    numMolecules(),
+    DistGeom::toAllInOneContribsDevicePtr(buffers, atomStartsDevice_.data()),
+    positions,
+    grad,
+    weights_,
+    torsionTermsEnabled_,
+    activeSystemMask,
+    stream);
+}
+
+}  // namespace nvMolKit

src/forcefields/aio_etkdg_batched_forcefield.h

@@ -0,0 +1,64 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef NVMOLKIT_AIO_ETKDG_BATCHED_FORCEFIELD_H
+#define NVMOLKIT_AIO_ETKDG_BATCHED_FORCEFIELD_H
+
+#include <variant>
+
+#include "src/forcefields/batched_forcefield.h"
+#include "src/forcefields/dist_geom.h"
+#include "src/forcefields/precision_workspace.h"
+#include "src/precision/precision_mode.h"
+
+namespace nvMolKit {
+
+//! \brief Batched-forcefield adapter for RDKit's all-in-one ETKDG coordinate refinement.
+//!
+//! Evaluates the 4D distance-geometry terms together with the all-in-one harmonic distance, linear angle,
+//! experimental torsion and planarity terms. The torsion and planarity terms start disabled, matching RDKit's first
+//! minimization phase.
+class AllInOneETKDGBatchedForcefield final : public BatchedForcefield, public SinglePrecisionBatchedForcefield {
+ public:
+  //! \param systemHost Flattened all-in-one host-side system description.
+  //! \param atomStartsHost Host-side atom offsets for the minimized systems.
+  //! \param weights Force constants shared by every system in the batch.
+  //! \param metadata Optional mapping from concrete systems back to logical molecules/conformers.
+  //! \param stream CUDA stream used for internal device allocations and uploads.
+  //! \param precision Arithmetic precision of the device buffers and kernels.
+  AllInOneETKDGBatchedForcefield(const DistGeom::BatchedAllInOneSystemHost& systemHost,
+                                 const std::vector<int>&                    atomStartsHost,
+                                 const DistGeom::AllInOneForceWeights&      weights,
+                                 BatchedForcefieldMetadata                  metadata  = {},
+                                 cudaStream_t                               stream    = nullptr,
+                                 PrecisionMode                              precision = PrecisionMode::FULL);
+
+  cudaError_t   computeEnergy(double*        energyOuts,
+                              const double*  positions,
+                              const uint8_t* activeSystemMask = nullptr,
+                              cudaStream_t   stream           = nullptr) override;
+  cudaError_t   computeGradients(double*        grad,
+                                 const double*  positions,
+                                 const uint8_t* activeSystemMask = nullptr,
+                                 cudaStream_t   stream           = nullptr) override;
+  cudaError_t   computeEnergy(float*, const float*, const uint8_t* = nullptr, cudaStream_t = nullptr) override;
+  cudaError_t   computeGradients(float*, const float*, const uint8_t* = nullptr, cudaStream_t = nullptr) override;
+  PrecisionMode precision() const override {
+    return std::holds_alternative<DistGeom::BatchedAllInOneDeviceBuffersSingle>(systemDevice_) ? PrecisionMode::SINGLE :
+                                                                                                 PrecisionMode::FULL;
+  }
+
+  //! Enables the experimental torsion and planarity terms for RDKit's second minimization phase.
+  void setTorsionTermsEnabled(bool enabled) { torsionTermsEnabled_ = enabled; }
+
+ private:
+  std::variant<DistGeom::BatchedAllInOneDeviceBuffers, DistGeom::BatchedAllInOneDeviceBuffersSingle> systemDevice_;
+  AsyncDeviceVector<int>                                                                             atomStartsDevice_;
+  SingleForcefieldConversionWorkspace                                                                singleConversion_;
+  DistGeom::AllInOneForceWeights                                                                     weights_;
+  bool torsionTermsEnabled_ = false;
+};
+
+}  // namespace nvMolKit
+
+#endif  // NVMOLKIT_AIO_ETKDG_BATCHED_FORCEFIELD_H