Add utilities to calculate Order Parameter

cmeutils/geometry.py

@@ -0,0 +1,87 @@
+import numpy as np
+from numpy.linalg import svd
+
+
+def get_plane_normal(points):
+    """Calculate the plane which best fits a cloud of points.
+
+    Best fit is calculated using numpy's Singular Value Decomposition.
+
+    Example
+    -------
+    To visualize the plane fit in a Jupyter notebook::
+
+        %matplotlib notebook
+        import matplotlib.pyplot as plt
+        import numpy as np
+
+        normal, ctr = get_plane_normal(points)
+        plt.figure()
+        ax = plt.subplot(111, projection='3d')
+        ax.scatter(points[:,0], points[:,1], points[:,2], color='b')
+
+        # plot plane
+        xlim = ax.get_xlim()
+        ylim = ax.get_ylim()
+        xx, yy = np.meshgrid(
+            np.linspace(xlim[0], xlim[1], 3), np.linspace(ylim[0], ylim[1], 3)
+        )
+
+        d = -ctr.dot(normal)
+        ax.scatter(ctr[0], ctr[1], ctr[2], color='r')
+        z = (-normal[0] * xx - normal[1] * yy - d) * 1 / normal[2]
+        ax.plot_surface(xx, yy, z, alpha=0.5)
+
+        ax.set_box_aspect(aspect = (1,1,1))
+        ax.set_xlabel('x')
+        ax.set_ylabel('y')
+        ax.set_zlabel('z')
+        plt.show()
+
+
+    Parameters
+    ----------
+    points : numpy.ndarray, shape (N,3)
+        Coordinates (x,y,z) through which to fit a plane. Must be at least 3
+        points.
+
+    Returns
+    -------
+    ctr : numpy.ndarray, shape (3,)
+        The center of the point cloud.
+    normal : numpy.ndarray, shape (3,)
+        The plane normal vector. A unit vector from the origin.
+    """
+    assert points.shape[0] >= 3, "Need at least 3 points to calculate a plane."
+    ctr = points.mean(axis=0)
+    shiftpoints = points - ctr
+    U, S, Vt = svd(shiftpoints.T @ shiftpoints)
+    normal = U[:, -1]
+    return ctr, normal
+
+
+def angle_between_vectors(u, v, min_angle=True):
+    """Calculate the angle between two vectors in degrees.
+
+    Parameters
+    ----------
+    u : np.ndarray, shape (3,)
+        Vector
+    v : np.ndarray, shape (3,)
+        Vector
+    min_angle : bool, default True
+        Whether to return the supplement if the angle is greater than 90
+        degrees. Useful for calculating the minimum angle between the normal
+        vectors of planes as direction doesn't matter.
+
+    Returns
+    -------
+    angle: float
+        Angle between u and v in degrees
+    """
+    angle = np.rad2deg(
+        np.arccos(u.dot(v) / (np.linalg.norm(u) * np.linalg.norm(v)))
+    )
+    if angle > 90 and min_angle:
+        return 180 - angle
+    return angle

cmeutils/gsd_utils.py

@@ -5,23 +5,25 @@
 
 
 def get_type_position(
-        typename,
-        gsd_file=None,
-        snap=None,
-        gsd_frame=-1,
-        images=False):
-    """
-    This function returns the  positions of a particular particle
-    type from a frame of a gsd trajectory file or from a snapshot.
+    typename,
+    gsd_file=None,
+    snap=None,
+    gsd_frame=-1,
+    images=False
+):
+    """Get the positions of a particle type.
+
+    This function returns the positions of a particular particle type from a
+    frame of a gsd trajectory file or from a snapshot.
     Pass in either a gsd file or a snapshot, but not both.
 
     Parameters
     ----------
     typename : str or list of str
-        Name of particles of which to get the positions
-        (found in gsd.hoomd.Snapshot.particles.types)
-        If you want the positions of multiple types, pass
-        in a list. Ex.) ['ca', 'c3']
+        Name of particles of which to get the positions (found in
+        gsd.hoomd.Snapshot.particles.types)
+        If you want the positions of multiple types, pass in a list
+        e.g., ['ca', 'c3']
     gsd_file : str, default None
         Filename of the gsd trajectory
     snap : gsd.hoomd.Snapshot, default None
@@ -35,8 +37,7 @@ def get_type_position(
     Returns
     -------
     numpy.ndarray(s)
-        Retruns a single array of positions or
-        arrays of positions and images
+        Returns an array of positions or arrays of positions and images
     """
     snap = _validate_inputs(gsd_file, snap, gsd_frame)
     if isinstance(typename, str):
@@ -63,8 +64,7 @@ def get_type_position(
 
 
 def get_all_types(gsd_file=None, snap=None, gsd_frame=-1):
-    """
-    Returns all particle types in a hoomd trajectory
+    """Return all particle types in a hoomd trajectory.
 
     Parameters
     ----------
@@ -101,7 +101,7 @@ def snap_molecule_cluster(gsd_file=None, snap=None, gsd_frame=-1):
 
     Returns
     -------
-    numpy array (N_particles,)
+    numpy.ndarray (N_particles,)
     """
     snap = _validate_inputs(gsd_file, snap, gsd_frame)
     system = freud.AABBQuery.from_system(snap)
@@ -133,3 +133,54 @@ def _validate_inputs(gsd_file, snap, gsd_frame):
     elif snap:
         assert isinstance(snap, gsd.hoomd.Snapshot)
     return snap
+
+
+def snap_delete_types(snap, delete_types):
+    """Create a new snapshot with certain particle types deleted.
+
+    Reads in a snapshot and writes the information (excluding delete_types) to
+    a new snapshot. Does not change the original snapshot.
+
+    Information written:
+        - particles (N, types, position, typeid, image)
+        - configuration (box)
+        - bonds (N, group)
+
+    Parameters
+    ----------
+    snap : gsd.hoomd.Snapshot
+        The snapshot to read in
+
+    Returns
+    -------
+    gsd.hoomd.Snapshot
+        The new snapshot with particles deleted.
+    """
+    new_snap = gsd.hoomd.Snapshot()
+    delete_ids = [snap.particles.types.index(i) for i in delete_types]
+    selection = np.where(~np.isin(snap.particles.typeid, delete_ids))[0]
+    new_snap.particles.N = len(selection)
+    new_snap.particles.types = [
+        i for i in snap.particles.types if i not in delete_types
+    ]
+    typeid_map = {
+        i:new_snap.particles.types.index(e)
+        for i, e in enumerate(snap.particles.types)
+        if e in new_snap.particles.types
+    }
+    new_snap.particles.position = snap.particles.position[selection]
+    new_snap.particles.image = snap.particles.image[selection]
+    new_snap.particles.typeid = np.vectorize(typeid_map.get)(
+        snap.particles.typeid[selection]
+    )
+    new_snap.configuration.box = snap.configuration.box
+    if snap.bonds.N > 0:
+        bonds = np.isin(snap.bonds.group, selection).all(axis=1)
+        if bonds.any():
+            inds = {e:i for i, e in enumerate(selection)}
+            new_snap.bonds.group = np.vectorize(inds.get)(
+                snap.bonds.group[bonds]
+            )
+            new_snap.bonds.N = len(new_snap.bonds.group)
+    new_snap.validate()
+    return new_snap

cmeutils/tests/base_test.py

@@ -31,6 +31,12 @@ def snap(self, gsdfile):
             snap = f[-1]
         return snap
 
+    @pytest.fixture
+    def snap_bond(self, gsdfile_bond):
+        with gsd.hoomd.open(name=gsdfile_bond, mode="rb") as f:
+            snap = f[-1]
+        return snap
+
 
 def create_frame(i, add_bonds, images, seed=42):
     np.random.seed(seed)
@@ -42,10 +48,10 @@ def create_frame(i, add_bonds, images, seed=42):
     s.particles.position = np.random.random(size=(5, 3))
     s.configuration.box = [3, 3, 3, 0, 0, 0]
     if add_bonds:
-        s.bonds.N = 2
-        s.bonds.types = ["AB"]
-        s.bonds.typeid = [0, 0]
-        s.bonds.group = [[0, 2], [1, 3]]
+        s.bonds.N = 3
+        s.bonds.types = ["AB", "BB"]
+        s.bonds.typeid = [0, 0, 1]
+        s.bonds.group = [[0, 2], [1, 3], [3, 4]]
     if images:
         s.particles.image = np.full(shape=(5,3), fill_value=i)
     else:
@@ -56,4 +62,9 @@ def create_frame(i, add_bonds, images, seed=42):
 
 def create_gsd(filename, add_bonds=False, images=False):
     with gsd.hoomd.open(name=filename, mode="wb") as f:
-        f.extend((create_frame(i, add_bonds=add_bonds, images=images) for i in range(10)))
+        f.extend(
+            [
+                create_frame(i, add_bonds=add_bonds, images=images)
+                for i in range(10)
+            ]
+        )

cmeutils/tests/test_geometry.py

@@ -0,0 +1,29 @@
+import numpy as np
+import pytest
+
+from base_test import BaseTest
+from cmeutils.geometry import get_plane_normal, angle_between_vectors
+
+
+class TestGeometry(BaseTest):
+    def test_get_plane_normal(self):
+        points = np.array(
+            [[ 1, 0, 0],
+             [ 0, 1, 0],
+             [-1, 0, 0],
+             [ 0,-1, 0]]
+        )
+        ctr, norm = get_plane_normal(points)
+        assert np.allclose(ctr, np.array([0, 0, 0]))
+        assert np.allclose(norm, np.array([0, 0, 1]))
+
+        with pytest.raises(AssertionError):
+            get_plane_normal(points[:2])
+
+    def test_angle_between_vectors(self):
+        assert np.isclose(
+            90, angle_between_vectors(np.array([1,0,0]),np.array([0,1,0]))
+        )
+        assert np.isclose(
+            0, angle_between_vectors(np.array([1,0,0]),np.array([-1,0,0]))
+        )

cmeutils/tests/test_gsd.py

@@ -2,8 +2,10 @@
 import pytest
 
 from base_test import BaseTest
-from cmeutils.gsd_utils import get_type_position, snap_molecule_cluster
-from cmeutils.gsd_utils import get_all_types, _validate_inputs
+from cmeutils.gsd_utils import (
+    get_type_position, snap_molecule_cluster, get_all_types, _validate_inputs,
+    snap_delete_types
+)
 
 
 class TestGSD(BaseTest):
@@ -43,4 +45,12 @@ def test_get_all_types(self, gsdfile):
 
     def test_snap_molecule_cluster(self, gsdfile_bond):
         cluster = snap_molecule_cluster(gsd_file=gsdfile_bond)
-        assert np.array_equal(cluster, [0, 1, 0, 1, 2])
+        assert np.array_equal(cluster, [1, 0, 1, 0, 0])
+
+    def test_snap_delete_types(self, snap):
+        new_snap = snap_delete_types(snap, "A")
+        assert "A" not in new_snap.particles.types
+
+    def test_snap_delete_types_bonded(self, snap_bond):
+        new_snap = snap_delete_types(snap_bond, "A")
+        assert "A" not in new_snap.particles.types