barnes_hut_solver.hpp

Go to the documentation of this file.

#pragma once

 
#include <span>
#include <string_view>
#include <vector>
 
#include "orrery/backend/executor.hpp"
#include "orrery/backend/worker_statistics.hpp"
#include "orrery/core/aligned_allocator.hpp"
#include "orrery/core/softening.hpp"
#include "orrery/core/types.hpp"
#include "orrery/core/vec3_span.hpp"
#include "orrery/solvers/direct_kernel.hpp"
#include "orrery/solvers/force_solver.hpp"
#include "orrery/solvers/interaction_count.hpp"
#include "orrery/solvers/morton.hpp"
#include "orrery/solvers/octree.hpp"
 
namespace orrery::solvers {

struct EvaluationTimings {
    backend::Duration ordering{};

    backend::Duration construction{};

    backend::Duration gathering{};

    backend::Duration traversal{};
};

class BarnesHutSolver final : public ForceSolver {
public:
    explicit BarnesHutSolver(TreeParameters parameters = {}) noexcept
        : parameters_(corrected_parameters(parameters)) {}

    BarnesHutSolver(TreeParameters parameters, core::Softening softening) noexcept
        : parameters_(corrected_parameters(parameters)), softening_(softening) {}

    BarnesHutSolver(TreeParameters parameters, core::Softening softening,
                    backend::Executor& executor) noexcept
        : parameters_(corrected_parameters(parameters)),
          softening_(softening),
          executor_(&executor) {}

    void evaluate(core::Vec3Span<const core::Real> positions, std::span<const core::Real> masses,
                  core::Vec3Span<core::Real> accelerations) override;
 
    [[nodiscard]] std::string_view name() const noexcept override { return "barnes-hut"; }
 
    [[nodiscard]] core::Softening softening() const noexcept override { return softening_; }
 
    [[nodiscard]] InteractionCount interaction_count() const noexcept override { return count_; }
 
    void reset_interaction_count() noexcept override { count_ = {}; }

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

    [[nodiscard]] const Octree& tree() const noexcept { return tree_; }

    [[nodiscard]] const EvaluationTimings& timings() const noexcept { return timings_; }

    [[nodiscard]] const backend::Executor* executor() const noexcept { return executor_; }

    void select_kernel(KernelKind kind) noexcept {
        kernel_ = kernel_available(kind) ? kind : KernelKind::kScalar;
    }

    [[nodiscard]] KernelKind kernel() const noexcept { return kernel_; }
 
private:
    using ComponentArray = std::vector<core::Real, core::AlignedAllocator<core::Real>>;

    void gather(core::Vec3Span<const core::Real> positions, std::span<const core::Real> masses);
 
    TreeParameters parameters_;
    core::Softening softening_;

    backend::Executor* executor_{nullptr};
 
    KernelKind kernel_{fastest_available_kernel()};
 
    MortonOrdering ordering_;
    Octree tree_;

    ComponentArray sorted_x_;
    ComponentArray sorted_y_;
    ComponentArray sorted_z_;
    ComponentArray sorted_mass_;
 
    InteractionCount count_;
    EvaluationTimings timings_;
};
 
} // namespace orrery::solvers