mirror of
https://github.com/dnlbauer/dotfiles.git
synced 2026-09-10 13:35:30 +00:00
87 lines
2.1 KiB
Python
87 lines
2.1 KiB
Python
'''
|
|
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
|