fluid2d.cpp

/*
2D fluid simulation using the Lattice-Boltzmann algorithm,
based on the Fluid Dynamics Simulation by Daniel Schroeder:
 - https://physics.weber.edu/schroeder/fluids/
 - http://physics.weber.edu/schroeder/javacourse/LatticeBoltzmann.pdf

Still need to check if the implementation is correct.

SDL is a prerequisite. Compile with the following commands:
g++ -c -I<include path> fluid2d.cpp;
g++ -L<library path> -lm -lSDL2 -lm fluid.o -o fluid
*/

#include <cmath>
#include <SDL2/SDL.h>
#include <iostream>
#include <stdint.h>


typedef double real;

static const int n = 100;
static const int m = 300;

struct Vec2 {
    Vec2(real a, real b): x(a), y(b) {}
    real x, y;
};

class Cell {
    real loc[9];
    public:
    Cell() {
        for (int i = 0; i < 9; i++) {
            loc[i] = 0.0;
        }
    }
    real operator() (int i, int j) const;
    void set(int i, int j, real val);
    void zero_all();
    void add(int i, int j, real val);
    real count() const;
    Vec2 average_velocity() const;
    friend Cell get_weights();
};

Cell get_weights();

void start(Cell lattices[2][n][m], int barrier[n][m]);
void step(Cell lattices[2][n][m], int ti);
void stream(Cell lattices[2][n][m], int barrier[n][m], int ti);



real dot(Vec2 const &v1, Vec2 const &v2) {
    return v1.x*v2.x + v1.y*v2.y;
}

real Cell::operator() (int i, int j) const {
    i = (i < 0)? 1-i: i;
    j = (j < 0)? 1-j: j;
    return loc[i*3 + j];
}

void Cell::set(int i, int j, real val) {
    i = (i < 0)? 1-i: i;
    j = (j < 0)? 1-j: j;
    loc[i*3 + j] = val;
}

void Cell::zero_all() {
    for (int i = 0; i < 9; i++) {
        loc[i] = 0.0;
    }
}

void Cell::add(int i, int j, real val) {
    i = (i < 0)? 1-i: i;
    j = (j < 0)? 1-j: j;
    loc[i*3 + j] += val;
}

real Cell::count() const {
    real count = 0;
    for (int i = 0; i < 9; i++) {
        count += loc[i];
    }
    return count;
}

Vec2 Cell::average_velocity() const {
    Vec2 v(0.0, 0.0);
    real count = 0.0;
    for (int i = -1; i < 2; i++) {
        for (int j = -1; j < 2; j++) {
            count += operator()(i, j);
            v.x += i*operator()(i, j);
            v.y += j*operator()(i, j);
        }
    }
    if (count > 0.0) {
        v.x /= count;
        v.y /= count;
    }
    return v;
}

Cell get_weights() {
    Cell w;
    real weights[] = {4.0/9.0, 1.0/9.0, 1.0/9.0,
                       1.0/9.0, 1.0/36.0, 1.0/36.0,
                       1.0/9.0, 1.0/36.0, 1.0/36.0};
    for (int i = 0; i < 9; i++) {
        w.loc[i] = weights[i];
    }
    return w;
}

void start(Cell lattices[2][n][m], int barrier[n][m]) {
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            if (barrier[i][j] == 0) {
                lattices[0][i][j].zero_all();
                lattices[1][i][j].zero_all();
                lattices[0][i][j].set(0, 0, 255.0);
                lattices[0][i][j].set(0, 1, 100.0);
                lattices[1][i][j].set(0, 0, 255.0);
                lattices[1][i][j].set(0, 1, 100.0);
            }
        }
    }
}

void step(Cell lattices[2][n][m], int ti) {
    ti %= 2;
    int tf = (ti + 1) % 2;
    Cell weights = get_weights();
    real omega = 1.0;
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            real density = lattices[ti][i][j].count();
            Vec2 mean_vel = lattices[ti][i][j].average_velocity();
            for (int x_dir = -1; x_dir < 2; x_dir++) {
                for (int y_dir = -1; y_dir < 2; y_dir++) {
                    Vec2 dir(x_dir, y_dir);
                    real factor = 1.0 + 3.0*dot(dir, mean_vel)
                                  + 4.5*dot(dir, mean_vel)*dot(dir, mean_vel)
                                  - (3.0/2.0)*dot(mean_vel, mean_vel);
                    real count_old = lattices[ti][i][j](x_dir, y_dir);
                    real count_now = density*weights(x_dir, y_dir)*factor;
                    real count_final = count_old + omega*(count_now - count_old);
                    lattices[tf][i][j].set(x_dir, y_dir, count_final);
                }
            }
        }
    }
}

void stream(Cell lattices[2][n][m], int barrier[n][m], int ti) {
    ti %= 2;
    int tf = (ti + 1) % 2;
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            lattices[tf][i][j].zero_all();
            if (barrier[i][j] == 0) {
                for (int y_dir = -1; y_dir < 2; y_dir++) {
                    for (int x_dir = -1; x_dir < 2; x_dir++) {
                        int i2 = i+y_dir;
                        int j2 = j+x_dir;
                        if (i2 == -1) i2 = n-1;
                        if (i2 == n) i2 = 0;
                        if (j2 == -1) j2 = m-1;
                        if (j2 == m) j2 = 0;
                        if (barrier[i2][j2]) 
                        lattices[tf][i][j].add(-y_dir, -x_dir, 
                                                lattices[ti][i][j](y_dir, x_dir));
                        else
                        lattices[tf][i][j].add(-y_dir, -x_dir, lattices[ti][i2][j2]
                                            (-y_dir, -x_dir));
                    }
                }
            }
        }
    }
}

int main() {
    int grid2screen = 4;
    int screen_width = grid2screen*m;
    int screen_height = grid2screen*n;
    bool is_paused = false;
    SDL_Window* window = NULL;
    if (SDL_Init(SDL_INIT_VIDEO) < 0) {
        fprintf(stderr, "Unable to initialize SDL2: %s\n", SDL_GetError());
        exit(1);
    }
    window = SDL_CreateWindow("Fluid",
                              SDL_WINDOWPOS_UNDEFINED, SDL_WINDOWPOS_UNDEFINED,
                              screen_width, screen_height,
                              SDL_WINDOW_SHOWN);
    if (window == NULL) {
        fprintf(stderr, "Unable to create window: %s\n", SDL_GetError());
        exit(1);
    }
    SDL_Renderer *renderer = SDL_CreateRenderer(window, -1, SDL_RENDERER_ACCELERATED);
    if (renderer == NULL) {
        fprintf(stderr, "Unable to create renderer: %s\n", SDL_GetError());
        exit(1);
    }
    Cell lattices[2][n][m] {{{}}};
    int barrier[n][m] = {{0}};
    SDL_Rect background_rect = {.x = 0, .y = 0, .w = screen_width, .h = screen_height};
    SDL_RenderClear(renderer);
    const Uint8 *keyboard_state;
    int returned = 0;
    int x = 0, y = 0;
    int delay_time = 1;
    int i = 0;
    int steps = 2;
    start(lattices, barrier);
    Cell weights = get_weights();
    do {
        SDL_SetRenderDrawColor(renderer, 0, 0, 0, 255);
        SDL_RenderDrawRect(renderer, &background_rect);
        SDL_PumpEvents();
        keyboard_state = SDL_GetKeyboardState(NULL);
        if (keyboard_state[SDL_SCANCODE_RETURN] || SDL_QuitRequested()) {
            returned = 1;
        }
        if (keyboard_state[SDL_SCANCODE_SPACE]) {
            is_paused = !is_paused;
            steps = (is_paused)? 0: 5;
        }
        if (keyboard_state[SDL_SCANCODE_R]) {
            start(lattices, barrier);
        }
        for (int g = 0; g < 2; g++) {
            for (int j = 0; j < steps; j++, i++) {
                if (j < steps-1) step(lattices, i);
                else stream(lattices, barrier, i);
            }
        }
        if (SDL_GetMouseState(&x, &y) & SDL_BUTTON(SDL_BUTTON_LEFT)) {
            int tmp = grid2screen;
            for (int p = 0; p < 5; p++) {
                for (int q = 0; q < 5; q++) {
                    barrier[(y+p)/tmp][(x+q)/tmp] = 1;
                    lattices[0][(y+p)/tmp][(x+q)/tmp].set(0, 0, 0.0);
                    lattices[1][(y+p)/tmp][(x+q)/tmp].set(0, 0, 0.0);
                }
            }
        }
        if (SDL_GetMouseState(&x, &y) & SDL_BUTTON(SDL_BUTTON_RIGHT)) {
            int tmp = grid2screen;
            for (int p = 0; p < 8; p++) {
                for (int q = 0; q < 8; q++) {
                    barrier[(y+p)/tmp][(x+q)/tmp] = 0;
                }
            }
        }
        for (int j = 0; j < n; j++) {
            for (int k = 0; k < m; k++) {
                real colour1 = lattices[1][j][k](0, 1);
                colour1 = (colour1 > 255.0)? 255.0: colour1;
                real colour2 = lattices[1][j][k](0, -1);
                colour2 = (colour2 > 255.0)? 255.0: colour2;
                real colour3 = lattices[1][j][k](0, 0);
                colour3 = (colour3 > 255.0)? 255.0: colour3;
                if (barrier[j][k]) {
                    SDL_SetRenderDrawColor(renderer, 0, 0, 0, 255);   
                } else {
                    SDL_SetRenderDrawColor(renderer, (uint8_t)colour3, 
                                        (uint8_t)colour3, (uint8_t)colour3, 255);
                }
                for (int p = 0; p < 4; p++) {
                    for (int q = 0; q < 4; q++) {
                        SDL_RenderDrawPoint(renderer, grid2screen*k+p, grid2screen*j+q);
                    }
                }
            }
        }
        SDL_RenderPresent(renderer);
        SDL_Delay(delay_time);
    } while(returned==0);
    SDL_DestroyRenderer(renderer);
    SDL_DestroyWindow(window);
    SDL_Quit();
    return 0;
}