diff --git a/tree_performance/benchmark.py b/tree_performance/benchmark.py index 4ab8c51..49fb7d7 100644 --- a/tree_performance/benchmark.py +++ b/tree_performance/benchmark.py @@ -5,8 +5,8 @@ import sys import os import pathlib -import benchmarks.numba -import benchmarks.numba_vectorised +import benchmarks.numba_recursive +import benchmarks.numba_recursive_vectorised import benchmarks.pythran import pandas as pd import numpy as np diff --git a/tree_performance/benchmarks/numba_iterative.py b/tree_performance/benchmarks/numba_iterative.py new file mode 100644 index 0000000..356fa5d --- /dev/null +++ b/tree_performance/benchmarks/numba_iterative.py @@ -0,0 +1,74 @@ +import tskit +import util +import functools +import numba +import numpy as np +import msprime + +VECTORISED = False + +@numba.njit() +def _iterative_hartigan_parsimony(genotypes, samples, left_child, parent, postorder, preorder, root): + num_nodes = len(postorder) + num_alleles = np.max(genotypes) + 1 + optimal_set = np.zeros((num_nodes, num_alleles), dtype=np.int8) + # Simple version assuming non missing data and one root + for allele, u in zip(genotypes, samples): + optimal_set[u, allele] = 1 + + allele_count = np.zeros((num_nodes, num_alleles), dtype=np.int32) + for node_j in postorder: + if left_child[node_j] != tskit.NULL: + max_allele_count = 0 + for allele_k in range(num_alleles): + if allele_count[node_j, allele_k] > max_allele_count: + max_allele_count = allele_count[node_j, allele_k] + for allele_k in range(num_alleles): + if allele_count[node_j, allele_k] == max_allele_count: + optimal_set[node_j, allele_k] = 1 + for allele_k in range(num_alleles): + allele_count[parent[node_j], allele_k] += optimal_set[node_j, allele_k] + + anc_index, max_val = -1, -1 + for i in range(len(optimal_set[root])): + if optimal_set[root, i] > max_val: + anc_index, max_val = i, optimal_set[root, i] + + state = np.zeros((num_nodes), dtype=np.int32) + state[:] = anc_index + mutations = 0 + for node_j in preorder: + state[node_j] = state[parent[node_j]] + if optimal_set[node_j, state[node_j]] == 0: + maxval, argmax = -1, -1 + for k in range(num_alleles): + if optimal_set[node_j, k] > maxval: + maxval = optimal_set[node_j, k] + argmax = k + state[node_j] = argmax + mutations += 1 + return mutations + +def iterative_hartigan_parsimony(tree, genotypes, alleles): + samples = tree.tree_sequence.samples() + left_child = tree.left_child_array + parent = tree.parent_array + postorder = tree.postorder() + preorder = tree.preorder() + return _iterative_hartigan_parsimony(genotypes, samples, left_child, parent, postorder, preorder, tree.root) + +def run(ts_path, max_sites): + ts = tskit.load(ts_path) + assert ts.num_trees == 1 + tree = ts.first() + return util.benchmark_python( + ts, + functools.partial(iterative_hartigan_parsimony, tree), + "py_numba_iterative", + max_sites=max_sites, + ) + +def warmup(): + ts = msprime.sim_ancestry(100, sequence_length=100000, random_seed=43) + genotypes = np.zeros(ts.num_samples, dtype=np.int8) + iterative_hartigan_parsimony(ts.first(), genotypes, ["0"]) \ No newline at end of file diff --git a/tree_performance/benchmarks/numba.py b/tree_performance/benchmarks/numba_recursive.py similarity index 98% rename from tree_performance/benchmarks/numba.py rename to tree_performance/benchmarks/numba_recursive.py index 1a4afe3..43e719c 100644 --- a/tree_performance/benchmarks/numba.py +++ b/tree_performance/benchmarks/numba_recursive.py @@ -75,7 +75,7 @@ def run(ts_path, max_sites): return util.benchmark_python( ts, functools.partial(numba_hartigan_parsimony, tree), - "py_numba", + "py_numba_recursive", max_sites=max_sites, ) diff --git a/tree_performance/benchmarks/numba_vectorised.py b/tree_performance/benchmarks/numba_recursive_vectorised.py similarity index 100% rename from tree_performance/benchmarks/numba_vectorised.py rename to tree_performance/benchmarks/numba_recursive_vectorised.py