Add helper plotting function, update requirements

Author: Jan Ludwiczak

deepcoil/utils/__init__.py

@@ -1 +1 @@
-from .utils import corr_seq, is_fasta, sharpen_preds
+from .utils import corr_seq, is_fasta, sharpen_preds, plot_preds

deepcoil/utils/utils.py

@@ -1,4 +1,6 @@
 import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
 from itertools import groupby
 from Bio import SeqIO
 from scipy.signal import find_peaks
@@ -54,3 +56,76 @@ def sharpen_preds(probs):
             sharp_probs[beg:end] = max(sharp_probs[beg:end])
 
     return sharp_probs
+
+
+def plot_preds(results, beg=0, end=-1, out_file=None):
+    """
+    Helper function for plotting DeepCoil results
+    :param results: results for given entry returned by DeepCoil
+    :param beg: (optional) beginning aa of the range to use for plotting (useful for long sequences to see only subset of results)
+    :param end: (optional) end aa of the range to use for plotting (useful for long sequences to see only subset of results)
+    :param out_file: (optional) if specified results will also be dumped to file
+    :return:
+    """
+
+    sharp_probs = sharpen_preds(results['cc'])[beg:end]
+    probs, hept = results['cc'][beg:end], results['hept'][beg:end, :]
+
+    fig, (ax2, ax1) = plt.subplots(2, gridspec_kw={'height_ratios': [1, 9]}, figsize=(9, 7))
+
+    # Plot probs and sharpened probs
+    ax1.plot(probs, linewidth=1, c='gray', linestyle='dashed')
+    ax1.plot(sharp_probs, linewidth=3, c='black')
+
+    # Set axis limits
+    ax1.set_ylim(0.01, 1)
+    ax1.set_xlim(0, len(probs))
+
+    # Show grid
+    ax1.grid(linestyle="dashed", color='gray', linewidth=0.5)
+
+    # Find and set appropriate spacing of xticks given the sequence length
+    spacings = [10, 25] + list(range(50, 1000, 50))
+    spacing = spacings[np.argmin([abs(10 - len([i for i in range(0, len(probs), spacing)])) for spacing in spacings])]
+    ticks = [i for i in range(0, len(probs), spacing)]
+    labels = [i + beg for i in range(0, len(probs), spacing)]
+    labels[0] = ''
+    ax1.set_xticks(ticks)
+    ax1.set_xticklabels(labels)
+    ax1.set_yticks([i / 10 for i in range(1, 10, 1)])
+    ax1.xaxis.set_ticks_position('none')
+    ax1.set_xlabel('Sequence position', fontsize=16)
+
+    # Parse a, d heptad annotations and plot
+    a, d = [], []
+    for i, pr in enumerate(hept):
+        if pr[1] > 0.2 or pr[2] > 0.2:
+            if pr[1] > pr[2]:
+                a.append(pr[1])
+                d.append(0)
+            else:
+                d.append(-pr[2])
+                a.append(0)
+        else:
+            a.append(0)
+            d.append(0)
+    kk = np.vstack((a, d))
+    sns.heatmap(np.asarray(kk), cmap='bwr', vmin=-1, vmax=1, cbar=None, ax=ax2)
+
+    # Hide bottom panel of the subplot
+    for name, spine in ax2.spines.items():
+        if name != 'bottom':
+            spine.set_visible(True)
+
+    # Show the ticks corresponding to the bottom panel
+    ax2.xaxis.grid(linestyle="dashed", color='gray', linewidth=0.5)
+    ax2.set_xticks([i for i in range(0, len(probs), 50)])
+    ax2.set_yticklabels(['a', 'd'], rotation=0, fontsize=12)
+    ax2.set_xticks(ticks)
+    ax2.set_xticklabels(labels)
+    ax2.xaxis.set_ticks_position('none')
+
+    plt.tight_layout()
+    plt.subplots_adjust(hspace=0)
+    if out_file:
+        plt.savefig(out_file, dpi=300)

requirements.txt

@@ -4,3 +4,4 @@ numpy<1.19.0,>=1.16.0
 tensorflow>=2.3.0
 allennlp>=0.9.0,<0.10.0
 torch>=1.2,<2.0
+seaborn>=0.10