parallel utility, fixed multi-outlet performance issue

.vscode/settings.json

@@ -4,6 +4,8 @@
         "neighbormatrix.c": "cpp",
         "stdbool.h": "c",
         "returncodes.h": "c",
-        "string.h": "c"
+        "string.h": "c",
+        "__locale": "c",
+        "bitset": "c"
     }
 }

PyOCN/__init__.py

@@ -1,16 +1,14 @@
 from .ocn import OCN
-from .plotting import plot_ocn_as_dag, plot_positional_digraph, plot_ocn_raster
-from .utils import net_type_to_dag, simulated_annealing_schedule, unwrap_digraph, get_subwatersheds, assign_subwatersheds
+from ._version import __version__
+
+from . import utils
+from . import plotting
+
 
 __all__ = [
     "OCN", 
-    "flowgrid", 
-    "plot_ocn_as_dag", 
-    "plot_positional_digraph", 
-    "plot_ocn_raster",
-    "net_type_to_dag",
-    "simulated_annealing_schedule",
-    "unwrap_digraph",
-    "get_subwatersheds",
-    "assign_subwatersheds",
-]
+    "utils",
+    "plotting",
+    "__version__",
+]
+

PyOCN/_flowgrid_convert.py

@@ -1,6 +1,3 @@
-#TODO: allow users to provide multiple DAGs that partition a space.
-#TODO: allow edge-wrapping (toroidal grids).
-#TODO: implement "every edge vertex is a root" option
 """
 Functions for converting between NetworkX graphs and FlowGrid_C structures.
 """
@@ -246,9 +243,6 @@ def direction_bit(pos1, pos2):
     p_c_graph.contents.resolution = float(resolution)
     p_c_graph.contents.nroots = len([n for n in G.nodes if G.out_degree(n) == 0])
     p_c_graph.contents.wrap = wrap
-
-    if p_c_graph.contents.nroots > 1:
-        warnings.warn(f"FlowGrid has {p_c_graph.contents.nroots} root nodes (nodes with no downstream). This will slow down certain operations.")
 
     # do not set energy
 
@@ -291,3 +285,10 @@ def validate_flowgrid(c_graph:_bindings.FlowGrid_C, verbose:bool=False) -> bool|
     except Exception as e:  # _digraph_to_flowgrid_c will destroy p_c_graph on failure
         return str(e)
     return "Graph is valid."
+
+__all__ = [
+    "to_digraph",
+    "from_digraph",
+    "validate_digraph",
+    "validate_flowgrid",
+]

PyOCN/_libocn_bindings.py

@@ -115,6 +115,14 @@ class FlowGrid_C(Structure):
 libocn.fg_cart_to_lin.argtypes = [CartPair_C, CartPair_C]
 libocn.fg_cart_to_lin.restype = linidx_t
 
+# CartPair fg_lin_to_cart(linidx_t a, CartPair dims);
+libocn.fg_lin_to_cart.argtypes = [linidx_t, CartPair_C]
+libocn.fg_lin_to_cart.restype = CartPair_C
+
+# Status fg_clockhand_to_lin(linidx_t *a_down, linidx_t a, clockhand_t down, CartPair dims, bool wrap)
+libocn.fg_clockhand_to_lin.argtypes = [POINTER(linidx_t), linidx_t, clockhand_t, CartPair_C, c_bool]
+libocn.fg_clockhand_to_lin.restype = Status
+
 # Status fg_get_lin(Vertex *out, FlowGrid *G, linidx_t a);
 libocn.fg_get_lin.argtypes = [POINTER(Vertex_C), POINTER(FlowGrid_C), linidx_t]
 libocn.fg_get_lin.restype = Status
@@ -135,6 +143,18 @@ class FlowGrid_C(Structure):
 libocn.fg_destroy.argtypes = [POINTER(FlowGrid_C)]
 libocn.fg_destroy.restype = Status
 
+# Status fg_find_upstream_neighbors(linidx_t upstream_indices[8], linidx_t *nupstream, FlowGrid *G, linidx_t a);
+libocn.fg_find_upstream_neighbors.argtypes = [linidx_t * 8, POINTER(linidx_t), POINTER(FlowGrid_C), linidx_t]
+libocn.fg_find_upstream_neighbors.restype = Status
+
+# Status fg_dfs_iterative(linidx_t upstream_indices[], linidx_t *nupstream, linidx_t idx_stack[], FlowGrid *G, linidx_t a);
+libocn.fg_dfs_iterative.argtypes = [POINTER(linidx_t), POINTER(linidx_t), POINTER(linidx_t), POINTER(FlowGrid_C), linidx_t]
+libocn.fg_dfs_iterative.restype = Status
+
+# Status flow_downstream(linidx_t downstream_indices[], linidx_t *ndownstream, FlowGrid *G, linidx_t a)
+libocn.flow_downstream.argtypes = [POINTER(linidx_t), POINTER(linidx_t), POINTER(FlowGrid_C), linidx_t]
+libocn.flow_downstream.restype = Status
+
 ##############################
 #     OCN.H EQUIVALENTS      #
 ##############################
@@ -143,12 +163,12 @@ class FlowGrid_C(Structure):
 libocn.ocn_compute_energy.argtypes = [POINTER(FlowGrid_C), c_double]
 libocn.ocn_compute_energy.restype = c_double
 
-# Status ocn_single_erosion_event(FlowGrid *G, double gamma, double temperature);
-libocn.ocn_single_erosion_event.argtypes = [POINTER(FlowGrid_C), c_double, c_double]
+# Status ocn_single_erosion_event(FlowGrid *G, double gamma, double temperature, bool calculate_full_energy);
+libocn.ocn_single_erosion_event.argtypes = [POINTER(FlowGrid_C), c_double, c_double, c_bool]
 libocn.ocn_single_erosion_event.restype = Status
 
-# Status ocn_outer_ocn_loop(FlowGrid *G, uint32_t niterations, double gamma, double *annealing_schedule);
-libocn.ocn_outer_ocn_loop.argtypes = [POINTER(FlowGrid_C), c_uint32, c_double, POINTER(c_double)]
+# Status ocn_outer_ocn_loop(FlowGrid *G, uint32_t niterations, double gamma, double *annealing_schedule, bool calculate_full_energy);
+libocn.ocn_outer_ocn_loop.argtypes = [POINTER(FlowGrid_C), c_uint32, c_double, POINTER(c_double), c_bool]
 libocn.ocn_outer_ocn_loop.restype = Status
 
 

PyOCN/_version.py

@@ -0,0 +1 @@
+__version__ = "1.4.20251029"

PyOCN/c_src/flowgrid.c

@@ -302,6 +302,82 @@ Status fg_check_for_cycles(FlowGrid *G, linidx_t a, uint8_t check_number){
     return SUCCESS;  // found root successfully, no cycles found
 }
 
+Status fg_find_upstream_neighbors(linidx_t upstream_indices[8], linidx_t *nupstream, FlowGrid *G, linidx_t a){
+    if (G == NULL || G->vertices == NULL || upstream_indices == NULL || nupstream == NULL) return NULL_POINTER_ERROR;
+
+    Status code;
+    Vertex vert;
+    code = fg_get_lin(&vert, G, a);
+    if (code != SUCCESS) return code;
+
+    // get the clockhand direction of all edges
+    *nupstream = 0;
+    for (clockhand_t dir = 0; dir < 8; dir++){
+        if ((vert.edges & (1u << dir)) && (vert.downstream != dir)){  // if there's an edge in this direction
+            linidx_t a_up;
+            code = fg_clockhand_to_lin(&a_up, a, dir, G->dims, G->wrap);  // get the upstream vertex index
+            if (code != SUCCESS) return code;
+            upstream_indices[*nupstream] = a_up;
+            (*nupstream)++;
+        }
+    }
+    return SUCCESS;
+}
+
+Status fg_dfs_iterative(linidx_t upstream_indices[], linidx_t *nupstream, linidx_t idx_stack[], FlowGrid *G, linidx_t a){
+    if (G == NULL || G->vertices == NULL || upstream_indices == NULL || nupstream == NULL || idx_stack == NULL) return NULL_POINTER_ERROR;
+    Status code;
+
+    // Seed idx_stack with immediate upstream of a
+    *nupstream = 0;
+    linidx_t local_upstream[8];
+    linidx_t nlocal_upstream = 0;
+    code = fg_find_upstream_neighbors(local_upstream, &nlocal_upstream, G, a);
+    if (code != SUCCESS) return code;
+    if (nlocal_upstream == 0) return SUCCESS;  // no upstream vertices found
+
+    linidx_t total = (linidx_t)G->dims.row * (linidx_t)G->dims.col;
+    if (nlocal_upstream >= total) return OOB_ERROR;
+    memcpy(idx_stack, local_upstream, nlocal_upstream * sizeof(linidx_t));
+
+    if (*nupstream >= total || total == 0) return OOB_ERROR;
+    linidx_t max_out = total - 1;
+    linidx_t top = nlocal_upstream;
+    while (top > 0){
+        // pop from stack
+        linidx_t u = idx_stack[--top];  
+        if (u >= total) return OOB_ERROR;
+
+        // Append to output buffer
+        if (*nupstream >= max_out) return OOB_ERROR;
+        upstream_indices[(*nupstream)++] = u;
+
+        // Push u's upstream neighbors
+        nlocal_upstream = 0;
+        code = fg_find_upstream_neighbors(local_upstream, &nlocal_upstream, G, u);
+        if (code != SUCCESS) return code;
+        if (top + nlocal_upstream >= total) return OOB_ERROR;
+        memcpy(idx_stack + top, local_upstream, nlocal_upstream * sizeof(linidx_t));
+        top += nlocal_upstream;
+    }
+    return SUCCESS;
+}
+
+Status flow_downstream(linidx_t downstream_indices[], linidx_t *ndownstream, FlowGrid *G, linidx_t a){
+    Vertex vert;
+    *ndownstream = 0;
+    do {
+        Status code = fg_get_lin(&vert, G, a);
+        if (code == OOB_ERROR) return OOB_ERROR;
+        downstream_indices[(*ndownstream)++] = a;
+
+        a = vert.adown;
+    } while (vert.downstream != IS_ROOT);
+
+    return SUCCESS;  
+}
+
+
 // vibe-coded display function
 const char E_ARROW = '-';
 const char S_ARROW = '|';

PyOCN/c_src/flowgrid.h

@@ -161,6 +161,36 @@ Status fg_change_vertex_outflow(FlowGrid *G, linidx_t a, clockhand_t down_new);
  */
 Status fg_check_for_cycles(FlowGrid *G, linidx_t a, uint8_t check_number);
 
+/**
+ * @brief Find all upstream neighbors of a given vertex.
+ * @param upstream_indices Array to store the linear indices of upstream neighbors.
+ * @param nupstream Pointer to store the number of upstream neighbors found.
+ * @param G Pointer to the FlowGrid.
+ * @param a The linear index of the vertex to find upstream neighbors for.
+ * @return Status code indicating success or failure
+ */
+Status fg_find_upstream_neighbors(linidx_t upstream_indices[8], linidx_t *nupstream, FlowGrid *G, linidx_t a);
+
+/**
+ * @brief Perform an iterative depth-first search to find all upstream vertices from a given starting vertex.
+ * @param upstream_indices Array to store the linear indices of all upstream vertices found. Must be preallocated to hold enough indices (ie G.dims.row * G.dims.col).
+ * @param nupstream Pointer to store the total number of upstream vertices found.
+ * @param idx_stack Preallocated stack array for DFS traversal. Must be large enough to hold all potential upstream vertices (ie G.dims.row * G.dims.col).
+ * @param G Pointer to the FlowGrid.
+ * @param a The linear index of the starting vertex.
+ * @return Status code indicating success or failure
+ */
+Status fg_dfs_iterative(linidx_t upstream_indices[], linidx_t *nupstream, linidx_t idx_stack[], FlowGrid *G, linidx_t a);
+
+/**
+ * @brief Follow the downstream path from a given vertex, recording each vertex along the path. Stops when a root node is reached. Includes root node. Does not include starting node.
+ * @param downstream_indices Array to store the linear indices of downstream vertices. Must be preallocated to hold enough indices (ie G.dims.row * G.dims.col).
+ * @param ndownstream Pointer to store the number of downstream vertices found.
+ * @param G Pointer to the FlowGrid.
+ * @param a The linear index of the starting vertex.
+ * @return Status code indicating success or failure
+ */
+Status flow_downstream(linidx_t downstream_indices[], linidx_t *ndownstream, FlowGrid *G, linidx_t a);
 /**
  * @brief Display the flowgrid in the terminal using ASCII or UTF-8 characters.
  * @param G Pointer to the FlowGrid to display.

PyOCN/c_src/ocn.c

@@ -37,7 +37,7 @@ static inline bool simulate_annealing(double energy_new, double energy_old, doub
  * @param a The linear index of the starting vertex.
  * @return Status code indicating success or failure
  */
-static inline Status update_drained_area(FlowGrid *G, drainedarea_t da_inc, linidx_t a){
+static Status update_drained_area(FlowGrid *G, drainedarea_t da_inc, linidx_t a){
     Vertex vert;
     do {
         Status code = fg_get_lin(&vert, G, a);
@@ -51,14 +51,18 @@ static inline Status update_drained_area(FlowGrid *G, drainedarea_t da_inc, lini
     return SUCCESS;  
 }
 
-double ocn_compute_energy(FlowGrid *G, double gamma){
+static inline double compute_energy(FlowGrid *G, double gamma){
     double energy = 0.0;
     for (linidx_t i = 0; i < (linidx_t)G->dims.row * (linidx_t)G->dims.col; i++){
         energy += pow(G->vertices[i].drained_area, gamma);
     }
     return energy;
 }
 
+double ocn_compute_energy(FlowGrid *G, double gamma){
+    return compute_energy(G, gamma);
+}
+
 /**
  * @brief Update the energy of the flowgrid along a single downstream path from a given vertex. Unsafe.
  * This function only works correctly if there is a single root in the flowgrid.
@@ -69,7 +73,7 @@ double ocn_compute_energy(FlowGrid *G, double gamma){
  * @param gamma The exponent used in the energy calculation.
  * @return Status code indicating success or failure
  */
-Status update_energy_single_root(FlowGrid *G, drainedarea_t da_inc, linidx_t a, double gamma){
+static Status update_energy_single_root(FlowGrid *G, drainedarea_t da_inc, linidx_t a, double gamma){
     Vertex vert;
     double energy_old = 0.0;
     double energy_new = 0.0;
@@ -86,7 +90,7 @@ Status update_energy_single_root(FlowGrid *G, drainedarea_t da_inc, linidx_t a,
     return SUCCESS;
 }
 
-Status ocn_single_erosion_event(FlowGrid *G, double gamma, double temperature){
+Status ocn_single_erosion_event(FlowGrid *G, double gamma, double temperature, bool calculate_full_energy){
     Status code;
 
     Vertex vert;
@@ -154,52 +158,37 @@ Status ocn_single_erosion_event(FlowGrid *G, double gamma, double temperature){
 
 
     mh_eval:
-    /*
-    TODO: PERFORMANCE ISSUE:
-    This function is supposed to update the energy of the flowgrid G after a 
-    change in drained area along the path starting at vertex a.
-
-    Simple but inefficient fix (current): recompute the *entire* energy of the flowgrid from scratch
-    each time this function is called.
-
-    More complex fix: find the set of all upstream vertices that flow into a and compute
-    their summed contribution to the energy. Pass this value (sum of (da^gamma) for all
-    upstream vertices) into this function, instead of just passing da_inc.
-    */
-    if ((G->nroots > 1) && (gamma < 1.0)){
-        // energy_old = ocn_compute_energy(G, gamma);  // recompute energy from scratch
+
+    if (calculate_full_energy){
+        energy_old = compute_energy(G, gamma);  // recompute energy from scratch
         update_drained_area(G, -da_inc, a_down_old);  // remove drainage from old path
         update_drained_area(G, da_inc, a_down_new);  // add drainage to new path
-        energy_new = ocn_compute_energy(G, gamma);  // recompute energy from scratch
+        energy_new = compute_energy(G, gamma);  // recompute energy from scratch
         if (simulate_annealing(energy_new, energy_old, temperature, &G->rng)){
             G->energy = energy_new;
             return SUCCESS;
         }
-        // reject swap: undo everything and try again
-        update_drained_area(G, da_inc, a_down_old);  // add removed drainage back to old path
-        update_drained_area(G, -da_inc, a_down_new);  // remove added drainage from new path
-        fg_change_vertex_outflow(G, a, down_old);  // undo the outflow change
-    } else {  // if there's only one root, we can use a more efficient method
+    } else {
         update_energy_single_root(G, -da_inc, a_down_old, gamma);  // remove drainage from old path and update energy
         update_energy_single_root(G, da_inc, a_down_new, gamma);  // add drainage to new path and update energy
         energy_new = G->energy;
         if (simulate_annealing(energy_new, energy_old, temperature, &G->rng)){
             return SUCCESS;
         }
-        // reject swap: undo everything and try again
-        update_energy_single_root(G, da_inc, a_down_old, gamma);  // add removed drainage back to old path and update energy
-        update_energy_single_root(G, -da_inc, a_down_new, gamma);  // remove added drainage from new path and update energy
-        fg_change_vertex_outflow(G, a, down_old);  // undo the outflow change
     }
 
-
+    // reject swap: undo everything and try again
+    update_drained_area(G, da_inc, a_down_old);  // add removed drainage back to old path and update energy
+    update_drained_area(G, -da_inc, a_down_new);  // remove added drainage from new path and update energy
+    fg_change_vertex_outflow(G, a, down_old);  // undo the outflow change
+    G->energy = energy_old;
     return EROSION_FAILURE;  // if we reach here, we failed to find a valid swap in many, many tries
 }
 
-Status ocn_outer_ocn_loop(FlowGrid *G, uint32_t niterations, double gamma, double *annealing_schedule){
+Status ocn_outer_ocn_loop(FlowGrid *G, uint32_t niterations, double gamma, double *annealing_schedule, bool calculate_full_energy){
     Status code;
     for (uint32_t i = 0; i < niterations; i++){
-        code = ocn_single_erosion_event(G, gamma, annealing_schedule[i]);
+        code = ocn_single_erosion_event(G, gamma, annealing_schedule[i], calculate_full_energy);
         if ((code != SUCCESS) && (code != EROSION_FAILURE)) return code;
     }
     return SUCCESS;

PyOCN/c_src/ocn.h

@@ -4,8 +4,6 @@
  * @brief Header file for OCN optimization.
  */
 
-//#TODO I'm worried that our method of choosing a new vertex if the last one is invalid introduces bias. Either choose a random direction to walk, or just try every vertex in random order.
-
 #ifndef OCN_H
 #define OCN_H
 
@@ -31,9 +29,10 @@ double ocn_compute_energy(FlowGrid *G, double gamma);
  * @param G Pointer to the FlowGrid.
  * @param gamma The exponent used in the energy calculation.
  * @param temperature The temperature parameter for the Metropolis-Hastings acceptance criterion.
+ * @param calculate_full_energy If true, the full energy of the graph is recalculated when considering the proposed change. If false, a more efficient incremental update is used. full_energy_recalc will be slower, but avoid accumulated numerical errors over many iterations.
  * @return Status code indicating success or failure.
  */
-Status ocn_single_erosion_event(FlowGrid *G, double gamma, double temperature);
+Status ocn_single_erosion_event(FlowGrid *G, double gamma, double temperature, bool calculate_full_energy);
 
 /**
  * @brief Perform multiple erosion events on the flowgrid.
@@ -43,8 +42,9 @@ Status ocn_single_erosion_event(FlowGrid *G, double gamma, double temperature);
  * @param gamma The exponent used in the energy calculation.
  * @param annealing_schedule An array of temperatures (ranging from 0-1) to use for each iteration. Length must be at least niterations.
  * @param wrap If true, allows wrapping around the edges of the grid (toroidal). If false, no wrapping is applied.
+ * @param calculate_full_energy If true, the full energy of the graph is recalculated after each proposed change. If false, a more efficient incremental update is used. full_energy_recalc will be slower, but avoid accumulated numerical errors over many iterations. 
  * @return Status code indicating success or failure
  */
-Status ocn_outer_ocn_loop(FlowGrid *G, uint32_t niterations, double gamma, double *annealing_schedule);
+Status ocn_outer_ocn_loop(FlowGrid *G, uint32_t niterations, double gamma, double *annealing_schedule, bool calculate_full_energy);
 
 #endif // OCN_H