WIP - export root AA sequences for VCF inputs

nextstrain · Dec 13, 2023 · 81eac92 · 81eac92
1 parent 96e93fb
commit 81eac92
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/augur/translate.py b/augur/translate.py
@@ -268,6 +268,15 @@ def assign_aa_vcf(tree, translations):
 
                 aa_muts[c.name]["aa_muts"][fname] = tmp
 
+    ## NEEDS TESTS ADDED WITH AA VARIATION AT ROOT!!!!
+    aa_muts[root.name]['aa_sequences'] = {}
+    for gene_name, gene_data in translations.items():
+        root_seq = list(gene_data['reference'])
+        for pos,alt in gene_data['sequences'][root.name].items():
+            # pos is 0-based, <class 'numpy.int64'>
+            root_seq[pos] = alt
+        aa_muts[root.name]['aa_sequences'][gene_name] = "".join(root_seq)
+
     return aa_muts
 
 def assign_aa_fasta(tree, translations):