def test_connector_add_downstream():
segment_1 = ChannelSegment([0, 1])
segment_2 = ChannelSegment([2, 0])
connector = ChannelSegmentConnector(segment_1)
connector.add(segment_2)
assert connector.root is segment_2
assert connector.root.count_segments(direction="upstream") == 1
assert connector.root.downstream is None
def test_connector_add_upstream():
segment = ChannelSegment([0, 1])
connector = ChannelSegmentConnector(segment)
assert connector.root is segment
assert len(connector.orphans) == 0
connector.add(ChannelSegment([1, 2]))
assert connector.root is segment
assert connector.root.count_segments(direction="upstream") == 1
assert connector.root.downstream is None
def test_create_xy_of_node_one_segement(nodes):
grid = RasterModelGrid((16, 64), xy_spacing=(2.0, 3.0))
network = ChannelSegmentConnector(nodes)
xy_of_node = create_xy_of_node(network.root, grid)
assert np.allclose(xy_of_node[:, 0], grid.x_of_node[nodes])
assert np.allclose(xy_of_node[:, 1], grid.y_of_node[nodes])
def test_reindex_network_nodes():
root = ChannelSegmentConnector([10, 11, 12], [12, 13], [12, 14], [14, 15]).root
reindex_network_nodes(root)
assert list(root.nodes) == [0, 1, 2]
assert list(root.upstream[0].nodes) == [2, 3]
assert list(root.upstream[1].nodes) == [2, 4]
assert list(root.upstream[1].upstream[0].nodes) == [4, 5]
def test_create_xy_of_node_with_branch():
grid = RasterModelGrid((3, 4), xy_spacing=(2.0, 3.0))
network = ChannelSegmentConnector([4, 5], [5, 2, 3], [5, 10, 11])
xy_of_node = create_xy_of_node(network.root, grid)
assert np.allclose(xy_of_node,
[[0, 3], [2, 3], [4, 0], [6, 0], [4, 6], [6, 6]])
def test_get_node_fields_one_field():
grid = RasterModelGrid((3, 4))
grid.at_node["foo"] = np.arange(12) * 10
network = ChannelSegmentConnector([0, 5], [5, 6, 7], [5, 9])
fields = get_node_fields(network.root, grid)
assert list(fields) == ["foo"]
assert_array_equal(fields["foo"], [0, 50, 60, 70, 90])
def test_xy_of_node_if_not_network_root():
grid = RasterModelGrid((3, 4), xy_spacing=(2.0, 3.0))
network = ChannelSegmentConnector([4, 5], [5, 2], [5, 10, 11], [2, 3], [2, 7])
base = network.root.upstream[0]
xy_of_node = create_xy_of_node(base, grid)
assert np.allclose(xy_of_node, [[4, 0], [6, 0], [6, 3]])
def test_get_node_fields_ignore_non_node_fields():
grid = RasterModelGrid((3, 4))
grid.add_empty("foo", at="node")
grid.add_empty("bar", at="node")
grid.add_empty("baz", at="link")
network = ChannelSegmentConnector([0, 5], [5, 6, 7], [5, 9])
fields = get_node_fields(network.root, grid, include="*")
assert sorted(list(fields)) == ["bar", "foo"]
def test_get_node_fields_exclude():
grid = RasterModelGrid((3, 4))
grid.add_empty("foo", at="node")
grid.add_empty("bar", at="node")
grid.add_empty("baz", at="node")
network = ChannelSegmentConnector([0, 5], [5, 6, 7], [5, 9])
expected = get_node_fields(network.root, grid, include="at_node:b*")
actual = get_node_fields(network.root, grid, exclude="at_node:f*")
assert actual.keys() == expected.keys()
for name in actual:
assert_array_equal(actual[name], expected[name])
def test_get_node_fields_include():
grid = RasterModelGrid((3, 4))
grid.at_node["foo"] = np.arange(12) * 10
grid.at_node["bar"] = np.arange(12) * 100
grid.at_node["baz"] = np.arange(12) * 1000
network = ChannelSegmentConnector([0, 5], [5, 6, 7], [5, 9])
fields = get_node_fields(network.root, grid, include="at_node:f*")
assert list(fields) == ["foo"]
assert_array_equal(fields["foo"], [0, 50, 60, 70, 90])
fields = get_node_fields(network.root, grid, include="at_node:b*")
assert sorted(list(fields)) == ["bar", "baz"]
assert_array_equal(fields["bar"], [0, 500, 600, 700, 900])
assert_array_equal(fields["baz"], [0, 5000, 6000, 7000, 9000])
def test_connector_add_orphan():
segment_1 = ChannelSegment([0, 1])
segment_2 = ChannelSegment([2, 3])
connector = ChannelSegmentConnector(segment_1)
connector.add(segment_2)
assert connector.root is segment_1
assert connector.root.count_segments(direction="upstream") == 0
assert connector.root.downstream is None
assert len(connector.orphans) == 1
assert connector.orphans == (segment_2,)
connector.add(ChannelSegment([1, 2]))
assert connector.root.count_segments(direction="upstream") == 2
assert connector.orphans == ()
def test_create_network_links():
root = ChannelSegmentConnector([0, 1, 2], [2, 3], [2, 4], [4, 5]).root
links = create_network_links(root)
assert links == [(0, 1), (1, 2), (2, 3), (2, 4), (4, 5)]