varvamp 1.1.3__py3-none-any.whl → 1.2.1__py3-none-any.whl

varvamp/scripts/reporting.py

@@ -16,13 +16,14 @@ from matplotlib.backends.backend_pdf import PdfPages
 # varVAMP
 from varvamp.scripts import primers
 from varvamp.scripts import config
+from varvamp.scripts import logging
 
 
-def write_fasta(path, seq_id, seq):
+def write_fasta(path, file_name, seq_id, seq):
     """
     write fasta files
     """
-    name = f"{seq_id}.fasta"
+    name = f"{file_name}.fasta"
     out = os.path.join(path, name)
     with open(out, 'w') as o:
         print(f">{seq_id}\n{seq}", file=o)
@@ -39,7 +40,7 @@ def write_alignment(path, alignment):
             print(f">{seq[0]}\n{seq[1]}", file=o)
 
 
-def write_regions_to_bed(primer_regions, path, mode=None):
+def write_regions_to_bed(primer_regions, scheme_name, path, mode=None):
     """
     write primer regions as bed file
     """
@@ -48,39 +49,43 @@ def write_regions_to_bed(primer_regions, path, mode=None):
         outfile = f"{path}probe_regions.bed"
     else:
         outfile = f"{path}primer_regions.bed"
-    counter = 0
 
     with open(outfile, 'w') as o:
-        for region in primer_regions:
+        for counter, region in enumerate(primer_regions):
             print(
-                "ambiguous_consensus",
+                f"{scheme_name}_consensus",
                 region[0],
                 region[1],
                 "REGION_"+str(counter),
                 sep="\t",
                 file=o
             )
-            counter += 1
 
 
-def write_primers_to_bed(outfile, primer_name, primer_properties, direction):
+def write_primers_to_bed(outfile, scheme_name, primer_name, primer_properties, numeric_value, direction, sequence=None):
     """
     write primers as bed file
     """
     with open(outfile, 'a') as o:
-        print(
-            "ambiguous_consensus",
+        # write header for primer bed
+        if os.path.getsize(outfile) == 0 and sequence is not None:
+            print("#chrom\tchromStart\tchromEnd\tprimer-name\tpool\tstrand\tprimer-sequence", file=o)
+        data = [f"{scheme_name}_consensus",
             primer_properties[1],  # start
             primer_properties[2],  # stop
             primer_name,
-            round(primer_properties[3], 1),  # penalty
-            direction,
+            numeric_value,  # can be pool or score
+            direction]
+        if sequence is not None:
+            data.append(sequence)
+        print(
+            *data,
             sep="\t",
             file=o
         )
 
 
-def write_all_primers(path, all_primers):
+def write_all_primers(path, scheme_name, all_primers):
     """
     write all primers that varVAMP designed as bed file
     """
@@ -88,7 +93,7 @@ def write_all_primers(path, all_primers):
 
     for direction in all_primers:
         for primer in all_primers[direction]:
-            write_primers_to_bed(outfile, primer, all_primers[direction][primer], direction)
+            write_primers_to_bed(outfile, scheme_name, primer, all_primers[direction][primer], round(all_primers[direction][primer][3], 2), direction)
 
 
 def get_permutations(seq):
@@ -121,7 +126,7 @@ def calc_mean_stats(permutations):
     return round(gc/len(permutations), 1), round(temp/len(permutations), 1)
 
 
-def write_qpcr_to_files(path, final_schemes, ambiguous_consensus):
+def write_qpcr_to_files(path, final_schemes, ambiguous_consensus, scheme_name, log_file):
     """
     write all relevant bed files and tsv file for the qPCR design
     """
@@ -134,32 +139,43 @@ def write_qpcr_to_files(path, final_schemes, ambiguous_consensus):
 
     with open(tsv_file, "w") as tsv, open(tsv_file_2, "w") as tsv2, open(amplicon_bed_file, "w") as bed, open(primer_fasta_file, "w") as fasta:
         print(
-            "qpcr_scheme\toligo_type\tstart\tstop\tseq\tsize\tgc_best\ttemp_best\tmean_gc\tmean_temp\tpenalty",
+            "qpcr_scheme\toligo_type\tstart\tstop\tseq\tsize\tgc_best\ttemp_best\tmean_gc\tmean_temp\tpenalty\toff_target_amplicons",
             file=tsv2
         )
         print(
-            "qpcr_scheme\tpenalty\tdeltaG\tlength\tstart\tstop\tseq",
+            "qpcr_scheme\toff_target_amplicons\tpenalty\tdeltaG\tlength\tstart\tstop\tseq",
             file=tsv
         )
-        for scheme in final_schemes:
+        for n, amp in enumerate(final_schemes):
+            amp_name = f"{scheme_name}_{n}"
             # write bed amplicon file
             print(
-                "ambiguous_consensus",
-                final_schemes[scheme]["LEFT"][1],
-                final_schemes[scheme]["RIGHT"][2],
-                scheme,
-                round(final_schemes[scheme]["penalty"], 1),
+                f"{scheme_name}_consensus",
+                amp["LEFT"][1],
+                amp["RIGHT"][2],
+                amp_name,
+                round(amp["penalty"], 1),
+                ".",
                 sep="\t",
                 file=bed
             )
             # write tsv
-            amplicon_start = final_schemes[scheme]["LEFT"][1]
-            amplicon_stop = final_schemes[scheme]["RIGHT"][2]
+            amplicon_start = amp["LEFT"][1]
+            amplicon_stop = amp["RIGHT"][2]
+            if "off_targets" in amp:
+                if amp["off_targets"]:
+                    amplicon_has_off_target = "Yes"
+                    write_BLAST_warning(amp_name, log_file)
+                else:
+                    amplicon_has_off_target = "No"
+            else:
+                amplicon_has_off_target = "n.d."
             amplicon_seq = ambiguous_consensus[amplicon_start:amplicon_stop]
             print(
-                scheme,
-                round(final_schemes[scheme]["penalty"], 1),
-                final_schemes[scheme]["deltaG"],
+                amp_name,
+                amplicon_has_off_target,
+                round(amp["penalty"], 1),
+                amp["deltaG"],
                 len(amplicon_seq),
                 amplicon_start + 1,
                 amplicon_stop,
@@ -168,11 +184,9 @@ def write_qpcr_to_files(path, final_schemes, ambiguous_consensus):
                 file=tsv
             )
             # write tsv2
-            for oligo_type in final_schemes[scheme]:
-                if oligo_type == "penalty" or oligo_type == "deltaG":
-                    continue
-                seq = ambiguous_consensus[final_schemes[scheme][oligo_type][1]:final_schemes[scheme][oligo_type][2]]
-                if oligo_type == "RIGHT" or all([oligo_type == "PROBE", final_schemes[scheme]["PROBE"][5] == "-"]):
+            for oligo_type in ["LEFT", "PROBE", "RIGHT"]:
+                seq = ambiguous_consensus[amp[oligo_type][1]:amp[oligo_type][2]]
+                if oligo_type == "RIGHT" or (oligo_type == "PROBE" and amp["PROBE"][5] == "-"):
                     seq = primers.rev_complement(seq)
                     direction = "-"
                 else:
@@ -180,34 +194,39 @@ def write_qpcr_to_files(path, final_schemes, ambiguous_consensus):
 
                 permutations = get_permutations(seq)
                 gc, temp = calc_mean_stats(permutations)
+                primer_name = f"{amp_name}_{oligo_type}"
 
                 print(
-                    scheme,
+                    amp_name,
                     oligo_type,
-                    final_schemes[scheme][oligo_type][1] + 1,
-                    final_schemes[scheme][oligo_type][2],
+                    amp[oligo_type][1] + 1,
+                    amp[oligo_type][2],
                     seq.upper(),
                     len(seq),
-                    round(primers.calc_gc(final_schemes[scheme][oligo_type][0]), 1),
-                    round(primers.calc_temp(final_schemes[scheme][oligo_type][0]), 1),
+                    round(primers.calc_gc(amp[oligo_type][0]), 1),
+                    round(primers.calc_temp(amp[oligo_type][0]), 1),
                     gc,
                     temp,
-                    round(final_schemes[scheme][oligo_type][3], 1),
+                    round(amp[oligo_type][3], 1),
+                    amplicon_has_off_target,
                     sep="\t",
                     file=tsv2
                 )
                 # write primer bed file
                 write_primers_to_bed(
                     primer_bed_file,
-                    f"{scheme}_{oligo_type}",
-                    final_schemes[scheme][oligo_type],
-                    direction
+                    scheme_name,
+                    primer_name,
+                    amp[oligo_type],
+                    round(amp[oligo_type][3], 2),
+                    direction,
+                    seq.upper()
                 )
                 # write fasta
-                print(f">{scheme}_{oligo_type}\n{seq.upper()}", file=fasta)
+                print(f">{primer_name}\n{seq.upper()}", file=fasta)
 
 
-def write_scheme_to_files(path, amplicon_scheme, ambiguous_consensus, mode):
+def write_scheme_to_files(path, amplicon_scheme, ambiguous_consensus, scheme_name, mode, log_file):
     """
     write all relevant bed files and a tsv file with all primer stats
     """
@@ -216,60 +235,65 @@ def write_scheme_to_files(path, amplicon_scheme, ambiguous_consensus, mode):
     amplicon_bed_file = os.path.join(path, "amplicons.bed")
     tabular_file = os.path.join(path, "primer_to_amplicon_assignment.tabular")
 
-    counter = 0
-
     # open files to write
     with open(tsv_file, "w") as tsv, open(amplicon_bed_file, "w") as bed, open(tabular_file, "w") as tabular:
         # write header for primer tsv
         print(
-            "amlicon_name\tamplicon_length\tprimer_name\talternate_primer_name\tpool\tstart\tstop\tseq\tsize\tgc_best\ttemp_best\tmean_gc\tmean_temp\tpenalty",
+            "amlicon_name\tamplicon_length\tprimer_name\tprimer_name_all_primers\tpool\tstart\tstop\tseq\tsize\tgc_best\ttemp_best\tmean_gc\tmean_temp\tpenalty\toff_target_amplicons",
             file=tsv
         )
-
-        for pool in amplicon_scheme:
+        amplicon_bed_records = []
+        primer_bed_records = []
+        primer_assignment_records = []
+        pools = {amp.get("pool", 0) for amp in amplicon_scheme}
+        for pool in pools:
             if mode == "single":
                 primer_fasta_file = os.path.join(path, "primers.fasta")
             else:
-                primer_fasta_file = os.path.join(path, f"primers_pool_{pool}.fasta")
+                primer_fasta_file = os.path.join(path, f"primers_pool_{pool+1}.fasta")
             with open(primer_fasta_file, "w") as primer_fasta:
-                for amp in amplicon_scheme[pool]:
+                for counter, amp in enumerate(amplicon_scheme[pool::len(pools)]):
                     # give a new amplicon name
-                    new_name = f"AMPLICON_{str(counter)}"
-                    counter += 1
+                    amplicon_index = counter*len(pools) + pool
+                    amp_name = f"{scheme_name}_{amplicon_index}"
                     # get left and right primers and their names
-                    primer_names = list(amplicon_scheme[pool][amp].keys())
-                    left = (primer_names[0], amplicon_scheme[pool][amp][primer_names[0]])
-                    right = (primer_names[1], amplicon_scheme[pool][amp][primer_names[1]])
-                    amp_length = right[1][2] - left[1][1]
+                    amp_length = amp["RIGHT"][2] - amp["LEFT"][1]
+                    if "off_targets" in amp:
+                        if amp["off_targets"]:
+                            amplicon_has_off_target = "Yes"
+                            write_BLAST_warning(amp_name, log_file)
+                        else:
+                            amplicon_has_off_target = "No"
+                    else:
+                        amplicon_has_off_target = "n.d."
                     # write amplicon bed
                     if mode == "tiled":
-                        bed_score = pool
+                        bed_score = pool+1
                     elif mode == "single":
-                        bed_score = round(left[1][3] + right[1][3], 1)
-                    print(
-                        "ambiguous_consensus",
-                        left[1][1],
-                        right[1][2],
-                        new_name,
-                        bed_score,
-                        sep="\t",
-                        file=bed
+                        bed_score = round(amp["LEFT"][3] + amp["RIGHT"][3], 1)
+                    amplicon_bed_records.append(
+                        (
+                            amp["LEFT"][1],
+                            amp["RIGHT"][2],
+                            amp_name,
+                            bed_score
+                        )
                     )
-                    # write primer assignments tabular file
-                    print(
-                        f"{new_name}_LEFT",
-                        f"{new_name}_RIGHT",
-                        sep="\t",
-                        file=tabular
+                    primer_assignment_records.append(
+                        (
+                            # will need amplicon_index for sorting
+                            amplicon_index,
+                            (f"{amp_name}_LEFT", f"{amp_name}_RIGHT")
+                        )
                     )
                     # write primer tsv and primer bed
-                    for direction, primer in [("+", left), ("-", right)]:
-                        seq = ambiguous_consensus[primer[1][1]:primer[1][2]]
+                    for direction, primer in [("+", amp["LEFT"]), ("-", amp["RIGHT"])]:
+                        seq = ambiguous_consensus[primer[1]:primer[2]]
                         if direction == "-":
                             seq = primers.rev_complement(seq)
-                            primer_name = f"{new_name}_RIGHT"
+                            primer_name = f"{amp_name}_RIGHT"
                         else:
-                            primer_name = f"{new_name}_LEFT"
+                            primer_name = f"{amp_name}_LEFT"
                         # write primers to fasta pool file
                         print(f">{primer_name}\n{seq.upper()}", file=primer_fasta)
                         # calc primer parameters for all permutations
@@ -277,30 +301,54 @@ def write_scheme_to_files(path, amplicon_scheme, ambiguous_consensus, mode):
                         gc, temp = calc_mean_stats(permutations)
                         # write tsv file
                         print(
-                            new_name,
+                            amp_name,
                             amp_length,
                             primer_name,
-                            primer[0],
-                            pool,
-                            primer[1][1] + 1,
-                            primer[1][2],
+                            primer[-1],
+                            pool+1,
+                            primer[1] + 1,
+                            primer[2],
                             seq.upper(),
-                            len(primer[1][0]),
-                            round(primers.calc_gc(primer[1][0]), 1),
-                            round(primers.calc_temp(primer[1][0]), 1),
+                            len(primer[0]),
+                            round(primers.calc_gc(primer[0]), 1),
+                            round(primers.calc_temp(primer[0]), 1),
                             gc,
                             temp,
-                            round(primer[1][3], 1),
+                            round(primer[3], 1),
+                            amplicon_has_off_target,
                             sep="\t",
                             file=tsv
                         )
-                        # write primer bed file
-                        write_primers_to_bed(
-                            primer_bed_file,
-                            primer_name,
-                            primer[1],
-                            direction
+                        primer_bed_records.append(
+                            (
+                                # will need amplicon_index for sorting
+                                amplicon_index,
+                                (primer_name, primer, pool+1, direction, seq.upper())
+                            )
                         )
+        # write amplicon bed with amplicons sorted by start position
+        for record in sorted(amplicon_bed_records, key=lambda x: x[0]):
+            print(
+                f"{scheme_name}_consensus",
+                *record,
+                ".",
+                sep="\t",
+                file=bed
+            )
+        # use sorting by amplicon index for primer assignment file
+        for record in sorted(primer_assignment_records):
+            print(
+                *record[1],
+                sep="\t",
+                file=tabular
+            )
+        # same for primer bed
+        for record in sorted(primer_bed_records):
+            write_primers_to_bed(
+                primer_bed_file,
+                scheme_name,
+                *record[1]
+            )
 
 
 def write_dimers(path, primer_dimers):
@@ -313,12 +361,12 @@ def write_dimers(path, primer_dimers):
             "pool\tprimer_name_1\tprimer_name_2\tdimer melting temp",
             file=tsv
         )
-        for dimers in primer_dimers:
+        for pool, primer1, primer2 in primer_dimers:
             print(
-                dimers[0][0],
-                dimers[0][2],
-                dimers[1][2],
-                round(primers.calc_dimer(dimers[0][3][0], dimers[1][3][0]).tm, 1),
+                pool+1,
+                primer1[1],
+                primer2[1],
+                round(primers.calc_dimer(primer1[2][0], primer2[2][0]).tm, 1),
                 sep="\t",
                 file=tsv
             )
@@ -368,7 +416,7 @@ def alignment_entropy(alignment_cleaned):
     return entropy_df
 
 
-def entropy_subplot(ax, alignment_cleaned):
+def entropy_subplot(ax, alignment_cleaned, scheme_name):
     """
     creates the entropy subplot
     """
@@ -380,7 +428,7 @@ def entropy_subplot(ax, alignment_cleaned):
     ax[0].set_ylim((0, 1))
     ax[0].set_xlim(0, max(entropy_df["position"]))
     ax[0].set_ylabel("normalized Shannon's entropy")
-    ax[0].set_title("final amplicon design")
+    ax[0].set_title(f"{scheme_name} amplicon design")
     ax[0].spines['top'].set_visible(False)
     ax[0].spines['right'].set_visible(False)
 
@@ -418,26 +466,23 @@ def amplicon_subplot(ax, amplicon_scheme):
     """
     creates the amplicon subplot
     """
-    counter = 0
-    for pool in amplicon_scheme:
-        for amp in amplicon_scheme[pool]:
-            if pool == 0:
-                position_amp = 0.7
-                position_text = 0.6
-            elif pool == 1:
-                position_amp = 0.6
-                position_text = 0.65
-            primer_names = [i for i in amplicon_scheme[pool][amp]]
-            left = amplicon_scheme[pool][amp][primer_names[0]]
-            right = amplicon_scheme[pool][amp][primer_names[1]]
-            # amplicons
-            ax[1].hlines(position_amp, left[1], right[2], linewidth=5)
-            # text
-            ax[1].text(right[2] - (right[2]-left[1])/2, position_text, str(counter), fontsize=8)
-            # primers
-            ax[1].hlines(position_amp, left[1], left[2], linewidth=5, color="red")
-            ax[1].hlines(position_amp, right[1], right[2], linewidth=5, color="red")
-            counter += 1
+    for counter, amp in enumerate(amplicon_scheme):
+        pool = amp.get("pool", 0)
+        if pool == 0:
+            position_amp = 0.7
+            position_text = 0.6
+        elif pool == 1:
+            position_amp = 0.6
+            position_text = 0.65
+        left = amp["LEFT"]
+        right = amp["RIGHT"]
+        # amplicons
+        ax[1].hlines(position_amp, left[1], right[2], linewidth=5)
+        # text
+        ax[1].text(right[2] - (right[2]-left[1])/2, position_text, str(counter), fontsize=8)
+        # primers
+        ax[1].hlines(position_amp, left[1], left[2], linewidth=5, color="red")
+        ax[1].hlines(position_amp, right[1], right[2], linewidth=5, color="red")
     # legends
     ax[1].hlines(position_amp, left[1]+config.PRIMER_SIZES[1], right[2]-config.PRIMER_SIZES[1], linewidth=5, label="amplicons")
     ax[1].hlines(position_amp, left[1], left[2], linewidth=5, color="red", label="primers")
@@ -447,12 +492,10 @@ def qpcr_subplot(ax, amplicon_scheme):
     """
     creates the qpcr subplot
     """
-    counter = 0
-
-    for scheme in amplicon_scheme:
-        left = amplicon_scheme[scheme]["LEFT"]
-        right = amplicon_scheme[scheme]["RIGHT"]
-        probe = amplicon_scheme[scheme]["PROBE"]
+    for counter, amp in enumerate(amplicon_scheme):
+        left = amp["LEFT"]
+        right = amp["RIGHT"]
+        probe = amp["PROBE"]
         # amplicons
         ax[1].hlines(0.8, left[1], right[2], linewidth=5)
         # text
@@ -463,14 +506,13 @@ def qpcr_subplot(ax, amplicon_scheme):
         # probe
         ax[1].hlines(0.75, probe[1], probe[2], linewidth=5, color="darkgrey")
 
-        counter += 1
     # legends
     ax[1].hlines(0.8, left[1]+config.PRIMER_SIZES[1], right[2]-config.PRIMER_SIZES[1], linewidth=5, label="amplicons")
     ax[1].hlines(0.8, left[1], left[2], linewidth=5, color="red", label="primers")
     ax[1].hlines(0.75, probe[1], probe[2], linewidth=5, color="darkgrey", label="probe")
 
 
-def varvamp_plot(path, alignment_cleaned, primer_regions, all_primers=None, amplicon_scheme=None, probe_regions=None):
+def varvamp_plot(path, alignment_cleaned, primer_regions, scheme_name, all_primers=None, amplicon_scheme=None, probe_regions=None):
     """
     creates overview plot for the amplicon design
     and per base coverage plots
@@ -483,7 +525,7 @@ def varvamp_plot(path, alignment_cleaned, primer_regions, all_primers=None, ampl
     fig, ax = plt.subplots(2, 1, figsize=[22, 6], squeeze=True, sharex=True, gridspec_kw={'height_ratios': [4, 1]})
     fig.subplots_adjust(hspace=0)
     # entropy plot
-    entropy_subplot(ax, alignment_cleaned)
+    entropy_subplot(ax, alignment_cleaned, scheme_name)
     # primer regions plot
     region_subplot(ax, primer_regions)
     # probe region plot for probes
@@ -509,48 +551,32 @@ def varvamp_plot(path, alignment_cleaned, primer_regions, all_primers=None, ampl
     plt.close()
 
 
-def get_SINGLE_TILED_primers_for_plot(amplicon_scheme):
+def get_primers_for_plot(amplicon_scheme, scheme_name, mode):
     """
     get the primers for per base pair plot (single, tiled)
     """
     amplicon_primers = []
 
-    for pool in amplicon_scheme:
-        for amp in amplicon_scheme[pool]:
-            primer_names = [i for i in amplicon_scheme[pool][amp]]
-            left = amplicon_scheme[pool][amp][primer_names[0]]
-            right = amplicon_scheme[pool][amp][primer_names[1]]
-            amplicon_primers.append((primer_names[0], left))
-            amplicon_primers.append((primer_names[1], right))
-
-    return amplicon_primers
-
-
-def get_QPCR_primers_for_plot(amplicon_schemes):
-    """
-    get the primers for per base pair plot (qpcr)
-    """
-    amplicon_primers = []
+    if mode == "SINGLE/TILED":
+        oligo_types = ["LEFT", "RIGHT"]
+    else:
+        oligo_types = ["PROBE", "LEFT", "RIGHT"]
 
-    for scheme in amplicon_schemes:
-        for type in amplicon_schemes[scheme]:
-            if type == "penalty" or type == "deltaG":
-                continue
-            primer_name = f"{scheme}_{type}"
-            amplicon_primers.append((primer_name, amplicon_schemes[scheme][type]))
+    for counter, amp in enumerate(amplicon_scheme):
+        for oligo_type in oligo_types:
+            primer_name = f"{scheme_name}_{counter}_{oligo_type}"
+            amplicon_primers.append((primer_name, amp[oligo_type]))
 
     return amplicon_primers
 
 
-def per_base_mismatch_plot(path, amplicon_scheme, threshold, mode="SINGLE/TILED"):
+def per_base_mismatch_plot(path, amplicon_scheme, threshold, scheme_name, mode="SINGLE/TILED"):
     """
     per base pair mismatch multiplot
     """
     out = os.path.join(path, "per_base_mismatches.pdf")
-    if mode == "SINGLE/TILED":
-        amplicon_primers = get_SINGLE_TILED_primers_for_plot(amplicon_scheme)
-    elif mode == "QPCR":
-        amplicon_primers = get_QPCR_primers_for_plot(amplicon_scheme)
+
+    amplicon_primers = get_primers_for_plot(amplicon_scheme, scheme_name, mode)
     # ini multi pdf
     with PdfPages(out) as pdf:
         # always print 4 primers to one page
@@ -581,3 +607,15 @@ def per_base_mismatch_plot(path, amplicon_scheme, threshold, mode="SINGLE/TILED"
             # - to pdf
             pdf.savefig(fig, bbox_inches='tight')
             plt.close()
+
+
+def write_BLAST_warning(amplicon_name, log_file):
+    """
+    for each primer pair that has potential unspecific amplicons
+    write warnings to file.
+    """
+    logging.raise_error(
+        f"{amplicon_name} could produce off-targets. No better amplicon in this area was found.",
+        log_file,
+        exit=False,
+    )