From 0f7c330b663d83ffb24b00e937ceae405f3c8d25 Mon Sep 17 00:00:00 2001
From: colganwi <colgan.william@gmail.com>
Date: Mon, 30 Oct 2023 11:42:24 -0400
Subject: [PATCH] newicks with branch lengths

---
 cassiopeia/data/utilities.py            |  2 +-
 test/data_tests/cassiopeia_tree_test.py | 32 +++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/cassiopeia/data/utilities.py b/cassiopeia/data/utilities.py
index 3aae157d..35cde4f4 100755
--- a/cassiopeia/data/utilities.py
+++ b/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
 
diff --git a/test/data_tests/cassiopeia_tree_test.py b/test/data_tests/cassiopeia_tree_test.py
index cc741e65..840576c6 100755
--- a/test/data_tests/cassiopeia_tree_test.py
+++ b/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(