From 3070e110b73ede7cf6858029cc1b1c496370bb8e Mon Sep 17 00:00:00 2001 From: jonas-fuchs Date: Fri, 17 May 2024 14:18:01 +0200 Subject: [PATCH] reformating and comments --- virheat/command.py | 21 ++++++++++----------- virheat/scripts/plotting.py | 11 +++++------ 2 files changed, 15 insertions(+), 17 deletions(-) diff --git a/virheat/command.py b/virheat/command.py index dd1fcf4..21257a5 100644 --- a/virheat/command.py +++ b/virheat/command.py @@ -224,16 +224,7 @@ 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 @@ -241,7 +232,15 @@ def main(sysargs=sys.argv[1:]): 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]) diff --git a/virheat/scripts/plotting.py b/virheat/scripts/plotting.py index 9cf4e65..cacc036 100644 --- a/virheat/scripts/plotting.py +++ b/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,7 +72,7 @@ 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)] @@ -81,14 +80,13 @@ def create_scores_vis(ax, genome_y_location, n_mutations, n_tracks, unique_mutat 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