''' 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