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 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_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_1():
head = np.array([0, 1, 2, 3], dtype=np.int32)
tail = np.array([1, 3, 1, 2], dtype=np.int32)
index = BiAdjacent(head, tail)
assert repr(index) == "BiAdjacent[m = 4, n = 4]"
i2 = index[2]
assert len(i2) == 2
assert list(i2) == [1, 3]
assert list(index[0]) == [1]
assert list(index[1]) == [0, 3, 2]
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 print_grid(g: GridGraph,
out=sys.stdout,
width=0.3,
g_attrs=["splines=ortho"]):
from graphidx.idx import BiAdjacent
n1, n2 = g.shape
indent = " "
neighs = BiAdjacent(g)
print("graph {", file=out)
for attr in g_attrs:
print(f"{indent}{attr};", file=out)
print(f"{indent}node [shape=circle, width={width}, fixedsize=true];",
file=out)
for i1 in range(n1):
for i2 in range(n2):
i = i1 + i2 * n1
print(f'{indent}n{i} [label="{i}", pos="{i2},-{i1}!"];', file=out)
for i in range(n1 * n2):
for j in neighs[i]:
if i < j:
print(f"{indent}n{i} -- n{j};", file=out)
print("}", file=out)
def test_grid3x2_idx():
grid = GridGraph(3, 2)
graph = BiAdjacent(grid)
assert graph.num_nodes() == 6
assert graph.num_edges() == 7
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)
def square():
head = np.array([0, 0, 1, 2])
tail = np.array([1, 2, 3, 3])
return BiAdjacent(head, tail)