Repository: cogent3/cogent3
Pull request #1583 (2023-09-21 10:38:48)

ENH: model app can now use a callback to get a tree
[NEW] The optional tree_func argument takes a callable which
    accepts an alignment and returns a corresponding cogent3 Tree.


Patch:
diff --git a/changelog.d/20230921_202221_Gavin.Huttley.md b/changelog.d/20230921_202221_Gavin.Huttley.md
new file mode 100644
index 000000000..6545abcff
--- /dev/null
+++ b/changelog.d/20230921_202221_Gavin.Huttley.md
@@ -0,0 +1,44 @@
+<!--
+A new scriv changelog fragment.
+
+Uncomment the section that is right (remove the HTML comment wrapper).
+-->
+
+<!--
+### Contributors
+
+- A bullet item for the Contributors category.
+
+-->
+### ENH
+
+- Added app.evo.model(tree_func: Optional[Callable]) argument. A callback
+  function assigned to tree_func is passed an alignment and returns a tree.
+  This could be done via loading a tree from disk for that alignment source,
+  or estimation of the tree using the alignment. The argument overrides the
+  other tree related arguments (tree, unique_trees).
+
+<!--
+### BUG
+
+- A bullet item for the BUG category.
+
+-->
+<!--
+### DOC
+
+- A bullet item for the DOC category.
+
+-->
+<!--
+### Deprecations
+
+- A bullet item for the Deprecations category.
+
+-->
+<!--
+### Discontinued
+
+- A bullet item for the Discontinued category.
+
+-->
diff --git a/src/cogent3/app/evo.py b/src/cogent3/app/evo.py
index 08988f5cb..9ab8afd02 100644
--- a/src/cogent3/app/evo.py
+++ b/src/cogent3/app/evo.py
@@ -55,6 +55,7 @@ class model:
         sm,
         tree=None,
         unique_trees=False,
+        tree_func=None,
         name=None,
         optimise_motif_probs=False,
         sm_args=None,
@@ -80,6 +81,9 @@ class model:
         unique_trees : bool
             whether to specify a unique tree per alignment. Only applies if
             number of sequences equals 3.
+        tree_func: callable
+            a callable that takes an alignment and returns a Tree instance.
+            Overrides tree and unique_tree settings.
         name : str
             name of the model
         optimise_motif_probs : bool
@@ -124,6 +128,11 @@ class model:
         with the optimised likelihood function. In the case of split_codons,
         the result object has a separate entry for each codon position.
         """
+        if tree_func:
+            assert callable(tree_func), "tree_func must be callable or None"
+            tree = None
+            unique_trees = False
+
         self._verbose = verbose
         self._lower = lower
         self._upper = upper
@@ -147,6 +156,7 @@ class model:
             split_codons = False
 
         self._tree = interpret_tree_arg(tree)
+        self._tree_func = tree_func
         self._lf_args = deepcopy(lf_args or {})
         if not name:
             name = sm.name or "unnamed model"
@@ -267,8 +277,16 @@ class model:
             return NotCompleted("ERROR", self, msg, source=aln)
 
         evaluation_limit = opt_args.get("max_evaluations", None)
-        if self._tree is None or self._unique_trees:
-            assert len(aln.names) == 3, "to model more than 3, you must provide a tree"
+        if callable(self._tree_func):
+            self._tree = self._tree_func(aln)
+        elif self._tree is None or self._unique_trees:
+            if len(aln.names) > 3:
+                return NotCompleted(
+                    "ERROR",
+                    self,
+                    message="to model more than 3, you must provide a tree",
+                    source=aln,
+                )
             self._tree = make_tree(tip_names=aln.names)
 
         result = model_result(


Test patch:
diff --git a/tests/test_app/test_evo.py b/tests/test_app/test_evo.py
index a26d4dec4..42a059989 100644
--- a/tests/test_app/test_evo.py
+++ b/tests/test_app/test_evo.py
@@ -5,6 +5,8 @@ from tempfile import TemporaryDirectory
 from unittest import TestCase
 from unittest.mock import MagicMock
 
+import pytest
+
 from numpy.testing import assert_allclose, assert_raises
 
 from cogent3 import (
@@ -32,22 +34,6 @@ data_dir = join(dirname(dirname(__file__)), "data")
 class TestModel(TestCase):
     basedir = "data"
 
-    def test_model_str(self):
-        """correct str representation"""
-        model = evo_app.model("HKY85", time_het="max")
-        got = " ".join(str(model).splitlines())
-        expect = (
-            "model(sm='HKY85', tree=None, unique_trees=False, "
-            "name=None, optimise_motif_probs=False, sm_args=None, lf_args=None, "
-            "time_het='max', param_rules=None, "
-            "opt_args=None, lower=1e-06, upper=50, split_codons=False, "
-            "show_progress=False, verbose=False)"
-        )
-        self.assertEqual(
-            got,
-            expect,
-        )
-
     def test_model_opt_mprob_arg(self):
         """argument controls optimisability of motif prob settings"""
         for mn in ("HKY85", "GN", "CNFGTR"):
@@ -281,25 +267,6 @@ class TestModel(TestCase):
         pval = str(result).splitlines()[4].split()[-1]
         self.assertTrue(re.search(r"[0-9\.]+e-\d+", pval) is not None)
 
-    def test_hypothesis_str(self):
-        """correct str representation"""
-        model1 = evo_app.model("HKY85")
-        model2 = evo_app.model("HKY85", name="hky85-max-het", time_het="max")
-        hyp = evo_app.hypothesis(model1, model2)
-        got = " ".join(str(hyp).splitlines())
-        expect = (
-            "hypothesis(null=model(sm='HKY85', tree=None, unique_trees=False, "
-            "name=None, optimise_motif_probs=False, sm_args=None, lf_args=None, "
-            "time_het=None, param_rules=None, opt_args=None, lower=1e-06, upper=50, "
-            "split_codons=False, show_progress=False, verbose=False), "
-            "alternates=(model(sm='HKY85', tree=None, unique_trees=False, "
-            "name='hky85-max-het', optimise_motif_probs=False, sm_args=None, lf_args=None, "
-            "time_het='max', param_rules=None, opt_args=None, lower=1e-06, upper=50,"
-            " split_codons=False, show_progress=False, verbose=False),),"
-            " sequential=True, init_alt=None)"
-        )
-        self.assertEqual(got, expect)
-
     def test_split_pos_model(self):
         """model with split codons, access .lf using codon position int"""
         _data = {
@@ -908,3 +875,46 @@ def test_get_app_tree_is_url():
     tree_url = "https://raw.githubusercontent.com/cogent3/cogent3/develop/tests/data/brca1_5.tree"
     mod = get_app("model", "F81", tree=tree_url)
     assert isinstance(mod, evo_app.model)
+
+
+def test_model_str():
+    """correct str representation"""
+    model = evo_app.model("HKY85", time_het="max")
+    got = " ".join(str(model).splitlines())
+    expect_start = "model(sm='HKY85', tree=None, unique_trees=False, "
+    assert got.startswith(expect_start)
+
+
+def test_hypothesis_str():
+    """correct str representation"""
+    model1 = evo_app.model("HKY85")
+    model2 = evo_app.model("HKY85", name="hky85-max-het", time_het="max")
+    hyp = evo_app.hypothesis(model1, model2)
+    got = " ".join(str(hyp).splitlines())
+    expect_start = "hypothesis(null=model(sm='HKY85', tree=None, unique_trees=False, "
+    assert got.startswith(expect_start)
+
+
+def test_model_tree_func(DATA_DIR):
+    dist_cal = get_app("fast_slow_dist", fast_calc="paralinear", moltype="dna")
+    est_tree = get_app("quick_tree", drop_invalid=True)
+    app = dist_cal + est_tree
+    model = get_app(
+        "model",
+        "HKY85",
+        tree_func=app,
+        opt_args=dict(max_evaluations=25, limit_action="ignore"),
+    )
+
+    aln = load_aligned_seqs(DATA_DIR / "brca1_5.paml", moltype="dna")
+    result = model(aln)
+    assert set(result.tree.get_tip_names()) == set(aln.names)
+
+
+def test_model_invalid_tree_func():
+    with pytest.raises(AssertionError):
+        get_app(
+            "model",
+            "HKY85",
+            tree_func="123",
+        )