Skip to content

Commit

Permalink
Merge pull request #31 from RedesignScience/fix_rg_bug
Browse files Browse the repository at this point in the history
Fixes bug in center-of-mass radii of gyration
  • Loading branch information
craabreu authored Oct 4, 2023
2 parents b97065e + 72c1c18 commit e49341a
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 34 deletions.
40 changes: 28 additions & 12 deletions cvpack/radius_of_gyration.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,26 @@
from .cvpack import AbstractCollectiveVariable


class RadiusOfGyration(openmm.CustomCentroidBondForce, AbstractCollectiveVariable):
class _RadiusOfGyrationBase(openmm.CustomCentroidBondForce, AbstractCollectiveVariable):
def __init__( # pylint: disable=too-many-arguments
self,
num_groups: int,
expression: str,
group: Sequence[int],
pbc: bool = False,
weighByMass: bool = False,
) -> None:
super().__init__(num_groups, expression)
for atom in group:
self.addGroup([atom])
if weighByMass:
self.addGroup(group)
else:
self.addGroup(group, [1] * len(group))
self.setUsesPeriodicBoundaryConditions(pbc)


class RadiusOfGyration(_RadiusOfGyrationBase):
"""
The radius of gyration of a group of :math:`n` atoms:
Expand All @@ -26,18 +45,18 @@ class RadiusOfGyration(openmm.CustomCentroidBondForce, AbstractCollectiveVariabl
{\\bf r}_i - {\\bf r}_c({\\bf r})
\\right\\|^2 }.
where :math:`{\\bf r}_c({\\bf r})` is the centroid of the group:
where :math:`{\\bf r}_c({\\bf r})` is the geometric center of the group:
.. math::
{\\bf r}_c({\\bf r}) = \\frac{1}{n} \\sum_{i=j}^n {\\bf r}_j
Optionally, the atoms can be weighted by their masses. In this case, the centroid is computed
as:
Optionally, the radius of gyration can be computed with respect to the center of
mass of the group. In this case, the geometric center is replaced by:
.. math::
{\\bf r}_c({\\bf r}) = \\frac{1}{M} \\sum_{i=1}^n m_i {\\bf r}_i
{\\bf r}_m({\\bf r}) = \\frac{1}{M} \\sum_{i=1}^n m_i {\\bf r}_i
where :math:`M = \\sum_{i=1}^n m_i` is the total mass of the group.
Expand All @@ -53,7 +72,8 @@ class RadiusOfGyration(openmm.CustomCentroidBondForce, AbstractCollectiveVariabl
pbc
Whether to use periodic boundary conditions
weighByMass
Whether to weigh the atoms by their masses
Whether to use the center of mass of the group instead of its geometric
center
Example
-------
Expand All @@ -78,10 +98,6 @@ def __init__(self, group: Sequence[int], pbc: bool = False, weighByMass: bool =
num_atoms = len(group)
num_groups = num_atoms + 1
sum_dist_sq = "+".join([f"distance(g{i+1}, g{num_atoms + 1})^2" for i in range(num_atoms)])
super().__init__(num_groups, f"sqrt(({sum_dist_sq})/{num_atoms})")
for atom in group:
self.addGroup([atom], [1])
self.addGroup(group, None if weighByMass else [1] * num_atoms)
self.addBond(list(range(num_groups)), [])
self.setUsesPeriodicBoundaryConditions(pbc)
super().__init__(num_groups, f"sqrt(({sum_dist_sq})/{num_atoms})", group, pbc, weighByMass)
self.addBond(list(range(num_groups)))
self._registerCV(mmunit.nanometers, group, pbc, weighByMass)
23 changes: 9 additions & 14 deletions cvpack/radius_of_gyration_sq.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,14 +10,12 @@

from typing import Sequence

import openmm

from cvpack import unit as mmunit

from .cvpack import AbstractCollectiveVariable
from .radius_of_gyration import _RadiusOfGyrationBase


class RadiusOfGyrationSq(openmm.CustomCentroidBondForce, AbstractCollectiveVariable):
class RadiusOfGyrationSq(_RadiusOfGyrationBase):
"""
The square of the radius of gyration of a group of :math:`n` atoms:
Expand All @@ -27,18 +25,18 @@ class RadiusOfGyrationSq(openmm.CustomCentroidBondForce, AbstractCollectiveVaria
{\\bf r}_i - {\\bf r}_c({\\bf r})
\\right\\|^2.
where :math:`{\\bf r}_c({\\bf r})` is the centroid of the group:
where :math:`{\\bf r}_c({\\bf r})` is the geometric center of the group:
.. math::
{\\bf r}_c({\\bf r}) = \\frac{1}{n} \\sum_{i=j}^n {\\bf r}_j
Optionally, the atoms can be weighted by their masses. In this case, the centroid is computed
as:
Optionally, the radius of gyration can be computed with respect to the center of
mass of the group. In this case, the geometric center is replaced by:
.. math::
{\\bf r}_c({\\bf r}) = \\frac{1}{M} \\sum_{i=1}^n m_i {\\bf r}_i
{\\bf r}_m({\\bf r}) = \\frac{1}{M} \\sum_{i=1}^n m_i {\\bf r}_i
where :math:`M = \\sum_{i=1}^n m_i` is the total mass of the group.
Expand All @@ -54,7 +52,8 @@ class RadiusOfGyrationSq(openmm.CustomCentroidBondForce, AbstractCollectiveVaria
pbc
Whether to use periodic boundary conditions
weighByMass
Whether to weigh the atoms by their masses
Whether to use the center of mass of the group instead of its geometric
center
Example
-------
Expand All @@ -77,11 +76,7 @@ class RadiusOfGyrationSq(openmm.CustomCentroidBondForce, AbstractCollectiveVaria

def __init__(self, group: Sequence[int], pbc: bool = False, weighByMass: bool = False) -> None:
num_atoms = len(group)
super().__init__(2, f"distance(g1, g2)^2/{num_atoms}")
for atom in group:
self.addGroup([atom], [1])
self.addGroup(group, None if weighByMass else [1] * num_atoms)
super().__init__(2, f"distance(g1, g2)^2/{num_atoms}", group, pbc, weighByMass)
for atom in group:
self.addBond([atom, num_atoms])
self.setUsesPeriodicBoundaryConditions(pbc)
self._registerCV(mmunit.nanometers**2, group, pbc, weighByMass)
36 changes: 28 additions & 8 deletions cvpack/tests/test_cvpack.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
from scipy.spatial.transform import Rotation

import cvpack
from cvpack import unit
from cvpack import serializer, unit


def test_cvpack_imported():
Expand Down Expand Up @@ -87,9 +87,9 @@ def perform_common_tests(

# Test serialization/deserialization
pipe = io.StringIO()
cvpack.serializer.serialize(collectiveVariable, pipe)
serializer.serialize(collectiveVariable, pipe)
pipe.seek(0)
new_cv = cvpack.serializer.deserialize(pipe)
new_cv = serializer.deserialize(pipe)
context.getSystem().addForce(new_cv)
context.reinitialize(preserveState=True)
value1 = collectiveVariable.getValue(context)
Expand Down Expand Up @@ -185,17 +185,27 @@ def test_radius_of_gyration():
"""
model = testsystems.AlanineDipeptideVacuum()
masses = np.array(
[unit.value_in_md_units(atom.element.mass) for atom in model.topology.atoms()]
)
positions = model.positions.value_in_unit(unit.nanometers)
centroid = positions.mean(axis=0)
rgsq = np.sum((positions - centroid) ** 2) / model.system.getNumParticles()
rg_cv = cvpack.RadiusOfGyration(range(model.system.getNumParticles()))
center_of_mass = np.einsum("i,ij->j", masses, positions) / np.sum(masses)
num_atoms = model.system.getNumParticles()
rgsq = np.sum((positions - centroid) ** 2) / num_atoms
rg_cv = cvpack.RadiusOfGyration(range(num_atoms))
weighted_rgsq = np.sum((positions - center_of_mass) ** 2) / num_atoms
weighted_rg_cv = cvpack.RadiusOfGyration(range(num_atoms), weighByMass=True)
model.system.addForce(rg_cv)
model.system.addForce(weighted_rg_cv)
integrator = openmm.CustomIntegrator(0)
platform = openmm.Platform.getPlatformByName("Reference")
context = openmm.Context(model.system, integrator, platform)
context.setPositions(model.positions)
rgval = rg_cv.getValue(context).value_in_unit(unit.nanometers)
assert rgval**2 == pytest.approx(rgsq)
weighted_rgval = weighted_rg_cv.getValue(context).value_in_unit(unit.nanometers)
assert weighted_rgval**2 == pytest.approx(weighted_rgsq)
perform_common_tests(rg_cv, context)


Expand All @@ -205,17 +215,27 @@ def test_radius_of_gyration_squared():
"""
model = testsystems.AlanineDipeptideVacuum()
masses = np.array(
[unit.value_in_md_units(atom.element.mass) for atom in model.topology.atoms()]
)
positions = model.positions.value_in_unit(unit.nanometers)
centroid = positions.mean(axis=0)
rgsq = np.sum((positions - centroid) ** 2) / model.system.getNumParticles()
rg_sq = cvpack.RadiusOfGyrationSq(range(model.system.getNumParticles()))
center_of_mass = np.einsum("i,ij->j", masses, positions) / np.sum(masses)
num_atoms = model.system.getNumParticles()
rgsq = np.sum((positions - centroid) ** 2) / num_atoms
rg_sq = cvpack.RadiusOfGyrationSq(range(num_atoms))
weighted_rgsq = np.sum((positions - center_of_mass) ** 2) / num_atoms
weighted_rg_sq = cvpack.RadiusOfGyrationSq(range(num_atoms), weighByMass=True)
model.system.addForce(rg_sq)
model.system.addForce(weighted_rg_sq)
integrator = openmm.CustomIntegrator(0)
platform = openmm.Platform.getPlatformByName("Reference")
context = openmm.Context(model.system, integrator, platform)
context.setPositions(model.positions)
rg_sq_value = rg_sq.getValue(context).value_in_unit(unit.nanometers**2)
rg_sq_value = unit.value_in_md_units(rg_sq.getValue(context))
assert rg_sq_value == pytest.approx(rgsq)
weighted_rg_sq_value = unit.value_in_md_units(weighted_rg_sq.getValue(context))
assert weighted_rg_sq_value == pytest.approx(weighted_rgsq)
perform_common_tests(rg_sq, context)


Expand Down

0 comments on commit e49341a

Please sign in to comment.