Fixes #56. K=1 are now plotted.

Author: ODiogoSilva

structure_threader/plotter/structplot.py

@@ -642,7 +642,7 @@ def plotk(self, kvals, output_dir):
         """
 
         # Get number of plots (confirm the kvals are valid before)
-        nplots = len([x for x in kvals if int(x) in self.kvals if x >= 1])
+        nplots = len([x for x in kvals if int(x) in self.kvals])
 
         # If no valid plots, issue an error
         if not nplots:
@@ -681,8 +681,15 @@ def plotk(self, kvals, output_dir):
             # Fetch PlotK object that will be plotted
             kobj = self.kvals[k]
 
+            # Transforms the qvals matrix when K = 1. If K > 2, use the
+            # original matrix
+            if len(kobj.qvals.shape) == 1:
+                qvals = [kobj.qvals.T]
+            else:
+                qvals = kobj.qvals.T
+
             # Iterate over each meanQ column (corresponding to each cluster)
-            for p, i in enumerate(kobj.qvals.T):
+            for p, i in enumerate(qvals):
 
                 # Create Bar trace for each cluster
                 current_bar = go.Bar(
@@ -819,8 +826,6 @@ def plotk_static(self, kval, output_dir, bw=False, use_ind=False):
         with the --ind option, use those labels instead of population labels
         """
 
-        qvalues = self.kvals[kval].qvals
-
         plt.style.use("ggplot")
 
         numinds = self.number_indv
@@ -833,7 +838,16 @@ def plotk_static(self, kval, output_dir, bw=False, use_ind=False):
         fig = plt.figure()
         axe = fig.add_subplot(111, xlim=(-.5, numinds - .5), ylim=(0, 1))
 
-        for i in range(qvalues.shape[1]):
+        # Transforms the qvals matrix when K = 1. If K > 2, use the
+        # original matrix
+        if len(self.kvals[kval].qvals.shape) == 1:
+            # This list comprehension ensures that the shape of the array
+            # is (i, 1), where i is the number of samples
+            qvalues = np.array([[x] for x in self.kvals[kval].qvals])
+        else:
+            qvalues = self.kvals[kval].qvals
+
+        for i in range(kval):
 
             # Determine color/pattern arguments
             kwargs = {}
@@ -894,10 +908,6 @@ def plotk_static(self, kval, output_dir, bw=False, use_ind=False):
         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)
@@ -930,12 +940,11 @@ def main(result_files, fmt, outdir, bestk=None, popfile=None, indfile=None,
     # Plot all K files individually
     for k, kobj in klist:
 
-        if k >= 1 and k in filter_k:
+        if k in filter_k:
             klist.plotk([k], outdir)
             klist.plotk_static(k, outdir, bw=bw, use_ind=use_ind)
 
     # If a sequence of multiple bestk is provided, plot all files in a single
     # plot
     if bestk:
-        bestk = [x for x in bestk if x >= 1]
         klist.plotk(bestk, outdir)

structure_threader/structure_threader.py

@@ -228,8 +228,6 @@ def create_plts(resultsdir, wrapped_prog, Ks, bestk, arg):
     :param resultsdir: path to results directory
     """
 
-    plt_list = [x for x in Ks if x != 1]  # Don't plot K=1
-
     outdir = os.path.join(resultsdir, "plots")
     if not os.path.exists(outdir):
         os.mkdir(outdir)
@@ -243,15 +241,15 @@ def create_plts(resultsdir, wrapped_prog, Ks, bestk, arg):
             file_to_plot = str(randrange(1, arg.replicates + 1))
         plt_files = [os.path.join(resultsdir, "str_K") + str(i) + "_rep" +
                      file_to_plot + "_f"
-                     for i in plt_list]
+                     for i in Ks]
     elif wrapped_prog == "maverick":
         plt_files = [os.path.join(os.path.join(resultsdir, "mav_K" + str(i)),
                                   "outputQmatrix_ind_K" + str(i) + ".csv")
-                     for i in plt_list]
+                     for i in Ks]
 
     else:
         plt_files = [os.path.join(resultsdir, "fS_run_K.") + str(i) + ".meanQ"
-                     for i in plt_list]
+                     for i in Ks]
 
     sp.main(plt_files, wrapped_prog, outdir, bestk=bestk, popfile=arg.popfile,
             indfile=arg.indfile, bw=arg.blacknwhite, use_ind=arg.use_ind)