Merge pull request #16 from jennyfothergill/feat/order_parameter

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/structure.py

@@ -5,6 +5,7 @@
 from rowan import vector_vector_rotation
 
 from cmeutils import gsd_utils
+from cmeutils.geometry import get_plane_normal, angle_between_vectors
 
 
 def get_quaternions(n_views = 20):
@@ -145,3 +146,95 @@ def gsd_rdf(
 
         normalization = post_filter / pre_filter if exclude_bonded else 1
         return rdf, normalization
+
+
+def order_parameter(aa_gsd, cg_gsd, mapping, r_max, a_max, large=6, start=-10):
+    """Calculate the order parameter of a system.
+
+    The order parameter is used to describe the proportion of structures in
+    "large" clusters. The clustering takes into account the distance between
+    centers and angles between planes of structures and only considers neighbors
+    within the angle and distance cutoffs. This clustering metric has been used
+    in our previous work to quantify ordering in perylene
+    (DOI:10.1021/acsomega.6b00371) and thiophenes in p3ht
+    (DOI:10.3390/polym10121305).
+
+    In an attempt to be more transferrable, this functions relies on GRiTS
+    (https://github.com/cmelab/grits) to handle the finding and mapping of
+    structures within the atomistic system.
+
+    Example::
+
+        from grits import CG_System
+        from grits.utils import amber_dict
+
+        gsdfile = "trajectory.gsd"
+        cg_gsdfile = "cg-trajectory.gsd"
+        system = CG_System(
+            gsdfile,
+            beads={"_B" : "c1cscc1"},
+            conversion_dict=amber_dict
+        )
+        mapping = system.mapping["_B...c1cscc1"]
+        system.save(cg_gsdfile)
+        order, _ = order_parameter(gsdfile, cg_gsdfile, mapping, r_max, a_max)
+
+    Parameters
+    ----------
+    aa_gsd : str
+        Path to the atomistic gsd file
+    cg_gsd : str
+        Path to the coarse-grain gsd file
+    mapping : list of numpy.ndarray
+        List of arrays containing indices of atomistic particles which map to
+        each coarse-grain bead
+    r_max : float
+        Cut-off distance for the order parameter analysis
+    a_max : float
+        Cut-off angle for the order parameter analysis
+    large : int, default 6
+        The number of "beads" needed for a cluster to be considered "large"
+    start : int, default -10
+        The starting frame of the gsdfiles.
+
+    Returns
+    -------
+    order : list of float
+        Order parameter for each frame analyzed from the trajectory.
+    cl_idx : list of numpy.ndarray
+        The cluster index of each coarse-grain bead. See
+        freud.cluster.Cluster.cluster_idx
+    """
+    order = []
+    cl_idx = []
+    with gsd.hoomd.open(aa_gsd) as aa_f, gsd.hoomd.open(cg_gsd) as cg_f:
+        for aa_snap, cg_snap in zip(aa_f[start:], cg_f[start:]):
+            f_box = freud.box.Box.from_box(aa_snap.configuration.box)
+            unwrap_xyz = f_box.unwrap(
+                aa_snap.particles.position, aa_snap.particles.image
+            )
+            aq = freud.locality.AABBQuery.from_system(cg_snap)
+            n_list = aq.query(
+                cg_snap.particles.position,
+                query_args={"exclude_ii":True, "r_max": r_max}
+            ).toNeighborList()
+
+            vec_point = [
+                get_plane_normal(unwrap_xyz[mapping[i]])[1]
+                for i in n_list.point_indices
+            ]
+            vec_querypoint = [
+                get_plane_normal(unwrap_xyz[mapping[i]])[1]
+                for i in n_list.query_point_indices
+            ]
+            n_list.filter([
+                angle_between_vectors(i,j) < a_max
+                for i,j in zip(vec_point, vec_querypoint)
+            ])
+            cl = freud.cluster.Cluster()
+            cl.compute(cg_snap, neighbors=n_list)
+            n_large = sum([len(i) for i in cl.cluster_keys if len(i) >= large])
+            n_total = cg_snap.particles.N
+            order.append(n_large / n_total)
+            cl_idx.append(cl.cluster_idx)
+    return order, cl_idx