Optimization by claude

Author: constantinpape

include/bioimage_cpp/graph/grid_features.hxx

@@ -2,17 +2,14 @@
 
 #include "bioimage_cpp/array_view.hxx"
 #include "bioimage_cpp/detail/edge_hash.hxx"
-#include "bioimage_cpp/detail/grid.hxx"
 #include "bioimage_cpp/graph/grid_graph.hxx"
 
-#include <algorithm>
 #include <array>
 #include <cmath>
 #include <cstddef>
 #include <cstdint>
 #include <stdexcept>
 #include <string>
-#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -21,13 +18,17 @@ namespace bioimage_cpp::graph {
 namespace detail_grid_features {
 
 template <std::size_t D>
-std::vector<std::ptrdiff_t> graph_shape_vector(const GridGraph<D> &graph) {
-    std::vector<std::ptrdiff_t> shape(D);
-    const auto &grid_shape = graph.shape();
+std::array<std::ptrdiff_t, D> offset_to_array(
+    const std::vector<std::ptrdiff_t> &offset
+) {
+    if (offset.size() != D) {
+        throw std::invalid_argument("each offset must have length matching graph ndim");
+    }
+    std::array<std::ptrdiff_t, D> result{};
     for (std::size_t axis = 0; axis < D; ++axis) {
-        shape[axis] = static_cast<std::ptrdiff_t>(grid_shape[axis]);
+        result[axis] = offset[axis];
     }
-    return shape;
+    return result;
 }
 
 inline std::ptrdiff_t l1_norm(const std::vector<std::ptrdiff_t> &offset) {
@@ -38,6 +39,16 @@ inline std::ptrdiff_t l1_norm(const std::vector<std::ptrdiff_t> &offset) {
     return result;
 }
 
+template <std::size_t D>
+std::ptrdiff_t l1_norm(const std::array<std::ptrdiff_t, D> &offset) {
+    std::ptrdiff_t result = 0;
+    for (std::size_t axis = 0; axis < D; ++axis) {
+        const auto value = offset[axis];
+        result += value < 0 ? -value : value;
+    }
+    return result;
+}
+
 template <std::size_t D>
 void require_graph_shape(
     const GridGraph<D> &graph,
@@ -87,30 +98,34 @@ void require_affinity_shape(
     }
 }
 
+// Two long-range offsets `a` and `b` produce the same set of edges iff
+// `b = -a` — each offset, restricted to its valid source region, is
+// internally injective, and `edge_key` canonicalizes endpoint order, so
+// only sign-flipped offset pairs collide. Checking pairs is O(n_offsets^2)
+// — a small constant in practice — and avoids sorting the full lifted
+// edge list at the end of `grid_affinity_features_with_lifted`.
 template <std::size_t D>
-std::uint64_t local_edge_id(
-    const GridGraph<D> &graph,
-    const std::uint64_t node,
-    const std::vector<std::ptrdiff_t> &offset
+void check_no_negated_long_range_offsets(
+    const std::vector<std::vector<std::ptrdiff_t>> &offsets
 ) {
-    std::size_t axis = D;
-    std::ptrdiff_t step = 0;
-    for (std::size_t candidate = 0; candidate < D; ++candidate) {
-        if (offset[candidate] != 0) {
-            axis = candidate;
-            step = offset[candidate];
-            break;
+    for (std::size_t i = 0; i < offsets.size(); ++i) {
+        if (l1_norm(offsets[i]) <= 1) continue;
+        for (std::size_t j = i + 1; j < offsets.size(); ++j) {
+            if (l1_norm(offsets[j]) <= 1) continue;
+            bool negated = true;
+            for (std::size_t axis = 0; axis < D; ++axis) {
+                if (offsets[i][axis] != -offsets[j][axis]) {
+                    negated = false;
+                    break;
+                }
+            }
+            if (negated) {
+                throw std::invalid_argument(
+                    "offsets produce duplicate long-range grid edges"
+                );
+            }
         }
     }
-    if (axis == D || (step != 1 && step != -1) || l1_norm(offset) != 1) {
-        throw std::invalid_argument("local offsets must have L1 norm 1");
-    }
-
-    auto pivot = graph.coordinates(node);
-    if (step < 0) {
-        --pivot[axis];
-    }
-    return graph.edge_id(axis, pivot);
 }
 
 template <std::size_t D>
@@ -126,7 +141,7 @@ std::array<std::uint64_t, D> uint64_strides_from_shape(
 }
 
 template <std::size_t D>
-std::size_t local_offset_axis(const std::vector<std::ptrdiff_t> &offset) {
+std::size_t local_offset_axis(const std::array<std::ptrdiff_t, D> &offset) {
     for (std::size_t axis = 0; axis < D; ++axis) {
         if (offset[axis] != 0) {
             return axis;
@@ -135,10 +150,40 @@ std::size_t local_offset_axis(const std::vector<std::ptrdiff_t> &offset) {
     return D;
 }
 
+// Template-recursive enumeration of (source_node, target_node, coordinate)
+// triples for a fixed grid offset. Flat node ids are computed incrementally
+// by addition only. With `Axis` known at compile time, the compiler unrolls
+// the entire nest for small D and inlines `callback`.
+template <std::size_t D, std::size_t Axis, class Callback>
+void for_each_valid_offset_link_impl(
+    const std::array<std::uint64_t, D> &begin,
+    const std::array<std::uint64_t, D> &end,
+    const std::array<std::uint64_t, D> &strides,
+    const std::int64_t delta,
+    std::array<std::uint64_t, D> &coordinate,
+    const std::uint64_t node,
+    Callback &&callback
+) {
+    if constexpr (Axis == D) {
+        const auto target =
+            static_cast<std::uint64_t>(static_cast<std::int64_t>(node) + delta);
+        callback(node, target, coordinate);
+    } else {
+        std::uint64_t axis_node = node + begin[Axis] * strides[Axis];
+        for (std::uint64_t coord = begin[Axis]; coord < end[Axis]; ++coord) {
+            coordinate[Axis] = coord;
+            for_each_valid_offset_link_impl<D, Axis + 1>(
+                begin, end, strides, delta, coordinate, axis_node, callback
+            );
+            axis_node += strides[Axis];
+        }
+    }
+}
+
 template <std::size_t D, class Callback>
 void for_each_valid_offset_link(
     const GridGraph<D> &graph,
-    const std::vector<std::ptrdiff_t> &offset,
+    const std::array<std::ptrdiff_t, D> &offset,
     Callback &&callback
 ) {
     const auto &shape = graph.shape();
@@ -158,42 +203,28 @@ void for_each_valid_offset_link(
         delta += static_cast<std::int64_t>(axis_offset) *
                  static_cast<std::int64_t>(strides[axis]);
     }
-
     std::array<std::uint64_t, D> coordinate{};
-    auto recurse = [&](auto &&self, const std::size_t axis, const std::uint64_t node) -> void {
-        if (axis == D) {
-            const auto target =
-                static_cast<std::uint64_t>(static_cast<std::int64_t>(node) + delta);
-            callback(node, target, coordinate);
-            return;
-        }
-        auto axis_node = node + begin[axis] * strides[axis];
-        for (std::uint64_t coord = begin[axis]; coord < end[axis]; ++coord) {
-            coordinate[axis] = coord;
-            self(self, axis + 1, axis_node);
-            axis_node += strides[axis];
-        }
-    };
-    recurse(recurse, 0, 0);
+    for_each_valid_offset_link_impl<D, 0>(
+        begin, end, strides, delta, coordinate, 0, callback
+    );
 }
 
+// Compute the local edge id from a source coordinate. The per-axis pivot
+// adjustment for negative offsets is hoisted into `pivot_adjustment` — a
+// constant per (offset, axis) the caller pre-computes once, so the inner
+// loop is a branch-free dot product.
 template <std::size_t D>
 std::uint64_t local_edge_id_from_source_coordinate(
-    const GridGraph<D> &graph,
-    const std::vector<std::ptrdiff_t> &offset,
+    const std::uint64_t edge_offset_base,
     const std::array<std::uint64_t, D> &source_coordinate,
-    const std::size_t axis,
-    const std::array<std::uint64_t, D> &edge_strides
+    const std::array<std::uint64_t, D> &edge_strides,
+    const std::uint64_t pivot_adjustment
 ) {
     std::uint64_t local_edge = 0;
-    for (std::size_t coord_axis = 0; coord_axis < D; ++coord_axis) {
-        auto pivot = source_coordinate[coord_axis];
-        if (coord_axis == axis && offset[axis] < 0) {
-            --pivot;
-        }
-        local_edge += pivot * edge_strides[coord_axis];
+    for (std::size_t axis = 0; axis < D; ++axis) {
+        local_edge += source_coordinate[axis] * edge_strides[axis];
     }
-    return graph.edge_offset(axis) + local_edge;
+    return edge_offset_base + local_edge - pivot_adjustment;
 }
 
 } // namespace detail_grid_features
@@ -210,8 +241,9 @@ void grid_boundary_features(
         throw std::invalid_argument("out shape must be (number_of_edges,)");
     }
 
-    for (std::uint64_t edge = 0; edge < graph.number_of_edges(); ++edge) {
-        const auto uv = graph.uv(edge);
+    const auto &uv_ids = graph.uv_ids();
+    for (std::size_t edge = 0; edge < uv_ids.size(); ++edge) {
+        const auto &uv = uv_ids[edge];
         out.data[edge] = 0.5 * (boundary_map.data[uv.first] + boundary_map.data[uv.second]);
     }
 }
@@ -245,24 +277,35 @@ void grid_local_affinity_features(
     }
 
     for (std::size_t channel = 0; channel < offsets.size(); ++channel) {
-        const auto axis = detail_grid_features::local_offset_axis<D>(offsets[channel]);
+        const auto offset_array =
+            detail_grid_features::offset_to_array<D>(offsets[channel]);
+        const auto axis = detail_grid_features::local_offset_axis<D>(offset_array);
         auto edge_shape = graph.shape();
         --edge_shape[axis];
-        const auto edge_strides = detail_grid_features::uint64_strides_from_shape(edge_shape);
-        const auto channel_offset = static_cast<std::uint64_t>(channel) * number_of_nodes;
+        const auto edge_strides =
+            detail_grid_features::uint64_strides_from_shape(edge_shape);
+        const auto edge_offset_base = graph.edge_offset(axis);
+        const std::uint64_t pivot_adjustment =
+            offset_array[axis] < 0 ? edge_strides[axis] : 0;
+        const auto channel_offset =
+            static_cast<std::uint64_t>(channel) * number_of_nodes;
         detail_grid_features::for_each_valid_offset_link<D>(
             graph,
-            offsets[channel],
+            offset_array,
             [&](const std::uint64_t node, const std::uint64_t, const auto &coordinate) {
-            const auto edge = detail_grid_features::local_edge_id_from_source_coordinate<D>(
-                graph, offsets[channel], coordinate, axis, edge_strides
-            );
-            if (valid_edges.data[edge] != 0) {
-                throw std::invalid_argument("offsets produce duplicate local grid edges");
+                const auto edge =
+                    detail_grid_features::local_edge_id_from_source_coordinate<D>(
+                        edge_offset_base, coordinate, edge_strides, pivot_adjustment
+                    );
+                if (valid_edges.data[edge] != 0) {
+                    throw std::invalid_argument(
+                        "offsets produce duplicate local grid edges"
+                    );
+                }
+                weights.data[edge] = affinities.data[channel_offset + node];
+                valid_edges.data[edge] = 1;
             }
-            weights.data[edge] = affinities.data[channel_offset + node];
-            valid_edges.data[edge] = 1;
-        });
+        );
     }
 }
 
@@ -281,6 +324,7 @@ GridLiftedAffinityFeatures grid_affinity_features_with_lifted(
     const std::vector<std::vector<std::ptrdiff_t>> &offsets
 ) {
     detail_grid_features::require_affinity_shape(graph, affinities, offsets);
+    detail_grid_features::check_no_negated_long_range_offsets<D>(offsets);
 
     const auto number_of_nodes = graph.number_of_nodes();
     GridLiftedAffinityFeatures result;
@@ -298,47 +342,52 @@ GridLiftedAffinityFeatures grid_affinity_features_with_lifted(
     result.lifted_offset_ids.reserve(static_cast<std::size_t>(lifted_capacity));
 
     for (std::size_t channel = 0; channel < offsets.size(); ++channel) {
-        const auto l1 = detail_grid_features::l1_norm(offsets[channel]);
-        const auto channel_offset = static_cast<std::uint64_t>(channel) * number_of_nodes;
+        const auto offset_array =
+            detail_grid_features::offset_to_array<D>(offsets[channel]);
+        const auto l1 = detail_grid_features::l1_norm(offset_array);
+        const auto channel_offset =
+            static_cast<std::uint64_t>(channel) * number_of_nodes;
         if (l1 == 1) {
-            const auto axis = detail_grid_features::local_offset_axis<D>(offsets[channel]);
+            const auto axis = detail_grid_features::local_offset_axis<D>(offset_array);
             auto edge_shape = graph.shape();
             --edge_shape[axis];
-            const auto edge_strides = detail_grid_features::uint64_strides_from_shape(edge_shape);
+            const auto edge_strides =
+                detail_grid_features::uint64_strides_from_shape(edge_shape);
+            const auto edge_offset_base = graph.edge_offset(axis);
+            const std::uint64_t pivot_adjustment =
+                offset_array[axis] < 0 ? edge_strides[axis] : 0;
             detail_grid_features::for_each_valid_offset_link<D>(
                 graph,
-                offsets[channel],
+                offset_array,
                 [&](const std::uint64_t node, const std::uint64_t, const auto &coordinate) {
-                const auto edge = detail_grid_features::local_edge_id_from_source_coordinate<D>(
-                    graph, offsets[channel], coordinate, axis, edge_strides
-                );
-                if (result.valid_local_edges[static_cast<std::size_t>(edge)] != 0) {
-                    throw std::invalid_argument("offsets produce duplicate local grid edges");
+                    const auto edge =
+                        detail_grid_features::local_edge_id_from_source_coordinate<D>(
+                            edge_offset_base, coordinate, edge_strides, pivot_adjustment
+                        );
+                    if (result.valid_local_edges[static_cast<std::size_t>(edge)] != 0) {
+                        throw std::invalid_argument(
+                            "offsets produce duplicate local grid edges"
+                        );
+                    }
+                    result.local_weights[static_cast<std::size_t>(edge)] =
+                        affinities.data[channel_offset + node];
+                    result.valid_local_edges[static_cast<std::size_t>(edge)] = 1;
                 }
-                result.local_weights[static_cast<std::size_t>(edge)] =
-                    affinities.data[channel_offset + node];
-                result.valid_local_edges[static_cast<std::size_t>(edge)] = 1;
-            });
+            );
         } else {
             detail_grid_features::for_each_valid_offset_link<D>(
                 graph,
-                offsets[channel],
+                offset_array,
                 [&](const std::uint64_t node, const std::uint64_t target, const auto &) {
-                const auto uv = bioimage_cpp::detail::edge_key(node, target);
-                result.lifted_uvs.push_back(uv);
-                result.lifted_weights.push_back(affinities.data[channel_offset + node]);
-                result.lifted_offset_ids.push_back(static_cast<std::uint64_t>(channel));
-            });
+                    const auto uv = bioimage_cpp::detail::edge_key(node, target);
+                    result.lifted_uvs.push_back(uv);
+                    result.lifted_weights.push_back(affinities.data[channel_offset + node]);
+                    result.lifted_offset_ids.push_back(static_cast<std::uint64_t>(channel));
+                }
+            );
         }
     }
 
-    auto sorted_lifted_uvs = result.lifted_uvs;
-    std::sort(sorted_lifted_uvs.begin(), sorted_lifted_uvs.end());
-    if (std::adjacent_find(sorted_lifted_uvs.begin(), sorted_lifted_uvs.end()) !=
-        sorted_lifted_uvs.end()) {
-        throw std::invalid_argument("offsets produce duplicate long-range grid edges");
-    }
-
     return result;
 }