Lots of improvements for I/O of cns tbl file and grouping of restraints.

Signed-off-by: Justin Lecher <[email protected]>

plot_noe.py

@@ -8,8 +8,9 @@
 '''
 
 from pymol import cmd, CmdException
+import re
 
-def plot_noe(filename, selection='', line_color='gray20', line_width='1.0', single=0, quiet=1):
+def plot_noe(filename, selection='', line_color='gray20', line_width='1.0', single=0, quiet=1, aria=0, per_atom=0, per_residue=1):
     """
 DESCRIPTION
 
@@ -25,6 +26,12 @@ def plot_noe(filename, selection='', line_color='gray20', line_width='1.0', sing
 
     line_width = float: The thickness of the NOE lines. {default: 1.0}
 
+    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
@@ -35,39 +42,78 @@ def plot_noe(filename, selection='', line_color='gray20', line_width='1.0', sing
     """
     from pymol.parsing import split
 
-    single, quiet = int(single), int(quiet)
+    single, quiet, aria, per_residue, per_atom = int(single), int(quiet), int(aria), int(per_residue), int(per_atom)
 
     count = 0
+    parent_id = None
+
+    rgx_restraint = re.compile(
+                        r"""
+                            \s*(?P<parent>\w*)\s+\(
+                            (segid\s+\"\s*(?P<src_segid>[\w\d]+)\"\s+and\s+)?
+                            resid\s+(?P<src_resid>\d+)\s+and\s+
+                            name\s+(?P<src_atom>[\w\d]+)
+                            \)\s+\(
+                            (segid\s+\"\s+(?P<dst_segid>[\w\d]+)\"\s+and\s+)?
+                            resid\s*(?P<dst_resid>\d+)\s+and\s+
+                            name\s*(?P<dst_atom>[\w\d]+)
+                            \)
+                            (\s+(?P<dist>[\d\.]+))?
+                            (.*id=(?P<ID>\d+))?
+                            .*
+                        """, re.X)
+
 
     for line in open(filename):
-        line = line.replace('(', ' (').replace(')', ') ').replace('#', '*')
+        restraint = rgx_restraint.match(line)
+        count += 1
 
-        a = filter(None, split(line, ' \t'))
-        if len(a) < 4 or not a[0].startswith('assi'):
-            continue
+        if restraint:
+            restraint = restraint.groupdict()
 
-        try:
-            float(a[3])
-        except ValueError:
-            continue
+            if restraint["src_segid"] == None:
+                restraint["src_segid"] = "A"
 
-        sele1, sele2 = a[1:3]
+            if restraint["dst_segid"] == None:
+                restraint["dst_segid"] = "A"
 
-        if selection:
-            sele1 = '(%s) and (%s)' % (selection, sele1)
-            sele2 = '(%s) and (%s)' % (selection, sele2)
+            if restraint["parent"][:4].upper() == "ASSI":
+                parent = restraint
 
-        label = single and "NOE" or ("NOE_%d" % (count + 1))
+            sele1 = 'segid %s and resi %s and name %s' % (restraint["src_segid"], restraint["src_resid"], restraint["src_atom"])
+            sele2 = 'segid %s and resi %s and name %s' % (restraint["dst_segid"], restraint["dst_resid"], restraint["dst_atom"])
 
-        try:
-            cmd.distance(label, sele1, sele2, quiet=quiet,
-                    width=line_width, gap=0, label=0)
-        except CmdException:
-            print 'FAILED: %s - %s' % (sele1, sele2)
-            continue
+            if single:
+                label = "NOE"
 
-        cmd.set("dash_color", line_color, label)
-        count += 1
+            elif aria and restraint["ID"]:
+                label = "NOE_" + restraint["ID"]
+            elif aria and parent and parent["ID"]:
+                label = "NOE_" + parent["ID"]
+
+            elif per_atom:
+                label = "NOE_%s_%s_%s" % (restraint["dst_segid"], restraint["dst_resid"], restraint["dst_atom"])
+            elif per_residue:
+                label = "NOE_%s_%s" % (restraint["dst_segid"], restraint["dst_resid"])
+
+            elif restraint["parent"][:4].upper() == "ASSI":
+                label = parent_id
+            else:
+                parent_id = "NOE_%d" % count
+                label = parent_id
+
+            try:
+                cmd.distance(label, sele1, sele2, quiet=quiet,
+                        width=line_width, gap=0, label=0)
+            except CmdException:
+                print 'FAILED: %s - %s' % (sele1, sele2)
+                continue
+
+            cmd.set("dash_color", line_color, label)
+        else:
+            print line
+
+    cmd.order("NOE*", "yes")
 
     if not quiet:
         print ' Info: Created distance objects for %d restraints' % (count)