Jos Stams fluid solver

solver.h

#ifndef SOLVERH
#define SOLVERH
#define IX(i,j) ((i)+(N+2)*(j))
#define SWAP(x0,x) {float * tmp=x0;x0=x;x=tmp;}
#define FOR_EACH_CELL for ( i=1 ; i<=N ; i++ ) { for ( j=1 ; j<=N ; j++ ) {
#define END_FOR }}

static void add_source ( int N, float * x, float * s, float dt )
{
	int i, size=(N+2)*(N+2);
	for ( i=0 ; i<size ; i++ ) x[i] += dt*s[i];
}

static void set_bnd ( int N, int b, float * x )
{
	int i;
	for ( i=1 ; i<=N ; i++ ) {
		x[IX(0  ,i)] = b==1 ? -x[IX(1,i)] : x[IX(1,i)];
		x[IX(N+1,i)] = b==1 ? -x[IX(N,i)] : x[IX(N,i)];
		x[IX(i,0  )] = b==2 ? -x[IX(i,1)] : x[IX(i,1)];
		x[IX(i,N+1)] = b==2 ? -x[IX(i,N)] : x[IX(i,N)];
	}
	x[IX(0  ,0  )] = 0.5f*(x[IX(1,0  )]+x[IX(0  ,1)]);
	x[IX(0  ,N+1)] = 0.5f*(x[IX(1,N+1)]+x[IX(0  ,N)]);
	x[IX(N+1,0  )] = 0.5f*(x[IX(N,0  )]+x[IX(N+1,1)]);
	x[IX(N+1,N+1)] = 0.5f*(x[IX(N,N+1)]+x[IX(N+1,N)]);
}

static void lin_solve ( int N, int b, float * x, float * x0, float a, float c )
{
	int i, j, k;

	for ( k=0 ; k<2 ; k++ ) {
		FOR_EACH_CELL
			x[IX(i,j)] = (x0[IX(i,j)] + a*(x[IX(i-1,j)]+x[IX(i+1,j)]+x[IX(i,j-1)]+x[IX(i,j+1)]))/c;
		END_FOR
		set_bnd ( N, b, x );
	}
}

static void diffuse ( int N, int b, float * x, float * x0, float diff, float dt )
{
	float a=dt*diff*N*N;
	lin_solve ( N, b, x, x0, a, 1+4*a );
}

static void advect ( int N, int b, float * d, float * d0, float * u, float * v, float dt )
{
	int i, j, i0, j0, i1, j1;
	float x, y, s0, t0, s1, t1, dt0;

	dt0 = dt*N;
	FOR_EACH_CELL
		x = i-dt0*u[IX(i,j)]; y = j-dt0*v[IX(i,j)];
		if (x<0.5f) x=0.5f; if (x>N+0.5f) x=N+0.5f; i0=(int)x; i1=i0+1;
		if (y<0.5f) y=0.5f; if (y>N+0.5f) y=N+0.5f; j0=(int)y; j1=j0+1;
		s1 = x-i0; s0 = 1-s1; t1 = y-j0; t0 = 1-t1;
		d[IX(i,j)] = s0*(t0*d0[IX(i0,j0)]+t1*d0[IX(i0,j1)])+
					 s1*(t0*d0[IX(i1,j0)]+t1*d0[IX(i1,j1)]);
	END_FOR
	set_bnd ( N, b, d );
}

static void project ( int N, float * u, float * v, float * p, float * div )
{
	int i, j;

	FOR_EACH_CELL
		div[IX(i,j)] = -0.5f*(u[IX(i+1,j)]-u[IX(i-1,j)]+v[IX(i,j+1)]-v[IX(i,j-1)])/N;
		p[IX(i,j)] = 0;
	END_FOR	
	set_bnd ( N, 0, div ); set_bnd ( N, 0, p );

	lin_solve ( N, 0, p, div, 1, 4 );

	FOR_EACH_CELL
		u[IX(i,j)] -= 0.5f*N*(p[IX(i+1,j)]-p[IX(i-1,j)]);
		v[IX(i,j)] -= 0.5f*N*(p[IX(i,j+1)]-p[IX(i,j-1)]);
	END_FOR
	set_bnd ( N, 1, u ); set_bnd ( N, 2, v );
}

static void dens_step ( int N, float * x, float * x0, float * u, float * v, float diff, float dt )
{
	add_source ( N, x, x0, dt );
	SWAP ( x0, x ); diffuse ( N, 0, x, x0, diff, dt );
	SWAP ( x0, x ); advect ( N, 0, x, x0, u, v, dt );
}

static void vel_step ( int N, float * u, float * v, float * u0, float * v0, float visc, float dt )
{
	add_source ( N, u, u0, dt ); add_source ( N, v, v0, dt );
	SWAP ( u0, u ); diffuse ( N, 1, u, u0, visc, dt );
	SWAP ( v0, v ); diffuse ( N, 2, v, v0, visc, dt );
	project ( N, u, v, u0, v0 );
	SWAP ( u0, u ); SWAP ( v0, v );
	advect ( N, 1, u, u0, u0, v0, dt ); advect ( N, 2, v, v0, u0, v0, dt );
	project ( N, u, v, u0, v0 );
}
#endif

testApp.cpp

#include "testApp.h"


static int allocate_data ( void )
{
	int size = (N+2)*(N+2);

	u			= (float *) malloc ( size*sizeof(float) );
	v			= (float *) malloc ( size*sizeof(float) );
	u_prev		= (float *) malloc ( size*sizeof(float) );
	v_prev		= (float *) malloc ( size*sizeof(float) );
	dens		= (float *) malloc ( size*sizeof(float) );	
	dens_prev	= (float *) malloc ( size*sizeof(float) );

	if ( !u || !v || !u_prev || !v_prev || !dens || !dens_prev ) {
		fprintf ( stderr, "cannot allocate data\n" );
		return ( 0 );
	}

	return ( 1 );
}

static void clear_data ( void )
{
	int i, size=(N+2)*(N+2);

	for ( i=0 ; i<size ; i++ ) {
		u[i] = v[i] = u_prev[i] = v_prev[i] = dens[i] = dens_prev[i] = 0.0f;
	}
}

static void draw_velocity ( void )
{
	int i, j;
	float x, y, h;

	h = 1.0f/N;
	glEnable(GL_BLEND);
	glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

	glLineWidth ( 3.0f );

	glBegin ( GL_LINES );

		for ( i=1 ; i<=N ; i++ ) {
			x = (i-0.5f)*h;
			for ( j=1 ; j<=N ; j++ ) {
				y = (j-0.5f)*h;
				glColor4f ( 1.0f, 0.0f, 0.0f,0.5f );
				glVertex2f ( x, y );
				glColor4f ( 1.0f, 1.0f, 0.0f,0.1f );
				glVertex2f ( x+u[IX(i,j)], y+v[IX(i,j)] );
			}
		}

	glEnd ();
}

static void draw_density ( void )
{
glEnable(GL_BLEND);
	glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);


	int i, j;
	float x, y, h, d00, d01, d10, d11;

	h = 1.0f/N;
	int size = (N+2)*(N+2);
	float max_col = 0.01;
	for(int k = 0; k < size; ++k) {
		if(dens[k] > max_col) {
			max_col = dens[k];
		}
	}
	float inv_col = 1.0f/max_col;
	glBegin ( GL_QUADS );
		for ( i=0 ; i<=N ; i++ ) {
			x = (i-0.5f)*h;
			for ( j=0 ; j<=N ; j++ ) {
				y = (j-0.5f)*h;

				d00 = dens[IX(i,j)] * inv_col;
				d01 = dens[IX(i,j+1)] * inv_col;
				d10 = dens[IX(i+1,j)] * inv_col;
				d11 = dens[IX(i+1,j+1)] * inv_col;
				d00 = v[IX(i,j)] * inv_col;
				d01 = v[IX(i, j+1)] * inv_col;
				d10 = u[IX(i+1, j)] * inv_col;
				glColor4f(d00*0.5,float(j)/N,1.0-d00, d01);
				
				glVertex2f (x, y );
				glVertex2f ( x+h, y );
				glVertex2f ( x+h, y+h );
				glVertex2f ( x, y+h );
			}
		}
	glEnd ();
}

testApp::testApp()
	:gif(ofGetWidth(), ofGetHeight(), 256)
{
}

//--------------------------------------------------------------
void testApp::setup(){	
	ofSetFrameRate(60);
	ofSetVerticalSync(true);
	ofBackground(22,33,44);
	prev_mx = 0;
	prev_my = 0;
	what = 1; // 1 = velocity, 2 = density
	N = 128;
	dt = 0.1f;
	diff = 0.00f;
	visc = 0.00f;
	force = 5.0f;
	source = 10.0f;
	allocate_data();
	clear_data();
	
	save_gif = false;
	add_gif = false;
	save_frame = false;
}

//--------------------------------------------------------------
void testApp::update(){
	ofSetWindowTitle(ofToString(ofGetFrameRate()));
	vel_step(N, u, v, u_prev, v_prev, visc, dt);
	dens_step(N, dens, dens_prev, u, v, diff, dt);
}

//--------------------------------------------------------------
void testApp::draw(){

	glScalef(ofGetWidth(), ofGetHeight(),0);
	draw_density();

	if(add_gif && ofGetFrameNum() % 10 == 0) {
		ofImage img;
		img.grabScreen(0,0,ofGetWidth(), ofGetHeight());
		gif.addFrame(img.getPixels());
	}
	
	if(save_frame) {
		ofImage img;
		img.grabScreen(0,0,ofGetWidth(), ofGetHeight());
		int num = ofGetWidth()*ofGetHeight() * 3;
		unsigned char* pix = new unsigned char[num];
		memcpy(pix, img.getPixels(), num*sizeof(unsigned char));
		frames.push_back(pix);
	}
}

//--------------------------------------------------------------
void testApp::keyPressed(int key){
	if(key == '1') {
		what = 1;
	}
	else if(key == '2') {
		what = 2;
	}
	else if(key == 's') {
		add_gif = false;
		string file = ofGetTimestampString("%m.gif");
		gif.save(ofToDataPath(file, true).c_str());
	}
	else if(key == 'a') {
		add_gif = !add_gif;
	}
	else if(key == 'f') {
		save_frame = !save_frame;
	}
	else if(key == 'g') {
		char fname[512];
		ofImage writer;
		int w = ofGetWidth();
		int h = ofGetHeight();
		writer.allocate(w, h, OF_IMAGE_COLOR);
		for(int i = 0; i < frames.size(); ++i) {
			sprintf(fname, "frame_%04d.bmp", i);
			writer.setFromPixels(frames[i], w, h, OF_IMAGE_COLOR);
			writer.saveImage(fname);
		}
	}
}

//--------------------------------------------------------------
void testApp::keyReleased(int key){

}

//--------------------------------------------------------------
void testApp::mouseMoved(int x, int y){

}

//--------------------------------------------------------------
void testApp::mouseDragged(int x, int y, int button) { 
	int i = float(x)/ofGetWidth() * N+1;
	int j = float(y)/ofGetHeight() * N+1;
	int size = (N+2)*(N+2);
	
	for(int k = 0; k < size; ++k) {
		u_prev[i] = v_prev[i] = dens_prev[i] = 0.0f;
	}
	
	if(what == 1) {
		if(prev_mx || prev_my != 0) {
			float fx = x-prev_mx;
			float fy = y-prev_my;
			u_prev[IX(i,j)] = force * fx;
			v_prev[IX(i,j)] = force * fy;
		}
	}
	else if(what == 2) {
		dens_prev[IX(i,j)] = source;
	}
	prev_mx = x;
	prev_my = y;
}

void testApp::mousePressed(int x, int y, int button){ }
void testApp::mouseReleased(int x, int y, int button){}
void testApp::windowResized(int w, int h){ }
void testApp::gotMessage(ofMessage msg){ }
void testApp::dragEvent(ofDragInfo dragInfo){ }

testApp.h

#pragma once

#include "ofMain.h"
#include "Gif.h"
using namespace roxlu;

#define IX(i,j) ((i)+(N+2)*(j))

static int N;
static float dt, diff, visc;
static float force, source;
static int dvel;

static float * u, * v, * u_prev, * v_prev;
static float * dens, * dens_prev;

#include "solver.h"


class testApp : public ofBaseApp{
	public:
		testApp();
		void setup();
		void update();
		void draw();
		
		void keyPressed(int key);
		void keyReleased(int key);
		void mouseMoved(int x, int y);
		void mouseDragged(int x, int y, int button);
		void mousePressed(int x, int y, int button);
		void mouseReleased(int x, int y, int button);
		void windowResized(int w, int h);
		void dragEvent(ofDragInfo dragInfo);
		void gotMessage(ofMessage msg);
		int prev_mx;
		int prev_my;
		int what;
		
		
		Gif gif;
		bool add_gif;
		bool save_gif;
		bool save_frame;
		
		vector<unsigned char*> frames;
};