implement barnes-hut particle repulsion using octree

src/octree.cpp

@@ -0,0 +1,148 @@
+#include "octree.hpp"
+#include "config.hpp"
+
+#include <raymath.h>
+
+auto Octree::CreateNode(const Vector3 &box_min, const Vector3 &box_max) -> int {
+  OctreeNode node;
+  node.box_min = box_min;
+  node.box_max = box_max;
+  nodes.push_back(node);
+
+  return nodes.size() - 1;
+}
+
+auto Octree::GetOctant(int node_idx, const Vector3 &pos) -> int {
+  OctreeNode &node = nodes[node_idx];
+  Vector3 center = Vector3((node.box_min.x + node.box_max.x) / 2.0,
+                           (node.box_min.y + node.box_max.y) / 2.0,
+                           (node.box_min.z + node.box_max.z) / 2.0);
+
+  // The octant is encoded as a 3-bit integer "zyx". The node area is split
+  // along all 3 axes, if a position is right of an axis, this bit is set to 1.
+  // If a position is right of the x-axis and y-axis and left of the z-axis, the
+  // encoded octant is "011".
+  int octant = 0;
+  if (pos.x >= center.x) {
+    octant |= 1;
+  }
+  if (pos.y >= center.y) {
+    octant |= 2;
+  }
+  if (pos.z >= center.z) {
+    octant |= 4;
+  }
+
+  return octant;
+}
+
+auto Octree::GetChildBounds(int node_idx, int octant)
+    -> std::pair<Vector3, Vector3> {
+  OctreeNode &node = nodes[node_idx];
+  Vector3 center = Vector3((node.box_min.x + node.box_max.x) / 2.0,
+                           (node.box_min.y + node.box_max.y) / 2.0,
+                           (node.box_min.z + node.box_max.z) / 2.0);
+
+  Vector3 min;
+  Vector3 max;
+
+  // If (octant & 1), the octant is to the right of the node region's x-axis
+  // (see GetOctant). This means the left bound is the x-axis and the right
+  // bound the node's region max.
+  min.x = (octant & 1) ? center.x : node.box_min.x;
+  max.x = (octant & 1) ? node.box_max.x : center.x;
+  min.y = (octant & 2) ? center.y : node.box_min.y;
+  max.y = (octant & 2) ? node.box_max.y : center.y;
+  min.z = (octant & 4) ? center.z : node.box_min.z;
+  max.z = (octant & 4) ? node.box_max.z : center.z;
+
+  return std::make_pair(min, max);
+}
+
+auto Octree::Insert(int node_idx, int mass_id, const Vector3 &pos, float mass)
+    -> void {
+  OctreeNode &node = nodes[node_idx];
+
+  if (node.leaf && node.mass_id == -1) {
+    // We can place the particle in the empty leaf
+    node.mass_id = mass_id;
+    node.mass_center = pos;
+    node.mass_total = mass;
+    return;
+  }
+
+  if (node.leaf) {
+    // The leaf is occupied, we need to subdivide
+    int existing_id = node.mass_id;
+    Vector3 existing_pos = node.mass_center;
+    float existing_mass = node.mass_total;
+    node.mass_id = -1;
+    node.leaf = false;
+
+    // Re-add the existing mass into a new empty leaf (see above)
+    int oct = GetOctant(node_idx, existing_pos);
+    if (node.children[oct] == -1) {
+      auto [min, max] = GetChildBounds(node_idx, oct);
+      node.children[oct] = CreateNode(min, max);
+    }
+    Insert(node.children[oct], existing_id, existing_pos, existing_mass);
+  }
+
+  // Insert the new mass
+  int oct = GetOctant(node_idx, pos);
+  if (nodes[node_idx].children[oct] == -1) {
+    auto [min, max] = GetChildBounds(node_idx, oct);
+    nodes[node_idx].children[oct] = CreateNode(min, max);
+  }
+  Insert(nodes[node_idx].children[oct], mass_id, pos, mass);
+
+  // Update the center of mass
+  node = nodes[node_idx];
+  float new_mass = node.mass_total + mass;
+  node.mass_center.x =
+      (node.mass_center.x * node.mass_total + pos.x) / new_mass;
+  node.mass_center.y =
+      (node.mass_center.y * node.mass_total + pos.y) / new_mass;
+  node.mass_center.z =
+      (node.mass_center.z * node.mass_total + pos.z) / new_mass;
+  node.mass_total = new_mass;
+}
+
+auto Octree::CalculateForce(int node_idx, const Vector3 &pos) -> Vector3 {
+  if (node_idx < 0) {
+    return Vector3Zero();
+  }
+
+  OctreeNode &node = nodes[node_idx];
+  if (node.mass_total == 0.0f) {
+    return Vector3Zero();
+  }
+
+  Vector3 diff = Vector3Subtract(pos, node.mass_center);
+  float dist_sq = diff.x * diff.x + diff.y * diff.y + diff.z * diff.z;
+
+  // Softening
+  dist_sq += SOFTENING;
+
+  float size = node.box_max.x - node.box_min.x;
+
+  // Barnes-Hut
+  if (node.leaf || (size * size / dist_sq) < (THETA * THETA)) {
+    float dist = std::sqrt(dist_sq);
+    float force_mag = REPULSION_FORCE * node.mass_total / dist_sq;
+
+    return Vector3Scale(diff, force_mag / dist);
+  }
+
+  // Collect child forces
+  Vector3 force = Vector3Zero();
+  for (int i = 0; i < 8; ++i) {
+    if (node.children[i] >= 0) {
+      Vector3 child_force = CalculateForce(node.children[i], pos);
+
+      force = Vector3Add(force, child_force);
+    }
+  }
+
+  return force;
+}

src/physics.cpp

@@ -2,13 +2,18 @@
 #include "config.hpp"
 
 #include <algorithm>
-#include <numeric>
+#include <cfloat>
 #include <raylib.h>
 #include <raymath.h>
+#include <tracy/Tracy.hpp>
 #include <unordered_map>
 #include <utility>
 #include <vector>
 
+#ifndef BARNES_HUT
+#include <numeric>
+#endif
+
 auto Mass::ClearForce() -> void { force = Vector3Zero(); }
 
 auto Mass::CalculateVelocity(const float delta_time) -> void {
@@ -121,7 +126,9 @@ auto MassSpringSystem::AddSpring(const State &massA, const State &massB,
 auto MassSpringSystem::Clear() -> void {
   masses.clear();
   springs.clear();
+#ifndef BARNES_HUT
   InvalidateGrid();
+#endif
 }
 
 auto MassSpringSystem::ClearForces() -> void {
@@ -136,8 +143,52 @@ auto MassSpringSystem::CalculateSpringForces() -> void {
   }
 }
 
+#ifdef BARNES_HUT
+auto MassSpringSystem::BuildOctree() -> void {
+  octree.nodes.clear();
+  octree.nodes.reserve(masses.size() * 2);
+
+  // Compute bounding box around all masses
+  Vector3 min = Vector3(FLT_MAX, FLT_MAX, FLT_MAX);
+  Vector3 max = Vector3(-FLT_MAX, -FLT_MAX, -FLT_MAX);
+  for (const auto &[state, mass] : masses) {
+    min.x = std::min(min.x, mass.position.x);
+    max.x = std::max(max.x, mass.position.x);
+    min.y = std::min(min.y, mass.position.y);
+    max.y = std::max(max.y, mass.position.y);
+    min.z = std::min(min.z, mass.position.z);
+    max.z = std::max(max.z, mass.position.z);
+  }
+
+  // Pad the bounding box
+  float pad = 1.0;
+  min = Vector3Subtract(min, Vector3Scale(Vector3One(), pad));
+  max = Vector3Add(max, Vector3Scale(Vector3One(), pad));
+
+  // Make it cubic (so subdivisions are balanced)
+  float max_extent = std::max({max.x - min.x, max.y - min.y, max.z - min.z});
+  max = Vector3Add(min, Vector3Scale(Vector3One(), max_extent));
+
+  // Root node spans the entire area
+  int root = octree.CreateNode(min, max);
+
+  // Use a vector of pointers to the masses, because we can't parallelize the
+  // range-based for loop over the masses unordered_map using OpenMP.
+  mass_pointers.clear();
+  mass_pointers.reserve(masses.size());
+  for (auto &[state, mass] : masses) {
+    mass_pointers.push_back(&mass);
+  }
+  for (int i = 0; i < mass_pointers.size(); ++i) {
+    octree.Insert(root, i, mass_pointers[i]->position, mass_pointers[i]->mass);
+  }
+}
+
+#else
+
 auto MassSpringSystem::BuildUniformGrid() -> void {
-  // Collect pointers to all masses
+  // Use a vector of pointers to masses, because we can't parallelize the
+  // range-based for loop over the masses unordered_map using OpenMP.
   mass_pointers.clear();
   mass_pointers.reserve(masses.size());
   for (auto &[state, mass] : masses) {
@@ -173,8 +224,24 @@ auto MassSpringSystem::BuildUniformGrid() -> void {
     cell_ids[i] = cell_id(mass_pointers[mass_indices[i]]->position);
   }
 }
+#endif
 
 auto MassSpringSystem::CalculateRepulsionForces() -> void {
+  ZoneScoped;
+#ifdef BARNES_HUT
+  BuildOctree();
+
+// Calculate forces using Barnes-Hut
+#pragma omp parallel for schedule(dynamic, 256)
+  for (int i = 0; i < mass_pointers.size(); ++i) {
+    int root = 0;
+    Vector3 force = octree.CalculateForce(root, mass_pointers[i]->position);
+
+    mass_pointers[i]->force = Vector3Add(mass_pointers[i]->force, force);
+  }
+
+#else
+
   // Refresh grid if necessary
   if (last_build >= REPULSION_GRID_REFRESH ||
       masses.size() != last_masses_count ||
@@ -186,8 +253,8 @@ auto MassSpringSystem::CalculateRepulsionForces() -> void {
   }
   last_build++;
 
-  // TODO: Use Barnes-Hut + Octree
-#pragma omp parallel for
+  // Calculate forces using uniform grid
+#pragma omp parallel for schedule(dynamic, 256)
   // Search the neighboring cells for each mass to calculate repulsion forces
   for (int i = 0; i < masses.size(); ++i) {
     Mass *mass = mass_pointers[mass_indices[i]];
@@ -241,6 +308,7 @@ auto MassSpringSystem::CalculateRepulsionForces() -> void {
 
     mass->force = Vector3Add(mass->force, force);
   }
+#endif
 }
 
 auto MassSpringSystem::VerletUpdate(float delta_time) -> void {
@@ -249,6 +317,7 @@ auto MassSpringSystem::VerletUpdate(float delta_time) -> void {
   }
 }
 
+#ifndef BARNES_HUT
 auto MassSpringSystem::InvalidateGrid() -> void {
   mass_pointers.clear();
   mass_indices.clear();
@@ -257,3 +326,4 @@ auto MassSpringSystem::InvalidateGrid() -> void {
   last_masses_count = 0;
   last_springs_count = 0;
 }
+#endif