From 5673000e33ea835390cf2b9cf7c0ba9a42d72342 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Wed, 5 Aug 2026 11:09:51 +0200 Subject: [PATCH] Further optimization of the distance transform --- development/distance/benchmark.py | 42 +- .../distance/distance_transform.hxx | 485 ++++++++++++------ tests/distance/test_distance_transform.py | 70 +++ 3 files changed, 427 insertions(+), 170 deletions(-) diff --git a/development/distance/benchmark.py b/development/distance/benchmark.py index e0c1c08..d48d946 100644 --- a/development/distance/benchmark.py +++ b/development/distance/benchmark.py @@ -9,7 +9,8 @@ outside the timing loop so per-call dtype conversion does not show up in the measurement of a particular library: -* bioimage_cpp.distance.distance_transform — uint8 mask +* bioimage_cpp.distance.distance_transform — uint8 mask +* bioimage_cpp.distance.distance_transform + indices — uint8 mask * bioimage_cpp.distance.vector_difference_transform — uint8 mask * vigra.filters.distanceTransform / vectorDistanceTransform — float32 mask * scipy.ndimage.distance_transform_edt — float32 mask @@ -43,7 +44,11 @@ LIBRARIES = ("bioimage_cpp", "vigra", "scipy") -OPERATIONS = ("distance_transform", "vector_difference_transform") +OPERATIONS = ( + "distance_transform", + "distance_transform_indices", + "vector_difference_transform", +) @dataclass(frozen=True) @@ -246,6 +251,21 @@ def fn(mask: np.ndarray) -> np.ndarray: return fn +def _bic_distance_indices(sampling: tuple[float, ...], n_threads: int): + from bioimage_cpp import distance + + def fn(mask: np.ndarray) -> np.ndarray: + distances, _ = distance.distance_transform( + mask, + sampling=sampling, + return_indices=True, + number_of_threads=n_threads, + ) + return distances + + return fn + + def _vigra_distance(sampling: tuple[float, ...]): _quiet_vigra_matplotlib_cache() import vigra.filters as vf @@ -300,6 +320,18 @@ def fn(mask: np.ndarray) -> np.ndarray: return fn +def _scipy_distance_indices(sampling: tuple[float, ...]): + from scipy import ndimage + + def fn(mask: np.ndarray) -> np.ndarray: + distances, _ = ndimage.distance_transform_edt( + mask, sampling=sampling, return_indices=True + ) + return distances + + return fn + + def _prepare_mask(library: str, base_mask: np.ndarray) -> np.ndarray: if library == "bioimage_cpp": # The Python wrapper fast-paths uint8 C-contiguous input. @@ -324,6 +356,10 @@ def build_adapters( "vigra": lambda: _vigra_distance(sampling), "scipy": lambda: _scipy_distance(sampling), }, + "distance_transform_indices": { + "bioimage_cpp": lambda: _bic_distance_indices(sampling, n_threads), + "scipy": lambda: _scipy_distance_indices(sampling), + }, "vector_difference_transform": { "bioimage_cpp": lambda: _bic_vector(sampling, n_threads), "vigra": lambda: _vigra_vector(sampling), @@ -394,7 +430,7 @@ def check_results( errors = {} for library, adapter in adapters.items(): result = np.asarray(adapter.fn(adapter.mask)) - if operation == "distance_transform": + if operation in ("distance_transform", "distance_transform_indices"): errors[library] = float( np.max(np.abs(result.astype(np.float32) - reference_distance)) ) diff --git a/include/bioimage_cpp/distance/distance_transform.hxx b/include/bioimage_cpp/distance/distance_transform.hxx index 9d5e71d..d4136f9 100644 --- a/include/bioimage_cpp/distance/distance_transform.hxx +++ b/include/bioimage_cpp/distance/distance_transform.hxx @@ -5,6 +5,7 @@ #include "bioimage_cpp/detail/profile.hxx" #include "bioimage_cpp/detail/threading.hxx" +#include #include #include #include @@ -159,6 +160,106 @@ inline void edt_1d_squared_iso(Edt1DWorkspace &ws, std::ptrdiff_t n) { edt_1d_squared_impl(ws, n, 1.0); } +// Exact first-axis transform for binary isotropic input. Equal-distance ties +// stay with the lower feature coordinate, as in edt_1d_squared_iso. +inline void binary_edt_1d_squared_iso( + Edt1DWorkspace &ws, + const std::uint8_t *input, + const std::ptrdiff_t base, + const std::ptrdiff_t stride, + const std::ptrdiff_t n +) { + std::ptrdiff_t number_of_features = 0; + for (std::ptrdiff_t i = 0; i < n; ++i) { + if (input[base + i * stride] == 0) { + ws.envelope_v[static_cast(number_of_features++)] = + static_cast(i); + } + } + + if (number_of_features == 0) { + for (std::ptrdiff_t i = 0; i < n; ++i) { + ws.distance[static_cast(i)] = kInfinity; + ws.source[static_cast(i)] = -1; + } + return; + } + if (number_of_features == n) { + for (std::ptrdiff_t i = 0; i < n; ++i) { + ws.distance[static_cast(i)] = 0.0; + ws.source[static_cast(i)] = static_cast(i); + } + return; + } + + if (number_of_features < 1 + (n - 1) / 8) { + std::ptrdiff_t feature = 0; + for (std::ptrdiff_t i = 0; i < n; ++i) { + while ( + feature + 1 < number_of_features && + 2 * static_cast(i) > + static_cast( + ws.envelope_v[static_cast(feature)] + ) + + static_cast( + ws.envelope_v[static_cast(feature + 1)] + ) + ) { + ++feature; + } + const auto nearest = ws.envelope_v[static_cast(feature)]; + const double diff = static_cast(i - nearest); + ws.distance[static_cast(i)] = diff * diff; + ws.source[static_cast(i)] = nearest; + } + return; + } + + std::ptrdiff_t begin = 0; + for (std::ptrdiff_t feature = 0; feature < number_of_features; ++feature) { + const auto nearest = ws.envelope_v[static_cast(feature)]; + std::ptrdiff_t end = n; + if (feature + 1 < number_of_features) { + const auto next = ws.envelope_v[static_cast(feature + 1)]; + end = static_cast( + (static_cast(nearest) + static_cast(next)) / + 2 + + 1 + ); + } + for (auto i = begin; i < end; ++i) { + const double diff = static_cast(i - nearest); + ws.distance[static_cast(i)] = diff * diff; + ws.source[static_cast(i)] = nearest; + } + begin = end; + } +} + +inline bool squared_distances_fit_exact_float( + const std::vector &shape +) { + constexpr std::uint64_t max_exact_integer = + std::uint64_t{1} << std::numeric_limits::digits; + constexpr std::uint64_t max_exact_delta = 4096; + std::uint64_t remaining = max_exact_integer; + for (const auto axis_size : shape) { + if (axis_size <= 0) { + return false; + } + const auto delta = static_cast(axis_size - 1); + if (delta > max_exact_delta) { + return false; + } + const auto squared = delta * delta; + if (squared > remaining) { + return false; + } + remaining -= squared; + } + return true; +} + inline void unravel( std::ptrdiff_t flat, const std::vector &strides, @@ -240,17 +341,28 @@ inline void detail_distance_transform_impl( BIOIMAGE_PROFILE_INIT(profiler) - // Detect the all-foreground (no background) case. SciPy reports distances - // and indices against a virtual background row at axis-0 coordinate -1; we - // mirror that convention so callers can switch between SciPy and us - // without surprises. + // Detect the all-foreground case and sample the binary fast-path density. + // SciPy reports all-foreground results against a virtual background row at + // axis-0 coordinate -1; we mirror that convention. + constexpr std::ptrdiff_t binary_sample_limit = 4096; bool has_background = use_initialized_squared; + std::ptrdiff_t sampled_background = 0; + std::ptrdiff_t sample_size = 0; if (!use_initialized_squared) { BIOIMAGE_PROFILE_SCOPE(profiler, "scan_for_bg") - for (std::ptrdiff_t i = 0; i < n; ++i) { + sample_size = std::min(n, binary_sample_limit); + for (std::ptrdiff_t i = 0; i < sample_size; ++i) { if (input.data[i] == 0) { + ++sampled_background; has_background = true; - break; + } + } + if (!has_background) { + for (std::ptrdiff_t i = sample_size; i < n; ++i) { + if (input.data[i] == 0) { + has_background = true; + break; + } } } } @@ -283,198 +395,237 @@ inline void detail_distance_transform_impl( return; } - // Squared sampled distance buffer. ndim per-axis feature-coord buffers - // (int32) replace the previous flat int64 feature index — this lets the - // output pass materialize indices/vectors without re-unraveling per pixel. - auto squared_distance = use_initialized_squared + const bool use_binary_first_sweep = + !use_initialized_squared && is_isotropic && + sampled_background >= (sample_size + 7) / 8; + const bool use_float_squared = + !use_initialized_squared && squared_distance_consumer == nullptr && + is_isotropic && detail::squared_distances_fit_exact_float(input.shape); + + std::unique_ptr squared_distance_double = use_initialized_squared ? std::move(initialized_squared_distance) - : std::make_unique_for_overwrite(static_cast(n)); - std::vector> feature_coord; - if (track_feature) { - feature_coord.resize(static_cast(ndim)); - for (auto &arr : feature_coord) { - arr.assign(static_cast(n), 0); - } + : nullptr; + std::unique_ptr squared_distance_float; + if (use_float_squared) { + squared_distance_float = std::make_unique_for_overwrite( + static_cast(n) + ); + } else if (squared_distance_double == nullptr) { + squared_distance_double = std::make_unique_for_overwrite( + static_cast(n) + ); } - for (std::ptrdiff_t ax = 0; ax < ndim; ++ax) { - BIOIMAGE_PROFILE_SCOPE(profiler, "sweep_axis") - const std::ptrdiff_t line_length = input.shape[static_cast(ax)]; - if (line_length <= 0) { - continue; - } - const std::ptrdiff_t stride = strides[static_cast(ax)]; - const std::ptrdiff_t inner_count = stride; - const std::ptrdiff_t axis_block = line_length * stride; - const std::ptrdiff_t outer_count = (axis_block == 0) ? 0 : n / axis_block; - const std::size_t n_lines = - static_cast(outer_count) * static_cast(inner_count); - if (n_lines == 0) { - continue; - } - const double sampling_ax = sampling[static_cast(ax)]; - const double squared_spacing = sampling_ax * sampling_ax; - const std::ptrdiff_t feature_axes_in = track_feature ? ax : 0; - - const auto process_line = [&](std::size_t line_id, detail::Edt1DWorkspace &ws) { - const auto outer_idx = static_cast(line_id) / inner_count; - const auto inner_idx = static_cast(line_id) % inner_count; - const std::ptrdiff_t base = outer_idx * axis_block + inner_idx; - ws.ensure(line_length, feature_axes_in); - - // The first-axis gather initializes the uninitialized squared- - // distance buffer directly from the input. Later axes gather the - // preceding sweep, avoiding a redundant full-volume init pass. - if (ax == 0 && !use_initialized_squared) { - for (std::ptrdiff_t i = 0; i < line_length; ++i) { - const auto index = static_cast(base + i * stride); - ws.f[static_cast(i)] = - input.data[index] == 0 ? 0.0 : detail::kInfinity; + + const auto run_sweeps = [&]( + SquaredDistance *squared_distance + ) { + // Indices can hold feature coordinates during the sweeps. Vector-only + // calls use one uninitialized coordinate buffer with the same layout. + std::vector feature_coord; + std::unique_ptr owned_feature_coord; + if (track_feature) { + feature_coord.resize(static_cast(ndim)); + if (want_indices) { + for (std::ptrdiff_t ax = 0; ax < ndim; ++ax) { + feature_coord[static_cast(ax)] = + outputs.indices.data + ax * n; } } else { - for (std::ptrdiff_t i = 0; i < line_length; ++i) { - const auto index = static_cast(base + i * stride); - ws.f[static_cast(i)] = squared_distance[index]; + owned_feature_coord = std::make_unique_for_overwrite( + static_cast(ndim) * static_cast(n) + ); + for (std::ptrdiff_t ax = 0; ax < ndim; ++ax) { + feature_coord[static_cast(ax)] = + owned_feature_coord.get() + ax * n; } } - // Gather already-tracked feature coords (axes < ax). - for (std::ptrdiff_t a = 0; a < feature_axes_in; ++a) { - const auto *src = feature_coord[static_cast(a)].data(); - auto *dst = ws.old_feature_coord.data() + a * line_length; - for (std::ptrdiff_t i = 0; i < line_length; ++i) { - dst[i] = src[base + i * stride]; - } - } - - if (is_isotropic) { - detail::edt_1d_squared_iso(ws, line_length); - } else { - detail::edt_1d_squared(ws, line_length, squared_spacing); - } + } - // Scatter squared distances back. - for (std::ptrdiff_t i = 0; i < line_length; ++i) { - squared_distance[static_cast(base + i * stride)] = - ws.distance[static_cast(i)]; + for (std::ptrdiff_t ax = 0; ax < ndim; ++ax) { + BIOIMAGE_PROFILE_SCOPE(profiler, "sweep_axis") + const std::ptrdiff_t line_length = input.shape[static_cast(ax)]; + if (line_length <= 0) { + continue; } - if (!track_feature) { - return; + const std::ptrdiff_t stride = strides[static_cast(ax)]; + const std::ptrdiff_t inner_count = stride; + const std::ptrdiff_t axis_block = line_length * stride; + const std::ptrdiff_t outer_count = (axis_block == 0) ? 0 : n / axis_block; + const std::size_t n_lines = + static_cast(outer_count) * static_cast(inner_count); + if (n_lines == 0) { + continue; } - // Scatter feature coords for axes < ax via source[i]. - for (std::ptrdiff_t a = 0; a < ax; ++a) { - auto *dst = feature_coord[static_cast(a)].data(); - const auto *src = ws.old_feature_coord.data() + a * line_length; - for (std::ptrdiff_t i = 0; i < line_length; ++i) { - const auto s = ws.source[static_cast(i)]; - if (s >= 0) { - dst[base + i * stride] = src[s]; + const double sampling_ax = sampling[static_cast(ax)]; + const double squared_spacing = sampling_ax * sampling_ax; + const std::ptrdiff_t feature_axes_in = track_feature ? ax : 0; + + const auto process_line = [&](std::size_t line_id, detail::Edt1DWorkspace &ws) { + const auto outer_idx = static_cast(line_id) / inner_count; + const auto inner_idx = static_cast(line_id) % inner_count; + const std::ptrdiff_t base = outer_idx * axis_block + inner_idx; + ws.ensure(line_length, feature_axes_in); + const bool use_binary_line = ax == 0 && use_binary_first_sweep; + + // The first-axis gather initializes the uninitialized squared- + // distance buffer directly from the input. Later axes gather the + // preceding sweep, avoiding a redundant full-volume init pass. + if (!use_binary_line) { + if (ax == 0 && !use_initialized_squared) { + for (std::ptrdiff_t i = 0; i < line_length; ++i) { + const auto index = static_cast(base + i * stride); + ws.f[static_cast(i)] = + input.data[index] == 0 ? 0.0 : detail::kInfinity; + } + } else { + for (std::ptrdiff_t i = 0; i < line_length; ++i) { + const auto index = static_cast(base + i * stride); + ws.f[static_cast(i)] = + static_cast(squared_distance[index]); + } } } - } - // Axis ax's feature coord is the parabola minimizer position itself. - { - auto *dst = feature_coord[static_cast(ax)].data(); + // Gather already-tracked feature coords (axes < ax). + for (std::ptrdiff_t a = 0; a < feature_axes_in; ++a) { + const auto *src = feature_coord[static_cast(a)]; + auto *dst = ws.old_feature_coord.data() + a * line_length; + for (std::ptrdiff_t i = 0; i < line_length; ++i) { + dst[i] = src[base + i * stride]; + } + } + + if (use_binary_line) { + detail::binary_edt_1d_squared_iso( + ws, input.data, base, stride, line_length + ); + } else if (is_isotropic) { + detail::edt_1d_squared_iso(ws, line_length); + } else { + detail::edt_1d_squared(ws, line_length, squared_spacing); + } + + // Scatter squared distances back. for (std::ptrdiff_t i = 0; i < line_length; ++i) { - const auto s = ws.source[static_cast(i)]; - if (s >= 0) { - dst[base + i * stride] = s; + squared_distance[static_cast(base + i * stride)] = + ws.distance[static_cast(i)]; + } + if (!track_feature) { + return; + } + // Scatter feature coords for axes < ax via source[i]. + for (std::ptrdiff_t a = 0; a < ax; ++a) { + auto *dst = feature_coord[static_cast(a)]; + const auto *src = ws.old_feature_coord.data() + a * line_length; + for (std::ptrdiff_t i = 0; i < line_length; ++i) { + const auto s = ws.source[static_cast(i)]; + if (s >= 0) { + dst[base + i * stride] = src[s]; + } } } - } - }; - - const auto resolved_threads = - bioimage_cpp::detail::normalize_thread_count(n_threads, n_lines); - if (resolved_threads <= 1) { - detail::Edt1DWorkspace ws; - for (std::size_t line_id = 0; line_id < n_lines; ++line_id) { - process_line(line_id, ws); - } - } else { - std::vector per_thread(resolved_threads); - bioimage_cpp::detail::parallel_for_chunks( - resolved_threads, - n_lines, - [&](std::size_t thread_id, std::size_t begin, std::size_t end) { - auto &ws = per_thread[thread_id]; - for (std::size_t line_id = begin; line_id < end; ++line_id) { - process_line(line_id, ws); + // Axis ax's feature coord is the parabola minimizer position itself. + { + auto *dst = feature_coord[static_cast(ax)]; + for (std::ptrdiff_t i = 0; i < line_length; ++i) { + const auto s = ws.source[static_cast(i)]; + if (s >= 0) { + dst[base + i * stride] = s; + } } } - ); - } - } - - // Output materialization. All three branches stream over flat indices in - // C-order with no integer divisions per pixel: indices come straight from - // per-axis feature buffers, vectors use an incremental coord counter. - if (want_distances) { - BIOIMAGE_PROFILE_SCOPE(profiler, "output_distances") - const auto output_threads = bioimage_cpp::detail::normalize_thread_count( - n_threads, static_cast(n) - ); - const auto write_distances = [&](const std::size_t begin, const std::size_t end) { - for (std::size_t i = begin; i < end; ++i) { - outputs.distances.data[i] = - static_cast(std::sqrt(squared_distance[i])); - } - }; - if (output_threads <= 1) { - write_distances(0, static_cast(n)); - } else { - bioimage_cpp::detail::parallel_for_chunks( - output_threads, - static_cast(n), - [&](const std::size_t, const std::size_t begin, const std::size_t end) { - write_distances(begin, end); + }; + + const auto resolved_threads = + bioimage_cpp::detail::normalize_thread_count(n_threads, n_lines); + if (resolved_threads <= 1) { + detail::Edt1DWorkspace ws; + for (std::size_t line_id = 0; line_id < n_lines; ++line_id) { + process_line(line_id, ws); } - ); + } else { + std::vector per_thread(resolved_threads); + bioimage_cpp::detail::parallel_for_chunks( + resolved_threads, + n_lines, + [&](std::size_t thread_id, std::size_t begin, std::size_t end) { + auto &ws = per_thread[thread_id]; + for (std::size_t line_id = begin; line_id < end; ++line_id) { + process_line(line_id, ws); + } + } + ); + } } - } - if (want_indices) { - BIOIMAGE_PROFILE_SCOPE(profiler, "output_indices") - for (std::ptrdiff_t ax = 0; ax < ndim; ++ax) { - const auto *src = feature_coord[static_cast(ax)].data(); - auto *dst = outputs.indices.data + ax * n; - for (std::ptrdiff_t i = 0; i < n; ++i) { - dst[i] = src[i]; + // Indices already contain their final values. Distances and vectors + // stream over flat indices in C-order without per-pixel division. + if (want_distances) { + BIOIMAGE_PROFILE_SCOPE(profiler, "output_distances") + const auto output_threads = bioimage_cpp::detail::normalize_thread_count( + n_threads, static_cast(n) + ); + const auto write_distances = [&](const std::size_t begin, const std::size_t end) { + for (std::size_t i = begin; i < end; ++i) { + outputs.distances.data[i] = + static_cast( + std::sqrt(static_cast(squared_distance[i])) + ); + } + }; + if (output_threads <= 1) { + write_distances(0, static_cast(n)); + } else { + bioimage_cpp::detail::parallel_for_chunks( + output_threads, + static_cast(n), + [&](const std::size_t, const std::size_t begin, const std::size_t end) { + write_distances(begin, end); + } + ); } } - } - if (want_vectors) { - BIOIMAGE_PROFILE_SCOPE(profiler, "output_vectors") - std::vector coord(static_cast(ndim), 0); - std::vector shape_i32(static_cast(ndim), 0); - for (std::ptrdiff_t ax = 0; ax < ndim; ++ax) { - shape_i32[static_cast(ax)] = - static_cast(input.shape[static_cast(ax)]); - } - for (std::ptrdiff_t i = 0; i < n; ++i) { - auto *dst = outputs.vectors.data + i * ndim; + if (want_vectors) { + BIOIMAGE_PROFILE_SCOPE(profiler, "output_vectors") + std::vector coord(static_cast(ndim), 0); + std::vector shape_i32(static_cast(ndim), 0); for (std::ptrdiff_t ax = 0; ax < ndim; ++ax) { - const double diff = - static_cast( - feature_coord[static_cast(ax)][static_cast(i)] - - coord[static_cast(ax)] - ) * - sampling[static_cast(ax)]; - dst[ax] = static_cast(diff); + shape_i32[static_cast(ax)] = + static_cast(input.shape[static_cast(ax)]); } - // Increment coord in C-order (innermost axis fastest). - for (std::ptrdiff_t ax = ndim - 1; ax >= 0; --ax) { - auto &c = coord[static_cast(ax)]; - if (++c < shape_i32[static_cast(ax)]) { - break; + for (std::ptrdiff_t i = 0; i < n; ++i) { + auto *dst = outputs.vectors.data + i * ndim; + for (std::ptrdiff_t ax = 0; ax < ndim; ++ax) { + const double diff = + static_cast( + feature_coord[static_cast(ax)][static_cast(i)] - + coord[static_cast(ax)] + ) * + sampling[static_cast(ax)]; + dst[ax] = static_cast(diff); + } + // Increment coord in C-order (innermost axis fastest). + for (std::ptrdiff_t ax = ndim - 1; ax >= 0; --ax) { + auto &c = coord[static_cast(ax)]; + if (++c < shape_i32[static_cast(ax)]) { + break; + } + c = 0; } - c = 0; } } + }; + + if (use_float_squared) { + run_sweeps(squared_distance_float.get()); + } else { + run_sweeps(squared_distance_double.get()); } + if (squared_distance_consumer != nullptr) { BIOIMAGE_PROFILE_SCOPE(profiler, "squared_distance_consumer") squared_distance_consumer( - squared_distance.get(), static_cast(n), + squared_distance_double.get(), static_cast(n), squared_distance_consumer_context ); } diff --git a/tests/distance/test_distance_transform.py b/tests/distance/test_distance_transform.py index effbf0b..d963110 100644 --- a/tests/distance/test_distance_transform.py +++ b/tests/distance/test_distance_transform.py @@ -15,6 +15,20 @@ def test_distance_transform_1d_matches_scipy(): np.testing.assert_allclose(got, ref) +def test_dense_binary_first_sweep_preserves_feature_ties(): + data = np.array([0, 1, 0, 0, 1, 1, 0, 1], dtype=np.uint8) + + got_dist, got_idx, got_vec = bic.distance.distance_transform( + data, return_indices=True, return_vectors=True + ) + ref_dist, ref_idx = ndimage.distance_transform_edt(data, return_indices=True) + + np.testing.assert_array_equal(got_dist, ref_dist.astype(np.float32)) + np.testing.assert_array_equal(got_idx, ref_idx) + expected_vec = (ref_idx[0] - np.arange(data.size)).astype(np.float32) + np.testing.assert_array_equal(got_vec[:, 0], expected_vec) + + @pytest.mark.parametrize("shape", [(7, 11), (4, 5, 6)]) def test_distance_transform_matches_scipy(shape): data = np.ones(shape, dtype=np.uint8) @@ -40,6 +54,18 @@ def test_anisotropic_sampling_matches_scipy(): np.testing.assert_allclose(got, ref, atol=1e-6) +@pytest.mark.parametrize("length", [4097, 4098]) +def test_isotropic_squared_buffer_exactness_boundary(length): + data = np.ones(length, dtype=np.uint8) + data[0] = 0 + + got_dist, got_idx = bic.distance.distance_transform(data, return_indices=True) + ref_dist, ref_idx = ndimage.distance_transform_edt(data, return_indices=True) + + np.testing.assert_array_equal(got_dist, ref_dist.astype(np.float32)) + np.testing.assert_array_equal(got_idx, ref_idx) + + def test_return_indices_matches_scipy_for_unique_nearest_background(): data = np.ones((5, 6), dtype=np.uint8) data[2, 3] = 0 @@ -175,6 +201,50 @@ def test_threaded_outputs_are_exact(sampling): np.testing.assert_array_equal(got, expected) +def test_dense_isotropic_fast_paths_match_scipy_and_threads(): + zz, yy, xx = np.indices((9, 11, 13)) + data = ((3 * zz + 5 * yy + 7 * xx) % 2 != 0).astype(np.uint8) + + sequential = bic.distance.distance_transform( + data, + return_indices=True, + return_vectors=True, + number_of_threads=1, + ) + threaded = bic.distance.distance_transform( + data, + return_indices=True, + return_vectors=True, + number_of_threads=4, + ) + ref_dist, ref_idx = ndimage.distance_transform_edt(data, return_indices=True) + + for got, expected in zip(threaded, sequential): + np.testing.assert_array_equal(got, expected) + np.testing.assert_array_equal(sequential[0], ref_dist.astype(np.float32)) + np.testing.assert_array_equal(sequential[1], ref_idx) + np.testing.assert_array_equal( + np.linalg.norm(sequential[2], axis=-1), sequential[0] + ) + + +def test_dense_vector_only_fast_path_uses_owned_feature_storage(): + yy, xx = np.indices((11, 13)) + data = ((3 * yy + 5 * xx) % 2 != 0).astype(np.uint8) + + vector_only = bic.distance.distance_transform( + data, + return_distances=False, + return_vectors=True, + ) + _, expected = bic.distance.distance_transform( + data, + return_vectors=True, + ) + + np.testing.assert_array_equal(vector_only, expected) + + def test_vector_difference_transform_unique_target(): data = np.ones((5, 6), dtype=np.uint8) data[2, 3] = 0