Implement mesh smoothing

Author: constantinpape

development/utils/benchmark_mesh_smoothing.py

@@ -0,0 +1,177 @@
+"""Benchmark Laplacian mesh smoothing: bioimage_cpp vs the nifty-based reference.
+
+The mesh is built by triangulating ``n`` random points sampled on the unit
+sphere via ``scipy.spatial.ConvexHull`` — a closed manifold of ``2n - 4``
+triangles is a fair workload for the algorithm. The Python reference is
+imported from ``_mesh_smoothing_reference.py`` and uses ``nifty`` for graph
+adjacency; if either library is missing, the corresponding column is skipped.
+
+Run::
+
+    python development/utils/benchmark_mesh_smoothing.py --n-vertices 5000 --iterations 5 --repeats 3
+    python development/utils/benchmark_mesh_smoothing.py --n-vertices 20000 --threads 1,2,4,0
+
+Notes
+-----
+The reference does in-place Gauss-Seidel smoothing after the first iteration
+(an aliasing accident in the Python source); bioimage_cpp does textbook Jacobi
+smoothing. Correctness against the reference is only checked for
+``iterations=1``, which is where the two implementations agree exactly.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import os
+import sys
+from dataclasses import dataclass
+from statistics import median
+from time import perf_counter
+
+import numpy as np
+
+import bioimage_cpp as bic
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument("--n-vertices", type=int, default=5000)
+    parser.add_argument("--iterations", type=int, default=5)
+    parser.add_argument("--repeats", type=int, default=3)
+    parser.add_argument("--warmup", type=int, default=1)
+    parser.add_argument("--seed", type=int, default=2026)
+    parser.add_argument(
+        "--threads",
+        type=str,
+        default="1,0",
+        help="Comma-separated list of n_threads values to benchmark (0=auto).",
+    )
+    parser.add_argument("--csv", type=str, default=None, help="Optional CSV output path.")
+    return parser.parse_args()
+
+
+@dataclass(frozen=True)
+class Mesh:
+    verts: np.ndarray
+    normals: np.ndarray
+    faces: np.ndarray
+
+
+def build_sphere_mesh(n_vertices: int, seed: int) -> Mesh:
+    from scipy.spatial import ConvexHull
+
+    rng = np.random.default_rng(seed)
+    raw = rng.standard_normal((n_vertices, 3))
+    points = (raw / np.linalg.norm(raw, axis=1, keepdims=True)).astype(np.float64)
+    hull = ConvexHull(points)
+    return Mesh(verts=points, normals=points.copy(), faces=hull.simplices.astype(np.int64))
+
+
+def time_call(fn, repeats: int, warmup: int) -> list[float]:
+    for _ in range(warmup):
+        fn()
+    times = []
+    for _ in range(repeats):
+        t0 = perf_counter()
+        fn()
+        times.append(perf_counter() - t0)
+    return times
+
+
+def import_reference():
+    here = os.path.dirname(os.path.abspath(__file__))
+    if here not in sys.path:
+        sys.path.insert(0, here)
+    try:
+        import nifty  # noqa: F401
+    except ImportError:
+        return None
+    try:
+        from _mesh_smoothing_reference import smooth_mesh as ref
+    except ImportError:
+        return None
+    return ref
+
+
+def main() -> int:
+    args = parse_args()
+    thread_values = [int(t) for t in args.threads.split(",") if t.strip()]
+
+    try:
+        mesh = build_sphere_mesh(args.n_vertices, args.seed)
+    except ImportError:
+        print("scipy is required to build the benchmark mesh.")
+        return 1
+
+    n_faces = mesh.faces.shape[0]
+    print(
+        f"Mesh: {args.n_vertices} vertices, {n_faces} faces "
+        f"({args.iterations} smoothing iterations, {args.repeats} repeats)"
+    )
+
+    rows: list[dict] = []
+
+    # bioimage_cpp variants (one per --threads value).
+    for n_threads in thread_values:
+        label = f"bioimage_cpp[n_threads={n_threads}]"
+
+        def run(verts=mesh.verts, normals=mesh.normals, faces=mesh.faces, n=n_threads):
+            bic.utils.smooth_mesh(verts, normals, faces, iterations=args.iterations, n_threads=n)
+
+        times = time_call(run, args.repeats, args.warmup)
+        rows.append(
+            {"impl": label, "median_s": median(times), "best_s": min(times)}
+        )
+
+    # Reference (single-threaded by definition).
+    reference = import_reference()
+    if reference is not None:
+        def run_ref():
+            reference(mesh.verts, mesh.normals, mesh.faces, args.iterations)
+
+        times = time_call(run_ref, args.repeats, args.warmup)
+        rows.append(
+            {"impl": "reference[nifty]", "median_s": median(times), "best_s": min(times)}
+        )
+    else:
+        print("(reference skipped: nifty not installed)")
+
+    # Correctness check against the reference at iterations=1 (only point of
+    # exact agreement; see module docstring).
+    if reference is not None:
+        ref_v, ref_n = reference(mesh.verts, mesh.normals, mesh.faces, 1)
+        ours_v, ours_n = bic.utils.smooth_mesh(mesh.verts, mesh.normals, mesh.faces, iterations=1)
+        v_max_diff = float(np.max(np.abs(ours_v - ref_v)))
+        n_max_diff = float(np.max(np.abs(ours_n - ref_n)))
+        print(f"Correctness vs reference (iterations=1): "
+              f"max|Δverts|={v_max_diff:.2e}, max|Δnormals|={n_max_diff:.2e}")
+
+    # Print result table.
+    print()
+    print(f"{'impl':<32} {'median (s)':>12} {'best (s)':>12} {'speedup':>10}")
+    baseline = None
+    for row in rows:
+        if row["impl"] == "reference[nifty]":
+            baseline = row["median_s"]
+            break
+    for row in rows:
+        speedup = baseline / row["median_s"] if baseline else float("nan")
+        print(
+            f"{row['impl']:<32} {row['median_s']:>12.4f} {row['best_s']:>12.4f} "
+            f"{speedup:>10.2f}x"
+        )
+
+    if args.csv:
+        with open(args.csv, "w", newline="") as f:
+            writer = csv.DictWriter(f, fieldnames=["impl", "median_s", "best_s"])
+            writer.writeheader()
+            for row in rows:
+                writer.writerow(row)
+        print(f"\nWrote {args.csv}")
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

include/bioimage_cpp/mesh_smoothing.hxx

@@ -0,0 +1,187 @@
+#pragma once
+
+#include "bioimage_cpp/array_view.hxx"
+#include "bioimage_cpp/detail/threading.hxx"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <numeric>
+#include <stdexcept>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace bioimage_cpp {
+
+namespace detail::mesh_smoothing {
+
+template <class I>
+struct Adjacency {
+    std::vector<std::size_t> offsets;
+    std::vector<I> neighbours;
+};
+
+template <class I>
+Adjacency<I> build_adjacency(const ConstArrayView<I> &faces, std::ptrdiff_t n_verts) {
+    const std::ptrdiff_t n_faces = faces.shape[0];
+
+    std::vector<std::pair<I, I>> edges;
+    edges.reserve(static_cast<std::size_t>(n_faces) * 3);
+
+    const auto signed_n_verts = static_cast<std::int64_t>(n_verts);
+    auto check_index = [&](const I value) {
+        const auto signed_value = static_cast<std::int64_t>(value);
+        if (signed_value < 0 || signed_value >= signed_n_verts) {
+            throw std::invalid_argument(
+                "faces contains index " + std::to_string(signed_value)
+                + " out of range [0, " + std::to_string(n_verts) + ")"
+            );
+        }
+    };
+
+    for (std::ptrdiff_t f = 0; f < n_faces; ++f) {
+        const I a = faces.data[f * 3 + 0];
+        const I b = faces.data[f * 3 + 1];
+        const I c = faces.data[f * 3 + 2];
+        check_index(a);
+        check_index(b);
+        check_index(c);
+        edges.emplace_back(std::min(a, b), std::max(a, b));
+        edges.emplace_back(std::min(b, c), std::max(b, c));
+        edges.emplace_back(std::min(a, c), std::max(a, c));
+    }
+
+    std::sort(edges.begin(), edges.end());
+    edges.erase(std::unique(edges.begin(), edges.end()), edges.end());
+
+    std::vector<std::size_t> offsets(static_cast<std::size_t>(n_verts) + 1, 0);
+    for (const auto &edge : edges) {
+        ++offsets[static_cast<std::size_t>(edge.first) + 1];
+        ++offsets[static_cast<std::size_t>(edge.second) + 1];
+    }
+    for (std::size_t i = 1; i < offsets.size(); ++i) {
+        offsets[i] += offsets[i - 1];
+    }
+
+    std::vector<I> neighbours(offsets.back());
+    std::vector<std::size_t> insert_pos(offsets.begin(), offsets.end() - 1);
+    for (const auto &edge : edges) {
+        neighbours[insert_pos[static_cast<std::size_t>(edge.first)]++] = edge.second;
+        neighbours[insert_pos[static_cast<std::size_t>(edge.second)]++] = edge.first;
+    }
+
+    return Adjacency<I>{std::move(offsets), std::move(neighbours)};
+}
+
+} // namespace detail::mesh_smoothing
+
+// Laplacian smoothing of a triangular mesh: each vertex (and corresponding
+// normal) is replaced by the mean of itself and its 1-ring neighbours,
+// repeated for `iterations` passes. `verts` and `normals` are (n_verts, dim);
+// `faces` is (n_faces, 3) with values in [0, n_verts). Adjacency is built once
+// and reused across iterations. For `iterations == 0`, inputs are copied to
+// outputs unchanged.
+template <class V, class I>
+void smooth_mesh(
+    const ConstArrayView<V> &verts,
+    const ConstArrayView<V> &normals,
+    const ConstArrayView<I> &faces,
+    std::size_t iterations,
+    int n_threads,
+    const ArrayView<V> &out_verts,
+    const ArrayView<V> &out_normals
+) {
+    if (verts.shape.size() != 2) {
+        throw std::invalid_argument(
+            "verts must have ndim=2, got ndim=" + std::to_string(verts.shape.size())
+        );
+    }
+    if (normals.shape != verts.shape) {
+        throw std::invalid_argument("normals shape must match verts shape");
+    }
+    if (faces.shape.size() != 2 || faces.shape[1] != 3) {
+        throw std::invalid_argument("faces must have shape (n_faces, 3)");
+    }
+    if (out_verts.shape != verts.shape) {
+        throw std::invalid_argument("out_verts shape must match verts shape");
+    }
+    if (out_normals.shape != verts.shape) {
+        throw std::invalid_argument("out_normals shape must match verts shape");
+    }
+
+    const std::ptrdiff_t n_verts = verts.shape[0];
+    const std::ptrdiff_t dim = verts.shape[1];
+    const std::size_t n_total = static_cast<std::size_t>(n_verts) * static_cast<std::size_t>(dim);
+
+    if (iterations == 0) {
+        std::copy(verts.data, verts.data + n_total, out_verts.data);
+        std::copy(normals.data, normals.data + n_total, out_normals.data);
+        return;
+    }
+
+    const auto adjacency = detail::mesh_smoothing::build_adjacency<I>(faces, n_verts);
+
+    std::vector<V> scratch_verts(n_total);
+    std::vector<V> scratch_normals(n_total);
+
+    // Choose initial write target so the final write lands in out_*.
+    // iter 0 reads from `verts`/`normals` (input); subsequent iters alternate
+    // between out_* and scratch. The last write must be to out_*, so when
+    // iterations is even, iter 0 writes scratch (and iter 1 writes out).
+    V *buf_a_verts = out_verts.data;
+    V *buf_a_normals = out_normals.data;
+    V *buf_b_verts = scratch_verts.data();
+    V *buf_b_normals = scratch_normals.data();
+    if (iterations % 2 == 0) {
+        std::swap(buf_a_verts, buf_b_verts);
+        std::swap(buf_a_normals, buf_b_normals);
+    }
+
+    const auto &offsets = adjacency.offsets;
+    const auto &neighbours = adjacency.neighbours;
+    const std::size_t threads = detail::normalize_thread_count(
+        n_threads < 0 ? 0 : static_cast<std::size_t>(n_threads),
+        static_cast<std::size_t>(n_verts)
+    );
+
+    auto smooth_pass = [&](const V *src_verts, const V *src_normals, V *dst_verts, V *dst_normals) {
+        detail::parallel_for_chunks(
+            threads,
+            static_cast<std::size_t>(n_verts),
+            [&](std::size_t, std::size_t begin, std::size_t end) {
+                for (std::size_t v = begin; v < end; ++v) {
+                    const std::size_t beg = offsets[v];
+                    const std::size_t fin = offsets[v + 1];
+                    const auto count = static_cast<V>(fin - beg + 1);
+                    const std::size_t row = v * static_cast<std::size_t>(dim);
+                    for (std::ptrdiff_t d = 0; d < dim; ++d) {
+                        V sum_v = src_verts[row + static_cast<std::size_t>(d)];
+                        V sum_n = src_normals[row + static_cast<std::size_t>(d)];
+                        for (std::size_t k = beg; k < fin; ++k) {
+                            const std::size_t nbr_row =
+                                static_cast<std::size_t>(neighbours[k]) * static_cast<std::size_t>(dim);
+                            sum_v += src_verts[nbr_row + static_cast<std::size_t>(d)];
+                            sum_n += src_normals[nbr_row + static_cast<std::size_t>(d)];
+                        }
+                        dst_verts[row + static_cast<std::size_t>(d)] = sum_v / count;
+                        dst_normals[row + static_cast<std::size_t>(d)] = sum_n / count;
+                    }
+                }
+            }
+        );
+    };
+
+    // First iteration: read from inputs, write to buf_a.
+    smooth_pass(verts.data, normals.data, buf_a_verts, buf_a_normals);
+
+    for (std::size_t it = 1; it < iterations; ++it) {
+        const V *src_verts = (it % 2 == 1) ? buf_a_verts : buf_b_verts;
+        const V *src_normals = (it % 2 == 1) ? buf_a_normals : buf_b_normals;
+        V *dst_verts = (it % 2 == 1) ? buf_b_verts : buf_a_verts;
+        V *dst_normals = (it % 2 == 1) ? buf_b_normals : buf_a_normals;
+        smooth_pass(src_verts, src_normals, dst_verts, dst_normals);
+    }
+}
+
+} // namespace bioimage_cpp