Source code for coarse_grained_conop

"""
.. codeauthor: Niklaus Johner <niklaus.johner@a3.epfl.ch>

This module contains some funcitons to connect coarse-grained martini structures
"""
from ost import *

__all__=('GenerateConnectivityDictFromITPFile','ConnectAtomsFromDict',\
         'ConnectMartiniProteinAtoms','ConnectMartiniAtomsHeuristic')

def GenerateConnectivityDictFromITPFile(filename,connectivity_dict={}):
  """
  This function reads in a martini itp file and generates a dicionary containing
  the connectivity information for the residues defined in the ITP file 
  """
  itp_file=open(filename,'r')
  for line in itp_file:
    s=line.strip()
    if len(s)==0:continue
    if s.startswith(';'):continue
    b=s.find('[')
    e=s.find(']')
    if b!=-1 and e!=-1 and e>b:
      block=s[b+1:e].strip()
      if block in ['moleculetype','atoms','bonds']:skip=False
      else: skip=True
      continue
    if skip:continue
    sp=line.split()
    if block=='moleculetype':
      molecule=sp[0]
      connectivity_dict[molecule]=[]
      atom_dict={}
      skip=True
      continue
    if block=='atoms':
      if len(sp)!=7:
        print 'problem for atom definition from',line
        continue
      atom_dict[sp[0]]=sp[4]
    if block=='bonds':
      if len(sp)!=5:
        print 'problem for bond definition from',line
        continue
      i1=sp[0]
      i2=sp[1]
      connectivity_dict[molecule].append([atom_dict[i1],atom_dict[i2]])
  return connectivity_dict

def ConnectAtomsFromDict(eh,connectivity_dict):
  """
  This function uses a connectivity dictionary to connect atoms of an Entity or EntityView
  """
  if type(eh)==mol.EntityView:edi=eh.GetHandle().EditXCS(mol.BUFFERED_EDIT)
  else:edi=eh.EditXCS(mol.BUFFERED_EDIT)
  for r in eh.residues:
    if not r.name in connectivity_dict:
      print r,'not in connectivity dict'
      continue
    bond_list=connectivity_dict[r.name]
    for bond in bond_list:
      a1=r.FindAtom(bond[0])
      a2=r.FindAtom(bond[1])
      edi.Connect(a1.handle,a2.handle)
  edi.ForceUpdate()
  return

  
def ConnectMartiniProteinAtoms(eh):
  """
  Heuristic connection builder for Martini proteins generated by the martinize.py script
  The backbone atoms have to be named BB
  """
  if type(eh)==mol.EntityView:edi=eh.GetHandle().EditXCS(mol.BUFFERED_EDIT)
  else:edi=eh.EditXCS(mol.BUFFERED_EDIT)
  r1=eh.residues[0]
  for r in eh.residues[1:]:
    a1=r1.FindAtom('BB')
    a2=r.FindAtom('BB')
    if not (a1.IsValid() and a2.IsValid()):continue
    if geom.Distance(a1.pos,a2.pos)<=10.0:edi.Connect(a1.handle,a2.handle)
    r1=r
  for r in eh.residues:
    #if r.name=='ILE':d=4.0
    #else:d=3.5
    d=4.0
    for a1 in r.atoms:
      for a2 in r.atoms:
        if geom.Distance(a1.pos,a2.pos)<=d:edi.Connect(a1.handle,a2.handle)
  edi.ForceUpdate()
  return

def ConnectMartiniAtomsHeuristic(eh,d=5.0):
  """
  This function simply connects atoms from the same resdiue and closer than the cutoff distance d
  """
  if type(eh)==mol.EntityView:edi=eh.GetHandle().EditXCS(mol.BUFFERED_EDIT)
  else:edi=eh.EditXCS(mol.BUFFERED_EDIT)
  for r in eh.residues:
    for i,a1 in enumerate(r.atoms):
      for j,a2 in enumerate(r.atoms):
        if i<=j:continue
        if geom.Distance(a1.pos,a2.pos)<=d:edi.Connect(a1.handle,a2.handle)
  edi.ForceUpdate()
  return