ParsePDB

import numpy as np 


def readPDB():
    atoms = []
    for line in open("1KNY.pdb.rtf"):
        list = line.split()
        id = list[0]
        if id == 'ATOM' or id == 'HETATM':
            coords = (list[6], list[7], list[8])
            coords = map(float,coords)
            atoms.append(coords)
    print np.array(atoms)

#If distance is less than 4.5 angstroms returns True, else returns False
def distance(atom1,atom2):
    if(np.sum((atom1-atom2)**2)<=20.25):
        return True
    else:
        return False

#Makes list or residues by parsing file
def makeRes():
    res = []
    current = ' '
    for line in open("1KNY.pdb.rtf"):
        list = line.split()
        if list[0] == 'ATOM'and list[3] != current:
            res.append(list[3])
            current = list[3]
    return res[0:15]

#makes a numpy array of the coordinates for all the atoms a certain residue
def makeCoords(id):
    atoms = []
    for line in open("1KNY.pdb.rtf"):
        list = line.split()
        if list [0] == 'ATOM' and list[3] == id:
            coords = (list[6],list[7],list[8])
            coords = map(float,coords)
            atoms.append(coords)
    return np.array(atoms)

#Compares two residues to see if there is a contact, returns 0 for false, 1 for true
def compRes(res1,res2):
   for i, x in enumerate(res1):
      for j,y in enumerate(res2):
         if(distance(x,y)):
            return 1
   return 0
  
def makeIdentity():
 
    res = makeRes()
    identity = [[0 for col in res] for row in res]
    for a, x in enumerate(res):
        r1 = makeCoords(x)
        for b, y in enumerate(res[a+1:]):
            r2 = makeCoords(y)
            identity[a][res.index(y)] = compRes(r1,r2)
    #        identity[res.index(y)][a] = compRes(r1,r2)
    identity = np.array(identity)
    identity = np.triu(identity, 4)
    np.savetxt('identityfullDiag.txt',np.array(identity),fmt = '%i')
    return identity 
  
def makeIDFile():
  identity = makeIdentity()
  start = 4
  file = open("idLine.txt","w")
  for i in identity:
    for j in i[start:]:
      file.write("%i" %(j))
    start += 1