diff --git a/rdkit_extensions/CMakeLists.txt b/rdkit_extensions/CMakeLists.txt index 4e896735..847ac774 100644 --- a/rdkit_extensions/CMakeLists.txt +++ b/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}) diff --git a/rdkit_extensions/dist_geom_flattened_builder.cpp b/rdkit_extensions/dist_geom_flattened_builder.cpp index 9e273171..984d5e92 100644 --- a/rdkit_extensions/dist_geom_flattened_builder.cpp +++ b/rdkit_extensions/dist_geom_flattened_builder.cpp @@ -6,6 +6,7 @@ #include #include +#include #include #include @@ -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(i)); + distTerms.idx2.push_back(static_cast(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, diff --git a/rdkit_extensions/dist_geom_flattened_builder.h b/rdkit_extensions/dist_geom_flattened_builder.h index e4541cab..8600cce9 100644 --- a/rdkit_extensions/dist_geom_flattened_builder.h +++ b/rdkit_extensions/dist_geom_flattened_builder.h @@ -8,6 +8,7 @@ #include #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 diff --git a/src/forcefields/CMakeLists.txt b/src/forcefields/CMakeLists.txt index 08985428..28bb2fb7 100644 --- a/src/forcefields/CMakeLists.txt +++ b/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) diff --git a/src/forcefields/aio_etkdg_batched_forcefield.cu b/src/forcefields/aio_etkdg_batched_forcefield.cu new file mode 100644 index 00000000..25d01af7 --- /dev/null +++ b/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& 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(), + stream); + } else { + DistGeom::sendAllInOneSystemToDevice(systemHost, + systemDevice_.emplace(), + 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(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(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(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(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(systemDevice_)) { + return cudaErrorInvalidValue; + } + const auto& buffers = std::get(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(systemDevice_)) { + return cudaErrorInvalidValue; + } + const auto& buffers = std::get(systemDevice_); + return DistGeom::launchBlockPerMolGradKernelAIO( + numMolecules(), + DistGeom::toAllInOneContribsDevicePtr(buffers, atomStartsDevice_.data()), + positions, + grad, + weights_, + torsionTermsEnabled_, + activeSystemMask, + stream); +} + +} // namespace nvMolKit diff --git a/src/forcefields/aio_etkdg_batched_forcefield.h b/src/forcefields/aio_etkdg_batched_forcefield.h new file mode 100644 index 00000000..cfbf0448 --- /dev/null +++ b/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 + +#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& 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(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 systemDevice_; + AsyncDeviceVector atomStartsDevice_; + SingleForcefieldConversionWorkspace singleConversion_; + DistGeom::AllInOneForceWeights weights_; + bool torsionTermsEnabled_ = false; +}; + +} // namespace nvMolKit + +#endif // NVMOLKIT_AIO_ETKDG_BATCHED_FORCEFIELD_H diff --git a/src/forcefields/dist_geom.cu b/src/forcefields/dist_geom.cu index a96d0b6c..4bffcaa5 100644 --- a/src/forcefields/dist_geom.cu +++ b/src/forcefields/dist_geom.cu @@ -1597,5 +1597,188 @@ BatchedIndices3DDevicePtr toBatchedIndices3DDevicePtr(const BatchedMolecular3DDe return toPointerStruct(molSystemDevice.indices, atomStarts); } +namespace { + +void appendOffsetIndices(std::vector& destination, const std::vector& source, const int offset) { + for (const int idx : source) { + destination.push_back(idx + offset); + } +} + +template void append(std::vector& destination, const std::vector& source) { + destination.insert(destination.end(), source.begin(), source.end()); +} + +void appendTermStart(std::vector& termStarts, const std::size_t numTerms) { + termStarts.push_back(termStarts.back() + static_cast(numTerms)); +} + +void upload(AsyncDeviceVector& destination, const std::vector& source, cudaStream_t stream) { + destination.setStream(stream); + destination.setFromVector(source); +} + +template +void upload(AsyncDeviceVector& destination, const std::vector& source, cudaStream_t stream) { + destination.setStream(stream); + uploadWithPrecisionConversion(destination, source); +} + +template +void sendAllInOneSystemToDeviceImpl(const BatchedAllInOneSystemHost& systemHost, + BatchedAllInOneDeviceBuffersT& systemDevice, + cudaStream_t stream) { + setStreamsImpl(systemDevice.distanceGeometry, stream); + sendContribsAndIndicesToDeviceImpl(systemHost.distanceGeometry, systemDevice.distanceGeometry); + + const auto& harmonicHost = systemHost.harmonicDistanceTerms; + auto& harmonicDevice = systemDevice.harmonicDistanceTerms; + upload(harmonicDevice.idx1, harmonicHost.idx1, stream); + upload(harmonicDevice.idx2, harmonicHost.idx2, stream); + upload(harmonicDevice.minLen, harmonicHost.minLen, stream); + upload(harmonicDevice.maxLen, harmonicHost.maxLen, stream); + upload(harmonicDevice.forceConstant, harmonicHost.forceConstant, stream); + + const auto& angleHost = systemHost.linearAngleTerms; + auto& angleDevice = systemDevice.linearAngleTerms; + upload(angleDevice.idx1, angleHost.idx1, stream); + upload(angleDevice.idx2, angleHost.idx2, stream); + upload(angleDevice.idx3, angleHost.idx3, stream); + upload(angleDevice.minAngle, angleHost.minAngle, stream); + upload(angleDevice.maxAngle, angleHost.maxAngle, stream); + + const auto& torsionHost = systemHost.experimentalTorsionTerms; + auto& torsionDevice = systemDevice.experimentalTorsionTerms; + upload(torsionDevice.idx1, torsionHost.idx1, stream); + upload(torsionDevice.idx2, torsionHost.idx2, stream); + upload(torsionDevice.idx3, torsionHost.idx3, stream); + upload(torsionDevice.idx4, torsionHost.idx4, stream); + upload(torsionDevice.forceConstants, torsionHost.forceConstants, stream); + upload(torsionDevice.signs, torsionHost.signs, stream); + + const auto& planarityHost = systemHost.planarityTerms; + auto& planarityDevice = systemDevice.planarityTerms; + upload(planarityDevice.idx1, planarityHost.idx1, stream); + upload(planarityDevice.idx2, planarityHost.idx2, stream); + upload(planarityDevice.idx3, planarityHost.idx3, stream); + upload(planarityDevice.idx4, planarityHost.idx4, stream); + + upload(systemDevice.harmonicDistanceTermStarts, systemHost.harmonicDistanceTermStarts, stream); + upload(systemDevice.linearAngleTermStarts, systemHost.linearAngleTermStarts, stream); + upload(systemDevice.experimentalTorsionTermStarts, systemHost.experimentalTorsionTermStarts, stream); + upload(systemDevice.planarityTermStarts, systemHost.planarityTermStarts, stream); +} + +template +AllInOneContribsDevicePtrT toAllInOneContribsDevicePtrImpl( + const BatchedAllInOneDeviceBuffersT& systemDevice, + const int* atomStarts) { + AllInOneContribsDevicePtrT dst; + dst.distanceGeometry = toPointerStruct(systemDevice.distanceGeometry.contribs); + dst.distanceGeometryIndices = toPointerStruct(systemDevice.distanceGeometry.indices); + dst.distanceGeometryIndices.atomStarts = atomStarts; + + const auto& harmonic = systemDevice.harmonicDistanceTerms; + dst.harmonicDistanceTerms = {harmonic.idx1.data(), + harmonic.idx2.data(), + harmonic.minLen.data(), + harmonic.maxLen.data(), + harmonic.forceConstant.data()}; + const auto& angle = systemDevice.linearAngleTerms; + dst.linearAngleTerms = {angle.idx1.data(), + angle.idx2.data(), + angle.idx3.data(), + angle.minAngle.data(), + angle.maxAngle.data()}; + const auto& torsion = systemDevice.experimentalTorsionTerms; + dst.experimentalTorsionTerms = {torsion.idx1.data(), + torsion.idx2.data(), + torsion.idx3.data(), + torsion.idx4.data(), + torsion.forceConstants.data(), + torsion.signs.data()}; + const auto& planarity = systemDevice.planarityTerms; + dst.planarityTerms = {planarity.idx1.data(), planarity.idx2.data(), planarity.idx3.data(), planarity.idx4.data()}; + + dst.harmonicDistanceTermStarts = systemDevice.harmonicDistanceTermStarts.data(); + dst.linearAngleTermStarts = systemDevice.linearAngleTermStarts.data(); + dst.experimentalTorsionTermStarts = systemDevice.experimentalTorsionTermStarts.data(); + dst.planarityTermStarts = systemDevice.planarityTermStarts.data(); + return dst; +} + +} // namespace + +void addMoleculeToAllInOneSystem(const AllInOneForceContribsHost& contribs, + const int numAtoms, + const std::vector& ctxAtomStarts, + BatchedAllInOneSystemHost& system, + BatchedForcefieldMetadata* metadata, + const int moleculeIdx, + const int conformerIdx) { + const int batchIdx = static_cast(system.harmonicDistanceTermStarts.size()) - 1; + const int atomOffset = ctxAtomStarts[batchIdx]; + addMoleculeToMolecularSystem(contribs.distanceGeometry, + numAtoms, + 4, + ctxAtomStarts, + system.distanceGeometry, + metadata, + moleculeIdx, + conformerIdx); + + const auto& harmonic = contribs.harmonicDistanceTerms; + appendOffsetIndices(system.harmonicDistanceTerms.idx1, harmonic.idx1, atomOffset); + appendOffsetIndices(system.harmonicDistanceTerms.idx2, harmonic.idx2, atomOffset); + append(system.harmonicDistanceTerms.minLen, harmonic.minLen); + append(system.harmonicDistanceTerms.maxLen, harmonic.maxLen); + append(system.harmonicDistanceTerms.forceConstant, harmonic.forceConstant); + appendTermStart(system.harmonicDistanceTermStarts, harmonic.idx1.size()); + + const auto& angle = contribs.linearAngleTerms; + appendOffsetIndices(system.linearAngleTerms.idx1, angle.idx1, atomOffset); + appendOffsetIndices(system.linearAngleTerms.idx2, angle.idx2, atomOffset); + appendOffsetIndices(system.linearAngleTerms.idx3, angle.idx3, atomOffset); + append(system.linearAngleTerms.minAngle, angle.minAngle); + append(system.linearAngleTerms.maxAngle, angle.maxAngle); + appendTermStart(system.linearAngleTermStarts, angle.idx1.size()); + + const auto& torsion = contribs.experimentalTorsionTerms; + appendOffsetIndices(system.experimentalTorsionTerms.idx1, torsion.idx1, atomOffset); + appendOffsetIndices(system.experimentalTorsionTerms.idx2, torsion.idx2, atomOffset); + appendOffsetIndices(system.experimentalTorsionTerms.idx3, torsion.idx3, atomOffset); + appendOffsetIndices(system.experimentalTorsionTerms.idx4, torsion.idx4, atomOffset); + append(system.experimentalTorsionTerms.forceConstants, torsion.forceConstants); + append(system.experimentalTorsionTerms.signs, torsion.signs); + appendTermStart(system.experimentalTorsionTermStarts, torsion.idx1.size()); + + const auto& planarity = contribs.planarityTerms; + appendOffsetIndices(system.planarityTerms.idx1, planarity.idx1, atomOffset); + appendOffsetIndices(system.planarityTerms.idx2, planarity.idx2, atomOffset); + appendOffsetIndices(system.planarityTerms.idx3, planarity.idx3, atomOffset); + appendOffsetIndices(system.planarityTerms.idx4, planarity.idx4, atomOffset); + appendTermStart(system.planarityTermStarts, planarity.idx1.size()); +} + +void sendAllInOneSystemToDevice(const BatchedAllInOneSystemHost& systemHost, + BatchedAllInOneDeviceBuffers& systemDevice, + cudaStream_t stream) { + sendAllInOneSystemToDeviceImpl(systemHost, systemDevice, stream); +} +void sendAllInOneSystemToDevice(const BatchedAllInOneSystemHost& systemHost, + BatchedAllInOneDeviceBuffersSingle& systemDevice, + cudaStream_t stream) { + sendAllInOneSystemToDeviceImpl(systemHost, systemDevice, stream); +} + +AllInOneContribsDevicePtr toAllInOneContribsDevicePtr(const BatchedAllInOneDeviceBuffers& systemDevice, + const int* atomStarts) { + return toAllInOneContribsDevicePtrImpl(systemDevice, atomStarts); +} +AllInOneContribsDevicePtrSingle toAllInOneContribsDevicePtr(const BatchedAllInOneDeviceBuffersSingle& systemDevice, + const int* atomStarts) { + return toAllInOneContribsDevicePtrImpl(systemDevice, atomStarts); +} + } // namespace DistGeom } // namespace nvMolKit diff --git a/src/forcefields/dist_geom.h b/src/forcefields/dist_geom.h index 79e59a19..053cff9e 100644 --- a/src/forcefields/dist_geom.h +++ b/src/forcefields/dist_geom.h @@ -129,6 +129,31 @@ struct Energy3DForceContribsHost { DistanceConstraintContribTerms longRangeDistTerms; }; +// ------------------------------ +// All-in-one ETKDG (AIO) terms +// ------------------------------ + +//! RDKit's all-in-one out-of-plane terms, 0.5 * k * chi^2, where chi is the angle in degrees between the idx2-idx4 +//! bond and the plane of idx1, idx2 and idx3. The force constant is AllInOneForceWeights::planarity. +struct PlanarityContribTerms { + std::vector idx1; + std::vector idx2; // Central atom + std::vector idx3; + std::vector idx4; +}; + +//! Terms of RDKit's all-in-one ETKDG force field for one molecule. The distance-geometry, harmonic distance and linear +//! angle terms are active in both minimization phases; the experimental torsion and planarity terms only in the second. +struct AllInOneForceContribsHost { + EnergyForceContribsHost distanceGeometry; + //! Flat-bottom 1-2 and 1-3 distance terms, evaluated on 4D distances. + DistanceConstraintContribTerms harmonicDistanceTerms; + //! Flat-bottom terms holding sp centers linear. The force constant is AllInOneForceWeights::linearAngle. + AngleConstraintContribTerms linearAngleTerms; + TorsionAngleContribTerms experimentalTorsionTerms; + PlanarityContribTerms planarityTerms; +}; + struct BatchedIndicesHost { //! Defines the start of each molecule's energy buffer region that will be added to then reduced. std::vector energyBufferStarts = {0}; @@ -173,6 +198,20 @@ struct BatchedMolecularSystem3DHost { BatchedIndices3DHost indices; }; +//! Flattened all-in-one ETKDG batch. Atom indices are offset into the batch, and each *TermStarts vector (size +//! n_molecules + 1) gives the range of a molecule's terms. +struct BatchedAllInOneSystemHost { + BatchedMolecularSystemHost distanceGeometry; + DistanceConstraintContribTerms harmonicDistanceTerms; + AngleConstraintContribTerms linearAngleTerms; + TorsionAngleContribTerms experimentalTorsionTerms; + PlanarityContribTerms planarityTerms; + std::vector harmonicDistanceTermStarts = {0}; + std::vector linearAngleTermStarts = {0}; + std::vector experimentalTorsionTermStarts = {0}; + std::vector planarityTermStarts = {0}; +}; + template struct DistViolationContribTermsDeviceT { nvMolKit::AsyncDeviceVector idx1; nvMolKit::AsyncDeviceVector idx2; @@ -343,6 +382,29 @@ using BatchedMolecularDeviceBuffersSingle = BatchedMolecularDeviceBuffersT; using BatchedMolecular3DDeviceBuffersSingle = BatchedMolecular3DDeviceBuffersT; +struct PlanarityContribTermsDevice { + nvMolKit::AsyncDeviceVector idx1; + nvMolKit::AsyncDeviceVector idx2; + nvMolKit::AsyncDeviceVector idx3; + nvMolKit::AsyncDeviceVector idx4; +}; + +//! Device copy of a BatchedAllInOneSystemHost. +template struct BatchedAllInOneDeviceBuffersT { + BatchedMolecularDeviceBuffersT distanceGeometry; + DistanceConstraintContribTermsDeviceT harmonicDistanceTerms; + AngleConstraintContribTermsDeviceT linearAngleTerms; + TorsionAngleContribTermsDeviceT experimentalTorsionTerms; + PlanarityContribTermsDevice planarityTerms; + nvMolKit::AsyncDeviceVector harmonicDistanceTermStarts; + nvMolKit::AsyncDeviceVector linearAngleTermStarts; + nvMolKit::AsyncDeviceVector experimentalTorsionTermStarts; + nvMolKit::AsyncDeviceVector planarityTermStarts; +}; + +using BatchedAllInOneDeviceBuffers = BatchedAllInOneDeviceBuffersT; +using BatchedAllInOneDeviceBuffersSingle = BatchedAllInOneDeviceBuffersT; + //! Set all DeviceVector streams for the batched molecular device buffers. void setStreams(BatchedMolecularDeviceBuffers& devBuffers, cudaStream_t stream); void setStreams(BatchedMolecularDeviceBuffersSingle& devBuffers, cudaStream_t stream); @@ -604,6 +666,29 @@ cudaError_t computeGradBlockPerMolETK(BatchedMolecular3DDeviceBuffers& const nvMolKit::AsyncDeviceVector& ctxPositionsDevice, const uint8_t* activeThisStage = nullptr, cudaStream_t stream = nullptr); + +//! Add a molecule's all-in-one ETKDG terms to the batch. ctxAtomStarts must already include the molecule. +void addMoleculeToAllInOneSystem(const AllInOneForceContribsHost& contribs, + int numAtoms, + const std::vector& ctxAtomStarts, + BatchedAllInOneSystemHost& system, + BatchedForcefieldMetadata* metadata = nullptr, + int moleculeIdx = -1, + int conformerIdx = -1); + +//! Set the buffer streams and upload an all-in-one ETKDG batch. +void sendAllInOneSystemToDevice(const BatchedAllInOneSystemHost& systemHost, + BatchedAllInOneDeviceBuffers& systemDevice, + cudaStream_t stream); +void sendAllInOneSystemToDevice(const BatchedAllInOneSystemHost& systemHost, + BatchedAllInOneDeviceBuffersSingle& systemDevice, + cudaStream_t stream); + +//! Create the pointer struct used by the all-in-one ETKDG block-per-mol kernels. +AllInOneContribsDevicePtr toAllInOneContribsDevicePtr(const BatchedAllInOneDeviceBuffers& systemDevice, + const int* atomStarts); +AllInOneContribsDevicePtrSingle toAllInOneContribsDevicePtr(const BatchedAllInOneDeviceBuffersSingle& systemDevice, + const int* atomStarts); } // namespace DistGeom } // namespace nvMolKit diff --git a/src/forcefields/dist_geom_kernels.cu b/src/forcefields/dist_geom_kernels.cu index c8c49095..91408850 100644 --- a/src/forcefields/dist_geom_kernels.cu +++ b/src/forcefields/dist_geom_kernels.cu @@ -1123,28 +1123,14 @@ __global__ void combinedEnergiesKernel(const Terms* terms, } } -template -__global__ void combinedGradKernel(const Terms* terms, - const BatchedIndicesDevicePtr* systemIndices, - const real* coords, - real* grad, - const real chiralWeight, - const real fourthDimWeight, - const uint8_t* activeThisStage) { - const int molIdx = blockIdx.x; - const int tid = threadIdx.x; - - if (activeThisStage != nullptr && activeThisStage[molIdx] == 0) { - return; - } - - const int atomStart = systemIndices->atomStarts[molIdx]; - const int atomEnd = systemIndices->atomStarts[molIdx + 1]; - const int numAtoms = atomEnd - atomStart; - +//! Zeroes a molecule's gradient, accumulates it with accumulate(molGrad), then stores it in grad. Accumulation runs in +//! shared memory when the molecule fits and directly in global memory otherwise. +template +__device__ void accumulateMolGrad(real* grad, const int atomStart, const int numAtoms, Accumulate&& accumulate) { constexpr int maxAtomSize = 256; __shared__ real accumGrad[maxAtomSize * 4]; // Support up to 4D + const int tid = threadIdx.x; const bool useSharedMem = numAtoms * dimension <= maxAtomSize * 4; real* molGradBase = useSharedMem ? accumGrad : grad + atomStart * dimension; @@ -1153,15 +1139,7 @@ __global__ void combinedGradKernel(const Terms* terms, } __syncthreads(); - const real* molCoords = coords + atomStart * dimension; - molGradDGDispatch(*terms, - *systemIndices, - molCoords, - molGradBase, - molIdx, - chiralWeight, - fourthDimWeight, - tid); + accumulate(molGradBase); __syncthreads(); if (useSharedMem) { @@ -1172,6 +1150,36 @@ __global__ void combinedGradKernel(const Terms* terms, } } +template +__global__ void combinedGradKernel(const Terms* terms, + const BatchedIndicesDevicePtr* systemIndices, + const real* coords, + real* grad, + const real chiralWeight, + const real fourthDimWeight, + const uint8_t* activeThisStage) { + const int molIdx = blockIdx.x; + const int tid = threadIdx.x; + + if (activeThisStage != nullptr && activeThisStage[molIdx] == 0) { + return; + } + + const int atomStart = systemIndices->atomStarts[molIdx]; + const int numAtoms = systemIndices->atomStarts[molIdx + 1] - atomStart; + const real* molCoords = coords + atomStart * dimension; + accumulateMolGrad(grad, atomStart, numAtoms, [&](real* molGrad) { + molGradDGDispatch(*terms, + *systemIndices, + molCoords, + molGrad, + molIdx, + chiralWeight, + fourthDimWeight, + tid); + }); +} + template cudaError_t launchBlockPerMolEnergyKernelImpl(int numMols, const Terms& terms, @@ -1491,6 +1499,201 @@ cudaError_t launchBlockPerMolGradKernelETK(int return launchBlockPerMolGradKernelETKImpl(numMols, terms, systemIndices, coords, grad, activeThisStage, stream); } +// All-in-one ETKDG (AIO) combined kernels +__device__ double molEnergyAIODispatch(const AllInOneContribsDevicePtr& terms, + const double* coords, + int molIdx, + const AllInOneForceWeights& weights, + bool includeTorsionTerms, + int tid) { + return fp64::molEnergyAIO(terms, coords, molIdx, weights, includeTorsionTerms, tid); +} +__device__ float molEnergyAIODispatch(const AllInOneContribsDevicePtrSingle& terms, + const float* coords, + int molIdx, + const AllInOneForceWeights& weights, + bool includeTorsionTerms, + int tid) { + return fp32::molEnergyAIO(terms, coords, molIdx, weights, includeTorsionTerms, tid); +} + +__device__ void molGradAIODispatch(const AllInOneContribsDevicePtr& terms, + const double* coords, + double* grad, + int molIdx, + const AllInOneForceWeights& weights, + bool includeTorsionTerms, + int tid) { + fp64::molGradAIO(terms, coords, grad, molIdx, weights, includeTorsionTerms, tid); +} +__device__ void molGradAIODispatch(const AllInOneContribsDevicePtrSingle& terms, + const float* coords, + float* grad, + int molIdx, + const AllInOneForceWeights& weights, + bool includeTorsionTerms, + int tid) { + fp32::molGradAIO(terms, coords, grad, molIdx, weights, includeTorsionTerms, tid); +} + +template +__global__ void combinedEnergiesKernelAIO(const Terms terms, + const real* coords, + real* energies, + const AllInOneForceWeights weights, + const bool includeTorsionTerms, + const uint8_t* activeThisStage) { + const int molIdx = blockIdx.x; + const int tid = threadIdx.x; + + if (activeThisStage != nullptr && activeThisStage[molIdx] == 0) { + if (tid == 0) { + energies[molIdx] = 0.0; + } + return; + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tempStorage; + + const int atomStart = terms.distanceGeometryIndices.atomStarts[molIdx]; + const real threadEnergy = + molEnergyAIODispatch(terms, coords + atomStart * 4, molIdx, weights, includeTorsionTerms, tid); + const reduceT blockEnergy = BlockReduce(tempStorage).Sum(static_cast(threadEnergy)); + + if (tid == 0) { + energies[molIdx] = static_cast(blockEnergy); + } +} + +template +__global__ void combinedGradKernelAIO(const Terms terms, + const real* coords, + real* grad, + const AllInOneForceWeights weights, + const bool includeTorsionTerms, + const uint8_t* activeThisStage) { + const int molIdx = blockIdx.x; + const int tid = threadIdx.x; + + if (activeThisStage != nullptr && activeThisStage[molIdx] == 0) { + return; + } + + const int atomStart = terms.distanceGeometryIndices.atomStarts[molIdx]; + const int numAtoms = terms.distanceGeometryIndices.atomStarts[molIdx + 1] - atomStart; + const real* molCoords = coords + atomStart * 4; + accumulateMolGrad<4>(grad, atomStart, numAtoms, [&](real* molGrad) { + molGradAIODispatch(terms, molCoords, molGrad, molIdx, weights, includeTorsionTerms, tid); + }); +} + +template +cudaError_t launchBlockPerMolEnergyKernelAIOImpl(const int numMols, + const Terms& terms, + const real* coords, + real* energies, + const AllInOneForceWeights& weights, + const bool includeTorsionTerms, + const uint8_t* activeThisStage, + cudaStream_t stream) { + if (numMols == 0) { + return cudaSuccess; + } + combinedEnergiesKernelAIO + <<>>(terms, coords, energies, weights, includeTorsionTerms, activeThisStage); + return cudaGetLastError(); +} + +template +cudaError_t launchBlockPerMolGradKernelAIOImpl(const int numMols, + const Terms& terms, + const real* coords, + real* grad, + const AllInOneForceWeights& weights, + const bool includeTorsionTerms, + const uint8_t* activeThisStage, + cudaStream_t stream) { + if (numMols == 0) { + return cudaSuccess; + } + combinedGradKernelAIO + <<>>(terms, coords, grad, weights, includeTorsionTerms, activeThisStage); + return cudaGetLastError(); +} + +cudaError_t launchBlockPerMolEnergyKernelAIO(const int numMols, + const AllInOneContribsDevicePtr& terms, + const double* coords, + double* energies, + const AllInOneForceWeights& weights, + const bool includeTorsionTerms, + const uint8_t* activeThisStage, + cudaStream_t stream) { + return launchBlockPerMolEnergyKernelAIOImpl(numMols, + terms, + coords, + energies, + weights, + includeTorsionTerms, + activeThisStage, + stream); +} + +cudaError_t launchBlockPerMolEnergyKernelAIO(const int numMols, + const AllInOneContribsDevicePtrSingle& terms, + const float* coords, + float* energies, + const AllInOneForceWeights& weights, + const bool includeTorsionTerms, + const uint8_t* activeThisStage, + cudaStream_t stream) { + return launchBlockPerMolEnergyKernelAIOImpl(numMols, + terms, + coords, + energies, + weights, + includeTorsionTerms, + activeThisStage, + stream); +} + +cudaError_t launchBlockPerMolGradKernelAIO(const int numMols, + const AllInOneContribsDevicePtr& terms, + const double* coords, + double* grad, + const AllInOneForceWeights& weights, + const bool includeTorsionTerms, + const uint8_t* activeThisStage, + cudaStream_t stream) { + return launchBlockPerMolGradKernelAIOImpl(numMols, + terms, + coords, + grad, + weights, + includeTorsionTerms, + activeThisStage, + stream); +} + +cudaError_t launchBlockPerMolGradKernelAIO(const int numMols, + const AllInOneContribsDevicePtrSingle& terms, + const float* coords, + float* grad, + const AllInOneForceWeights& weights, + const bool includeTorsionTerms, + const uint8_t* activeThisStage, + cudaStream_t stream) { + return launchBlockPerMolGradKernelAIOImpl(numMols, + terms, + coords, + grad, + weights, + includeTorsionTerms, + activeThisStage, + stream); +} + template __global__ void planarEnergyKernelETK(const Terms* terms, const BatchedIndices3DDevicePtr* indices, diff --git a/src/forcefields/dist_geom_kernels.h b/src/forcefields/dist_geom_kernels.h index 6999a6f3..82125100 100644 --- a/src/forcefields/dist_geom_kernels.h +++ b/src/forcefields/dist_geom_kernels.h @@ -428,6 +428,79 @@ cudaError_t launchPlanarEnergyKernelETK(int float* energies, const uint8_t* activeThisStage = nullptr, cudaStream_t stream = 0); + +// All-in-one ETKDG (AIO) block-per-molecule kernels + +struct PlanarityContribTermsDevicePtr { + const int* idx1; + const int* idx2; + const int* idx3; + const int* idx4; +}; + +//! Device view of an all-in-one ETKDG batch: the 4D distance-geometry terms plus the all-in-one terms, each with +//! per-molecule CSR offsets. +template struct AllInOneContribsDevicePtrT { + EnergyForceContribsDevicePtrT distanceGeometry; + BatchedIndicesDevicePtr distanceGeometryIndices; + DistanceConstraintContribTermsDevicePtrT harmonicDistanceTerms; + AngleConstraintContribTermsDevicePtrT linearAngleTerms; + TorsionAngleContribTermsDevicePtrT experimentalTorsionTerms; + PlanarityContribTermsDevicePtr planarityTerms; + const int* harmonicDistanceTermStarts; + const int* linearAngleTermStarts; + const int* experimentalTorsionTermStarts; + const int* planarityTermStarts; +}; + +using AllInOneContribsDevicePtr = AllInOneContribsDevicePtrT; +using AllInOneContribsDevicePtrSingle = AllInOneContribsDevicePtrT; + +//! Force constants shared by every system in an all-in-one ETKDG batch. +struct AllInOneForceWeights { + double chiral; + double fourthDim; + double linearAngle; + double planarity; +}; + +//! Computes all-in-one ETKDG energies, one block per molecule. Experimental torsion and planarity terms are included +//! only when includeTorsionTerms is set, matching RDKit's second minimization phase. Inactive systems get zero energy. +cudaError_t launchBlockPerMolEnergyKernelAIO(int numMols, + const AllInOneContribsDevicePtr& terms, + const double* coords, + double* energies, + const AllInOneForceWeights& weights, + bool includeTorsionTerms, + const uint8_t* activeThisStage = nullptr, + cudaStream_t stream = 0); +cudaError_t launchBlockPerMolEnergyKernelAIO(int numMols, + const AllInOneContribsDevicePtrSingle& terms, + const float* coords, + float* energies, + const AllInOneForceWeights& weights, + bool includeTorsionTerms, + const uint8_t* activeThisStage = nullptr, + cudaStream_t stream = 0); + +//! Computes all-in-one ETKDG gradients, one block per molecule. Gradients of active systems are overwritten; inactive +//! systems are left untouched. +cudaError_t launchBlockPerMolGradKernelAIO(int numMols, + const AllInOneContribsDevicePtr& terms, + const double* coords, + double* grad, + const AllInOneForceWeights& weights, + bool includeTorsionTerms, + const uint8_t* activeThisStage = nullptr, + cudaStream_t stream = 0); +cudaError_t launchBlockPerMolGradKernelAIO(int numMols, + const AllInOneContribsDevicePtrSingle& terms, + const float* coords, + float* grad, + const AllInOneForceWeights& weights, + bool includeTorsionTerms, + const uint8_t* activeThisStage = nullptr, + cudaStream_t stream = 0); } // namespace DistGeom } // namespace nvMolKit diff --git a/src/forcefields/dist_geom_kernels_device.cuh b/src/forcefields/dist_geom_kernels_device.cuh index d42ac398..c2041ebb 100644 --- a/src/forcefields/dist_geom_kernels_device.cuh +++ b/src/forcefields/dist_geom_kernels_device.cuh @@ -357,6 +357,8 @@ static __device__ __forceinline__ DG_REAL inversionEnergy(const DG_REAL* pos, return forceConstant * (C0 + C1 * sinY + C2 * cos2W); } +//! The distance is taken over distanceDim coordinates; the gradient only moves xyz, as in RDKit. +template static __device__ __forceinline__ DG_REAL distanceConstraintEnergy(const DG_REAL* pos, const int idx1, const int idx2, @@ -365,7 +367,7 @@ static __device__ __forceinline__ DG_REAL distanceConstraintEnergy(const DG_REAL const DG_REAL forceConstant) { const int posIdx1 = idx1 * 4; const int posIdx2 = idx2 * 4; - const DG_REAL distance2 = distanceSquaredPosIdx(pos, posIdx1, posIdx2, 3); + const DG_REAL distance2 = distanceSquaredPosIdx(pos, posIdx1, posIdx2); const DG_REAL minLen2 = minLen * minLen; const DG_REAL maxLen2 = maxLen * maxLen; @@ -688,6 +690,7 @@ static __device__ __forceinline__ void inversionGrad(const DG_REAL* pos, atomicAdd(&grad[posIdx4 + 2], dE_dW * tg4z); } +template static __device__ __forceinline__ void distanceConstraintGrad(const DG_REAL* pos, const int idx1, const int idx2, @@ -701,7 +704,7 @@ static __device__ __forceinline__ void distanceConstraintGrad(const DG_REAL* pos const int posIdx2 = idx2 * 4; // Calculate squared distance - const DG_REAL distance2 = distanceSquaredPosIdx(pos, posIdx1, posIdx2, 3); + const DG_REAL distance2 = distanceSquaredPosIdx(pos, posIdx1, posIdx2); // Check if distance is outside bounds and calculate preFactor DG_REAL preFactor; @@ -1338,3 +1341,239 @@ static __device__ __inline__ void molGradETK(const Energy3DForceContribsDevicePt } } } + +// ---------------------------- +// All-in-one ETKDG (AIO) terms +// ---------------------------- + +//! Unit bond vectors from the central atom to atoms 1, 3 and 4 of a planarity term, their inverse lengths, and the +//! sine of the out-of-plane angle between the bond to atom 4 and the plane normal (-rJI) x rJK. Returns false when a +//! bond or the plane normal is degenerate. +static __device__ __forceinline__ bool planarityGeometry(const DG_REAL* pos, + const int posIdx1, + const int posIdx2, + const int posIdx3, + const int posIdx4, + DG_REAL (&rJI)[3], + DG_REAL (&rJK)[3], + DG_REAL (&rJL)[3], + DG_REAL& invJI, + DG_REAL& invJK, + DG_REAL& invJL, + DG_REAL& sinChi) { + for (int i = 0; i < 3; ++i) { + rJI[i] = pos[posIdx1 + i] - pos[posIdx2 + i]; + rJK[i] = pos[posIdx3 + i] - pos[posIdx2 + i]; + rJL[i] = pos[posIdx4 + i] - pos[posIdx2 + i]; + } + const DG_REAL dJI2 = dotProduct(rJI[0], rJI[1], rJI[2], rJI[0], rJI[1], rJI[2]); + const DG_REAL dJK2 = dotProduct(rJK[0], rJK[1], rJK[2], rJK[0], rJK[1], rJK[2]); + const DG_REAL dJL2 = dotProduct(rJL[0], rJL[1], rJL[2], rJL[0], rJL[1], rJL[2]); + if (DG_IS_ZERO(dJI2) || DG_IS_ZERO(dJK2) || DG_IS_ZERO(dJL2)) { + return false; + } + invJI = DG_RSQRT(dJI2); + invJK = DG_RSQRT(dJK2); + invJL = DG_RSQRT(dJL2); + for (int i = 0; i < 3; ++i) { + rJI[i] *= invJI; + rJK[i] *= invJK; + rJL[i] *= invJL; + } + + DG_REAL nx, ny, nz; + crossProduct(-rJI[0], -rJI[1], -rJI[2], rJK[0], rJK[1], rJK[2], nx, ny, nz); + const DG_REAL normalLength2 = nx * nx + ny * ny + nz * nz; + if (DG_IS_ZERO(normalLength2)) { + return false; + } + sinChi = clamp(dotProduct(rJL[0], rJL[1], rJL[2], nx, ny, nz) * DG_RSQRT(normalLength2), -DG_REAL(1.0), DG_REAL(1.0)); + return true; +} + +static __device__ __forceinline__ DG_REAL planarityEnergy(const DG_REAL* pos, + const int idx1, + const int idx2, + const int idx3, + const int idx4, + const DG_REAL forceConstant) { + DG_REAL rJI[3], rJK[3], rJL[3], invJI, invJK, invJL, sinChi; + if (!planarityGeometry(pos, idx1 * 4, idx2 * 4, idx3 * 4, idx4 * 4, rJI, rJK, rJL, invJI, invJK, invJL, sinChi)) { + return DG_REAL(0.0); + } + const DG_REAL chi = RAD2DEG * DG_ASIN(sinChi); + return DG_REAL(0.5) * forceConstant * chi * chi; +} + +static __device__ __forceinline__ void planarityGrad(const DG_REAL* pos, + const int idx1, + const int idx2, + const int idx3, + const int idx4, + const DG_REAL forceConstant, + DG_REAL* grad) { + const int posIdx1 = idx1 * 4; + const int posIdx2 = idx2 * 4; + const int posIdx3 = idx3 * 4; + const int posIdx4 = idx4 * 4; + + DG_REAL rJI[3], rJK[3], rJL[3], invJI, invJK, invJL, sinChi; + if (!planarityGeometry(pos, posIdx1, posIdx2, posIdx3, posIdx4, rJI, rJK, rJL, invJI, invJK, invJL, sinChi)) { + return; + } + const DG_REAL cosChi = DG_FMAX(DG_SQRT(DG_FMAX(DG_REAL(1.0) - sinChi * sinChi, DG_REAL(0.0))), DG_REAL(1.0e-8)); + const DG_REAL dEdChi = forceConstant * RAD2DEG * DG_ASIN(sinChi) * RAD2DEG; + const DG_REAL cosTheta = + clamp(dotProduct(rJI[0], rJI[1], rJI[2], rJK[0], rJK[1], rJK[2]), -DG_REAL(1.0), DG_REAL(1.0)); + const DG_REAL sinThetaSq = DG_FMAX(DG_REAL(1.0) - cosTheta * cosTheta, DG_REAL(1.0e-8)); + const DG_REAL sinTheta = DG_FMAX(DG_SQRT(sinThetaSq), DG_REAL(1.0e-8)); + + DG_REAL t1[3], t2[3], t3[3]; + crossProduct(rJL[0], rJL[1], rJL[2], rJK[0], rJK[1], rJK[2], t1[0], t1[1], t1[2]); + crossProduct(rJI[0], rJI[1], rJI[2], rJL[0], rJL[1], rJL[2], t2[0], t2[1], t2[2]); + crossProduct(rJK[0], rJK[1], rJK[2], rJI[0], rJI[1], rJI[2], t3[0], t3[1], t3[2]); + + const DG_REAL inverseTerm1 = DG_REAL(1.0) / (cosChi * sinTheta); + const DG_REAL term2 = sinChi / (cosChi * sinThetaSq); + const DG_REAL tanChi = sinChi / cosChi; + for (int i = 0; i < 3; ++i) { + const DG_REAL tg1 = dEdChi * (t1[i] * inverseTerm1 - (rJI[i] - rJK[i] * cosTheta) * term2) * invJI; + const DG_REAL tg3 = dEdChi * (t2[i] * inverseTerm1 - (rJK[i] - rJI[i] * cosTheta) * term2) * invJK; + const DG_REAL tg4 = dEdChi * (t3[i] * inverseTerm1 - rJL[i] * tanChi) * invJL; + atomicAdd(&grad[posIdx1 + i], tg1); + atomicAdd(&grad[posIdx2 + i], -(tg1 + tg3 + tg4)); + atomicAdd(&grad[posIdx3 + i], tg3); + atomicAdd(&grad[posIdx4 + i], tg4); + } +} + +//! Per-thread share of a molecule's all-in-one ETKDG energy. molCoords points at the molecule's first atom. +static __device__ __inline__ DG_REAL molEnergyAIO(const AllInOneContribsDevicePtr& terms, + const DG_REAL* molCoords, + const int molIdx, + const AllInOneForceWeights& weights, + const bool includeTorsionTerms, + const int tid) { + DG_REAL energy = molEnergyDG<4>(terms.distanceGeometry, + terms.distanceGeometryIndices, + molCoords, + molIdx, + static_cast(weights.chiral), + static_cast(weights.fourthDim), + tid); + const int atomStart = terms.distanceGeometryIndices.atomStarts[molIdx]; + + const auto& harmonic = terms.harmonicDistanceTerms; + for (int i = terms.harmonicDistanceTermStarts[molIdx] + tid; i < terms.harmonicDistanceTermStarts[molIdx + 1]; + i += blockDim.x) { + energy += distanceConstraintEnergy<4>(molCoords, + harmonic.idx1[i] - atomStart, + harmonic.idx2[i] - atomStart, + harmonic.minLen[i], + harmonic.maxLen[i], + harmonic.forceConstant[i]); + } + const auto& angle = terms.linearAngleTerms; + for (int i = terms.linearAngleTermStarts[molIdx] + tid; i < terms.linearAngleTermStarts[molIdx + 1]; + i += blockDim.x) { + energy += angleConstraintEnergy(molCoords, + angle.idx1[i] - atomStart, + angle.idx2[i] - atomStart, + angle.idx3[i] - atomStart, + angle.minAngle[i], + angle.maxAngle[i], + static_cast(weights.linearAngle)); + } + if (!includeTorsionTerms) { + return energy; + } + const auto& torsion = terms.experimentalTorsionTerms; + for (int i = terms.experimentalTorsionTermStarts[molIdx] + tid; i < terms.experimentalTorsionTermStarts[molIdx + 1]; + i += blockDim.x) { + energy += torsionAngleEnergy(molCoords, + torsion.idx1[i] - atomStart, + torsion.idx2[i] - atomStart, + torsion.idx3[i] - atomStart, + torsion.idx4[i] - atomStart, + &torsion.forceConstants[i * 6], + &torsion.signs[i * 6]); + } + const auto& planarity = terms.planarityTerms; + for (int i = terms.planarityTermStarts[molIdx] + tid; i < terms.planarityTermStarts[molIdx + 1]; i += blockDim.x) { + energy += planarityEnergy(molCoords, + planarity.idx1[i] - atomStart, + planarity.idx2[i] - atomStart, + planarity.idx3[i] - atomStart, + planarity.idx4[i] - atomStart, + static_cast(weights.planarity)); + } + return energy; +} + +//! Accumulates a molecule's all-in-one ETKDG gradient into molGrad, which points at the molecule's first atom. +static __device__ __inline__ void molGradAIO(const AllInOneContribsDevicePtr& terms, + const DG_REAL* molCoords, + DG_REAL* molGrad, + const int molIdx, + const AllInOneForceWeights& weights, + const bool includeTorsionTerms, + const int tid) { + molGradDG<4>(terms.distanceGeometry, + terms.distanceGeometryIndices, + molCoords, + molGrad, + molIdx, + static_cast(weights.chiral), + static_cast(weights.fourthDim), + tid); + const int atomStart = terms.distanceGeometryIndices.atomStarts[molIdx]; + + const auto& harmonic = terms.harmonicDistanceTerms; + for (int i = terms.harmonicDistanceTermStarts[molIdx] + tid; i < terms.harmonicDistanceTermStarts[molIdx + 1]; + i += blockDim.x) { + distanceConstraintGrad<4>(molCoords, + harmonic.idx1[i] - atomStart, + harmonic.idx2[i] - atomStart, + harmonic.minLen[i], + harmonic.maxLen[i], + harmonic.forceConstant[i], + molGrad); + } + const auto& angle = terms.linearAngleTerms; + for (int i = terms.linearAngleTermStarts[molIdx] + tid; i < terms.linearAngleTermStarts[molIdx + 1]; + i += blockDim.x) { + angleConstraintGrad(molCoords, + angle.idx1[i] - atomStart, + angle.idx2[i] - atomStart, + angle.idx3[i] - atomStart, + angle.minAngle[i], + angle.maxAngle[i], + static_cast(weights.linearAngle), + molGrad); + } + if (!includeTorsionTerms) { + return; + } + const auto& torsion = terms.experimentalTorsionTerms; + for (int i = terms.experimentalTorsionTermStarts[molIdx] + tid; i < terms.experimentalTorsionTermStarts[molIdx + 1]; + i += blockDim.x) { + torsionAngleGrad(molCoords, + torsion.idx1[i] - atomStart, + torsion.idx2[i] - atomStart, + torsion.idx3[i] - atomStart, + torsion.idx4[i] - atomStart, + &torsion.forceConstants[i * 6], + &torsion.signs[i * 6], + molGrad); + } + const auto& planarity = terms.planarityTerms; + for (int i = terms.planarityTermStarts[molIdx] + tid; i < terms.planarityTermStarts[molIdx + 1]; i += blockDim.x) { + planarityGrad(molCoords, + planarity.idx1[i] - atomStart, + planarity.idx2[i] - atomStart, + planarity.idx3[i] - atomStart, + planarity.idx4[i] - atomStart, + static_cast(weights.planarity), + molGrad); + } +} diff --git a/src/forcefields/dist_geom_kernels_device_dispatch.cuh b/src/forcefields/dist_geom_kernels_device_dispatch.cuh index 91bd9914..061d74e3 100644 --- a/src/forcefields/dist_geom_kernels_device_dispatch.cuh +++ b/src/forcefields/dist_geom_kernels_device_dispatch.cuh @@ -20,6 +20,7 @@ namespace DistGeom { namespace fp64 { using EnergyForceContribsDevicePtr = DistGeom::EnergyForceContribsDevicePtr; using Energy3DForceContribsDevicePtr = DistGeom::Energy3DForceContribsDevicePtr; +using AllInOneContribsDevicePtr = DistGeom::AllInOneContribsDevicePtr; constexpr double RAD2DEG = 180.0 / M_PI; #define DG_REAL double #define DG_PARAM_REAL double @@ -28,6 +29,7 @@ constexpr double RAD2DEG = 180.0 / M_PI; #define DG_SQRTF sqrtf #define DG_SQRT sqrt #define DG_ACOS acos +#define DG_ASIN asin #define DG_FMAX fmax #define DG_IS_ZERO isDoubleZero #include "src/forcefields/dist_geom_kernels_device.cuh" @@ -38,6 +40,7 @@ constexpr double RAD2DEG = 180.0 / M_PI; #undef DG_SQRTF #undef DG_SQRT #undef DG_ACOS +#undef DG_ASIN #undef DG_FMAX #undef DG_IS_ZERO } // namespace fp64 @@ -45,6 +48,7 @@ constexpr double RAD2DEG = 180.0 / M_PI; namespace fp32 { using EnergyForceContribsDevicePtr = DistGeom::EnergyForceContribsDevicePtrSingle; using Energy3DForceContribsDevicePtr = DistGeom::Energy3DForceContribsDevicePtrSingle; +using AllInOneContribsDevicePtr = DistGeom::AllInOneContribsDevicePtrSingle; constexpr float RAD2DEG = 180.0f / static_cast(M_PI); #define DG_REAL float #define DG_PARAM_REAL float @@ -53,6 +57,7 @@ constexpr float RAD2DEG = 180.0f / static_cast(M_PI); #define DG_SQRTF sqrtf #define DG_SQRT sqrtf #define DG_ACOS acosf +#define DG_ASIN asinf #define DG_FMAX fmaxf #define DG_IS_ZERO isFloatZero #include "src/forcefields/dist_geom_kernels_device.cuh" @@ -63,6 +68,7 @@ constexpr float RAD2DEG = 180.0f / static_cast(M_PI); #undef DG_SQRTF #undef DG_SQRT #undef DG_ACOS +#undef DG_ASIN #undef DG_FMAX #undef DG_IS_ZERO } // namespace fp32 diff --git a/src/minimizer/bfgs_types.h b/src/minimizer/bfgs_types.h index b550e7a5..fdb1e7f4 100644 --- a/src/minimizer/bfgs_types.h +++ b/src/minimizer/bfgs_types.h @@ -20,10 +20,11 @@ namespace nvMolKit { /// Force field type enum for BFGS minimization enum class ForceFieldType { - MMFF = 0, - ETK = 1, - DG = 2, - UFF = 3 + MMFF = 0, + ETK = 1, + DG = 2, + UFF = 3, + AIO_ETKDG = 4 }; /// Debug level for BFGS minimization diff --git a/src/versions.h.in b/src/versions.h.in index 844963d2..a8df7170 100644 --- a/src/versions.h.in +++ b/src/versions.h.in @@ -13,4 +13,12 @@ // versions. Consumers select the codepath with `#if RDKIT_NEW_FLAG_API`. #define RDKIT_NEW_FLAG_API (RDKIT_VERSION_MAJOR > 2024 || (RDKIT_VERSION_MAJOR == 2024 && RDKIT_VERSION_MINOR >= 9)) +// RDKit's all-in-one ETKDG coordinate refinement was released in 2026.03.4. +// Keep this as a preprocessor macro because EmbedParameters::useLegacyImplementation +// and CrystalFFDetails::forceConsts do not exist in older supported RDKit headers. +#define RDKIT_AIO_ETKDG_API \ + (RDKIT_VERSION_MAJOR > 2026 || \ + (RDKIT_VERSION_MAJOR == 2026 && \ + (RDKIT_VERSION_MINOR > 3 || (RDKIT_VERSION_MINOR == 3 && RDKIT_VERSION_PATCH >= 4)))) + #endif // NVMOLKIT_VERSIONS_H \ No newline at end of file diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 82662a86..94271d01 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -196,6 +196,18 @@ target_link_libraries( test_utils embedder_utils) +add_executable(test_aio_etkdg_forcefield test_aio_etkdg_forcefield.cu) +target_link_libraries( + test_aio_etkdg_forcefield + PRIVATE ${RDKit_LIBS} + rdkit_dist_geom_flattened + aio_etkdg_batched_forcefield + dist_geom + device_vector + test_utils + embedder_utils + nvmolkit_versions) + add_executable(test_rdkit_bounds_matrix test_rdkit_bounds_matrix.cpp) target_link_libraries( test_rdkit_bounds_matrix PRIVATE rdkit_bounds_matrix symmetricEigenSolver @@ -466,6 +478,7 @@ set(TEST_LIST test_device_timings test_device_vector test_distgeom_kernels + test_aio_etkdg_forcefield test_etkdg test_etkdg_device_output test_ff_device_output diff --git a/tests/test_aio_etkdg_forcefield.cu b/tests/test_aio_etkdg_forcefield.cu new file mode 100644 index 00000000..d25d25dd --- /dev/null +++ b/tests/test_aio_etkdg_forcefield.cu @@ -0,0 +1,651 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include "versions.h" + +#if RDKIT_AIO_ETKDG_API + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "rdkit_extensions/dist_geom_flattened_builder.h" +#include "src/embedder_utils.h" +#include "src/forcefields/aio_etkdg_batched_forcefield.h" +#include "src/forcefields/dist_geom.h" +#include "src/forcefields/dist_geom_kernels.h" +#include "tests/test_utils.h" + +namespace { + +using nvMolKit::AsyncDeviceVector; +using nvMolKit::PrecisionMode; +namespace DG = nvMolKit::DistGeom; + +//! Molecules with linear, planar and chiral centers in addition to the SDF test sets. +const std::vector kCoverageSmiles = + {"CC#N", "C#CC(=O)NC", "c1ccccc1C(=O)O", "C[C@H](O)F", "F/C=C/C#CC", "CC(C)(C)C(=O)N[C@@H](C)C(=O)O"}; + +//! One molecule prepared for the all-in-one force field: RDKit embedder state and perturbed 4D coordinates. +struct AllInOneMolecule { + nvMolKit::detail::EmbedArgs eargs; + std::vector> pointOwners; + RDGeom::PointPtrVect points; + std::vector positions; +}; + +std::vector> coverageMolecules() { + std::vector> mols; + for (const auto& smiles : kCoverageSmiles) { + std::unique_ptr mol(RDKit::SmilesToMol(smiles)); + mols.emplace_back(RDKit::MolOps::addHs(*mol)); + } + return mols; +} + +std::vector> sdfMolecules(const std::string& fileName, const int count) { + std::vector> mols; + getMols(getTestDataFolderPath() + "/" + fileName, mols, count); + return mols; +} + +//! Prepares each molecule with the all-in-one force constants. Coordinates start from the input conformer when there +//! is one, are perturbed, and get a nonzero fourth dimension. +std::vector prepareMolecules(const std::vector>& mols, + const RDKit::DGeomHelpers::EmbedParameters& params) { + std::mt19937 generator(1729); + std::uniform_real_distribution perturbation(-0.3, 0.3); + std::uniform_real_distribution randomCoordinate(-1.5, 1.5); + + std::vector prepared(mols.size()); + for (std::size_t molIdx = 0; molIdx < mols.size(); ++molIdx) { + auto& mol = *mols[molIdx]; + auto& molecule = prepared[molIdx]; + molecule.eargs.dim = 4; + EXPECT_TRUE(nvMolKit::DGeomHelpers::prepareEmbedderArgs(mol, params, molecule.eargs)); + molecule.eargs.etkdgDetails.forceConsts = ForceFields::CrystalFF::ETKDGForceConsts::AIO::Cosine; + molecule.eargs.etkdgDetails.distMat = RDKit::MolOps::getDistanceMat(mol); + + for (unsigned int atomIdx = 0; atomIdx < mol.getNumAtoms(); ++atomIdx) { + auto& point = *molecule.pointOwners.emplace_back(std::make_unique(4)); + for (int dim = 0; dim < 3; ++dim) { + point[dim] = mol.getNumConformers() > 0 ? + mol.getConformer().getAtomPos(atomIdx)[dim] + perturbation(generator) : + randomCoordinate(generator); + } + point[3] = perturbation(generator); + molecule.points.push_back(&point); + for (int dim = 0; dim < 4; ++dim) { + molecule.positions.push_back(point[dim]); + } + } + } + return prepared; +} + +DG::AllInOneForceContribsHost buildContribs(const AllInOneMolecule& molecule, + const RDKit::DGeomHelpers::EmbedParameters& params) { + return DG::constructAllInOneForceFieldContribs(molecule.eargs.dim, + *molecule.eargs.mmat, + molecule.eargs.chiralCenters, + molecule.eargs.etkdgDetails, + params.useExpTorsionAnglePrefs, + params.useBasicKnowledge); +} + +//! RDKit's all-in-one force field for the given minimization phase, built as in RDKit's minimizeAllInOne. +std::unique_ptr referenceForcefield(AllInOneMolecule& molecule, + const RDKit::DGeomHelpers::EmbedParameters& params, + const bool torsionTermsEnabled) { + std::unique_ptr field( + ::DistGeom::constructAllInOneForceField(*molecule.eargs.mmat, + molecule.points, + molecule.eargs.etkdgDetails, + &molecule.eargs.chiralCenters)); + if (torsionTermsEnabled) { + ::DistGeom::addTorsionTerms(field.get(), + molecule.eargs.etkdgDetails, + params.useBasicKnowledge, + params.useExpTorsionAnglePrefs); + } + field->initialize(); + return field; +} + +//! RDKit 2026.09's TorsionAngleContribs::getGrad mis-evaluates the six-fold term (fixed upstream in rdkit/rdkit#9699), +//! so the reference torsion gradient is a central difference of RDKit's torsion energy instead. +void replaceTorsionGradient(const ForceFields::ForceField& field, + std::vector& positions, + std::vector& gradient) { + constexpr double step = 1e-6; + for (const auto& contrib : field.contribs()) { + const auto* torsions = dynamic_cast(contrib.get()); + if (torsions == nullptr) { + continue; + } + std::vector rdkitGradient(gradient.size(), 0.0); + torsions->getGrad(positions.data(), rdkitGradient.data()); + for (std::size_t coord = 0; coord < positions.size(); ++coord) { + const double original = positions[coord]; + positions[coord] = original + step; + const double plus = torsions->getEnergy(positions.data()); + positions[coord] = original - step; + const double minus = torsions->getEnergy(positions.data()); + positions[coord] = original; + gradient[coord] += (plus - minus) / (2.0 * step) - rdkitGradient[coord]; + } + } +} + +struct Evaluation { + std::vector energies; + std::vector> gradients; +}; + +Evaluation evaluateReference(std::vector& molecules, + const RDKit::DGeomHelpers::EmbedParameters& params, + const bool torsionTermsEnabled, + const std::vector& active, + const std::function& filterContribs = {}) { + Evaluation reference; + for (std::size_t molIdx = 0; molIdx < molecules.size(); ++molIdx) { + auto& molecule = molecules[molIdx]; + std::vector gradient(molecule.positions.size(), 0.0); + double energy = 0.0; + if (active[molIdx]) { + auto field = referenceForcefield(molecule, params, torsionTermsEnabled); + if (filterContribs) { + filterContribs(*field); + } + energy = field->calcEnergy(molecule.positions.data()); + field->calcGrad(molecule.positions.data(), gradient.data()); + replaceTorsionGradient(*field, molecule.positions, gradient); + } + reference.energies.push_back(energy); + reference.gradients.push_back(std::move(gradient)); + } + return reference; +} + +struct Batch { + DG::BatchedAllInOneSystemHost system; + std::vector atomStarts = {0}; + std::vector positions; +}; + +Batch buildBatch(const std::vector& molecules, + const std::vector& contribs) { + Batch batch; + for (std::size_t molIdx = 0; molIdx < molecules.size(); ++molIdx) { + const auto& positions = molecules[molIdx].positions; + const int numAtoms = static_cast(positions.size() / 4); + batch.atomStarts.push_back(batch.atomStarts.back() + numAtoms); + batch.positions.insert(batch.positions.end(), positions.begin(), positions.end()); + DG::addMoleculeToAllInOneSystem(contribs[molIdx], numAtoms, batch.atomStarts, batch.system); + } + return batch; +} + +Evaluation splitByMolecule(std::vector energies, const std::vector& gradient, const Batch& batch) { + Evaluation result; + result.energies = std::move(energies); + for (std::size_t molIdx = 0; molIdx + 1 < batch.atomStarts.size(); ++molIdx) { + result.gradients.emplace_back(gradient.begin() + 4 * batch.atomStarts[molIdx], + gradient.begin() + 4 * batch.atomStarts[molIdx + 1]); + } + return result; +} + +constexpr double kGradientFill = 7.0; + +//! Calls the block-per-molecule kernels directly. Gradients start at kGradientFill so untouched systems are visible. +template +Evaluation evaluateKernels(const Batch& batch, + const DG::AllInOneForceWeights& weights, + const bool torsionTermsEnabled, + const std::vector& active) { + using Buffers = std::conditional_t, + DG::BatchedAllInOneDeviceBuffers, + DG::BatchedAllInOneDeviceBuffersSingle>; + Buffers buffers; + DG::sendAllInOneSystemToDevice(batch.system, buffers, nullptr); + const int numMols = static_cast(batch.atomStarts.size()) - 1; + AsyncDeviceVector atomStartsDevice; + AsyncDeviceVector positionsDevice; + AsyncDeviceVector energiesDevice(numMols); + AsyncDeviceVector gradientDevice; + AsyncDeviceVector activeDevice; + atomStartsDevice.setFromVector(batch.atomStarts); + positionsDevice.setFromVector(std::vector(batch.positions.begin(), batch.positions.end())); + gradientDevice.setFromVector(std::vector(batch.positions.size(), static_cast(kGradientFill))); + activeDevice.setFromVector(active); + + const auto terms = DG::toAllInOneContribsDevicePtr(buffers, atomStartsDevice.data()); + CHECK_CUDA_RETURN(DG::launchBlockPerMolEnergyKernelAIO(numMols, + terms, + positionsDevice.data(), + energiesDevice.data(), + weights, + torsionTermsEnabled, + activeDevice.data())); + CHECK_CUDA_RETURN(DG::launchBlockPerMolGradKernelAIO(numMols, + terms, + positionsDevice.data(), + gradientDevice.data(), + weights, + torsionTermsEnabled, + activeDevice.data())); + std::vector energies(numMols); + std::vector gradient(batch.positions.size()); + energiesDevice.copyToHost(energies); + gradientDevice.copyToHost(gradient); + CHECK_CUDA_RETURN(cudaDeviceSynchronize()); + return splitByMolecule(std::vector(energies.begin(), energies.end()), + std::vector(gradient.begin(), gradient.end()), + batch); +} + +//! Calls the batched forcefield through its double-precision interface. +Evaluation evaluateForcefield(nvMolKit::AllInOneETKDGBatchedForcefield& forcefield, + const Batch& batch, + const std::vector& active) { + const int numMols = static_cast(batch.atomStarts.size()) - 1; + AsyncDeviceVector positionsDevice; + AsyncDeviceVector energiesDevice(numMols); + AsyncDeviceVector gradientDevice(batch.positions.size()); + AsyncDeviceVector activeDevice; + positionsDevice.setFromVector(batch.positions); + activeDevice.setFromVector(active); + energiesDevice.zero(); + gradientDevice.zero(); + CHECK_CUDA_RETURN(forcefield.computeEnergy(energiesDevice.data(), positionsDevice.data(), activeDevice.data())); + CHECK_CUDA_RETURN(forcefield.computeGradients(gradientDevice.data(), positionsDevice.data(), activeDevice.data())); + std::vector energies(numMols); + std::vector gradient(batch.positions.size()); + energiesDevice.copyToHost(energies); + gradientDevice.copyToHost(gradient); + CHECK_CUDA_RETURN(cudaDeviceSynchronize()); + return splitByMolecule(std::move(energies), gradient, batch); +} + +//! Tolerances relative to max(1, |expected|). +struct Tolerance { + double energy; + double gradient; +}; + +Tolerance toleranceFor(const PrecisionMode precision) { + return precision == PrecisionMode::FULL ? Tolerance{1e-5, 1e-4} : Tolerance{2e-3, 2e-3}; +} + +bool withinTolerance(const double got, const double want, const double tolerance) { + return std::abs(got - want) <= tolerance * std::max(1.0, std::abs(want)); +} + +void expectMatches(const Evaluation& got, + const Evaluation& want, + const std::vector& active, + const Tolerance& tolerance, + const double inactiveGradient) { + ASSERT_EQ(got.energies.size(), want.energies.size()); + int energyFailures = 0; + int gradientFailures = 0; + for (std::size_t molIdx = 0; molIdx < want.energies.size(); ++molIdx) { + if (!withinTolerance(got.energies[molIdx], want.energies[molIdx], tolerance.energy)) { + ADD_FAILURE() << "molecule " << molIdx << " energy " << got.energies[molIdx] << " vs " << want.energies[molIdx]; + if (++energyFailures > 5) { + return; + } + } + for (std::size_t coord = 0; coord < want.gradients[molIdx].size(); ++coord) { + const double expected = active[molIdx] ? want.gradients[molIdx][coord] : inactiveGradient; + if (!withinTolerance(got.gradients[molIdx][coord], expected, tolerance.gradient)) { + ADD_FAILURE() << "molecule " << molIdx << " coordinate " << coord << " gradient " + << got.gradients[molIdx][coord] << " vs " << expected; + if (++gradientFailures > 5) { + return; + } + break; + } + } + } +} + +std::vector randomActiveMask(const std::size_t size) { + std::mt19937 generator(42); + std::uniform_int_distribution<> coin(0, 1); + std::vector active(size); + for (auto& flag : active) { + flag = static_cast(coin(generator)); + } + active.front() = 1; + return active; +} + +// --------------------------------------- +// Kernel tests: one term family at a time +// --------------------------------------- + +enum class AllInOneTerm { + DistViolation, + ChiralViolation, + FourthDim, + HarmonicDistance, + LinearAngle, + ExperimentalTorsion, + Planarity +}; + +std::string termName(const AllInOneTerm term) { + switch (term) { + case AllInOneTerm::DistViolation: + return "DistViolation"; + case AllInOneTerm::ChiralViolation: + return "ChiralViolation"; + case AllInOneTerm::FourthDim: + return "FourthDim"; + case AllInOneTerm::HarmonicDistance: + return "HarmonicDistance"; + case AllInOneTerm::LinearAngle: + return "LinearAngle"; + case AllInOneTerm::ExperimentalTorsion: + return "ExperimentalTorsion"; + case AllInOneTerm::Planarity: + return "Planarity"; + } + return "Unknown"; +} + +bool isTermContrib(const ForceFields::ForceFieldContrib* contrib, const AllInOneTerm term) { + switch (term) { + case AllInOneTerm::DistViolation: + return dynamic_cast(contrib) != nullptr; + case AllInOneTerm::ChiralViolation: + return dynamic_cast(contrib) != nullptr; + case AllInOneTerm::FourthDim: + return dynamic_cast(contrib) != nullptr; + case AllInOneTerm::HarmonicDistance: + return dynamic_cast(contrib) != nullptr; + case AllInOneTerm::LinearAngle: + return dynamic_cast(contrib) != nullptr; + case AllInOneTerm::ExperimentalTorsion: + return dynamic_cast(contrib) != nullptr; + case AllInOneTerm::Planarity: + return dynamic_cast(contrib) != nullptr; + } + return false; +} + +//! Clears every term family except one, returning how many terms of that family remain. +std::size_t keepOnlyTerm(DG::AllInOneForceContribsHost& contribs, const AllInOneTerm term) { + auto& dg = contribs.distanceGeometry; + if (term != AllInOneTerm::DistViolation) { + dg.distTerms = {}; + } + if (term != AllInOneTerm::ChiralViolation) { + dg.chiralTerms = {}; + } + if (term != AllInOneTerm::FourthDim) { + dg.fourthTerms = {}; + } + if (term != AllInOneTerm::HarmonicDistance) { + contribs.harmonicDistanceTerms = {}; + } + if (term != AllInOneTerm::LinearAngle) { + contribs.linearAngleTerms = {}; + } + if (term != AllInOneTerm::ExperimentalTorsion) { + contribs.experimentalTorsionTerms = {}; + } + if (term != AllInOneTerm::Planarity) { + contribs.planarityTerms = {}; + } + return dg.distTerms.idx1.size() + dg.chiralTerms.idx1.size() + dg.fourthTerms.idx.size() + + contribs.harmonicDistanceTerms.idx1.size() + contribs.linearAngleTerms.idx1.size() + + contribs.experimentalTorsionTerms.idx1.size() + contribs.planarityTerms.idx1.size(); +} + +using KernelTestParams = std::tuple; + +class AllInOneKernelTermTest : public ::testing::TestWithParam {}; + +TEST_P(AllInOneKernelTermTest, MatchesRdkitInBothPhases) { + const auto [term, precision] = GetParam(); + const auto params = RDKit::DGeomHelpers::ETKDGv3; + + auto mols = sdfMolecules("MMFF94_dative.sdf", 50); + for (auto& mol : sdfMolecules("MMFF94_hypervalent.sdf", 50)) { + mols.push_back(std::move(mol)); + } + for (auto& mol : coverageMolecules()) { + mols.push_back(std::move(mol)); + } + auto molecules = prepareMolecules(mols, params); + + std::vector contribs; + std::size_t numTerms = 0; + for (const auto& molecule : molecules) { + numTerms += keepOnlyTerm(contribs.emplace_back(buildContribs(molecule, params)), term); + } + ASSERT_GT(numTerms, 0u) << "test molecules have no " << termName(term) << " terms"; + + const auto batch = buildBatch(molecules, contribs); + const auto weights = DG::allInOneForceWeights(molecules.front().eargs.etkdgDetails); + const auto active = randomActiveMask(molecules.size()); + for (const bool torsionTermsEnabled : {false, true}) { + SCOPED_TRACE("torsionTermsEnabled=" + std::to_string(torsionTermsEnabled)); + const auto want = evaluateReference(molecules, params, torsionTermsEnabled, active, [term](auto& field) { + std::erase_if(field.contribs(), [term](const auto& contrib) { return !isTermContrib(contrib.get(), term); }); + }); + const auto got = precision == PrecisionMode::FULL ? + evaluateKernels(batch, weights, torsionTermsEnabled, active) : + evaluateKernels(batch, weights, torsionTermsEnabled, active); + expectMatches(got, want, active, toleranceFor(precision), kGradientFill); + } +} + +INSTANTIATE_TEST_SUITE_P(AllTerms, + AllInOneKernelTermTest, + ::testing::Combine(::testing::Values(AllInOneTerm::DistViolation, + AllInOneTerm::ChiralViolation, + AllInOneTerm::FourthDim, + AllInOneTerm::HarmonicDistance, + AllInOneTerm::LinearAngle, + AllInOneTerm::ExperimentalTorsion, + AllInOneTerm::Planarity), + ::testing::Values(PrecisionMode::FULL, PrecisionMode::SINGLE)), + [](const ::testing::TestParamInfo& info) { + return termName(std::get<0>(info.param)) + + (std::get<1>(info.param) == PrecisionMode::FULL ? "_Double" : "_Float"); + }); + +TEST(AllInOneKernel, MoleculesTooLargeForSharedMemoryAccumulateInGlobalMemory) { + const auto params = RDKit::DGeomHelpers::ETKDGv3; + std::vector> mols; + for (const auto& smiles : {std::string(100, 'C') + "C#N", std::string("CC#N")}) { + std::unique_ptr mol(RDKit::SmilesToMol(smiles)); + mols.emplace_back(RDKit::MolOps::addHs(*mol)); + } + ASSERT_GT(mols.front()->getNumAtoms(), 256u); + auto molecules = prepareMolecules(mols, params); + std::vector contribs; + for (const auto& molecule : molecules) { + contribs.push_back(buildContribs(molecule, params)); + } + const auto batch = buildBatch(molecules, contribs); + const auto weights = DG::allInOneForceWeights(molecules.front().eargs.etkdgDetails); + const std::vector active(molecules.size(), 1); + const auto want = evaluateReference(molecules, params, true, active); + expectMatches(evaluateKernels(batch, weights, true, active), + want, + active, + toleranceFor(PrecisionMode::FULL), + 0.0); + expectMatches(evaluateKernels(batch, weights, true, active), + want, + active, + toleranceFor(PrecisionMode::SINGLE), + 0.0); +} + +// --------------------------------------------- +// Batched forcefield tests: all terms together +// --------------------------------------------- + +using ForcefieldTestParams = std::tuple; + +class AllInOneForcefieldTest : public ::testing::TestWithParam {}; + +TEST_P(AllInOneForcefieldTest, BatchMatchesRdkitInBothPhases) { + const auto [option, fileName, precision] = GetParam(); + const auto params = getETKDGOption(option); + const auto mols = sdfMolecules(fileName, 100); + auto molecules = prepareMolecules(mols, params); + + std::vector contribs; + for (const auto& molecule : molecules) { + contribs.push_back(buildContribs(molecule, params)); + } + const auto batch = buildBatch(molecules, contribs); + const auto active = randomActiveMask(molecules.size()); + nvMolKit::AllInOneETKDGBatchedForcefield forcefield(batch.system, + batch.atomStarts, + DG::allInOneForceWeights(molecules.front().eargs.etkdgDetails), + {}, + nullptr, + precision); + EXPECT_EQ(forcefield.precision(), precision); + EXPECT_EQ(forcefield.dataDim(), 4); + for (const bool torsionTermsEnabled : {false, true}) { + SCOPED_TRACE("torsionTermsEnabled=" + std::to_string(torsionTermsEnabled)); + forcefield.setTorsionTermsEnabled(torsionTermsEnabled); + expectMatches(evaluateForcefield(forcefield, batch, active), + evaluateReference(molecules, params, torsionTermsEnabled, active), + active, + toleranceFor(precision), + 0.0); + } +} + +INSTANTIATE_TEST_SUITE_P( + ETKDGOptions, + AllInOneForcefieldTest, + ::testing::Combine(::testing::Values(ETKDGOption::ETKDGv3, ETKDGOption::ETKDGv2, ETKDGOption::ETDG, ETKDGOption::KDG), + ::testing::Values("MMFF94_dative.sdf", "MMFF94_hypervalent.sdf"), + ::testing::Values(PrecisionMode::FULL, PrecisionMode::SINGLE)), + [](const ::testing::TestParamInfo& info) { + const auto& fileName = std::get<1>(info.param); + return getETKDGOptionName(std::get<0>(info.param)) + "_" + fileName.substr(0, fileName.find('.')) + + (std::get<2>(info.param) == PrecisionMode::FULL ? "_Double" : "_Float"); + }); + +struct EdgeTestCase { + std::vector smiles; + std::vector active; +}; + +class AllInOneForcefieldEdgeTest : public ::testing::TestWithParam> {}; + +TEST_P(AllInOneForcefieldEdgeTest, OneAndTwoAtomMolecules) { + const auto& [testCase, precision] = GetParam(); + const auto params = RDKit::DGeomHelpers::ETKDGv3; + std::vector> mols; + for (const auto& smiles : testCase.smiles) { + mols.emplace_back(RDKit::SmilesToMol(smiles)); + } + auto molecules = prepareMolecules(mols, params); + std::vector contribs; + for (const auto& molecule : molecules) { + contribs.push_back(buildContribs(molecule, params)); + } + const auto batch = buildBatch(molecules, contribs); + nvMolKit::AllInOneETKDGBatchedForcefield forcefield(batch.system, + batch.atomStarts, + DG::allInOneForceWeights(molecules.front().eargs.etkdgDetails), + {}, + nullptr, + precision); + forcefield.setTorsionTermsEnabled(true); + expectMatches(evaluateForcefield(forcefield, batch, testCase.active), + evaluateReference(molecules, params, true, testCase.active), + testCase.active, + toleranceFor(precision), + 0.0); +} + +INSTANTIATE_TEST_SUITE_P(OneTwoAtoms, + AllInOneForcefieldEdgeTest, + ::testing::Combine(::testing::Values( + EdgeTestCase{ + {"C", "O"}, + { 1, 1} +}, + EdgeTestCase{{"C", "CC"}, {0, 1}}, + EdgeTestCase{{"CO", "N"}, {1, 0}}, + EdgeTestCase{{"CC", "CN"}, {1, 1}}), + ::testing::Values(PrecisionMode::FULL, PrecisionMode::SINGLE))); + +TEST(AllInOneForcefield, SinglePrecisionInterfaceEvaluatesNativelyAndRejectsFullPrecisionBuffers) { + const auto params = RDKit::DGeomHelpers::ETKDGv3; + auto mols = coverageMolecules(); + auto molecules = prepareMolecules(mols, params); + + std::vector contribs; + for (const auto& molecule : molecules) { + contribs.push_back(buildContribs(molecule, params)); + } + const auto batch = buildBatch(molecules, contribs); + const auto weights = DG::allInOneForceWeights(molecules.front().eargs.etkdgDetails); + const int numMols = static_cast(molecules.size()); + + AsyncDeviceVector positionsDevice; + AsyncDeviceVector energiesDevice(numMols); + AsyncDeviceVector gradientDevice(batch.positions.size()); + positionsDevice.setFromVector(std::vector(batch.positions.begin(), batch.positions.end())); + + nvMolKit::AllInOneETKDGBatchedForcefield full(batch.system, batch.atomStarts, weights); + EXPECT_EQ(full.computeEnergy(energiesDevice.data(), positionsDevice.data()), cudaErrorInvalidValue); + EXPECT_EQ(full.computeGradients(gradientDevice.data(), positionsDevice.data()), cudaErrorInvalidValue); + + nvMolKit::AllInOneETKDGBatchedForcefield single(batch.system, + batch.atomStarts, + weights, + {}, + nullptr, + PrecisionMode::SINGLE); + single.setTorsionTermsEnabled(true); + gradientDevice.zero(); + CHECK_CUDA_RETURN(single.computeEnergy(energiesDevice.data(), positionsDevice.data())); + CHECK_CUDA_RETURN(single.computeGradients(gradientDevice.data(), positionsDevice.data())); + std::vector energies(numMols); + std::vector gradient(batch.positions.size()); + energiesDevice.copyToHost(energies); + gradientDevice.copyToHost(gradient); + CHECK_CUDA_RETURN(cudaDeviceSynchronize()); + + const std::vector active(numMols, 1); + expectMatches(splitByMolecule(std::vector(energies.begin(), energies.end()), + std::vector(gradient.begin(), gradient.end()), + batch), + evaluateReference(molecules, params, true, active), + active, + toleranceFor(PrecisionMode::SINGLE), + 0.0); +} + +} // namespace + +#endif // RDKIT_AIO_ETKDG_API