Implement Barnes-Hut algorithm for electron crystal simulation

examples/electron_crystal.cpp

@@ -17,6 +17,7 @@
 #include <psemek/ui/grid_layout.hpp>
 #include <psemek/ui/event_interceptor.hpp>
 #include <psemek/util/to_string.hpp>
+#include <psemek/util/recursive.hpp>
 
 #include <vector>
 
@@ -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<float, 2> 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, 0, step_name_label);
-	layout->set(1, 1, potential_value_label);
+	layout->set(1, 1, step_value_label);
-	layout->set(1, 2, potential_slider);
+	layout->set(1, 2, step_slider);
-	layout->set(2, 0, step_name_label);
+	layout->set(2, 0, precision_name_label);
-	layout->set(2, 1, step_value_label);
+	layout->set(2, 1, precision_value_label);
-	layout->set(2, 2, step_slider);
+	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<float, 2> center = geom::point<float, 2>::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<float, 2>::zero();
 
+	std::vector<node> nodes;
+	std::uint32_t root;
+
+	auto add_node = util::recursive([&](auto && self, geom::box<float, 2> 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<float, 2>{{{bbox[0].min, middle_x}, {bbox[1].min, middle_y}}}, begin, left_mid_it);
+			nodes[id].children[0][1] = self(geom::box<float, 2>{{{bbox[0].min, middle_x}, {middle_y, bbox[1].max}}}, left_mid_it, mid_it);
+			nodes[id].children[1][0] = self(geom::box<float, 2>{{{middle_x, bbox[0].max}, {bbox[1].min, middle_y}}}, mid_it, right_mid_it);
+			nodes[id].children[1][1] = self(geom::box<float, 2>{{{middle_x, bbox[0].max}, {middle_y, bbox[1].max}}}, right_mid_it, end);
+
+			geom::vector<float, 2> 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<float, 2>{{{-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<geom::vector<float, 2>> 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)
-		{
+			delta[i] += force_at(points[i], root);
-			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;
-			}
-		}
 
 		for (std::size_t i = 0; i < points.size(); ++i)
 			points[i] += delta[i] * step * dt;