plot_noe.py

"""
http://pymolwiki.org/index.php/plot_noe

(c) 2011-09-23 by Justin L Lorieau
(c) 2012-10-31 by Thomas Holder

License: BSD-2-Clause
"""

from __future__ import print_function

from pymol import cmd, CmdException
import re
import shlex


def plot_noe(filename, line_color=None, line_width='1.0', advanced_coloring=0, single=0, quiet=1, aria=0, per_atom=0,
             per_residue=1):
    """
DESCRIPTION

    A function for plotting XPLOR NOE restraints on a structure

ARGUMENTS

    filename = string: The filename of the NOE retraint file in XPLOR NIH format.

    line_color = string: The color for the NOE lines. {default: yellow}

    line_width = float: The thickness of the NOE lines. {default: 1.0}

    advanced_coloring = color restraints by distance.

    single = string: create a single object for all restraints.

    aria = integer: Name NOEs after Aria IDs.

    per_atom: Group NOEs on atom basis

    per_residue: Group NOEs on residue basis (default)

NOE Restraint Format Example

    assign (residue 5 and name HB#) (residue 21 and name HA) 3.0 0.7 0.7

EXAMPLE

    PyMOL> plot_noe noe_short.tbl
    """

    single, quiet, aria, per_residue, per_atom, advanced_coloring = \
        int(single), int(quiet), int(aria), int(per_residue), int(per_atom), int(advanced_coloring)

    count = 0

    if advanced_coloring == 1:
        cmd.set("group_auto_mode", 2)

    rgx_assi = re.compile("\s*[asignASIGN]+\s+.*")
    rgx_or = re.compile("\s*[orOR]+\s+.*")
    rgx_bracket = re.compile("(\(|\))")
    rgx_tick = re.compile("(\w)\'")

    restraint_line = None
    restraint_block = []
    restraints_blocks = []

    for line in open(filename):
        if rgx_assi.match(line):
            if restraint_line:
                restraint_block.append(rgx_tick.sub("\g<1>^", rgx_bracket.sub(" ", restraint_line)))
                restraints_blocks.append(restraint_block)
            restraint_block = []
            restraint_line = line
        elif rgx_or.match(line):
            restraint_block.append(rgx_tick.sub("\g<1>^", rgx_bracket.sub(" ", restraint_line)))
            restraint_line = line
        else:
            restraint_line += line

    restraint_block.append(rgx_tick.sub("\g<1>^", rgx_bracket.sub(" ", restraint_line)))
    restraints_blocks.append(restraint_block)

    for restraint_block in restraints_blocks:
        restraints = {"connections": []}

        try:
            shlex_split = shlex.split(restraint_block[0])
        except ValueError:
            print('Cannot process "%s"' % restraint_block[0])
        try:
            restraints["distance"] = float(shlex_split[11].strip())
        except IndexError:
            print('Cannot process "%s"' % "".join(restraint_block))
            continue
        except ValueError:
            try:
                restraints["distance"] = float(shlex_split[17].strip())
            except ValueError:
                print("Failed to extract distance, setting to 1A")
                restraints["distance"] = 1

        if len(shlex_split) > 20:
            restraints["ID"] = shlex_split[26][3:-1].strip()
        else:
            restraints["ID"] = str(count + 1)

        for restraint_line in restraint_block:
            distance = {}
            try:
                shlex_split = shlex.split(restraint_line)

                if len(shlex_split) > 6:
                    if shlex_split[1].strip().lower() == shlex_split[9].strip().lower() == "segid":
                        distance["src_segid"] = shlex_split[2].strip()
                        distance["src_resid"] = shlex_split[5].strip()
                        distance["src_atom"] = shlex_split[8].strip().replace("^", "'")
                        distance["dst_segid"] = shlex_split[10].strip()
                        distance["dst_resid"] = shlex_split[13].strip()
                        distance["dst_atom"] = shlex_split[16].strip().replace("^", "'")
                    else:
                        distance["src_segid"] = "A"
                        distance["src_resid"] = shlex_split[2].strip()
                        distance["src_atom"] = shlex_split[5].strip().replace("^", "'")
                        distance["dst_segid"] = "A"
                        distance["dst_resid"] = shlex_split[10].strip()
                        distance["dst_atom"] = shlex_split[13].strip().replace("^", "'")

                else:
                    print('Cannot process "%s"' % restraint_line)
                    continue
            except IndexError:
                print('Cannot process "%s"' % restraint_line)
                continue
            except ValueError:
                print('Cannot process "%s"' % restraint_line)
                continue

            restraints["connections"].append(distance)

        count += _draw_restraint(
            restraints,
            line_color,
            line_width,
            single,
            quiet,
            per_residue,
            per_atom,
            advanced_coloring
        )

    if not quiet:
        print(' Info: Created distance objects for %d restraints' % count)

    cmd.order("NOE*", "yes")


def _draw_restraint(restraints, line_color, line_width, single, quiet, per_residue, per_atom, advanced_coloring):
    cnt = 0
    for restraint in restraints["connections"]:
        cnt += 1
        sele1 = 'segid %s and resi %s and name %s' % (
            restraint["src_segid"], restraint["src_resid"], restraint["src_atom"].replace("#", "*")
        )
        sele2 = 'segid %s and resi %s and name %s' % (
            restraint["dst_segid"], restraint["dst_resid"], restraint["dst_atom"].replace("#", "*")
        )

        if line_color is None:
            if advanced_coloring == 1:
                if restraints["distance"] < 2 or restraints["distance"] > 6:
                    line_color = "red"
                elif 5 < restraints["distance"] < 6:
                    line_color = "orange"
                else:
                    line_color = "splitpea"
            else:
                line_color = "yellow"

        if per_atom or per_residue:
            for x, y, z in [
                [restraint["dst_segid"], restraint["dst_resid"], restraint["dst_atom"]],
                [restraint["src_segid"], restraint["src_resid"], restraint["src_atom"]]
            ]:
                if per_atom:
                    label = "NOE_%s_%s_%s" % (x, y, z)
                elif per_residue:
                    label = "NOE_%s_%s" % (x, y)
                if advanced_coloring == 1:
                    label += ".%s_%i" % (restraints["ID"], cnt)

                try:
                    cmd.distance(label, sele1, sele2, quiet=quiet,
                                 width=line_width, gap=0, label=0)
                    _color_restraint(label, line_color)
                except CmdException:
                    print('FAILED: %s - %s' % (sele1, sele2))
                    print(restraints["ID"])
                    return 0

        else:
            if single:
                label = "NOE"

            else:
                label = "NOE_" + restraints["ID"]

            try:
                cmd.distance(label, sele1, sele2, quiet=quiet,
                             width=line_width, gap=0, label=0)

                _color_restraint(label, line_color)
            except CmdException:
                print('FAILED: %s - %s' % (sele1, sele2))
                print(restraints["ID"])
                return 0

    return 1


def _color_restraint(object_name, color="yellow"):
    cmd.set("dash_color", color, object_name)


cmd.extend("plot_noe", plot_noe)