Added option to generate greyscale plots

Author: ODiogoSilva

structure_threader/plotter/structplot.py

@@ -797,9 +797,13 @@ def plotk(self, kvals, output_dir):
         with open(filepath, "w") as fh:
             fh.write(ploty_html(pdiv))
 
-    def plotk_static(self, kval, output_dir):
+    def plotk_static(self, kval, output_dir, bw=False):
         """
-
+        Generates a structure plot in svg format.
+        :param kval: (int) Must match the K value from self.kvals
+        :param output_dir: (string) Path of the plot file
+        :param bw: If True, plots will be generated with patterns instead of
+        colors to distinguish k groups.
         """
 
         qvalues = self.kvals[kval].qvals
@@ -817,19 +821,36 @@ def plotk_static(self, kval, output_dir):
         axe = fig.add_subplot(111, xlim=(-.5, numinds - .5), ylim=(0, 1))
 
         for i in range(qvalues.shape[1]):
-            # Get bar color. If K exceeds the 12 colors, generate random color
-            try:
-                clr = clist[i]
-            except IndexError:
-                clr = np.random.rand(3, 1)
+
+            # Determine color/pattern arguments
+            kwargs = {}
+
+            # Use colors todistinguish k groups
+            if not bw:
+                # Get bar color. If K exceeds the 12 colors, generate random
+                # color
+                try:
+                    clr = clist[i]
+                except IndexError:
+                    clr = np.random.rand(3, 1)
+
+                kwargs["facecolor"] = clr
+                kwargs["edgecolor"] = "grey"
+
+            else:
+                grey_rgb = [(float((i + 1)) / (float(qvalues.shape[1]) + 1))
+                            for _ in range(3)]
+
+                kwargs["facecolor"] = grey_rgb
+                kwargs["edgecolor"] = "white"
 
             if i == 0:
-                axe.bar(range(numinds), qvalues[:, i], facecolor=clr,
-                        edgecolor="grey", width=1)
+                axe.bar(range(numinds), qvalues[:, i],
+                        width=1, label="K {}".format(i + 1), **kwargs)
                 former_q = qvalues[:, i]
             else:
                 axe.bar(range(numinds), qvalues[:, i], bottom=former_q,
-                        facecolor=clr, edgecolor="grey", width=1)
+                        width=1, label="K {}".format(i + 1), **kwargs)
                 former_q = former_q + qvalues[:, i]
 
         # Annotate population info
@@ -860,15 +881,23 @@ def plotk_static(self, kval, output_dir):
         plt.yticks([])
         plt.xticks([])
 
+        # Add k legend
+        legend = plt.legend(bbox_to_anchor=(1.2, .5), loc=7, borderaxespad=0.)
+        legend.get_frame().set_facecolor("white")
+
         kfile = self.kvals[kval].file_path
         filename = splitext(basename(kfile))[0]
         filepath = join(output_dir, filename)
 
         plt.savefig("{}.svg".format(filepath), bbox_inches="tight")
 
+        # Clear plot object
+        plt.clf()
+        plt.close()
+
 
 def main(result_files, fmt, outdir, bestk=None, popfile=None, indfile=None,
-         filter_k=None):
+         filter_k=None, bw=False):
     """
     Wrapper function that generates one plot for each K value.
     :return:
@@ -890,7 +919,7 @@ def main(result_files, fmt, outdir, bestk=None, popfile=None, indfile=None,
 
         if k >= 1 and k in filter_k:
             klist.plotk([k], outdir)
-            klist.plotk_static(k, outdir)
+            klist.plotk_static(k, outdir, bw=bw)
 
     # If a sequence of multiple bestk is provided, plot all files in a single
     # plot

structure_threader/structure_threader.py

@@ -254,7 +254,7 @@ def create_plts(resultsdir, wrapped_prog, Ks, bestk, arg):
                      for i in plt_list]
 
     sp.main(plt_files, wrapped_prog, outdir, bestk=bestk, popfile=arg.popfile,
-            indfile=arg.indfile)
+            indfile=arg.indfile, bw=arg.blacknwhite)
 
 
 def maverick_merger(outdir, k_list, no_tests):
@@ -424,11 +424,16 @@ def argument_parser(args):
                            help="Override 'K' values from the given list"
                            " to be ploteted in the combined figure.",
                            metavar="'2 4 5'", default=None)
+    plot_opts.add_argument("-bw", dest="blacknwhite",
+                            action="store_const", const=True,
+                            help="Set this flag to draw greyscale plots"
+                                 " instead of colored ones.")
 
     # ####################### PLOT ARGUMENTS ##################################
     # Group definitions
 
     main_opts = plot_parser.add_argument_group("Main plotting options")
+    extra_opts = plot_parser.add_argument_group("Extra plotting options")
     sort_opts = plot_parser.add_argument_group("Plot sorting options")
 
     # Group options
@@ -454,6 +459,11 @@ def argument_parser(args):
                                 " the current working directory"
                                 " will be used.")
 
+    extra_opts.add_argument("-bw", dest="blacknwhite",
+                            action="store_const", const=True,
+                            help="Set this flag to draw greyscale plots"
+                                 " instead of colored ones.")
+
     sort_opts_ex.add_argument("--pop", dest="popfile", type=str,
                               required=False,
                               help="File with population information.",
@@ -597,7 +607,7 @@ def main():
         bestk = [int(x) for x in arg.bestk]
 
         sp.main(infiles, arg.format, arg.outpath, bestk, popfile=arg.popfile,
-                indfile=arg.indfile, filter_k=bestk)
+                indfile=arg.indfile, filter_k=bestk, bw=arg.blacknwhite)
 
 
 if __name__ == "__main__":