Improves public API (#71)

* Improved getEffectiveMass example

* Removed setUnit method

* Removed getArgument method from public API

* Added example for getValue method

* Included addToSystem method to all CVs

* Refactored tests and doctests

* Fixed code formatting

cvpack/atomic_function.py

@@ -93,8 +93,8 @@ class AtomicFunction(openmm.CustomCompoundBondForce, BaseCustomFunction):
         ...     k = 1000 * unit.kilojoules_per_mole/unit.radian**2,
         ...     theta0 = [np.pi/2, np.pi/3] * unit.radian,
         ... )
-        >>> [model.system.addForce(f) for f in [angle1, angle2, colvar]]
-        [5, 6, 7]
+        >>> for cv in [angle1, angle2, colvar]:
+        ...     cv.addToSystem(model.system)
         >>> integrator = openmm.VerletIntegrator(0)
         >>> platform = openmm.Platform.getPlatformByName('Reference')
         >>> context = openmm.Context(model.system, integrator, platform)

cvpack/attraction_strength.py

@@ -128,8 +128,7 @@ class AttractionStrength(openmm.CustomNonbondedForce, BaseCollectiveVariable):
     >>> host = [a.index for a in model.topology.atoms() if a.residue.name == "CUC"]
     >>> forces = {f.getName(): f for f in model.system.getForces()}
     >>> cv1 = cvpack.AttractionStrength(guest, host, forces["NonbondedForce"])
-    >>> _ = cv1.setUnusedForceGroup(0, model.system)
-    >>> _ = model.system.addForce(cv1)
+    >>> cv1.addToSystem(model.system)
     >>> platform = openmm.Platform.getPlatformByName("Reference")
     >>> integrator = openmm.VerletIntegrator(1.0 * mmunit.femtoseconds)
     >>> context = openmm.Context(model.system, integrator, platform)
@@ -139,15 +138,13 @@ class AttractionStrength(openmm.CustomNonbondedForce, BaseCollectiveVariable):
 
     >>> water = [a.index for a in model.topology.atoms() if a.residue.name == "HOH"]
     >>> cv2 = cvpack.AttractionStrength(guest, water, forces["NonbondedForce"])
-    >>> _ = cv2.setUnusedForceGroup(0, model.system)
-    >>> _ = model.system.addForce(cv2)
+    >>> cv2.addToSystem(model.system)
     >>> context.reinitialize(preserveState=True)
     >>> print(cv2.getValue(context))
     2063.3... dimensionless
 
     >>> cv3 = cvpack.AttractionStrength(guest, host, forces["NonbondedForce"], water)
-    >>> _ = cv3.setUnusedForceGroup(0, model.system)
-    >>> _ = model.system.addForce(cv3)
+    >>> cv3.addToSystem(model.system)
     >>> context.reinitialize(preserveState=True)
     >>> print(cv3.getValue(context))
     2849.17... dimensionless
@@ -157,8 +154,7 @@ class AttractionStrength(openmm.CustomNonbondedForce, BaseCollectiveVariable):
     >>> cv4 = cvpack.AttractionStrength(
     ...     guest, host, forces["NonbondedForce"], water, contrastScaling=0.5
     ... )
-    >>> _ = cv4.setUnusedForceGroup(0, model.system)
-    >>> _ = model.system.addForce(cv4)
+    >>> cv4.addToSystem(model.system)
     >>> context.reinitialize(preserveState=True)
     >>> print(cv4.getValue(context))
     3880.8... dimensionless

cvpack/centroid_function.py

@@ -121,10 +121,7 @@ class CentroidFunction(openmm.CustomCentroidBondForce, BaseCustomFunction):
         >>> res_coord = cvpack.ResidueCoordination(
         ...     residues[115:124], residues[126:135], stepFunction="step(1-x)"
         ... )
-        >>> res_coord.setUnusedForceGroup(0, model.system)
-        1
-        >>> model.system.addForce(res_coord)
-        6
+        >>> res_coord.addToSystem(model.system)
         >>> integrator = openmm.VerletIntegrator(0)
         >>> platform = openmm.Platform.getPlatformByName('Reference')
         >>> context = openmm.Context(model.system, integrator, platform)
@@ -142,10 +139,7 @@ class CentroidFunction(openmm.CustomCentroidBondForce, BaseCustomFunction):
         ...    atoms[115:124] + atoms[126:135],
         ...    list(it.product(range(9), range(9, 18))),
         ... )
-        >>> colvar.setUnusedForceGroup(0, model.system)
-        2
-        >>> model.system.addForce(colvar)
-        7
+        >>> colvar.addToSystem(model.system)
         >>> context.reinitialize(preserveState=True)
         >>> print(colvar.getValue(context))
         33.0 dimensionless

cvpack/composite_rmsd.py

@@ -105,10 +105,7 @@ class CompositeRMSD(CompositeRMSDForce, BaseCollectiveVariable):
     ...     )
     ... except ImportError:
     ...     pytest.skip("openmm-cpp-forces is not installed")
-    >>> composite_rmsd.setUnusedForceGroup(0, model.system)
-    1
-    >>> model.system.addForce(composite_rmsd)
-    5
+    >>> composite_rmsd.addToSystem(model.system)
     >>> context = mm.Context(
     ...     model.system,
     ...     mm.VerletIntegrator(1.0 * unit.femtoseconds),

cvpack/cvpack.py

@@ -128,9 +128,9 @@ def _registerCV(
         """
         cls = self.__class__
         self.setName(cls.__name__)
-        self.setUnit(unit)
+        self._unit = unit
         self._mass_unit = mmunit.dalton * (mmunit.nanometers / self.getUnit()) ** 2
-        arguments, _ = self.getArguments()
+        arguments, _ = self._getArguments()
         self._args = dict(zip(arguments, args))
         self._args.update(kwargs)
 
@@ -149,21 +149,22 @@ def _registerPeriod(self, period: float) -> None:
         self._period = period
 
     @classmethod
-    def getArguments(cls) -> t.Tuple[collections.OrderedDict, collections.OrderedDict]:
+    def _getArguments(cls) -> t.Tuple[collections.OrderedDict, collections.OrderedDict]:
         """
         Inspect the arguments needed for constructing an instance of this collective
         variable.
 
         Returns
         -------
+        OrderedDict
             A dictionary with the type annotations of all arguments
-
+        OrderedDict
             A dictionary with the default values of optional arguments
 
         Example
         -------
             >>> import cvpack
-            >>> args, defaults = cvpack.RadiusOfGyration.getArguments()
+            >>> args, defaults = cvpack.RadiusOfGyration._getArguments()
             >>> for name, annotation in args.items():
             ...     print(f"{name}: {annotation}")
             group: typing.Iterable[int]
@@ -180,16 +181,32 @@ def getArguments(cls) -> t.Tuple[collections.OrderedDict, collections.OrderedDic
                 defaults[name] = parameter.default
         return arguments, defaults
 
-    def setUnit(self, unit: mmunit.Unit) -> None:
+    def _setUnusedForceGroup(self, system: openmm.System) -> None:
         """
-        Set the unit of measurement of this collective variable.
+        Set the force group of this collective variable to the one at a given position
+        in the ascending ordered list of unused force groups in an :OpenMM:`System`.
+
+        .. note::
+
+            Evaluating a collective variable (see :meth:`getValue`) or computing its
+            effective mass (see :meth:`getEffectiveMass`) is more efficient when the
+            collective variable is the only force in its own force group.
 
         Parameters
         ----------
-            unit
-                The unit of measurement of this collective variable
+            system
+                The system to search for unused force groups
+
+        Raises
+        ------
+            RuntimeError
+                If all force groups are already in use
         """
-        self._unit = unit
+        used_groups = {force.getForceGroup() for force in system.getForces()}
+        new_group = next(filter(lambda i: i not in used_groups, range(32)), None)
+        if new_group is None:
+            raise RuntimeError("All force groups are already in use.")
+        self.setForceGroup(new_group)
 
     def getUnit(self) -> mmunit.Unit:
         """
@@ -209,42 +226,23 @@ def getPeriod(self) -> t.Optional[mmunit.SerializableQuantity]:
             return None
         return mmunit.SerializableQuantity(self._period, self.getUnit())
 
-    def setUnusedForceGroup(self, position: int, system: openmm.System) -> int:
+    def addToSystem(
+        self, system: openmm.System, setUnusedForceGroup: bool = True
+    ) -> None:
         """
-        Set the force group of this collective variable to the one at a given position
-        in the ascending ordered list of unused force groups in an :OpenMM:`System`.
-
-        .. note::
-
-            Evaluating a collective variable (see :meth:`getValue`) or computing its
-            effective mass (see :meth:`getEffectiveMass`) is more efficient when the
-            collective variable is the only force in its own force group.
+        Add this collective variable to an :OpenMM:`System`.
 
         Parameters
         ----------
-            position
-                The position of the force group in the ascending ordered list of unused
-                force groups in the system
-            system
-                The system to search for unused force groups
-
-        Returns
-        -------
-            The index of the force group that was set
-
-        Raises
-        ------
-            RuntimeError
-                If all force groups are already in use
+        system
+            The system to which this collective variable should be added
+        setUnusedForceGroup
+            If True, the force group of this collective variable will be set to the
+            first available force group in the system
         """
-        free_groups = sorted(
-            set(range(32)) - {force.getForceGroup() for force in system.getForces()}
-        )
-        if not free_groups:
-            raise RuntimeError("All force groups are already in use.")
-        new_group = free_groups[position]
-        self.setForceGroup(new_group)
-        return new_group
+        if setUnusedForceGroup:
+            self._setUnusedForceGroup(system)
+        system.addForce(self)
 
     def getValue(self, context: openmm.Context) -> mmunit.Quantity:
         """
@@ -257,12 +255,43 @@ def getValue(self, context: openmm.Context) -> mmunit.Quantity:
 
         Parameters
         ----------
-            context
-                The context at which this collective variable should be evaluated
+        context
+            The context at which this collective variable should be evaluated
 
         Returns
         -------
+        unit.Quantity
             The value of this collective variable at the given context
+
+
+        Example
+        -------
+        In this example, we compute the values of the backbone dihedral angles and
+        the radius of gyration of an alanine dipeptide molecule in water:
+
+        >>> import cvpack
+        >>> import openmm
+        >>> from openmmtools import testsystems
+        >>> model = testsystems.AlanineDipeptideExplicit()
+        >>> top = model.mdtraj_topology
+        >>> backbone_atoms = top.select("name N C CA and resid 1 2")
+        >>> phi = cvpack.Torsion(*backbone_atoms[0:4])
+        >>> psi = cvpack.Torsion(*backbone_atoms[1:5])
+        >>> radius_of_gyration = cvpack.RadiusOfGyration(
+        ...     top.select('not water')
+        ... )
+        >>> for cv in [phi, psi, radius_of_gyration]:
+        ...     cv.addToSystem(model.system)
+        >>> context = openmm.Context(
+        ...     model.system, openmm.VerletIntegrator(0)
+        ... )
+        >>> context.setPositions(model.positions)
+        >>> print(phi.getValue(context))
+        3.1415... rad
+        >>> print(psi.getValue(context))
+        3.1415... rad
+        >>> print(radius_of_gyration.getValue(context))
+        0.29514... nm
         """
         state = get_single_force_state(self, context, getEnergy=True)
         value = value_in_md_units(state.getPotentialEnergy())
@@ -291,36 +320,41 @@ def getEffectiveMass(self, context: openmm.Context) -> mmunit.Quantity:
 
         Parameters
         ----------
-            context
-                The context at which this collective variable's effective mass should be
-                evaluated
+        context
+            The context at which this collective variable's effective mass should be
+            evaluated
 
         Returns
         -------
+        unit.Quantity
             The effective mass of this collective variable at the given context
 
         Example
         -------
+            In this example, we compute the effective masses of the backbone dihedral
+            angles and the radius of gyration of an alanine dipeptide molecule in water:
+
             >>> import cvpack
             >>> import openmm
             >>> from openmmtools import testsystems
-            >>> model = testsystems.AlanineDipeptideImplicit()
-            >>> peptide = [
-            ...     a.index
-            ...     for a in model.topology.atoms()
-            ...     if a.residue.name != 'HOH'
-            ... ]
-            >>> radius_of_gyration = cvpack.RadiusOfGyration(peptide)
-            >>> radius_of_gyration.setForceGroup(1)
-            >>> radius_of_gyration.setUnusedForceGroup(0, model.system)
-            1
-            >>> model.system.addForce(radius_of_gyration)
-            6
-            >>> platform = openmm.Platform.getPlatformByName('Reference')
+            >>> model = testsystems.AlanineDipeptideExplicit()
+            >>> top = model.mdtraj_topology
+            >>> backbone_atoms = top.select("name N C CA and resid 1 2")
+            >>> phi = cvpack.Torsion(*backbone_atoms[0:4])
+            >>> psi = cvpack.Torsion(*backbone_atoms[1:5])
+            >>> radius_of_gyration = cvpack.RadiusOfGyration(
+            ...     top.select('not water')
+            ... )
+            >>> for cv in [phi, psi, radius_of_gyration]:
+            ...     cv.addToSystem(model.system)
             >>> context = openmm.Context(
-            ...     model.system,openmm.VerletIntegrator(0), platform
+            ...     model.system, openmm.VerletIntegrator(0)
             ... )
             >>> context.setPositions(model.positions)
+            >>> print(phi.getEffectiveMass(context))
+            0.05119... nm**2 Da/(rad**2)
+            >>> print(psi.getEffectiveMass(context))
+            0.05186... nm**2 Da/(rad**2)
             >>> print(radius_of_gyration.getEffectiveMass(context))
             30.946... Da
         """

cvpack/helix_angle_content.py

@@ -84,10 +84,7 @@ class HelixAngleContent(openmm.CustomAngleForce, BaseCollectiveVariable):
         >>> print(*[r.name for r in residues])  # doctest: +ELLIPSIS
         LYS ASP GLU ... ILE LEU ARG
         >>> helix_content = cvpack.HelixAngleContent(residues)
-        >>> helix_content.setUnusedForceGroup(0, model.system)
-        1
-        >>> model.system.addForce(helix_content)
-        6
+        >>> helix_content.addToSystem(model.system)
         >>> platform = openmm.Platform.getPlatformByName('Reference')
         >>> integrator = openmm.VerletIntegrator(0)
         >>> context = openmm.Context(model.system, integrator, platform)

cvpack/helix_hbond_content.py

@@ -72,10 +72,7 @@ class HelixHBondContent(openmm.CustomBondForce, BaseCollectiveVariable):
         >>> print(*[r.name for r in residues])
         LYS ASP GLU ... ILE LEU ARG
         >>> helix_content = cvpack.HelixHBondContent(residues)
-        >>> helix_content.setUnusedForceGroup(0, model.system)
-        1
-        >>> model.system.addForce(helix_content)
-        6
+        >>> helix_content.addToSystem(model.system)
         >>> platform = openmm.Platform.getPlatformByName('Reference')
         >>> integrator = openmm.VerletIntegrator(0)
         >>> context = openmm.Context(model.system, integrator, platform)

cvpack/helix_rmsd_content.py

@@ -111,10 +111,7 @@ class HelixRMSDContent(BaseRMSDContent):
         ... )
         >>> helix_content.getNumResidueBlocks()
         16
-        >>> helix_content.setUnusedForceGroup(0, model.system)
-        1
-        >>> model.system.addForce(helix_content)
-        6
+        >>> helix_content.addToSystem(model.system)
         >>> platform = openmm.Platform.getPlatformByName('Reference')
         >>> integrator = openmm.VerletIntegrator(0)
         >>> context = openmm.Context(model.system, integrator, platform)

cvpack/helix_torsion_content.py

@@ -94,10 +94,7 @@ class HelixTorsionContent(openmm.CustomTorsionForce, BaseCollectiveVariable):
         >>> print(*[r.name for r in residues])
         LYS ASP GLU ... ILE LEU ARG
         >>> helix_content = cvpack.HelixTorsionContent(residues)
-        >>> helix_content.setUnusedForceGroup(0, model.system)
-        1
-        >>> model.system.addForce(helix_content)
-        6
+        >>> helix_content.addToSystem(model.system)
         >>> platform = openmm.Platform.getPlatformByName('Reference')
         >>> integrator = openmm.VerletIntegrator(0)
         >>> context = openmm.Context(model.system, integrator, platform)

cvpack/number_of_contacts.py

@@ -100,10 +100,7 @@ class NumberOfContacts(openmm.CustomNonbondedForce, BaseCollectiveVariable):
     ...     forces["NonbondedForce"],
     ...     stepFunction="step(1-x)",
     ... )
-    >>> nc.setUnusedForceGroup(0, model.system)
-    1
-    >>> model.system.addForce(nc)
-    5
+    >>> nc.addToSystem(model.system)
     >>> platform = openmm.Platform.getPlatformByName("Reference")
     >>> integrator = openmm.VerletIntegrator(1.0 * mmunit.femtoseconds)
     >>> context = openmm.Context(model.system, integrator, platform)
@@ -117,10 +114,7 @@ class NumberOfContacts(openmm.CustomNonbondedForce, BaseCollectiveVariable):
     ...     stepFunction="step(1-x)",
     ...     reference=context,
     ... )
-    >>> nc_normalized.setUnusedForceGroup(0, model.system)
-    2
-    >>> model.system.addForce(nc_normalized)
-    6
+    >>> nc_normalized.addToSystem(model.system)
     >>> context.reinitialize(preserveState=True)
     >>> print(nc_normalized.getValue(context))
     0.99999... dimensionless