Create a helper function for basin recognition

This function accepts an array representing the height map of an area
and accumulates pixel count through steepest path downward.
This is a naive implementation that takes quite some time to complete
for larger maps.

diff --git a/povodi.py b/povodi.py
new file mode 100644 (file)
index 0000000..86477ef
--- /dev/null
+++ b/povodi.py
@@ -0,0 +1,44 @@
+import numpy as np
+
+def    is_noncalced_max(height_map, calced, multi_ind):
+       low_x = max(0, multi_ind[0] - 1)
+       high_x = min(height_map.shape[0] - 1, multi_ind[0] + 1)
+       low_y = max(0, multi_ind[1] - 1)
+       high_y = min(height_map.shape[1] - 1, multi_ind[1] + 1)
+       masked_subarr = np.ma.masked_array(height_map[low_x:high_x + 1, low_y:high_y + 1],
+                       mask=calced[low_x:high_x + 1, low_y:high_y + 1])
+       return (masked_subarr.max() == height_map[multi_ind])
+
+def    get_recipient(height_map, multi_ind):
+       sub_x = 1
+       sub_y = 1
+       high_x = min(height_map.shape[0] - 1, multi_ind[0] + 1)
+       low_x = multi_ind[0] - 1
+       if (low_x < 0):
+               sub_x = 0
+               low_x = 0
+       high_y = min(height_map.shape[1] - 1, multi_ind[1] + 1)
+       low_y = multi_ind[1] - 1
+       if (low_y < 0):
+               sub_y = 0
+               low_y = 0
+       sub_arr = height_map[low_x:high_x + 1, low_y:high_y + 1]
+       min_ind_small = np.unravel_index(sub_arr.argmin(), sub_arr.shape)
+       return ((min_ind_small[0] + multi_ind[0] - sub_x, min_ind_small[1] + multi_ind[1] - sub_y))
+
+def povodi(height_map):
+       height_map = np.array(height_map)
+       calced = height_map.astype('bool')
+       calced.fill(0)
+       result = height_map.astype("uint32")
+       result.fill(1)
+       while (not(calced.all())):
+               print(calced.sum() / calced.size)
+               it = np.nditer(height_map, flags=['multi_index'])
+               for el in it:
+                       if (not(calced[it.multi_index]) and is_noncalced_max(height_map, calced, it.multi_index)):
+                               calced[it.multi_index] = 1
+                               recipient_ind = get_recipient(height_map, it.multi_index)
+                               if (recipient_ind != it.multi_index):
+                                       result[recipient_ind] += result[it.multi_index]
+       return (result)