surfel.hpp

 
#ifndef UFO_CORE_SURFEL_HPP
#define UFO_CORE_SURFEL_HPP
 
// UFO
#include <ufo/numeric/mat.hpp>
#include <ufo/numeric/vec.hpp>
 
// STL
#include <algorithm>
#include <array>
#include <concepts>
#include <cstdint>
#include <format>
#include <ostream>
#include <ranges>
 
namespace ufo
{
class Surfel
{
 public:
    //
    // Constructors
    //
 
    constexpr Surfel() = default;
 
    constexpr Surfel(Vec<3, float> sum, std::array<float, 6> sum_squares,
                     std::uint32_t num_points)
        : sum_squares_(sum_squares), sum_(sum), num_points_(num_points)
    {
    }
 
    constexpr Surfel(Vec<3, float> position, Mat<3, 3, float> scatter = {})
        : sum_squares_{scatter[0][0], scatter[0][1], scatter[0][2],
                       scatter[1][1], scatter[1][2], scatter[2][2]}
        , sum_(position)
        , num_points_(1)
    {
    }
 
    template <std::input_iterator InputIt>
      requires std::constructible_from<Vec<3, float>, std::iter_value_t<InputIt>>
    constexpr Surfel(InputIt first, InputIt last)
    {
        add(first, last);
    }
 
    template <std::ranges::input_range Points>
      requires std::constructible_from<Vec<3, float>, std::ranges::range_value_t<Points>>
    constexpr Surfel(Points const& points) : Surfel(std::cbegin(points), std::cend(points))
    {
    }
 
    constexpr Surfel(std::initializer_list<Vec<3, float>> points)
        : Surfel(begin(points), end(points))
    {
    }
 
    [[nodiscard]] constexpr bool operator==(Surfel const&) const noexcept = default;
 
    //
    // Empty
    //
 
    [[nodiscard]] constexpr bool empty() const noexcept { return 0 == num_points_; }
 
    //
    // Add
    //
 
    constexpr Surfel& operator+=(Surfel const& rhs) noexcept
    {
        add(rhs);
        return *this;
    }
 
    constexpr Surfel& operator+=(Vec<3, float> rhs) noexcept
    {
        add(rhs);
        return *this;
    }
 
    [[nodiscard]] friend Surfel operator+(Surfel lhs, Surfel const& rhs) noexcept
    {
        return lhs += rhs;
    }
 
    [[nodiscard]] friend Surfel operator+(Surfel lhs, Vec<3, float> rhs) noexcept
    {
        return lhs += rhs;
    }
 
    constexpr void add(Surfel const& surfel) noexcept
    {
        if (empty()) {
            *this = surfel;
            return;
        }
        add(toDouble(surfel.sum_squares_), cast<double>(surfel.sum_), surfel.num_points_);
    }
 
    constexpr void add(Vec<3, float> point) noexcept
    {
        add(std::array<double, 6>{}, cast<double>(point), 1);
    }
 
    template <std::input_iterator InputIt>
      requires std::constructible_from<Vec<3, float>, std::iter_value_t<InputIt>>
    constexpr void add(InputIt first, InputIt last)
    {
        if (first == last) {
            return;
        }
        auto [ss, s, n] = batchScatter(first, last);
        add(ss, s, n);
    }
 
    template <std::ranges::input_range Points>
      requires std::constructible_from<Vec<3, float>, std::ranges::range_value_t<Points>>
    constexpr void add(Points const& points)
    {
        add(begin(points), end(points));
    }
 
    constexpr void add(std::initializer_list<Vec<3, float>> points) noexcept
    {
        add(begin(points), end(points));
    }
 
    //
    // Remove
    //
 
    constexpr Surfel& operator-=(Surfel const& rhs) noexcept
    {
        remove(rhs);
        return *this;
    }
 
    constexpr Surfel& operator-=(Vec<3, float> rhs) noexcept
    {
        remove(rhs);
        return *this;
    }
 
    [[nodiscard]] friend Surfel operator-(Surfel lhs, Surfel const& rhs) noexcept
    {
        return lhs -= rhs;
    }
 
    [[nodiscard]] friend Surfel operator-(Surfel lhs, Vec<3, float> rhs) noexcept
    {
        return lhs -= rhs;
    }
 
    constexpr void remove(Surfel const& surfel) noexcept
    {
        remove(toDouble(surfel.sum_squares_), cast<double>(surfel.sum_), surfel.num_points_);
    }
 
    constexpr void remove(Vec<3, float> point) noexcept
    {
        remove(std::array<double, 6>{}, cast<double>(point), 1);
    }
 
    template <std::input_iterator InputIt>
      requires std::constructible_from<Vec<3, float>, std::iter_value_t<InputIt>>
    constexpr void remove(InputIt first, InputIt last)
    {
        if (first == last) {
            return;
        }
        auto [ss, s, n] = batchScatter(first, last);
        remove(ss, s, n);
    }
 
    template <std::ranges::input_range Points>
      requires std::constructible_from<Vec<3, float>, std::ranges::range_value_t<Points>>
    constexpr void remove(Points const& points)
    {
        remove(begin(points), end(points));
    }
 
    constexpr void remove(std::initializer_list<Vec<3, float>> points) noexcept
    {
        remove(begin(points), end(points));
    }
 
    //
    // Clear
    //
 
    constexpr void clear() noexcept
    {
        sum_squares_ = {};
        sum_         = {};
        num_points_  = {};
    }
 
    //
    // Get mean
    //
 
    [[nodiscard]] constexpr Vec<3, double> mean() const noexcept
    {
        return cast<double>(sum_) / static_cast<double>(num_points_);
    }
 
    //
    // Get covariance
    //
 
    [[nodiscard]] constexpr Mat<3, 3, double> covariance() const noexcept
    {
        double const n = static_cast<double>(num_points_ - 1);
        auto const   s = [&](std::size_t i) noexcept {
      return static_cast<double>(sum_squares_[i]) / n;
        };
        return Mat<3, 3, double>{s(0), s(1), s(2), s(1), s(3), s(4), s(2), s(4), s(5)};
    }
 
    //
    // Get normal
    //
 
    [[nodiscard]] constexpr Vec3d normal() const noexcept
    {
        return eigenVectors(covariance())[0];
    }
 
    //
    // Get planarity
    //
 
    [[nodiscard]] constexpr double planarity() const noexcept
    {
        auto const e = eigenValues(covariance());
        return 2.0 * (e[1] - e[0]) / (e[0] + e[1] + e[2]);
    }
 
    //
    // Get num points
    //
 
    [[nodiscard]] constexpr std::uint32_t numPoints() const noexcept { return num_points_; }
 
    //
    // Get sum
    //
 
    [[nodiscard]] constexpr Vec<3, float> sum() const noexcept { return sum_; }
 
    //
    // Get sum squares
    //
 
    [[nodiscard]] constexpr std::array<float, 6> sumSquares() const noexcept
    {
        return sum_squares_;
    }
 
 protected:
    //
    // Helpers
    //
 
    [[nodiscard]] static constexpr std::array<double, 6> toDouble(
        std::array<float, 6> const& arr) noexcept
    {
        std::array<double, 6> result;
        std::ranges::transform(arr, result.begin(),
                               [](float v) { return static_cast<double>(v); });
        return result;
    }
 
    static constexpr void toFloat(std::array<double, 6> const& src,
                                  std::array<float, 6>&        dst) noexcept
    {
        std::ranges::transform(src, dst.begin(),
                               [](double v) { return static_cast<float>(v); });
    }
 
    static constexpr void addOuter(std::array<double, 6>& ss, Vec<3, double> const& v,
                                   double scale) noexcept
    {
        ss[0] += v[0] * v[0] * scale;
        ss[1] += v[0] * v[1] * scale;
        ss[2] += v[0] * v[2] * scale;
        ss[3] += v[1] * v[1] * scale;
        ss[4] += v[1] * v[2] * scale;
        ss[5] += v[2] * v[2] * scale;
    }
 
    template <std::input_iterator InputIt>
      requires std::constructible_from<Vec<3, double>, std::iter_value_t<InputIt>>
    [[nodiscard]] static constexpr std::tuple<std::array<double, 6>, Vec<3, double>,
                                              std::uint32_t>
    batchScatter(InputIt first, InputIt last) noexcept
    {
        std::array<double, 6> ss{};
        Vec<3, double>        s{};
        std::uint32_t         n{};
 
        for (; first != last; ++first) {
            Vec<3, double> const p(*first);
            ss[0] += p[0] * p[0];
            ss[1] += p[0] * p[1];
            ss[2] += p[0] * p[2];
            ss[3] += p[1] * p[1];
            ss[4] += p[1] * p[2];
            ss[5] += p[2] * p[2];
            s += p;
            ++n;
        }
 
        double const nd = static_cast<double>(n);
        ss[0] -= s[0] * s[0] / nd;
        ss[1] -= s[0] * s[1] / nd;
        ss[2] -= s[0] * s[2] / nd;
        ss[3] -= s[1] * s[1] / nd;
        ss[4] -= s[1] * s[2] / nd;
        ss[5] -= s[2] * s[2] / nd;
 
        return {ss, s, n};
    }
 
    //
    // Private add/remove
    //
 
    constexpr void add(std::array<double, 6> ss, Vec<3, double> const& sum,
                       std::uint32_t num_points)
    {
        if (empty()) {
            toFloat(ss, sum_squares_);
            sum_        = cast<float>(sum);
            num_points_ = num_points;
            return;
        }
 
        double const         n   = static_cast<double>(num_points_);
        double const         n_o = static_cast<double>(num_points);
        Vec<3, double> const s(sum_);
 
        std::ranges::transform(ss, sum_squares_, ss.begin(),
                               [](double a, float b) { return a + static_cast<double>(b); });
        addOuter(ss, s * n_o - sum * n, 1.0 / (n * n_o * (n + n_o)));
 
        toFloat(ss, sum_squares_);
        sum_ = cast<float>(s + sum);
        num_points_ += num_points;
    }
 
    constexpr void remove(std::array<double, 6> ss, Vec<3, double> const& sum,
                          std::uint32_t num_points)
    {
        if (num_points >= num_points_) {
            clear();
            return;
        }
 
        double const         n_a   = static_cast<double>(num_points_ - num_points);
        double const         n_b   = static_cast<double>(num_points);
        Vec<3, double> const sum_a = cast<double>(sum_) - sum;
 
        std::ranges::transform(sum_squares_, ss, ss.begin(),
                               [](float a, double b) { return static_cast<double>(a) - b; });
        addOuter(ss, sum_a * n_b - sum * n_a, -1.0 / (n_a * n_b * (n_a + n_b)));
 
        toFloat(ss, sum_squares_);
        sum_ = cast<float>(sum_a);
        num_points_ -= num_points;
    }
 
    //
    // Data members
    //
 
    std::array<float, 6> sum_squares_{};
 
    Vec<3, float> sum_{};
 
    std::uint32_t num_points_{};
};
 
}  // namespace ufo
 
template <>
struct std::formatter<ufo::Surfel> {
    constexpr auto parse(std::format_parse_context& ctx) const { return ctx.begin(); }
 
    auto format(ufo::Surfel const& s, std::format_context& ctx) const
    {
        if (s.empty()) {
            return std::format_to(ctx.out(), "Surfel{{empty}}");
        }
        auto const m = s.mean();
        return std::format_to(ctx.out(), "Surfel{{n={}, mean=({}, {}, {})}}", s.numPoints(),
                              m[0], m[1], m[2]);
    }
};
 
namespace ufo
{
inline std::ostream& operator<<(std::ostream& out, Surfel const& s)
{
    return out << std::format("{}", s);
}
}  // namespace ufo
#endif  // UFO_CORE_SURFEL_HPP