kmer annotation supported

hiruna72 · Apr 27, 2024 · a60fa61 · a60fa61
1 parent 9326d61
commit a60fa61
Showing 1 changed file with 82 additions and 5 deletions.
diff --git a/src/plot.py b/src/plot.py
@@ -29,6 +29,7 @@
 # start_kmer is always 0based closed
 # end_kmer is always 0based open
 
+KMER_LENGTH = 9
 FIXED_BASE_WIDTH = 10
 FIXED_INSERTION_WIDTH = 10
 BASE_LIMIT = 1000
@@ -65,6 +66,13 @@ def plot_function(p, read_id, signal_tuple, sig_algn_data, fasta_sequence, base_
     sample_label_colors = []
     location_plot = 0
 
+    kmer_x = []
+    kmer_y = []
+    kmer_label = []
+    kmer_label_position = y_max
+    kmer_start = sig_algn_data["start_kmer"] + draw_data["base_shift"]
+    kmer_length = draw_data["kmer_length"]
+
     move_x = []
     move_y = []
     move_label = []
@@ -106,6 +114,13 @@ def plot_function(p, read_id, signal_tuple, sig_algn_data, fasta_sequence, base_
                 base_label.append(label)
                 base_label_colors.append('red')
 
+                if kmer_start+kmer_length < len(fasta_sequence) and kmer_start > -1:
+                    kmer_text = "{}'{}'{}".format(fasta_sequence[kmer_start:base_index],fasta_sequence[base_index],fasta_sequence[base_index+1:kmer_start+kmer_length])
+                    kmer_label.append(kmer_text)
+                    kmer_x.append(prev_loc)
+                    kmer_y.append(kmer_label_position)
+                kmer_start += 1
+
                 prev_loc += draw_data["fixed_base_width"]
                 prev_x_cord += draw_data["fixed_base_width"]
                 base_index += 1
@@ -157,6 +172,14 @@ def plot_function(p, read_id, signal_tuple, sig_algn_data, fasta_sequence, base_
             label = str(base) + "\n" + str(base_index + 1)
             base_label.append(label)
             base_label_colors.append('black')
+
+            if kmer_start+kmer_length < len(fasta_sequence) and kmer_start > -1:
+                kmer_text = "{}'{}'{}".format(fasta_sequence[kmer_start:base_index],fasta_sequence[base_index],fasta_sequence[base_index+1:kmer_start+kmer_length])
+                kmer_label.append(kmer_text)
+                kmer_x.append(previous_location)
+                kmer_y.append(kmer_label_position)
+            kmer_start += 1
+
             for j in range(0, n_samples):
                 sample_label_colors.append('red')
 
@@ -182,16 +205,24 @@ def plot_function(p, read_id, signal_tuple, sig_algn_data, fasta_sequence, base_
     toggle_bases = Toggle(label="base", button_type="primary", active=True, height=30, width=60)
     toggle_bases.js_link('active', base_annotation_labels, 'visible')
 
+    kmer_annotation = ColumnDataSource(data=dict(kmer_x=kmer_x, kmer_y=kmer_y, kmer_label=kmer_label))
+    kmer_annotation_labels = LabelSet(x='kmer_x', y='kmer_y', text='kmer_label', x_offset=5, y_offset=5, source=kmer_annotation, text_font_size="9pt", angle=45, angle_units='deg')
+
+    btn_kmer_label = "{}-mer".format(kmer_length)
+    toggle_kmers = Toggle(label=btn_kmer_label, button_type="primary", active=True, height=30, width=60)
+    toggle_kmers.js_link('active', kmer_annotation_labels, 'visible')
+
     move_annotation = ColumnDataSource(data=dict(move_x=move_x, move_y=move_y, move_label=move_label, colors=move_label_colors))
     move_annotation_labels = LabelSet(x='move_x', y='move_y', text='move_label', x_offset=5, y_offset=5, source=move_annotation, text_font_size="9pt", text_color='colors', visible=False)
 
-    toggle_moves = Toggle(label="num samples", button_type="primary", active=False, height=30, width=60)
+    toggle_moves = Toggle(label="num samples", button_type="danger", active=False, height=30, width=60)
     toggle_moves.js_link('active', move_annotation_labels, 'visible')
 
     source = ColumnDataSource(data=dict(x=x[:x_coordinate], y=y[:x_coordinate], x_real=x_real[:x_coordinate]))
     p.quad(top=y_max, bottom=y_min, left=base_box_details['left'], right=base_box_details['right'], color=base_box_details['fill_color'], alpha=0.75, visible=draw_data['no_colours'])
     p.add_glyph(line_segment_source, glyph)
     p.add_layout(base_annotation_labels)
+    p.add_layout(kmer_annotation_labels)
     p.add_layout(move_annotation_labels)
 
     p.line('x', 'y', name="sig_plot_line", line_width=2, source=source)
@@ -232,14 +263,21 @@ def plot_function(p, read_id, signal_tuple, sig_algn_data, fasta_sequence, base_
     """)
     p.x_range.js_on_change('start', x_callback_base_annotation)
 
+    x_callback_kmer_annotation = CustomJS(args=dict(kmer_annotation_labels=kmer_annotation_labels, init_font_size=kmer_annotation_labels.text_font_size[:-2], init_xrange=draw_data['xrange']), code="""
+    let xzoom = (init_font_size * init_xrange) / (cb_obj.end - cb_obj.start);
+    kmer_annotation_labels['text_font_size'] = String(xzoom) + 'pt';
+    """)
+    p.x_range.js_on_change('start', x_callback_kmer_annotation)
+
+
     x_callback_move_annotation = CustomJS(args=dict(move_annotation_labels=move_annotation_labels, init_font_size=base_annotation_labels.text_font_size[:-2], init_xrange=draw_data['xrange']), code="""
     let xzoom = (init_font_size * init_xrange) / (cb_obj.end - cb_obj.start);
     move_annotation_labels['text_font_size'] = String(xzoom) + 'pt';
     """)
     p.x_range.js_on_change('start', x_callback_move_annotation)
     message_browser = Div(text="Bokeh version: 3.1.1 (Google Chrome is recommended)", width=400, height=30)
 
-    layout_ = p, row(toggle_bases, toggle_samples, Spacer(width=10), toggle_moves, Spacer(width=40), message_browser)
+    layout_ = p, row(toggle_bases, toggle_kmers, toggle_samples, Spacer(width=10), toggle_moves, Spacer(width=40), message_browser)
     return layout_
 def plot_function_fixed_width(p, read_id, signal_tuple, sig_algn_data, fasta_sequence, base_limit, draw_data):
     x = signal_tuple[0]
@@ -258,6 +296,13 @@ def plot_function_fixed_width(p, read_id, signal_tuple, sig_algn_data, fasta_seq
     base_label_colors = []
     sample_label_colors = []
 
+    kmer_x = []
+    kmer_y = []
+    kmer_label = []
+    kmer_label_position = y_max
+    kmer_start = sig_algn_data["start_kmer"] + draw_data["base_shift"]
+    kmer_length = draw_data["kmer_length"]
+
     move_x = []
     move_y = []
     move_label = []
@@ -299,6 +344,13 @@ def plot_function_fixed_width(p, read_id, signal_tuple, sig_algn_data, fasta_seq
                 base_label.append(label)
                 base_label_colors.append('red')
 
+                if kmer_start+kmer_length < len(fasta_sequence) and kmer_start > -1:
+                    kmer_text = "{}'{}'{}".format(fasta_sequence[kmer_start:base_index],fasta_sequence[base_index],fasta_sequence[base_index+1:kmer_start+kmer_length])
+                    kmer_label.append(kmer_text)
+                    kmer_x.append(prev_loc)
+                    kmer_y.append(kmer_label_position)
+                kmer_start += 1
+
                 prev_loc += draw_data["fixed_base_width"]
                 prev_x_cord += draw_data["fixed_base_width"]
                 base_index += 1
@@ -364,6 +416,13 @@ def plot_function_fixed_width(p, read_id, signal_tuple, sig_algn_data, fasta_seq
             base_label.append(label)
             base_label_colors.append('black')
 
+            if kmer_start+kmer_length < len(fasta_sequence) and kmer_start > -1:
+                kmer_text = "{}'{}'{}".format(fasta_sequence[kmer_start:base_index],fasta_sequence[base_index],fasta_sequence[base_index+1:kmer_start+kmer_length])
+                kmer_label.append(kmer_text)
+                kmer_x.append(previous_location)
+                kmer_y.append(kmer_label_position)
+            kmer_start += 1
+
             move_x.append(previous_location+draw_data["fixed_base_width"]/2)
             move_y.append(label_position/2)
             move_label.append(str(n_samples))
@@ -388,10 +447,17 @@ def plot_function_fixed_width(p, read_id, signal_tuple, sig_algn_data, fasta_seq
     toggle_bases = Toggle(label="base", button_type="primary", active=True, height=30, width=60)
     toggle_bases.js_link('active', base_annotation_labels, 'visible')
 
+    kmer_annotation = ColumnDataSource(data=dict(kmer_x=kmer_x, kmer_y=kmer_y, kmer_label=kmer_label))
+    kmer_annotation_labels = LabelSet(x='kmer_x', y='kmer_y', text='kmer_label', x_offset=5, y_offset=5, source=kmer_annotation, text_font_size="9pt", angle=45, angle_units='deg')
+
+    btn_kmer_label = "{}-mer".format(kmer_length)
+    toggle_kmers = Toggle(label=btn_kmer_label, button_type="primary", active=True, height=30, width=60)
+    toggle_kmers.js_link('active', kmer_annotation_labels, 'visible')
+
     move_annotation = ColumnDataSource(data=dict(move_x=move_x, move_y=move_y, move_label=move_label, colors=move_label_colors))
     move_annotation_labels = LabelSet(x='move_x', y='move_y', text='move_label', x_offset=5, y_offset=5, source=move_annotation, text_font_size="9pt", text_color='colors', visible=False)
 
-    toggle_moves = Toggle(label="num samples", button_type="primary", active=False, height=30, width=60)
+    toggle_moves = Toggle(label="num samples", button_type="danger", active=False, height=30, width=60)
     toggle_moves.js_link('active', move_annotation_labels, 'visible')
 
     fixed_width_x = fixed_width_x[1:]
@@ -400,6 +466,7 @@ def plot_function_fixed_width(p, read_id, signal_tuple, sig_algn_data, fasta_seq
     p.quad(top=y_max, bottom=y_min, left=base_box_details['left'], right=base_box_details['right'], color=base_box_details['fill_color'], alpha=0.75, visible=draw_data['no_colours'])
     p.add_glyph(line_segment_source, glyph)
     p.add_layout(base_annotation_labels)
+    p.add_layout(kmer_annotation_labels)
     p.add_layout(move_annotation_labels)
 
     p.line('x', 'y', name="sig_plot_line", line_width=2, source=source)
@@ -453,14 +520,20 @@ def plot_function_fixed_width(p, read_id, signal_tuple, sig_algn_data, fasta_seq
     """)
     p.x_range.js_on_change('start', x_callback_base_annotation)
 
+    x_callback_kmer_annotation = CustomJS(args=dict(kmer_annotation_labels=kmer_annotation_labels, init_font_size=kmer_annotation_labels.text_font_size[:-2], init_xrange=draw_data['xrange']), code="""
+    let xzoom = (init_font_size * init_xrange) / (cb_obj.end - cb_obj.start);
+    kmer_annotation_labels['text_font_size'] = String(xzoom) + 'pt';
+    """)
+    p.x_range.js_on_change('start', x_callback_kmer_annotation)
+
     x_callback_move_annotation = CustomJS(args=dict(move_annotation_labels=move_annotation_labels, init_font_size=base_annotation_labels.text_font_size[:-2], init_xrange=draw_data['xrange']), code="""
     let xzoom = (init_font_size * init_xrange) / (cb_obj.end - cb_obj.start);
     move_annotation_labels['text_font_size'] = String(xzoom) + 'pt';
     """)
     p.x_range.js_on_change('start', x_callback_move_annotation)
     message_browser = Div(text="Bokeh version: 3.1.1 (Google Chrome is recommended)", width=400, height=30)
 
-    layout_ = p, row(toggle_bases, toggle_samples, Spacer(width=10), toggle_moves, Spacer(width=40), message_browser)
+    layout_ = p, row(toggle_bases, toggle_kmers, toggle_samples, Spacer(width=10), toggle_moves, Spacer(width=40), message_browser)
     return layout_
 
 def run(args):
@@ -547,6 +620,7 @@ def run(args):
     draw_data["fixed_width"] = args.fixed_width
     draw_data["sig_dir"] = "->"
     draw_data["bed_labels"] = args.print_bed_labels
+    draw_data["kmer_length"] = args.kmer_length
     if args.plot_reverse:
         draw_data["sig_dir"] = "<-"
 
@@ -560,7 +634,9 @@ def run(args):
     kmer_correction = 0
     if args.profile != "":
         kmer_correction = -1*(plot_utils.search_for_profile_base_shift(args.profile)[0] + plot_utils.search_for_profile_base_shift(args.profile)[1])
-
+    if args.kmer_length < abs(draw_data["base_shift"]):
+        print("Info: increased the kmer length to {} match with the base shift".format(abs(draw_data["base_shift"])+1))
+        draw_data["kmer_length"] = abs(draw_data["base_shift"]) + 1
 
     if use_paf == 1 and plot_sig_ref_flag == 0:
         print("Info: Signal to read method using PAF ...")
@@ -1220,6 +1296,7 @@ def argparser():
     parser.add_argument('--base_limit', required=False, type=int, default=BASE_LIMIT, help="maximum number of bases to plot")
     parser.add_argument('--region', required=False, type=str, default="", help="[start-end] 1-based closed interval region to plot. For SAM/BAM eg: chr1:6811428-6811467 or chr1:6,811,428-6,811,467. For PAF eg:100-200.")
     parser.add_argument('--tag_name', required=False, type=str, default="", help="a tag name to easily identify the plot")
+    parser.add_argument('-k', '--kmer_length', required=False, type=int, default=KMER_LENGTH, help="kmer length")
     parser.add_argument('--plot_reverse', required=False, action='store_true', help="plot only the reverse mapped reads.")
     parser.add_argument('--rna', required=False, action='store_true', help="specify for RNA reads")
     parser.add_argument('--sig_ref', required=False, action='store_true', help="plot signal to reference mapping")