add d4/d8 option

xtopo/algos/flow_routing.py

@@ -14,14 +14,23 @@ def reset_receivers(receivers, nnodes):
 
 
 @numba.njit
-def compute_receivers_d8(receivers, dist2receivers, elevation,
+def compute_receivers(receivers, dist2receivers, elevation,
                          active_nodes,
-                         nx, ny, dx, dy):
-    # queen (D8) neighbor lookup
-    dr = np.array([-1, -1, -1, 0, 0, 1, 1, 1], dtype=np.intp)
-    dc = np.array([-1, 0, 1, -1, 1, -1, 0, 1], dtype=np.intp)
+                         nx, ny, dx, dy, directions = 8):
+
+    if directions == 8:
+        # queen (D8) neighbor lookup
+        dr = (-1, -1, -1, 0, 0, 1, 1, 1)
+        dc = (-1, 0, 1, -1, 1, -1, 0, 1)
+
+    else:
+        directions = 4
+        # king (D4) neighbor lookup
+        dr = (-1,  0, 0, 1, 0, 0, 0, 0)
+        dc = ( 0, -1, 1, 0, 0, 0, 0, 0)
+
+    length = [math.sqrt((dy * dr[i])**2 + (dx * dc[i])**2) for i in range(8)]
 
-    length = np.sqrt((dy * dr)**2 + (dx * dc)**2)
 
     tiny = np.finfo(elevation.dtype).tiny
 
@@ -34,7 +43,7 @@ def compute_receivers_d8(receivers, dist2receivers, elevation,
                 
             slope_max = tiny
 
-            for k in range(8):
+            for k in range(directions):
                 if r + dr[k] >=0 and r + dr[k]<ny and c + dc[k] >=0 and c + dc[k] <nx:
                     ineighbor = (r + dr[k]) * nx + (c + dc[k])
                     slope = (elevation[inode] - elevation[ineighbor]) / length[k]
@@ -124,7 +133,7 @@ def compute_pits(pits, outlets, active_nodes, nbasins):
 @numba.njit
 def _connect_basins(conn_basins, conn_nodes, conn_weights,
                     nbasins, basins, outlets, receivers, stack,
-                    active_nodes, elevation, nx, ny):
+                    active_nodes, elevation, nx, ny, directions = 8):
     """Connect adjacent basins together through their lowest pass.
 
     Creates an (undirected) graph of basins and their connections.
@@ -173,8 +182,8 @@ def _connect_basins(conn_basins, conn_nodes, conn_weights,
     iactive = False
 
     # king (D4) neighbor lookup
-    dr = (0, -1, 1, 0)
-    dc = (-1, 0, 0, 1)
+    dr = (0, -1, 1, 0, 1, 1, -1, -1)
+    dc = (-1, 0, 0, 1, 1, -1, 1, -1)
 
     for istack in stack:
         irec = receivers[istack]
@@ -201,7 +210,7 @@ def _connect_basins(conn_basins, conn_nodes, conn_weights,
             r = istack // nx
             c = istack % nx
 
-            for k in range(4):
+            for k in range(directions):
                 kr = r + dr[k]
                 kc = c + dc[k]
 
@@ -437,7 +446,7 @@ def _update_pits_receivers(receivers, dist2receivers, outlets,
 
 
 @numba.njit
-def _sediment_lakes(sills, basin_stack, basin_parent, basins, water_height, elevation, sediment_slope, nx, ny, dx, dy):
+def _sediment_lakes(sills, basin_stack, basin_parent, basins, water_height, elevation, sediment_slope, nx, ny, dx, dy, directions = 8):
     
     """ Parse the basins from sea to crests, filling each of the basin
     by looking at each of the basin nodes, sorted by distance to the sill
@@ -475,7 +484,7 @@ def _sediment_lakes(sills, basin_stack, basin_parent, basins, water_height, elev
             parse_begin +=1
             
 
-            for nb_dir_i in range (8):
+            for nb_dir_i in range (directions):
 
                 kr = i_node // nx + dr[nb_dir_i]
                 kc = i_node %  nx + dc[nb_dir_i]
@@ -505,7 +514,9 @@ def correct_flowrouting(receivers, dist2receivers, ndonors, donors,
                         active_nodes, elevation, nx, ny, dx, dy,
                         method='mst_linear', sediment_lakes = False, 
                         sediment_slope = 1e-10, water_height=None, 
-                        reorder = True):
+                        reorder = True,
+                        randomize = True,
+                        directions = 8):
     """Ensure that no flow is captured in sinks.
 
     If needed, update `receivers`, `dist2receivers`, `ndonors`,
@@ -527,14 +538,14 @@ def correct_flowrouting(receivers, dist2receivers, ndonors, donors,
     nconn, basin0 = _connect_basins(
         conn_basins, conn_nodes, conn_weights,
         nbasins, basins, outlets, receivers, stack,
-        active_nodes, elevation, nx, ny)
+        active_nodes, elevation, nx, ny, directions)
     
     #resize conn arrays, applying a random shuffle:
     #  In case of basins with equal height sills (flat area or 1 sized localminima),
     #  the resulting connection order in the mst will depend on the order of conn_basins
     #  Another way to see that is that MST is not unique when some edges has the same weigth,
     #  in this case we chose the most 'natural' random one
-    perm = np.random.permutation(nconn)
+    perm = np.random.permutation(nconn) if randomize else np.arange(nconn)
     conn_basins = conn_basins[perm]
     conn_nodes = conn_nodes[perm]
     conn_weights = conn_weights[perm]
@@ -567,17 +578,18 @@ def correct_flowrouting(receivers, dist2receivers, ndonors, donors,
             water_height = np.empty_like(elevation)
             
         water_height[:] = elevation[:]
-        _sediment_lakes(sills, basin_stack[:basin_stack_size], basin_parent, basins, water_height, elevation, sediment_slope, nx, ny, dx, dy)
+        _sediment_lakes(sills, basin_stack[:basin_stack_size], basin_parent, basins, water_height, elevation, sediment_slope, nx, ny, dx, dy, directions)
         
         # the water height obtained from _sediment_lakes is a very regular geometric shape.
         # we apply some ramdomness to obtin a more natural drainage pattern,
         # without adding any other local minima
-        water_height += 0.99 * sediment_slope * np.random.rand(nx*ny) * (water_height > elevation)
+        if randomize:
+            water_height += 0.99 * sediment_slope * np.random.rand(nx*ny) * (water_height > elevation)
         
         if reorder:
             reset_receivers(receivers, nnodes)
-            compute_receivers_d8(receivers, dist2receivers,
-                                 water_height, active_nodes, nx, ny, dx, dy)
+            compute_receivers(receivers, dist2receivers,
+                                 water_height, active_nodes, nx, ny, dx, dy, directions)
                 
     else:
         _update_pits_receivers(receivers, dist2receivers, outlets,
@@ -590,7 +602,10 @@ def correct_flowrouting(receivers, dist2receivers, ndonors, donors,
         
     
     
-def compute_water_level(elevation, nx, ny, dx, dy, water_height, active_nodes = None, method='mst_linear', sediment_slope = 1e-10):
+def compute_water_level(elevation, nx, ny, dx, dy, water_height,
+                        active_nodes = None, method='mst_linear',
+                        sediment_slope = 1e-10, randomize = True,
+                        directions = 8):
     
     """ Gives, in 'water_height', the elevation of water for a dem 'elevation' of size (nx, ny).
         The water level can be given a small slope 'sediment_slope' from the sill, in which case
@@ -620,8 +635,8 @@ def compute_water_level(elevation, nx, ny, dx, dy, water_height, active_nodes =
         active_nodes = mask.flatten()
     
     reset_receivers(receivers, nnodes)
-    compute_receivers_d8(receivers, dist2receivers,
-                                      elevation, active_nodes, nx, ny, dx, dy)
+    compute_receivers(receivers, dist2receivers,
+                                      elevation, active_nodes, nx, ny, dx, dy, directions)
     compute_donors(ndonors, donors, receivers, nnodes)
     compute_stack(stack, ndonors, donors, receivers, nnodes)
 
@@ -636,4 +651,8 @@ def compute_water_level(elevation, nx, ny, dx, dy, water_height, active_nodes =
                          method=method, sediment_lakes = True, 
                          sediment_slope = sediment_slope,
                          water_height = water_height,
-                         reorder = False)    
+                         reorder = False,
+                         randomize = randomize,
+                         directions = directions)
+
+