Sync upstream TMD v0.3.0#5
Merged
Merged
Conversation
…md-industries#64) * Use custom OpenMM templates for placeholder atoms * Avoids the implicit requirement that water forcefields have to be from amber14 * Bump to version 0.2.1 * Fix linting * Switch to using Uranium as the atom type * Please don't simulate Uranium with TMD * Fix typo
* Use Amber AM1BCCs as the base charges then add correctable charge corrections. * The aromaticity model doesn't match, but for training this should be a fine starting point. OpenFF prefers MDL which is the standard for TMD as well. * Clean up tests related to OE
* Was not performing conversion of units correctly
* Remove quotes around CMAKE vars * No reason to do this * Add lib directory as additional rpath * Include shared objects in package * Fix links to old repo * Add docs * Include versions in shared objects * PR Feedback: Remove build dependecy * Not strictly necessary
* Remove restriction that waters are at the end of the PDB * Make waters contiguous to ensure water sampler is setup correctly * Probably can make the water sampler more flexible. May need to be revisited for performance later
* Similar to FlatBottomBond except it restraints particles to positions * Useful for minimization
* Turn on heavy matches heavy by default * Not to be merged until full benchmarks are back * Fix failing graph building tests * Update hashes
* Remove stale README comment * Bump to 0.2.3
* Batch Bonded Potentials + Integrators (#8) * Support Batching in Harmonic Bond * Lots of boilplate to handle changes to the API * execute_dim is janky, but need to figure out how to handle this correctly * Enforce potentials must have batch size 1 in context * To be fixed later on * Add EnergyAccumulator utility class * Modifies all the potentials to use it instead * Reduces boilerplate and makes it easier to batch everything * Batch PeriodicTorsions * Fix tests * Adds num_atoms to all other valence terms * HarmonicAngle Batch Support * ChiralAtomRestraint Batch Support * Adds some jank around params, since they are flat for ChiralAtomRestraint * Improve test robustness * ChiralBondRestraint Batch Support * FlatBottomBond Batch Support * Fix up the tests * LogFlatBottomBond Batch support * NonbondedPairList Batching Support * Should batching of JAX methods be supported? * NonbondedPrecomputedPairlist Batching support * Enforce idxs more consistently * Fixes test * Fix failing tests * Exclusions were flat * Initial Batching of Vacuum simulations * Expand test to cover langevin integrator * Fix up some typing * Batch Neighborlist and Nonbonded Potentials (tmd-industries#38) * Swap ordering of row/col in set_idxs_device * Fix number of atoms for precomputed pairlist * Fix up the integrators * Rename batch_size to num_systems in older code * Batch Movers (tmd-industries#45) * WIP * WIP: Fix up the barostat * Expose num systems to move signature * Remove memcheck for tests * OOMing * Return scale factors * Fix up the existing tests * Fix memory leak in energy accumulator * Support batching in TIBD/BD ExchangeMover * Update tests * Fix batching for single topology * Add verification that the size of the system matches the potentials * Add tests for exchange movers * Reduce memory usage during CI * Fix test * Tweak memory usage in CI * Probably need more ram still * Reduce number of workers further * Fix a neighborlist test * Reduce vacuum tests memory usage * Add batching benchmarks * Appears slower than MPS. The issues with rebuilding the nblist constnatly may be a problem for global. * Reduce memory consumption further * Reduce the memory usage * Disable float64 tests * Remove batch simulations from memcheck * Absurd how much memory these things are taking * Disable memory checks under x-dist * Mixture of tests impacts RSS * Undo change * Fix failing test * Re-enable parallelism in memory tests * Adds Batch Support for Local MD (tmd-industries#70) * WIP * WIP * Initial change, it must be broken * WIP * WIP * REVERT * Clean up * Fixes a bug in the setup of the free * Fix issue with construction of system idxs * Fix linting * Skip testing large batches in CI with memcheck on * Remove batch tests from memcheck * Causing ooms * Skip memory check in failing test * Running enough RBFE at this stage that it doesn't seem like we will be caught unaware. But running with pytest-xdist has inconsistent memory behavior * Fix resetting of nonbonded ixn group * More memory allocated than strictly needed, but that is fine for now * Batch Flat Bottom Restraint (tmd-industries#76) * Adds Amber AM1CCC Forcefield (tmd-industries#69) * Use Amber AM1BCCs as the base charges then add correctable charge corrections. * The aromaticity model doesn't match, but for training this should be a fine starting point. OpenFF prefers MDL which is the standard for TMD as well. * Clean up tests related to OE * Add CCC refitting test (tmd-industries#71) * Fixes incorrect predictions in CSVs (tmd-industries#72) * Was not performing conversion of units correctly * Bump to 0.2.2 (tmd-industries#73) * CMake Fixes + RPath (tmd-industries#67) * Remove quotes around CMAKE vars * No reason to do this * Add lib directory as additional rpath * Include shared objects in package * Fix links to old repo * Add docs * Include versions in shared objects * PR Feedback: Remove build dependecy * Not strictly necessary * Ensure that water atoms are contiguous after prep (tmd-industries#74) * Remove restriction that waters are at the end of the PDB * Make waters contiguous to ensure water sampler is setup correctly * Probably can make the water sampler more flexible. May need to be revisited for performance later * Add FlatBottomRestraint (tmd-industries#75) * Similar to FlatBottomBond except it restraints particles to positions * Useful for minimization * Batch HREX (tmd-industries#77) * WIP * WIP * Initial change, it must be broken * WIP * WIP * REVERT * Clean up * Fixes a bug in the setup of the free * Fix issue with construction of system idxs * Fix linting * Skip testing large batches in CI with memcheck on * Adds Amber AM1CCC Forcefield (tmd-industries#69) * Use Amber AM1BCCs as the base charges then add correctable charge corrections. * The aromaticity model doesn't match, but for training this should be a fine starting point. OpenFF prefers MDL which is the standard for TMD as well. * Clean up tests related to OE * Remove batch tests from memcheck * Causing ooms * Skip memory check in failing test * Running enough RBFE at this stage that it doesn't seem like we will be caught unaware. But running with pytest-xdist has inconsistent memory behavior * Add CCC refitting test (tmd-industries#71) * Fix resetting of nonbonded ixn group * More memory allocated than strictly needed, but that is fine for now * JANK * REMOVESYNC * jank * Fixes incorrect predictions in CSVs (tmd-industries#72) * Was not performing conversion of units correctly * Bump to 0.2.2 (tmd-industries#73) * CMake Fixes + RPath (tmd-industries#67) * Remove quotes around CMAKE vars * No reason to do this * Add lib directory as additional rpath * Include shared objects in package * Fix links to old repo * Add docs * Include versions in shared objects * PR Feedback: Remove build dependecy * Not strictly necessary * Ensure that water atoms are contiguous after prep (tmd-industries#74) * Remove restriction that waters are at the end of the PDB * Make waters contiguous to ensure water sampler is setup correctly * Probably can make the water sampler more flexible. May need to be revisited for performance later * Add FlatBottomRestraint (tmd-industries#75) * Similar to FlatBottomBond except it restraints particles to positions * Useful for minimization * jank * Working * Further * Add tests of batching on and off * Fix up water sampling * Revert "REMOVESYNC" This reverts commit 41cef11. * Fix the tests * Fix environment variable and examples * Update hashes * Update hashes * Fix output dirs * Turn on heavy matches heavy by default (tmd-industries#66) * Turn on heavy matches heavy by default * Not to be merged until full benchmarks are back * Fix failing graph building tests * Update hashes * Update shas * Update missed sha * Loosen tolerance on failing nightly test * Reduce memory consumption of energy accumulation in the batched mode (tmd-industries#80) * Batch Bisection (tmd-industries#78) * Bisection initial implementation * Lots of duplicate code * Handle bisection where more than 2 initial lambdas provided * Disable batching by default * Enabled by the TMD_BATCH_MODE=on flag * Update vacuum sha * Fix bug in bisection * Was running 1 more iteration than it should * Fix test * Update hashes
* Bump to 0.3.0 * Add note in documentation about batching
Merges tmd-industries/tmd v0.3.0 into the Genesis fork, preserving our customizations: - Python 3.10 compatibility (TypeAlias, TypeVar for Self, relaxed version constraints) - jax[cuda12] dependency bundling - GCC 14 / CUDA 12.9 build fixes (-Wno-error=double-promotion, -Wno-error=template-id-cdtor) - CMake CUDA_ARCHITECTURES quoting fix Upstream additions (v0.2.0 -> v0.3.0): - Batched MD simulation support (TMD_BATCH_MODE env var) - FlatBottomRestraint potential - Amber AM1CCC forcefield - Heavy-matches-heavy atom mapping by default - CMake RPath fixes - Water atom ordering fix in system prep - CSV prediction bug fix - Extensive CUDA kernel batching support across all potentials Single conflict in CMakeLists.txt resolved by combining our warning suppression flags with upstream's CUDA_FLAGS quoting change. Made-with: Cursor
Our codebase is on Python 3.12 now, so the 3.10 backcompat workarounds are no longer needed. Reverts to upstream's PEP 695 type statements and typing.Self usage. Made-with: Cursor
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
3 participants
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.
Summary
typestatements andtyping.SelfUpstream changes included
Fork-specific changes
TypeAlias→ PEP 695typestatements,TypeVarworkaround →typing.Self, bumppython_requires/CMake/conda to 3.12Test plan
All tests ran locally and on Nucleus: https://nucleus.playground.genesistherapeutics.ai/job_manager/job/1b191c24-1f1c-4d75-983d-a1a5760d8d13