Repository navigation
Support All-in-one ETKDG pipeline - #386
Merged
Merged
Conversation
Contributor
 [Medium impact] The PR appears safe to merge; no actionable issue was established.
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Runs RDKit's all-in-one coordinate refinement when useLegacyImplementation=false (RDKit 2026.03.4+): a batched stage that minimizes the distance-geometry terms for each molecule's heavy-atom count, then adds the experimental torsion and planarity terms for up to eight 100-iteration cycles, followed by the all-in-one energy, K-term, clash and sp2 geometry checks. The stage evaluates in the minimizer's precision and caches each molecule's terms across batches like the DG/ETK stages. Batched BFGS accepts optional per-system iteration limits so the first phase stops each molecule at its own heavy-atom count, as RDKit does. Using the batch maximum instead let strained molecules converge that RDKit rejects and retries. The K-term check builds the terms of RDKit's construct3DImproperForceField and evaluates them through the ETK batch buffers. It applies the total-energy limit that RDKit 2026.09 tests before summing, so RDKit never enforces it. For all-in-one refinement, prepareEmbedderArgs selects RDKit's all-in-one force constants before building the torsions, which RDKit scales by them, and keeps its own copy of the topological distance matrix, which the all-in-one force field uses to find 1-2 and 1-3 pairs. The copy is shared by conformer clones; the matrix RDKit caches on the molecule can be replaced while they still use it. All-in-one refinement runs on the batched minimizer; an explicit PER_MOLECULE backend raises. Tests compare the stage against RDKit's all-in-one minimization from the same coordinates (acceptance counts and minimized energies, float and double, four ETKDG variants), check that per-system iteration limits stop each system exactly where a separately limited run does, check the supported and rejected parameter combinations, and the Python RDKit comparisons now cover both implementations.
scal444
force-pushed
the
feat/aio-etkdg-pipeline
branch
from
October 29, 2026 19:03
d0875cc to
c3cc1cd
Compare
Collaborator
Author
All-in-one refinement previously forced every run onto the batched minimizer: an explicit PER_MOLECULE backend raised, and HYBRID silently ran batched even for small molecules, which HYBRID normally sends to the per-molecule kernels. The fused per-molecule BFGS kernel gains an all-in-one force field type that evaluates the existing molEnergyAIO/molGradAIO device functions, with the all-in-one weights, the torsion phase, and optional per-system iteration limits. A system that reaches its own limit counts as finished, as in the batched minimizer, so the first phase still stops each molecule at its heavy-atom count. The force field settings now travel as one struct, and the four energy and gradient dispatch sites share one helper each. The all-in-one stage runs both phases through minimizeWithAIO when the backend resolves to PER_MOLECULE and computes the final energy check with the block-per-molecule energy kernel in the minimizer's precision, as the batched path does. The stage test compares both backends against RDKit, the iteration limit test now covers the per-molecule backend, and an end-to-end test embeds small molecules and one above the shared-memory limit on every backend, in batches small enough that HYBRID sends the small molecules to the per-molecule kernels. The embedMolecules docstring drops the backend restriction.
- Each all-in-one phase minimizes only molecules whose energy is above RDKit's ERROR_TOL, and the second phase drops molecules as they converge, as RDKit's per-molecule minimizeAllInOne loop does. Before, every active molecule was re-minimized in all eight second-phase cycles whenever any molecule in the batch needed another, so its result depended on its batchmates. Both BFGS backends; the minimizer exposes its per-system statuses for this. - The clash check stores each distinct bounds matrix's lower triangle once and generates atom pairs in the kernel, one block per system, instead of uploading every pair of every conformer. - The K-term check caches its terms per molecule across batches. - The all-in-one path checks chiral centers once, as RDKit does; the final pass-through check stays on the legacy path. - The bounds matrix and triangle smoothing tests use nvMolKit's initETKDG instead of stale copies. Tests check on both backends and precisions that a molecule's result does not depend on its batchmates and that a system already below ERROR_TOL is not moved, and the RDKit reference minimization now gates its phases on ERROR_TOL too. They also cover the clash check across systems of several sizes, two sharing one bounds matrix, against a host evaluation, and the K-term cache.
scal444
force-pushed
the
feat/aio-etkdg-pipeline
branch
from
October 29, 2026 21:36
c3cc1cd to
8aa7f6f
Compare
evasnow1992
approved these changes
Oct 29, 2026
evasnow1992
left a comment
evasnow1992
left a comment
Collaborator
There was a problem hiding this comment.
Changes look good to me!
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
We previously added kernels and setup support, now wiring into the ETKDG driver.
Also consolidates some duplicated initETKDG functionality