nopaste

commit 18010940faba6c5d20fbf19e37bcb3faf79f0b4b
Author: Thomas Sibley <tsibley@fredhutch.org>
Date:   Fri Dec 7 12:33:17 2018 -0800

    titers tree: Add a --branch-filter option allowing for a custom function
    
    Sidney and I talked about this as a way to support a custom branch
    "criterium" function.  In her dengue work, she only includes
    interserotype branches in the model.

diff --git a/augur/titers.py b/augur/titers.py
index 25d6dbdf..16ec3a58 100644
--- a/augur/titers.py
+++ b/augur/titers.py
@@ -6,7 +6,8 @@ import json, os, sys
 import numpy as np
 from collections import defaultdict
 from Bio import Phylo
-from .utils import read_metadata, read_node_data, write_json
+from functools import partial
+from .utils import read_metadata, read_node_data, write_json, load_function
 
 
 def register_arguments(parser):
@@ -14,9 +15,11 @@ def register_arguments(parser):
     sub_model = subparsers.add_parser('sub', help='substitution model')
     tree_model = subparsers.add_parser('tree', help='tree model')
 
+    branch_filter = partial(load_function, "branch_filter")
 
     tree_model.add_argument('--titers', required=True, type=str, 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('--branch-filter', type=branch_filter, help="an optional file containing a Python function named branch_filter which takes a tree node and returns a boolean indicating if the branch leading to the node should be included in the model (e.g. a branch may want to be excluded if amino acid mutations map to the branch)")
     tree_model.add_argument('--output', '-o', type=str, help='JSON file to save titer model')
     tree_model.set_defaults(model = 'tree')
 
@@ -71,7 +74,7 @@ def infer_tree_model(args):
     T = Phylo.read(args.tree, 'newick')
     from .titer_model import TreeModel
     TM_tree = TreeModel(T, args.titers)
-    TM_tree.prepare()
+    TM_tree.prepare(criterium = args.branch_filter)
     TM_tree.train()
 
     # export the tree model
diff --git a/augur/utils.py b/augur/utils.py
index 8855f633..79b8514c 100644
--- a/augur/utils.py
+++ b/augur/utils.py
@@ -6,6 +6,8 @@ from treetime.utils import numeric_date
 from collections import defaultdict
 from pkg_resources import resource_stream
 from io import TextIOWrapper
+from inspect import isfunction
+from pathlib import Path
 
 def myopen(fname, mode):
     if fname.endswith('.gz'):
@@ -451,3 +453,50 @@ def nthreads_value(value):
         return int(value)
     except ValueError:
         raise argparse.ArgumentTypeError("'%s' is not an integer or the word 'auto'" % value) from None
+
+
+def load_function(name, filename):
+    """
+    Runs the Python code in *filename* and returns the function named *name*.
+
+    This is used for commands to load arbitrary user-supplied functions from
+    files specified on the command-line.
+    """
+    # Read the whole file into a string
+    with Path(filename).open(encoding = "utf-8") as file:
+        source = file.read()
+
+    # Compile the source string for exec(), without inheriting any __future__
+    # flags from our own context.
+    try:
+        code = compile(source, filename, "exec", dont_inherit = True)
+
+    except (SyntaxError, ValueError) as error:
+        print("Error loading %s() from %s: %s" % (name, filename, error),
+            end  = "\n\n",
+            file = sys.stderr)
+        raise error from None
+
+    # Execute the code in an isolated scope, which we'll inspect afterwards.
+    scope = {}
+
+    try:
+        exec(code, scope)
+
+    except Exception as error:
+        print("Error loading %s() from %s: %s" % (name, filename, error),
+            end  = "\n\n",
+            file = sys.stderr)
+        raise error from None
+
+    # Extract the function we desire.
+    try:
+        function = scope[name]
+
+    except KeyError:
+        raise RuntimeError("No function named %s in %s" % (name, filename)) from None
+
+    assert isfunction(function), \
+        "%s is a not a function" % function
+
+    return function