diff --git a/cvpack/atomic_function.py b/cvpack/atomic_function.py index 0f117985..edc4a538 100644 --- a/cvpack/atomic_function.py +++ b/cvpack/atomic_function.py @@ -9,7 +9,7 @@ from __future__ import annotations -from typing import Sequence, Union +import typing as t import numpy as np import openmm @@ -27,9 +27,31 @@ from .cvpack import AbstractCollectiveVariable +def _add_parameters( + force: openmm.Force, + size: int, + **parameters: t.Union[mmunit.ScalarQuantity, t.Sequence[mmunit.ScalarQuantity]], +) -> t.List[t.Union[float, t.Sequence[mmunit.ScalarQuantity]]]: + perbond_parameters = [] + for name, data in parameters.items(): + if isinstance(data, mmunit.Quantity): + data = data.value_in_unit_system(mmunit.md_unit_system) + if isinstance(data, t.Sequence): + if len(data) != size: + raise ValueError( + f"The length of parameter {name} is {len(data)}. Should be {size}." + ) + force.addPerBondParameter(name) + perbond_parameters.append(data) + else: + force.addGlobalParameter(name, data) + return perbond_parameters + + class AtomicFunction(openmm.CustomCompoundBondForce, AbstractCollectiveVariable): """ - A generic function of the coordinates of `m` groups of `n` atoms: + A generic function of the coordinates of atoms split into `m` groups of `n` + atoms each: .. math:: @@ -38,7 +60,7 @@ class AtomicFunction(openmm.CustomCompoundBondForce, AbstractCollectiveVariable) \\right) where :math:`F` is a user-defined function and :math:`{\\bf r}_{i,j}` is the - position of the :math:`j`-th atom of the :math:`i`-th group. + coordinate of the :math:`j`-th atom of the :math:`i`-th group. The function :math:`F` is defined as a string and can be any expression supported by :OpenMM:`CustomCompoundBondForce`. If it contains named parameters, they must be @@ -48,38 +70,37 @@ class AtomicFunction(openmm.CustomCompoundBondForce, AbstractCollectiveVariable) Parameters ---------- - atomsPerGroup - The number of atoms in each group - function - The function to be evaluated. It must be a valid - :OpenMM:`CustomCompoundBondForce` expression - atoms - The indices of the atoms in each group, passed as a 2D array-like object of - shape `(m, n)`, where `m` is the number of groups and `n` is the number of - atoms per group, or as a 1D array-like object of length `m*n`, where the - first `n` elements are the indices of the atoms in the first group, the next - `n` elements are the indices of the atoms in the second group, and so on. - unit - The unit of measurement of the collective variable. It must be compatible - with the MD unit system (mass in `daltons`, distance in `nanometers`, time - in `picoseconds`, temperature in `kelvin`, energy in `kilojoules_per_mol`, - angle in `radians`). If the collective variables does not have a unit, use - `unit.dimensionless` - pbc - Whether to use periodic boundary conditions + function + The function to be evaluated. It must be a valid + :OpenMM:`CustomCompoundBondForce` expression + groups + The indices of the atoms in each group, passed as a 2D array-like object of + shape `(m, n)`, where `m` is the number of groups and `n` is the number of + atoms per group. If a 1D object is passed, it is assumed that `m` is 1 and + `n` is the length of the object. + unit + The unit of measurement of the collective variable. It must be compatible + with the MD unit system (mass in `daltons`, distance in `nanometers`, time + in `picoseconds`, temperature in `kelvin`, energy in `kilojoules_per_mol`, + angle in `radians`). If the collective variables does not have a unit, use + `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 - groups. In the latter case, each entry of the array is used for the - corresponding group of atoms. + **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 + groups. In the latter case, each entry of the array is used for the + corresponding group of atoms. Raises ------ - ValueError - If the collective variable is not compatible with the MD unit system + ValueError + If the groups are not specified as a 1D or 2D array-like object + ValueError + If the unit of the collective variable is not compatible with the MD unit system Example ------- @@ -92,10 +113,9 @@ class AtomicFunction(openmm.CustomCompoundBondForce, AbstractCollectiveVariable) >>> angle1 = cvpack.Angle(0, 5, 10) >>> angle2 = cvpack.Angle(10, 15, 20) >>> colvar = cvpack.AtomicFunction( - ... 3, ... "(k/2)*(angle(p1, p2, p3) - theta0)^2", - ... [[0, 5, 10], [10, 15, 20]], ... unit.kilojoules_per_mole, + ... [[0, 5, 10], [10, 15, 20]], ... k = 1000 * unit.kilojoules_per_mole/unit.radian**2, ... theta0 = [np.pi/2, np.pi/3] * unit.radian, ... ) @@ -113,50 +133,31 @@ class AtomicFunction(openmm.CustomCompoundBondForce, AbstractCollectiveVariable) 429.479 kJ/mol """ - @mmunit.convert_quantities def __init__( # pylint: disable=too-many-arguments self, - atomsPerGroup: int, function: str, - indices: ArrayLike, unit: mmunit.Unit, - pbc: bool = False, - **parameters: Union[mmunit.ScalarQuantity, Sequence[mmunit.ScalarQuantity]], + groups: ArrayLike, + pbc: bool = True, + **parameters: t.Union[mmunit.ScalarQuantity, t.Sequence[mmunit.ScalarQuantity]], ) -> None: - groups = np.asarray(indices, dtype=np.int32).reshape(-1, atomsPerGroup) - super().__init__(atomsPerGroup, function) - perbond_parameters = [] - for name, data in parameters.items(): - 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) + groups = np.atleast_2d(groups) + num_groups, atoms_per_group, *others = groups.shape + if others: + raise ValueError("Array `groups` cannot have more than 2 dimensions") + super().__init__(atoms_per_group, function) + perbond_parameters = _add_parameters(self, num_groups, **parameters) for group, *values in zip(groups, *perbond_parameters): self.addBond(group, values) self.setUsesPeriodicBoundaryConditions(pbc) self._checkUnitCompatibility(unit) - self._registerCV( - unit, - atomsPerGroup, - function, - groups, - mmunit.SerializableUnit(unit), - pbc, - **parameters, - ) + unit = mmunit.SerializableUnit(unit) + self._registerCV(unit, function, unit, groups, pbc, **parameters) @classmethod def _fromCustomForce( cls, - force: Union[ + force: t.Union[ CustomAngleForce, CustomBondForce, CustomCompoundBondForce, @@ -206,7 +207,8 @@ def _fromCustomForce( atoms.append(indices) for name, value in zip(per_item_parameter_names, per_item_parameters): parameters[name].append(value) - return cls(number, function, atoms, unit, pbc, **parameters) + atoms = np.asarray(atoms).reshape(-1, number) + return cls(function, unit, atoms, pbc, **parameters) @classmethod def _fromHarmonicBondForce( @@ -225,7 +227,7 @@ def _fromHarmonicBondForce( atoms.append(indices) parameters["r0"].append(length) parameters["k"].append(k) - return cls(2, "(k/2)*(distance(p1, p2)-r0)^2", atoms, unit, pbc, **parameters) + return cls("(k/2)*(distance(p1, p2)-r0)^2", unit, atoms, pbc, **parameters) @classmethod def _fromHarmonicAngleForce( @@ -244,9 +246,7 @@ def _fromHarmonicAngleForce( atoms.append(indices) parameters["theta0"].append(angle) parameters["k"].append(k) - return cls( - 3, "(k/2)*(angle(p1, p2, p3)-theta0)^2", atoms, unit, pbc, **parameters - ) + return cls("(k/2)*(angle(p1, p2, p3)-theta0)^2", unit, atoms, pbc, **parameters) @classmethod def _fromPeriodicTorsionForce( @@ -267,10 +267,9 @@ def _fromPeriodicTorsionForce( parameters["phase"].append(phase) parameters["k"].append(k) return cls( - 4, "k*(1 + cos(periodicity*dihedral(p1, p2, p3, p4) - phase))", - atoms, unit, + atoms, pbc, **parameters, ) diff --git a/cvpack/centroid_function.py b/cvpack/centroid_function.py index 2e320456..143ab646 100644 --- a/cvpack/centroid_function.py +++ b/cvpack/centroid_function.py @@ -7,123 +7,177 @@ """ -from typing import Sequence +import typing as t +import numpy as np import openmm +from numpy.typing import ArrayLike from cvpack import unit as mmunit +from .atomic_function import _add_parameters from .cvpack import AbstractCollectiveVariable -from .unit import SerializableUnit class CentroidFunction(openmm.CustomCentroidBondForce, AbstractCollectiveVariable): """ - A generic function of the centroids of `n` groups of atoms: + A generic function of the centroids of :math:`m \\times n` atoms groups split + into `m` collections of `n` groups each: .. math:: - f({\\bf r}) = F({\\bf R}_1, {\\bf R}_2, \\dots, {\\bf R}_n) + f({\\bf r}) = \\sum_{i=1}^m F\\Big( + {\\bf g}^i_1({\\bf r}), + {\\bf g}^i_2({\\bf r}), + \\dots, + {\\bf g}^i_n({\\bf r}) + \\Big) - where :math:`F` is a user-defined function and :math:`{\\bf R}_i` is the centroid of - the :math:`i`-th group of atoms. The function :math:`F` is defined as a string and - can be any valid :OpenMM:`CustomCentroidBondForce` expression. + where :math:`F` is a user-defined function and :math:`{\\bf g}^i_1({\\bf r})` is the + centroid of the :math:`j`-th group of atoms of the :math:`i`-th collection of + groups. + + The function :math:`F` is defined as a string and can be any expression supported + by :OpenMM:`CustomCentroidBondForce`. If it contains named parameters, they must + be passed as keyword arguments to the :class:`CentroidFunction` constructor. The + parameters can be scalars or arrays of length :math:`m`. In the latter case, each + value will be assigned to the corresponding collection of atom groups. The centroid of a group of atoms is defined as: .. math:: - {\\bf R}_i({\\bf r}) = \\frac{1}{n_i} \\sum_{j=1}^{n_i} {\\bf r}_{j, i} + {\\bf g}_j({\\bf r}) = \\frac{1}{N_j} \\sum_{k=1}^{N_j} {\\bf r}_{k,j} - where :math:`n_i` is the number of atoms in group :math:`i` and - :math:`{\\bf r}_{j, i}` is the position of the :math:`j`-th atom in group :math:`i`. - Optionally, the centroid can be weighted by the mass of each atom in the group. In - this case, it is defined as: + where :math:`N_j` is the number of atoms in group :math:`j` and + :math:`{\\bf r}_{k,j}` is the coordinate of the :math:`k`-th atom of group + :math:`j`. Optionally, the centroid can be weighted by the mass of each atom + in the group. In this case, it is redefined as: .. math:: - {\\bf R}_i({\\bf r}) = \\frac{1}{M_i} \\sum_{j=1}^{n_i} m_{j, i} {\\bf r}_{j, i} + {\\bf g}_j({\\bf r}) = \\frac{1}{M_j} \\sum_{k=1}^{N_j} m_{k,j} {\\bf r}_{k,j} - where :math:`M_i` is the total mass of group :math:`i` and :math:`m_{j, i}` is the - mass of the :math:`j`-th atom in group :math:`i`. + where :math:`M_j` is the total mass of atom group :math:`j` and :math:`m_{k,j}` is + the mass of the :math:`k`-th atom in group :math:`j`. Parameters ---------- - function - The function to be evaluated. It must be a valid - :OpenMM:`CustomCentroidBondForce` expression - groups - The groups of atoms to be used in the function. Each group must be a list of - atom indices - unit - The unit of measurement of the collective variable. It must be compatible - with the MD unit system (mass in `daltons`, distance in `nanometers`, time - in `picoseconds`, temperature in `kelvin`, energy in `kilojoules_per_mol`, - angle in `radians`). If the collective variables does not have a unit, use - `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 + function + The function to be evaluated. It must be a valid + :OpenMM:`CustomCentroidBondForce` expression + groups + The groups of atoms to be used in the function. Each group must be specified + as a list of atom indices with arbitrary length + collections + The indices of the groups in each collection, passed as a 2D array-like object + of shape `(m, n)`, where `m` is the number of collections and `n` is the number + groups per collection. If a 1D object is passed, it is assumed that `m` is 1 and + `n` is the length of the object. + unit + The unit of measurement of the collective variable. It must be compatible + with the MD unit system (mass in `daltons`, distance in `nanometers`, time + in `picoseconds`, temperature in `kelvin`, energy in `kilojoules_per_mol`, + angle in `radians`). If the collective variables does not have a unit, use + `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. + **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 + group collections. In the latter case, each entry of the array is used for + the corresponding collection of groups. Raises ------ - ValueError - If the collective variable is not compatible with the MD unit system + ValueError + If the collections are not specified as a 1D or 2D array-like object + ValueError + If group indices are out of bounds + ValueError + If the unit of the collective variable is not compatible with the MD unit system Example ------- >>> import cvpack + >>> import itertools as it + >>> import numpy as np >>> import openmm >>> from openmm import unit >>> from openmmtools import testsystems - >>> model = testsystems.AlanineDipeptideVacuum() - >>> num_atoms = model.system.getNumParticles() - >>> atoms = list(range(num_atoms)) - >>> 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 = testsystems.LysozymeImplicit() + >>> residues = list(model.topology.residues()) + >>> atoms = [[a.index for a in r.atoms()] for r in residues] + + Compute the residue coordination between two helices: + + >>> res_coord = cvpack.ResidueCoordination( + ... residues[115:124], residues[126:135], stepFunction="step(1-x)" ... ) - >>> colvar.setUnusedForceGroup(0, model.system) + >>> res_coord.setUnusedForceGroup(0, model.system) 1 - >>> model.system.addForce(colvar) - 5 + >>> model.system.addForce(res_coord) + 6 >>> integrator = openmm.VerletIntegrator(0) >>> platform = openmm.Platform.getPlatformByName('Reference') >>> context = openmm.Context(model.system, integrator, platform) >>> context.setPositions(model.positions) - >>> print(colvar.getValue(context, digits=6)) - 0.2951431 nm + >>> print(res_coord.getValue(context)) + 33.0 dimensionless + + Recompute the residue coordination using the centroid function: + + >>> groups = [atoms[115:124], atoms[126:135]] + >>> collections = list(it.product(range(9), range(9, 18))) + >>> colvar = cvpack.CentroidFunction( + ... "step(1 - distance(g1, g2))", + ... unit.dimensionless, + ... 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 + >>> context.reinitialize(preserveState=True) + >>> print(colvar.getValue(context)) + 33.0 dimensionless """ def __init__( # pylint: disable=too-many-arguments self, function: str, - groups: Sequence[Sequence[int]], unit: mmunit.Unit, - pbc: bool = False, - weighByMass: bool = False, + groups: t.Sequence[t.Sequence[int]], + collections: t.Optional[ArrayLike] = None, + pbc: bool = True, + weighByMass: bool = True, **parameters: mmunit.ScalarQuantity, ) -> None: + collections = np.atleast_2d( + np.arange(len(groups)) if collections is None else collections + ) + num_collections, groups_per_collection, *others = collections.shape + if others: + raise ValueError("Array `collections` cannot have more than 2 dimensions") num_groups = len(groups) - super().__init__(num_groups, function) + if np.any(collections < 0) or np.any(collections >= num_groups): + raise ValueError("Group index out of bounds") + super().__init__(groups_per_collection, 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.addGroup(group, *([] if weighByMass else [[1] * len(group)])) + perbond_parameters = _add_parameters(self, num_collections, **parameters) + for collection, *values in zip(collections, *perbond_parameters): + self.addBond(collection, values) self.setUsesPeriodicBoundaryConditions(pbc) self._checkUnitCompatibility(unit) - self._registerCV(unit, function, groups, SerializableUnit(unit), pbc) + unit = mmunit.SerializableUnit(unit) + self._registerCV( + unit, function, unit, groups, collections, pbc, weighByMass, **parameters + ) diff --git a/cvpack/tests/test_cvpack.py b/cvpack/tests/test_cvpack.py index 8d1927b1..186f8a03 100644 --- a/cvpack/tests/test_cvpack.py +++ b/cvpack/tests/test_cvpack.py @@ -613,11 +613,11 @@ def test_atomic_function(): np.random.shuffle(atoms) function = "+".join(f"distance(p{i+1}, p{i+2})" for i in range(num_atoms - 1)) with pytest.raises(ValueError) as excinfo: - colvar = cvpack.AtomicFunction(num_atoms, function, atoms, unit.angstrom) + colvar = cvpack.AtomicFunction(function, unit.angstrom, atoms) assert ( str(excinfo.value) == "Unit angstrom is not compatible with the MD unit system." ) - colvar = cvpack.AtomicFunction(num_atoms, function, atoms, unit.nanometers) + colvar = cvpack.AtomicFunction(function, unit.nanometers, atoms) colvar.setUnusedForceGroup(0, model.system) model.system.addForce(colvar) context = openmm.Context( @@ -651,15 +651,14 @@ def test_centroid_function(): groups = np.reshape(atoms[: 3 * num_groups], (num_groups, 3)) function = "+".join(f"distance(g{i+1}, g{i+2})" for i in range(num_groups - 1)) with pytest.raises(ValueError) as excinfo: - colvar = cvpack.CentroidFunction(function, groups, unit.angstrom) + colvar = cvpack.CentroidFunction(function, unit.angstrom, groups) assert ( str(excinfo.value) == "Unit angstrom is not compatible with the MD unit system." ) colvar = cvpack.CentroidFunction( - function, groups, unit.nanometers, weighByMass=False + function, unit.nanometers, groups, weighByMass=False ) colvar.setUnusedForceGroup(0, model.system) - colvar.setUnusedForceGroup(0, model.system) model.system.addForce(colvar) context = openmm.Context( model.system,