def test_children3():
t = Tree.from_prufer([1, 2])
assert t.root == 3
assert (t.parent == [1, 2, 3, 3]).all()
list(t.children[0]) == []
for i in range(1, 4):
assert list(t.children[i]) == [i - 1]
def test_grid_5x3():
graph = BiAdjacent(GridGraph(5, 3))
assert graph.num_nodes() == 5 * 3
pi = random_spanning_tree(graph, seed=0)
assert all(pi == [1, 2, 3, 8, 9, 6, 1, 8, 9, 9, 5, 6, 7, 14, 9])
t = Tree(pi)
assert t.root == 9
assert t.degree.sum() == 2 * (t.n - 1)
def test_grid_2x2():
graph = BiAdjacent(GridGraph(2, 2))
pi = random_spanning_tree(graph, seed=0)
if os.getenv("show"):
Tree(pi).show()
np.random.seed(0)
y = np.random.randn(len(pi))
x = tree_dp(y=y, parent=pi, lam=0.2)
def test_grid_3x3():
graph = BiAdjacent(GridGraph(3, 3))
pi = random_spanning_tree(graph, seed=0)
if os.getenv("show"):
Tree(pi).show()
assert all(pi == [3, 0, 1, 6, 7, 8, 6, 6, 7])
np.random.seed(0)
y = np.random.randn(len(pi))
x = tree_dp(y=y, parent=pi, lam=0.2)
def load_tree(treeh5):
"""Return root, parent from a tree instance stored in hdf5"""
import h5py
with h5py.File(treeh5, "r") as io:
parent = io["parent"][()]
if parent.max() >= len(parent):
parent -= 1
return Tree(parent)
def test_grid_200x200():
graph = BiAdjacent(GridGraph(200, 200))
pi = random_spanning_tree(graph, seed=0)
t = Tree(pi)
cidx = ChildrenIndex(pi)
assert t.degree.sum() == 2 * (t.n - 1)
def test_tree_from_prufer5():
p = [a for i in range(5) for a in [i] * 2]
assert len(p) == 10
t = Tree.from_prufer(p)
assert t.n == 12
def test_tree_from_prufer3():
p = np.array([2, 2, 3], np.int32)
t = Tree.from_prufer(p)
assert t.root == 4
assert (t.parent == [2, 2, 3, 4, 4]).all()
def test_tree_from_prufer3():
p = np.array([2, 2, 3], np.int32)
t = Tree.from_prufer(p)
assert t.root == 4
assert (t.parent == [2, 2, 3, 4, 4]).all()
def test_tree_from_prufer5():
p = [a for i in range(5) for a in [i] * 2]
assert len(p) == 10
t = Tree.from_prufer(p)
assert t.n == 12
if __name__ == '__main__':
import argparse
p = argparse.ArgumentParser(description=__doc__)
p.add_argument('-s',
'--seed',
type=int,
default=2018,
help='Random seed [default=2018]')
args = p.parse_args()
children_spec = np.array([2, 2, 3])
t = Tree.generate(children_spec, seed=args.seed)
print(f"Generated tree with {t.n} nodes")
t.show()
def test_rtree(n=5, seed=2018, eps=1e-14):
"""Random tree creation"""
t = Tree.random(n, seed=seed)
assert t.n == n
def test_degree3():
t = Tree.from_prufer([1, 2])
assert all(t.degree == [1, 2, 2, 1]), str(t.degree)