FluidSolver.hx

//A backup of http://lists.motion-twin.com/pipermail/haxe/2009-July/027010.html
/**
 * Memory Layout:
 * Block size b = DATA_SIZE * (N + 2) * (M + 2)
 * 1. d			0  ...  b-1
 * 2. d_prev	b  ... 2b-1
 * 3. u			2b ... 3b-1
 * 4. u_prev	3b ... 4b-1
 * 5. v			4b ... 5b-1
 * 6. v_prev	5b ... 6b-1
 * Total length: 6b
 * @author Simon Krajewski
 */

package saf.physics;
import flash.Memory;
import flash.utils.ByteArray;

class FluidSolver 
{
	static public inline var DATA_SIZE:UInt = 8;
	static public inline var K:UInt = 4;
	
	public var D_PREV:UInt;
	public var D:UInt;
	public var U:UInt;
	public var U_PREV:UInt;
	public var V:UInt;
	public var V_PREV:UInt;
	public var N:Int;
	public var M:Int;

	public var _timer:phx.Timer;
	
	private var _blockSize:UInt;
	private var _bN:Int;
	private var _bM:Int;
	private var _bN2:Int;
	private var _Nhalf:Float;
	private var _Mhalf:Float;
	
	public function new(N:Int, M:Int = 0)
	{
		this.N = N;
		
		_timer = new phx.Timer();
				
		if (M == 0)
			M = N;
		else
			this.M = M;
			
		_blockSize = DATA_SIZE * (N + 2) * (M + 2);
		var ba = new ByteArray();
		ba.length = _blockSize * 6;
		Memory.select(ba);

		_bN = N * DATA_SIZE;
		_bM = M * DATA_SIZE;
		_bN2 = (N + 2) * DATA_SIZE;
		_Nhalf = N / 2;
		_Mhalf = M / 2;
		
		D 		= 0;
		D_PREV 	= _blockSize;
		U 		= _blockSize * 2;
		U_PREV 	= _blockSize * 3;
		V 		= _blockSize * 4;
		V_PREV 	= _blockSize * 5;
		
		var i = 0;
		while (i < _blockSize * 6)
		{
			Memory.setDouble(i, 0.0);
			i += DATA_SIZE;
		}
	}
	
	public function addSource(x:UInt, x0:UInt, dt:Float):Void
	{
		var i = x;
		var t = x + _blockSize;
		while (i < t)
		{
			Memory.setDouble(i, Memory.getDouble(i) + Memory.getDouble(x0 + i - x) * dt);
			i += DATA_SIZE;
		}
	}
	
	private inline function setBnd (b:Int, x:Int):Void
	{
	}

	public function linSolve(b:Int, x:UInt, x0:UInt, a:Float, c:Float):Void
	{
		var i;
		var cInv = 1 / c;
		var t = x + _blockSize - _bN2 - DATA_SIZE;
		
		for (k in 0...K)
		{
			i = x + DATA_SIZE + _bN2;
			while (i < t)
			{
				Memory.setDouble(i,
					(
						Memory.getDouble(x0 + i - x)
					  + a * (Memory.getDouble(i - DATA_SIZE) + Memory.getDouble(i + DATA_SIZE) + Memory.getDouble(i - _bN2) + Memory.getDouble(i + _bN2))
					) * cInv
				);
					
				i += DATA_SIZE;
			}
		}
		setBnd(b, x);
	}
	
	public inline function diffuse (b:Int, x:Int, x0:Int, diff:Float, dt:Float):Void
	{
		var a:Float = dt * diff * N * M;
		linSolve(b, x, x0, a, 1 + 4 * a);
	}
	
	public function advect (b:Int, d:Int, d0:Int, u:Int, v:Int, dt:Float):Void
	{
		var i0, j0, i1, j1:Int;
		var x, y, s0, t0, s1, t1:Float;
		var dt0h:Float = dt * N;
		var dt0v:Float = dt * M;
		
		var i, j;

		i = 1;
		while (i <= N)
		{
			j = 1;
			while (j <= M)
			{
				x = i - dt0h * Memory.getDouble(u + IX(i, j) * DATA_SIZE);
				y = j - dt0v * Memory.getDouble(v + IX(i, j) * DATA_SIZE);
				if (x < 0.5)
					x = 0.5;
				if (x > N + 0.5)
					x = N + 0.5;
				i0 = cast(x);
				i1 = i0 + 1;
				if (y < 0.5)
					y = 0.5;
				if (y > M + 0.5)
					y = M + 0.5;
				j0 = cast(y);
				j1 = j0 + 1;
				s1 = x - i0; s0 = 1 - s1; t1 = y - j0; t0 = 1 - t1;
				
				Memory.setDouble(d + IX(i, j) * DATA_SIZE,
					s0 * (t0 * Memory.getDouble(d0 + IX(i0, j0) * DATA_SIZE) + t1 * Memory.getDouble(d0 + IX(i0, j1) * DATA_SIZE))
				  +	s1 * (t0 * Memory.getDouble(d0 + IX(i1, j0) * DATA_SIZE) + t1 * Memory.getDouble(d0 + IX(i1, j1) * DATA_SIZE))
				);
				j++;
			}
			i++;
		}
		setBnd (b, d);
	}
	
	public function project(u:Int, v:Int, p:Int, div:Int):Void
	{
		var i;
		var nmInv = -1 / (N + M);
		
		i = DATA_SIZE + _bN2;
		
		while (i < _blockSize - DATA_SIZE - _bN2)
		{
			Memory.setDouble(div + i,
				nmInv * (Memory.getDouble(u + i + DATA_SIZE) - Memory.getDouble(u + i - DATA_SIZE) + Memory.getDouble(v + i + _bN2) - Memory.getDouble(v + i -_bN2))
			);
			Memory.setDouble(p + i, 0.0);
			i += DATA_SIZE;
		}

		setBnd (0, div);
		setBnd (0, p);

		linSolve(0, p, div, 1, 4 );

		i = DATA_SIZE + _bN2;
	
		while (i < _blockSize - DATA_SIZE - _bN2)
		{
			Memory.setDouble(u + i,
				Memory.getDouble(u + i) -
				_Nhalf * (Memory.getDouble(p + i + DATA_SIZE) - Memory.getDouble(p + i - DATA_SIZE))
			);
			Memory.setDouble(v + i,
				Memory.getDouble(v + i) -
				_Mhalf * (Memory.getDouble(p + i + _bN2) - Memory.getDouble(p + i - _bN2))
			);			
			i += DATA_SIZE;
		}
		
		setBnd (1, u);
		setBnd (2, v);
	}
	
	public function densityStep(diff:Float, dt:Float)
	{
		densStep(D, D_PREV, U, V, diff, dt);
	}
	
	public function velocityStep(visc:Float, dt:Float)
	{
		velStep(U, V, U_PREV, V_PREV, visc, dt);
	}
	
	public function densStep(x:Int, x0:Int, u:Int, v:Int, diff:Float, dt:Float):Void
	{
		_timer.start("dens");
		_timer.start("dens_add_source");
		addSource(x, x0, dt );
		_timer.stop();
		_timer.start("dens_diffuse");
		diffuse (0, x0, x, diff, dt );
		_timer.stop();
		_timer.start("dens_advect");
		advect (0, x, x0, u, v, dt );
		_timer.stop();
		_timer.stop();
	}

	public function velStep (u:Int, v:Int, u0:Int, v0:Int, visc:Float, dt:Float):Void
	{	
		_timer.start("vel");
		_timer.start("vel_add_source");
		addSource (u, u0, dt );
		addSource (v, v0, dt );
		_timer.stop();
		_timer.start("vel_diffuse");
		diffuse (1, u0, u, visc, dt );
		diffuse (2, v0, v, visc, dt );
		_timer.stop();
		_timer.start("vel_project1");
		project (u0, v0, u, v );
		_timer.stop();
		_timer.start("vel_advect");
		advect (1, u, u0, u0, v0, dt );
		advect (2, v, v0, u0, v0, dt );
		_timer.stop();
		_timer.start("vel_project2");
		project (u, v, u0, v0 );
		_timer.stop();
		_timer.stop();
	}
	
	public inline function reset(field:Int, value:Float = 0.0)
	{
		var i = 0;
		while (i < _blockSize)
		{
			Memory.setDouble(field + i, value);
			i += DATA_SIZE;
		}
	}
	
	public inline function setValue(field:Int, x:Int, y:Int, value:Float) { Memory.setDouble(field + IX(x, y) * DATA_SIZE, value); }
	public inline function getValue(field:Int, x:Int, y:Int) { return Memory.getDouble(field + IX(x, y) * DATA_SIZE); }
	public inline function IX(x:Int, y:Int):Int { return x + y * (N + 2); }
}