Local branching metrics fixes (#96)

* black format tests

* revise author info for pypi

* revise `local_branching` method, and add a unit test

Author: wsdewitt

Makefile

@@ -9,7 +9,7 @@ test:
 	gctree test
 
 format:
-	black gctree
+	black gctree tests
 	docformatter --in-place gctree/*.py
 
 lint:

gctree/branching_processes.py

@@ -851,8 +851,7 @@ def support(
                     compatibility_ += weights[i] if weights is not None else 1
             node.support = compatibility_ if compatibility else support
 
-    @np.errstate(all="raise")
-    def local_branching(self, tau=1, tau0=0.1):
+    def local_branching(self, tau=1, tau0=0.1, infinite_root_branch=True):
         r"""Add local branching statistics (Neher et al. 2014) as tree node
         features to the ETE tree attribute.
         After execution, all nodes will have new features ``LBI``
@@ -862,45 +861,48 @@ def local_branching(self, tau=1, tau0=0.1):
         Args:
             tau: decay timescale for exponential filter
             tau0: effective branch length for branches with zero mutations
+            infinite_root_branch: calculate assuming the root node has an infinite branch
         """
         # the fixed integral contribution for clonal cells indicated by abundance annotations
         clone_contribution = tau * (1 - np.exp(-tau0 / tau))
 
-        # post-order traversal to populate downward integral for each node
+        # post-order traversal to populate downward integrals for each node
         for node in self.tree.traverse(strategy="postorder"):
             if node.is_leaf():
-                node.add_feature(
-                    "LB_down",
-                    node.abundance * clone_contribution if node.abundance > 1 else 0,
-                )
+                node.LB_down = {
+                    node: node.abundance * clone_contribution
+                    if node.abundance > 1
+                    else 0
+                }
             else:
-                node.add_feature(
-                    "LB_down",
-                    node.abundance * clone_contribution
-                    + sum(
-                        tau * (1 - np.exp(-child.dist / tau))
-                        + np.exp(-child.dist / tau) * child.LB_down
-                        for child in node.children
-                    ),
-                )
+                node.LB_down = {node: node.abundance * clone_contribution}
+                for child in node.children:
+                    node.LB_down[child] = tau * (
+                        1 - np.exp(-child.dist / tau)
+                    ) + np.exp(-child.dist / tau) * sum(child.LB_down.values())
 
         # pre-order traversal to populate upward integral for each node
         for node in self.tree.traverse(strategy="preorder"):
             if node.is_root():
-                node.add_feature("LB_up", 0)
+                # integral corresponding to infinite branch above root node
+                node.LB_up = tau if infinite_root_branch else 0
             else:
-                node.add_feature(
-                    "LB_up",
-                    tau * (1 - np.exp(-node.dist / tau))
-                    + np.exp(-node.dist / tau) * (node.up.LB_up + node.up.LB_down),
+                node.LB_up = tau * (1 - np.exp(-node.dist / tau)) + np.exp(
+                    -node.dist / tau
+                ) * (
+                    node.up.LB_up
+                    + sum(
+                        node.up.LB_down[message]
+                        for message in node.up.LB_down
+                        if message != node
+                    )
                 )
 
-        # finally, compute LBI (LBR) as the sum (ratio) of upward and downward integrals at each node
+        # finally, compute LBI (LBR) as the sum (ratio) of downward and upward integrals at each node
         for node in self.tree.traverse():
-            node.add_feature("LBI", node.LB_down + node.LB_up)
-            node.add_feature(
-                "LBR", node.LB_down / node.LB_up if not node.is_root() else np.nan
-            )
+            node_LB_down_total = sum(node.LB_down.values())
+            node.LBI = node_LB_down_total + node.LB_up
+            node.LBR = node_LB_down_total / node.LB_up
 
 
 def _requires_dag(func):

setup.py

@@ -9,8 +9,8 @@
     name="gctree",
     version=versioneer.get_version(),
     cmdclass=versioneer.get_cmdclass(),
-    author="William DeWitt",
-    author_email="wsdewitt@gmail.com",
+    author="Matsen Group",
+    author_email="ematsen@gmail.com",
     description="phylogenetic inference of genotype-collapsed trees",
     long_description=long_description,
     long_description_content_type="text/markdown",

tests/test_likelihoods.py

@@ -149,7 +149,10 @@ def test_newlikelihoods():
 
             assert ll_isclose(oldfll, newfll)
             assert ll_isclose(oldfll, newfctreell)
-        for dagll, treell in zip(sorted(newforest._forest.weight_count(**ll_dagfuncs).elements()), sorted(ctree.ll(p, q) for ctree in newforest)):
+        for dagll, treell in zip(
+            sorted(newforest._forest.weight_count(**ll_dagfuncs).elements()),
+            sorted(ctree.ll(p, q) for ctree in newforest),
+        ):
             assert np.isclose(dagll, treell)
 
 
@@ -180,7 +183,7 @@ def test_validate_ll_genotype():
         for c in range(c_max):
             for m in range(m_max):
                 if c > 0 or m > 1:
-                    with np.errstate(all='raise'):
+                    with np.errstate(all="raise"):
                         true_res = OldCollapsedTree._ll_genotype(c, m, *params)
                         res = bp.CollapsedTree._ll_genotype(c, m, *params)
                     assert np.isclose(true_res[0], res[0])
@@ -193,5 +196,5 @@ def test_recursion_depth():
     recursion depth issues"""
     bp.CollapsedTree._ll_genotype.cache_clear()
     bp.CollapsedTree._max_ll_cache = {}
-    with np.errstate(all='raise'):
+    with np.errstate(all="raise"):
         bp.CollapsedTree._ll_genotype(2, 500, 0.4, 0.6)

tests/test_local_branching.py

@@ -0,0 +1,85 @@
+import ete3
+import gctree
+import numpy as np
+
+tree = ete3.TreeNode(name="naive", dist=0)
+tree.abundance = 0
+tree.sequence = "A"
+tree.isotype = {}
+
+child1 = ete3.TreeNode(name="seq1", dist=1)
+child1.abundance = 1
+child1.sequence = "C"
+child1.isotype = {}
+
+child2 = ete3.TreeNode(name="seq2", dist=2)
+child2.abundance = 2
+child2.sequence = "G"
+child2.isotype = {}
+
+grandchild = ete3.TreeNode(name="seq3", dist=3)
+grandchild.abundance = 3
+grandchild.sequence = "T"
+grandchild.isotype = {}
+
+tree.add_child(child1)
+tree.add_child(child2)
+child1.add_child(grandchild)
+
+ctree = gctree.CollapsedTree()
+ctree.tree = tree
+
+τ = 1
+τ0 = 0.1
+
+ctree.local_branching(tau=τ, tau0=τ0)
+
+# dummy integrated branch length for exploded clonal abundances
+clone_contribution = τ * (1 - np.exp(-τ0 / τ))
+
+LB_down = {
+    tree: {
+        tree: 0,
+        child1: τ * (1 - np.exp(-(child1.dist + grandchild.dist) / τ))
+        + np.exp(-child1.dist / τ) * child1.abundance * clone_contribution
+        + np.exp(-(child1.dist + grandchild.dist) / τ)
+        * grandchild.abundance
+        * clone_contribution,
+        child2: τ * (1 - np.exp(-child2.dist / τ))
+        + np.exp(-child2.dist / τ) * child2.abundance * clone_contribution,
+    },
+    child1: {
+        child1: child1.abundance * clone_contribution,
+        grandchild: τ * (1 - np.exp(-grandchild.dist / τ))
+        + np.exp(-grandchild.dist / τ) * grandchild.abundance * clone_contribution,
+    },
+    child2: {child2: child2.abundance * clone_contribution},
+    grandchild: {grandchild: grandchild.abundance * clone_contribution},
+}
+
+LB_up = {tree: τ}
+LB_up[child1] = τ * (1 - np.exp(-child1.dist / τ)) + np.exp(-child1.dist / τ) * (
+    LB_up[tree]
+    + sum(LB_down[tree][message] for message in LB_down[tree] if message != child1)
+)
+LB_up[child2] = τ * (1 - np.exp(-child2.dist / τ)) + np.exp(-child2.dist / τ) * (
+    LB_up[tree]
+    + sum(LB_down[tree][message] for message in LB_down[tree] if message != child2)
+)
+LB_up[grandchild] = τ * (1 - np.exp(-grandchild.dist / τ)) + np.exp(
+    -grandchild.dist / τ
+) * (
+    LB_up[child1]
+    + sum(
+        LB_down[child1][message] for message in LB_down[child1] if message != grandchild
+    )
+)
+
+LBI = {node: LB_up[node] + sum(LB_down[node].values()) for node in tree.traverse()}
+LBR = {node: sum(LB_down[node].values()) / LB_up[node] for node in tree.traverse()}
+
+for node in ctree.tree.traverse():
+    assert LB_up[node] == node.LB_up
+    assert LB_down[node] == node.LB_down
+    assert LBI[node] == node.LBI
+    assert LBR[node] == node.LBR