Allow users to specify a titer attribute prefix

CHANGES.md

@@ -2,6 +2,11 @@
 
 ## __NEXT__
 
+### Features
+
+* titers: Allow users to specify a custom prefix for attributes in the JSON output (e.g., `cTiter` can be changed to `custom_prefix_cTiter`). [#1106][] (@huddlej)
+
+[#1106]: https://github.com/nextstrain/augur/pull/1106
 
 ## 19.1.0 (14 December 2022)
 

augur/titers.py

@@ -22,6 +22,7 @@ def register_parser(parent_subparsers):
     tree_model.add_argument('--titers', nargs='+', type=str, required=True, help="file with titer measurements")
     tree_model.add_argument('--tree', '-t', type=str, required=True, help="tree to perform fit titer model to")
     tree_model.add_argument('--allow-empty-model', action="store_true", help="allow model to be empty")
+    tree_model.add_argument('--attribute-prefix', default="", help="prefix for node attributes in the JSON output including cumulative titer drop ('cTiter') and per-branch titer drop ('dTiter'). Set a prefix to disambiguate annotations from multiple tree model JSONs in the final Auspice JSON.")
     tree_model.add_argument('--output', '-o', type=str, required=True, help='JSON file to save titer model')
     tree_model.set_defaults(
         __command__ = infer_tree_model
@@ -33,6 +34,7 @@ def register_parser(parent_subparsers):
     sub_model.add_argument('--gene-names', nargs='+', type=str, required=True, help="names of the sequences in the alignment, same order assumed")
     sub_model.add_argument('--tree', '-t', type=str, help="optional tree to annotate fit titer model to")
     sub_model.add_argument('--allow-empty-model', action="store_true", help="allow model to be empty")
+    sub_model.add_argument('--attribute-prefix', default="", help="prefix for node attributes in the JSON output including cumulative titer drop ('cTiterSub') and per-substitution titer drop ('dTiterSub'). Set a prefix to disambiguate annotations from multiple substitution model JSONs in the final Auspice JSON.")
     sub_model.add_argument('--output', '-o', type=str, required=True, help='JSON file to save titer model')
     sub_model.set_defaults(
         __command__ = infer_substitution_model
@@ -75,8 +77,8 @@ def run(args):
             # Store annotations for export to JSON.
             nodes = {
                 node.name: {
-                    "dTiterSub": node.dTiterSub,
-                    "cTiterSub": node.cTiterSub
+                    f"{args.attribute_prefix}dTiterSub": node.dTiterSub,
+                    f"{args.attribute_prefix}cTiterSub": node.cTiterSub
                 }
                 for node in tree.find_clades()
             }
@@ -99,11 +101,18 @@ def run(args):
         try:
             TM_tree.prepare()
             TM_tree.train()
-            tree_model = {'titers':TM_tree.compile_titers(),
-                          'potency':TM_tree.compile_potencies(),
-                          'avidity':TM_tree.compile_virus_effects(),
-                          'nodes':{n.name:{"dTiter": n.dTiter, "cTiter":n.cTiter}
-                                      for n in T.find_clades()}}
+            tree_model = {
+                'titers':TM_tree.compile_titers(),
+                'potency':TM_tree.compile_potencies(),
+                'avidity':TM_tree.compile_virus_effects(),
+                'nodes':{
+                    n.name: {
+                        f"{args.attribute_prefix}dTiter": n.dTiter,
+                        f"{args.attribute_prefix}cTiter": n.cTiter
+                    }
+                    for n in T.find_clades()
+                }
+            }
         except InsufficientDataException:
             print("Unable to train tree model.", file=sys.stderr)
             if args.allow_empty_model:

tests/functional/titers/cram/titers-sub-with-tree-and-custom-prefix.t

@@ -0,0 +1,20 @@
+Setup
+
+  $ pushd "$TESTDIR" > /dev/null
+  $ export AUGUR="${AUGUR:-../../../../bin/augur}"
+
+Test titer substitution model with alignment and tree inputs and a custom prefix for the node data attributes in the output.
+
+  $ ${AUGUR} titers sub \
+  >   --tree ../data/tree.nwk \
+  >   --titers ../data/titers.tsv \
+  >   --alignment ../data/aa_seq_HA1.fasta \
+  >   --gene-names HA1 \
+  >   --attribute-prefix custom_prefix_ \
+  >   --output $TMP/titers-sub.json > /dev/null
+  Read titers from ../data/titers.tsv, found:
+   --- 61 strains
+   --- 15 data sources
+   --- 232 total measurements
+  $ grep custom_prefix_cTiterSub $TMP/titers-sub.json | wc -l
+  \s*120 (re)

tests/functional/titers/cram/titers-sub-with-tree.t

@@ -0,0 +1,19 @@
+Setup
+
+  $ pushd "$TESTDIR" > /dev/null
+  $ export AUGUR="${AUGUR:-../../../../bin/augur}"
+
+Test titer substitution model with alignment and tree inputs.
+
+  $ ${AUGUR} titers sub \
+  >   --tree ../data/tree.nwk \
+  >   --titers ../data/titers.tsv \
+  >   --alignment ../data/aa_seq_HA1.fasta \
+  >   --gene-names HA1 \
+  >   --output $TMP/titers-sub.json > /dev/null
+  Read titers from ../data/titers.tsv, found:
+   --- 61 strains
+   --- 15 data sources
+   --- 232 total measurements
+  $ grep cTiterSub $TMP/titers-sub.json | wc -l
+  \s*120 (re)

tests/functional/titers/cram/titers-tree-with-custom-prefix.t

@@ -0,0 +1,18 @@
+Setup
+
+  $ pushd "$TESTDIR" > /dev/null
+  $ export AUGUR="${AUGUR:-../../../../bin/augur}"
+
+Test titer tree model with a custom prefix for the node data attributes in the output.
+
+  $ ${AUGUR} titers tree \
+  >   --tree ../data/tree.nwk \
+  >   --titers ../data/titers.tsv \
+  >   --attribute-prefix custom_prefix_ \
+  >   --output $TMP/titers-tree.json > /dev/null
+  Read titers from ../data/titers.tsv, found:
+   --- 61 strains
+   --- 15 data sources
+   --- 232 total measurements
+  $ grep custom_prefix_cTiter $TMP/titers-tree.json | wc -l
+  \s*120 (re)

tests/functional/titers/cram/titers-tree.t

@@ -0,0 +1,17 @@
+Setup
+
+  $ pushd "$TESTDIR" > /dev/null
+  $ export AUGUR="${AUGUR:-../../../../bin/augur}"
+
+Test titer tree model.
+
+  $ ${AUGUR} titers tree \
+  >   --tree ../data/tree.nwk \
+  >   --titers ../data/titers.tsv \
+  >   --output $TMP/titers-tree.json > /dev/null
+  Read titers from ../data/titers.tsv, found:
+   --- 61 strains
+   --- 15 data sources
+   --- 232 total measurements
+  $ grep cTiter $TMP/titers-tree.json | wc -l
+  \s*120 (re)