Fixes RMSD force management

docs/readme.rst

@@ -6,11 +6,11 @@ UFEDMM is an OpenMM extension for applying the Unified Free Energy Dynamics meth
 
 * Free software: MIT license
 
+============
 Installation
 ============
 
-UFEDMM is available as a conda package installable from the
-[craabreu](https://anaconda.org/craabreu) channel.
+UFEDMM is available as a conda package installable from the `craabreu <https://anaconda.org/craabreu>`_ conda channel.
 To install it, either run:
 
 .. code-block:: bash
@@ -29,15 +29,10 @@ To use UFEDMM in your own Python script or Jupyter notebook, import it as follow
 
     import ufedmm
 
-Documentation
-=============
-
-https://craabreu.github.io/ufedmm/
-
+============
 Contributing
 ============
 
-Feedback, bug reports, and feature requests are welcome. They should be
-submitted to https://github.com/craabreu/ufedmm/issues.
-
-Pull requests are greatly appreciated!
+Feedback, bug reports, and feature requests are welcome. They should be submitted to
+UFEDMM's `issue tracker <https://github.com/craabreu/ufedmm/issues>`_. Pull requests are
+alo greatly appreciated!

ufedmm/tests/test_ufedmm.py

@@ -5,12 +5,14 @@
 # Import package, test suite, and other packages as needed
 import io
 import sys
+from copy import deepcopy
 
 import openmm
 import pytest
 from openmm import unit
 
 import ufedmm
+import numpy as np
 
 
 def ufed_model(
@@ -171,3 +173,25 @@ def test_well_tempered_metadynamics():
     simulation.step(100)
     energy = simulation.context.getState(getEnergy=True).getPotentialEnergy()
     assert energy / energy.unit == pytest.approx(153.17, abs=0.1)
+
+
+def test_rmsd_forces():
+    model, ufed = ufed_model(constraints=None)
+    rmsd = openmm.RMSDForce(  # Level 1
+        model.positions, np.arange(model.system.getNumParticles())
+    )
+    cvforce = openmm.CustomCVForce("rmsd + cvforce")
+    cvforce.addCollectiveVariable("rmsd", deepcopy(rmsd))  # Level 2
+    inner_cvforce = openmm.CustomCVForce("rmsd")
+    inner_cvforce.addCollectiveVariable("rmsd", deepcopy(rmsd))  # Level 3
+    cvforce.addCollectiveVariable("cvforce", inner_cvforce)
+    model.system.addForce(rmsd)
+    model.system.addForce(cvforce)
+    integrator = ufedmm.MiddleMassiveNHCIntegrator(
+        300 * unit.kelvin, 10 * unit.femtoseconds, 1 * unit.femtoseconds
+    )
+    platform = openmm.Platform.getPlatformByName("Reference")
+    simulation = ufed.simulation(model.topology, model.system, integrator, platform)
+    simulation.context.setPositions(model.positions)
+    simulation.context.setVelocitiesToTemperature(300 * unit.kelvin, 1234)
+    simulation.step(1)

ufedmm/ufedmm.py

@@ -47,22 +47,27 @@ def _standardized(quantity):
 def _update_RMSD_forces(system):
     N = system.getNumParticles()
 
-    def update_RMSDForce(force):
+    def _update_RMSDForce(force):
         positions = force.getReferencePositions()._value
         if len(positions) >= N:
             positions = positions[:N]
         else:
             positions += [openmm.Vec3(0, 0, 0)] * (N - len(positions))
         force.setReferencePositions(positions)
 
+    def _update_RMSD_forces_in_cvforce(cvforce):
+        for index in range(cvforce.getNumCollectiveVariables()):
+            force = cvforce.getCollectiveVariable(index)
+            if isinstance(force, openmm.RMSDForce):
+                _update_RMSDForce(force)
+            elif isinstance(force, openmm.CustomCVForce):
+                _update_RMSD_forces_in_cvforce(force)
+
     for force in system.getForces():
         if isinstance(force, openmm.RMSDForce):
-            update_RMSDForce(force)
+            _update_RMSDForce(force)
         elif isinstance(force, openmm.CustomCVForce):
-            for index in range(force.getNumCollectiveVariables()):
-                cv = force.getCollectiveVariable(index)
-                if isinstance(cv, openmm.RMSDForce):
-                    update_RMSDForce(cv)
+            _update_RMSD_forces_in_cvforce(force)
 
 
 class CollectiveVariable(object):