def test_res_contacts(self):
contacts, num_frames = ct.parse_contacts(self.input_lines, set(['sb']))
rcontacts = ct.res_contacts(contacts)
self.assertEqual(len(rcontacts), 2)
self.assertEqual(rcontacts[0][0], 0)
self.assertEqual(rcontacts[1][0], 1)
self.assertEqual(rcontacts[1][1], "A:ARG:76")
self.assertEqual(rcontacts[1][2], "A:GLU:82")
def test_parse_contacts(self):
contacts, num_frames = ct.parse_contacts(self.input_lines, set(['sb']))
self.assertEqual(len(contacts), 3)
self.assertEqual(num_frames, 2)
self.assertTrue(all([c[1] == 'sb' for c in contacts]))
self.assertEqual(type(contacts[0][0]), int)
self.assertEqual(type(contacts[0][1]), str)
self.assertEqual(type(contacts[0][2]), str)
self.assertEqual(type(contacts[0][3]), str)
required = parser.add_argument_group('required arguments')
parser._action_groups.append(optional) # added this line
required.add_argument('--input',
required=True,
type=ap.FileType('r'),
metavar='FILE',
help='A contact-file or frequency file')
required.add_argument('--output',
required=False,
metavar='FILE',
type=ap.FileType('w'),
help='The name of the output contact-file')
args = parser.parse_args()
contacts, total_frames = parse_contacts(args.input)
# Build a dict mapping residue ids to a set of all its atom names
representative_atoms = defaultdict(list)
for c in contacts:
resi1 = c[2][0:c[2].rindex(":")]
name1 = c[2][c[2].rindex(":") + 1:]
representative_atoms[resi1].append(name1)
resi2 = c[3][0:c[3].rindex(":")]
name2 = c[3][c[3].rindex(":") + 1:]
representative_atoms[resi2].append(name2)
# Reduce the sets of atoms to a single representative atoms
representative_atoms = {
resi: "CA" if "CA" in names else names[0]
for (resi, names) in representative_atoms.items()
help='A contact-file or frequency file')
optional.add_argument('--output',
required=False,
metavar='FILE',
type=ap.FileType('w'),
help='The name of the output contact-file')
optional.add_argument('--min_frequency',
required=False,
metavar='FLOAT',
default=0.6,
type=float,
help='Minimum residue frequency')
args = parser.parse_args()
contacts, num_frames = parse_contacts(args.input)
contacts = multi_to_single_contact(contacts, args.min_frequency * num_frames)
contacts = [str(c[0]) + "\t" + "\t".join(c[1:]) for c in contacts]
# Write to output
if args.output:
args.output.write("\n".join(contacts))
args.output.close()
print("Wrote residue contact file to " + args.output.name)
else:
print("\n".join(contacts))
__license__ = "Apache License 2.0"
__maintainer__ = "Rasmus Fonseca"
__email__ = "*****@*****.**"
def main(argv=None):
# Parse command line arguments
import argparse as ap
parser = ap.ArgumentParser(description=__doc__, formatter_class=ap.RawTextHelpFormatter)
optional = parser._action_groups.pop()
required = parser.add_argument_group('required arguments')
parser._action_groups.append(optional) # added this line
required.add_argument('--input',
required=True,
type=ap.FileType('r'),
metavar='FILE',
help='A contact-file generated by get_dynamic_contacts.py or get_static_contacts.py')
required.add_argument('--bridge',
required=True,
type=str,
metavar='REGEX',
help='Regular expression matching any atom to be included as a bridge')
required.add_argument('--bridges_only',
required=False,
type=bool,
metavar='BOOL',
default=False,
help='Indicates whether to output non-bridged interactions as well as the bridges')
required.add_argument('--output',
required=False,
metavar='FILE',
type=ap.FileType('w'),
help='The name of the output contact-file')
args = parser.parse_args(argv)
contacts, total_frames = parse_contacts(args.input)
bridges_only = args.bridges_only
# Build the bridge_neighbor datastructure which for each frame has a dictionary mapping bridging-residues to
# non-bridging neighbors. Also collects contacts in `bridged_contacts` that are not part of any bridges unless
# `bridges_only` has been enabled
bridge_neighbors = [defaultdict(list) for _ in range(total_frames)]
bridge_pattern = re.compile(args.bridge)
bridged_contacts = []
for contact in contacts:
frame = contact[0]
a1_match = bridge_pattern.match(contact[2])
a2_match = bridge_pattern.match(contact[3])
if a1_match and not a2_match:
a1_res = ":".join(contact[2].split(":")[0:3])
bridge_neighbors[frame][a1_res].append(contact[3])
elif a2_match:
a2_res = ":".join(contact[3].split(":")[0:3])
bridge_neighbors[frame][a2_res].append(contact[2])
elif not bridges_only:
bridged_contacts.append(contact)
# Based on the neighbor-lists in `bridge_neighbors`, add atom pairs to `bridged_contacts`
for frame, bridge_map in enumerate(bridge_neighbors):
for bridge_res in bridge_map:
for a1, a2 in combinations(bridge_map[bridge_res], 2):
bridged_contacts.append([frame, 'br', a1, a2, bridge_res])
# Sort the contacts and convert them to strings
from operator import itemgetter
bridged_contacts.sort(key=itemgetter(0))
for contact in bridged_contacts:
contact[0] = str(contact[0])
bridged_contacts = ["\t".join(contact) for contact in bridged_contacts]
# Write to output
if args.output:
args.output.write("# total_frames:%d\n" % total_frames)
args.output.write("\n".join(bridged_contacts))
args.output.close()
print("Wrote residue contact file to " + args.output.name)
else:
print("\n".join(bridged_contacts))