Merge pull request #20 from cov-lineages/dev

merge in dev

Author: rmcolq

scorpio/__init__.py

@@ -1,2 +1,2 @@
 _program = "scorpio"
-__version__ = "0.3.5"
+__version__ = "0.3.7"

scorpio/__main__.py

@@ -99,13 +99,17 @@ def main(sysargs = sys.argv[1:]):
                                  ' and their coordinates'
     )
     subparser_haplotype.add_argument(
-        "--ref-char", dest="ref_char", default='-', required=False,
+        "--ref-char", dest="ref_char", required=False,
         help="Symbol to use to represent reference allele"
     )
     subparser_haplotype.add_argument(
         "--output-counts", dest="output_counts", action="store_true",
         help="Save a file per constellation of ref, alt and other counts"
     )
+    subparser_haplotype.add_argument(
+        "--append-genotypes", dest="append_genotypes", action="store_true",
+        help="Output a column per variant with the call"
+    )
     subparser_haplotype.set_defaults(func=scorpio.subcommands.haplotype.run)
 
     # _______________________________  report  __________________________________#
@@ -224,6 +228,14 @@ def main(sysargs = sys.argv[1:]):
         if "call_all" in args and args.call_all and args.long:
             print("Cannot provide long format summary file with multiple calls, ignoring --long\n")
 
+        if "append_genotypes" in args and args.append_genotypes and not args.ref_char:
+            args.ref_char = None
+        elif "ref_char" in args and not args.ref_char:
+            args.ref_char = '-'
+
+        if "append_genotypes" in args and args.append_genotypes and not args.output_counts:
+            args.output_counts = True
+
     """
     Exit with help menu if no args supplied
     """

scorpio/scripts/extract_definitions.py

@@ -124,6 +124,8 @@ def translate_if_possible(nuc_start, nuc_ref, nuc_alt, feature_dict, reference_s
     #print(query_seq[nuc_start-5: nuc_end+5])
 
     for feature in feature_dict:
+        if len(feature_dict[feature]) > 2:
+            continue # ignore nsp definitions
         if feature_dict[feature][0] <= nuc_start <= feature_dict[feature][1]:
             start, end = nuc_start, nuc_end
             while (start - feature_dict[feature][0]) % 3 != 0:

scorpio/scripts/type_constellations.py

@@ -552,18 +552,20 @@ def generate_barcode(record_seq, variant_list, ref_char=None, ins_char="?", oth_
     counts['support'] = round(counts['alt'] / float(counts['alt'] + counts['ref'] + counts['ambig'] + counts['oth']), 4)
     counts['conflict'] = round(counts['ref'] / float(counts['alt'] + counts['ref'] + counts['ambig'] + counts['oth']), 4)
 
-    return ''.join(barcode_list), counts
+    return barcode_list, counts
 
 
 def type_constellations(in_fasta, list_constellation_files, constellation_names, out_csv, reference_json, ref_char=None,
-                        output_counts=False, label=None):
+                        output_counts=False, label=None, append_genotypes=False):
     reference_seq, features_dict = load_feature_coordinates(reference_json)
 
     constellation_dict = {}
     for constellation_file in list_constellation_files:
         constellation, variants, ignore, mrca_lineage, incompatible_lineage_calls = parse_variants_in(reference_seq, features_dict, constellation_file, label=label)
         if not constellation:
             continue
+        if constellation_names and constellation not in constellation_names:
+            continue
         if len(variants) > 0:
             constellation_dict[constellation] = variants
             print("Found file %s for constellation %s containing %i variants" % (
@@ -579,7 +581,10 @@ def type_constellations(in_fasta, list_constellation_files, constellation_names,
         for constellation in constellation_dict:
             clean_name = re.sub("[^a-zA-Z0-9_\-.]", "_", constellation)
             counts_out[constellation] = open("%s.%s_counts.csv" % (out_csv.replace(".csv", ""), clean_name), "w")
-            counts_out[constellation].write("query,ref_count,alt_count,ambig_count,other_count,support,conflict\n")
+            columns = ["query,ref_count,alt_count,ambig_count,other_count,support,conflict"]
+            if append_genotypes:
+                columns.extend([var["name"] for var in constellation_dict[constellation]])
+            counts_out[constellation].write("%s\n" % ','.join(columns))
 
     with open(in_fasta, "r") as f:
         for record in SeqIO.parse(f, "fasta"):
@@ -590,11 +595,14 @@ def type_constellations(in_fasta, list_constellation_files, constellation_names,
 
             out_list = [record.id]
             for constellation in constellation_dict:
-                barcode, counts = generate_barcode(record.seq, constellation_dict[constellation], ref_char)
+                barcode_list, counts = generate_barcode(record.seq, constellation_dict[constellation], ref_char)
                 if output_counts:
-                    counts_out[constellation].write("%s,%i,%i,%i,%i,%f,%f\n" % (record.id, counts['ref'], counts['alt'],
-                                                 counts['ambig'], counts['oth'], counts['support'], counts['conflict']))
-                out_list.append(barcode)
+                    columns = ["%s,%i,%i,%i,%i,%f,%f" % (record.id, counts['ref'], counts['alt'],
+                               counts['ambig'], counts['oth'], counts['support'], counts['conflict'])]
+                    if append_genotypes:
+                        columns.extend(barcode_list)
+                    counts_out[constellation].write("%s\n" % ','.join(columns))
+                out_list.append(''.join(barcode_list))
 
             variants_out.write("%s\n" % ",".join(out_list))
 

scorpio/subcommands/haplotype.py

@@ -11,4 +11,5 @@ def run(options):
                         options.reference_json,
                         options.ref_char,
                         options.output_counts,
-                        options.label)
+                        options.label,
+                        options.append_genotypes)

scorpio/tests/type_constellations_test.py

@@ -276,20 +276,24 @@ def test_generate_barcode():
                      {"ref": 0, "alt": 1, "ambig": 0, "oth": 5, "support": 0.1667, "conflict": 0.0}]
 
     for i in range(len(seqs)):
-        barcode, counts = generate_barcode(seqs[i], variants, ref_char="-", ins_char="?", oth_char="X")
+        barcode_list, counts = generate_barcode(seqs[i], variants, ref_char="-", ins_char="?", oth_char="X")
+        barcode = ''.join(barcode_list)
         print(i, barcode, counts)
         assert barcode == expect_barcode_dash[i]
         assert counts == expect_counts[i]
 
-        barcode, counts = generate_barcode(seqs[i], variants, ref_char=None, ins_char="?", oth_char="X")
+        barcode_list, counts = generate_barcode(seqs[i], variants, ref_char=None, ins_char="?", oth_char="X")
+        barcode = ''.join(barcode_list)
         print(i, barcode, counts)
         assert barcode == expect_barcode_ref[i]
 
-        barcode, counts = generate_barcode(seqs[i], variants, ref_char=None, ins_char="?", oth_char=None)
+        barcode_list, counts = generate_barcode(seqs[i], variants, ref_char=None, ins_char="?", oth_char=None)
+        barcode = ''.join(barcode_list)
         print(i, barcode, counts)
         assert barcode == expect_barcode_ref_oth[i]
 
-        barcode, counts = generate_barcode(seqs[i], variants, ref_char="-", ins_char="$", oth_char="X")
+        barcode_list, counts = generate_barcode(seqs[i], variants, ref_char="-", ins_char="$", oth_char="X")
+        barcode = ''.join(barcode_list)
         print(i, barcode, counts)
         assert barcode == expect_barcode_ins[i]