reorganize dotfiles

profiles/pymol/.pymol/startup/center_of_mass.py

@@ -0,0 +1,86 @@
+'''
+See more here: http://www.pymolwiki.org/index.php/center_of_mass
+
+DESCRIPTION
+
+   Places a pseudoatom at the center of mass
+
+   Author: Sean Law
+   Michigan State University
+   slaw (at) msu . edu
+
+SEE ALSO
+
+   pseudoatom, get_com
+'''
+
+from __future__ import print_function
+from pymol import cmd
+
+
+def com(selection, state=None, mass=None, object=None, quiet=1, **kwargs):
+    quiet = int(quiet)
+    if (object == None):
+        try:
+            object = cmd.get_legal_name(selection)
+            object = cmd.get_unused_name(object + "_COM", 0)
+        except AttributeError:
+            object = 'COM'
+    cmd.delete(object)
+
+    if (state != None):
+        x, y, z = get_com(selection, mass=mass, quiet=quiet)
+        if not quiet:
+            print("%f %f %f" % (x, y, z))
+        cmd.pseudoatom(object, pos=[x, y, z], **kwargs)
+        cmd.show("spheres", object)
+    else:
+        for i in range(cmd.count_states()):
+            x, y, z = get_com(selection, mass=mass, state=i + 1, quiet=quiet)
+            if not quiet:
+                print("State %d:%f %f %f" % (i + 1, x, y, z))
+            cmd.pseudoatom(object, pos=[x, y, z], state=i + 1, **kwargs)
+            cmd.show("spheres", 'last ' + object)
+
+cmd.extend("com", com)
+
+
+def get_com(selection, state=1, mass=None, quiet=1):
+    """
+ DESCRIPTION
+
+    Calculates the center of mass
+
+    Author: Sean Law
+    Michigan State University
+    slaw (at) msu . edu
+    """
+    quiet = int(quiet)
+
+    totmass = 0.0
+    if mass != None and not quiet:
+        print("Calculating mass-weighted COM")
+
+    state = int(state)
+    model = cmd.get_model(selection, state)
+    x, y, z = 0, 0, 0
+    for a in model.atom:
+        if (mass != None):
+            m = a.get_mass()
+            x += a.coord[0] * m
+            y += a.coord[1] * m
+            z += a.coord[2] * m
+            totmass += m
+        else:
+            x += a.coord[0]
+            y += a.coord[1]
+            z += a.coord[2]
+
+    if (mass != None):
+        return x / totmass, y / totmass, z / totmass
+    else:
+        return x / len(model.atom), y / len(model.atom), z / len(model.atom)
+
+cmd.extend("get_com", get_com)
+
+# vi:expandtab:sw=3