feat(openmm,jax): add Urey-Bradley term support (#116)#191
Merged
ericchansen merged 2 commits intomasterfrom Mar 30, 2026
Merged
feat(openmm,jax): add Urey-Bradley term support (#116)#191ericchansen merged 2 commits intomasterfrom
ericchansen merged 2 commits intomasterfrom
Conversation
Add Urey-Bradley (UB) 1-3 distance interactions for CHARMM force fields. UB terms are optional per-angle parameters (ub_force_constant, ub_equilibrium) appended to the tail of the parameter vector for backward compatibility. Changes: - AngleParam: add optional ub_force_constant and ub_equilibrium fields - ForceField: extend param vector, bounds, labels, and indices for UB tail - OpenMM backend: add HarmonicBondForce for UB terms, update_forcefield support - JAX backend: add _ub_energy function, compile UB into JIT energy closure - _jax_common: add n_vdw_types param and 'ub' offset to compute_param_offsets - Tests: 29 new tests covering FF, OpenMM, JAX, and cross-engine parity Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
Adds Urey–Bradley (1–3 distance) term support for CHARMM-style force fields across the ForceField model plus OpenMM and JAX backends, with tests verifying energy behavior and cross-engine parity.
Changes:
- Extend
AngleParamandForceFieldparam-vector handling to include optional UB parameters appended after vdW. - Add UB term evaluation to OpenMM (via
HarmonicBondForce) and to the JAX engine energy function/topology. - Add a comprehensive UB-focused test suite and update parameter-cycling expectations.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
q2mm/models/forcefield.py |
Adds optional UB fields to AngleParam and appends UB params to the param vector + indices/bounds/labels. |
q2mm/backends/mm/openmm.py |
Builds and updates UB terms using HarmonicBondForce on 1–3 atom pairs. |
q2mm/backends/mm/jax_engine.py |
Implements _ub_energy, adds UB topology mapping, and includes UB energy in both harmonic/MM3 branches. |
q2mm/backends/mm/_jax_common.py |
Extends param-offset computation to account for vdW block size and adds "ub" offset. |
test/test_urey_bradley.py |
New tests for UB param-vector integration, OpenMM/JAX energy, and parity. |
test/test_parameter_cycling.py |
Updates expected parameter-type keys to include ub_k/ub_eq. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
- Check both ub_force_constant and ub_equilibrium are non-None before creating or updating Urey-Bradley terms - Raise ValueError if only one is set (configuration error) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
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.
Summary
Add Urey-Bradley (1-3 distance) term support for CHARMM force fields. Closes #116.
Changes
q2mm/models/forcefield.pyub_force_constant/ub_equilibriumonAngleParam, extended param vector with UB tail sectionq2mm/backends/mm/openmm.pyHarmonicBondForceon 1-3 atom pairs, with parameter update supportq2mm/backends/mm/jax_engine.py_ub_energy()function, UB topology inJaxHandle, wired into both harmonic and MM3 branchesq2mm/backends/mm/_jax_common.pycompute_param_offsets()gainsn_vdw_typesparam and"ub"offset keyDesign
_ub_anglesproperty filters to angles with UB params setE_UB = k_UB * (r_13 - r_13_0)^2(harmonic bond between outer angle atoms)Testing