Add UFF FIRE wrapper integration

nvmolkit/CMakeLists.txt

@@ -53,7 +53,8 @@ installpythontarget(_mmffOptimization ./)
 add_library(_uffOptimization MODULE uffOptimization.cpp)
 target_link_libraries(_uffOptimization PUBLIC ${Boost_LIBRARIES}
                                               ${PYTHON_LIBRARIES})
-target_link_libraries(_uffOptimization PRIVATE ${RDKit_LIBS} bfgs_uff ff_utils)
+target_link_libraries(_uffOptimization PRIVATE ${RDKit_LIBS} uff_minimize
+                                               ff_utils)
 target_include_directories(_uffOptimization PUBLIC ${Python_INCLUDE_DIRS})
 target_include_directories(_uffOptimization SYSTEM PUBLIC ${Boost_INCLUDE_DIRS})
 installpythontarget(_uffOptimization ./)
@@ -65,7 +66,7 @@ target_link_libraries(
   _batchedForcefield
   PRIVATE ${RDKit_LIBS}
           mmff_minimize
-          bfgs_uff
+          uff_minimize
           mmff_batched_forcefield
           uff_batched_forcefield
           batched_forcefield

nvmolkit/batchedForcefield.cpp

@@ -32,8 +32,8 @@
 #include "src/forcefields/mmff_properties.h"
 #include "src/forcefields/uff_batched_forcefield.h"
 #include "src/hardware_options.h"
-#include "src/minimizer/bfgs_uff.h"
 #include "src/minimizer/mmff_minimize.h"
+#include "src/minimizer/uff_minimize.h"
 #include "src/utils/device_vector.h"
 
 namespace bp = boost::python;

nvmolkit/uffOptimization.cpp

@@ -17,7 +17,7 @@
 
 #include "nvmolkit/boost_python_utils.h"
 #include "nvmolkit/device_result_python.h"
-#include "src/minimizer/bfgs_uff.h"
+#include "src/minimizer/uff_minimize.h"
 
 namespace bp = boost::python;
 

src/minimizer/CMakeLists.txt

@@ -55,13 +55,14 @@ target_link_libraries(
           OpenMP::OpenMP_CXX
           openmp_helpers
           cccl_interface)
-add_library(bfgs_uff bfgs_uff.cpp)
+add_library(uff_minimize uff_minimize.cpp)
 target_link_libraries(
-  bfgs_uff
-  PUBLIC bfgs conformer_types
+  uff_minimize
+  PUBLIC bfgs fire_minimizer conformer_types
   PRIVATE ${RDKit_LIBS}
           bfgs_common
           ff_device_collect
+          ff_utils
           uff_batched_forcefield
           rdkit_uff_flattened
           forcefield_constraints

src/minimizer/bfgs_uff.cpp → src/minimizer/uff_minimize.cpp

@@ -13,7 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "src/minimizer/bfgs_uff.h"
+#include "src/minimizer/uff_minimize.h"
 
 #include <GraphMol/ROMol.h>
 #include <omp.h>
@@ -262,4 +262,203 @@ std::vector<std::vector<double>> UFFOptimizeMoleculesConfsBfgs(std::vector<RDKit
     .energies;
 }
 
+UFFMinimizeResult UFFMinimizeMoleculesConfsFire(
+  std::vector<RDKit::ROMol*>&                                  mols,
+  const int                                                    maxIters,
+  const FireOptions&                                           fireOptions,
+  const std::vector<double>&                                   vdwThresholds,
+  const std::vector<bool>&                                     ignoreInterfragInteractions,
+  const std::vector<ForceFieldConstraints::PerMolConstraints>& constraints,
+  const BatchHardwareOptions&                                  perfOptions,
+  const CoordinateOutput                                       output,
+  int                                                          targetGpu) {
+  ScopedNvtxRange fullRange("FIRE UFF Minimize Molecules Confs");
+
+  if (vdwThresholds.size() != mols.size()) {
+    throw std::invalid_argument("Expected one vdw threshold per molecule");
+  }
+  if (ignoreInterfragInteractions.size() != mols.size()) {
+    throw std::invalid_argument("Expected one interfragment interaction flag per molecule");
+  }
+  if (!constraints.empty() && constraints.size() != mols.size()) {
+    throw std::invalid_argument("Expected one PerMolConstraints entry per molecule");
+  }
+
+  const bool deviceOutput = output == CoordinateOutput::DEVICE;
+
+  auto ctx = setupBatchExecution(perfOptions);
+  if (deviceOutput) {
+    if (targetGpu < 0) {
+      targetGpu = ctx.devicesPerThread.empty() ? 0 : ctx.devicesPerThread.front();
+    }
+    if (std::find(ctx.devicesPerThread.begin(), ctx.devicesPerThread.end(), targetGpu) == ctx.devicesPerThread.end()) {
+      throw std::invalid_argument(
+        "targetGpu " + std::to_string(targetGpu) +
+        " is not in the configured set of execution GPUs; pass it via perfOptions.gpuIds first.");
+    }
+  }
+
+  std::vector<std::vector<double>> moleculeEnergies;
+  const auto                       allConformers = flattenConformers(mols, moleculeEnergies);
+
+  std::vector<std::vector<int8_t>> moleculeConverged(mols.size());
+  for (size_t i = 0; i < mols.size(); ++i) {
+    moleculeConverged[i].resize(moleculeEnergies[i].size(), 0);
+  }
+
+  const size_t totalConformers    = allConformers.size();
+  const size_t effectiveBatchSize = ctx.batchSize == 0 ? totalConformers : ctx.batchSize;
+  if (totalConformers == 0) {
+    if (deviceOutput) {
+      std::vector<detail::DeviceCoordCollector> emptyCollectors;
+      return {{}, {}, detail::finalizeOnTarget(emptyCollectors, targetGpu, static_cast<int>(mols.size()))};
+    }
+    return {moleculeEnergies, moleculeConverged, std::nullopt};
+  }
+
+  std::vector<ThreadLocalBuffers>           threadBuffers(ctx.numThreads);
+  std::vector<detail::DeviceCoordCollector> deviceCollectors(deviceOutput ? ctx.numThreads : 0);
+  if (deviceOutput) {
+    for (int threadId = 0; threadId < ctx.numThreads; ++threadId) {
+      auto& collector  = deviceCollectors[threadId];
+      collector.gpuId  = ctx.devicesPerThread[threadId];
+      collector.stream = ctx.streamPool[threadId].stream();
+      collector.positions.setStream(collector.stream);
+      collector.energies.setStream(collector.stream);
+      collector.converged.setStream(collector.stream);
+    }
+  }
+  detail::OpenMPExceptionRegistry exceptionHandler;
+#pragma omp parallel for num_threads(ctx.numThreads) schedule(dynamic) default(none) \
+  shared(allConformers,                                                              \
+           moleculeEnergies,                                                         \
+           moleculeConverged,                                                        \
+           totalConformers,                                                          \
+           effectiveBatchSize,                                                       \
+           maxIters,                                                                 \
+           fireOptions,                                                              \
+           vdwThresholds,                                                            \
+           ignoreInterfragInteractions,                                              \
+           constraints,                                                              \
+           ctx,                                                                      \
+           threadBuffers,                                                            \
+           deviceCollectors,                                                         \
+           deviceOutput,                                                             \
+           exceptionHandler)
+  for (size_t batchStart = 0; batchStart < totalConformers; batchStart += effectiveBatchSize) {
+    try {
+      ScopedNvtxRange singleBatchRange("OpenMP loop thread");
+      ScopedNvtxRange setupBatchRange("OpenMP loop preprocessing");
+
+      const int                            threadId = omp_get_thread_num();
+      const WithDevice                     dev(ctx.devicesPerThread[threadId]);
+      const size_t                         batchEnd = std::min(batchStart + effectiveBatchSize, totalConformers);
+      std::vector<nvMolKit::ConformerInfo> batchConformers(allConformers.begin() + batchStart,
+                                                           allConformers.begin() + batchEnd);
+
+      cudaStream_t streamPtr = ctx.streamPool[threadId].stream();
+
+      BatchedMolecularSystemHost systemHost;
+      BatchedForcefieldMetadata  metadata;
+      std::vector<uint32_t>      conformerAtomStarts;
+      std::vector<double>        massesPerAtom;
+      uint32_t                   currentAtomOffset = 0;
+      std::vector<double>        pos;
+
+      for (const auto& confInfo : batchConformers) {
+        const uint32_t numAtoms = confInfo.mol->getNumAtoms();
+        for (uint32_t atomIdx = 0; atomIdx < numAtoms; ++atomIdx) {
+          massesPerAtom.push_back(confInfo.mol->getAtomWithIdx(atomIdx)->getMass());
+        }
+        conformerAtomStarts.push_back(currentAtomOffset);
+        currentAtomOffset += numAtoms;
+
+        nvMolKit::confPosToVect(*confInfo.conformer, pos);
+        auto ffParams = constructForcefieldContribs(*confInfo.mol,
+                                                    vdwThresholds[confInfo.molIdx],
+                                                    confInfo.conformerId,
+                                                    ignoreInterfragInteractions[confInfo.molIdx]);
+        if (!constraints.empty()) {
+          constraints[confInfo.molIdx].applyTo(ffParams, pos);
+        }
+        addMoleculeToBatch(ffParams, pos, systemHost, metadata, confInfo.molIdx, static_cast<int>(confInfo.confIdx));
+      }
+
+      auto& buffers = threadBuffers[threadId];
+      buffers.ensureCapacity(systemHost.positions.size(), batchConformers.size());
+      std::copy(systemHost.positions.begin(), systemHost.positions.end(), buffers.initialPositions.begin());
+
+      UFFBatchedForcefield      forcefield(systemHost, metadata, streamPtr);
+      AsyncDeviceVector<double> positionsDevice;
+      AsyncDeviceVector<double> gradDevice;
+      AsyncDeviceVector<double> energyOutsDevice;
+      positionsDevice.setStream(streamPtr);
+      gradDevice.setStream(streamPtr);
+      energyOutsDevice.setStream(streamPtr);
+      positionsDevice.resize(systemHost.positions.size());
+      positionsDevice.copyFromHost(buffers.initialPositions.data(), systemHost.positions.size());
+      gradDevice.resize(systemHost.positions.size());
+      gradDevice.zero();
+      energyOutsDevice.resize(batchConformers.size());
+      energyOutsDevice.zero();
+
+      FireBatchMinimizer fireMinimizer(/*dataDim=*/3,
+                                       fireOptions,
+                                       streamPtr,
+                                       /*debugMode=*/false,
+                                       FireBackend::BATCHED);
+      if (fireOptions.useMass) {
+        fireMinimizer.setMasses(massesPerAtom);
+      }
+      setupBatchRange.pop();
+      fireMinimizer.minimize(maxIters, fireOptions.gradTol, forcefield, positionsDevice, gradDevice, energyOutsDevice);
+      forcefield.computeEnergy(energyOutsDevice.data(), positionsDevice.data(), nullptr, streamPtr);
+
+      if (deviceOutput) {
+        detail::appendBatch(batchConformers,
+                            positionsDevice,
+                            energyOutsDevice,
+                            fireMinimizer.statuses(),
+                            deviceCollectors[threadId]);
+      } else {
+        ScopedNvtxRange finalizeBatchRange("OpenMP loop finalizing batch");
+        positionsDevice.copyToHost(buffers.positions.data(), positionsDevice.size());
+        energyOutsDevice.copyToHost(buffers.energies.data(), energyOutsDevice.size());
+        std::vector<uint8_t> statusesHost(batchConformers.size());
+        fireMinimizer.statuses().copyToHost(statusesHost.data(), batchConformers.size());
+        cudaStreamSynchronize(streamPtr);
+
+        writeBackResults(batchConformers, conformerAtomStarts, buffers, moleculeEnergies);
+        for (size_t i = 0; i < batchConformers.size(); ++i) {
+          const auto& confInfo                                 = batchConformers[i];
+          moleculeConverged[confInfo.molIdx][confInfo.confIdx] = static_cast<int8_t>(statusesHost[i] == 0);
+        }
+      }
+    } catch (...) {
+      exceptionHandler.store(std::current_exception());
+    }
+  }
+  exceptionHandler.rethrow();
+  if (deviceOutput) {
+    return {{}, {}, detail::finalizeOnTarget(deviceCollectors, targetGpu, static_cast<int>(mols.size()))};
+  }
+  return {moleculeEnergies, moleculeConverged, std::nullopt};
+}
+
+std::vector<std::vector<double>> UFFOptimizeMoleculesConfsFire(std::vector<RDKit::ROMol*>& mols,
+                                                               const int                   maxIters,
+                                                               const FireOptions&          fireOptions,
+                                                               const std::vector<double>&  vdwThresholds,
+                                                               const std::vector<bool>&    ignoreInterfragInteractions,
+                                                               const BatchHardwareOptions& perfOptions) {
+  return UFFMinimizeMoleculesConfsFire(mols,
+                                       maxIters,
+                                       fireOptions,
+                                       vdwThresholds,
+                                       ignoreInterfragInteractions,
+                                       {},
+                                       perfOptions)
+    .energies;
+}
+
 }  // namespace nvMolKit::UFF

src/minimizer/bfgs_uff.h → src/minimizer/uff_minimize.h

@@ -13,8 +13,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef NVMOLKIT_BFGS_UFF_H
-#define NVMOLKIT_BFGS_UFF_H
+#ifndef NVMOLKIT_UFF_MINIMIZE_H
+#define NVMOLKIT_UFF_MINIMIZE_H
 
 #include <cstdint>
 #include <optional>
@@ -24,6 +24,7 @@
 #include "src/forcefields/forcefield_constraints.h"
 #include "src/hardware_options.h"
 #include "src/minimizer/bfgs_minimize.h"
+#include "src/minimizer/fire_minimizer.h"
 
 namespace RDKit {
 class ROMol;
@@ -37,6 +38,13 @@ std::vector<std::vector<double>> UFFOptimizeMoleculesConfsBfgs(std::vector<RDKit
                                                                const std::vector<bool>&    ignoreInterfragInteractions,
                                                                const BatchHardwareOptions& perfOptions = {});
 
+std::vector<std::vector<double>> UFFOptimizeMoleculesConfsFire(std::vector<RDKit::ROMol*>& mols,
+                                                               int                         maxIters,
+                                                               const FireOptions&          fireOptions,
+                                                               const std::vector<double>&  vdwThresholds,
+                                                               const std::vector<bool>&    ignoreInterfragInteractions,
+                                                               const BatchHardwareOptions& perfOptions = {});
+
 //! \brief Result from constraint-aware UFF minimization.
 //!
 //! In CoordinateOutput::RDKIT_CONFORMERS mode, @ref energies and @ref converged are populated and
@@ -87,6 +95,34 @@ UFFMinimizeResult UFFMinimizeMoleculesConfs(
   int                                                          targetGpu   = -1,
   const DeviceCoordResult*                                     deviceInput = nullptr);
 
+//! \brief Minimize UFF energies with FIRE 2.0 and report per-conformer convergence.
+//! \param mols Molecules whose conformers provide the initial coordinates. In
+//!             RDKIT_CONFORMERS mode, optimized coordinates are written back in place.
+//! \param maxIters Maximum FIRE iterations.
+//! \param fireOptions FIRE integration and convergence settings. A conformer converges when
+//!                    its gradient norm satisfies @ref FireOptions::gradTol. The UFF wrapper
+//!                    uses the batched backend and can additionally use energy-plateau detection.
+//! \param vdwThresholds Per-molecule VDW cutoff distances.
+//! \param ignoreInterfragInteractions Per-molecule interfragment interaction flags.
+//! \param constraints Per-molecule constraint specifications (empty = no constraints).
+//! \param perfOptions Hardware and batching configuration.
+//! \param output Whether to write coordinates back into RDKit conformers (default) or return
+//!               them on-device as a DeviceCoordResult.
+//! \param targetGpu In DEVICE mode, the GPU to consolidate the result onto. -1 selects the first
+//!                  configured execution GPU (or device 0).
+//! \return Final energies and convergence flags in RDKIT_CONFORMERS mode, or a device-resident
+//!         result in DEVICE mode. A convergence flag is false when @p maxIters is reached first.
+UFFMinimizeResult UFFMinimizeMoleculesConfsFire(
+  std::vector<RDKit::ROMol*>&                                  mols,
+  int                                                          maxIters,
+  const FireOptions&                                           fireOptions,
+  const std::vector<double>&                                   vdwThresholds,
+  const std::vector<bool>&                                     ignoreInterfragInteractions,
+  const std::vector<ForceFieldConstraints::PerMolConstraints>& constraints = {},
+  const BatchHardwareOptions&                                  perfOptions = {},
+  CoordinateOutput                                             output      = CoordinateOutput::RDKIT_CONFORMERS,
+  int                                                          targetGpu   = -1);
+
 }  // namespace nvMolKit::UFF
 
-#endif  // NVMOLKIT_BFGS_UFF_H
+#endif  // NVMOLKIT_UFF_MINIMIZE_H

tests/CMakeLists.txt

@@ -105,7 +105,12 @@ target_link_libraries(
   test_uff
   PRIVATE ${RDKit_LIBS}
           uff_batched_forcefield
+          uff_minimize
+          bfgs_common
           bfgs
+          ff_device_collect
+          fire_minimizer
+          forcefield_constraints
           uff
           rdkit_uff_flattened
           device_vector
@@ -189,7 +194,7 @@ add_executable(test_ff_device_output test_ff_device_output.cu)
 target_link_libraries(
   test_ff_device_output
   PRIVATE mmff_minimize
-          bfgs_uff
+          uff_minimize
           bfgs_common
           mmff
           uff

tests/test_ff_device_output.cu

@@ -23,8 +23,8 @@
 #include <vector>
 
 #include "src/conformer/device_coord_result.h"
-#include "src/minimizer/bfgs_uff.h"
 #include "src/minimizer/mmff_minimize.h"
+#include "src/minimizer/uff_minimize.h"
 #include "src/utils/cuda_error_check.h"
 #include "src/utils/device.h"