main.cpp

Go to the documentation of this file.

/// @file main.cpp
/// @brief Real-time raylib visualisation of the 2-D TDSE solver.
///
/// Visualisation:
///   Default mode   -- probability density |psi|^2  coloured by quantum phase
///                   (HSV: hue = phase, brightness = prob, saturation = 1)
///   Mode key D     -- probability density only (black -> cyan -> white hot-map)
///   Potential overlay drawn at 30% opacity over the wavefunction
///
/// Controls:
///   SPACE          pause / resume
///   R              reset wavepacket (click to reposition)
///   P              cycle potentials  (Free -> Barrier -> DoubleSlit -> Harmonic -> CircularWell)
///   E              compute eigenmodes (Lanczos, runs in current thread ~1-2 s)
///   1-8            load eigenmode k (after E)
///   D              toggle density-only vs phase colour map
///   Left click     place wavepacket at cursor position
///   Scroll up/dn   increase / decrease initial momentum (kx)
///   +  / -         substeps per frame
///   ESC / Q        quit
 
#include "raylib.h"
#include "tdse_solver.hpp"
 
#include <vector>
#include <cmath>
#include <cstdio>
#include <algorithm>
 
//  Colour maps
 
/// Phase-amplitude HSV colour: hue = phase, value = sqrt(prob) / max_amp
static Color phase_color(double prob, double phase, double max_prob) {
    if (max_prob < 1e-20) return {0, 0, 0, 255};
    float amp = static_cast<float>(std::sqrt(prob / max_prob));
    amp = std::min(1.0f, amp);
 
    // HSV: hue from phase, full saturation, brightness from amplitude
    float hue = static_cast<float>((phase + M_PI) / (2.0 * M_PI)) * 360.0f;
    float s = 0.95f;
    float v = amp;
 
    // HSV to RGB
    float h6 = hue / 60.0f;
    int   hi = static_cast<int>(h6) % 6;
    float f  = h6 - static_cast<int>(h6);
    float p  = v * (1.0f - s);
    float q  = v * (1.0f - s * f);
    float t  = v * (1.0f - s * (1.0f - f));
    float r, g, b;
    switch (hi) {
        case 0: r=v; g=t; b=p; break;
        case 1: r=q; g=v; b=p; break;
        case 2: r=p; g=v; b=t; break;
        case 3: r=p; g=q; b=v; break;
        case 4: r=t; g=p; b=v; break;
        default:r=v; g=p; b=q; break;
    }
    return {static_cast<unsigned char>(r*255),
            static_cast<unsigned char>(g*255),
            static_cast<unsigned char>(b*255), 255};
}
 
/// Density-only: black -> teal -> white
static Color density_color(double prob, double max_prob) {
    if (max_prob < 1e-20) return {0, 0, 0, 255};
    float t = static_cast<float>(prob / max_prob);
    t = std::min(1.0f, std::pow(t, 0.45f)); // gamma for better range
    // black -> cyan(0,255,255) -> white
    if (t < 0.5f) {
        float u = t * 2.0f;
        return {0, static_cast<unsigned char>(u * 255),
                static_cast<unsigned char>(u * 255), 255};
    } else {
        float u = (t - 0.5f) * 2.0f;
        return {static_cast<unsigned char>(u * 255), 255, 255, 255};
    }
}
 
//  main
 
int main(int argc, char* argv[]) {
    int N_arg = (argc > 1) ? std::atoi(argv[1]) : 256;
    if (N_arg < 32 || N_arg > 512) {
        std::fprintf(stderr, "N must be in [32, 512]; got %d\n", N_arg);
        return 1;
    }
    const int N = N_arg;
 
    const double L  = 10.0;
    const double dt = 0.004;
    tdse::TDSESolver solver(N, L, dt);
    solver.set_potential(tdse::Potential::DoubleSlit);
    solver.init_gaussian(L * 0.2, L * 0.5, 5.0, 0.0, L * 0.06);
 
    tdse::Potential pot_order[] = {
        tdse::Potential::Free,
        tdse::Potential::Barrier,
        tdse::Potential::DoubleSlit,
        tdse::Potential::Harmonic,
        tdse::Potential::CircularWell,
    };
    int pot_idx = 2;  // start at DoubleSlit
    double kx_init = 5.0;
 
    const int WIN = 900;
    SetConfigFlags(FLAG_MSAA_4X_HINT);
    InitWindow(WIN, WIN, "TDSE  -- 2D Quantum Wavepacket  [P pot  E eigenmodes  1-8 mode  SPACE pause  R reset  LClick place  Scroll kx]");
    SetTargetFPS(60);
 
    std::vector<Color> pixels(N * N);
    Image img = {pixels.data(), N, N, 1, PIXELFORMAT_UNCOMPRESSED_R8G8B8A8};
    Texture2D tex = LoadTextureFromImage(img);
 
    // Potential overlay texture (computed once per potential change)
    std::vector<Color> pot_pixels(N * N);
    Image pot_img = {pot_pixels.data(), N, N, 1, PIXELFORMAT_UNCOMPRESSED_R8G8B8A8};
    Texture2D pot_tex = LoadTextureFromImage(pot_img);
    bool pot_dirty = true;
 
    bool  paused       = false;
    int   substeps     = 1;
    float sim_time     = 0.0f;
    bool  phase_mode   = true;   // true = phase-coloured, false = density-only
    int   active_mode  = -1;     // which eigenmode is displayed (-1 = wavepacket)
    bool  computing_modes = false;
 
    auto rebuild_potential_tex = [&]() {
        double V0 = 5000.0;
        for (int i = 0; i < N; ++i) {
            for (int j = 0; j < N; ++j) {
                double v = solver.V[solver.at(i, j)];
                unsigned char a = (v > V0 * 0.5) ? 200 : 0;
                // White wall with partial alpha
                pot_pixels[j * N + i] = {220, 220, 255, a};
            }
        }
        UpdateTexture(pot_tex, pot_pixels.data());
        pot_dirty = false;
    };
    rebuild_potential_tex();
 
    while (!WindowShouldClose()) {
 
        if (IsKeyPressed(KEY_SPACE))  paused = !paused;
        if (IsKeyPressed(KEY_Q))      break;
        if (IsKeyPressed(KEY_D))      phase_mode = !phase_mode;
        if (IsKeyPressed(KEY_EQUAL) || IsKeyPressed(KEY_KP_ADD))
            substeps = std::min(substeps + 1, 8);
        if (IsKeyPressed(KEY_MINUS) || IsKeyPressed(KEY_KP_SUBTRACT))
            substeps = std::max(substeps - 1, 1);
 
        // Scroll to change momentum
        float scroll = GetMouseWheelMove();
        if (std::abs(scroll) > 0.1f) {
            kx_init = std::clamp(kx_init + scroll * 0.5, 0.5, 20.0);
        }
 
        // Cycle potential
        if (IsKeyPressed(KEY_P)) {
            pot_idx = (pot_idx + 1) % 5;
            solver.set_potential(pot_order[pot_idx]);
            // Reset wavepacket to left side, aimed right
            solver.init_gaussian(L * 0.2, L * 0.5, kx_init, 0.0, L * 0.06);
            sim_time   = 0.0f;
            active_mode = -1;
            pot_dirty  = true;
        }
 
        // Reset
        if (IsKeyPressed(KEY_R)) {
            solver.init_gaussian(L * 0.2, L * 0.5, kx_init, 0.0, L * 0.06);
            sim_time    = 0.0f;
            active_mode = -1;
        }
 
        // Place wavepacket at click position
        if (IsMouseButtonPressed(MOUSE_LEFT_BUTTON)) {
            Vector2 mp = GetMousePosition();
            double x0 = (mp.x / WIN) * L;
            double y0 = (mp.y / WIN) * L;  // screen y goes down = physical y down
            solver.init_gaussian(x0, y0, kx_init, 0.0, L * 0.06);
            active_mode = -1;
            sim_time    = 0.0f;
        }
 
        // Compute eigenmodes (blocks for 1-3s)
        if (IsKeyPressed(KEY_E)) {
            computing_modes = true;
            solver.compute_modes(8);
            computing_modes = false;
            // Update stats
            solver.stats.n_modes = static_cast<int>(solver.modes.size());
        }
 
        // Load eigenmode
        for (int k = 0; k < 8; ++k) {
            if (IsKeyPressed(KEY_ONE + k)) {
                if (solver.modes_ready && k < static_cast<int>(solver.modes.size())) {
                    solver.set_mode(k);
                    active_mode = k;
                    sim_time    = 0.0f;
                }
            }
        }
 
        if (!paused) {
            for (int s = 0; s < substeps; ++s) {
                solver.step();
                sim_time += static_cast<float>(dt);
            }
        }
 
        if (pot_dirty) rebuild_potential_tex();
 
        // Find max probability for colour normalisation
        double max_prob = 1e-20;
        for (int i = 0; i < N; ++i)
            for (int j = 0; j < N; ++j)
                max_prob = std::max(max_prob, solver.prob(i, j));
 
        // Screen: row j = 0 at top -> physical y = (j+1)*h = small -> flip
        for (int i = 0; i < N; ++i) {
            for (int j = 0; j < N; ++j) {
                double p = solver.prob(i, j);
                int px = j * N + i;            // texture: column = x, row = y (top-to-bottom)
                // Flip y: screen row 0 = top = physical y = L
                int screen_j = N - 1 - j;
                px = screen_j * N + i;
 
                if (phase_mode)
                    pixels[px] = phase_color(p, solver.phase_ang(i, j), max_prob);
                else
                    pixels[px] = density_color(p, max_prob);
            }
        }
        UpdateTexture(tex, pixels.data());
 
        BeginDrawing();
        ClearBackground(BLACK);
 
        // Wavefunction
        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);
 
        // Potential walls overlay
        DrawTexturePro(pot_tex, src, dst, {0, 0}, 0.0f, WHITE);
 
        DrawFPS(WIN - 90, 8);
        const int FS  = 17;
        const int FS2 = 14;
        int y = 8;
 
        DrawText(TextFormat("Grid %dx%d   t = %.3f   substeps %d  [+/-]",
            N, N, sim_time, substeps), 8, y, FS, RAYWHITE);
        y += FS + 3;
 
        DrawText(TextFormat("Potential: %s  [P]", tdse::potential_name(pot_order[pot_idx])),
            8, y, FS, {100, 220, 255, 255});
        y += FS + 3;
 
        DrawText(TextFormat("kx = %.1f  [scroll]  |  Colourmap: %s  [D]",
            kx_init, phase_mode ? "phase-HSV" : "density"),
            8, y, FS2, {180, 255, 180, 255});
        y += FS2 + 3;
 
        // Step stats
        DrawText(TextFormat("step %.2f ms   norm %.5f   E = %.4f",
            solver.stats.step_ms,
            solver.stats.norm,
            solver.stats.energy),
            8, y, FS2, {220, 200, 100, 255});
        y += FS2 + 6;
 
        // Eigenmodes panel
        if (solver.modes_ready) {
            DrawText(TextFormat("Eigenmodes [%d]  (1-8 to view):", solver.stats.n_modes),
                8, y, FS2, {255, 200, 100, 255});
            y += FS2 + 2;
            int show = std::min(8, static_cast<int>(solver.modes.size()));
            for (int k = 0; k < show; ++k) {
                const auto& m = solver.modes[k];
                Color c = (k == active_mode) ? Color{255, 255, 100, 255}
                                             : Color{180, 180, 180, 255};
                if (m.exact_energy >= 0)
                    DrawText(TextFormat("  [%d] E=%.4f  exact=%.4f", k+1, m.energy, m.exact_energy),
                        8, y, FS2, c);
                else
                    DrawText(TextFormat("  [%d] E=%.4f", k+1, m.energy), 8, y, FS2, c);
                y += FS2 + 1;
            }
        } else {
            DrawText("Press E to compute eigenmodes (Lanczos)", 8, y, FS2, {180, 180, 180, 255});
            y += FS2 + 4;
        }
 
        if (active_mode >= 0)
            DrawText(TextFormat("Eigenmode %d  (SPACE to evolve, R to reset)", active_mode + 1),
                8, WIN - 40, FS2, {255, 255, 100, 255});
 
        if (paused)
            DrawText("PAUSED", WIN/2 - 48, WIN/2, 32, YELLOW);
 
        DrawText("SPACE pause  P pot  E modes  1-8 mode  R reset  LClick place  D color  ESC quit",
            8, WIN - 20, 12, {130, 130, 130, 200});
 
        EndDrawing();
 
        // Update norm/energy once per frame (cheap)
        solver.stats.norm   = solver.compute_norm();
        solver.stats.energy = solver.compute_energy();
    }
 
    UnloadTexture(tex);
    UnloadTexture(pot_tex);
    CloseWindow();
    return 0;
}