newick_play.py

import time

import numpy as np
import numba
import tskit

@numba.njit(cache=True)
def decode(s):
    stack = np.zeros(10000, dtype=np.int32)
    stack_p = np.int32(0)
    need_to_write_popped = False
    node_num = np.int32(0)
    parents = np.full((3000000), -1, dtype=np.int32)
    children = np.full((3000000), -1, dtype=np.int32)
    node_depth = np.zeros(3000000, dtype=np.int32)
    edge_count = np.int32(0)

    for c in s:
        if c == 32: #" "
            continue
        elif c == 40: #"("
            stack_p += 1
            stack[stack_p] = node_num
            node_num+=1
        elif c == 41: #")"
            if need_to_write_popped:
                parents[edge_count] = stack[stack_p]
                children[edge_count] = stack[stack_p + 1]
                node_depth[stack[stack_p + 1]] = stack_p
            else:
                parents[edge_count] = stack[stack_p]
                children[edge_count] = node_num
                node_depth[node_num] = stack_p
                node_num += 1
            edge_count += 1
            stack_p-=1
            need_to_write_popped = True
        elif c == 44: #","
            if need_to_write_popped:
                parents[edge_count] = stack[stack_p]
                children[edge_count] = stack[stack_p + 1]
                node_depth[stack[stack_p + 1]] = stack_p
                need_to_write_popped = False
            else:
                parents[edge_count] = stack[stack_p]
                children[edge_count] = node_num
                node_depth[node_num] = stack_p
                node_num += 1
            edge_count += 1
        else: #In label
            pass
    return node_num, children[:node_num-1], parents[:node_num-1], node_depth[:node_num]

def from_newick(string):
    str_arr = np.frombuffer(string.encode("utf-8"), dtype=np.uint8)
    num_nodes, children, parents, node_depth = decode(str_arr)
    max_depth = np.max(node_depth)
    node_depth = (-node_depth) + max_depth
    tc = tskit.TableCollection(1)
    tc.nodes.set_columns(flags=np.zeros(num_nodes,dtype=np.uint32), time=node_depth)
    tc.edges.set_columns(
        left=np.zeros_like(children,dtype=np.uint32),
        right=np.ones_like(children,dtype=np.uint32),
        child=children,
        parent=parents
    )
    tc.sort()
    return tc.tree_sequence()


with open("big.newick","r") as f:
    string = f.read()

t = time.perf_counter()
from_newick(string)
print(time.perf_counter()-t)