def average_tree(fname: str, lam: float, rep: int):
with h5py.File(fname) as io:
y = io["y"][()]
grid_opt = fname.replace(".img", f"_lam{lam}.grid_opt")
if path.exists(grid_opt):
with h5py.File(grid_opt) as io:
print(grid_opt)
grid_opt = io["xgrid"][()]
else:
grid_opt = None
graph = BiAdjacent(GridGraph(*y.shape))
yvec = y.reshape(-1, order='F')
xsol = list()
for s in range(rep):
print("seed =", s)
pi = random_spanning_tree(graph, seed=s)
assert len(pi) == len(yvec)
assert pi.max() < len(pi)
assert pi.min() >= 0
x = treelas.tree_dp(y=yvec, parent=pi, lam=8 * lam, verbose=False)
xsol.append(x.reshape(*y.shape, order='F'))
xopt = np.mean(xsol, axis=0)
return xsol, xopt, grid_opt
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 test_belgium():
url = "https://www.cc.gatech.edu/dimacs10/archive/data/streets/" + \
"belgium.osm.graph.bz2"
fname = "belgium.bz2"
if not path.exists(fname):
print("downloading", url)
urlretrieve(url, fname)
idx = gio.parse_dimacs10(fname)
parent = graphidx.random_spanning_tree(idx, seed=2020)
cidx = ChildrenIndex(parent)
assert list(cidx[cidx.root]) == [188343, 205091]
def test_clique5():
edges = zip(*combinations(range(5), 2))
head, tail = map(lambda l: np.array(l, dtype=np.int32), edges)
assert (head < tail).all()
index = BiAdjacent(head, tail)
assert list(index[4]) == [0, 1, 2, 3]
parent = random_spanning_tree(index)
assert (parent >= 0).all()
assert (parent < 5).all()
assert find_root(parent) >= 0
assert (parent == [3, 4, 1, 4, 4]).all(), 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_grid(width):
graph = BiAdjacent(GridGraph(width, width))
pi = random_spanning_tree(graph, seed=0)
np.random.seed(0)
y = np.random.randn(len(pi))
x = tree_dp(y=y, parent=pi, lam=0.2)