Repository: cogent3/cogent3
Pull request #1797 (2024-03-29 04:37:18)

ENH: improve accessibility of ENS statistic from likelihood functions


Patch:
diff --git a/changelog.d/20240329_152534_Gavin.Huttley.md b/changelog.d/20240329_152534_Gavin.Huttley.md
new file mode 100644
index 000000000..779132349
--- /dev/null
+++ b/changelog.d/20240329_152534_Gavin.Huttley.md
@@ -0,0 +1,48 @@
+<!--
+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
+
+- methods on likelihood function objects that work with continuous-time
+  Markov models no longer fail if the model also has discrete-time edges.
+  These include lf.get_lengths_as_ens(), lf.get_annotated_tree(),
+  lf.get_paralinear_metric().
+- Added new lf.get_ens_tree(). This returns trees with the expected number
+  of substitutions as the branch length. On discrete-time edges the branch
+  length is set to None. Think of this tree as the true "evolutionary tree".
+  Thanks to Von Bing Yap for suggesting this!
+
+<!--
+### 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/result.py b/src/cogent3/app/result.py
index a2d23cbc8..aacdc75e1 100644
--- a/src/cogent3/app/result.py
+++ b/src/cogent3/app/result.py
@@ -341,6 +341,7 @@ class model_result(generic_result):
     @property
     def tree(self):
         """an annotated tree with 'ENS' set as the branch length
+
         Note
         ----
         In the case of a discrete time process, length is 'paralinear'"""
diff --git a/src/cogent3/evolve/likelihood_function.py b/src/cogent3/evolve/likelihood_function.py
index 64aed6862..46225c116 100644
--- a/src/cogent3/evolve/likelihood_function.py
+++ b/src/cogent3/evolve/likelihood_function.py
@@ -1,5 +1,6 @@
 import json
 import random
+import typing
 
 from collections import defaultdict
 from copy import deepcopy
@@ -8,6 +9,7 @@ import numpy
 
 from cogent3._version import __version__
 from cogent3.core.alignment import ArrayAlignment
+from cogent3.core.tree import PhyloNode
 from cogent3.evolve import substitution_model
 from cogent3.evolve.simulate import AlignmentEvolver, random_sequence
 from cogent3.maths.matrix_exponential_integration import expected_number_subs
@@ -17,6 +19,7 @@ from cogent3.maths.measure import (
     paralinear_discrete_time,
 )
 from cogent3.recalculation.definition import ParameterController
+from cogent3.recalculation.scope import InvalidScopeError
 from cogent3.util import table
 from cogent3.util.dict_array import DictArrayTemplate
 from cogent3.util.misc import adjusted_gt_minprob, get_object_provenance
@@ -376,10 +379,8 @@ class LikelihoodFunction(ParameterController):
                 length = self.get_param_value("length", edge=name, **kw)
                 array *= length
         except KeyError as err:
-            if err[0] == "Q" and name != "Q":
-                raise RuntimeError("rate matrix not known by this model")
-            else:
-                raise
+            raise InvalidScopeError from err
+
         return DictArrayTemplate(self._motifs, self._motifs).wrap(array)
 
     def _getLikelihoodValuesSummedAcrossAnyBins(self, locus=None):
@@ -564,7 +565,7 @@ class LikelihoodFunction(ParameterController):
         results = [title, lnL, nfp] + results if lnL else [title, nfp] + results
         return "\n".join(map(str, results))
 
-    def get_annotated_tree(self, length_as=None):
+    def get_annotated_tree(self, length_as: typing.Optional[str] = None) -> PhyloNode:
         """returns tree with model attributes on node.params
 
         length_as : str or None
@@ -582,15 +583,15 @@ class LikelihoodFunction(ParameterController):
 
         is_discrete = isinstance(self.model, DiscreteSubstitutionModel)
 
-        if is_discrete and not length_as == "paralinear":
+        if is_discrete and length_as != "paralinear":
             raise ValueError(f"{length_as} invalid for discrete time process")
 
         assert length_as in ("ENS", "paralinear", None)
         d = self.get_param_value_dict(["edge"])
         lengths = d.pop("length", None)
         mprobs = self.get_motif_probs_by_node()
-        if not is_discrete:
-            ens = self.get_lengths_as_ens(motif_probs=mprobs)
+
+        ens = {} if is_discrete else self.get_lengths_as_ens(motif_probs=mprobs)
 
         plin = self.get_paralinear_metric(motif_probs=mprobs)
         if length_as == "ENS":
@@ -604,20 +605,46 @@ class LikelihoodFunction(ParameterController):
                 edge.params["mprobs"] = mprobs[edge.name].to_dict()
                 continue
 
-            if not is_discrete:
-                edge.params["ENS"] = ens[edge.name]
-
+            edge.params["ENS"] = ens.get(edge.name)
             edge.params["length"] = lengths[edge.name]
             edge.params["paralinear"] = plin[edge.name]
             edge.params["mprobs"] = mprobs[edge.name].to_dict()
             for par in d:
-                val = d[par][edge.name]
+                val = d[par].get(edge.name)
                 if par == length_as:
                     val = ens[edge.name]
                 edge.params[par] = val
 
         return tree
 
+    def get_ens_tree(self) -> PhyloNode:
+        """returns tree with length as ENS
+
+        Notes
+        -----
+        The paralinear distance is added to node.params["paralinear"].
+
+        If it's a discrete-time model, branch lengths are set to None.
+
+        For a stationary model, branch lengths will be unchanged from
+        those values displayed in the statistics tables.
+        """
+        from cogent3.evolve.ns_substitution_model import (
+            DiscreteSubstitutionModel,
+        )
+
+        mprobs = self.get_motif_probs_by_node()
+        if isinstance(self.model, DiscreteSubstitutionModel):
+            raise TypeError("cannot get ENS for discrete-time models")
+
+        ens = self.get_lengths_as_ens(motif_probs=mprobs)
+
+        tree = self._tree.deepcopy()
+        for edge in tree.get_edge_vector(include_root=False):
+            edge.params["length"] = ens[edge.name]
+
+        return tree
+
     def get_motif_probs(self, edge=None, bin=None, locus=None, position=None):
         """
         Parameters
@@ -702,23 +729,22 @@ class LikelihoodFunction(ParameterController):
         motif_probs : dict or DictArray
             an item for each edge of the tree. Computed if not provided.
         """
-        from cogent3.evolve.ns_substitution_model import (
-            DiscreteSubstitutionModel,
-        )
-
-        is_discrete = isinstance(self.model, DiscreteSubstitutionModel)
-
         if motif_probs is None:
             motif_probs = self.get_motif_probs_by_node()
+
         plin = {}
         for edge in self.tree.get_edge_vector(include_root=False):
             parent_name = edge.parent.name
             pi = motif_probs[parent_name]
             P = self.get_psub_for_edge(edge.name)
-            if is_discrete:
+            try:
+                Q = self.get_rate_matrix_for_edge(edge.name, calibrated=False)
+            except InvalidScopeError:
+                Q = None
+
+            if Q is None:
                 para = paralinear_discrete_time(P.array, pi.array)
             else:
-                Q = self.get_rate_matrix_for_edge(edge.name, calibrated=False)
                 para = paralinear_continuous_time(P.array, pi.array, Q.array)
 
             plin[edge.name] = para
@@ -729,16 +755,31 @@ class LikelihoodFunction(ParameterController):
         """returns {edge.name: ens, ...} where ens is the expected number of substitutions
 
         for a stationary Markov process, this is just branch length
+
         Parameters
         ----------
         motif_probs : dict or DictArray
             an item for each edge of the tree. Computed if not provided.
         """
+        from cogent3.evolve.ns_substitution_model import (
+            DiscreteSubstitutionModel,
+        )
+
         if motif_probs is None:
             motif_probs = self.get_motif_probs_by_node()
 
         edge_parent = self.tree.child_parent_map()
-        lengths = {e: self.get_param_value("length", edge=e) for e in edge_parent}
+        lengths = {}
+        for e in edge_parent:
+            try:
+                length = self.get_param_value("length", edge=e)
+            except (InvalidScopeError, KeyError):
+                length = None
+            lengths[e] = length
+
+        if isinstance(self.model, DiscreteSubstitutionModel):
+            return lengths
+
         if not isinstance(self.model, substitution_model.Stationary):
             ens = {}
             for e in edge_parent:


Test patch:
diff --git a/tests/test_evolve/test_likelihood_function.py b/tests/test_evolve/test_likelihood_function.py
index 0a3d8acc3..11a3fa8be 100644
--- a/tests/test_evolve/test_likelihood_function.py
+++ b/tests/test_evolve/test_likelihood_function.py
@@ -2261,8 +2261,8 @@ def test_simulate_alignment3():
     assert re.sub("[ATCG]", "x", simulated.to_dict()["a"]) == "x??xxxxxx?"
 
 
-def test_get_lengths_as_ens_matches_manual_calc():
-    """lengths as ENS for a non-reversible model matches a manual calculation"""
+@pytest.fixture(scope="session")
+def gn_mod():
     mprobs = numpy.array([0.1, 0.2, 0.3, 0.4])
     # unequal lengths
     lf = GN().make_likelihood_function(make_tree("(a:0.1,b:0.1,c:0.1)"))
@@ -2272,10 +2272,15 @@ def test_get_lengths_as_ens_matches_manual_calc():
     for par_name in lf.model.get_param_list():
         lf.set_param_rule(par_name, init=init)
         init += 0.1
+    return lf
+
 
-    len_dict = lf.get_lengths_as_ens()
+def test_get_lengths_as_ens_matches_manual_calc(gn_mod):
+    """lengths as ENS for a non-reversible model matches a manual calculation"""
+    mprobs = numpy.array([0.1, 0.2, 0.3, 0.4])
+    len_dict = gn_mod.get_lengths_as_ens()
     expect = expected_number_subs(
-        mprobs, lf.get_rate_matrix_for_edge("a", calibrated=False), 1.0
+        mprobs, gn_mod.get_rate_matrix_for_edge("a", calibrated=False), 1.0
     )
     assert_allclose(len_dict["a"], expect)
 
@@ -2289,4 +2294,125 @@ def test_expm_zero_lengths(expm, DATA_DIR):
     sm = GTR()
     lf = sm.make_likelihood_function(tree=tree, expm=expm)
     lf.set_alignment(aln)
-    lf.optimise(max_evaluations=5, show_progress=False)
+    lf.optimise(max_evaluations=5, show_progress=False, limit_action="ignore")
+
+
+def test_lengths_as_ens_model_mix(DATA_DIR):
+    aln = load_aligned_seqs((DATA_DIR / "brca1_5.paml"), moltype="dna")
+    names = ["Mouse", "Human", "HowlerMon"]
+    aln = aln.take_seqs(names)
+    tree = make_tree((f"({names[0]}:0.0,{names[1]}:0.0,{names[2]}:0.0)"))
+    sm = GTR()
+    lf = sm.make_likelihood_function(tree=tree, discrete_edges=["Mouse"])
+    lf.set_alignment(aln)
+    lengths = lf.get_lengths_as_ens()
+    assert lengths["Mouse"] is None
+    assert lengths["Human"] == lengths["HowlerMon"] == 1
+
+
+@pytest.mark.parametrize("discrete_time", (False, True)[1:])
+def test_get_annotated_tree_model_mix(DATA_DIR, discrete_time):
+    aln = load_aligned_seqs((DATA_DIR / "brca1_5.paml"), moltype="dna")
+    names = ["Mouse", "Human", "HowlerMon"]
+    aln = aln.take_seqs(names)
+    tree = make_tree((f"({names[0]}:0.0,{names[1]}:0.0,{names[2]}:0.0)"))
+    sm = GTR()
+    lf = sm.make_likelihood_function(
+        tree=tree, discrete_edges=["Mouse"] if discrete_time else None
+    )
+    lf.set_alignment(aln)
+    lf.optimise(max_evaluations=5, show_progress=False, limit_action="ignore")
+    anno_tree = lf.get_annotated_tree(length_as="ENS")
+    assert isinstance(anno_tree, type(tree))
+
+
+@pytest.fixture(scope="session")
+def mixed_model(DATA_DIR):
+    aln = load_aligned_seqs((DATA_DIR / "brca1_5.paml"), moltype="dna")
+    names = ["Mouse", "Human", "HowlerMon"]
+    aln = aln.take_seqs(names)
+    tree = make_tree((f"({names[0]}:0.0,{names[1]}:0.0,{names[2]}:0.0)"))
+    sm = GTR()
+    lf = sm.make_likelihood_function(tree=tree, discrete_edges=["Mouse"])
+    lf.set_alignment(aln)
+    return lf
+
+
+@pytest.fixture(scope="session")
+def bh_model(DATA_DIR):
+    aln = load_aligned_seqs((DATA_DIR / "brca1_5.paml"), moltype="dna")
+    names = ["Mouse", "Human", "HowlerMon"]
+    aln = aln.take_seqs(names)
+    tree = make_tree((f"({names[0]}:0.0,{names[1]}:0.0,{names[2]}:0.0)"))
+    sm = get_model("BH")
+    lf = sm.make_likelihood_function(tree=tree)
+    lf.set_alignment(aln)
+    return lf
+
+
+def test_rate_matrix_for_edge_mixed_discrete(mixed_model):
+    from cogent3.recalculation.scope import InvalidScopeError
+    from cogent3.util.dict_array import DictArray
+
+    with pytest.raises(InvalidScopeError):
+        mixed_model.get_rate_matrix_for_edge("Mouse")
+
+    Q = mixed_model.get_rate_matrix_for_edge("Human")
+
+    assert isinstance(Q, DictArray)
+
+
+def test_get_paralinear_metric_mixed(mixed_model):
+    # should not fail with a mixed discrete, continuous-time model
+    dists = mixed_model.get_paralinear_metric()
+    # making sure have no None's
+    assert None not in set(dists.values())
+
+
+def test_get_lengths_as_ens_mixed(mixed_model):
+    # should not fail with a mixed discrete, continuous-time model
+    lengths = mixed_model.get_lengths_as_ens()
+    # and the length for discrete-time edge Mouse should be None
+    assert lengths["Mouse"] is None
+    # comparing against default value
+    assert lengths["Human"] == lengths["HowlerMon"] == 1
+
+
+def test_get_lengths_as_ens_discrete(bh_model):
+    # should not fail with a discrete-time model
+    lengths = bh_model.get_lengths_as_ens()
+    assert lengths["Human"] is lengths["HowlerMon"] is lengths["Mouse"] is None
+
+
+def test_get_annotated_tree_mixed(mixed_model):
+    # should not fail with a mixed discrete, continuous-time model
+    tree = mixed_model.get_annotated_tree(length_as="ENS")
+    # and the length for discrete-time edge Mouse should be None
+    mouse = tree.get_node_matching_name("Mouse")
+    assert mouse.length is None
+    assert mouse.params["ENS"] is None
+    assert mouse.params["paralinear"] is not None
+
+
+def test_get_ens_tree_mixed(mixed_model):
+    # should not fail with a mixed discrete, continuous-time model
+    tree = mixed_model.get_ens_tree()
+    # and the length for discrete-time edge Mouse should be None
+    mouse = tree.get_node_matching_name("Mouse")
+    assert mouse.length is None
+
+
+def test_get_ens_tree_gn(gn_mod):
+    # should not fail with a mixed discrete, continuous-time model
+    tree = gn_mod.get_ens_tree()
+    ens = gn_mod.get_lengths_as_ens()
+    tree_lengths = {
+        edge.name: edge.length for edge in tree.postorder() if edge.name != "root"
+    }
+    assert ens == tree_lengths
+
+
+def test_get_ens_tree_discrete(bh_model):
+    # should fail with a discrete-time model
+    with pytest.raises(TypeError):
+        bh_model.get_ens_tree()