Merge branch 'master' of github.com:Kortemme-Lab/pull_into_place

pull_into_place/structures.py

@@ -226,12 +226,15 @@ def read_and_calculate(workspace, pdb_paths):
 
                 # Keep track of this model's sequence.
 
-                if line.startswith('ATOM'): 
-                    if residue_id != last_residue_id:
+                if residue_id != last_residue_id:
+                    if line.startswith('ATOM'): 
                         one_letter_code = residue_type_3to1_map.get(residue_name, 'X')
                         sequence += one_letter_code
                         sequence_map[residue_id] = one_letter_code
                         last_residue_id = residue_id
+                    elif line.startswith('HETATM'):
+                        sequence_map[residue_id] = 'Lig'
+                        last_residue_id = residue_id
 
                 # Save the coordinate for this atom.  This will be used later 
                 # to calculate restraint distances.
@@ -295,6 +298,9 @@ def parse_restraints(path):
     parsers = {
             'CoordinateConstraint': CoordinateRestraint,
             'AtomPairConstraint': AtomPairRestraint,
+            'AtomPair': AtomPairRestraint,
+            'Dihedral': DihedralRestraint,
+            'NamedAngle': AngleRestraint,
     }
 
     with open(path) as file:
@@ -395,6 +401,53 @@ def xyz_to_array(xyz):
     """
     return np.array([float(x) for x in xyz])
 
+def angle(array_of_xyzs):
+    """
+    Calculates angle between three coordinate points (I could not find a package
+    that does this but if one exists that would probably be better). Used for Angle constraints. 
+    """
+    ab = array_of_xyzs[0] - array_of_xyzs[1]
+    cb = array_of_xyzs[2] - array_of_xyzs[1]
+    return np.arccos((np.dot(ab,cb)) / (np.sqrt(ab[0]**2 + ab[1]**2  \
+        + ab[2]**2) * np.sqrt(cb[0]**2 + cb[1]**2 + cb[2]**2)))
+
+def dihedral(array_of_xyzs):
+    """
+    Calculates dihedral angle between four coordinate points. Used for
+    dihedral constraints. 
+    """
+    p1 = array_of_xyzs[0]
+    p2 = array_of_xyzs[1]
+    p3 = array_of_xyzs[2]
+    p4 = array_of_xyzs[3]
+
+    vector1 = -1.0 * (p2 - p1)
+    vector2 = p3 - p2
+    vector3 = p4 - p3
+
+    # Normalize vector so as not to influence magnitude of vector
+    # rejections
+    vector2 /= np.linalg.norm(vector2)
+
+    # Vector rejections:
+    # v = projection of vector1 onto plane perpendicular to vector2
+    #   = vector1 - component that aligns with vector2
+    # w = projection of vector3 onto plane perpendicular to vector2
+    #   = vector3 = component that aligns with vector2
+    v = vector1 - np.dot(vector1, vector2) * vector2
+    w = vector3 - np.dot(vector3, vector2) * vector2
+
+    # Angle between v and w in a plane that is the torsion angle
+    # v and w not normalized explicitly, but we use tan so that doesn't
+    # matter
+    x = np.dot(v, w)
+    y = np.dot(np.cross(vector2, v), w)
+    principle_dihedral = np.arctan2(y,x)
+    # I'm leaving this variable explicit because it should be clear that
+    # we need to make sure there are no non-principle angles that better
+    # satisfy the constraint for some reason, but that needs to be done
+    # outside this function. 
+    return principle_dihedral
 
 def find_pareto_front(metrics, metadata, columns, depth=1, epsilon=1e-7, progress=None):
     """
@@ -582,11 +635,47 @@ def __init__(self, args):
         self.atom_names = [args[0], args[2]]
         self.residue_ids = [int(args[i]) for i in (1,3)]
         self.atom_pair = zip(self.atom_names, self.residue_ids)
+        self.ideal_distance = float(args[5])
 
     def distance_from_ideal(self, atom_xyzs):
         coords = [atom_xyzs[x] for x in self.atom_pair]
-        return euclidean(*coords)
+        return euclidean(*coords) - self.ideal_distance
 
+class DihedralRestraint(object):
+
+    def __init__(self, args):
+        self.atom_names = [args[0],args[2],args[4],args[6]]
+        self.residue_ids = [int(args[i]) for i in (1,3,5,7)]
+        self.atoms = zip(self.atom_names, self.residue_ids)
+        self.ideal_dihedral = float(args[9])
+
+    def distance_from_ideal(self, atom_xyzs):
+        coords = [atom_xyzs[x] for x in self.atoms]
+        measured_dihedral = dihedral(coords)
+        # Make sure we don't get the wrong number because of
+        # non-principle angles:
+        dihedrals = [abs(measured_dihedral - self.ideal_dihedral), abs(measured_dihedral + (2 * \
+                np.pi) - self.ideal_dihedral), abs(measured_dihedral - (2 * np.pi) - \
+                    self.ideal_dihedral)]
+        return min(dihedrals) 
+
+class AngleRestraint(object):
+
+    def __init__(self, args):
+        self.atom_names = [args[0],args[2],args[4],args[6]]
+        self.residue_ids = [int(args[i]) for i in (1,3,5)]
+        self.atoms = zip(self.atom_names, self.residue_ids)
+        self.ideal_angle = float(args[7])
+
+    def distance_from_ideal(self, atom_xyzs):
+        coords = [atom_xyzs[x] for x in self.atoms]
+        measured_angle = angle(coords)
+        # Make sure we don't get the wrong number because of
+        # non-principle angles:
+        angles = [abs(measured_angle - self.ideal_angle), abs(measured_angle + (2 * np.pi) -
+                self.ideal_angle), abs(measured_angle - (2 * np.pi) -
+                self.ideal_angle)]
+        return min(angles)
 
 class Design (object):
     """