center_of_mass.py

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