get_contact_bridges.py

#!/usr/bin/env python

"""
Forms bridges between pairs of atoms that share an interaction with a residue
matched by a user-specified regex. The output is a new list contact file where
interactions to atoms matching the regex have been replaced with bridges.

For example, running this app with the regex "A:CA:.*" and the contact list:
    0       vdw     A:ASP:52:C      A:CA:201:CA
    0       vdw     A:CA:201:CA     A:GLN:53:N
    0       vdw     A:ASP:52:N      A:PHE:48:O
will output the contact list
    0       br      A:ASP:52:C      A:GLN:53:N
    0       vdw     A:ASP:52:N      A:PHE:48:O
"""

from collections import defaultdict
import re
from itertools import combinations

from contact_calc.transformations import parse_contacts


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))


if __name__ == "__main__":
    main()


__license__ = "Apache License 2.0"
__maintainer__ = "Rasmus Fonseca"
__email__ = "fonseca.rasmus@gmail.com"