Skip to content

Backport fix from PR #195 #196

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lohedges merged 1 commit into from
Aug 26, 2026
Merged

Backport fix from PR #195 #196

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 18 additions & 24 deletions src/somd2/_utils/_schedules.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,13 +154,19 @@ def ring_break_morph():
Three stages: potential_swap → restraints_off → morph.

During restraints_off the Morse restraint ramps off (morse_soft: 1→0) while
the ring-break softcore simultaneously ramps on (alpha: 1→0, kappa: 0→1),
equations mirror ring-break so that ``ring_break_morph().reverse()`` is the
providing a smooth handover with no gap between the two forces. The ring-make
correct schedule for the ring-making direction (used by
:func:`reverse_ring_break_morph`). Because ring_break_morph is only used for
ring-breaking perturbations (no ring-make force present), the ring-make
equations have no effect on forward simulations.
the ring-break softcore LJ simultaneously ramps on (alpha: 1→0), providing a
smooth handover with no gap between the two forces.

Coulomb is decoupled from the LJ and driven by its own coul_kappa lever,
which is held at zero through both bonded stages and ramps 0→1 during morph
only, once the softcore LJ has already separated the pair.

The ring-make equations mirror ring-break so that
``ring_break_morph().reverse()`` is the correct schedule for the ring-making
direction (used by :func:`reverse_ring_break_morph`). Because
ring_break_morph is only used for ring-breaking perturbations (no ring-make
force present), the ring-make equations have no effect on forward
simulations.

Returns
-------
Expand All @@ -172,8 +178,8 @@ def ring_break_morph():

s = _LambdaSchedule.standard_morph()

# restraints_off [1/3, 2/3): Morse ramps off while ring-break softcore ramps
# on simultaneously (alpha: 1→0, kappa: 0→1). Bonded terms (angles, torsions)
# restraints_off [1/3, 2/3): Morse ramps off while the ring-break softcore LJ
# ramps on simultaneously (alpha: 1→0). Bonded terms (angles, torsions)
# interpolate initial→final over the same stage. ring-make mirrors ring-break
# so that after .reverse(), the ring-make softcore ramps off as morse_soft ramps
# on in the reversed restraints_off stage, correct for ring-making perturbations.
Expand Down Expand Up @@ -205,15 +211,9 @@ def ring_break_morph():
s.set_equation(
stage="restraints_off", force="ring-break", lever="alpha", equation=1 - s.lam()
)
s.set_equation(
stage="restraints_off", force="ring-break", lever="kappa", equation=s.lam()
)
s.set_equation(
stage="restraints_off", force="ring-make", lever="alpha", equation=1 - s.lam()
)
s.set_equation(
stage="restraints_off", force="ring-make", lever="kappa", equation=s.lam()
)

s.prepend_stage("potential_swap", s.initial())
s.set_equation(stage="potential_swap", lever="morse_hard", equation=1 - s.lam())
Expand All @@ -237,15 +237,11 @@ def ring_break_morph():
s.set_equation(
stage="potential_swap", force="ring-break", lever="alpha", equation=1
)
s.set_equation(
stage="potential_swap", force="ring-break", lever="kappa", equation=0
)
s.set_equation(stage="potential_swap", force="ring-make", lever="alpha", equation=1)
s.set_equation(stage="potential_swap", force="ring-make", lever="kappa", equation=0)

# morph [2/3, 1]: standard nonbonded morphing with ring-break/ring-make fixed
# at fully open (kappa=1, alpha=0). ring-make mirrors ring-break so .reverse()
# gives kappa=1 at lam=0 of the reversed morph stage (ring-making start).
# at fully open (alpha=0). ring-make mirrors ring-break so .reverse() gives
# alpha=0 at lam=0 of the reversed morph stage (ring-making start).
s.set_equation(stage="morph", lever="morse_hard", equation=0)
s.set_equation(stage="morph", lever="morse_soft", equation=0)
s.set_equation(stage="morph", lever="bond_k", equation=s.final())
Expand All @@ -255,9 +251,7 @@ def ring_break_morph():
s.set_equation(stage="morph", lever="torsion_k", equation=s.final())
s.set_equation(stage="morph", lever="torsion_phase", equation=s.final())
s.set_equation(stage="morph", force="ring-break", lever="alpha", equation=0)
s.set_equation(stage="morph", force="ring-break", lever="kappa", equation=1)
s.set_equation(stage="morph", force="ring-make", lever="alpha", equation=0)
s.set_equation(stage="morph", force="ring-make", lever="kappa", equation=1)

# coul_kappa: zero through both bonded stages so the CLJ exception carries no
# charge while atoms are at covalent distances; ramps 0→1 in morph only once
Expand Down Expand Up @@ -296,7 +290,7 @@ def reverse_ring_break_morph():
This schedule is correct for two equivalent use-cases:

1. A ring-making perturbation run with ``swap_end_states=False``: the
ring-make softcore force (kappa=1 at λ=0, ramping to 0) is controlled
ring-make softcore force (alpha=0 at λ=0, ramping to 1) is controlled
directly by the ring-make lever equations.
2. A ring-breaking perturbation run with ``swap_end_states=True`` (the
runner reverses the schedule automatically, yielding the same effective
Expand Down
169 changes: 80 additions & 89 deletions tests/schedules/test_ring_break.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
reason="openmm support is not available",
)

# Energy threshold (kcal/mol) for the "active" state: kappa=1 should give
# Energy threshold (kcal/mol) for the "active" state: alpha=0 should give
# clearly non-zero CustomBondForce energy.
_ACTIVE_THRESHOLD = 0.1

Expand Down Expand Up @@ -122,59 +122,60 @@ def test_reverse_has_ring_make_not_ring_break(reverse_dynamics):
)


# ── schedule kappa/alpha tests ────────────────────────────────────────────────
# ── schedule alpha/coul_kappa tests ───────────────────────────────────────────
#
# These tests verify kappa and alpha values by calling schedule.morph() directly,
# using the same initial/final values that lambdalever passes in production.
# They are completely independent of Sire's energy formula and will continue to
# work correctly regardless of changes to the softcore implementation.

# ring_break_morph kappa/alpha points (3 equal stages: [0,1/3), [1/3,2/3), [2/3,1]):
# λ=0.00 potential_swap start kappa=0, alpha=1
# λ=0.15 potential_swap mid kappa=0, alpha=1
# λ=1/3 restraints_off start kappa=0, alpha=1 (within-stage lam=0)
# λ=0.45 restraints_off mid kappa=0.35, alpha=0.65 (within-stage lam=0.35)
# λ=0.50 restraints_off mid kappa=0.5, alpha=0.5 (within-stage lam=0.5)
# λ=0.55 restraints_off mid kappa=0.65, alpha=0.35 (within-stage lam=0.65)
# λ=0.60 restraints_off near end kappa=0.8, alpha=0.2 (within-stage lam=0.8)
# λ=2/3 morph start kappa=1, alpha=0
# λ=0.85 morph mid kappa=1, alpha=0
# λ=1.00 morph end kappa=1, alpha=0
_FWD_KAPPA_ALPHA = [
(0.00, 0.0, 1.0),
(0.15, 0.0, 1.0),
(1 / 3, 0.0, 1.0),
(0.45, 0.35, 0.65),
(0.50, 0.5, 0.5),
(0.55, 0.65, 0.35),
(0.60, 0.8, 0.2),
(2 / 3, 1.0, 0.0),
(0.85, 1.0, 0.0),
(1.00, 1.0, 0.0),
# These tests verify alpha and coul_kappa values by calling schedule.morph()
# directly, using the same initial/final values that lambdalever passes in
# production. They are completely independent of Sire's energy formula and will
# continue to work correctly regardless of changes to the softcore
# implementation.

# ring_break_morph alpha points (3 equal stages: [0,1/3), [1/3,2/3), [2/3,1]):
# λ=0.00 potential_swap start alpha=1
# λ=0.15 potential_swap mid alpha=1
# λ=1/3 restraints_off start alpha=1 (within-stage lam=0)
# λ=0.45 restraints_off mid alpha=0.65 (within-stage lam=0.35)
# λ=0.50 restraints_off mid alpha=0.5 (within-stage lam=0.5)
# λ=0.55 restraints_off mid alpha=0.35 (within-stage lam=0.65)
# λ=0.60 restraints_off near end alpha=0.2 (within-stage lam=0.8)
# λ=2/3 morph start alpha=0
# λ=0.85 morph mid alpha=0
# λ=1.00 morph end alpha=0
_FWD_ALPHA = [
(0.00, 1.0),
(0.15, 1.0),
(1 / 3, 1.0),
(0.45, 0.65),
(0.50, 0.5),
(0.55, 0.35),
(0.60, 0.2),
(2 / 3, 0.0),
(0.85, 0.0),
(1.00, 0.0),
]

# reverse_ring_break_morph ring-make kappa/alpha points (mirror of forward):
# λ=0.00 reversed morph start kappa=1, alpha=0
# λ=0.15 reversed morph mid kappa=1, alpha=0
# λ=1/3 reversed restraints_off start kappa=1, alpha=0 (within-stage lam=0)
# λ=0.45 reversed restraints_off mid kappa=0.65, alpha=0.35
# λ=0.50 reversed restraints_off mid kappa=0.5, alpha=0.5
# λ=0.55 reversed restraints_off mid kappa=0.35, alpha=0.65
# λ=0.60 reversed restraints_off near end kappa=0.2, alpha=0.8
# λ=2/3 reversed potential_swap start kappa=0, alpha=1
# λ=0.85 reversed potential_swap mid kappa=0, alpha=1
# λ=1.00 reversed potential_swap end kappa=0, alpha=1
_REV_KAPPA_ALPHA = [
(0.00, 1.0, 0.0),
(0.15, 1.0, 0.0),
(1 / 3, 1.0, 0.0),
(0.45, 0.65, 0.35),
(0.50, 0.5, 0.5),
(0.55, 0.35, 0.65),
(0.60, 0.2, 0.8),
(2 / 3, 0.0, 1.0),
(0.85, 0.0, 1.0),
(1.00, 0.0, 1.0),
# reverse_ring_break_morph ring-make alpha points (mirror of forward):
# λ=0.00 reversed morph start alpha=0
# λ=0.15 reversed morph mid alpha=0
# λ=1/3 reversed restraints_off start alpha=0 (within-stage lam=0)
# λ=0.45 reversed restraints_off mid alpha=0.35
# λ=0.50 reversed restraints_off mid alpha=0.5
# λ=0.55 reversed restraints_off mid alpha=0.65
# λ=0.60 reversed restraints_off near end alpha=0.8
# λ=2/3 reversed potential_swap start alpha=1
# λ=0.85 reversed potential_swap mid alpha=1
# λ=1.00 reversed potential_swap end alpha=1
_REV_ALPHA = [
(0.00, 0.0),
(0.15, 0.0),
(1 / 3, 0.0),
(0.45, 0.35),
(0.50, 0.5),
(0.55, 0.65),
(0.60, 0.8),
(2 / 3, 1.0),
(0.85, 1.0),
(1.00, 1.0),
]

# ring_break_morph coul_kappa points (initial=0, final=1):
Expand Down Expand Up @@ -213,43 +214,35 @@ def test_reverse_has_ring_make_not_ring_break(reverse_dynamics):
]


@pytest.mark.parametrize("lam,expected_kappa,expected_alpha", _FWD_KAPPA_ALPHA)
def test_ring_break_morph_schedule(lam, expected_kappa, expected_alpha):
@pytest.mark.parametrize("lam,expected_alpha", _FWD_ALPHA)
def test_ring_break_morph_schedule(lam, expected_alpha):
"""
ring_break_morph() produces the correct ring-break kappa and alpha at each λ.
ring_break_morph() produces the correct ring-break alpha at each λ.

Uses lambdalever's initial/final values (kappa: 0→1, alpha: 1→0) to ensure
the test matches production behaviour exactly.
Uses lambdalever's initial/final values (alpha: 1→0) to ensure the test
matches production behaviour exactly.
"""
from somd2._utils._schedules import ring_break_morph

s = ring_break_morph()
kappa = s.morph("ring-break", "kappa", 0.0, 1.0, lam)
alpha = s.morph("ring-break", "alpha", 1.0, 0.0, lam)
assert abs(kappa - expected_kappa) < 1e-10, (
f"ring-break kappa={kappa:.8f} at λ={lam:.4f}, expected {expected_kappa}"
)
assert abs(alpha - expected_alpha) < 1e-10, (
f"ring-break alpha={alpha:.8f} at λ={lam:.4f}, expected {expected_alpha}"
)


@pytest.mark.parametrize("lam,expected_kappa,expected_alpha", _REV_KAPPA_ALPHA)
def test_reverse_ring_break_morph_schedule(lam, expected_kappa, expected_alpha):
@pytest.mark.parametrize("lam,expected_alpha", _REV_ALPHA)
def test_reverse_ring_break_morph_schedule(lam, expected_alpha):
"""
reverse_ring_break_morph() produces the correct ring-make kappa and alpha at each λ.
reverse_ring_break_morph() produces the correct ring-make alpha at each λ.

Uses lambdalever's initial/final values (kappa: 1→0, alpha: 0→1) to ensure
the test matches production behaviour exactly.
Uses lambdalever's initial/final values (alpha: 0→1) to ensure the test
matches production behaviour exactly.
"""
from somd2._utils._schedules import reverse_ring_break_morph

s = reverse_ring_break_morph()
kappa = s.morph("ring-make", "kappa", 1.0, 0.0, lam)
alpha = s.morph("ring-make", "alpha", 0.0, 1.0, lam)
assert abs(kappa - expected_kappa) < 1e-10, (
f"ring-make kappa={kappa:.8f} at λ={lam:.4f}, expected {expected_kappa}"
)
assert abs(alpha - expected_alpha) < 1e-10, (
f"ring-make alpha={alpha:.8f} at λ={lam:.4f}, expected {expected_alpha}"
)
Expand Down Expand Up @@ -297,30 +290,30 @@ def test_reverse_ring_break_morph_coul_kappa(lam, expected_coul_kappa):
@pytest.mark.parametrize("lam", [2 / 3, 1.0])
def test_ring_break_active_in_morph(forward_dynamics, lam):
"""
Ring-break energy is clearly non-zero (kappa=1) throughout the morph stage.
Ring-break energy is clearly non-zero (alpha=0) throughout the morph stage.
"""
e = _force_energy_kcal(forward_dynamics, lam, "ring-break")
assert abs(e) > _ACTIVE_THRESHOLD, (
f"ring-break energy {e:.4f} kcal/mol at λ={lam:.4f} is below active "
f"threshold {_ACTIVE_THRESHOLD} kcal/mol (kappa should be 1)"
f"threshold {_ACTIVE_THRESHOLD} kcal/mol (alpha should be 0)"
)


def test_ring_make_active_at_lambda_zero(reverse_dynamics):
"""
Ring-make energy is non-zero at λ=0: the morph stage fixes kappa=1
Ring-make energy is non-zero at λ=0: the morph stage fixes alpha=0
so the ring-make interaction is fully on from the start.
"""
e = _force_energy_kcal(reverse_dynamics, 0.0, "ring-make")
assert abs(e) > _ACTIVE_THRESHOLD, (
f"ring-make energy {e:.4f} kcal/mol at λ=0 is below active threshold "
f"{_ACTIVE_THRESHOLD} kcal/mol (kappa should be 1 in morph stage)"
f"{_ACTIVE_THRESHOLD} kcal/mol (alpha should be 0 in morph stage)"
)


def test_ring_make_inactive_at_lambda_one(reverse_dynamics):
"""
Ring-make energy is near-zero at λ=1 (potential_swap end, kappa=0).
Ring-make energy is near-zero at λ=1 (potential_swap end, alpha=1).

At λ=1 the system is at the ring-open end state; the hard-hard correction
term in the CustomBondForce is small because the pair is at nonbonded
Expand All @@ -329,25 +322,25 @@ def test_ring_make_inactive_at_lambda_one(reverse_dynamics):
e = _force_energy_kcal(reverse_dynamics, 1.0, "ring-make")
assert abs(e) < _ACTIVE_THRESHOLD, (
f"ring-make energy {e:.4f} kcal/mol at λ=1 exceeds threshold "
f"{_ACTIVE_THRESHOLD} kcal/mol (kappa should be 0)"
f"{_ACTIVE_THRESHOLD} kcal/mol (alpha should be 1)"
)


# ── energy symmetry tests ─────────────────────────────────────────────────────
#
# The invariant ring_break_morph().reverse() == reverse_ring_break_morph() means
# that the softcore kappa/alpha values at (forward, λ) and (reverse, 1-λ) are
# equal. Both forces act on the same bond (the original ring_breaking_bond,
# which swap_end_states=True maps to ring_making_pairs), so the energies must
# also match. The hard-hard correction appears identically on both sides and
# that the softcore alpha values at (forward, λ) and (reverse, 1-λ) are equal.
# Both forces act on the same bond (the original ring_breaking_bond, which
# swap_end_states=True maps to ring_making_pairs), so the energies must also
# match. The hard-hard correction appears identically on both sides and
# cancels in the comparison, making this test robust to formula changes.
#
# Test points span zero and non-zero energy regions:
# λ=0.0 → forward kappa=0, reverse at 1-λ=1.0 kappa=0 (both ≈0)
# λ=0.55 → forward restraints_off (kappa=0.65), reverse restraints_off at 0.45 (kappa=0.65)
# λ=2/3 → forward morph start (kappa=1), reverse restraints_off start at 1/3 (kappa=1)
# λ=0.85 → forward morph (kappa=1), reverse reversed-morph at 0.15 (kappa=1)
# λ=1.0 → forward morph end (kappa=1), reverse at 0.0 reversed-morph (kappa=1)
# λ=0.0 → forward alpha=1, reverse at 1-λ=1.0 alpha=1 (both ≈0)
# λ=0.55 → forward restraints_off (alpha=0.35), reverse restraints_off at 0.45 (alpha=0.35)
# λ=2/3 → forward morph start (alpha=0), reverse restraints_off start at 1/3 (alpha=0)
# λ=0.85 → forward morph (alpha=0), reverse reversed-morph at 0.15 (alpha=0)
# λ=1.0 → forward morph end (alpha=0), reverse at 0.0 reversed-morph (alpha=0)


@pytest.mark.parametrize("lam", [0.0, 0.55, 2 / 3, 0.85, 1.0])
Expand All @@ -356,7 +349,7 @@ def test_energy_symmetry_forward_reverse(forward_dynamics, reverse_dynamics, lam
Single-point energy symmetry: E_ring_break_forward(λ) == E_ring_make_reverse(1-λ).

Verifies that reverse_ring_break_morph() == ring_break_morph().reverse() and
that the mirrored kappa/alpha produce identical corrections on the same bond.
that the mirrored alpha values produce identical corrections on the same bond.
"""
e_fwd = _force_energy_kcal(forward_dynamics, lam, "ring-break")
e_rev = _force_energy_kcal(reverse_dynamics, 1.0 - lam, "ring-make")
Expand All @@ -373,8 +366,8 @@ def test_schedule_symmetry():
reverse_ring_break_morph() must equal ring_break_morph().reverse().

Checks that the simplified implementation produces identical schedules by
comparing kappa values at a dense grid of lambda points using the default
initial/final values that lambdalever passes for ring-break kappa.
comparing lever values at a dense grid of lambda points using the default
initial/final values that lambdalever passes for each lever.
"""
from somd2._utils._schedules import ring_break_morph, reverse_ring_break_morph

Expand All @@ -385,10 +378,8 @@ def test_schedule_symmetry():
test_lambdas = [i / 20 for i in range(21)]
for lam in test_lambdas:
for force, lever, init, fin in [
("ring-break", "kappa", 0.0, 1.0),
("ring-break", "alpha", 1.0, 0.0),
("ring-break", "coul_kappa", 0.0, 1.0),
("ring-make", "kappa", 1.0, 0.0),
("ring-make", "alpha", 0.0, 1.0),
("ring-make", "coul_kappa", 1.0, 0.0),
]:
Expand Down