Adds named parameter option to CentroidFunction (#36)

* Included unit compatibility check

* CentroidFunction accepts named parameters

* Fixed code formatting

cvpack/atomic_function.py

@@ -67,6 +67,9 @@ class AtomicFunction(openmm.CustomCompoundBondForce, AbstractCollectiveVariable)
             `unit.dimensionless`
         pbc
             Whether to use periodic boundary conditions
+
+    Keyword Args
+    ------------
         **parameters
             The named parameters of the function. Each parameter can be a scalar
             quantity or a 1D array-like object of length `m`, where `m` is the number of
@@ -98,9 +101,9 @@ class AtomicFunction(openmm.CustomCompoundBondForce, AbstractCollectiveVariable)
         ... )
         >>> [model.system.addForce(f) for f in [angle1, angle2, colvar]]
         [5, 6, 7]
-        >>> integrator =openmm.VerletIntegrator(0)
-        >>> platform =openmm.Platform.getPlatformByName('Reference')
-        >>> context =openmm.Context(model.system, integrator, platform)
+        >>> integrator = openmm.VerletIntegrator(0)
+        >>> platform = openmm.Platform.getPlatformByName('Reference')
+        >>> context = openmm.Context(model.system, integrator, platform)
         >>> context.setPositions(model.positions)
         >>> theta1 = angle1.getValue(context).value_in_unit(openmm.unit.radian)
         >>> theta2 = angle2.getValue(context).value_in_unit(openmm.unit.radian)
@@ -127,15 +130,19 @@ def __init__(  # pylint: disable=too-many-arguments
             if isinstance(data, mmunit.Quantity):
                 data = data.value_in_unit_system(mmunit.md_unit_system)
             if isinstance(data, Sequence):
+                if len(data) != groups.shape[0]:
+                    raise ValueError(
+                        f"The length of the parameter {name} "
+                        f"must be equal to {groups.shape[0]}."
+                    )
                 self.addPerBondParameter(name)
                 perbond_parameters.append(data)
             else:
                 self.addGlobalParameter(name, data)
         for group, *values in zip(groups, *perbond_parameters):
             self.addBond(group, values)
         self.setUsesPeriodicBoundaryConditions(pbc)
-        if mmunit.Quantity(1, unit).value_in_unit_system(mmunit.md_unit_system) != 1:
-            raise ValueError(f"Unit {unit} is not compatible with the MD unit system.")
+        self._checkUnitCompatibility(unit)
         self._registerCV(
             unit,
             atomsPerGroup,

cvpack/centroid_function.py

@@ -63,6 +63,15 @@ class CentroidFunction(openmm.CustomCentroidBondForce, AbstractCollectiveVariabl
             `dimensionless`
         pbc
             Whether to use periodic boundary conditions
+        weighByMass
+            Whether to define the centroid as the center of mass of the group instead of
+            the geometric center
+
+    Keyword Args
+    ------------
+        **parameters
+            The named parameters of the function. If the specified value has units, it
+            will be converted to the MD unit system.
 
     Raises
     ------
@@ -78,20 +87,21 @@ class CentroidFunction(openmm.CustomCentroidBondForce, AbstractCollectiveVariabl
         >>> model = testsystems.AlanineDipeptideVacuum()
         >>> num_atoms = model.system.getNumParticles()
         >>> atoms = list(range(num_atoms))
-        >>> rg = cvpack.RadiusOfGyration(atoms)
-        >>> definitions = [f'd{i+1} = distance(g{i+1}, g{num_atoms+1})' for i in atoms]
-        >>> sum_dist_sq = "+".join(f'd{i+1}^2' for i in atoms)
-        >>> function = ";".join([f"sqrt(({sum_dist_sq})/{num_atoms})"] + definitions)
-        >>> groups = [[i] for i in atoms] + [atoms]
-        >>> colvar = cvpack.CentroidFunction(function, groups, unit.nanometers)
-        >>> [model.system.addForce(f) for f in [rg, colvar]]
-        [5, 6]
-        >>> integrator =openmm.VerletIntegrator(0)
-        >>> platform =openmm.Platform.getPlatformByName('Reference')
-        >>> context =openmm.Context(model.system, integrator, platform)
+        >>> groups = [[i] for i in atoms]  # Each atom is a group
+        >>> groups.append(atoms)  # The whole molecule is a group
+        >>> sum_dist_sq = "+".join(
+        ...     f'distance(g{i+1}, g{num_atoms+1})^2' for i in atoms
+        ... )
+        >>> function = f"sqrt(({sum_dist_sq})/n)"  # The radius of gyration
+        >>> colvar = cvpack.CentroidFunction(
+        ...     function, groups, unit.nanometers, n=num_atoms,
+        ... )
+        >>> model.system.addForce(colvar)
+        5
+        >>> integrator = openmm.VerletIntegrator(0)
+        >>> platform = openmm.Platform.getPlatformByName('Reference')
+        >>> context = openmm.Context(model.system, integrator, platform)
         >>> context.setPositions(model.positions)
-        >>> print(rg.getValue(context, digits=6))
-        0.2951431 nm
         >>> print(colvar.getValue(context, digits=6))
         0.2951431 nm
     """
@@ -103,13 +113,15 @@ def __init__(  # pylint: disable=too-many-arguments
         unit: mmunit.Unit,
         pbc: bool = False,
         weighByMass: bool = False,
+        **parameters: mmunit.ScalarQuantity,
     ) -> None:
         num_groups = len(groups)
         super().__init__(num_groups, function)
         for group in groups:
             self.addGroup(group, None if weighByMass else [1] * len(group))
+        for name, value in parameters.items():
+            self.addGlobalParameter(name, value)
         self.addBond(list(range(num_groups)), [])
         self.setUsesPeriodicBoundaryConditions(pbc)
-        if mmunit.Quantity(1, unit).value_in_unit_system(mmunit.md_unit_system) != 1:
-            raise ValueError(f"Unit {unit} is not compatible with the MD unit system.")
+        self._checkUnitCompatibility(unit)
         self._registerCV(unit, function, groups, SerializableUnit(unit), pbc)

cvpack/cvpack.py

@@ -131,6 +131,22 @@ def _precisionRound(self, number: float, digits: Optional[int] = None) -> float:
         power = f"{number:e}".split("e")[1]
         return round(number, -(int(power) - digits))
 
+    @staticmethod
+    def _checkUnitCompatibility(unit: mmunit.Unit) -> None:
+        """
+        Check if the given unit is compatible with the MD unit system.
+
+        Parameters
+        ----------
+            unit
+                The unit to check
+        """
+        if not np.isclose(
+            mmunit.Quantity(1.0, unit).value_in_unit_system(mmunit.md_unit_system),
+            1.0,
+        ):
+            raise ValueError(f"Unit {unit} is not compatible with the MD unit system.")
+
     @classmethod
     def getArguments(cls) -> Tuple[OrderedDict, OrderedDict]:
         """