Repository: cogent3/cogent3
Pull request #1572 (2023-09-15 02:36:23)

BUG: fixed error in calculating ENS


Patch:
diff --git a/changelog.d/20230915_112525_Gavin.Huttley.md b/changelog.d/20230915_112525_Gavin.Huttley.md
new file mode 100644
index 000000000..92f7817b7
--- /dev/null
+++ b/changelog.d/20230915_112525_Gavin.Huttley.md
@@ -0,0 +1,41 @@
+<!--
+A new scriv changelog fragment.
+
+Uncomment the section that is right (remove the HTML comment wrapper).
+-->
+
+### Contributors
+
+- YapengLang for reporting the bug in computing the ENS for
+  non-stationary processes.
+
+<!--
+### ENH
+
+- A bullet item for the ENH category.
+
+-->
+
+### BUG
+
+- The calculation of the expected number of substitutions on a tree branch
+  was incorrect. It was not using the motif probs from the parent node.
+
+<!--
+### 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/core/tree.py b/src/cogent3/core/tree.py
index 1c022fe5a..84efc41ab 100644
--- a/src/cogent3/core/tree.py
+++ b/src/cogent3/core/tree.py
@@ -24,6 +24,8 @@ Definition of relevant terms or abbreviations:
        from a node
     -  stem: the edge immediately preceeding a clade
 """
+from __future__ import annotations
+
 import json
 import numbers
 import re
@@ -1617,6 +1619,10 @@ class TreeNode(object):
 
         return lin_rajan_moret(self, tree2)
 
+    def child_parent_map(self) -> dict[str, str]:
+        """return dict of {<child name>: <parent name>, ...}"""
+        return {e.name: e.parent.name for e in self.postorder(include_self=False)}
+
 
 class PhyloNode(TreeNode):
     def __init__(self, *args, **kwargs):
diff --git a/src/cogent3/evolve/likelihood_function.py b/src/cogent3/evolve/likelihood_function.py
index 27b2ed0c3..c4e4badb3 100644
--- a/src/cogent3/evolve/likelihood_function.py
+++ b/src/cogent3/evolve/likelihood_function.py
@@ -736,13 +736,16 @@ class LikelihoodFunction(ParameterController):
         """
         if motif_probs is None:
             motif_probs = self.get_motif_probs_by_node()
-        node_names = [n for n in self.tree.get_node_names() if n != "root"]
-        lengths = {e: self.get_param_value("length", edge=e) for e in node_names}
+
+        edge_parent = self.tree.child_parent_map()
+        lengths = {e: self.get_param_value("length", edge=e) for e in edge_parent}
         if not isinstance(self.model, substitution_model.Stationary):
             ens = {}
-            for e in node_names:
+            for e in edge_parent:
                 Q = self.get_rate_matrix_for_edge(e)
-                length = expected_number_subs(motif_probs[e], Q, lengths[e])
+                length = expected_number_subs(
+                    motif_probs[edge_parent[e]], Q, lengths[e]
+                )
                 ens[e] = length
 
             lengths = ens


Test patch:
diff --git a/tests/test_core/test_tree.py b/tests/test_core/test_tree.py
index 6d8d8647c..0acfc36d3 100644
--- a/tests/test_core/test_tree.py
+++ b/tests/test_core/test_tree.py
@@ -3,10 +3,11 @@
 import json
 import os
 import pathlib
+import random
 
 from copy import copy, deepcopy
 from tempfile import TemporaryDirectory
-from unittest import TestCase, main
+from unittest import TestCase
 
 import pytest
 
@@ -2210,3 +2211,16 @@ def test_lrm_method():
     b = make_tree(treestring="(1,((((2,3),4),5),((6,7),((8,9),(10,11)))),12);")
     distance = a.lin_rajan_moret(b)
     assert distance == 8
+
+
+@pytest.mark.parametrize(
+    "nwk", ("(((g,b)gb,(c,d)cd),(e,f),a)", "(a,b,c)", ("(a,(b,(c,d)cd))"))
+)
+def test_child_parent_map(nwk):
+    tree = make_tree(nwk)
+    child_2_parent = tree.child_parent_map()
+    all_edges = tree.get_edge_vector(include_root=False)
+    assert len(child_2_parent) == len(all_edges)
+    assert child_2_parent["a"] == "root"
+    edge = random.choice(all_edges)
+    assert child_2_parent[edge.name] == edge.parent.name
diff --git a/tests/test_evolve/test_likelihood_function.py b/tests/test_evolve/test_likelihood_function.py
index d8b35aa81..2aa77a2d7 100644
--- a/tests/test_evolve/test_likelihood_function.py
+++ b/tests/test_evolve/test_likelihood_function.py
@@ -1,4 +1,3 @@
-#!/usr/bin/env python
 """
 Some tests for the likelihood function class.
 
@@ -14,7 +13,7 @@ import json
 import os
 import warnings
 
-from unittest import TestCase, main
+from unittest import TestCase
 
 import numpy
 
@@ -41,6 +40,7 @@ from cogent3.evolve.models import (
     ssGN,
 )
 from cogent3.evolve.ns_substitution_model import GeneralStationary
+from cogent3.maths.matrix_exponential_integration import expected_number_subs
 from cogent3.maths.matrix_exponentiation import PadeExponentiator as expm
 from cogent3.maths.stats.information_criteria import aic, bic
 
@@ -1343,10 +1343,10 @@ DogFaced     root      1.0000    1.0000
 
     def test_get_lengths_as_ens_equal(self):
         """lengths equals ENS for a time-reversible model"""
-        moprobs = numpy.array([0.1, 0.2, 0.3, 0.4])
+        mprobs = numpy.array([0.1, 0.2, 0.3, 0.4])
         length = 0.1
         lf = HKY85().make_likelihood_function(make_tree(tip_names=["a", "b", "c"]))
-        lf.set_motif_probs(moprobs)
+        lf.set_motif_probs(mprobs)
         lf.set_param_rule("kappa", init=1)
         lf.set_param_rule("length", edge="a", init=length)
         len_dict = lf.get_lengths_as_ens()
@@ -1354,10 +1354,10 @@ DogFaced     root      1.0000    1.0000
 
     def test_get_lengths_as_ens_not_equal(self):
         """lengths do not equal ENS for a non-reversible model"""
-        moprobs = numpy.array([0.1, 0.2, 0.3, 0.4])
+        mprobs = numpy.array([0.1, 0.2, 0.3, 0.4])
         length = 0.1
         lf = GN().make_likelihood_function(make_tree(tip_names=["a", "b", "c"]))
-        lf.set_motif_probs(moprobs)
+        lf.set_motif_probs(mprobs)
         lf.set_param_rule("length", init=length)
         # setting arbitrary values for GN rate terms
         init = 0.1
@@ -2260,5 +2260,20 @@ def test_simulate_alignment3():
     assert re.sub("[ATCG]", "x", simulated.to_dict()["a"]) == "x??xxxxxx?"
 
 
-if __name__ == "__main__":
-    main()
+def test_get_lengths_as_ens_matches_manual_calc():
+    """lengths as ENS for a non-reversible model matches a manual calculation"""
+    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)"))
+    lf.set_motif_probs(mprobs)
+    # setting arbitrary values for GN rate terms
+    init = 0.1
+    for par_name in lf.model.get_param_list():
+        lf.set_param_rule(par_name, init=init)
+        init += 0.1
+
+    len_dict = lf.get_lengths_as_ens()
+    expect = expected_number_subs(
+        mprobs, lf.get_rate_matrix_for_edge("a", calibrated=False), 1.0
+    )
+    assert_allclose(len_dict["a"], expect)
diff --git a/tests/test_parse/test_tree.py b/tests/test_parse/test_tree.py
index da1dc16c4..3e21e22b5 100644
--- a/tests/test_parse/test_tree.py
+++ b/tests/test_parse/test_tree.py
@@ -1,7 +1,5 @@
 """Unit tests for tree parsers.
 """
-import pathlib
-
 from unittest import TestCase
 
 from cogent3.core.tree import PhyloNode