Add ploting feature to standalone version and correct plotting internals to correctly handle multiple plots in a row

Author: Jan Ludwiczak

bin/deepcoil

@@ -4,7 +4,7 @@ import os
 import argparse
 from Bio import SeqIO
 from deepcoil import DeepCoil
-from deepcoil.utils import is_fasta, sharpen_preds
+from deepcoil.utils import is_fasta, sharpen_preds, plot_preds
 
 parser = argparse.ArgumentParser(description='DeepCoil')
 parser.add_argument('-i',
@@ -18,6 +18,9 @@ parser.add_argument('-out_path',
 parser.add_argument('--gpu',
                     help='Use GPU',
                     action='store_true')
+parser.add_argument('--plot',
+                    help='Plot predictions. Images will be stored in the path defined by -out_path',
+                    action='store_true')
 args = parser.parse_args()
 
 # Check if input file exists
@@ -64,4 +67,7 @@ for entry in out_keys:
     for aa, cc_prob, cc_prob_raw, a_prob, d_prob in zip(data[entry], cc_pred, cc_pred_raw, hept_pred[:, 1], hept_pred[:, 2]):
         f.write('{0}\t{1:.3f}\t{2:.3f}\t{3:.3f}\t{4:.3f}\n'.format(aa, float(cc_prob), float(cc_prob_raw), float(a_prob), float(d_prob)))
     f.close()
+if args.plot:
+    for entry in out_keys:
+        plot_preds(preds[entry], out_file=f'{args.out_path}/{entry}.png')
 print("Done!")

deepcoil/utils/utils.py

@@ -128,4 +128,5 @@ def plot_preds(results, beg=0, end=-1, out_file=None):
     plt.tight_layout()
     plt.subplots_adjust(hspace=0)
     if out_file:
-        plt.savefig(out_file, dpi=300)
+        plt.savefig(out_file, dpi=300)
+        plt.close()