def test_get_nodes_in_watersheds():
endpoints = np.array([7, 7, 7, 5, 5, 5, 7, 7, 7, 22, 5, 5, 13, 13, 13, 22, 22,
23, 13, 19, 19, 22, 22, 23, 25, 25, 25, 28, 28, 29])
combined_minimums = [{5}, {7, 13, 19, 25}, {22, 23, 28, 29}]
unique = np.array([5, 7, 13, 19, 22, 23, 25, 28, 29])
counts = np.array([5, 6, 3, 2, 3, 5, 2, 2, 1])
result_watersheds = [np.array([3, 4, 5, 10, 11]),
np.array([0, 1, 2, 6, 7, 8, 12, 13, 14, 18, 19, 20, 24, 25, 26]),
np.array([9, 15, 16, 17, 21, 22, 23, 27, 28, 29])]
watersheds = trap_analysis.get_nodes_in_watersheds(endpoints, combined_minimums)
# Sort the elements of each watershed as the compare method is more conservative.
for i in range(len(watersheds)):
watersheds[i] = np.sort(watersheds[i])
are_equal = compare_methods.compare_two_lists_of_arrays(watersheds, result_watersheds)
assert are_equal
import time
saved_files = '/home/shomea/a/anderovo/Dropbox/watershedLargeFiles/'
file_name = saved_files + 'anders_hoh.tiff'
"""
Save the watershed using pickle and numpy save.
"""
landscape = load_geotiff.get_landscape_tyrifjorden(file_name)
# Get downslope neighbors
downslope_neighbors = util.get_downslope_indices(landscape.num_of_nodes_x, landscape.num_of_nodes_y,
landscape.coordinates[:, 2])
# Get endpoints
endpoints = trap_analysis.get_node_endpoints(landscape.num_of_nodes_x, landscape.num_of_nodes_y, downslope_neighbors)
# Get minimums in each watershed
minimum_indices = np.where(downslope_neighbors == -1)[0]
minimums_in_each_watershed = sorted(trap_analysis.get_minimums_in_watersheds(minimum_indices, landscape.num_of_nodes_x,
landscape.num_of_nodes_y))
# Get indices leading to endpoints
indices_leading_to_endpoints = trap_analysis.get_indices_leading_to_endpoints(endpoints)
# Get the nodes in the watersheds. Save to file.
nodes_in_watersheds = trap_analysis.get_nodes_in_watersheds(endpoints, minimums_in_each_watershed)
cPickle.dump(nodes_in_watersheds, open('nodesInWatershedsStandard.pkl', 'wb'))
