reformating and comments

virheat/command.py

@@ -224,24 +224,23 @@ def main(sysargs=sys.argv[1:]):
                          unique_mutations, reference_name, n_scores)
     plotting.create_mutation_legend(mutation_set, min_y_location, n_samples, n_scores)
     plotting.create_colorbar(args.threshold, cmap_cells, min_y_location, n_samples, ax)
-    # plot scores as track below the genome track
-    if args.scores:
-        score_count = 1
-        for score_params in args.scores:
-            scores_file, pos_col, score_col, score_name = score_params
-            unique_scores = data_prep.extract_scores(unique_mutations, scores_file, pos_col, score_col)
-            track_created = plotting.create_scores_vis(ax, genome_y_location, n_mutations, n_tracks, unique_scores, start, stop, score_count, score_name)
-            if track_created:
-                score_count += 1
-
+    # plot gene track
     if args.gff3_path is not None:
         if genes_with_mutations:
             # distinct colors for the genes
             cmap_genes = plt.get_cmap('tab20', len(genes_with_mutations))
             colors_genes = [cmap_genes(i) for i in range(len(genes_with_mutations))]
             # plot gene track
             plotting.create_gene_vis(ax, genes_with_mutations, n_mutations, y_size, n_tracks, start, stop, min_y_location, genome_y_location, colors_genes, n_scores)
-
+    # plot scores as track below the gene/genome track
+    if args.scores:
+        score_count = 1
+        for score_params in args.scores:
+            scores_file, pos_col, score_col, score_name = score_params
+            unique_scores = data_prep.extract_scores(unique_mutations, scores_file, pos_col, score_col)
+            track_created = plotting.create_scores_vis(ax, genome_y_location, n_mutations, n_tracks, unique_scores, start, stop, score_count, score_name)
+            if track_created:  # only creates a track if it finds mutations to annotate
+                score_count += 1
     # create output folder
     if not os.path.exists(args.input[1]):
         os.makedirs(args.input[1])

virheat/scripts/plotting.py

@@ -9,7 +9,7 @@
 # LIBS
 import matplotlib.pyplot as plt
 import matplotlib.patches as patches
-import matplotlib.colors as mcolors
+
 
 def create_heatmap(ax, frequency_array, cmap):
     """
@@ -49,7 +49,6 @@ def create_genome_vis(ax, genome_y_location, n_mutations, unique_mutations, star
                             ec="none", fc="lightgrey"
                         )
     )
-
     # create mutation lines on the genome rectangle and the mapping to the respective cells
     x_start = 0
     length = stop - start
@@ -73,22 +72,21 @@ def create_scores_vis(ax, genome_y_location, n_mutations, n_tracks, unique_mutat
     create the scores rectangles, mappings to the reference
     """
     score_set = []
-    y_zero = -genome_y_location - (n_tracks + score_count + 1) * (genome_y_location / 2)
+    y_zero = -genome_y_location - (n_tracks + score_count + 1) * (genome_y_location / 2)  # zero line
     length = stop - start
 
     # create list of tuples [(nt pos, score)]
     for mutation in unique_mutations:
         mutation_attributes = mutation.split("_")
         if not np.isnan(float(mutation_attributes[5])):
             score_set.append((int(mutation_attributes[0]), float(mutation_attributes[5])))
-
+    # check if there is something to plot
     if score_set:
         # create zero line and score name on the left
         plt.axhline(y=y_zero, color='black', linestyle='-', linewidth=0.5)
         ax.text(-0.5, y_zero, score_name, ha='right', va='center')
-        # define normalization multiplier
+        # define normalization multiplier for the height of the score v lines
         multiplier = max([abs(score[1]) for score in score_set]) / abs((y_zero + genome_y_location / 4) - y_zero)
-
         # create score lines on the score rectangle and the mapping to the respective mutation lines
         for score in score_set:
             # define x value
@@ -100,6 +98,7 @@ def create_scores_vis(ax, genome_y_location, n_mutations, n_tracks, unique_mutat
             else:
                 plt.vlines(x=mutation_x_location, ymin=y_zero, ymax=y_zero + score[1] / multiplier, color="blue", linestyle='-')
         return True
+    # if not the track is not created
     else:
         print("\033[31m\033[1mERROR:\033[0m Seems like there are no scores in the score set '{}' corresponding to the plotted mutation positions.".format(score_name))
         return False