Add a new option --chronumental_add_inferred_date (#608)

* Add a new option --chronumental_add_inferred_date to specify the name of a new metadata 'column' for date inferred by Chronumental.  refs #607

taxoniumtools/src/taxoniumtools/usher_to_taxonium.py

@@ -28,6 +28,7 @@ def do_processing(input_file,
                   chronumental_date_output=None,
                   chronumental_tree_output=None,
                   chronumental_reference_node=None,
+                  chronumental_add_inferred_date=None,
                   config_file=None,
                   title=None,
                   overlay_html=None,
@@ -82,7 +83,10 @@ def do_processing(input_file,
             metadata_file=metadata_file,
             chronumental_steps=chronumental_steps,
             chronumental_date_output=chronumental_date_output,
-            chronumental_tree_output=chronumental_tree_output)
+            chronumental_tree_output=chronumental_tree_output,
+            chronumental_add_inferred_date=chronumental_add_inferred_date,
+            metadata_dict=metadata_dict,
+            metadata_cols=metadata_cols)
 
     print("Ladderizing tree..")
     mat.tree.ladderize(ascending=False)
@@ -236,6 +240,12 @@ def get_parser():
         help=
         "A reference node to be used for Chronumental. This should be earlier in the outbreak and have a good defined date. If not set the oldest sample will be automatically picked by Chronumental.",
         default=None)
+    parser.add_argument(
+        "--chronumental_add_inferred_date",
+        type=str,
+        help=
+        "A new metadata-column-like name to be added for display with the value of Chronumental's inferred date for each sample.",
+        default=None)
     parser.add_argument(
         '-j',
         "--config_json",
@@ -304,26 +314,28 @@ def main():
     parser = get_parser()
 
     args = parser.parse_args()
-    do_processing(args.input,
-                  args.output,
-                  metadata_file=args.metadata,
-                  genbank_file=args.genbank,
-                  chronumental_enabled=args.chronumental,
-                  chronumental_steps=args.chronumental_steps,
-                  columns=args.columns,
-                  chronumental_date_output=args.chronumental_date_output,
-                  chronumental_tree_output=args.chronumental_tree_output,
-                  chronumental_reference_node=args.chronumental_reference_node,
-                  config_file=args.config_json,
-                  title=args.title,
-                  overlay_html=args.overlay_html,
-                  remove_after_pipe=args.remove_after_pipe,
-                  clade_types=args.clade_types,
-                  name_internal_nodes=args.name_internal_nodes,
-                  shear=args.shear,
-                  shear_threshold=args.shear_threshold,
-                  only_variable_sites=args.only_variable_sites,
-                  key_column=args.key_column)
+    do_processing(
+        args.input,
+        args.output,
+        metadata_file=args.metadata,
+        genbank_file=args.genbank,
+        chronumental_enabled=args.chronumental,
+        chronumental_steps=args.chronumental_steps,
+        columns=args.columns,
+        chronumental_date_output=args.chronumental_date_output,
+        chronumental_tree_output=args.chronumental_tree_output,
+        chronumental_reference_node=args.chronumental_reference_node,
+        chronumental_add_inferred_date=args.chronumental_add_inferred_date,
+        config_file=args.config_json,
+        title=args.title,
+        overlay_html=args.overlay_html,
+        remove_after_pipe=args.remove_after_pipe,
+        clade_types=args.clade_types,
+        name_internal_nodes=args.name_internal_nodes,
+        shear=args.shear,
+        shear_threshold=args.shear_threshold,
+        only_variable_sites=args.only_variable_sites,
+        key_column=args.key_column)
 
 
 if __name__ == "__main__":

taxoniumtools/src/taxoniumtools/utils.py

@@ -1,7 +1,7 @@
 from alive_progress import alive_it, alive_bar
 import pandas as pd
 import warnings
-import os, tempfile, sys
+import os, tempfile, sys, errno
 import treeswift
 import shutil
 from . import ushertools
@@ -34,7 +34,7 @@ def read_metadata(metadata_file, columns, key_column):
             )
 
         metadata_dict = metadata.to_dict("index")
-        metadata_cols = metadata.columns
+        metadata_cols = metadata.columns.tolist()
         del metadata
         print("Metadata loaded")
         return metadata_dict, metadata_cols
@@ -46,7 +46,8 @@ def read_metadata(metadata_file, columns, key_column):
 
 def do_chronumental(mat, chronumental_reference_node, metadata_file,
                     chronumental_steps, chronumental_date_output,
-                    chronumental_tree_output):
+                    chronumental_tree_output, chronumental_add_inferred_date,
+                    metadata_dict, metadata_cols):
     chronumental_is_available = os.system(
         "which chronumental > /dev/null") == 0
     if not chronumental_is_available:
@@ -93,6 +94,42 @@ def do_chronumental(mat, chronumental_reference_node, metadata_file,
         del time_tree
         del time_tree_iter
 
+        if chronumental_add_inferred_date:
+            print(
+                f"Adding chronumental inferred date as metadata-like item {chronumental_add_inferred_date}"
+            )
+            if chronumental_date_output:
+                date_output_filename = chronumental_date_output
+            else:
+                metadata_dir = os.path.dirname(metadata_file)
+                metadata_base = os.path.basename(metadata_file)
+                date_output_filename = f"chronumental_dates_{metadata_base}.tsv"
+                if metadata_dir:
+                    date_output_filename = metadata_dir + "/" + date_output_filename
+                if not os.path.exists(date_output_filename):
+                    raise FileNotFoundError(
+                        errno.ENOENT,
+                        f"Can't find default date output file in the expected location (try specifying a file name with --chronumental_date_output)",
+                        date_output_filename)
+            metadata_cols.append(chronumental_add_inferred_date)
+            inferred_dates = pd.read_csv(
+                date_output_filename,
+                sep="\t" if metadata_file.endswith(".tsv")
+                or metadata_file.endswith(".tsv.gz") else ",",
+                usecols=['strain', 'predicted_date'])
+            for idx, row in inferred_dates.iterrows():
+                node_name = row['strain']
+                inferred_date = row['predicted_date']
+                # Add inferred_date even if a node (e.g. internal node) has no metadata
+                if node_name not in metadata_dict:
+                    metadata_dict[node_name] = {
+                        col: ""
+                        for col in metadata_cols
+                    }
+                metadata_dict[node_name][
+                    chronumental_add_inferred_date] = inferred_date
+            del inferred_dates
+
 
 def set_x_coords(root, chronumental_enabled):
     """ Set x coordinates for the tree"""