Skip to content

Commit

Permalink
Preprocessing units when registering CVs
Browse files Browse the repository at this point in the history
  • Loading branch information
craabreu committed Mar 22, 2024
1 parent 3c67a3a commit f7b83df
Show file tree
Hide file tree
Showing 4 changed files with 123 additions and 124 deletions.
101 changes: 51 additions & 50 deletions cvpack/cvpack.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
from openmm import unit as mmunit

from .serializer import Serializable
from .units import Quantity, value_in_md_units
from .units import Quantity, preprocess_units, value_in_md_units
from .utils import compute_effective_mass, get_single_force_state


Expand Down Expand Up @@ -103,6 +103,7 @@ def __copy__(self):
def __deepcopy__(self, memo):
return yaml.safe_load(yaml.safe_dump(self))

@preprocess_units
def _registerCV(
self,
unit: mmunit.Unit,
Expand All @@ -116,14 +117,14 @@ def _registerCV(
Parameters
----------
unit
The unit of measurement of this collective variable. It must be a unit
in the MD unit system (mass in Da, distance in nm, time in ps,
temperature in K, energy in kJ/mol, angle in rad).
args
The arguments needed to construct this collective variable
kwargs
The keyword arguments needed to construct this collective variable
unit
The unit of measurement of this collective variable. It must be a unit
in the MD unit system (mass in Da, distance in nm, time in ps,
temperature in K, energy in kJ/mol, angle in rad).
args
The arguments needed to construct this collective variable
kwargs
The keyword arguments needed to construct this collective variable
"""
cls = self.__class__
self.setName(cls.__name__)
Expand All @@ -142,8 +143,8 @@ def _registerPeriod(self, period: float) -> None:
Parameters
----------
period
The period of this collective variable
period
The period of this collective variable
"""
self._period = period

Expand All @@ -162,15 +163,15 @@ def _getArguments(cls) -> t.Tuple[collections.OrderedDict, collections.OrderedDi
Example
-------
>>> import cvpack
>>> args, defaults = cvpack.RadiusOfGyration._getArguments()
>>> for name, annotation in args.items():
... print(f"{name}: {annotation}")
group: typing.Iterable[int]
pbc: <class 'bool'>
weighByMass: <class 'bool'>
>>> print(*defaults.items())
('pbc', False) ('weighByMass', False)
>>> import cvpack
>>> args, defaults = cvpack.RadiusOfGyration._getArguments()
>>> for name, annotation in args.items():
... print(f"{name}: {annotation}")
group: typing.Iterable[int]
pbc: <class 'bool'>
weighByMass: <class 'bool'>
>>> print(*defaults.items())
('pbc', False) ('weighByMass', False)
"""
arguments = collections.OrderedDict()
defaults = collections.OrderedDict()
Expand All @@ -193,13 +194,13 @@ def _setUnusedForceGroup(self, system: openmm.System) -> None:
Parameters
----------
system
The system to search for unused force groups
system
The system to search for unused force groups
Raises
------
RuntimeError
If all force groups are already in use
RuntimeError
If all force groups are already in use
"""
used_groups = {force.getForceGroup() for force in system.getForces()}
new_group = next(filter(lambda i: i not in used_groups, range(32)), None)
Expand Down Expand Up @@ -330,31 +331,31 @@ def getEffectiveMass(self, context: openmm.Context) -> mmunit.Quantity:
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.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.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
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.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.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
"""
return mmunit.Quantity(compute_effective_mass(self, context), self._mass_unit)
39 changes: 19 additions & 20 deletions cvpack/radius_of_gyration.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,29 +46,28 @@ class RadiusOfGyration(BaseRadiusOfGyration):
Parameters
----------
group
The indices of the atoms in the group
pbc
Whether to use periodic boundary conditions
weighByMass
Whether to use the center of mass of the group instead of its geometric
center
group
The indices of the atoms in the group
pbc
Whether to use periodic boundary conditions
weighByMass
Whether to use the center of mass of the group instead of its geometric center
Example
-------
>>> import cvpack
>>> import openmm
>>> from openmmtools import testsystems
>>> model = testsystems.AlanineDipeptideVacuum()
>>> num_atoms = model.system.getNumParticles()
>>> radius_of_gyration = cvpack.RadiusOfGyration(list(range(num_atoms)))
>>> radius_of_gyration.addToSystem(model.system)
>>> platform = openmm.Platform.getPlatformByName('Reference')
>>> integrator = openmm.VerletIntegrator(0)
>>> context = openmm.Context(model.system, integrator, platform)
>>> context.setPositions(model.positions)
>>> print(radius_of_gyration.getValue(context))
0.2951... nm
>>> import cvpack
>>> import openmm
>>> from openmmtools import testsystems
>>> model = testsystems.AlanineDipeptideVacuum()
>>> num_atoms = model.system.getNumParticles()
>>> radius_of_gyration = cvpack.RadiusOfGyration(list(range(num_atoms)))
>>> radius_of_gyration.addToSystem(model.system)
>>> platform = openmm.Platform.getPlatformByName('Reference')
>>> integrator = openmm.VerletIntegrator(0)
>>> context = openmm.Context(model.system, integrator, platform)
>>> context.setPositions(model.positions)
>>> print(radius_of_gyration.getValue(context))
0.2951... nm
"""

Expand Down
39 changes: 19 additions & 20 deletions cvpack/radius_of_gyration_sq.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,29 +48,28 @@ class RadiusOfGyrationSq(BaseRadiusOfGyration):
Parameters
----------
group
The indices of the atoms in the group
pbc
Whether to use periodic boundary conditions
weighByMass
Whether to use the center of mass of the group instead of its geometric
center
group
The indices of the atoms in the group
pbc
Whether to use periodic boundary conditions
weighByMass
Whether to use the center of mass of the group instead of its geometric center
Example
-------
>>> import cvpack
>>> import openmm
>>> from openmmtools import testsystems
>>> model = testsystems.AlanineDipeptideVacuum()
>>> num_atoms = model.system.getNumParticles()
>>> rgsq = cvpack.RadiusOfGyrationSq(list(range(num_atoms)))
>>> rgsq.addToSystem(model.system)
>>> platform = openmm.Platform.getPlatformByName('Reference')
>>> integrator = openmm.VerletIntegrator(0)
>>> context = openmm.Context(model.system, integrator, platform)
>>> context.setPositions(model.positions)
>>> print(rgsq.getValue(context)) # doctest: +ELLIPSIS
0.0871... nm**2
>>> import cvpack
>>> import openmm
>>> from openmmtools import testsystems
>>> model = testsystems.AlanineDipeptideVacuum()
>>> num_atoms = model.system.getNumParticles()
>>> rgsq = cvpack.RadiusOfGyrationSq(list(range(num_atoms)))
>>> rgsq.addToSystem(model.system)
>>> platform = openmm.Platform.getPlatformByName('Reference')
>>> integrator = openmm.VerletIntegrator(0)
>>> context = openmm.Context(model.system, integrator, platform)
>>> context.setPositions(model.positions)
>>> print(rgsq.getValue(context)) # doctest: +ELLIPSIS
0.0871... nm**2
"""

Expand Down
68 changes: 34 additions & 34 deletions cvpack/torsion_similarity.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,43 +36,43 @@ class TorsionSimilarity(openmm.CustomCompoundBondForce, BaseCollectiveVariable):
Parameters
----------
firstList
A list of :math:`n` tuples of four atom indices defining the first torsion
angle in each pair.
secondList
A list of :math:`n` tuples of four atom indices defining the second torsion
angle in each pair.
pbc
Whether to use periodic boundary conditions
firstList
A list of :math:`n` tuples of four atom indices defining the first torsion
angle in each pair.
secondList
A list of :math:`n` tuples of four atom indices defining the second torsion
angle in each pair.
pbc
Whether to use periodic boundary conditions
Example
-------
>>> import cvpack
>>> import mdtraj
>>> import openmm
>>> import openmm.unit as mmunit
>>> from openmmtools import testsystems
>>> model = testsystems.LysozymeImplicit()
>>> traj = mdtraj.Trajectory(
... model.positions, mdtraj.Topology.from_openmm(model.topology)
... )
>>> phi_atoms, _ = mdtraj.compute_phi(traj)
>>> valid_atoms = traj.top.select("resid 59 to 79 and backbone")
>>> phi_atoms = [
... phi
... for phi in phi_atoms
... if all(atom in valid_atoms for atom in phi)
... ]
>>> torsion_similarity = cvpack.TorsionSimilarity(
... phi_atoms[1:], phi_atoms[:-1]
... )
>>> torsion_similarity.addToSystem(model.system)
>>> integrator = openmm.VerletIntegrator(0)
>>> platform = openmm.Platform.getPlatformByName("Reference")
>>> context = openmm.Context(model.system, integrator, platform)
>>> context.setPositions(model.positions)
>>> print(torsion_similarity.getValue(context))
18.659... dimensionless
>>> import cvpack
>>> import mdtraj
>>> import openmm
>>> import openmm.unit as mmunit
>>> from openmmtools import testsystems
>>> model = testsystems.LysozymeImplicit()
>>> traj = mdtraj.Trajectory(
... model.positions, mdtraj.Topology.from_openmm(model.topology)
... )
>>> phi_atoms, _ = mdtraj.compute_phi(traj)
>>> valid_atoms = traj.top.select("resid 59 to 79 and backbone")
>>> phi_atoms = [
... phi
... for phi in phi_atoms
... if all(atom in valid_atoms for atom in phi)
... ]
>>> torsion_similarity = cvpack.TorsionSimilarity(
... phi_atoms[1:], phi_atoms[:-1]
... )
>>> torsion_similarity.addToSystem(model.system)
>>> integrator = openmm.VerletIntegrator(0)
>>> platform = openmm.Platform.getPlatformByName("Reference")
>>> context = openmm.Context(model.system, integrator, platform)
>>> context.setPositions(model.positions)
>>> print(torsion_similarity.getValue(context))
18.659... dimensionless
"""

def __init__(
Expand Down

0 comments on commit f7b83df

Please sign in to comment.