Repository: biocore/improved-octo-waddle (language: python)

Issue:
newick writing support
@kwcantrell, sorry to bother -- able to once over this one as well?

Note I just added GH actions and getting that working in master. Once thats good I'll merge the commits here


Discussion / hints:
Ugh, bit array isn't coming through on actions. Unfortunately I'm out of time today to poke. WIll fix as soon as I can. Code works locally at least... :)
@kwcantrell, okay I think I fixed this finally... able to once over the newick writing pieces?

Fix patch:
diff --git a/bp/_bp.pxd b/bp/_bp.pxd
index da6127b..81efdfb 100644
--- a/bp/_bp.pxd
+++ b/bp/_bp.pxd
@@ -5,6 +5,7 @@ from bp._ba cimport BIT_ARRAY
 
 ctypedef np.npy_intp SIZE_t
 ctypedef np.npy_uint32 UINT32_t
+ctypedef np.npy_int32 INT32_t
 ctypedef np.npy_float64 DOUBLE_t
 ctypedef np.npy_uint8 BOOL_t
 
@@ -34,6 +35,8 @@ cdef class BP:
         SIZE_t[:] _k_index_1
         np.ndarray _names
         np.ndarray _lengths
+        np.ndarray _edges
+        np.ndarray _edge_lookup
         mM _rmm
         SIZE_t size
 
@@ -57,13 +60,16 @@ cdef class BP:
     cpdef SIZE_t root(self) nogil
     cdef int scan_block_forward(self, int i, int k, int b, int d) nogil
     cdef int scan_block_backward(self, int i, int k, int b, int d) nogil
-   
+    cdef void _set_edges(self, np.ndarray[INT32_t, ndim=1] edges)
+
     # TODO: evalute down the road what methods should be cdef. There is a 
     # performance cost for cpdef, so for high use functions, it may make sense
     # to punt down to cdef.
     # http://notes-on-cython.readthedocs.io/en/latest/fibo_speed.html
     cpdef inline unicode name(self, SIZE_t i)
     cpdef inline DOUBLE_t length(self, SIZE_t i)
+    cpdef inline INT32_t edge(self, SIZE_t i)
+    cpdef SIZE_t edge_from_number(self, INT32_t n)
     cpdef SIZE_t rmq(self, SIZE_t i, SIZE_t j) nogil
     cpdef SIZE_t rMq(self, SIZE_t i, SIZE_t j) nogil
     cpdef SIZE_t postorderselect(self, SIZE_t k) nogil
diff --git a/bp/_bp.pyx b/bp/_bp.pyx
index a4d3c02..19aa0b7 100644
--- a/bp/_bp.pyx
+++ b/bp/_bp.pyx
@@ -25,9 +25,11 @@ np.import_array()
 cdef extern from "limits.h":
     int INT_MAX
 
+
 DOUBLE = np.float64
 SIZE = np.intp
 BOOL = np.uint8
+INT32 = np.int32
 
 
 cdef inline int min(int a, int b) nogil:
@@ -184,7 +186,8 @@ cdef class BP:
 
     def __cinit__(self, np.ndarray[BOOL_t, ndim=1] B, 
                   np.ndarray[DOUBLE_t, ndim=1] lengths=None,
-                  np.ndarray[object, ndim=1] names=None):
+                  np.ndarray[object, ndim=1] names=None,
+                  np.ndarray[INT32_t, ndim=1] edges=None):
         cdef SIZE_t i
         cdef SIZE_t size
         cdef SIZE_t[:] _e_index
@@ -194,6 +197,8 @@ cdef class BP:
         cdef SIZE_t[:] _r_index_1
         cdef np.ndarray[object, ndim=1] _names
         cdef np.ndarray[DOUBLE_t, ndim=1] _lengths
+        cdef np.ndarray[INT32_t, ndim=1] _edges
+        cdef np.ndarray[SIZE_t, ndim=1] _edge_lookup
 
         # the tree is only valid if it is balanaced
         assert B.sum() == (float(B.size) / 2)
@@ -213,6 +218,12 @@ cdef class BP:
         else:
             self._lengths = np.zeros(self.B.size, dtype=DOUBLE)
 
+        if edges is not None:
+            self._set_edges(edges)
+        else:
+            self._edges = np.full(self.B.size, 0, dtype=INT32)
+            self._edge_lookup = None
+
         # construct a rank index. These operations are performed frequently,
         # and easy to cache at a relatively minor memory expense
         #TODO: leverage rmm tree, and calculate rank on the fly
@@ -253,6 +264,28 @@ cdef class BP:
 
     def set_lengths(self, np.ndarray[DOUBLE_t, ndim=1] lengths):
         self._lengths = lengths
+   
+    cdef void _set_edges(self, np.ndarray[INT32_t, ndim=1] edges):
+        cdef:
+            int i, n
+            INT32_t edge
+            np.ndarray[SIZE_t, ndim=1] _edge_lookup
+            np.ndarray[BOOL_t, ndim=1] b
+
+        b = self.B
+        n = b.size
+       
+        _edge_lookup = np.full(n, 0, dtype=SIZE)
+        for i in range(n):
+            if b[i] == 1:
+                edge = edges[i]
+                _edge_lookup[edge] = i
+        
+        self._edge_lookup = _edge_lookup
+        self._edges = edges
+
+    def set_edges(self, np.ndarray[INT32_t, ndim=1] edges):
+        self._set_edges(edges)
 
     cpdef inline unicode name(self, SIZE_t i):
         return self._names[i]
@@ -260,6 +293,12 @@ cdef class BP:
     cpdef inline DOUBLE_t length(self, SIZE_t i):
         return self._lengths[i]
 
+    cpdef inline INT32_t edge(self, SIZE_t i):
+        return self._edges[i]
+
+    cpdef SIZE_t edge_from_number(self, INT32_t n):
+        return self._edge_lookup[n]
+
     cdef inline SIZE_t rank(self, SIZE_t t, SIZE_t i) nogil:
         """Determine the rank order of the ith bit t
 
diff --git a/bp/_conv.pyx b/bp/_conv.pyx
index f32cc48..ef99c44 100644
--- a/bp/_conv.pyx
+++ b/bp/_conv.pyx
@@ -27,6 +27,7 @@ def to_skbio_treenode(BP bp):
         node_idx = bp.preorderselect(i)
         nodes[i].name = bp.name(node_idx)
         nodes[i].length = bp.length(node_idx)
+        nodes[i].edge_num = bp.edge(node_idx)
 
         if node_idx != bp.root():
             # preorder starts at 1 annoyingly
diff --git a/bp/_io.pyx b/bp/_io.pyx
index 09d917a..3022e2d 100644
--- a/bp/_io.pyx
+++ b/bp/_io.pyx
@@ -9,31 +9,69 @@ cimport cython
 np.import_array()
 
 
+cdef inline np.double_t length_from_edge(unicode token):
+    cdef:
+        Py_ssize_t split_idx
+
+    # 0.12345{0123} -> 0.12345
+    split_idx = token.find('{')
+    if split_idx == -1:
+        return np.double(token)
+    else:
+        return np.double(token[:split_idx])
+
+
+cdef inline np.int32_t number_from_edge(unicode token):
+    cdef:
+        Py_ssize_t split_idx
+
+    # 0.12345{0123} -> 0123
+    split_idx = token.find('{')
+    if split_idx == -1:
+        return 0
+    else:
+        return np.int32(token[split_idx + 1:-1])
+
+
 cdef void _set_node_metadata(np.uint32_t ptr, unicode token,
                              np.ndarray[object, ndim=1] names, 
-                             np.ndarray[np.double_t, ndim=1] lengths):
+                             np.ndarray[np.double_t, ndim=1] lengths,
+                             np.ndarray[np.int32_t, ndim=1] edges):
     """Inplace update of names and lengths given token details"""
     cdef:
         np.double_t length
+        np.int32_t edge
         Py_ssize_t split_idx, i
-        unicode name
+        unicode name, token_parsed
 
     name = None
     length = 0.0
+    edge = 0
 
+    # NOTE: there is likely some fat to trim in this method. we do a lot
+    # of work per token, we could probably do that work smarter. as is,
+    # the changes to support edge numbers increase parsing ~20%, which
+    # is annoying but probably not a critical
     if token[0] == u':':
-        length = np.double(token[1:])
-    elif ':' in token:
+        token_parsed = token[1:]
+        length = length_from_edge(token_parsed)
+        edge = number_from_edge(token_parsed)
+    elif u':' in token:
         split_idx = token.rfind(':')
         name = token[:split_idx]
-        length = np.double(token[split_idx + 1:])
-        name = name.strip("'")
+        token_parsed = token[split_idx + 1:]
+        length = length_from_edge(token_parsed)
+        edge = number_from_edge(token_parsed)
+        name = name.strip("'").strip()
+    elif u'{' in token:
+        # strip as " {123}" is valid?
+        edge = np.int32(token.strip()[1:-1])
     else:
-        name = token.replace("'", "").replace('"', "")
-        pass
+        name = token.replace("'", "").replace('"', "").strip()
 
     names[ptr] = name
     lengths[ptr] = length
+    edges[ptr] = edge
 
 
 def write_newick(BP tree, object output):
@@ -90,12 +128,14 @@ cpdef parse_newick(unicode data):
         unicode token, last_token
         np.ndarray[object, ndim=1] names
         np.ndarray[np.double_t, ndim=1] lengths
+        np.ndarray[np.int32_t, ndim=1] edges
 
     datalen = len(data)
     topology = _newick_to_bp(data)
 
     names = np.full(len(topology.B), None, dtype=object)
     lengths = np.zeros(len(topology.B), dtype=np.double)
+    edges = np.full(len(topology.B), 0, dtype=np.int32) 
 
     ptr = 0
     token_ptr = _ctoken(data, datalen, 0)
@@ -120,7 +160,7 @@ cpdef parse_newick(unicode data):
             lag = 0
 
             open_ptr = topology.open(ptr)
-            _set_node_metadata(open_ptr, token, names, lengths)
+            _set_node_metadata(open_ptr, token, names, lengths, edges)
 
             if topology.isleaf(ptr):
                 ptr += 2
@@ -134,6 +174,8 @@ cpdef parse_newick(unicode data):
 
     topology.set_names(names)
     topology.set_lengths(lengths)
+    topology.set_edges(edges)
+
     return topology
 
 


Test patch:
diff --git a/bp/tests/test_bp.py b/bp/tests/test_bp.py
index ac317b4..a0b89e8 100644
--- a/bp/tests/test_bp.py
+++ b/bp/tests/test_bp.py
@@ -15,6 +15,12 @@
 import numpy.testing as npt
 
 from bp import BP, parse_newick
+import bp.tests.test_bp_cy as tbc
+
+
+for name in dir(tbc):
+    if name.startswith('test_'):
+        getattr(tbc, name)()
 
 
 class BPTests(TestCase):
diff --git a/bp/tests/test_conv.py b/bp/tests/test_conv.py
index ae69e5c..cfcd0a2 100644
--- a/bp/tests/test_conv.py
+++ b/bp/tests/test_conv.py
@@ -15,6 +15,20 @@ def setUp(self):
         self.bp = parse_newick(self.tstr)
         self.sktn = skbio.TreeNode.read(StringIO(self.tstr))
 
+    def test_to_skbio_treenode_with_edge_numbers(self):
+        # from https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0031009
+        # but without edge labels
+        # 0 1 2 3 4 5 6 7 8 9
+        # 1 1 1 0 1 0 0 1 0 0
+        in_ = '((A:.01{0}, B:.01{1})D:.01{3}, C:.01{4}) {5};'
+        obs = parse_newick(in_)
+        obs_sk = to_skbio_treenode(obs)
+        self.assertEqual(obs_sk.find('A').edge_num, 0)
+        self.assertEqual(obs_sk.find('B').edge_num, 1)
+        self.assertEqual(obs_sk.find('D').edge_num, 3)
+        self.assertEqual(obs_sk.find('C').edge_num, 4)
+        self.assertEqual(obs_sk.edge_num, 5)
+
     def test_to_skbio_treenode(self):
         obs = to_skbio_treenode(self.bp)
         for o, e in zip(obs.traverse(), self.sktn.traverse()):
diff --git a/bp/tests/test_io.py b/bp/tests/test_io.py
index 01a3705..d74de10 100644
--- a/bp/tests/test_io.py
+++ b/bp/tests/test_io.py
@@ -8,6 +8,29 @@
 
 
 class NewickTests(TestCase):
+    def test_parse_newick_simple_edge_numbers(self):
+        # from https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0031009
+        # but without edge labels
+        # 0 1 2 3 4 5 6 7 8 9
+        # 1 1 1 0 1 0 0 1 0 0
+        in_ = '((A:.01{0}, B:.01{1})D:.01{3}, C:.01{4}) {5};'
+        exp_sk = '((A:.01, B:.01)D:.01, C:.01);'  # skbio doesn't know about edge numbers
+        obs = parse_newick(in_)
+        obs_sk = to_skbio_treenode(obs)
+        exp_sk = skbio.TreeNode.read([exp_sk])
+        self.assertEqual(obs_sk.compare_rfd(exp_sk), 0)
+
+        self.assertEqual(obs.edge(2), 0)
+        self.assertEqual(obs.edge(4), 1)
+        self.assertEqual(obs.edge(1), 3)
+        self.assertEqual(obs.edge(7), 4)
+        self.assertEqual(obs.edge(0), 5)
+        self.assertEqual(obs.edge_from_number(0), 2)
+        self.assertEqual(obs.edge_from_number(1), 4)
+        self.assertEqual(obs.edge_from_number(3), 1)
+        self.assertEqual(obs.edge_from_number(4), 7)
+        self.assertEqual(obs.edge_from_number(5), 0)
+
     def _compare_newick(self, obs, exp):
         a = skbio.TreeNode.read([obs])
         b = skbio.TreeNode.read([exp])