octree.hpp

Go to the documentation of this file.

#pragma once

 
#include <cstdint>
#include <span>
#include <vector>
 
#include "orrery/backend/executor.hpp"
#include "orrery/core/types.hpp"
#include "orrery/core/vec3.hpp"
#include "orrery/core/vec3_span.hpp"
#include "orrery/solvers/morton.hpp"
 
namespace orrery::solvers {

struct Quadrupole {
    core::Real xx{};
    core::Real xy{};
    core::Real xz{};
    core::Real yy{};
    core::Real yz{};
    core::Real zz{};
};

struct TreeNode {
    core::Real mass{};

    core::Vec3 centre_of_mass;

    core::Real acceptance_radius_squared{};

    core::Index next{};

    core::Index first_particle{};

    core::Index particle_count{};
 
    [[nodiscard]] constexpr bool leaf() const noexcept { return particle_count > 0; }
};

struct TreeParameters {
    core::Real opening_angle{static_cast<core::Real>(0.5)};

    core::Index leaf_capacity{32};

    bool quadrupole{false};
};

[[nodiscard]] TreeParameters corrected_parameters(TreeParameters parameters) noexcept;

class Octree {
public:
    void build(core::Vec3Span<const core::Real> positions, std::span<const core::Real> masses,
               std::span<const MortonKey> keys, const BoundingCube& cube,
               const TreeParameters& parameters, backend::Executor* executor);

    [[nodiscard]] std::span<const TreeNode> nodes() const noexcept { return nodes_; }

    [[nodiscard]] std::span<const Quadrupole> quadrupoles() const noexcept { return quadrupoles_; }
 
    [[nodiscard]] bool empty() const noexcept { return nodes_.empty(); }

    [[nodiscard]] const TreeParameters& parameters() const noexcept { return parameters_; }

    [[nodiscard]] core::Index leaf_count() const noexcept { return leaf_count_; }

    [[nodiscard]] unsigned depth() const noexcept { return depth_; }
 
private:
    struct Subtree {
        core::Index begin{};
        core::Index end{};
        core::Index top_index{};
        unsigned level{};
        core::Vec3 centre;
        core::Real size{};
        std::vector<TreeNode> nodes;
        std::vector<Quadrupole> quadrupoles;

        core::Index leaves{};
        unsigned depth{};
    };

    struct BuildInput {
        core::Vec3Span<const core::Real> positions;
        std::span<const core::Real> masses;
        std::span<const MortonKey> keys;
    };

    struct TopCell {
        core::Vec3 centre;
        core::Real size{};
    };

    void descend(const BuildInput& input, core::Index begin, core::Index end, unsigned level,
                 core::Vec3 centre, core::Real size, core::Index threshold);

    void build_subtree(const BuildInput& input, Subtree& subtree) const;

    void emit(const BuildInput& input, Subtree& sink, core::Index begin, core::Index end,
              unsigned level, core::Vec3 centre, core::Real size) const;

    [[nodiscard]] Subtree& next_subtree();

    [[nodiscard]] core::Real acceptance_radius_squared(core::Real size,
                                                       core::Real offset) const noexcept;

    void assemble();
 
    std::vector<TreeNode> nodes_;
    std::vector<Quadrupole> quadrupoles_;
 
    // The five members below hold nothing between builds except their storage,
    // which is the point of them. A solver evaluates millions of times on
    // configurations of one size, so a tree that allocated its working buffers
    // on every build would spend a measurable part of every force evaluation in
    // the allocator and would return the same pages to it immediately
    // afterwards.

    std::vector<Subtree> subtrees_;

    core::Index subtree_count_{};

    std::vector<TreeNode> assembled_;
    std::vector<Quadrupole> assembled_quadrupoles_;

    std::vector<core::Index> shift_;

    std::vector<TopCell> top_cells_;
 
    TreeParameters parameters_;
    core::Index leaf_count_{};
    unsigned depth_{};
};
 
} // namespace orrery::solvers