space.cpp

Go to the documentation of this file.

/**
 * \file space.cpp
 * \brief Implementation of the n-body simulation space, handling physics and
 *        collisions.
 * \author Le Bars, Yoann
 *
 * This file is part of the pure C++ benchmark.
 *
 * This program is free software: you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
 * more details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */
 
#include "space.hpp"
 
#include <omp.h>
 
#include <Eigen/StdVector>
#include <QDebug>
#include <algorithm>
#include <chrono>
#include <cmath>
#include <random>
 
#include "pcg_random.hpp"
 
// Type alias for profiling clock (used throughout the file).
using ProfilingClock = std::chrono::high_resolution_clock;
 
Model::Space::Space(const Configuration::SimulationConfig& config)
    : dt_(config.max_dt),
      previous_dt_(config.max_dt),
      G_(config.universal_g),
      epsilon_(config.epsilon),
      coeffRestitution_(config.coeff_restitution),
      coeffFriction_(config.coeff_friction),
      coeffStaticFriction_(config.coeff_static_friction),
      positionalCorrectionFactor_(config.positional_correction_factor),
      linearDamping_(config.linear_damping),
      angularDamping_(config.angular_damping),
      diagnosticsEnabled_(config.diagnostics_enabled),
      logFreq_(config.log_freq) {
    initializeBodies(config.n_bodies, config.dens, config.seed);
    dt_ = previous_dt_ = config.max_dt;  // Set initial dt after initialisation
}
 
namespace {
    // Bring Eigen types into this namespace for the PlacementGrid.
    using Eigen::Vector3d;
    using Eigen::Vector3i;
 
    /**
     * \brief A simple hash function for Eigen::Vector3i to allow its use as a
     *        key in std::unordered_map.
     */
    struct Vector3iHash {
        std::size_t operator()(const Vector3i& v) const {
            // A common way to combine hashes for vector types.
            std::hash<int> hasher;
            return hasher(v.x()) ^ (hasher(v.y()) << 1) ^ (hasher(v.z()) << 2);
        }
    };
 
    /**
     * \brief A spatial grid for Poisson Disk Sampling (Bridson's algorithm).
     *
     * This grid is used to accelerate overlap checks during Poisson Disk
     * Sampling. Each cell stores the index of a point (or -1 if empty).
     */
    class PoissonDiskGrid {
       public:
        /// \brief The grid data: grid_[i][j][k] = point index or -1 if empty.
        std::vector<std::vector<std::vector<int>>> grid_;
 
        /// \brief The size of each grid cell.
        double cell_size_;
 
        /// \brief The dimensions of the grid (number of cells in each
        /// dimension).
        Vector3i grid_shape_;
 
        /// \brief The domain dimensions (side length of the cube).
        Vector3d dims_;
 
        PoissonDiskGrid(double in_cell_size, const Vector3d& in_dims)
            : cell_size_(in_cell_size), dims_(in_dims) {
            // Calculate grid shape (number of cells in each dimension)
            grid_shape_ = Vector3i(
                static_cast<int>(std::ceil(in_dims.x() / in_cell_size)),
                static_cast<int>(std::ceil(in_dims.y() / in_cell_size)),
                static_cast<int>(std::ceil(in_dims.z() / in_cell_size)));
 
            // Initialize grid with -1 (empty)
            grid_.resize(grid_shape_.x());
            for (int i = 0; i < grid_shape_.x(); ++i) {
                grid_[i].resize(grid_shape_.y());
                for (int j = 0; j < grid_shape_.y(); ++j) {
                    grid_[i][j].resize(grid_shape_.z(), -1);
                }
            }
        }
 
        /**
         * \brief Gets the grid cell coordinates for a position.
         * \param pos The 3D position.
         * \return The cell coordinates as a Vector3i.
         */
        Vector3i getCellCoords(const Vector3d& pos) const {
            // Convert position to grid coordinates (shift to positive domain)
            const Vector3d shifted = pos + dims_ / 2.0;
            return Vector3i(static_cast<int>(shifted.x() / cell_size_),
                            static_cast<int>(shifted.y() / cell_size_),
                            static_cast<int>(shifted.z() / cell_size_));
        }
 
        /**
         * \brief Checks if a position is valid (within domain and not
         * overlapping).
         * \param pos The candidate position.
         * \param min_dist_sq The minimum squared distance required.
         * \param existing_points The vector of existing points.
         * \return True if the position is valid.
         */
        bool isValid(const Vector3d& pos, double min_dist_sq,
                     const std::vector<Vector3d>& existing_points) const {
            // Check if within domain
            if (std::abs(pos.x()) >= dims_.x() / 2.0 ||
                std::abs(pos.y()) >= dims_.y() / 2.0 ||
                std::abs(pos.z()) >= dims_.z() / 2.0) {
                return false;
            }
 
            // Get cell coordinates
            const Vector3i cell = getCellCoords(pos);
 
            // Check bounds
            if (cell.x() < 0 || cell.x() >= grid_shape_.x() || cell.y() < 0 ||
                cell.y() >= grid_shape_.y() || cell.z() < 0 ||
                cell.z() >= grid_shape_.z()) {
                return false;
            }
 
            // Check if cell is already occupied
            if (grid_[cell.x()][cell.y()][cell.z()] != -1) {
                return false;
            }
 
            // Check 3x3x3 neighborhood for nearby points
            for (int i = -1; i <= 1; ++i) {
                for (int j = -1; j <= 1; ++j) {
                    for (int k = -1; k <= 1; ++k) {
                        const int x = cell.x() + i;
                        const int y = cell.y() + j;
                        const int z = cell.z() + k;
 
                        if (x >= 0 && x < grid_shape_.x() && y >= 0 &&
                            y < grid_shape_.y() && z >= 0 &&
                            z < grid_shape_.z()) {
                            const int point_idx = grid_[x][y][z];
                            if (point_idx != -1) {
                                const double dist_sq =
                                    (pos - existing_points[point_idx])
                                        .squaredNorm();
                                if (dist_sq < min_dist_sq) {
                                    return false;
                                }
                            }
                        }
                    }
                }
            }
 
            return true;
        }
 
        /**
         * \brief Adds a point to the grid.
         * \param point_idx The index of the point in the points array.
         * \param pos The position of the point.
         */
        void add(std::size_t point_idx, const Vector3d& pos) {
            const Vector3i cell = getCellCoords(pos);
            if (cell.x() >= 0 && cell.x() < grid_shape_.x() && cell.y() >= 0 &&
                cell.y() < grid_shape_.y() && cell.z() >= 0 &&
                cell.z() < grid_shape_.z()) {
                grid_[cell.x()][cell.y()][cell.z()] =
                    static_cast<int>(point_idx);
            }
        }
    };
}  // namespace
 
Model::Space::~Space() {
    // Clean up the OpenMP locks.
    for (std::size_t i = 0; i < graphLocks_.size(); ++i) {
        omp_destroy_lock(&graphLocks_[i]);
    }
}
 
void Model::Space::initializeBodies(std::size_t n, double dens,
                                    unsigned int seed) {
    // Validate input parameters
    if (n == 0) {
        throw std::invalid_argument("Number of bodies must be greater than 0");
    }
    if (dens <= 0.0) {
        throw std::invalid_argument("Density must be greater than 0");
    }
    if (!std::isfinite(dens)) {
        throw std::invalid_argument("Density must be a finite number");
    }
 
    if (seed == 0) {
        /* Random device to create a seed. */
        std::random_device rd;
        seed = rd();
    }
    // Use the PCG generator instead of mt19937
    Rng::Pcg32 gen(seed);  // NOLINT
    // Random values for bodies' masses.
    std::uniform_real_distribution<> massDist(MIN_BODY_MASS, MAX_BODY_MASS);
 
    /* Pre-generate all random properties in vectors, mirroring the
       vectorised approach in the Python version. */
    std::vector<double> masses(n);
    std::vector<double> radii(n);
    std::vector<Vector3d> velocities(n);
    std::vector<Vector3d> angularVelocities(n);
    std::uniform_real_distribution<> vDis(MIN_INITIAL_VELOCITY,
                                          MAX_INITIAL_VELOCITY);
    std::uniform_real_distribution<> omegaDis(-MAX_INITIAL_ANGULAR_VELOCITY,
                                              MAX_INITIAL_ANGULAR_VELOCITY);
    for (std::size_t i = 0; i < n; ++i) {
        masses[i] = massDist(gen);
        radii[i] = masses[i] * dens;
        velocities[i] = Vector3d(vDis(gen), vDis(gen), vDis(gen));
        angularVelocities[i] =
            Vector3d(omegaDis(gen), omegaDis(gen), omegaDis(gen));
    }
 
    // --- Poisson Disk Sampling for Body Placement ---
    /* Use Poisson Disk Sampling (Bridson's algorithm) to generate
       non-overlapping positions, matching the Python implementation.
       This produces a more uniform and stable initial distribution. */
    const double avgMass = (MIN_BODY_MASS + MAX_BODY_MASS) / 2.0;
    const double avgRadius = avgMass * dens;
    const double placementScale =
        std::cbrt(static_cast<double>(n)) * avgRadius * PLACEMENT_SCALE_FACTOR;
    const Vector3d dims(2.0 * placementScale, 2.0 * placementScale,
                        2.0 * placementScale);
 
    // Find maximum radius for minimum distance calculation
    const double maxRadius = *std::max_element(radii.begin(), radii.end());
    const double minDist = 2.0 * maxRadius;  // Minimum distance between centers
    const double minDistSq = minDist * minDist;
 
    // Cell size for acceleration grid
    const double cellSize = minDist / std::sqrt(3.0);
    PoissonDiskGrid grid(cellSize, dims);
 
    // Generate positions using Poisson Disk Sampling
    std::vector<Vector3d> positions;
    positions.reserve(n);
    std::vector<std::size_t> activeList;
    constexpr int kRejectionSamples = 30;
 
    // Start with a single random point
    std::uniform_real_distribution<> initialDist(-placementScale,
                                                 placementScale);
    Vector3d initialPoint(initialDist(gen), initialDist(gen), initialDist(gen));
    positions.push_back(initialPoint);
    grid.add(0, initialPoint);
    activeList.push_back(0);
 
    // Generate subsequent points
    constexpr double TWO_PI = 2.0 * 3.14159265358979323846;
    std::uniform_real_distribution<> angleDist(0.0, TWO_PI);
    std::uniform_real_distribution<> radiusDist(minDist, 2.0 * minDist);
 
    while (!activeList.empty() && positions.size() < n) {
        // Choose a random active point
        std::uniform_int_distribution<std::size_t> activeDist(
            0, activeList.size() - 1);
        const std::size_t activeIdx = activeList[activeDist(gen)];
        const Vector3d& activePoint = positions[activeIdx];
        bool foundCandidate = false;
 
        // Try to generate a valid candidate around the active point
        for (int attempt = 0; attempt < kRejectionSamples; ++attempt) {
            // Generate random point in annulus around active point
            // Using spherical coordinates
            const double theta = angleDist(gen);  // Azimuthal angle
            const double phi = angleDist(gen);    // Polar angle
            const double radius = radiusDist(gen);
 
            Vector3d offset(radius * std::sin(phi) * std::cos(theta),
                            radius * std::sin(phi) * std::sin(theta),
                            radius * std::cos(phi));
            const Vector3d candidate = activePoint + offset;
 
            // Check if candidate is valid (within domain and non-overlapping)
            if (grid.isValid(candidate, minDistSq, positions)) {
                positions.push_back(candidate);
                grid.add(positions.size() - 1, candidate);
                activeList.push_back(positions.size() - 1);
                foundCandidate = true;
                break;
            }
        }
 
        // If no valid candidate found, remove from active list
        if (!foundCandidate) {
            activeList.erase(
                std::remove(activeList.begin(), activeList.end(), activeIdx),
                activeList.end());
        }
    }
 
    // Handle case where not all bodies could be placed
    const std::size_t nPlaced = positions.size();
    if (nPlaced < n) {
        qWarning().noquote()
            << "Failed to generate" << n
            << "non-overlapping positions. The density may be too high. Using"
            << nPlaced << "bodies instead.";
    }
 
    // Place bodies using the generated positions
    bodies_.reserve(nPlaced);
    for (std::size_t i = 0; i < nPlaced; ++i) {
        bodies_.emplaceBack(masses[i], radii[i], positions[i], velocities[i],
                            angularVelocities[i]);
    }
 
    // Initialise the collision graph, its locks, and the spatial grid.
    // Use nPlaced instead of n in case not all bodies could be placed.
    const std::size_t numBodies = nPlaced;
    collisionGraph_.resize(numBodies);
    graphLocks_.resize(numBodies);
    for (std::size_t i = 0; i < numBodies; ++i) {
        omp_init_lock(&graphLocks_[i]);
    }
    /* Initialise thread-local collision pairs storage.
       This size is fixed for the lifetime of the Space object. */
    thread_local_collision_pairs_.resize(omp_get_max_threads());
 
    /* Pre-allocate some memory for the thread-local collision vectors to reduce
       reallocations during the simulation loop. The exact number is a
       heuristic. */
    const std::size_t reserve_size = numBodies / omp_get_max_threads();
    for (auto& vec : thread_local_collision_pairs_) {
        vec.reserve(reserve_size);
    }
 
    /* To avoid duplicating force calculation logic, we run one step of the
       dynamics computation. This correctly calculates initial accelerations
       and handles any initial collisions.
       To prevent numerical explosions from bodies starting too close, we run a
       few “settling” steps with a very small time step. This allows the system
       to reach a more stable state before the main simulation begins. */
    const double originalDt = dt_;  // Save the real time step.
    kdTree_ = std::make_unique<KDTree>(bodies_);
    previous_dt_ = epsilon_;  // Use a tiny time step for settling.
 
    for (int i = 0; i < SETTLING_STEPS; ++i) {
        computeDynamics(0);  // Iteration count doesn't matter for settling
        // For settling, we force the next dt to be the same small value.
        dt_ = previous_dt_;
    }
 
    dt_ = previous_dt_ = originalDt;  // Restore the real time step.
}
 
std::tuple<Model::Vector3dVec, Model::QuaterniondVec>
Model::Space::getAllBodyTransforms() const {
    const std::size_t numBodies = n();
    Vector3dVec positions;
    QuaterniondVec quaternions;
 
    positions.reserve(numBodies);
    quaternions.reserve(numBodies);
 
    positions.assign(bodies_.x_.begin(), bodies_.x_.begin() + numBodies);
    quaternions.assign(bodies_.q_.begin(), bodies_.q_.begin() + numBodies);
 
    return {positions, quaternions};
}
 
std::tuple<Model::Vector3dVec, Model::QuaterniondVec, Model::Vector3dVec,
           Model::Vector3dVec>
Model::Space::getDataForDisplay() const {
    const std::size_t numBodies = n();
    Vector3dVec positions;
    QuaterniondVec quaternions;
    Vector3dVec torques;
    Vector3dVec alphas;
 
    positions.assign(bodies_.x_.begin(), bodies_.x_.begin() + numBodies);
    quaternions.assign(bodies_.q_.begin(), bodies_.q_.begin() + numBodies);
    torques.assign(bodies_.torque_.begin(),
                   bodies_.torque_.begin() + numBodies);
    alphas.assign(bodies_.alpha_.begin(), bodies_.alpha_.begin() + numBodies);
 
    return {positions, quaternions, torques, alphas};
}
 
void Model::Space::printProfilingReport() const {
    const double totalTime = profIntegration_ + profCollisionGraph_ +
                             profCollisionResponse_ + profForceCalculation_;
 
    if (totalTime < epsilon_) {
        qDebug().noquote() << "\n--- Profiling Report (No data) ---";
        return;
    }
 
    const QString msg =
        QString(
            "\n--- Profiling Report ---\nIntegration:          %1s "
            "(%2%%)\nCollision Graph:      %3s (%4%%)\nCollision Response:   "
            "%5s (%6%%)\nForce Calculation:    %7s "
            "(%8%%)\n------------------------\nTotal Instrumented Time:%9s")
            .arg(profIntegration_, 9, 'f', 4)
            .arg(profIntegration_ / totalTime * 100.0, 5, 'f', 1)
            .arg(profCollisionGraph_, 9, 'f', 4)
            .arg(profCollisionGraph_ / totalTime * 100.0, 5, 'f', 1)
            .arg(profCollisionResponse_, 9, 'f', 4)
            .arg(profCollisionResponse_ / totalTime * 100.0, 5, 'f', 1)
            .arg(profForceCalculation_, 9, 'f', 4)
            .arg(profForceCalculation_ / totalTime * 100.0, 5, 'f', 1)
            .arg(totalTime, 9, 'f', 4);
    qDebug().noquote() << msg;
}
 
std::tuple<double, double, double> Model::Space::calculateSystemEnergy() const {
    double transKe = 0.0;
    double rotKe = 0.0;
    double potential = 0.0;
 
    // --- Kinetic Energy Calculation (Vectorized) ---
    // This loop is perfectly parallel, as each body's KE is independent.
#pragma omp parallel for reduction(+ : transKe, rotKe)
    for (std::size_t i = 0; i < n(); ++i) {
        const auto bi = body(i);
        // Linear kinetic energy: 0.5 * m * v^2
        transKe += 0.5 * bi.m() * bi.v().squaredNorm();
        /* Rotational kinetic energy: 0.5 * I * w^2.
           Since I = 1.0 / getInverseInertia(), this is 0.5 * (1 /
           getInverseInertia()) * w^2. */
        if (bi.getInverseInertia() > 0) {
            rotKe +=
                0.5 * (1.0 / bi.getInverseInertia()) * bi.omega().squaredNorm();
        }
    }
 
    // --- Gravitational Potential Energy Calculation ---
    /* This O(N^2) calculation is kept separate to allow for better parallel
       load balancing of its triangular workload. This is a reduction over
       `potential`. */
#pragma omp parallel for reduction(+ : potential) schedule(dynamic)
    for (std::size_t i = 0; i < n(); ++i) {
        const auto bi = body(i);
        /* To avoid double-counting, we only calculate for pairs (i, j) where
           j > i. */
        for (std::size_t j = i + 1; j < n(); ++j) {
            const auto bj = body(j);
            const double dSq = (bi.x() - bj.x()).squaredNorm();
            if (dSq > epsilon_ * epsilon_) {
                potential -= G_ * bi.m() * bj.m() / std::sqrt(dSq);
            }
        }
    }
 
    return {transKe, rotKe, potential};
}
 
void Model::Space::resolveInterpenetration(const CollisionContext& ctx) {
    /* The total correction is distributed between the two bodies based on
       their inverse mass. */
    const double totalInvMass =
        ctx.b1.getInverseMass() + ctx.b2.getInverseMass();
    if (totalInvMass < epsilon_) {
        return;
    }
 
    // A small percentage of the overlap is corrected to avoid jittering.
    constexpr double PERCENT = 0.8;  // 80% correction
    const Vector3d correction =
        (ctx.overlap_ * PERCENT / totalInvMass) * ctx.nVec_;
 
    if (ctx.b1.getInverseMass() > 0) {                         // NOLINT
        ctx.b1.addToX(-correction * ctx.b1.getInverseMass());  // NOLINT
    }
    if (ctx.b2.getInverseMass() > 0) {                        // NOLINT
        ctx.b2.addToX(correction * ctx.b2.getInverseMass());  // NOLINT
    }
}
 
double Model::Space::getEffectiveMass(const CollisionContext& ctx,
                                      const Vector3d& direction_vec) {
    // Rotational inertia contribution for body 1: (r1 x n)^2 / I1
    const double term1 = (ctx.r1Vec_.cross(direction_vec)).squaredNorm() *
                         ctx.b1.getInverseInertia();  // NOLINT
    // Rotational inertia contribution for body 2: (r2 x n)^2 / I2
    const double term2 = (ctx.r2Vec_.cross(direction_vec)).squaredNorm() *
                         ctx.b2.getInverseInertia();  // NOLINT
 
    const double effectiveMassInv = ctx.b1.getInverseMass() +
                                    ctx.b2.getInverseMass() + term1 +
                                    term2;  // NOLINT
    return effectiveMassInv > epsilon_ ? 1.0 / effectiveMassInv : 0.0;
}
 
std::tuple<Model::Vector3d, double, double>
Model::Space::applyRestitutionImpulse(CollisionContext& ctx) {
    const double vRelNormal = ctx.vRel_.dot(ctx.nVec_);
    if (vRelNormal > 0) {
        // Bodies are already moving apart, no restitution needed.
        return {Vector3d::Zero(), 0.0, 0.0};
    }
    const double effectiveMassN = getEffectiveMass(ctx, ctx.nVec_);
 
    // --- Baumgarte Stabilisation ---
    // Calculate the positional error correction impulse (bias).
    const double bias_jn = (positionalCorrectionFactor_ / previous_dt_) *
                           std::max(0.0, ctx.overlap_ - 0.01) * effectiveMassN;
 
    /* Impulse for restitution (velocity change), based only on the normal
       velocity. */
    const double restitution_jn =
        -(1.0 + coeffRestitution_) * vRelNormal * effectiveMassN;
 
    /* Total impulse magnitude including positional correction.
       The physical restitution impulse cannot be attractive (negative). */
    const double jn = std::max(0.0, restitution_jn) + bias_jn;
 
    return {jn * ctx.nVec_, jn, bias_jn};  // NOLINT
}
 
Model::Vector3d Model::Space::applyFrictionImpulse(const CollisionContext& ctx,
                                                   double jn,
                                                   const Vector3d& impulseN) {
    /* Calculate the tangential component of the INITIAL relative velocity.
       Since impulses are applied simultaneously, the friction impulse must
       oppose the initial tangential velocity to ensure energy dissipation.
       The work done is: W = impulseT · vT_initial, and we need W < 0. */
    const Vector3d vT_initial =
        ctx.vRel_ - ctx.vRel_.dot(ctx.nVec_) * ctx.nVec_;
    const double vT_initialNorm = vT_initial.norm();
 
    // If there’s no tangential velocity, there’s no friction.
    if (vT_initialNorm < epsilon_) {
        return Vector3d::Zero();
    }
 
    // Direction of tangential motion (based on initial relative velocity)
    const Vector3d tVec = vT_initial / vT_initialNorm;
    const double effectiveMassT = getEffectiveMass(ctx, tVec);  // NOLINT
 
    /* Calculate the relative velocity after the restitution impulse is applied.
       This is used to determine the magnitude of friction needed. */
    const Vector3d delta_v1 = -impulseN * ctx.b1.getInverseMass();
    const Vector3d delta_v2 = impulseN * ctx.b2.getInverseMass();
 
    // Calculate the torques generated by the restitution impulse
    const Vector3d torque1 = ctx.r1Vec_.cross(-impulseN);
    const Vector3d torque2 = ctx.r2Vec_.cross(impulseN);
 
    // Calculate the change in angular velocities
    const Vector3d delta_omega1 = ctx.b1.getInverseInertia() * torque1;
    const Vector3d delta_omega2 = ctx.b2.getInverseInertia() * torque2;
 
    // Calculate the relative velocity after restitution
    const Vector3d vRelAfterRestitution =
        ((ctx.b2.v() + delta_v2) +
         (ctx.b2.omega() + delta_omega2).cross(ctx.r2Vec_)) -
        ((ctx.b1.v() + delta_v1) +
         (ctx.b1.omega() + delta_omega1).cross(ctx.r1Vec_));
 
    /* Project the tangential velocity after restitution onto the initial
       tangential direction to determine the magnitude needed. */
    const Vector3d vTAfterRestitution =
        vRelAfterRestitution - vRelAfterRestitution.dot(ctx.nVec_) * ctx.nVec_;
    const double vTAfterRestitutionAlongT = vTAfterRestitution.dot(tVec);
 
    /* Calculate the magnitude of the impulse required to stop the tangential
       motion along the initial tangential direction. */
    const double jtRequiredMagnitude =
        std::abs(vTAfterRestitutionAlongT) * effectiveMassT;
 
    // --- Coulomb’s Law of Friction ---
    /* The friction limit should be based on the physical normal force, not
       including the positional correction bias. Calculate the physical
       restitution impulse magnitude (without bias) for the friction limit. */
    const double vRelNormal = ctx.vRel_.dot(ctx.nVec_);
    const double effectiveMassN = getEffectiveMass(ctx, ctx.nVec_);
    const double restitution_jn =
        (vRelNormal <= 0) ? std::max(0.0, -(1.0 + coeffRestitution_) *
                                              vRelNormal * effectiveMassN)
                          : 0.0;
    /* Use only the physical restitution impulse for friction limit, not the
       bias. However, if jn_physical is very small (e.g., when vRelNormal ≈ 0),
       use the provided jn as a fallback to allow friction testing in cases
       where the normal velocity component is negligible. */
    const double jn_physical = restitution_jn;
    const double jn_for_friction = (jn_physical < epsilon_) ? jn : jn_physical;
    const double staticFrictionLimit = coeffStaticFriction_ * jn_for_friction;
 
    double jt;  // The magnitude of the tangential impulse (signed).
    if (jtRequiredMagnitude < staticFrictionLimit) {
        /* Static friction: the impulse magnitude is exactly what's required
           to stop the motion. The impulse must oppose the direction of
           motion. */
        jt = jtRequiredMagnitude;
    } else {
        /* Kinetic friction: the impulse magnitude is the kinetic friction
           limit. The impulse must oppose the direction of motion.
           Use jn_for_friction which falls back to jn if jn_physical is too
           small. */
        jt = coeffFriction_ * jn_for_friction;
    }
    /* Return the impulse. tVec points in the direction of INITIAL tangential
       motion. jt is positive and represents the magnitude needed.
       The impulse must oppose the initial tangential motion to ensure energy
       dissipation. The work done by friction is: W = impulseT · vT_initial.
       For energy dissipation, we need W < 0, so impulseT must oppose
       vT_initial. Since tVec = vT_initial / |vT_initial|, we need impulseT =
       -jt * tVec to ensure it opposes the initial motion direction. This
       guarantees W = (-jt * tVec) · (|vT_initial| * tVec) = -jt * |vT_initial|
       < 0. */
    return -jt * tVec;
}
 
std::optional<Model::CollisionContext> Model::Space::createCollisionContext(
    BodyProxy& b1, BodyProxy& b2) {
    const Vector3d dVec = b2.x() - b1.x();
    const double distSq = dVec.squaredNorm();
    const double rSum = b1.r() + b2.r();
 
    if (distSq >= rSum * rSum || distSq < epsilon_ * epsilon_) {
        return std::nullopt;
    }
 
    const double dist = std::sqrt(distSq);
    const Vector3d nVec = dVec / dist;
    const Vector3d r1Vec = b1.r() * nVec;
    const Vector3d r2Vec = -b2.r() * nVec;
    const Vector3d vRel =
        (b2.v() + b2.omega().cross(r2Vec)) - (b1.v() + b1.omega().cross(r1Vec));
    const double overlap = (b1.r() + b2.r()) - dist;
 
    return CollisionContext{b1,    b2,   nVec,   r1Vec,
                            r2Vec, vRel, overlap};  // NOLINT
}
 
void Model::Space::handleCollision(BodyProxy& b1, BodyProxy& b2) {
    auto context_opt = createCollisionContext(b1, b2);
    if (!context_opt) {
        return;
    }
    CollisionContext& ctx = *context_opt;
 
    // Resolve interpenetration by moving bodies apart. This modifies body
    // positions, which invalidates the collision context (vRel_, overlap_,
    // r1Vec_, r2Vec_, nVec_ all depend on positions).
    resolveInterpenetration(ctx);
 
    // Recreate the collision context with updated positions to ensure accurate
    // impulse calculations. This matches the approach used in the unit tests.
    auto context_opt_updated = createCollisionContext(b1, b2);
    if (!context_opt_updated) {
        return;
    }
    CollisionContext& ctx_updated = *context_opt_updated;
 
    const auto [impulseN, jn, bias_jn] = applyRestitutionImpulse(ctx_updated);
 
    // The friction impulse is only applied if there is a contact impulse.
    if (jn > 0) {
        const Vector3d impulseT =
            applyFrictionImpulse(ctx_updated, jn, impulseN);
        const Vector3d totalImpulse = impulseN + impulseT;
        const Vector3d torque1 =
            ctx_updated.b1.applyImpulse(-totalImpulse, ctx_updated.r1Vec_);
        const Vector3d torque2 =
            ctx_updated.b2.applyImpulse(totalImpulse, ctx_updated.r2Vec_);
        ctx_updated.b1.accumulateTorque(torque1);
        ctx_updated.b2.accumulateTorque(torque2);
    }
}
 
void Model::Space::colorGraph(
    const std::vector<std::vector<std::size_t>>& in_graph) {
    // Implementation of a parallel greedy graph colouring algorithm.
    // See space.hpp for detailed documentation and examples.
    const std::size_t numBodies = n();
    colors_.assign(numBodies, 0);  // Reset colors to 0
    std::vector<bool> conflicts(numBodies, true);
    bool hasConflicts = true;
 
    // A lock to protect writes to colors_ (only needed for thread-safe writes)
    omp_lock_t color_lock;
    omp_init_lock(&color_lock);
    while (hasConflicts) {
        hasConflicts = false;  // Assume no conflicts until one is found.
 
        // Calculate max_color once before the parallel loop to avoid
        // recalculating it N times inside the lock.
        // Note: This is a lower bound; usedColors will be resized if needed.
        const int max_color = *std::max_element(colors_.begin(), colors_.end());
 
        // Pass 1: Colour nodes in parallel based on current neighbour colours.
#pragma omp parallel for
        for (std::size_t i = 0; i < numBodies; ++i) {
            if (!conflicts[i]) continue;
 
            // Read neighbour colours under lock to avoid data races.
            // We need to synchronize reads because other threads may be writing
            // to colors_ concurrently.
            std::vector<int> neighbourColours;
            neighbourColours.reserve(in_graph[i].size());
            omp_set_lock(&color_lock);
            for (const auto& neighbourIdx : in_graph[i]) {
                neighbourColours.push_back(colors_[neighbourIdx]);
            }
            omp_unset_lock(&color_lock);
 
            // Create usedColors vector locally for each thread (no lock needed)
            // Start with max_color + 2, but resize if we encounter larger
            // colors
            std::vector<bool> usedColors(max_color + 2, false);
            for (const int neighbourColour : neighbourColours) {
                // Resize if necessary to handle colors larger than initial
                // max_color
                if (neighbourColour >= static_cast<int>(usedColors.size())) {
                    usedColors.resize(neighbourColour + 2, false);
                }
                usedColors[neighbourColour] = true;
            }
 
            int newColor = 0;  // Find the smallest available colour.
            while (newColor < static_cast<int>(usedColors.size()) &&
                   usedColors[newColor]) {
                newColor++;
            }
            // If we've exhausted all colors in usedColors, we need a new one
            // (this should be rare, but can happen in race conditions)
            if (newColor >= static_cast<int>(usedColors.size())) {
                // This color is guaranteed to be available since it's beyond
                // all currently used colors
            }
 
            // Lock when writing to colors_ (minimal lock time)
            omp_set_lock(&color_lock);
            colors_[i] = newColor;
            omp_unset_lock(&color_lock);
        }
 
        // Pass 2: Detect conflicts in parallel
        // Mark both bodies in a conflicting pair to ensure both are recolored.
        // The condition i < neighbourIdx was used to avoid double-counting,
        // but it prevented higher-indexed bodies from being marked, breaking
        // the algorithm. We now mark both bodies when they share a color.
#pragma omp parallel for reduction(| : hasConflicts)
        for (std::size_t i = 0; i < numBodies; ++i) {
            conflicts[i] = false;
            for (const auto& neighbourIdx : in_graph[i]) {
                if (colors_[i] == colors_[neighbourIdx]) {
                    conflicts[i] = true;
                    hasConflicts = true;
                    break;  // One conflict is enough to mark this body
                }
            }
        }
    }
 
    omp_destroy_lock(&color_lock);
}
 
void Model::Space::logSystemEnergy(std::size_t iteration) {
    if (!diagnosticsEnabled_ || iteration == 0 || iteration % logFreq_ != 0) {
        return;
    }
 
    const auto [transKe, rotKe, potential] = calculateSystemEnergy();
    const double totalEnergy = transKe + rotKe + potential;
 
    // Single string so the logger does not insert spaces between tokens.
    const QString msg =
        QString(
            "\n--- Iteration %1: System Energy ---\n  Translational KE: %2\n  "
            "Rotational KE:    %3\n  Potential Energy: %4\n  Total Energy:     "
            "%5\n  Time Step (dt):   %6")
            .arg(iteration)
            .arg(transKe, 12, 'e', 4)
            .arg(rotKe, 12, 'e', 4)
            .arg(potential, 12, 'e', 4)
            .arg(totalEnergy, 12, 'e', 4)
            .arg(getPreviousTimeStep(), 12, 'e', 4);
    qDebug().noquote() << msg;
 
    // Check for energy increase, which indicates instability.
    if (totalEnergy > lastTotalEnergy_ + epsilon_) {
        qWarning().noquote()
            << "\n========================================"
            << "\n  Total energy increased by "
            << QString::asprintf("%.4e", totalEnergy - lastTotalEnergy_)
            << "\n========================================";
    }
    lastTotalEnergy_ = totalEnergy;
}
 
void Model::Space::buildCollisionGraph() {
    // Implementation of two-phase collision detection (broad-phase +
    // narrow-phase). See space.hpp for detailed documentation and examples.
    const std::size_t numBodies = n();
    // Defensive check: Ensure the collision graph is sized correctly.
    if (collisionGraph_.size() != numBodies) {
        collisionGraph_.resize(numBodies);
    }
 
    // Clear previous frame's data.
#pragma omp parallel for
    for (std::size_t i = 0; i < numBodies; ++i) {
        collisionGraph_[i].clear();
    }
 
    // Clear thread-local collision pairs storage. This must be done separately
    // because thread_local_collision_pairs_ is indexed by thread ID, not body
    // index.
#pragma omp parallel for
    for (int tid = 0;
         tid < static_cast<int>(thread_local_collision_pairs_.size()); ++tid) {
        thread_local_collision_pairs_[tid].clear();
    }
 
    // Rebuild the k-d tree with the new positions for broad-phase detection.
    kdTree_ = std::make_unique<KDTree>(bodies_);
 
    // Find the maximum radius for a safe search radius.
    // Cache it to avoid recalculating every frame (radii don't change during
    // simulation, only positions do). Recalculate only if cache is invalid.
    if (cached_max_radius_ == 0.0) {
        cached_max_radius_ = 0.0;
#pragma omp parallel for reduction(max : cached_max_radius_)
        for (std::size_t i = 0; i < numBodies; ++i) {
            cached_max_radius_ = std::max(cached_max_radius_, bodies_.r_[i]);
        }
    }
    const double max_radius = cached_max_radius_;
 
    // Find all potentially colliding pairs using thread-local storage.
#pragma omp parallel
    {
        const int tid = omp_get_thread_num();
        // Pre-allocate results vector to avoid repeated allocations.
        // Reserve space for a reasonable number of neighbours (heuristic: 50).
#if defined(NANOFLANN_VERSION) && NANOFLANN_VERSION < 0x150
        std::vector<std::pair<std::size_t, double>> results;
#else
        std::vector<nanoflann::ResultItem<std::size_t, double>> results;
#endif
        results.reserve(50);
 
#pragma omp for schedule(dynamic)
        for (std::size_t i = 0; i < numBodies; ++i) {
            // Access SoA data directly to avoid BodyProxy overhead.
            const double r_i = bodies_.r_[i];
            const Vector3d& x_i = bodies_.x_[i];
            const double search_radius = r_i + max_radius;
 
            // Clear results vector for reuse (capacity is preserved).
            results.clear();
            kdTree_->radiusSearch(x_i.data(), search_radius, results);
 
            // Narrow-Phase: Check actual collisions.
            // Access SoA data directly to avoid repeated BodyProxy creation.
            for (const auto& result : results) {
                const std::size_t j = result.first;
                if (i >= j) continue;
 
                const double rSum = r_i + bodies_.r_[j];
                const double distSq = (x_i - bodies_.x_[j]).squaredNorm();
                if (distSq < rSum * rSum) {
                    thread_local_collision_pairs_[tid].push_back({i, j});
                }
            }
        }
    }
 
    // Merge thread-local results into the global graph.
    for (const auto& localPairs : thread_local_collision_pairs_) {
        for (const auto& pair : localPairs) {
            collisionGraph_[pair.first].emplace_back(pair.second);
            collisionGraph_[pair.second].emplace_back(pair.first);
        }
    }
}
 
void Model::Space::computeDynamics(std::size_t iteration) {
    // --- 1. Diagnostics ---
    logSystemEnergy(iteration);
 
    // --- 2. Integration (Part 1) ---
    /* Update positions based on the previous frame's state.
       Capture dt_ before it may be modified by updateAdaptiveTimeStep().
       Both integratePart1 and integratePart2 must use the same dt for the
       Velocity Verlet scheme to be correct. */
    const double current_dt = dt_;
    {
        const auto startTime = ProfilingClock::now();
        bodies_.integratePart1(dt_, previous_dt_);
        profIntegration_ +=
            std::chrono::duration<double>(ProfilingClock::now() - startTime)
                .count();
    }
 
    // --- 3. Collision Detection ---
    {
        const auto startTime = ProfilingClock::now();
        buildCollisionGraph();
        profCollisionGraph_ +=
            std::chrono::duration<double>(ProfilingClock::now() - startTime)
                .count();
    }
 
    // --- 4. Collision Response ---
    {
        const auto startTime = ProfilingClock::now();
        resolveCollisions();
        profCollisionResponse_ +=
            std::chrono::duration<double>(ProfilingClock::now() - startTime)
                .count();
    }
 
    // --- 5. Adaptive Time Stepping ---
    updateAdaptiveTimeStep();
 
    // --- 6. Force Calculation ---
    {
        const auto startTime = ProfilingClock::now();
        applyGravity();
        profForceCalculation_ +=
            std::chrono::duration<double>(ProfilingClock::now() - startTime)
                .count();
    }
 
    // --- 7. Integration (Part 2) ---
    /* Complete the velocity update using the new accelerations.
       Use the same dt that was used in integratePart1 to maintain Velocity
       Verlet correctness. */
    bodies_.integratePart2(current_dt);
 
    // --- 8. Damping ---
    applyDamping();
}
 
void Model::Space::resolveCollisions() {
    const std::size_t numBodies = n();
    if (numBodies == 0) return;
 
    // Colour the graph to partition collisions into independent sets.
    colorGraph(collisionGraph_);
    const int numColours =
        colors_.empty() ? 0
                        : *std::max_element(colors_.begin(), colors_.end()) + 1;
 
    // Process collisions in parallel, one colour at a time.
    for (int c = 0; c < numColours; ++c) {
#pragma omp parallel for schedule(dynamic)
        for (std::size_t i = 0; i < numBodies; ++i) {
            if (colors_[i] != c) continue;
 
            for (const auto& j : collisionGraph_[i]) {
                // To avoid double-processing, only handle pair (i, j) if i < j.
                if (i < j) {
                    auto b1 = body(i);
                    auto b2 = body(j);
                    handleCollision(b1, b2);
                }
            }
        }
    }
}
 
void Model::Space::applyDamping() {
    const std::size_t numBodies = n();
    // Use direct SoA access to avoid BodyProxy overhead, similar to
    // buildCollisionGraph() optimization.
#pragma omp parallel for
    for (std::size_t i = 0; i < numBodies; ++i) {
        bodies_.v_[i] *= (1.0 - linearDamping_);
        bodies_.omega_[i] *= (1.0 - angularDamping_);
    }
}
 
void Model::Space::updateAdaptiveTimeStep() {
    // Store the dt we just used for this frame's integration.
    previous_dt_ = dt_;
 
    // Find the maximum squared acceleration in the system.
    double maxASq = 0.;
#pragma omp parallel for reduction(max : maxASq) schedule(static)
    for (std::size_t i = 0; i < n(); ++i) {
        const double aSq = body(i).a().squaredNorm();
        maxASq = std::max(maxASq, aSq);
    }
 
    // Find the maximum squared velocity (CFL condition).
    double maxVSq = 0.0;
    std::size_t fastestBodyIdx = 0;
    for (std::size_t i = 0; i < n(); ++i) {
        const double vSq = body(i).v().squaredNorm();
        if (vSq > maxVSq) {
            maxVSq = vSq;
            fastestBodyIdx = i;
        }
    }
 
    // Find the maximum squared angular velocity and angular acceleration.
    double maxOmegaSq = 0.;
    double maxAlphaSq = 0.;
#pragma omp parallel for reduction(max : maxOmegaSq, maxAlphaSq) \
    schedule(static)
    for (std::size_t i = 0; i < n(); ++i) {
        const double omegaSq = body(i).omega().squaredNorm();
        const double alphaSq = body(i).alpha().squaredNorm();
        maxOmegaSq = std::max(maxOmegaSq, omegaSq);
        maxAlphaSq = std::max(maxAlphaSq, alphaSq);
    }
 
    // Adapt dt for the next frame based on the current maximum acceleration.
    double dtA;
    if (maxASq > epsilon_ * epsilon_) {
        dtA = std::sqrt(2.0 * TARGET_DX / std::sqrt(maxASq));
    } else {
        dtA = MAX_DT;
    }
 
    double dtV;
    if (maxVSq > epsilon_ * epsilon_) {
        dtV = (0.5 * body(fastestBodyIdx).r()) / std::sqrt(maxVSq);
    } else {
        dtV = MAX_DT;
    }
 
    // Rotational constraints: prevent bodies from rotating too far in a single
    // step. Max allowed rotation in radians (approx. 5.7 degrees), matching
    // Python implementation.
    constexpr double MAX_ANGLE = 0.1;
 
    double dtOmega;
    if (maxOmegaSq > epsilon_ * epsilon_) {
        // Angular velocity constraint: omega * dt < max_angle
        // => dt < max_angle / sqrt(omega²)
        // => dt² < max_angle² / omega²
        dtOmega = MAX_ANGLE / std::sqrt(maxOmegaSq);
    } else {
        dtOmega = MAX_DT;
    }
 
    double dtAlpha;
    if (maxAlphaSq > epsilon_ * epsilon_) {
        // Angular acceleration constraint: 0.5 * alpha * dt² < max_angle
        // => dt² < 2 * max_angle / sqrt(alpha²)
        // => dt < sqrt(2 * max_angle / sqrt(alpha²))
        dtAlpha = std::sqrt(2.0 * MAX_ANGLE / std::sqrt(maxAlphaSq));
    } else {
        dtAlpha = MAX_DT;
    }
 
    // Combine all constraints and apply damping to smooth the change.
    const double targetDt = std::min({MAX_DT, dtA, dtV, dtOmega, dtAlpha});
    dt_ = dt_ * (1.0 - DT_DAMPING_FACTOR) + targetDt * DT_DAMPING_FACTOR;
}
 
void Model::Space::applyGravity() {
#pragma omp parallel for
    for (std::size_t i = 0; i < n(); ++i) {
        body(i).resetAccelerations();
        /* Note: Do not reset torque here, as it may contain torques accumulated
           from collision responses. The torque will be converted to angular
           acceleration below, and then reset for the next frame. */
    }
 
#pragma omp parallel for schedule(dynamic)
    for (std::size_t i = 0; i < n(); ++i) {
        Vector3d totalForceVec = Vector3d::Zero();
        for (std::size_t j = 0; j < n(); ++j) {
            if (i == j) continue;
            const Vector3d dv = body(j).x() - body(i).x();
            const double dSq = dv.squaredNorm();
            if (dSq > epsilon_ * epsilon_) {
                const double force_mag =
                    G_ * body(j).m() / (dSq * std::sqrt(dSq));
                totalForceVec += dv * force_mag;
            }
        }
        body(i).setAcceleration(totalForceVec);
    }
 
    /* Convert accumulated torque to angular acceleration.
       This includes torques from collision responses accumulated in
       handleCollision(). */
#pragma omp parallel for
    for (std::size_t i = 0; i < n(); ++i) {
        body(i).updateAlphaFromTorque();
    }
 
    /* Reset torque after converting to angular acceleration, preparing for
       the next frame. */
#pragma omp parallel for
    for (std::size_t i = 0; i < n(); ++i) {
        body(i).resetTorque();
    }
}