Barnes-Hut \$n\$-body simulation (3D) in C++

Question

I have implemented the Barnes-Hut algorithm for \$n\$-body-simulations (in this case "sort-of" 3D-gravity - see below). I'd appreciate any comments for improving my code (especially concerning speedups). I used the pseudo-code of this site https://beltoforion.de/en/barnes-hut-galaxy-simulator/ to implement it (nice for the "basic layout").

The code is a little bit messy. I use an array of the tree-struct of size 3 * NumParticles (I think this is enough and for most cases additional memory allocation should not be neccessary - although I should probably implement it). TreeBase is just the pointer to the first element of this array, Tree in the function head is a pointer to some child (of a child of a child...), which is at some place (TreeBase.Oct[n].Oct[m].Oct[...]) in the array. TreeCount is a global variable holding the next free element in the array.

struct BHtree
{
    unsigned int NrParticles;
    Vec3 centerOfMass;
    double mass;
    Vec3 Pos;
    double size;
    BHtree* Oct[8];
    unsigned int Okt[8]; // Ignore these last 3 lines (I think, I need them later for implementing this on GPU)
    unsigned int NodeID;
    int Parent;
};

Building the tree:

void InitTree(BHtree* Tree, double* Particles, unsigned int count)
{
    double minX = Particles[0], maxX = Particles[0], minY = Particles[1], maxY = Particles[1], minZ = Particles[2], maxZ = Particles[2], dia;
    for (int n = 0; n < count; n++)
    {
        if (Particles[n * 3 + 0] < minX)
            minX = Particles[n * 3 + 0];
        if (Particles[n * 3 + 0] > maxX)
            maxX = Particles[n * 3 + 0];
        if (Particles[n * 3 + 1] < minY)
            minY = Particles[n * 3 + 1];
        if (Particles[n * 3 + 1] > maxY)
            maxY = Particles[n * 3 + 1];
        if (Particles[n * 3 + 2] < minZ)
            minZ = Particles[n * 3 + 2];
        if (Particles[n * 3 + 2] > maxZ)
            maxZ = Particles[n * 3 + 2];
    }
    dia = maxX - minX;
    if ((maxY - minY) > dia)
        dia = maxY - minY;
    if ((maxZ - minZ) > dia)
        dia = maxZ - minZ;
    Tree[0].NrParticles = 0;
    Tree[0].Pos.x = (maxX + minX) / 2.0;
    Tree[0].Pos.y = (maxY + minY) / 2.0;
    Tree[0].Pos.z = (maxZ + minZ) / 2.0;
    Tree[0].size = dia;
    for (int n = 0; n < 8; n++)
    {
        Tree[0].Oct[n] = nullptr;
        Tree[0].Okt[n] = 0;
    }
    TreeCount = 1; // Next free space
    Tree->Parent = -1;
    Tree->NodeID = 0;
}

unsigned char GetOct(Vec3 TreePos, double TreeSize, Vec3 Particle)
{
    unsigned char oct = 0;
    if (Particle.x > TreePos.x)
        oct++;
    if (Particle.y > TreePos.y)
        oct += 2;
    if (Particle.z > TreePos.z)
        oct += 4;

    return oct;
}

BHtree* CreateSubNode(BHtree* Tree, unsigned char oct)
{
    BHtree* SubNode = &TreeBase[TreeCount];

    //Initialize subnode
    SubNode[0].NodeID = TreeCount++;
    SubNode[0].NrParticles = 0;
    SubNode[0].size = Tree[0].size * 0.5;

    if (oct & 0b001)
        SubNode[0].Pos.x = Tree->Pos.x + 0.5 * SubNode[0].size;
    else
        SubNode[0].Pos.x = Tree->Pos.x - 0.5 * SubNode[0].size;
    if (oct & 0b010)
        SubNode[0].Pos.y = Tree->Pos.y + 0.5 * SubNode[0].size;
    else
        SubNode[0].Pos.y = Tree->Pos.y - 0.5 * SubNode[0].size;
    if (oct & 0b100)
        SubNode[0].Pos.z = Tree->Pos.z + 0.5 * SubNode[0].size;
    else
        SubNode[0].Pos.z = Tree->Pos.z - 0.5 * SubNode[0].size;

    for (int n = 0; (n < 8); n++)
    {
        SubNode[0].Oct[n] = nullptr;
        SubNode[0].Okt[n] = 0;
    }
    Tree[0].Oct[oct] = SubNode;
    Tree[0].Okt[oct] = TreeCount - 1;

    SubNode[0].NrParticles = 0;
    SubNode[0].Parent = Tree[0].NodeID;
    return SubNode;
}

void InsertToNode(BHtree* Tree, Vec3 Particle, double mass)
{
    unsigned char oct;
    BHtree* NewNode;
    if (Tree[0].NrParticles > 1)
    {
        oct = GetOct(Tree[0].Pos, Tree[0].size, Particle);

        if (Tree[0].Oct[oct] == nullptr) // If subnode does not exist, create new one
            NewNode = CreateSubNode(Tree, oct);
        else
            NewNode = Tree[0].Oct[oct];

        InsertToNode(NewNode, Particle, mass);
    }
    else if (Tree[0].NrParticles == 1)
    {
        // existing particle
        Vec3 Particle2 = { Tree[0].centerOfMass.x, Tree[0].centerOfMass.y, Tree[0].centerOfMass.z };
        oct = GetOct(Tree[0].Pos, Tree[0].size, Particle2);

        if (Tree[0].Oct[oct] == nullptr) // If subnode does not exist, create new one
            NewNode = CreateSubNode(Tree, oct);
        else
            NewNode = Tree[0].Oct[oct];
        InsertToNode(NewNode, Particle2, Tree->mass);

        // new particle
        oct = GetOct(Tree[0].Pos, Tree[0].size, Particle);

        if (Tree[0].Oct[oct] == nullptr) // If subnode does not exist, create new one
            NewNode = CreateSubNode(Tree, oct);
        else
            NewNode = Tree[0].Oct[oct];
        InsertToNode(NewNode, Particle, mass);
    }
    else
    {
        Tree[0].centerOfMass = Particle;
        Tree[0].mass = mass;
    }

    Tree[0].NrParticles++;
}

void ComputeMassDistribution(BHtree* Tree)
{
    if (Tree[0].NrParticles > 1)
    {
        Tree[0].centerOfMass = { 0, 0, 0 };
        Tree[0].mass = 0;
        for (int n = 0; n < 8; n++)
        {
            if (Tree[0].Oct[n] != nullptr)
            {
                ComputeMassDistribution(Tree[0].Oct[n]);
                Tree[0].mass += Tree[0].Oct[n]->mass;
                Tree[0].centerOfMass.x += Tree[0].Oct[n]->mass * Tree[0].Oct[n]->centerOfMass.x;
                Tree[0].centerOfMass.y += Tree[0].Oct[n]->mass * Tree[0].Oct[n]->centerOfMass.y;
                Tree[0].centerOfMass.z += Tree[0].Oct[n]->mass * Tree[0].Oct[n]->centerOfMass.z;
            }
        }
        Tree->centerOfMass.x /= Tree->mass;
        Tree->centerOfMass.y /= Tree->mass;
        Tree->centerOfMass.z /= Tree->mass;
    }
}

The actual force calculation:

Vec3 CalcForceLJ(BHtree* Tree, Vec3 Particle, double delta)
{
    Vec3 a, diff;
    double dist3, dist2;
    a.x = a.y = a.z = 0;

    if (Tree[0].NrParticles == 1)
    {
        if (Particle.x != Tree[0].centerOfMass.x || Particle.y != Tree[0].centerOfMass.y || Particle.z != Tree[0].centerOfMass.z)
        {
            diff.x = Tree[0].centerOfMass.x - Particle.x;
            diff.y = Tree[0].centerOfMass.y - Particle.y;
            diff.z = Tree[0].centerOfMass.z - Particle.z;
            dist2 = diff.x * diff.x + diff.y * diff.y + diff.z * diff.z;
            dist3 = Tree[0].mass / (sqrt(dist2) * dist2) - 0.1 * Tree[0].mass / (dist2 * dist2);

            a.x = diff.x * (dist3);
            a.y = diff.y * (dist3);
            a.z = diff.z * (dist3);
        }
    }
    else
    {
        diff.x = Tree[0].centerOfMass.x - Particle.x;
        diff.y = Tree[0].centerOfMass.y - Particle.y;
        diff.z = Tree[0].centerOfMass.z - Particle.z;
        dist2 = diff.x * diff.x + diff.y * diff.y + diff.z * diff.z;
        if ((Tree[0].size * Tree[0].size / dist2) < (delta * delta))
        {
            dist3 = Tree[0].mass / (sqrt(dist2) * dist2) - 0.1 * Tree[0].mass / (dist2 * dist2);
            a.x = diff.x * dist3;
            a.y = diff.y * dist3;
            a.z = diff.z * dist3;
        }
        else
        {
            Vec3 aTemp;
            for (int n = 0; n < 8; n++)
            {
                if (Tree[0].Oct[n] != nullptr)
                    aTemp = CalcForceLJ(Tree[0].Oct[n], Particle, delta, NumRec);
                else
                    aTemp.x = aTemp.y = aTemp.z = 0;
                a.x += aTemp.x; a.y += aTemp.y; a.z += aTemp.z;
            }
        }
    }
    return a;
}

Note to CalcForceLJ(): I'm using a Lennard-Jones-potential here, not exact gravity. \$1/r^2 - 0.1/r^3\$ in this case (actually \$-1/r^2 + 0.1/r^3\$) is (almost) exactly gravity for "bigger" \$r\$ and is repulsive for smaller \$r\$. When putting in a little friction, everything forms a "Lennard-Jones-ball" after a while. I can then later (without friction) let those "balls" circle each other, which gives nicely what closely rotating stars would do. I will later also do a CalcForce like \$1/(r^2 + 0.001)\$ (0.001 = gravitational softening to avoid the integrator with variable step-size go crazy if particles get to close to each other).

I'm using Boost/odeint (with variable step-size) to integrate the particle system using this as an ODE-function:

void odefun0LJ(std::vector<double> x, std::vector<double>& dxdt, const double)
{
    InitTree(TreeBase, &x[0], NumParticles);
    for (int n = 0; n < NumParticles; n++)
    {
        InsertToNode(TreeBase, { x[n * 3 + 0], x[n * 3 + 1], x[n * 3 + 2] }, 1.0);
    }
    ComputeMassDistribution(TreeBase);

    for (int n = 0; n < NumParticles; n++)
    {
        Vec3 a = CalcForceLJ(TreeBase, { x[n * 3 + 0], x[n * 3 + 1], x[n * 3 + 2] }, 0.5);
    }
}

Thank you in advance for any comments on my code.

Edit:

I removed some lines of code that are irrelavant for the algorithm.

Here are some more declarations not in the code above (I think, I got them all...):

Vec3 is a simple struct

Vec3 {double x, y, z}

std::vector<double> x and std::vector<double> x&dxdt are passed through like this by Boost/odeint. x is defined as x[0], x[3], x[6] ... x-Position of Particle 1, 2, 3..., x[1], x[4], x[7]... y-Position and x[2], x[4], x[8]... z-Position

NumParticles is a global unsigned int.

Answer 1

Use the power of C++

A lot of what you are doing does not make good use of all the facilities C++ provides you. In fact, apart from the std::vector parameters in odefun0LJ(), it looks like plain C code. C++ allows you to write much more generic code that removes a lot of code duplication.

To start with, use a numerical library that provides you with proper mathematical vector types, like Eigen, or if you want to target the GPU as well, perhaps GLM. For example, with GLM, CalculateForceLJ() could look like this:

glm::vec3 CalcForceLJ(BHtree* Tree, glm::vec3 Particle, double delta)
{
    auto diff = Tree->centerOfMass - Particle;
    auto dist = glm::length(diff);

    // Leaf node
    if (Tree->NrParticles == 1) {
        if (dist != 0) {
            auto lennardJones = std::pow(dist, -2) - 0.1 * std::pow(dist, -3);
            return Tree->mass * diff * lennardJones;
        } else {
            return {};
        }
    }

    // We are far away enough to just use its center of mass as an approximation
    if (dist * delta > Tree->size) {
        auto lennardJones = std::pow(dist, -2) - 0.1 * std::pow(dist, -3);
        return Tree->mass * diff * lennardJones;
    }

    // Otherwise, recurse down the octree
    glm::vec3 force = {};

    for (int n = 0; n < 8; n++)
        if (Tree->Oct[n])
            force += Tree->Oct[n];

    return force;
}

Notice how we no longer need to deal with x, y and z individually.

Another improvement would be to have safer memory management; instead of hoping that an array of 3 * NumParticles is enough, you can use a std::vector or std::deque to store the tree nodes, or alternatively make Oct an array of std::unique_ptrs.

Naming things

You are not very consistent in how you name things. First, I see both camelCase and PascalCase being used. Choose one way to write function and variable names and stick with it. If you want a recommendation: use camelCase. Some people prefer to use a different style for their type names, so it is easier to distinguish those from variables and functions. But again, if you do this, do it consistently.

Apart from that, some names are not accurate. Consider:

• BHtree: this is not actually a whole tree, it just is one node in the tree. So a better name might be BHNode (but see below).
• Oct: prefer using the plural for arrays. Furthermore, it's more common in trees to talk about child nodes, so children might be a better name for this array. Use octant for the index into the array children.
• Particles: those are actually just the positions of the particles.
• odefun0LJ(): it's a function, so fun is already redundant, I don't know what the 0 means, and this does not implement an ODE, it just implements one step for the integration. So perhaps better would be to just name it step().

Use -> instead of [0].

If you have a pointer to some object, and it's just a single object, or it's one element of an array but you don't care about the other elements, then use -> instead of [0]. to dereference it.

Make BHtree a proper class

Consider making BHtree a class along with member functions to manipulate the tree. Also, BHtree itself should represent the whole tree, use a separate class to represent nodes in the tree. There are various ways to go about this. You can just start slowly, converting your existing functions to member functions of the appropriate class. Doing this will clean up your code, for example you will not have to write Tree-> or Tree[0]. as much.