Merge pull request #412 from OpenBioSim/feature_openmm_vsite_water

Add support for 4- and 5-point water models with OpenMM

Author: lohedges

corelib/src/libs/SireCAS/lambdaschedule.cpp

@@ -65,7 +65,7 @@ static RegisterMetaType<LambdaSchedule> r_schedule;
 
 QDataStream &operator<<(QDataStream &ds, const LambdaSchedule &schedule)
 {
-    writeHeader(ds, r_schedule, 3);
+    writeHeader(ds, r_schedule, 4);
 
     SharedDataStream sds(ds);
 
@@ -75,6 +75,7 @@ QDataStream &operator<<(QDataStream &ds, const LambdaSchedule &schedule)
         << schedule.default_equations
         << schedule.stage_equations
         << schedule.mol_schedules
+        << schedule.coupled_levers
         << static_cast<const Property &>(schedule);
 
     return ds;
@@ -91,21 +92,30 @@ QDataStream &operator>>(QDataStream &ds, LambdaSchedule &schedule)
 {
     VersionID v = readHeader(ds, r_schedule);
 
-    if (v == 1 or v == 2 or v == 3)
+    if (v == 1 or v == 2 or v == 3 or v == 4)
     {
         SharedDataStream sds(ds);
 
         sds >> schedule.constant_values;
 
-        if (v == 3)
+        if (v >= 3)
             sds >> schedule.force_names;
 
         sds >> schedule.lever_names >> schedule.stage_names >>
             schedule.default_equations >> schedule.stage_equations;
 
-        if (v == 2 or v == 3)
+        if (v >= 2)
             sds >> schedule.mol_schedules;
 
+        if (v >= 4)
+            sds >> schedule.coupled_levers;
+        else
+        {
+            // Populate the default coupling for streams written before v4
+            schedule.coupled_levers[_get_lever_name("cmap", "cmap_grid")] =
+                _get_lever_name("torsion", "torsion_k");
+        }
+
         if (v < 3)
         {
             // need to make sure that the lever names are namespaced
@@ -146,13 +156,17 @@ QDataStream &operator>>(QDataStream &ds, LambdaSchedule &schedule)
         }
     }
     else
-        throw version_error(v, "1, 2, 3", r_schedule, CODELOC);
+        throw version_error(v, "1, 2, 3, 4", r_schedule, CODELOC);
 
     return ds;
 }
 
 LambdaSchedule::LambdaSchedule() : ConcreteProperty<LambdaSchedule, Property>()
 {
+    // By default, cmap_grid falls back to torsion_k so that customising
+    // torsion scaling automatically keeps the CMAP correction in sync.
+    coupled_levers[_get_lever_name("cmap", "cmap_grid")] =
+        _get_lever_name("torsion", "torsion_k");
 }
 
 LambdaSchedule::LambdaSchedule(const LambdaSchedule &other)
@@ -162,7 +176,8 @@ LambdaSchedule::LambdaSchedule(const LambdaSchedule &other)
       force_names(other.force_names),
       lever_names(other.lever_names), stage_names(other.stage_names),
       default_equations(other.default_equations),
-      stage_equations(other.stage_equations)
+      stage_equations(other.stage_equations),
+      coupled_levers(other.coupled_levers)
 {
 }
 
@@ -1006,6 +1021,32 @@ void LambdaSchedule::removeEquation(const QString &stage,
     this->stage_equations[idx].remove(_get_lever_name(force, lever));
 }
 
+/** Couple the lever 'force::lever' to 'fallback_force::fallback_lever'.
+ *  If no custom equation has been set for 'force::lever' at any stage,
+ *  the equation for 'fallback_force::fallback_lever' will be used instead
+ *  of the stage default. This is a single level of indirection — the
+ *  fallback lever is not itself followed further.
+ *
+ *  By default, 'cmap::cmap_grid' is coupled to 'torsion::torsion_k' so
+ *  that custom torsion schedules automatically keep the CMAP correction
+ *  in sync.
+ */
+void LambdaSchedule::coupleLever(const QString &force, const QString &lever,
+                                  const QString &fallback_force,
+                                  const QString &fallback_lever)
+{
+    coupled_levers[_get_lever_name(force, lever)] =
+        _get_lever_name(fallback_force, fallback_lever);
+}
+
+/** Remove any coupling for the lever 'force::lever', reverting it to
+ *  use the stage default equation when no custom equation is set.
+ */
+void LambdaSchedule::removeCoupledLever(const QString &force, const QString &lever)
+{
+    coupled_levers.remove(_get_lever_name(force, lever));
+}
+
 /** Return whether or not the specified 'lever' in the specified 'force'
  *  at the specified 'stage' has a custom equation set for it
  */
@@ -1077,6 +1118,32 @@ SireCAS::Expression LambdaSchedule::_getEquation(int stage,
         return it.value();
     }
 
+    // Check coupled levers: if this lever is coupled to another, try to
+    // find a custom equation for the coupled lever before falling back to
+    // the stage default. This is a single level of indirection (no recursion)
+    // to prevent loops.
+    auto coupled_it = this->coupled_levers.find(lever_name);
+
+    if (coupled_it != this->coupled_levers.end())
+    {
+        const auto &coupled_name = coupled_it.value();
+        const int sep = coupled_name.indexOf("::");
+        const QString coupled_force = sep >= 0 ? coupled_name.left(sep) : "*";
+        const QString coupled_lever_part = sep >= 0 ? coupled_name.mid(sep + 2) : coupled_name;
+
+        it = equations.find(coupled_name);
+        if (it != equations.end())
+            return it.value();
+
+        it = equations.find(_get_lever_name(coupled_force, "*"));
+        if (it != equations.end())
+            return it.value();
+
+        it = equations.find(_get_lever_name("*", coupled_lever_part));
+        if (it != equations.end())
+            return it.value();
+    }
+
     // we don't have any match, so return the default equation for this stage
     return this->default_equations[stage];
 }

corelib/src/libs/SireCAS/lambdaschedule.h

@@ -164,6 +164,12 @@ namespace SireCAS
                                       const QString &force = "*",
                                       const QString &lever = "*") const;
 
+        void coupleLever(const QString &force, const QString &lever,
+                         const QString &fallback_force,
+                         const QString &fallback_lever);
+
+        void removeCoupledLever(const QString &force, const QString &lever);
+
         SireCAS::Expression getEquation(const QString &stage = "*",
                                         const QString &force = "*",
                                         const QString &lever = "*") const;
@@ -248,6 +254,12 @@ namespace SireCAS
             particular stage */
         QVector<QHash<QString, SireCAS::Expression>> stage_equations;
 
+        /** Coupled lever fallbacks: if a lever (force::lever) has no custom
+            equation set, use the equation for the paired lever instead of
+            falling straight back to the stage default. Stored as
+            _get_lever_name(force, lever) → _get_lever_name(fb_force, fb_lever). */
+        QHash<QString, QString> coupled_levers;
+
         /** The symbol used to represent the lambda value */
         static SireCAS::Symbol lambda_symbol;
 

doc/source/changelog.rst

@@ -29,6 +29,10 @@ organisation on `GitHub <https://github.com/openbiosim/sire>`__.
 
 * Fix hang in ``sire.load`` function when shared GROMACS topology path is missing.
 
+* Add support for 4- and 5-point water models in the OpenMM conversion layer.
+
+* Add functionality for coupling one lambda lever to another.
+
 `2025.4.0 <https://github.com/openbiosim/sire/compare/2025.3.0...2025.4.0>`__ - February 2026
 ---------------------------------------------------------------------------------------------
 

doc/source/tutorial/part07/02_levers.rst

@@ -159,6 +159,31 @@ In this case, as we have a perturbable system, the Force objects used are;
   It uses parameters that are controlled by the ``charge``, ``sigma``, ``epsilon``,
   ``alpha``, ``kappa``, ``charge_scale`` and ``lj_scale`` levers.
 
+For systems that use a force field with CMAP backbone torsion corrections (such
+as AMBER ff19SB), an additional Force object is present:
+
+* ``cmap``: `OpenMM::CMAPTorsionForce <http://docs.openmm.org/latest/api-c++/generated/CMAPTorsionForce.html>`__.
+  This models the cross-term energy correction (CMAP) for backbone φ/ψ dihedral
+  pairs. It uses a parameter controlled by the ``cmap_grid`` lever, which
+  interpolates the two-dimensional energy grid from its λ=0 value to its λ=1 value.
+
+.. note::
+
+   By default, ``cmap::cmap_grid`` is coupled to ``torsion::torsion_k``:
+   if you set a custom equation for ``torsion_k``, the same equation will
+   automatically be used for ``cmap_grid``, keeping the CMAP correction
+   in sync with the torsion perturbation. You can override this by
+   setting an explicit equation for ``cmap::cmap_grid``::
+
+      # freeze CMAP at its λ=0 grid while still perturbing torsions
+      s.set_equation(stage="morph", force="cmap", lever="cmap_grid",
+                     equation=s.initial())
+
+   To remove the coupling entirely (so that ``cmap_grid`` falls back to
+   the stage default independently of ``torsion_k``)::
+
+      s.remove_coupled_lever(force="cmap", lever="cmap_grid")
+
 Some levers, like ``bond_length``, are used only by a single Force object.
 However, others, like ``charge``, are used by multiple Force objects.
 

src/sire/mol/_dynamics.py

@@ -492,11 +492,11 @@ def _exit_dynamics_block(
                         zip(lambda_windows, rest2_scale_factors)
                     ):
                         if lambda_value != sim_lambda_value:
+                            key = f"{lambda_value:.5f}"
                             if (
                                 not has_lambda_index
                                 or abs(lambda_index - i) <= num_energy_neighbours
                             ):
-                                key = f"{lambda_value:.5f}"
                                 self._omm_mols.set_lambda(
                                     lambda_value,
                                     rest2_scale=rest2_scale,

tests/cas/test_lambdaschedule.py

@@ -154,3 +154,105 @@ def test_lambdaschedule():
 def test_has_force_specific_equation(force, lever, contained):
     l = sr.cas.LambdaSchedule.standard_decouple()
     assert l.has_force_specific_equation("decouple", force, lever) == contained
+
+
+def test_coupled_lever_default():
+    """cmap::cmap_grid should follow torsion::torsion_k by default."""
+    l = sr.cas.LambdaSchedule.standard_morph()
+    morph_equation = (1 - l.lam()) * l.initial() + l.lam() * l.final()
+
+    # With no custom equations, both should return the stage default.
+    _assert_same_equation(
+        l.lam(),
+        l.get_equation(stage="morph", force="torsion", lever="torsion_k"),
+        morph_equation,
+    )
+    _assert_same_equation(
+        l.lam(),
+        l.get_equation(stage="morph", force="cmap", lever="cmap_grid"),
+        morph_equation,
+    )
+
+
+def test_coupled_lever_follows_torsion_k():
+    """Setting a custom torsion_k equation should automatically apply to cmap_grid."""
+    l = sr.cas.LambdaSchedule.standard_morph()
+    custom_eq = l.lam() ** 2 * l.final() + (1 - l.lam() ** 2) * l.initial()
+
+    l.set_equation(
+        stage="morph", force="torsion", lever="torsion_k", equation=custom_eq
+    )
+
+    # cmap_grid should now follow the custom torsion_k equation via coupling.
+    _assert_same_equation(
+        l.lam(),
+        l.get_equation(stage="morph", force="cmap", lever="cmap_grid"),
+        custom_eq,
+    )
+
+
+def test_coupled_lever_explicit_override():
+    """An explicit cmap_grid equation should take precedence over the coupling."""
+    l = sr.cas.LambdaSchedule.standard_morph()
+    torsion_eq = l.lam() ** 2 * l.final() + (1 - l.lam() ** 2) * l.initial()
+    cmap_eq = l.initial()  # freeze CMAP at λ=0
+
+    l.set_equation(
+        stage="morph", force="torsion", lever="torsion_k", equation=torsion_eq
+    )
+    l.set_equation(stage="morph", force="cmap", lever="cmap_grid", equation=cmap_eq)
+
+    # cmap_grid should use its own explicit equation, not torsion_k's.
+    _assert_same_equation(
+        l.lam(),
+        l.get_equation(stage="morph", force="cmap", lever="cmap_grid"),
+        cmap_eq,
+    )
+    # torsion_k should be unaffected.
+    _assert_same_equation(
+        l.lam(),
+        l.get_equation(stage="morph", force="torsion", lever="torsion_k"),
+        torsion_eq,
+    )
+
+
+def test_remove_coupled_lever():
+    """Removing the coupling makes cmap_grid fall back to the stage default."""
+    l = sr.cas.LambdaSchedule.standard_morph()
+    morph_equation = (1 - l.lam()) * l.initial() + l.lam() * l.final()
+    custom_eq = l.lam() ** 2 * l.final() + (1 - l.lam() ** 2) * l.initial()
+
+    l.set_equation(
+        stage="morph", force="torsion", lever="torsion_k", equation=custom_eq
+    )
+    l.remove_coupled_lever(force="cmap", lever="cmap_grid")
+
+    # cmap_grid should now use the stage default, not follow torsion_k.
+    _assert_same_equation(
+        l.lam(),
+        l.get_equation(stage="morph", force="cmap", lever="cmap_grid"),
+        morph_equation,
+    )
+
+
+def test_couple_lever_custom():
+    """coupleLever can set an arbitrary coupling between levers."""
+    l = sr.cas.LambdaSchedule.standard_morph()
+    custom_eq = l.lam() ** 2 * l.final() + (1 - l.lam() ** 2) * l.initial()
+
+    # Couple bond_k to torsion_k (unusual, but should work).
+    l.couple_lever(
+        force="bond",
+        lever="bond_k",
+        fallback_force="torsion",
+        fallback_lever="torsion_k",
+    )
+    l.set_equation(
+        stage="morph", force="torsion", lever="torsion_k", equation=custom_eq
+    )
+
+    _assert_same_equation(
+        l.lam(),
+        l.get_equation(stage="morph", force="bond", lever="bond_k"),
+        custom_eq,
+    )