Pymol（1.8.6）作图技巧之创建质心

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