main.cpp

Go to the documentation of this file.

/// @file apps/nbody/main.cpp
/// @brief Gravitational N-body simulation — raylib visualiser.
///
/// Demonstrates num::ode_verlet (symplectic) vs num::ode_rk4 on three
/// classical scenarios, plus a galaxy-collapse scenario with body merging.
///
/// Controls:
///   1 / 2 / 3 / 4  -- switch scenario
///   V           -- toggle Verlet ↔ RK4
///   SPACE       -- pause / resume
///   R           -- reset current scenario
///   +/-         -- increase / decrease simulation speed
///   F           -- toggle trail fade
 
#include "nbody.hpp"
#include <raylib.h>
#include <deque>
#include <vector>
#include <random>
#include <cmath>
#include <cstdio>
#include <algorithm>
#include <string>
 
// ── NBodySim implementation ───────────────────────────────────────────────────
 
namespace nbody {
 
// Color mapped from body mass: blue-white (low) → yellow (mid) → orange-red (high)
static uint32_t mass_color(double m) {
    const double t = std::min(m / 5.0, 1.0);
    uint8_t r_, g_, b_;
    if (t < 0.5) {
        float u = float(t * 2.0);
        r_ = uint8_t(120 + u * 135);   // 120 → 255
        g_ = uint8_t(200 + u *  38);   // 200 → 238
        b_ = uint8_t(255 - u *  85);   // 255 → 170
    } else {
        float u = float((t - 0.5) * 2.0);
        r_ = 255;
        g_ = uint8_t(238 - u * 136);   // 238 → 102
        b_ = uint8_t(170 - u * 136);   // 170 →  34
    }
    return (uint32_t(r_) << 24) | (uint32_t(g_) << 16) | (uint32_t(b_) << 8) | 0xFFu;
}
 
// Display radius (pixels) and physics radius for a given mass
static constexpr float DISP_SCALE = 6.5f;   // pixels for mass = 1.0
static constexpr float PHYS_SCALE = 0.05f;  // phys units for mass = 1.0
static float disp_r(double m) { return DISP_SCALE * std::cbrt(float(m)); }
static float phys_r(double m) { return PHYS_SCALE * std::cbrt(float(m)); }
 
void NBodySim::make_accel(const num::Vector& pos, num::Vector& acc) const {
    const int nb = n();
    std::fill(acc.begin(), acc.end(), 0.0);
    for (int i = 0; i < nb; ++i) {
        for (int j = 0; j < nb; ++j) {
            if (i == j) continue;
            double dx = pos[2*j]   - pos[2*i];
            double dy = pos[2*j+1] - pos[2*i+1];
            double r2 = dx*dx + dy*dy + soft*soft;
            double r3 = r2 * std::sqrt(r2);
            double fac = G * bodies[j].mass / r3;
            acc[2*i]   += fac * dx;
            acc[2*i+1] += fac * dy;
        }
    }
}
 
double NBodySim::kinetic_energy() const {
    double KE = 0.0;
    for (int i = 0; i < n(); ++i)
        KE += 0.5 * bodies[i].mass * (v[2*i]*v[2*i] + v[2*i+1]*v[2*i+1]);
    return KE;
}
 
double NBodySim::potential_energy() const {
    double PE = 0.0;
    for (int i = 0; i < n(); ++i) {
        for (int j = i+1; j < n(); ++j) {
            double dx = q[2*j] - q[2*i], dy = q[2*j+1] - q[2*i+1];
            double r  = std::sqrt(dx*dx + dy*dy + soft*soft);
            PE -= G * bodies[i].mass * bodies[j].mass / r;
        }
    }
    return PE;
}
 
void NBodySim::step(double dt) {
    if (use_verlet) {
        auto accel = [this](const num::Vector& pos, num::Vector& acc) {
            make_accel(pos, acc);
        };
        auto res = num::ode_verlet(accel, q, v, t, t + dt, dt);
        q = std::move(res.q);
        v = std::move(res.v);
    } else {
        const int n2 = 2 * n();
        num::Vector y(2 * n2);
        for (int i = 0; i < n2; ++i) { y[i] = q[i]; y[n2+i] = v[i]; }
 
        auto rhs = [this, n2](num::real /*tt*/, const num::Vector& yy, num::Vector& dydt) {
            num::Vector pos(n2), acc(n2);
            for (int i = 0; i < n2; ++i) pos[i] = yy[i];
            make_accel(pos, acc);
            for (int i = 0; i < n2; ++i) {
                dydt[i]    = yy[n2+i];
                dydt[n2+i] = acc[i];
            }
        };
 
        auto res = num::ode_rk4(rhs, std::move(y), t, t + dt, dt);
        for (int i = 0; i < n2; ++i) { q[i] = res.y[i]; v[i] = res.y[n2+i]; }
    }
    t += dt;
}
 
std::vector<std::pair<int,int>> NBodySim::check_merges() {
    std::vector<std::pair<int,int>> ops;
    int i = 0;
    while (i < n() && n() > 1) {
        bool merged = false;
        for (int j = i + 1; j < n(); ++j) {
            double dx  = q[2*i] - q[2*j];
            double dy  = q[2*i+1] - q[2*j+1];
            double mrd = bodies[i].phys_radius + bodies[j].phys_radius;
            if (dx*dx + dy*dy >= mrd*mrd) continue;
 
            // Merge j into i — conserve mass and momentum
            double mi = bodies[i].mass, mj = bodies[j].mass, mt = mi + mj;
            q[2*i]   = (mi*q[2*i]   + mj*q[2*j])   / mt;
            q[2*i+1] = (mi*q[2*i+1] + mj*q[2*j+1]) / mt;
            v[2*i]   = (mi*v[2*i]   + mj*v[2*j])   / mt;
            v[2*i+1] = (mi*v[2*i+1] + mj*v[2*j+1]) / mt;
            bodies[i].mass           = mt;
            bodies[i].display_radius = disp_r(mt);
            bodies[i].phys_radius    = phys_r(mt);
            bodies[i].color          = mass_color(mt);
 
            // Swap j with the last body and shrink arrays
            int last = n() - 1;
            ops.push_back({j, last});
            if (j != last) {
                std::swap(bodies[j], bodies[last]);
                q[2*j]   = q[2*last];   q[2*j+1] = q[2*last+1];
                v[2*j]   = v[2*last];   v[2*j+1] = v[2*last+1];
            }
            bodies.pop_back();
 
            // Rebuild q and v at the new (smaller) size
            int nn = n();
            num::Vector nq(2*nn, 0.0), nv(2*nn, 0.0);
            for (int k = 0; k < 2*nn; ++k) { nq[k] = q[k]; nv[k] = v[k]; }
            q = std::move(nq);
            v = std::move(nv);
 
            merged = true;
            break;  // restart j-scan for body i (it may now touch another body)
        }
        if (!merged) ++i;
    }
    return ops;
}
 
void NBodySim::reset(Scenario s) {
    t = 0.0;
    G = 1.0;
    enable_merges = false;
 
    if (s == Scenario::Figure8) {
        soft = 1e-4;
        bodies = {
            {1.0, 10.0f, 0xFF5555FF},
            {1.0, 10.0f, 0x55FF55FF},
            {1.0, 10.0f, 0x5599FFFF},
        };
        q = { -0.97000436,  0.24308753,
               0.97000436, -0.24308753,
               0.0,         0.0        };
        v = {  0.46620368,  0.43236573,
               0.46620368,  0.43236573,
              -0.93240737, -0.86473146 };
 
    } else if (s == Scenario::SolarSystem) {
        soft = 1e-3;
        const double Msun = 1000.0;
        bodies = {
            {Msun, 18.0f, 0xFFDD33FF},
            { 1.0,  4.0f, 0xAAAAAAFF},
            { 1.0,  6.0f, 0x44AAFFFF},
            { 1.0,  5.0f, 0xFF6644FF},
            {10.0, 11.0f, 0xFFAA44FF},
        };
        const double radii[] = {0.0, 0.22, 0.40, 0.60, 1.10};
        const int nb = 5;
        q = num::Vector(2*nb, 0.0);
        v = num::Vector(2*nb, 0.0);
        for (int i = 1; i < nb; ++i) {
            q[2*i]   = radii[i];
            v[2*i+1] = std::sqrt(G * Msun / radii[i]);
        }
 
    } else if (s == Scenario::BinaryPlus) {
        soft = 1e-3;
        const double Mstar = 50.0;
        bodies = {
            {Mstar, 13.0f, 0xFF7744FF},
            {Mstar, 13.0f, 0x44AAFFFF},
            { 0.01,  5.0f, 0xAAFF88FF},
        };
        const double sep   = 0.3;
        const double omega = std::sqrt(G * 2.0 * Mstar / (sep*sep*sep));
        double vs  = omega * sep * 0.5;
        double r_tp = 1.5;
        double v_tp = std::sqrt(G * 2.0 * Mstar / r_tp) * 0.85;
        q = { -sep*0.5, 0.0,
               sep*0.5, 0.0,
               r_tp,    0.0 };
        v = { 0.0, -vs,
              0.0,  vs,
              0.0,  v_tp };
 
    } else { // Galaxy
        enable_merges = true;
        soft = 0.10;
 
        const int    N_gal  = 50;
        const double R_disk = 3.0;   // initial spread radius
 
        bodies.clear();
        q = num::Vector(2 * N_gal, 0.0);
        v = num::Vector(2 * N_gal, 0.0);
 
        std::mt19937 rng(20240101);
        std::uniform_real_distribution<double> uni(0.0, 1.0);
 
        // Estimate total mass to seed tangential speeds
        // With masses uniform in [0.3, 2.0], mean ≈ 1.15, M_total ≈ 57.5
        const double M_total_est = N_gal * 1.15;
 
        for (int i = 0; i < N_gal; ++i) {
            // Radius: slightly concentrated toward center
            double r     = R_disk * std::pow(uni(rng), 0.65);
            double theta = uni(rng) * 2.0 * M_PI;
 
            // Mass: random uniform [0.3, 2.0]
            double m = 0.3 + 1.7 * uni(rng);
 
            // Tangential velocity: ~35% of circular velocity at this radius
            // This gives highly elliptical orbits that cross each other
            double vcirc = (r > 0.05) ? std::sqrt(G * M_total_est / r) : 0.0;
            double vtang = vcirc * (0.25 + 0.20 * uni(rng));  // 25–45% of circular
 
            // Small random perturbation so orbits have different phases
            double vrand_x = (2.0*uni(rng) - 1.0) * 0.3;
            double vrand_y = (2.0*uni(rng) - 1.0) * 0.3;
 
            q[2*i]   = r * std::cos(theta);
            q[2*i+1] = r * std::sin(theta);
            // Tangential direction (counterclockwise): (-sin θ, cos θ)
            v[2*i]   = -vtang * std::sin(theta) + vrand_x;
            v[2*i+1] =  vtang * std::cos(theta) + vrand_y;
 
            bodies.push_back({ m, disp_r(m), mass_color(m), phys_r(m) });
        }
    }
 
    E0 = total_energy();
}
 
} // namespace nbody
 
// ── Renderer ──────────────────────────────────────────────────────────────────
 
static constexpr int   SW          = 1000;
static constexpr int   SH          = 720;
static constexpr int   PANEL_W     = 240;
static constexpr int   SIM_W       = SW - PANEL_W;
static constexpr int   MAX_TRAIL   = 600;
static constexpr float BASE_DT     = 1.0f / 120.0f;
 
struct Trail {
    std::deque<Vector2> pts;
    void push(float x, float y) {
        pts.push_front({x, y});
        if ((int)pts.size() > MAX_TRAIL) pts.pop_back();
    }
};
 
static Vector2 sim_to_screen(double sx, double sy,
                              float cx, float cy, float scale) {
    return { cx + static_cast<float>(sx) * scale,
             cy - static_cast<float>(sy) * scale };
}
 
static Color unpack_color(uint32_t c) {
    return { uint8_t(c >> 24), uint8_t(c >> 16), uint8_t(c >> 8), uint8_t(c) };
}
 
static float scene_scale(nbody::Scenario s) {
    switch (s) {
        case nbody::Scenario::Figure8:     return 240.0f;
        case nbody::Scenario::SolarSystem: return 240.0f;
        case nbody::Scenario::BinaryPlus:  return 200.0f;
        case nbody::Scenario::Galaxy:      return 130.0f;
        default:                           return 200.0f;
    }
}
 
int main() {
    SetConfigFlags(FLAG_MSAA_4X_HINT | FLAG_VSYNC_HINT);
    InitWindow(SW, SH, "N-Body Gravitational Simulation");
    SetTargetFPS(120);
 
    nbody::NBodySim sim;
    nbody::Scenario current = nbody::Scenario::Figure8;
    sim.reset(current);
 
    std::vector<Trail> trails(sim.n());
 
    bool paused        = false;
    bool trail_fade    = true;
    int  speed_mult    = 4;
    float scale        = scene_scale(current);
    float cx           = SIM_W * 0.5f;
    float cy           = SH   * 0.5f;
 
    std::deque<float> energy_hist;
    constexpr int ENERGY_HIST = 400;
 
    while (!WindowShouldClose()) {
        // ── Input ─────────────────────────────────────────────────────────────
        auto switch_scenario = [&](nbody::Scenario s) {
            current = s;
            sim.reset(s);
            trails.assign(sim.n(), Trail{});
            scale = scene_scale(s);
            energy_hist.clear();
        };
 
        if (IsKeyPressed(KEY_ONE))   switch_scenario(nbody::Scenario::Figure8);
        if (IsKeyPressed(KEY_TWO))   switch_scenario(nbody::Scenario::SolarSystem);
        if (IsKeyPressed(KEY_THREE)) switch_scenario(nbody::Scenario::BinaryPlus);
        if (IsKeyPressed(KEY_FOUR))  switch_scenario(nbody::Scenario::Galaxy);
        if (IsKeyPressed(KEY_R))     switch_scenario(current);
        if (IsKeyPressed(KEY_SPACE)) paused = !paused;
        if (IsKeyPressed(KEY_V)) {
            sim.use_verlet = !sim.use_verlet;
            sim.reset(current);
            trails.assign(sim.n(), Trail{});
            energy_hist.clear();
        }
        if (IsKeyPressed(KEY_F)) trail_fade = !trail_fade;
        if (IsKeyPressed(KEY_EQUAL) || IsKeyPressed(KEY_KP_ADD))
            speed_mult = std::min(speed_mult + 1, 16);
        if (IsKeyPressed(KEY_MINUS) || IsKeyPressed(KEY_KP_SUBTRACT))
            speed_mult = std::max(speed_mult - 1, 1);
 
        // ── Simulate ──────────────────────────────────────────────────────────
        if (!paused) {
            double dt    = static_cast<double>(BASE_DT) / speed_mult;
            int    inner = std::max(1, speed_mult);
 
            for (int s = 0; s < inner; ++s) {
                sim.step(dt);
 
                // Merge detection — mirror every swap-and-pop on the trails vector
                if (sim.enable_merges) {
                    auto ops = sim.check_merges();
                    for (auto& [j, last_] : ops) {
                        if (j != last_) std::swap(trails[j], trails[last_]);
                        trails.pop_back();
                    }
                }
            }
 
            // Record trail positions
            for (int i = 0; i < sim.n(); ++i) {
                auto sc = sim_to_screen(sim.q[2*i], sim.q[2*i+1], cx, cy, scale);
                if (sc.x > 0 && sc.x < SIM_W && sc.y > 0 && sc.y < SH)
                    trails[i].push(sc.x, sc.y);
            }
 
            float drift = static_cast<float>(sim.energy_drift() * 100.0);
            energy_hist.push_front(drift);
            if ((int)energy_hist.size() > ENERGY_HIST) energy_hist.pop_back();
        }
 
        // ── Draw ──────────────────────────────────────────────────────────────
        BeginDrawing();
        ClearBackground({13, 17, 23, 255});
 
        BeginScissorMode(0, 0, SIM_W, SH);
 
        // Trails
        for (int i = 0; i < sim.n(); ++i) {
            const auto& tr = trails[i];
            Color col = unpack_color(sim.bodies[i].color);
            int len = static_cast<int>(tr.pts.size());
            for (int k = 0; k + 1 < len; ++k) {
                if (trail_fade) {
                    float alpha = 1.0f - static_cast<float>(k) / MAX_TRAIL;
                    col.a = static_cast<uint8_t>(alpha * alpha * 200);
                }
                DrawLineV(tr.pts[k], tr.pts[k+1], col);
            }
        }
 
        // Bodies
        for (int i = 0; i < sim.n(); ++i) {
            auto sc = sim_to_screen(sim.q[2*i], sim.q[2*i+1], cx, cy, scale);
            Color col = unpack_color(sim.bodies[i].color);
            float r = sim.bodies[i].display_radius;
            DrawCircleV(sc, r * 2.2f, {col.r, col.g, col.b, 40});
            DrawCircleV(sc, r * 1.4f, {col.r, col.g, col.b, 100});
            DrawCircleV(sc, r,        col);
        }
 
        EndScissorMode();
 
        DrawRectangle(SIM_W, 0, 1, SH, {48, 54, 61, 255});
 
        // ── Info panel ────────────────────────────────────────────────────────
        const int PX = SIM_W + 14;
        int py = 18;
        char buf[64];
 
        auto label = [&](const char* txt, Color c = {139,148,158,255}) {
            DrawText(txt, PX, py, 13, c);
            py += 18;
        };
        auto value = [&](const char* key, const char* val) {
            DrawText(key, PX, py, 13, {139,148,158,255});
            DrawText(val, PX + 90, py, 13, {201,209,217,255});
            py += 18;
        };
 
        DrawText(scenario_name(current), PX, py, 14, {88,166,255,255});
        py += 22;
        DrawLine(PX, py, SW-8, py, {48,54,61,255}); py += 10;
 
        snprintf(buf, sizeof(buf), "%.2f", sim.t);
        value("time",       buf);
        snprintf(buf, sizeof(buf), "%dx", speed_mult);
        value("speed",      buf);
        value("integrator", sim.use_verlet ? "Verlet (symp)" : "RK4");
        snprintf(buf, sizeof(buf), "%d", sim.n());
        value("bodies",     buf);
        py += 6;
 
        // Energy drift (skip for Galaxy — merges make it meaningless)
        DrawLine(PX, py, SW-8, py, {48,54,61,255}); py += 10;
        if (!sim.enable_merges) {
            label("Energy drift (%)");
            snprintf(buf, sizeof(buf), "%+.2e", sim.energy_drift() * 100.0);
            Color drift_col = std::abs(sim.energy_drift()) < 1e-4
                              ? Color{63,185,80,255}
                              : std::abs(sim.energy_drift()) < 1e-2
                                ? Color{255,166,64,255}
                                : Color{248,81,73,255};
            DrawText(buf, PX, py, 14, drift_col);
            py += 22;
 
            if (!energy_hist.empty()) {
                float plot_h = 70.0f;
                float plot_w = static_cast<float>(PANEL_W - 28);
                float plot_x = static_cast<float>(PX);
                float plot_y = static_cast<float>(py);
                DrawRectangle((int)plot_x, (int)plot_y, (int)plot_w, (int)plot_h, {22,27,34,255});
                DrawRectangleLines((int)plot_x, (int)plot_y, (int)plot_w, (int)plot_h, {48,54,61,255});
 
                float max_abs = 0.01f;
                for (float vv : energy_hist) max_abs = std::max(max_abs, std::abs(vv));
                max_abs *= 1.2f;
 
                int sz = static_cast<int>(energy_hist.size());
                for (int k = 0; k + 1 < sz && k + 1 < (int)plot_w; ++k) {
                    float x0 = plot_x + plot_w - k - 1;
                    float x1 = plot_x + plot_w - k - 2;
                    float y0 = plot_y + plot_h*0.5f - energy_hist[k]   / max_abs * (plot_h*0.45f);
                    float y1 = plot_y + plot_h*0.5f - energy_hist[k+1] / max_abs * (plot_h*0.45f);
                    DrawLineV({x0,y0}, {x1,y1}, drift_col);
                }
                DrawLine((int)plot_x, (int)(plot_y+plot_h*0.5f),
                         (int)(plot_x+plot_w), (int)(plot_y+plot_h*0.5f), {48,54,61,255});
                py += static_cast<int>(plot_h) + 10;
            }
        } else {
            // Galaxy: show total mass instead of energy drift
            double total_m = 0.0;
            for (int i = 0; i < sim.n(); ++i) total_m += sim.bodies[i].mass;
            label("Merging simulation");
            snprintf(buf, sizeof(buf), "%.1f", total_m);
            value("total mass", buf);
            if (sim.n() == 1) {
                label("** Fully collapsed! **", {255, 166, 64, 255});
            }
        }
 
        // Controls
        py += 4;
        DrawLine(PX, py, SW-8, py, {48,54,61,255}); py += 10;
        label("Controls", {88,166,255,255});
        label("1  Figure-8");
        label("2  Solar system");
        label("3  Binary+");
        label("4  Galaxy collapse");
        label("V  Verlet/RK4");
        label("SPACE  pause");
        label("R  reset");
        label("+/-  speed");
        label("F  trail fade");
 
        if (paused) {
            DrawRectangle(SIM_W/2 - 50, SH/2 - 16, 100, 32, {0,0,0,160});
            DrawText("PAUSED", SIM_W/2 - 38, SH/2 - 10, 20, {201,209,217,255});
        }
 
        DrawFPS(8, 8);
        EndDrawing();
    }
 
    CloseWindow();
    return 0;
}