B1 Balance index (Python)

[jeromekelleher]

Add a definition of the B1 balance index following the pattern established for Sackin index in #2246. Requires #2249

See #2245 for information.

cc @jeremyguez

[jeremyguez]

I wrote two definitions of B1 (first one is using list comprehension). Any thoughts @jeromekelleher?

def b1_index_definition(tree):
    return (sum
        (1/max([tree.path_length(n, leaf)
        for leaf in tree.leaves(n)])
        for n in tree.nodes()
        if n != tree.root and tree.is_internal(n))
    )

def b1_index_definition2(tree):
    total = 0
    for n in tree.nodes():
        if n != tree.root and tree.is_internal(n):
            depth_list = []
            for leaf in tree.leaves(n):
                depth_list.append(tree.path_length(n, leaf))
            total += 1/max(depth_list)
    return(total)

Do you have ideas for the method implementation? Maybe by going postorder we can have something more efficient?

[jeromekelleher]

I really like the list comprehension version @jeremyguez, very neat.

You're right, there is a very clean and efficient implementation of this using a postorder traversal. Do you want to have a go at figuring it out? It took me a while to see it and it is really nice, so I don't want to spoil it for you (unless you want me to!) 😄

[jeremyguez]

Thanks!

Yes, I'll give it a try! Here's what I got:

def b1_index(self):
    depth_max = np.zeros(self.tree_sequence.num_nodes, dtype=np.int32)
    total = 0
    for u in self.nodes(order="postorder"):
        if self.is_internal(u) and u != self.root:
            depth_max[u] = max([depth_max[v] for v in self.children(u)])+1
            total += 1/depth_max[u]
        elif self.is_leaf(u):
            depth_max[u] = 0    
    return total

[jeromekelleher]

Very nice! That's basically what I got:

def b1_index(tree):
    max_path_length = np.zeros(tree.tree_sequence.num_nodes, dtype=int)
    total = 0
    for u in tree.postorder():
        if tree.parent(u) != tskit.NULL and tree.is_internal(u):
            max_path_length[u] = 1 + max(max_path_length[v] for v in tree.children(u))
            total += 1 / max_path_length[u]
    return total

The only real difference is that I'm testing to see if u is a root rather than the root, so that this will work for multirooted trees.

[jeremyguez]

Yes, I forgot about multirooted trees!
And the elif test for leaves was really not necessary 😁

[jeromekelleher]

Closed in #2281