/home/ferreirafm/bin/geometry_mod.py

homeferreirafmbingeometry_mod.py.txt

#!/usr/bin/env python                                                                                                                            # Copyright (C) 2012, Joao Rodrigues (anaryin@gmail.com) and Ezgi Karaca (ezzgikaraca@gmail.com)
# This code is part of the Biopython distribution and governed by its
# license.  Please see the LICENSE file that should have been included
# as part of this package.

"""
Module with assorted geometrical functions on
macromolecules.
"""

from Bio.PDB import Entity
from Bio.PDB.PDBParser import PDBParser
from math import sqrt
from numpy import *
import sys
import glob

# Globals
parser = PDBParser()

# Functions
def center_of_mass(entity, geometric=False):
    """
    Returns gravitic [default] or geometric center of mass of an Entity.
    Geometric assumes all masses are equal (geometric=True)
    """
    # Structure, Model, Chain, Residue
    if isinstance(entity, Entity.Entity):
        atom_list = entity.get_atoms()
    # List of Atoms
    elif hasattr(entity, '__iter__') and [x for x in entity if x.level == 'A']:
        atom_list = entity
    else: # Some other weirdo object
        raise ValueError("Center of Mass can only be calculated from the following objects:\n"
                            "Structure, Model, Chain, Residue, list of Atoms.")
    class COM:
        def __init__(self,coord):
	    self.coord=coord
       
    positions = [ [], [], [] ] # [ [X1, X2, ..] , [Y1, Y2, ...] , [Z1, Z2, ...] ]
    masses = []
    
    for atom in atom_list:
        masses.append(atom.mass)
        
        for i, coord in enumerate(atom.coord.tolist()):
            positions[i].append(coord)

     If there is a single atom with undefined mass complain loudly.
    if 'ukn' in set(masses) and not geometric:
        raise ValueError("Some Atoms don't have an element assigned.\n"
                         "Try adding them manually or calculate the geometrical center of mass instead.")
    
    if geometric:
	com = COM([sum(coord_list)/len(masses) for coord_list in positions])
        return com
    else:       
        w_pos = [ [], [], [] ]
        for atom_index, atom_mass in enumerate(masses):
            w_pos[0].append(positions[0][atom_index]*atom_mass)
            w_pos[1].append(positions[1][atom_index]*atom_mass)
            w_pos[2].append(positions[2][atom_index]*atom_mass)
	com = COM([sum(coord_list)/sum(masses) for coord_list in w_pos])
        return com
        
def calculate_shape_param(structure):
  
  """
  Calculates the gyration tensor of a structure.  
  Returns a tuple containing shape parameters:
  
    ((a,b,c), rg, A)

    (a,b,c) - dimensions of the structure in each 3D direction
    rg 	  - radius of gyration of the structure
    A     - anisotropy value
  """

  com = center_of_mass(structure, True)
  cx, cy, cz = com.coord

  n_atoms = 0
  tensor_xx, tensor_xy, tensor_xz = 0, 0, 0
  tensor_yx, tensor_yy, tensor_yz = 0, 0, 0
  tensor_zx, tensor_zy, tensor_zz = 0, 0, 0

  for atom in structure.get_atoms():
    ax, ay, az = atom.coord
    tensor_xx += (ax-cx)*(ax-cx)
    tensor_yx += (ax-cx)*(ay-cy)
    tensor_xz += (ax-cx)*(az-cz)
    tensor_yy += (ay-cy)*(ay-cy)
    tensor_yz += (ay-cy)*(az-cz)
    tensor_zz += (az-cz)*(az-cz)
    n_atoms += 1

  gy_tensor =  mat([[tensor_xx, tensor_yx, tensor_xz], [tensor_yx, tensor_yy, tensor_yz], [tensor_xz, tensor_yz, tensor_zz]])
  gy_tensor = (1.0/n_atoms) * gy_tensor
  
  D,V = linalg.eig(gy_tensor)
  [a, b, c] = sorted(sqrt(5 * D))
  rg = sqrt(sum(D))
  
  l = average([D[0],D[1],D[2]])
  A = (((D[0] - l)**2 + (D[1] - l)**2 + (D[2] - l)**2) / l**2) * 1/6
  S = (((D[0] - l) * (D[1] - l) * (D[2] - l))/ l**3)
  
  print "%s" % '#Dimensions(a,b,c) #Rg #Anisotropy'
  print "%.2f" % round(2*a,2), round(2*b,2), round(2*c,2) , round(rg,2) , round(A,2)  # I doubled axes here!

  return ((a,b,c),rg,A)


def calculate_moment_of_intertia_tensor(structure):
  """
  Calculates the moment of inertia tensor from the molecule.
  Returns a numpy matrix.
  """
  com = center_of_mass(structure, False)
  cx, cy, cz = com.coord

  n_atoms = 0
  tensor_xx, tensor_xy, tensor_xz = 0, 0, 0
  tensor_yx, tensor_yy, tensor_yz = 0, 0, 0
  tensor_zx, tensor_zy, tensor_zz = 0, 0, 0
  s_mass = sum([a.mass for a in structure.get_atoms()])

  for atom in structure.get_atoms():
    ax, ay, az = atom.coord
    tensor_xx += ((ay-cy)**2 + (az-cz)**2)*atom.mass
    tensor_xy += -1*(ax-cx)*(ay-cy)*atom.mass
    tensor_xz += -1*(ax-cx)*(az-cz)*atom.mass
    tensor_yy += ((ax-cx)**2 + (az-cz)**2)*atom.mass
    tensor_yz += -1*(ay-cy)*(az-cz)*atom.mass
    tensor_zz += ((ax-cx)**2 + (ay-cy)**2)*atom.mass

  in_tensor =  mat([[tensor_xx, tensor_xy, tensor_xz], [tensor_xy, tensor_yy, tensor_yz], [tensor_xz, 
tensor_yz, tensor_zz]])
  D,V = linalg.eig(in_tensor)
  
  a = sqrt((5/(2*s_mass)) * (D[0] - D[1] + D[2]))
  b = sqrt((5/(2*s_mass)) * (D[2] - D[0] + D[1]))
  c = sqrt((5/(2*s_mass)) * (D[1] - D[2] + D[0]))

  return sorted([a, b, c])

def maxD(pdbf):
    """
    Calculates the shape maximum dimension
    """
    d = 0
    Dmax = 0
    struct = parser.get_structure('tmp', pdbf)
    for chain1 in struct[0]:
        for res1 in chain1:
            for atom1 in res1:
                x1, y1, z1 = atom1.coord
                for chain2 in struct[0]:
                    for res2 in chain2:
                        for atom2 in res2:
                            x2, y2, z2 = atom2.coord
                            d =  sqrt( (x1 - x2)**2 + (y1 - y2)**2 + (z1 - z2)**2)
                            if d >= Dmax:
                                Dmax = d
    return Dmax
                 

# Main Function
def main():
    args = sys.argv[1:]
    if args[0] == '--calc':
        for pdbf in sorted(glob.iglob('*.pdb')):
                struct = parser.get_structure('tmp', pdbf)
                mod = struct[0]
                chain = mod['A']
                print pdbf
                calculate_shape_param(chain)
                print 'Dmax: %.2f\n' % maxD(pdbf)
                
    elif args[0] == '--max':
            maxD(args[1])    



# call main function
if __name__ == '__main__':
    main()