Add edge projection to grid graph (#49)

Author: constantinpape

MIGRATION_GUIDE.md

@@ -899,6 +899,18 @@ Important differences:
   `grid_affinity_features_with_lifted(...)`, which returns local graph weights
   plus explicit long-range `uv_ids` and weights for lifted multicut or mutex
   watershed style workflows.
+- Nifty's `projectEdgeIdsToPixels` and `projectEdgeIdsToPixelsWithOffsets` map
+  to `graph.project_edge_ids_to_pixels()` and
+  `graph.project_edge_ids_to_pixels_with_offsets(offsets, *, strides=None, mask=None)`
+  on `GridGraph2D` / `GridGraph3D`. The basic form returns an `int64` array of
+  shape `(ndim, *graph.shape)` with each grid edge id written at its pivot
+  pixel and `-1` elsewhere. The offsets form returns
+  `(array, n_valid)`: an `int64` array of shape `(len(offsets), *graph.shape)`
+  whose non-`-1` entries are a sequential counter over the in-bounds (and
+  filter-accepted) targets, plus the total count. `strides` keeps only coords
+  aligned along every axis; `mask` keeps only coords where a boolean array
+  of shape `(len(offsets), *graph.shape)` is true. Like in nifty, `strides`
+  and `mask` are mutually exclusive — passing both raises `ValueError`.
 - The three grid feature functions preserve `float32` and `float64` input
   dtype end-to-end (no internal copy to `float64`); other dtypes are
   promoted to `float64`. Output weight arrays match the input dtype.

include/bioimage_cpp/graph/grid_edge_projection.hxx

@@ -0,0 +1,276 @@
+#pragma once
+
+#include "bioimage_cpp/array_view.hxx"
+#include "bioimage_cpp/graph/grid_graph.hxx"
+
+#include <array>
+#include <cstddef>
+#include <cstdint>
+#include <stdexcept>
+#include <vector>
+
+namespace bioimage_cpp::graph {
+
+namespace detail_grid_edge_projection {
+
+// Walk a sub-shape in C-order, computing a flat output offset incrementally
+// using `strides` (the strides of the *full* containing array, not of the
+// sub-shape itself). Calls `callback(flat)` at every leaf. With `Axis` known
+// at compile time the compiler unrolls the entire nest for D = 2 / D = 3.
+template <std::size_t Axis, std::size_t D, class Callback>
+void enumerate_in_c_order(
+    const std::array<std::uint64_t, D> &subshape,
+    const std::array<std::ptrdiff_t, D> &strides,
+    std::ptrdiff_t base,
+    Callback &&callback
+) {
+    if constexpr (Axis == D) {
+        callback(base);
+    } else {
+        for (std::uint64_t i = 0; i < subshape[Axis]; ++i) {
+            enumerate_in_c_order<Axis + 1, D>(subshape, strides, base, callback);
+            base += strides[Axis];
+        }
+    }
+}
+
+// Walk every coordinate in C-order over `shape`, tracking whether
+// `coord + offset` is in bounds, and call `callback(in_bounds, coord)` at the
+// leaf. `coord` is filled in-place and re-used across calls so the sampler
+// can read it without copying.
+template <std::size_t Axis, std::size_t D, class Callback>
+void enumerate_with_offset(
+    const std::array<std::uint64_t, D> &shape,
+    const std::array<std::ptrdiff_t, D> &offset,
+    std::array<std::ptrdiff_t, D> &coord,
+    bool in_bounds,
+    Callback &&callback
+) {
+    if constexpr (Axis == D) {
+        callback(in_bounds, coord);
+    } else {
+        for (std::uint64_t c = 0; c < shape[Axis]; ++c) {
+            coord[Axis] = static_cast<std::ptrdiff_t>(c);
+            const std::ptrdiff_t target = coord[Axis] + offset[Axis];
+            const bool axis_ok =
+                target >= 0 &&
+                target < static_cast<std::ptrdiff_t>(shape[Axis]);
+            enumerate_with_offset<Axis + 1, D>(
+                shape, offset, coord, in_bounds && axis_ok, callback
+            );
+        }
+    }
+}
+
+template <std::size_t D>
+std::array<std::ptrdiff_t, D> full_c_strides(
+    const std::array<std::uint64_t, D> &shape
+) {
+    std::array<std::ptrdiff_t, D> strides{};
+    std::ptrdiff_t s = 1;
+    for (std::size_t i = D; i-- > 0; ) {
+        strides[i] = s;
+        s *= static_cast<std::ptrdiff_t>(shape[i]);
+    }
+    return strides;
+}
+
+template <std::size_t D>
+std::ptrdiff_t shape_product(const std::array<std::uint64_t, D> &shape) {
+    std::ptrdiff_t p = 1;
+    for (std::size_t i = 0; i < D; ++i) {
+        p *= static_cast<std::ptrdiff_t>(shape[i]);
+    }
+    return p;
+}
+
+template <std::size_t D>
+void require_projection_output_shape(
+    const std::array<std::uint64_t, D> &graph_shape,
+    const std::vector<std::ptrdiff_t> &out_shape,
+    const std::ptrdiff_t expected_leading,
+    const char *argument_name
+) {
+    if (out_shape.size() != D + 1) {
+        throw std::invalid_argument(
+            std::string(argument_name) + " must have ndim == graph.ndim + 1"
+        );
+    }
+    if (out_shape[0] != expected_leading) {
+        throw std::invalid_argument(
+            std::string(argument_name) + " leading dimension does not match expected length"
+        );
+    }
+    for (std::size_t d = 0; d < D; ++d) {
+        if (out_shape[d + 1] != static_cast<std::ptrdiff_t>(graph_shape[d])) {
+            throw std::invalid_argument(
+                std::string(argument_name) + " spatial shape must match graph shape"
+            );
+        }
+    }
+}
+
+} // namespace detail_grid_edge_projection
+
+// Writes graph edge ids into `output` of shape (D, *graph.shape) and dtype
+// int64. The caller pre-fills `output` with -1; this function then sets
+// output[axis, c_0, ..., c_{D-1}] = edge_id for every edge whose spanning
+// axis is `axis` and whose pivot (smaller-endpoint) coordinate is
+// (c_0, ..., c_{D-1}). Slots where c_axis == shape[axis] - 1 are not visited
+// and keep their pre-filled -1.
+//
+// Exploits the fact that the GridGraph stores edges axis-major + C-order
+// within each axis: the k-th edge of axis `d` corresponds exactly to the
+// k-th pivot coordinate of the sub-shape (s_0, ..., s_d - 1, ..., s_{D-1})
+// walked in C-order. So we walk that sub-shape with the *full* output
+// strides and emit consecutive edge ids — no per-edge coordinate division.
+template <std::size_t D>
+void project_edge_ids_to_pixels(
+    const GridGraph<D> &graph,
+    const ArrayView<std::int64_t> &output
+) {
+    namespace dgep = detail_grid_edge_projection;
+    const auto &shape = graph.shape();
+    dgep::require_projection_output_shape<D>(
+        shape, output.shape, static_cast<std::ptrdiff_t>(D), "output"
+    );
+
+    const auto strides = dgep::full_c_strides<D>(shape);
+    const std::ptrdiff_t plane = dgep::shape_product<D>(shape);
+
+    auto *data = output.data;
+    std::int64_t edge_id = 0;
+    for (std::size_t axis = 0; axis < D; ++axis) {
+        auto pivot_shape = shape;
+        if (pivot_shape[axis] == 0) {
+            continue;
+        }
+        --pivot_shape[axis];
+        const std::ptrdiff_t base =
+            static_cast<std::ptrdiff_t>(axis) * plane;
+        dgep::enumerate_in_c_order<0, D>(
+            pivot_shape, strides, base,
+            [data, &edge_id](std::ptrdiff_t flat) {
+                data[flat] = edge_id++;
+            }
+        );
+    }
+}
+
+// Enumerates "lifted" edges defined by a set of per-channel grid offsets.
+// Walks (offset_idx, *coord) in C-order over (offsets.size(), *graph.shape).
+// For each coord whose target `coord + offsets[offset_idx]` stays in bounds
+// AND the sampler accepts, writes a sequential counter starting at 0;
+// rejected slots get -1. Returns the total number of valid entries written.
+//
+// The counter is NOT a graph edge id — it indexes into the implicit array
+// of lifted edges derived from an affinity volume.
+template <std::size_t D, class Sampler>
+std::uint64_t project_edge_ids_to_pixels_with_offsets_impl(
+    const GridGraph<D> &graph,
+    const std::vector<std::array<std::ptrdiff_t, D>> &offsets,
+    Sampler &&sampler,
+    const ArrayView<std::int64_t> &output
+) {
+    namespace dgep = detail_grid_edge_projection;
+    const auto &shape = graph.shape();
+    dgep::require_projection_output_shape<D>(
+        shape, output.shape, static_cast<std::ptrdiff_t>(offsets.size()), "output"
+    );
+
+    auto *out = output.data;
+    std::int64_t edge_id = 0;
+    std::size_t out_idx = 0;
+    std::array<std::ptrdiff_t, D> coord{};
+    for (std::size_t off_idx = 0; off_idx < offsets.size(); ++off_idx) {
+        const auto &off = offsets[off_idx];
+        dgep::enumerate_with_offset<0, D>(
+            shape, off, coord, /*in_bounds=*/true,
+            [&](bool in_bounds, const std::array<std::ptrdiff_t, D> &c) {
+                if (in_bounds && sampler(off_idx, c)) {
+                    out[out_idx++] = edge_id++;
+                } else {
+                    out[out_idx++] = -1;
+                }
+            }
+        );
+    }
+    return static_cast<std::uint64_t>(edge_id);
+}
+
+// Convenience overload: no filter.
+template <std::size_t D>
+std::uint64_t project_edge_ids_to_pixels_with_offsets(
+    const GridGraph<D> &graph,
+    const std::vector<std::array<std::ptrdiff_t, D>> &offsets,
+    const ArrayView<std::int64_t> &output
+) {
+    return project_edge_ids_to_pixels_with_offsets_impl<D>(
+        graph, offsets,
+        [](std::size_t, const std::array<std::ptrdiff_t, D> &) { return true; },
+        output
+    );
+}
+
+// Convenience overload: only coords aligned with `strides` along every axis.
+template <std::size_t D>
+std::uint64_t project_edge_ids_to_pixels_with_offsets(
+    const GridGraph<D> &graph,
+    const std::vector<std::array<std::ptrdiff_t, D>> &offsets,
+    const std::array<std::ptrdiff_t, D> &strides,
+    const ArrayView<std::int64_t> &output
+) {
+    for (std::size_t d = 0; d < D; ++d) {
+        if (strides[d] <= 0) {
+            throw std::invalid_argument("strides must be positive");
+        }
+    }
+    return project_edge_ids_to_pixels_with_offsets_impl<D>(
+        graph, offsets,
+        [&strides](std::size_t, const std::array<std::ptrdiff_t, D> &c) {
+            for (std::size_t d = 0; d < D; ++d) {
+                if (c[d] % strides[d] != 0) {
+                    return false;
+                }
+            }
+            return true;
+        },
+        output
+    );
+}
+
+// Convenience overload: only coords where the (n_offsets, *graph.shape)
+// `mask` is non-zero.
+template <std::size_t D>
+std::uint64_t project_edge_ids_to_pixels_with_offsets(
+    const GridGraph<D> &graph,
+    const std::vector<std::array<std::ptrdiff_t, D>> &offsets,
+    const ConstArrayView<std::uint8_t> &mask,
+    const ArrayView<std::int64_t> &output
+) {
+    namespace dgep = detail_grid_edge_projection;
+    const auto &shape = graph.shape();
+    dgep::require_projection_output_shape<D>(
+        shape, mask.shape, static_cast<std::ptrdiff_t>(offsets.size()), "mask"
+    );
+
+    const auto mask_spatial_strides = dgep::full_c_strides<D>(shape);
+    const std::ptrdiff_t mask_plane = dgep::shape_product<D>(shape);
+    const auto *mdata = mask.data;
+    return project_edge_ids_to_pixels_with_offsets_impl<D>(
+        graph, offsets,
+        [mdata, mask_spatial_strides, mask_plane](
+            std::size_t off_idx, const std::array<std::ptrdiff_t, D> &c
+        ) {
+            std::ptrdiff_t flat =
+                static_cast<std::ptrdiff_t>(off_idx) * mask_plane;
+            for (std::size_t d = 0; d < D; ++d) {
+                flat += c[d] * mask_spatial_strides[d];
+            }
+            return mdata[flat] != 0;
+        },
+        output
+    );
+}
+
+} // namespace bioimage_cpp::graph