main.cpp

Go to the documentation of this file.

/// @file main.cpp
/// @brief Real-time raylib visualisation of the 2-D Navier-Stokes demo.
///
/// Render layers (back -> front):
///   1. Vorticity texture (NxN)   -- red = +omega (CCW), blue = -omega (CW), black = 0
///   2. Particle tracers          -- 3 000 massless particles, white with trail fade
///   3. HUD                       -- solver stats, fps, controls
///
/// Keyboard controls:
///   SPACE      pause / resume
///   R          reset simulation
///   +  / -     increase / decrease substeps per frame
///   [ / ]      halve / double visualisation omega scale
///   ESC / Q    quit
///
/// Command-line:
///   ./ns_demo [N]     e.g.  ./ns_demo 512   (default 256)
 
#include "raylib.h"
#include "ns_solver.hpp"
 
#include <vector>
#include <cmath>
#include <cstdlib>
#include <algorithm>
 
//  Colour map: diverging blue-black-red
 
static Color vorticity_color(float t) {
    // t in [-1, 1]:  -1 -> deep blue,  0 -> black,  +1 -> deep red
    t = std::max(-1.0f, std::min(1.0f, t));
    if (t >= 0.0f) {
        auto v = static_cast<unsigned char>(255 * t * t);
        return {v, 0, static_cast<unsigned char>(v * 0.1f), 255};
    } else {
        auto v = static_cast<unsigned char>(255 * t * t);
        return {static_cast<unsigned char>(v * 0.1f), 0, v, 255};
    }
}
 
//  Particle tracer
 
struct Particle {
    float x = 0, y = 0;  // physical coords in [0, 1]
    int   age = 0;
};
 
static void advect_particle(Particle& p, const ns::NSSolver& s, float dt_f) {
    float ux = static_cast<float>(s.interp_u(p.x, p.y));
    float vy = static_cast<float>(s.interp_v(p.x, p.y));
    p.x = std::fmod(p.x + ux * dt_f + 10.0f, 1.0f);
    p.y = std::fmod(p.y + vy * dt_f + 10.0f, 1.0f);
    ++p.age;
}
 
static void respawn(Particle& p) {
    p.x   = static_cast<float>(rand()) / RAND_MAX;
    p.y   = static_cast<float>(rand()) / RAND_MAX;
    p.age = 0;
}
 
//  main
 
int main(int argc, char* argv[]) {
    int N_arg = (argc > 1) ? std::atoi(argv[1]) : 256;
    if (N_arg < 32 || N_arg > 2048) {
        fprintf(stderr, "N must be in [32, 2048]; got %d\n", N_arg);
        return 1;
    }
    const ns::idx N  = static_cast<ns::idx>(N_arg);
    const ns::real dt = 0.5 / static_cast<ns::real>(N);   // CFL ~= 0.5
 
    ns::NSSolver solver(N, dt, /*nu=*/0.0);
    solver.init_shear_layer();
 
    const int NPART   = 3000;
    const int MAX_AGE = 300;
    std::vector<Particle> particles(NPART);
    for (auto& p : particles) respawn(p);
 
    const int WIN = 900;
    SetConfigFlags(FLAG_MSAA_4X_HINT);
    InitWindow(WIN, WIN, "NS Demo  -- 2D Incompressible Flow  [SPACE pause  R reset  +/- speed  [/] scale]");
    SetTargetFPS(60);
 
    std::vector<Color> pixels(N * N);
    Image img = {
        pixels.data(),
        static_cast<int>(N), static_cast<int>(N),
        1, PIXELFORMAT_UNCOMPRESSED_R8G8B8A8
    };
    Texture2D tex = LoadTextureFromImage(img);
 
    bool  paused    = false;
    int   substeps  = 1;
    float sim_time  = 0.0f;
    float omega_scale = 20.0f;   // vorticity display range; auto-adjusts
 
    while (!WindowShouldClose()) {
 
        if (IsKeyPressed(KEY_SPACE)) paused = !paused;
        if (IsKeyPressed(KEY_R)) {
            solver.init_shear_layer();
            sim_time = 0.0f;
            for (auto& p : particles) respawn(p);
        }
        if (IsKeyPressed(KEY_EQUAL) || IsKeyPressed(KEY_KP_ADD))
            substeps = std::min(substeps + 1, 16);
        if (IsKeyPressed(KEY_MINUS) || IsKeyPressed(KEY_KP_SUBTRACT))
            substeps = std::max(substeps - 1, 1);
        if (IsKeyPressed(KEY_LEFT_BRACKET))
            omega_scale *= 0.5f;
        if (IsKeyPressed(KEY_RIGHT_BRACKET))
            omega_scale *= 2.0f;
 
        if (!paused) {
            for (int s = 0; s < substeps; ++s) {
                solver.step();
                sim_time += static_cast<float>(dt);
            }
 
            // Advect particles (use total elapsed sim dt for this frame)
            float frame_dt = static_cast<float>(dt) * substeps;
            for (auto& p : particles) {
                advect_particle(p, solver, frame_dt);
                if (p.age > MAX_AGE) respawn(p);
            }
        }
 
        // Compute max vorticity this frame for auto-range (lightweight, O(N^2))
        float cur_max = 1.0f;
        for (ns::idx i = 0; i < N; ++i)
            for (ns::idx j = 0; j < N; ++j)
                cur_max = std::max(cur_max, std::abs(
                    static_cast<float>(solver.vorticity(i, j))));
 
        // Exponential moving average to prevent scale jitter
        omega_scale = 0.98f * omega_scale + 0.02f * cur_max;
 
        // Fill pixel buffer
        // Texture origin is top-left; grid j=0 is bottom -> flip j
        for (ns::idx i = 0; i < N; ++i) {
            for (ns::idx j = 0; j < N; ++j) {
                float w = static_cast<float>(solver.vorticity(i, j)) / omega_scale;
                pixels[i + (N - 1 - j) * N] = vorticity_color(w);
            }
        }
        UpdateTexture(tex, pixels.data());
 
        BeginDrawing();
        ClearBackground(BLACK);
 
        // Vorticity field (stretched to full window)
        Rectangle src = {0, 0, static_cast<float>(N), static_cast<float>(N)};
        Rectangle dst = {0, 0, static_cast<float>(WIN), static_cast<float>(WIN)};
        DrawTexturePro(tex, src, dst, {0, 0}, 0.0f, WHITE);
 
        // Particles
        for (const auto& p : particles) {
            float px = p.x * WIN;
            float py = (1.0f - p.y) * WIN;  // flip y for screen coords
            float alpha = static_cast<float>(MAX_AGE - p.age) / MAX_AGE;
            unsigned char a = static_cast<unsigned char>(180 * alpha);
            DrawCircleV({px, py}, 1.5f, {255, 255, 255, a});
        }
 
        const int FS  = 18;
        const int FS2 = 16;
        int y = 8;
        const auto& st = solver.stats;
 
        DrawFPS(WIN - 90, 8);
 
        DrawText(TextFormat("Grid  %d x %d  (%.1fM cells)",
            N_arg, N_arg, (float)(N * N) / 1e6f), 8, y, FS, RAYWHITE);
        y += FS + 4;
 
        DrawText(TextFormat("Time  %.3f    substeps/frame %d  [+/-]",
            sim_time, substeps), 8, y, FS, RAYWHITE);
        y += FS + 4;
 
        DrawText(TextFormat("CG    %zu iters   res %.2e   %5.1f ms",
            st.cg_iters, st.cg_residual, st.pressure_ms),
            8, y, FS, {100, 220, 255, 255});
        y += FS + 4;
 
        DrawText(TextFormat("Advect %4.1f ms   Project %4.1f ms   Total %5.1f ms",
            st.advect_ms, st.project_ms, st.total_ms),
            8, y, FS, {180, 255, 180, 255});
        y += FS + 4;
 
        DrawText(TextFormat("omega scale %.1f  [[ ]]", omega_scale),
            8, y, FS2, GRAY);
 
        if (paused)
            DrawText("PAUSED", WIN / 2 - 50, WIN / 2, 32, YELLOW);
 
        DrawText("SPACE pause  R reset  +/- speed  [/] omega scale  ESC quit",
            8, WIN - 22, 14, {140, 140, 140, 200});
 
        EndDrawing();
    }
 
    UnloadTexture(tex);
    CloseWindow();
    return 0;
}