Barycenter Calculation in OpenMM

[svarner9]

I am not sure if this is a known issue, but I noticed that CVs which rely on the barycenter has a small issue in OpenMM.

To my knowledge, OpenMM wraps particle coordinates internally when using PBCs. As a result, when an atom group sits directly on the boundary, the barycenter is calculated to be in the center of the box, rather than on the edge. (See the attached image).

I think a way to remedy this would be to use the angular method mentioned in this Wikipedia article: https://en.wikipedia.org/wiki/Center_of_mass. I am going to implement this for myself, but would recommend looking into it for the package.

Thank you! :)

[svarner9]

I added this CV to my PySAGES coordinates.py file and it seems to work well as far as I can tell:

class ComponentPBC(AxisCV):
    def __init__(self, indices, axis, cell):
        super().__init__(indices, axis)
        self.cell = cell

    @property
    def function(self):
        return lambda rs: barycenter_pbc(rs, self.cell)[self.axis]

def barycenter_pbc(rs, cell):
    cell = np.diag(cell) # Only works for rectangular cells
    thetas = 2*np.pi*rs/cell
    xis = np.cos(thetas)
    zetas = np.sin(thetas)
    xibar = np.sum(xis, axis=0)/rs.shape[0]
    zetabar = np.sum(zetas, axis=0)/rs.shape[0]
    thetabar = np.arctan2(-zetabar, -xibar) + np.pi

    return cell*thetabar/(2*np.pi)

[p-zubieta]

Thanks for reporting this @svarner9! Just for completeness, can you share an example of the original CV you used that was problematic?

[svarner9]

Any CV that utilizes the barycenter will be incorrect I think.

I also think even for the CVs that I wrote, there are some edge cases where the center of mass will be incorrect, because OpenMM will not wrap the particle coordinates of any atom in a given molecule unless all of the atoms are outside the box. So you can have atoms that are beyond the boundary, but not yet wrapped. In this case, I think one needs to do the wrapping within the CV, before transforming to the unit circle (I haven't fully thought this out though).

The original CVs before changing them were:

class Component(AxisCV):
    """
    Use a specific Cartesian component of the center of mass of the group of atom selected
    via the indices.

    Parameters
    ----------
    indices: list[int], list[tuple(int)]
       Select atom groups via indices. From each group the barycenter is calculated.
    axis: int
       Cartesian coordinate axis component `0` (X), `1` (Y), `2` (Z) that is requested as CV.
    """

    @property
    def function(self):
        return lambda rs: barycenter(rs)[self.axis]

and

class Distance(TwoPointCV):
    """
    Use the distance of atom groups selected via the indices as collective variable.

    Parameters
    ----------
    indices: list[int], list[tuple(int)]
       Select atom groups via indices. (2 Groups required)
    """

    @property
    def function(self):
        if len(self.groups) == 0:
            return distance
        return lambda r1, r2: distance(barycenter(r1), barycenter(r2))


def distance(r1, r2):
    """
    Returns the distance between two points in space or
    between the barycenters of two groups of points in space.

    Parameters
    ----------
    r1: jax.Array
        Array containing the position in space of the first point or group of points.
    r2: jax.Array
        Array containing the position in space of the second point or group of points.

    Returns
    -------
    distance: float
        Distance between the two points.
    """

    return linalg.norm(r1 - r2)

I made new CVs for both of these that calculate the correct COM but also calculate the distance correctly with the minimum image convention. User beware, I am not an expert nor a PySAGES dev, please check these for yourself before you implement or use them:

class ComponentPBC(AxisCV):
    def __init__(self, indices, axis, cell):
        super().__init__(indices, axis)
        self.cell = cell

    @property
    def function(self):
        return lambda rs: barycenter_pbc(rs, self.cell)[self.axis]

def barycenter_pbc(rs, cell):
    cell = np.diag(cell) # Only works for rectangular cells
    thetas = 2*np.pi*rs/cell
    xis = np.cos(thetas)
    zetas = np.sin(thetas)
    xibar = np.sum(xis, axis=0)/rs.shape[0]
    zetabar = np.sum(zetas, axis=0)/rs.shape[0]
    thetabar = np.arctan2(-zetabar, -xibar) + np.pi

    return cell*thetabar/(2*np.pi)

and

class DistancePBC(TwoPointCV):
    """
    Use the distance of atom groups selected via the indices as collective variable, with PBCs.

    Parameters
    ----------
    indices: list[int], list[tuple(int)]
       Select atom groups via indices. (2 Groups required)
    cell: jax.Array
        The unit cell dimensions for handling PBCs.
    """

    def __init__(self, indices, cell):
        super().__init__(indices)
        self.cell = cell

    @property
    def function(self):
        if len(self.groups) == 0:
            return distance
        return lambda r1, r2: dist_mic(barycenter_pbc(r1,self.cell), barycenter_pbc(r2,self.cell), self.cell)

def barycenter_pbc(rs, cell):
    cell = np.diag(cell) # Only works for rectangular cells
    thetas = 2*np.pi*rs/cell
    xis = np.cos(thetas)
    zetas = np.sin(thetas)
    xibar = np.sum(xis, axis=0)/rs.shape[0]
    zetabar = np.sum(zetas, axis=0)/rs.shape[0]
    thetabar = np.arctan2(-zetabar, -xibar) + np.pi

    return cell*thetabar/(2*np.pi)

def dist_mic(r1, r2, cell):
    """Calculate minimum image convention (MIC) distance in non-cubic cell."""
    cell_inv = np.linalg.inv(cell)
    d = r1 - r2
    d_scaled = d @ cell_inv
    d_scaled -= np.round(d_scaled)
    d_mic = d_scaled @ cell
    return np.linalg.norm(d_mic)

Note that these CVs require you to pass the cell dimension up front, so these will not work for NPT simulations.