Analyze the neighborhood around special residues in a set of PDB files

analyze_neighboring_residues.py

#!/usr/bin/env python
# -*- coding: UTF-8 -*-
from __future__ import division
from marshal import dump, load
from random import sample
from glob import glob
from pdb import *
import sys


if len(sys.argv) <= 1:
  print("Please provide a datafile")
  sys.exit(1)
residue_data = sys.argv[1]

print("Reading residue data...")
residues_file = open(residue_data, 'rb')
all_residues  = load(residues_file)
residues_file.close()
num_residues  = len(all_residues)

#############################################################################

targets  = ['ASP', 'ASN', 'GLU', 'GLN', 'VAL']
ss_codes = [None, 'loop', 'helix', 'sheet']

for target in targets:
  for ss_code in ss_codes:
    ss = 'overall'
    if ss_code != None: ss = ss_code

    print("Calculating %s (%s)..." % (target, ss))
    real_stdout = sys.stdout
    sys.stdout = open("ratios_%s_%s.txt" % (target.lower(), ss), 'w+')
    print("Results: %s (%s)" % (target, ss))

    base_total, base_counts = 0, {}
    k01_total, k01_counts   = 0, {}
    k02_total, k02_counts   = 0, {}
    k03_total, k03_counts   = 0, {}
    k04_total, k04_counts   = 0, {}
    k05_total, k05_counts   = 0, {}
    k06_total, k06_counts   = 0, {}
    k07_total, k07_counts   = 0, {}
    k08_total, k08_counts   = 0, {}
    k09_total, k09_counts   = 0, {}
    k10_total, k10_counts   = 0, {}
    print("")

    for i, residue in enumerate(all_residues):
      name = residue['resName']

      if ss_code == None or ss_code == residue['ssCode']:
        try:
          base_counts[name] += 1
        except:
          base_counts[name] = 1
        finally:
          base_total += 1

      if i-1 >= 0 and target == all_residues[i-1]['resName']:
        if ss_code == None or ss_code == all_residues[i-1]['ssCode']:
          try:
            k01_counts[name] += 1
          except:
            k01_counts[name] = 1
          finally:
            k01_total += 1

      if i-2 >= 0 and target == all_residues[i-2]['resName']:
        if ss_code == None or ss_code == all_residues[i-2]['ssCode']:
          try:
            k02_counts[name] += 1
          except:
            k02_counts[name] = 1
          finally:
            k02_total += 1

      if i-3 >= 0 and target == all_residues[i-3]['resName']:
        if ss_code == None or ss_code == all_residues[i-3]['ssCode']:
          try:
            k03_counts[name] += 1
          except:
            k03_counts[name] = 1
          finally:
            k03_total += 1

      if i-4 >= 0 and target == all_residues[i-4]['resName']:
        if ss_code == None or ss_code == all_residues[i-4]['ssCode']:
          try:
            k04_counts[name] += 1
          except:
            k04_counts[name] = 1
          finally:
            k04_total += 1

      if i-5 >= 0 and target == all_residues[i-5]['resName']:
        if ss_code == None or ss_code == all_residues[i-5]['ssCode']:
          try:
            k05_counts[name] += 1
          except:
            k05_counts[name] = 1
          finally:
            k05_total += 1

      if i-6 >= 0 and target == all_residues[i-6]['resName']:
        if ss_code == None or ss_code == all_residues[i-6]['ssCode']:
          try:
            k06_counts[name] += 1
          except:
            k06_counts[name] = 1
          finally:
            k06_total += 1

      if i-7 >= 0 and target == all_residues[i-7]['resName']:
        if ss_code == None or ss_code == all_residues[i-7]['ssCode']:
          try:
            k07_counts[name] += 1
          except:
            k07_counts[name] = 1
          finally:
            k07_total += 1

      if i-8 >= 0 and target == all_residues[i-8]['resName']:
        if ss_code == None or ss_code == all_residues[i-8]['ssCode']:
          try:
            k08_counts[name] += 1
          except:
            k08_counts[name] = 1
          finally:
            k08_total += 1

      if i-9 >= 0 and target == all_residues[i-9]['resName']:
        if ss_code == None or ss_code == all_residues[i-9]['ssCode']:
          try:
            k09_counts[name] += 1
          except:
            k09_counts[name] = 1
          finally:
            k09_total += 1

      if i-10 >= 0 and target == all_residues[i-10]['resName']:
        if ss_code == None or ss_code == all_residues[i-10]['ssCode']:
          try:
            k10_counts[name] += 1
          except:
            k10_counts[name] = 1
          finally:
            k10_total += 1


      if i+1 < num_residues and target == all_residues[i+1]['resName']:
        if ss_code == None or ss_code == all_residues[i+1]['ssCode']:
          try:
            k01_counts[name] += 1
          except:
            k01_counts[name] = 1
          finally:
            k01_total += 1

      if i+2 < num_residues and target == all_residues[i+2]['resName']:
        if ss_code == None or ss_code == all_residues[i+2]['ssCode']:
          try:
            k02_counts[name] += 1
          except:
            k02_counts[name] = 1
          finally:
            k02_total += 1

      if i+3 < num_residues and target == all_residues[i+3]['resName']:
        if ss_code == None or ss_code == all_residues[i+3]['ssCode']:
          try:
            k03_counts[name] += 1
          except:
            k03_counts[name] = 1
          finally:
            k03_total += 1

      if i+4 < num_residues and target == all_residues[i+4]['resName']:
        if ss_code == None or ss_code == all_residues[i+4]['ssCode']:
          try:
            k04_counts[name] += 1
          except:
            k04_counts[name] = 1
          finally:
            k04_total += 1

      if i+5 < num_residues and target == all_residues[i+5]['resName']:
        if ss_code == None or ss_code == all_residues[i+5]['ssCode']:
          try:
            k05_counts[name] += 1
          except:
            k05_counts[name] = 1
          finally:
            k05_total += 1

      if i+6 < num_residues and target == all_residues[i+6]['resName']:
        if ss_code == None or ss_code == all_residues[i+6]['ssCode']:
          try:
            k06_counts[name] += 1
          except:
            k06_counts[name] = 1
          finally:
            k06_total += 1

      if i+7 < num_residues and target == all_residues[i+7]['resName']:
        if ss_code == None or ss_code == all_residues[i+7]['ssCode']:
          try:
            k07_counts[name] += 1
          except:
            k07_counts[name] = 1
          finally:
            k07_total += 1

      if i+8 < num_residues and target == all_residues[i+8]['resName']:
        if ss_code == None or ss_code == all_residues[i+8]['ssCode']:
          try:
            k08_counts[name] += 1
          except:
            k08_counts[name] = 1
          finally:
            k08_total += 1

      if i+9 < num_residues and target == all_residues[i+9]['resName']:
        if ss_code == None or ss_code == all_residues[i+9]['ssCode']:
          try:
            k09_counts[name] += 1
          except:
            k09_counts[name] = 1
          finally:
            k09_total += 1

      if i+10 < num_residues and target == all_residues[i+10]['resName']:
        if ss_code == None or ss_code == all_residues[i+10]['ssCode']:
          try:
            k10_counts[name] += 1
          except:
            k10_counts[name] = 1
          finally:
            k10_total += 1


    def print_ratios(ratios):
      for name in ratios:
        print("%s\t%.3f" % (name, ratios[name]))

    sec_struct = 'overall'
    if ss_code != None:
      sec_struct = "in a %s" % ss_code

    base_ratios = {}
    for name in base_counts:
      base_ratios[name] = base_counts[name] / base_total
    print("baseline:")
    print_ratios(base_ratios)
    print("")

    # if ss_code != None:
    k01_ratios = {}
    for name in k01_counts:
      k01_ratios[name] = k01_counts[name] / k01_total
    print("k=1:")
    print_ratios(k01_ratios)
    print("")

    k02_ratios = {}
    for name in k02_counts:
      k02_ratios[name] = k02_counts[name] / k02_total
    print("k=2:")
    print_ratios(k02_ratios)
    print("")

    k03_ratios = {}
    for name in k03_counts:
      k03_ratios[name] = k03_counts[name] / k03_total
    print("k=3:")
    print_ratios(k03_ratios)
    print("")

    k04_ratios = {}
    for name in k04_counts:
      k04_ratios[name] = k04_counts[name] / k04_total
    print("k=4:")
    print_ratios(k04_ratios)
    print("")

    k05_ratios = {}
    for name in k05_counts:
      k05_ratios[name] = k05_counts[name] / k05_total
    print("k=5:")
    print_ratios(k05_ratios)
    print("")

    k06_ratios = {}
    for name in k06_counts:
      k06_ratios[name] = k06_counts[name] / k06_total
    print("k=6:")
    print_ratios(k06_ratios)
    print("")

    k07_ratios = {}
    for name in k07_counts:
      k07_ratios[name] = k07_counts[name] / k07_total
    print("k=7:")
    print_ratios(k07_ratios)
    print("")

    k08_ratios = {}
    for name in k08_counts:
      k08_ratios[name] = k08_counts[name] / k08_total
    print("k=8:")
    print_ratios(k08_ratios)
    print("")

    k09_ratios = {}
    for name in k09_counts:
      k09_ratios[name] = k09_counts[name] / k09_total
    print("k=9:")
    print_ratios(k09_ratios)
    print("")

    k10_ratios = {}
    for name in k10_counts:
      k10_ratios[name] = k10_counts[name] / k10_total
    print("k=10:")
    print_ratios(k10_ratios)

    sys.stdout.close()
    sys.stdout = real_stdout