Merge branch 'master' into single-thread-hybrid

cassiopeia/data/utilities.py

@@ -109,7 +109,7 @@ def ete3_to_networkx(tree: ete3.Tree) -> nx.DiGraph:
         if n.is_root():
             continue
 
-        g.add_edge(n.up.name, n.name)
+        g.add_edge(n.up.name, n.name, length=n.dist)
 
     return g
 

test/data_tests/cassiopeia_tree_test.py

@@ -137,6 +137,20 @@ def test_newick_to_networkx(self):
         for e in test_edges:
             self.assertIn(e, expected_edges)
 
+    def test_newick_to_networkx_with_branch_lengths(self):
+
+        nwk = "(node1:0.2,node2:0.3)node0;"
+
+        network = data_utilities.newick_to_networkx(nwk)
+
+        test_edges = network.edges(data=True)
+        expected_edges = [
+            ("node0", "node1", {"length": 0.2}),
+            ("node0", "node2", {"length": 0.3}),
+        ]
+        for e in test_edges:
+            self.assertIn(e, expected_edges)
+
     def test_newick_constructor(self):
 
         tree = cas.data.CassiopeiaTree(
@@ -181,6 +195,24 @@ def test_newick_constructor(self):
         for n in obs_nodes:
             self.assertIn(n, expected_nodes)
 
+    def test_newick_constructor_with_branch_lengths(self):
+
+        nwk = "(node1:0.2,node2:0.3)node0;"
+
+        tree = cas.data.CassiopeiaTree(tree = nwk)
+
+        network = tree.get_tree_topology()
+        test_edges = network.edges(data=True)
+        expected_edges = [
+            ("node0", "node1", {"length": 0.2}),
+            ("node0", "node2", {"length": 0.3}),
+        ]
+        for e in test_edges:
+            self.assertIn(e, expected_edges)
+
+        expected_times = {"node0": 0.0, "node1": 0.2, "node2": 0.3}
+        self.assertEqual(tree.get_times(), expected_times)
+
     def test_networkx_constructor(self):
 
         tree = cas.data.CassiopeiaTree(