Added option to show individual sample names, even when the population is provided

Author: ODiogoSilva

structure_threader/plotter/structplot.py

@@ -601,6 +601,9 @@ def _parse_indfile(self, indfile):
                     self._sort_qvals_pop(indarray=indarray)
                     indarray = indarray[indarray[:, 1].argsort()]
 
+                self.indv = indarray[:, 0]
+                self.number_indv = len(self.indv)
+
                 # Set self.pops attribute
                 pop_counts = Counter(indarray[:, 1])
                 pop_sums = np.cumsum([np.count_nonzero(indarray == x)
@@ -797,13 +800,15 @@ def plotk(self, kvals, output_dir):
         with open(filepath, "w") as fh:
             fh.write(ploty_html(pdiv))
 
-    def plotk_static(self, kval, output_dir, bw=False):
+    def plotk_static(self, kval, output_dir, bw=False, use_ind=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
+        :param bw: (bool) If True, plots will be generated with patterns instead of
         colors to distinguish k groups.
+        :param show_ind: (bool) If True, and if individual labels were provided
+        with the --ind option, use those labels instead of population labels
         """
 
         qvalues = self.kvals[kval].qvals
@@ -854,8 +859,7 @@ def plotk_static(self, kval, output_dir, bw=False):
                 former_q = former_q + qvalues[:, i]
 
         # Annotate population info
-        if self.pops:
-
+        if self.pops and not use_ind:
             pop_lines = list(OrderedDict.fromkeys(
                 [x for y in self.pops_xrange for x in y]))[1:-1]
 
@@ -872,7 +876,7 @@ def plotk_static(self, kval, output_dir, bw=False):
 
             for pos in range(self.number_indv):
                 axe.text(pos, -0.05, self.indv[pos],
-                         rotation=45, va="top", ha="right", fontsize=10)
+                         rotation=45, va="top", ha="right", fontsize=8)
 
         for axis in ["top", "bottom", "left", "right"]:
             axe.spines[axis].set_linewidth(2)
@@ -897,7 +901,7 @@ def plotk_static(self, kval, output_dir, bw=False):
 
 
 def main(result_files, fmt, outdir, bestk=None, popfile=None, indfile=None,
-         filter_k=None, bw=False):
+         filter_k=None, bw=False, use_ind=False):
     """
     Wrapper function that generates one plot for each K value.
     :return:
@@ -919,7 +923,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, bw=bw)
+            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

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, bw=arg.blacknwhite)
+            indfile=arg.indfile, bw=arg.blacknwhite, use_ind=arg.use_ind)
 
 
 def maverick_merger(outdir, k_list, no_tests):
@@ -428,6 +428,11 @@ def argument_parser(args):
                             action="store_const", const=True,
                             help="Set this flag to draw greyscale plots"
                                  " instead of colored ones.")
+    plot_opts.add_argument("--use-ind-labels", dest="use_ind",
+                           action="store_const", const=True,
+                           help="Use the individual labels in the "
+                                "structure plot instead of population"
+                                " labels")
 
     # ####################### PLOT ARGUMENTS ##################################
     # Group definitions
@@ -463,6 +468,11 @@ def argument_parser(args):
                             action="store_const", const=True,
                             help="Set this flag to draw greyscale plots"
                                  " instead of colored ones.")
+    extra_opts.add_argument("--use-ind-labels", dest="use_ind",
+                            action="store_const", const=True,
+                            help="Use the individual labels in the "
+                                 "structure plot instead of population"
+                                 " labels")
 
     sort_opts_ex.add_argument("--pop", dest="popfile", type=str,
                               required=False,
@@ -607,7 +617,8 @@ 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, bw=arg.blacknwhite)
+                indfile=arg.indfile, filter_k=bestk, bw=arg.blacknwhite,
+                use_ind=arg.use_ind)
 
 
 if __name__ == "__main__":