From 13548f54ab69abbede0c61ccc6839c51b07b6fd4 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Mon, 29 Jun 2026 10:16:54 -0400 Subject: [PATCH] Add UFF FIRE wrapper integration --- nvmolkit/CMakeLists.txt | 5 +- nvmolkit/batchedForcefield.cpp | 2 +- nvmolkit/uffOptimization.cpp | 2 +- src/minimizer/CMakeLists.txt | 7 +- .../{bfgs_uff.cpp => uff_minimize.cpp} | 201 +++++++++++++++++- src/minimizer/{bfgs_uff.h => uff_minimize.h} | 42 +++- tests/CMakeLists.txt | 7 +- tests/test_ff_device_output.cu | 2 +- tests/test_uff.cu | 49 +++++ 9 files changed, 304 insertions(+), 13 deletions(-) rename src/minimizer/{bfgs_uff.cpp => uff_minimize.cpp} (55%) rename src/minimizer/{bfgs_uff.h => uff_minimize.h} (65%) diff --git a/nvmolkit/CMakeLists.txt b/nvmolkit/CMakeLists.txt index 816e8d64..50eac782 100644 --- a/nvmolkit/CMakeLists.txt +++ b/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 diff --git a/nvmolkit/batchedForcefield.cpp b/nvmolkit/batchedForcefield.cpp index e52ed3a4..4b8f41e5 100644 --- a/nvmolkit/batchedForcefield.cpp +++ b/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; diff --git a/nvmolkit/uffOptimization.cpp b/nvmolkit/uffOptimization.cpp index f1f11e91..9356552a 100644 --- a/nvmolkit/uffOptimization.cpp +++ b/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; diff --git a/src/minimizer/CMakeLists.txt b/src/minimizer/CMakeLists.txt index ec293953..7e7f3b8d 100644 --- a/src/minimizer/CMakeLists.txt +++ b/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 diff --git a/src/minimizer/bfgs_uff.cpp b/src/minimizer/uff_minimize.cpp similarity index 55% rename from src/minimizer/bfgs_uff.cpp rename to src/minimizer/uff_minimize.cpp index af5b08b4..f11e293d 100644 --- a/src/minimizer/bfgs_uff.cpp +++ b/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 #include @@ -262,4 +262,203 @@ std::vector> UFFOptimizeMoleculesConfsBfgs(std::vector& mols, + const int maxIters, + const FireOptions& fireOptions, + const std::vector& vdwThresholds, + const std::vector& ignoreInterfragInteractions, + const std::vector& 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> moleculeEnergies; + const auto allConformers = flattenConformers(mols, moleculeEnergies); + + std::vector> 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 emptyCollectors; + return {{}, {}, detail::finalizeOnTarget(emptyCollectors, targetGpu, static_cast(mols.size()))}; + } + return {moleculeEnergies, moleculeConverged, std::nullopt}; + } + + std::vector threadBuffers(ctx.numThreads); + std::vector 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 batchConformers(allConformers.begin() + batchStart, + allConformers.begin() + batchEnd); + + cudaStream_t streamPtr = ctx.streamPool[threadId].stream(); + + BatchedMolecularSystemHost systemHost; + BatchedForcefieldMetadata metadata; + std::vector conformerAtomStarts; + std::vector massesPerAtom; + uint32_t currentAtomOffset = 0; + std::vector 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(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 positionsDevice; + AsyncDeviceVector gradDevice; + AsyncDeviceVector 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 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(statusesHost[i] == 0); + } + } + } catch (...) { + exceptionHandler.store(std::current_exception()); + } + } + exceptionHandler.rethrow(); + if (deviceOutput) { + return {{}, {}, detail::finalizeOnTarget(deviceCollectors, targetGpu, static_cast(mols.size()))}; + } + return {moleculeEnergies, moleculeConverged, std::nullopt}; +} + +std::vector> UFFOptimizeMoleculesConfsFire(std::vector& mols, + const int maxIters, + const FireOptions& fireOptions, + const std::vector& vdwThresholds, + const std::vector& ignoreInterfragInteractions, + const BatchHardwareOptions& perfOptions) { + return UFFMinimizeMoleculesConfsFire(mols, + maxIters, + fireOptions, + vdwThresholds, + ignoreInterfragInteractions, + {}, + perfOptions) + .energies; +} + } // namespace nvMolKit::UFF diff --git a/src/minimizer/bfgs_uff.h b/src/minimizer/uff_minimize.h similarity index 65% rename from src/minimizer/bfgs_uff.h rename to src/minimizer/uff_minimize.h index e842b953..fcdfe88b 100644 --- a/src/minimizer/bfgs_uff.h +++ b/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 #include @@ -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> UFFOptimizeMoleculesConfsBfgs(std::vector& ignoreInterfragInteractions, const BatchHardwareOptions& perfOptions = {}); +std::vector> UFFOptimizeMoleculesConfsFire(std::vector& mols, + int maxIters, + const FireOptions& fireOptions, + const std::vector& vdwThresholds, + const std::vector& 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& mols, + int maxIters, + const FireOptions& fireOptions, + const std::vector& vdwThresholds, + const std::vector& ignoreInterfragInteractions, + const std::vector& 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 diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 658dbfef..e6cc92df 100644 --- a/tests/CMakeLists.txt +++ b/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 diff --git a/tests/test_ff_device_output.cu b/tests/test_ff_device_output.cu index 6b77af7c..02142c3c 100644 --- a/tests/test_ff_device_output.cu +++ b/tests/test_ff_device_output.cu @@ -23,8 +23,8 @@ #include #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" diff --git a/tests/test_uff.cu b/tests/test_uff.cu index b325a304..2be7abc3 100644 --- a/tests/test_uff.cu +++ b/tests/test_uff.cu @@ -28,6 +28,7 @@ #include #include +#include #include #include #include @@ -39,6 +40,7 @@ #include "src/forcefields/uff.h" #include "src/forcefields/uff_batched_forcefield.h" #include "src/minimizer/bfgs_minimize.h" +#include "src/minimizer/uff_minimize.h" #include "src/utils/device_vector.h" #include "tests/test_utils.h" @@ -291,6 +293,53 @@ TEST(UFFMinimizer, BatchMinimizerMatchesRDKitFinalEnergies) { } } +TEST(UFFMinimizer, FireWrapperImprovesAndApproachesRDKitFinalEnergies) { + const std::string sdfPath = getTestDataFolderPath() + "/MMFF94_dative.sdf"; + std::vector> mols; + getMols(sdfPath, mols, 2); + ASSERT_FALSE(mols.empty()); + + std::vector molPtrs; + std::vector initialEnergies; + std::vector referenceFinalEnergies; + molPtrs.reserve(mols.size()); + initialEnergies.reserve(mols.size()); + referenceFinalEnergies.reserve(mols.size()); + for (auto& molPtr : mols) { + auto referenceMol = std::make_unique(*molPtr); + auto initialFF = buildReferenceForceField(*referenceMol); + initialEnergies.push_back(initialFF->calcEnergy()); + + const auto optimizeResult = RDKit::UFF::UFFOptimizeMolecule(*referenceMol, 1000, 100.0, -1, true); + referenceFinalEnergies.push_back(optimizeResult.second); + molPtrs.push_back(molPtr.get()); + } + + nvMolKit::FireOptions options{}; + options.useMass = false; + options.stuckDetectionEnabled = false; + options.gradTol = 1e-3; + options.dtInit = 0.05; + options.dMax = 0.2; + + const std::vector vdwThresholds(molPtrs.size(), 100.0); + const std::vector ignoreInterfragInteractions(molPtrs.size(), true); + const auto result = UFFMinimizeMoleculesConfsFire(molPtrs, + /*maxIters=*/10000, + options, + vdwThresholds, + ignoreInterfragInteractions); + + ASSERT_FALSE(result.device.has_value()); + ASSERT_EQ(result.energies.size(), referenceFinalEnergies.size()); + for (size_t i = 0; i < result.energies.size(); ++i) { + ASSERT_EQ(result.energies[i].size(), 1u); + EXPECT_LT(result.energies[i][0], initialEnergies[i]) << "molecule " << i; + const double tolerance = 2.0 + 0.05 * std::abs(referenceFinalEnergies[i]); + EXPECT_NEAR(result.energies[i][0], referenceFinalEnergies[i], tolerance) << "molecule " << i; + } +} + namespace { double getConstraintEnergyViaForcefield(const EnergyForceContribsHost& contribs, const std::vector& positions) {