From 0434a7d9045eaaf718a83d1ea7868897e4c47a43 Mon Sep 17 00:00:00 2001 From: lisyarus Date: Thu, 8 Dec 2022 17:41:48 +0300 Subject: [PATCH] Implement Barnes-Hut algorithm for electron crystal simulation --- examples/electron_crystal.cpp | 175 +++++++++++++++++++++++++++------- 1 file changed, 140 insertions(+), 35 deletions(-) diff --git a/examples/electron_crystal.cpp b/examples/electron_crystal.cpp index f35179ba..de18c294 100644 --- a/examples/electron_crystal.cpp +++ b/examples/electron_crystal.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include @@ -35,6 +36,7 @@ struct main_scene float potential = 1000.f; float step = 1e-4f; + float multipole_threshold = 2.f; random::generator rng{random::device{}}; @@ -77,7 +79,7 @@ main_scene::main_scene(ui::controller & ui_controller) count_value_label->set_halign(ui::label::halignment::center); auto count_slider = element_factory.make_slider(); - count_slider->set_value_range({1, 2000}, false); + count_slider->set_value_range({1, 5000}, false); count_slider->on_value_changed([this, count_value_label](int value){ count_value_label->set_text(util::to_string(value)); @@ -85,7 +87,7 @@ main_scene::main_scene(ui::controller & ui_controller) points.resize(value); else { - float radius = std::max(0.5f, std::sqrt(points.size() / 1000.f)); + float radius = std::max(0.5f, std::sqrt((points.size() + value) / 1000.f)); random::uniform_sphere_point_distribution d({0.f, 0.f}, radius * 1.25f); while (points.size() < value) @@ -94,21 +96,6 @@ main_scene::main_scene(ui::controller & ui_controller) }); count_slider->set_value(100); - auto potential_name_label = element_factory.make_label("Potential:"); - potential_name_label->set_valign(ui::label::valignment::center); - potential_name_label->set_halign(ui::label::halignment::left); - auto potential_value_label = element_factory.make_label(""); - potential_value_label->set_valign(ui::label::valignment::center); - potential_value_label->set_halign(ui::label::halignment::center); - - auto potential_slider = element_factory.make_slider(); - potential_slider->set_value_range({0, 100}, false); - potential_slider->on_value_changed([this, potential_value_label](int value){ - potential = value * 10.f; - potential_value_label->set_text(util::to_string(potential)); - }); - potential_slider->set_value(100); - auto step_name_label = element_factory.make_label("Step:"); step_name_label->set_valign(ui::label::valignment::center); step_name_label->set_halign(ui::label::halignment::left); @@ -124,23 +111,38 @@ main_scene::main_scene(ui::controller & ui_controller) }); step_slider->set_value(11); + auto precision_name_label = element_factory.make_label("Precision:"); + precision_name_label->set_valign(ui::label::valignment::center); + precision_name_label->set_halign(ui::label::halignment::left); + auto precision_value_label = element_factory.make_label(""); + precision_value_label->set_valign(ui::label::valignment::center); + precision_value_label->set_halign(ui::label::halignment::center); + + auto precision_slider = element_factory.make_slider(); + precision_slider->set_value_range({5, 50}, false); + precision_slider->on_value_changed([this, precision_value_label](int value){ + multipole_threshold = value * 0.1f; + precision_value_label->set_text(util::to_string(std::setprecision(1), std::fixed, multipole_threshold)); + }); + precision_slider->set_value(15); + layout->set_size(3, 3); layout->set_column_weight(0, 0.5f); layout->set_column_weight(1, 0.5f); layout->set(0, 0, count_name_label); layout->set(0, 1, count_value_label); layout->set(0, 2, count_slider); - layout->set(1, 0, potential_name_label); - layout->set(1, 1, potential_value_label); - layout->set(1, 2, potential_slider); - layout->set(2, 0, step_name_label); - layout->set(2, 1, step_value_label); - layout->set(2, 2, step_slider); + layout->set(1, 0, step_name_label); + layout->set(1, 1, step_value_label); + layout->set(1, 2, step_slider); + layout->set(2, 0, precision_name_label); + layout->set(2, 1, precision_value_label); + layout->set(2, 2, precision_slider); ui::style style; style.font = ui::make_default_9x12_font(); style.text_scale = 2; - style.bg_color = gfx::color_rgba{0, 0, 127, 255}; + style.bg_color = gfx::color_rgba{127, 127, 191, 255}; style.fg_color = gfx::color_rgba{255, 255, 255, 255}; style.action_color = gfx::color_rgba{0, 0, 255, 255}; style.highlight_color = gfx::color_rgba{0, 255, 255, 255}; @@ -151,6 +153,17 @@ main_scene::main_scene(ui::controller & ui_controller) set_ui(root); } +struct node +{ + static constexpr std::uint32_t null = -1; + + float size; + float mass = 0.f; + geom::point center = geom::point::zero(); + + std::uint32_t children[2][2] {{null, null}, {null, null}}; +}; + void main_scene::update() { ui_scene::update(); @@ -159,6 +172,93 @@ void main_scene::update() const int iterations = 1; auto const origin = geom::point::zero(); + std::vector nodes; + std::uint32_t root; + + auto add_node = util::recursive([&](auto && self, geom::box const & bbox, auto begin, auto end) -> std::uint32_t + { + if (begin == end) + return node::null; + + std::uint32_t id = nodes.size(); + nodes.emplace_back().size = bbox[0].length(); + + if (end - begin > 1) + { + auto middle_x = bbox[0].center(); + auto middle_y = bbox[1].center(); + + auto mid_it = std::partition(begin, end, [middle_x](auto const & p){ return p[0] < middle_x; }); + auto left_mid_it = std::partition(begin, mid_it, [middle_y](auto const & p){ return p[1] < middle_y; }); + auto right_mid_it = std::partition(mid_it, end, [middle_y](auto const & p){ return p[1] < middle_y; }); + + nodes[id].children[0][0] = self(geom::box{{{bbox[0].min, middle_x}, {bbox[1].min, middle_y}}}, begin, left_mid_it); + nodes[id].children[0][1] = self(geom::box{{{bbox[0].min, middle_x}, {middle_y, bbox[1].max}}}, left_mid_it, mid_it); + nodes[id].children[1][0] = self(geom::box{{{middle_x, bbox[0].max}, {bbox[1].min, middle_y}}}, mid_it, right_mid_it); + nodes[id].children[1][1] = self(geom::box{{{middle_x, bbox[0].max}, {middle_y, bbox[1].max}}}, right_mid_it, end); + + geom::vector sum{0.f, 0.f}; + for (int x : {0, 1}) + { + for (int y : {0, 1}) + { + auto cid = nodes[id].children[x][y]; + if (cid != node::null) + { + nodes[id].mass += nodes[cid].mass; + sum += (nodes[cid].center - origin) * nodes[cid].mass; + } + } + } + + nodes[id].center = origin + sum / nodes[id].mass; + } + else + { + nodes[id].mass = 1.f; + nodes[id].center = *begin; + } + + return id; + }); + + { + float max_size = 0.f; + for (auto const & p : points) + geom::make_max(max_size, geom::distance(p, origin)); + + prof::profiler prof("tree"); + root = add_node(geom::box{{{-max_size, max_size}, {-max_size, max_size}}}, points.begin(), points.end()); + } + + auto force_at = util::recursive([&](auto && self, auto const & p, auto id){ + auto const & n = nodes[id]; + + if (n.center == p) + return geom::vector{0.f, 0.f}; + + auto d = p - n.center; + + auto l = geom::length(d); + + if (n.mass == 1.f || l > n.size * multipole_threshold) + return n.mass * d / std::pow(l, 3.f); + + geom::vector sum{0.f, 0.f}; + + for (int x : {0, 1}) + { + for (int y : {0, 1}) + { + auto cid = nodes[id].children[x][y]; + if (cid != node::null) + sum += self(p, cid); + } + } + + return sum; + }); + std::vector> delta; for (int iteration = 0; iteration < iterations; ++iteration) @@ -169,18 +269,23 @@ void main_scene::update() for (std::size_t i = 0; i < points.size(); ++i) delta[i] += 2.f * potential * (origin - points[i]); + // Exact O(n^2) method +// for (std::size_t i = 0; i < points.size(); ++i) +// { +// for (std::size_t j = i + 1; j < points.size(); ++j) +// { +// auto d = points[i] - points[j]; + +// d /= std::pow(geom::length(d), 3.f); + +// delta[i] += d; +// delta[j] -= d; +// } +// } + + // Barnes-Hut algorithm for (std::size_t i = 0; i < points.size(); ++i) - { - for (std::size_t j = i + 1; j < points.size(); ++j) - { - auto d = points[i] - points[j]; - - d /= std::pow(geom::length(d), 3.f); - - delta[i] += d; - delta[j] -= d; - } - } + delta[i] += force_at(points[i], root); for (std::size_t i = 0; i < points.size(); ++i) points[i] += delta[i] * step * dt;