Merge pull request #117 from nextstrain/keep-reference-outliers

Keep reference strains in tree even if they are outliers

scripts/flag_outliers.py

@@ -90,6 +90,7 @@ def prepare_tree(T):
     parser.add_argument('--reroot', action="store_true", help="reroot the tree")
     parser.add_argument('--optimize', action="store_true", help="optimize sigma and mu")
     parser.add_argument('--dates', type=str, help='csv/tsv file with dates for each sequence')
+    parser.add_argument('--keep-strains', type=str, help='a list of strains to keep in the output tree regardless of outlier status (i.e., reference strains that need to be retained in the build)')
     parser.add_argument('--output-outliers', type=str, help='file for outliers')
     parser.add_argument('--output-tree', type=str, help='file for pruned tree')
 
@@ -153,9 +154,14 @@ def prepare_tree(T):
         df.to_csv(args.output_outliers, index=False, sep='\t')
 
     if args.output_tree:
+        keep_strains = set()
+        if args.keep_strains:
+            with open(args.keep_strains, "r", encoding="utf-8") as fh:
+                keep_strains = {line.strip() for line in fh}
+
         from Bio import Phylo
         T = tt.tree
         for r, row in df.iterrows():
-            if row['diagnosis']!='bad_date':
+            if row['diagnosis']!='bad_date' and row["sequence"] not in keep_strains:
                 T.prune(row['sequence'])
         Phylo.write(T, args.output_tree, 'newick')

workflow/snakemake_rules/core.smk

@@ -146,13 +146,16 @@ rule prune_outliers:
     output:
         tree = build_dir + "/{build_name}/{segment}/tree_without_outgroup_clean.nwk",
         outliers = build_dir + "/{build_name}/{segment}/outliers.tsv"
+    params:
+        keep_strains_argument=lambda wildcards: "--keep-strains " + config["builds"][wildcards.build_name]["include"] if "include" in config["builds"][wildcards.build_name] else "",
     shell:
         """
         python3 scripts/flag_outliers.py \
             --tree {input.tree:q} \
             --aln {input.aln} \
             --dates {input.metadata} \
             --cutoff 4.0 \
+            {params.keep_strains_argument} \
             --output-tree {output.tree:q} --output-outliers {output.outliers} 2>&1 | tee {log}
         """