Benchmark node traversals

bench_traversals.py

import itertools
import sys
from time import perf_counter
from matplotlib import pyplot
from scipy import stats
import numpy as np
import tqdm

sys.path.append('/home/benj/projects/tskit/python')
import tskit
import msprime


def timeit(func):
    t1_start = perf_counter()
    func()
    return perf_counter() - t1_start


def loop(ts, order, num_trees):
    for t in itertools.islice(ts.trees(), num_trees):
        for u in t.nodes(order=order):
            pass


NUM_TREES = 20
NUM_TRIALS = 5

trees = {n: msprime.simulate(n, recombination_rate=10, random_seed=42) for n in
         [20000, 50000, 100000, 200000]}

time_fig = pyplot.figure()
time_ax = time_fig.add_subplot(1, 1, 1)

for order in ["preorder",
              "inorder",
              "postorder",
              "levelorder",
              #"breadthfirst", same as levelorder
              "timeheap",
              "timesimple"]:
    times = np.array([
        [timeit(lambda: loop(ts, order, NUM_TREES)) for i in range(NUM_TRIALS)]
        for ts in tqdm.tqdm(trees.values(), desc=f'{order} - time')])
    time_ax.errorbar(trees.keys(), np.mean(times, axis=1), yerr=stats.sem(times, axis=1), label=order)
time_ax.set_xlabel('# samples')
time_ax.set_ylabel('seconds')
time_fig.legend()
time_fig.savefig('time.png')