First version of coordinate generatign processor

vermouth/processors/__init__.py

@@ -44,3 +44,4 @@
 from .sort_molecule_atoms import SortMoleculeAtoms
 from .merge_all_molecules import MergeAllMolecules
 from .annotate_mut_mod import AnnotateMutMod
+from .generate_coordinates import GenerateCoordinates

vermouth/processors/generate_coordinates.py

@@ -21,19 +21,20 @@
 import networkx as nx
 
 from .processor import Processor
+from .. import selectors
 
 
 def mindist(A, B):
     return ((A[:, None] - B[None])**2).sum(axis=2).min(axis=1)
 
 
-def get_missing_atoms(molecule):
+def get_atoms_missing_coords(molecule):
     '''
     Determine particles without coordinates in molecule.
     '''
     return [
         ix for ix in molecule
-        if molecule.nodes[ix].get('position') is None
+        if not selectors.selector_has_position(molecule.nodes[ix])
     ]
 
 
@@ -44,7 +45,7 @@ def get_anchors(molecule, indices):
     '''
     return {
         a for ix in indices for a in molecule[ix].keys()
-        if molecule.nodes[a].get('position') is not None
+        if not selectors.selector_has_position(molecule.nodes[a]) is not None
     }
 
 
@@ -95,7 +96,7 @@ def triple_cone(n=24, distance=0.125, angle=109.4):
     return P, S, T
 
 
-def _out(molecule, anchor, target=None, control=None, distance=0.125):
+def _out(molecule, anchor, base, distance=0.125):
     '''
     Generate coordinates for atom bonded to anchor
     using resultant vector of known substituents.
@@ -107,9 +108,13 @@ def _out(molecule, anchor, target=None, control=None, distance=0.125):
      b1
        \
     b2--a-->X
-       /
+       /  u
      b3
 
+    X := position to determine
+    a := position of anchoring particle
+    B = [b1, b2, ...] := positions of base particles
+
     Parameters
     ----------
 
@@ -118,25 +123,14 @@ def _out(molecule, anchor, target=None, control=None, distance=0.125):
         None
     '''
     a = molecule.nodes[anchor]['position']
-    if control == None:
-        B = [
-            molecule.nodes[ix]['position']
-            for ix, v in molecule[anchor].items()
-            if v != {}
-        ]
-    else:
-        B = [ molecule.nodes[ix]['position'] for ix in control ]
+    B = [ molecule.nodes[ix]['position'] for ix in base ]
     B = B - a
     u = -B.sum(axis=0)
     u *= distance / ((u ** 2).sum() ** 0.5)
-    pos = a + u
-    if target is not None:
-        molecule.nodes[target]['position'] = pos
-
-    return pos
+    return a + u
 
 
-def _chiral(molecule, anchor, up=None, down=None, control=None, distance=0.125):
+def _chiral(molecule, anchor, base, distance=(0.125, 0.125)):
     '''
     Generate coordinates for atoms bonded to chiral center
     based on two known substituents.
@@ -164,32 +158,19 @@ def _chiral(molecule, anchor, up=None, down=None, control=None, distance=0.125):
         None
     '''
     a = molecule.nodes[anchor]['position']
-    if control is None:
-        B = [
-            molecule.nodes[ix]['position']
-            for ix, v in molecule[anchor].items()
-            if v != {}
-        ]
-    else:
-        B = [ molecule.nodes[ix]['position'] for ix in control ]
+    B = [ molecule.nodes[ix]['position'] for ix in base ]
     B = B - a
     u = -0.5 * (B[0] + B[1])
     v = np.cross(B[0], B[1])
     # Set length of v to length to half distance between subs
     v = 0.5 * ((B[0] - B[1])**2).sum()**0.5 * v / (v ** 2).sum()**0.5
     # Normalization of vector sum
     n = 1 / ((u + v)**2).sum()**0.5
-    # Duck type way to see if we got one or two distances
-    try:
-        rup, rdown = distance
-    except TypeError:
-        rup = rdown = distance
+
+    rup, rdown = distance
+
     pup = a + rup * n * (u + v)
     pdown = a + rdown * n * (u - v)
-    if up is not None:
-        molecule.nodes[up]['position'] = pup
-    if down is not None:
-        molecule.nodes[down]['position'] = pdown
 
     return pup, pdown
 
@@ -214,7 +195,7 @@ def extrude(self, distance=0.125, coords=None, contact=0.5):
 
             if valence > 2 and len(b) == valence - 1:
                 # Trivial chiral/flat/bipyramidal/octahedral/...
-                pos = _out(mol, a, target[0], control=b)
+                mol.nodes[target[0]]['position'] = _out(mol, anchor=a, base=b)
                 for k in self.limbs:
                     k.discard(target[0])
                 # Simply update this segment and return it
@@ -226,7 +207,9 @@ def extrude(self, distance=0.125, coords=None, contact=0.5):
                 # Trivial chiral - except for the actual chirality
                 # a simple 'up' or 'down' flag would be helpful
                 # amino acid side chain is 'up'
-                up, down = _chiral(mol, a, target[0], target[1], control=b)
+                up, down = _chiral(mol, anchor=a, base=b, up=target[0], down=target[1])
+                mol.nodes[target[0]]['position'] = up
+                mol.nodes[target[1]]['position'] = down
                 # Simply update this segment and return it
                 for k in self.limbs:
                     k.discard(target[0])
@@ -326,7 +309,7 @@ def run_molecule(self, molecule):
         """
 
         # Missing atoms - those without 'positions'
-        missing = get_missing_atoms(molecule)
+        missing = get_atoms_missing_coords(molecule)
 
         prev = len(molecule)
         while prev - len(missing) > 0:
@@ -356,42 +339,8 @@ def run_molecule(self, molecule):
                     segments.extend(r)
 
             prev = len(missing)
-            missing = get_missing_atoms(molecule)
+            missing = get_atoms_missing_coords(molecule)
             #print(len(missing))
 
         return molecule
 
-
-def main(args):
-    '''
-    Main function for running processor with a PDB file as input (testing).
-    '''
-    pdb = args[1]
-
-    system = vermouth.System()
-    vermouth.PDBInput(pdb).run_system(system)
-    system.force_field = vermouth.forcefield.get_native_force_field('universal')
-
-    flow = [
-        vermouth.MakeBonds(allow_name=True, allow_dist=True, fudge=1),
-        vermouth.RepairGraph(delete_unknown=True, include_graph=True),
-        vermouth.SortMoleculeAtoms(),
-        GenerateCoordinates()
-    ]
-
-    for process in flow:
-        process.run_system(system)
-
-    vermouth.pdb.write_pdb(system, args[2], nan_missing_pos=True)
-
-    return 0
-
-
-if __name__ == '__main__':
-    import sys
-    import vermouth
-    import vermouth.forcefield
-    import vermouth.map_input
-
-    exco = main(sys.argv)
-    sys.exit(exco)