def test_get_watersheds_basic(): num_of_cols = 3 num_of_rows = 3 heights = np.array([0, 1, 2, 1, 2, 3, 2, 3, 1]) result_watersheds = [np.array([0, 1, 2, 3, 4, 6]), np.array([5, 7, 8])] watersheds = trap_analysis.get_watersheds(heights, num_of_cols, num_of_rows) # 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
def test_get_watersheds_advanced(): num_of_cols = 6 num_of_rows = 5 heights = np.array([5, 7, 8, 7, 6, 0, 7, 2, 10, 10, 7, 6, 7, 2, 4, 5, 5, 4, 7, 7, 3.9, 4, 0, 0, 6, 5, 4, 4, 0, 0]) 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_watersheds(heights, num_of_cols, num_of_rows) # 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