Move prep_alleles_table out of plot function in Core

edilytics · Aug 9, 2024 · 0402f25 · 0402f25
1 parent fa05bd5
commit 0402f25
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/CRISPResso2/CRISPRessoCORE.py b/CRISPResso2/CRISPRessoCORE.py
@@ -4462,6 +4462,7 @@ def count_alternate_alleles(sub_base_vectors, ref_name, ref_sequence, ref_total_
                     if not args.expand_allele_plots_by_quantification:
                         df_to_plot = df_alleles_around_cut.groupby(['Aligned_Sequence', 'Reference_Sequence']).sum().reset_index().set_index('Aligned_Sequence')
                         df_to_plot.sort_values(by=['#Reads', 'Aligned_Sequence', 'Reference_Sequence'], inplace=True, ascending=[False, True, True])
+                    df_alleles_prep = CRISPRessoPlot.prep_alleles_table(df_to_plot, ref_seq_around_cut, args.max_rows_alleles_around_cut_to_plot, args.min_frequency_alleles_around_cut_to_plot)
 
                     new_sgRNA_intervals = []
                     #adjust coordinates of sgRNAs
@@ -4470,11 +4471,9 @@ def count_alternate_alleles(sub_base_vectors, ref_name, ref_sequence, ref_total_
                         new_sgRNA_intervals += [(int_start - new_sel_cols_start - 1, int_end - new_sel_cols_start - 1)]
                     plot_9_input = {
                         'reference_seq': ref_seq_around_cut,
-                        'df_alleles': df_to_plot,
+                        'df_alleles_prep': df_alleles_prep,
                         'fig_filename_root': fig_filename_root,
                         'custom_colors': custom_config["colors"],
-                        'MIN_FREQUENCY': args.min_frequency_alleles_around_cut_to_plot,
-                        'MAX_N_ROWS': args.max_rows_alleles_around_cut_to_plot,
                         'SAVE_ALSO_PNG': save_png,
                         'plot_cut_point': plot_cut_point,
                         'sgRNA_intervals': new_sgRNA_intervals,

diff --git a/CRISPResso2/CRISPRessoPlot.py b/CRISPResso2/CRISPRessoPlot.py
@@ -3082,7 +3082,20 @@ def plot_alleles_heatmap_hist(reference_seq,fig_filename_root,X,annot,y_labels,i
         plt.savefig(fig_filename_root+'.png', bbox_inches='tight', bbox_extra_artists=(lgd,), pad_inches=0.1)
     plt.close()
 
-def plot_alleles_table(reference_seq,df_alleles,fig_filename_root,custom_colors,MIN_FREQUENCY=0.5,MAX_N_ROWS=100,SAVE_ALSO_PNG=False,plot_cut_point=True,cut_point_ind=None,sgRNA_intervals=None,sgRNA_names=None,sgRNA_mismatches=None,annotate_wildtype_allele='****',**kwargs):
+def plot_alleles_table(
+    reference_seq,
+    df_alleles_prep,
+    fig_filename_root,
+    custom_colors,
+    SAVE_ALSO_PNG=False,
+    plot_cut_point=True,
+    cut_point_ind=None,
+    sgRNA_intervals=None,
+    sgRNA_names=None,
+    sgRNA_mismatches=None,
+    annotate_wildtype_allele='****',
+    **kwargs,
+):
     """
     plots an allele table for a dataframe with allele frequencies
     input:
@@ -3100,7 +3113,7 @@ def plot_alleles_table(reference_seq,df_alleles,fig_filename_root,custom_colors,
     custom_colors: dict of colors to plot (e.g. colors['A'] = (1,0,0,0.4) # red,blue,green,alpha )
     annotate_wildtype_allele: string to add to the end of the wildtype allele (e.g. ** or '')
     """
-    X, annot, y_labels, insertion_dict, per_element_annot_kws, is_reference = prep_alleles_table(df_alleles, reference_seq, MAX_N_ROWS, MIN_FREQUENCY)
+    X, annot, y_labels, insertion_dict, per_element_annot_kws, is_reference = df_alleles_prep
     if annotate_wildtype_allele != '':
         for ix, is_ref in enumerate(is_reference):
             if is_ref: