1import numpy as np
2from skimage.graph._graph import pixel_graph, central_pixel
3
4mask = np.array([[1, 0, 0], [0, 1, 1], [0, 1, 0]], dtype=bool)
5image = np.random.default_rng().random(mask.shape)
6
7
8def test_small_graph():
9    g, n = pixel_graph(mask, connectivity=2)
10    assert g.shape == (4, 4)
11    assert len(g.data) == 8
12    np.testing.assert_allclose(np.unique(g.data), [1, np.sqrt(2)])
13    np.testing.assert_array_equal(n, [0, 4, 5, 7])
14
15
16def test_central_pixel():
17    g, n = pixel_graph(mask, connectivity=2)
18    px, ds = central_pixel(g, n, shape=mask.shape)
19    np.testing.assert_array_equal(px, (1, 1))
20    s2 = np.sqrt(2)
21    np.testing.assert_allclose(ds, [s2*3 + 2, s2 + 2, s2*2 + 2, s2*2 + 2])
22
23    # test raveled coordinate
24    px, _ = central_pixel(g, n)
25    assert px == 4
26
27    # test no nodes given
28    px, _ = central_pixel(g)
29    assert px == 1
30
31
32def test_edge_function():
33    def edge_func(values_src, values_dst, distances):
34        return np.abs(values_src - values_dst) + distances
35
36    g, n = pixel_graph(
37            image, mask=mask, connectivity=2, edge_function=edge_func
38            )
39    s2 = np.sqrt(2)
40    np.testing.assert_allclose(g[0, 1], np.abs(image[0, 0] - image[1, 1]) + s2)
41    np.testing.assert_allclose(g[1, 2], np.abs(image[1, 1] - image[1, 2]) + 1)
42    np.testing.assert_array_equal(n, [0, 4, 5, 7])
43
44
45def test_default_edge_func():
46    g, n = pixel_graph(image, spacing=np.array([0.78, 0.78]))
47    num_edges = len(g.data) // 2  # each edge appears in both directions
48    assert num_edges == 12  # lattice in a (3, 3) grid
49    np.testing.assert_almost_equal(
50            g[0, 1], 0.78 * np.abs(image[0, 0] - image[0, 1])
51            )
52    np.testing.assert_array_equal(n, np.arange(image.size))
53
54
55if __name__ == '__main__':
56    test_edge_function()
57