3 files changed, 37 insertions, 10 deletions

diff --git a/src/shader/code/collide.glsl b/src/shader/code/collide.glsl
index cef2549..cbddfa9 100644
--- a/src/shader/code/collide.glsl
+++ b/src/shader/code/collide.glsl
@@ -7,8 +7,6 @@ layout (std430, binding=1) buffer bufferCollide{ float collideCells[]; };
 layout (std430, binding=2) buffer bufferStream{  float streamCells[]; };
 layout (std430, binding=3) buffer bufferFluid{   float fluidCells[]; };
 
-const float displayAmplifier = 10.;
-
 /// LBM constants
 
 uniform uint nX;
@@ -26,8 +24,8 @@ const float omega = 1/tau;
 
 /// Vector utilities
 
-float comp(int x, int y, vec2 v) {
-	return x*v.x + y*v.y;
+float comp(int i, int j, vec2 v) {
+	return i*v.x + j*v.y;
 }
 
 float sq(float x) {
@@ -68,9 +66,9 @@ void set(uint x, uint y, int i, int j, float v) {
 
 void setFluid(uint x, uint y, vec2 v, float d) {
 	const uint idx = indexOfFluidVertex(x, y);
-	fluidCells[idx + 0] = float(x) - nX/2;
-	fluidCells[idx + 1] = float(y) - nY/2;
-	fluidCells[idx + 2] = displayAmplifier * norm(v);
+	fluidCells[idx + 0] = v.x;
+	fluidCells[idx + 1] = v.y;
+	fluidCells[idx + 2] = norm(v);
 }
 
 /// Moments
@@ -101,6 +99,10 @@ void main() {
 	const uint x = gl_GlobalInvocationID.x;
 	const uint y = gl_GlobalInvocationID.y;
 
+	if ( !(x < nX && y < nY) ) {
+		return;
+	}
+
 	const float d = density(x,y);
 	const vec2  v = velocity(x,y,d);
 
@@ -112,6 +114,5 @@ void main() {
 			set(x,y,i,j, get(x,y,i,j) + omega * (eq - get(x,y,i,j)));
 		}
 	}
-
 }
 )";
diff --git a/src/shader/code/stream.glsl b/src/shader/code/stream.glsl
index 982b207..9b77c85 100644
--- a/src/shader/code/stream.glsl
+++ b/src/shader/code/stream.glsl
@@ -29,6 +29,10 @@ void main() {
 	const uint x = gl_GlobalInvocationID.x;
 	const uint y = gl_GlobalInvocationID.y;
 
+	if ( !(x < nX && y < nY) ) {
+		return;
+	}
+
 	if ( x != 0 && x != nX-1 && y != 0 && y != nY-1 ) {
 		for ( int i = -1; i <= 1; ++i ) {
 			for ( int j = -1; j <= 1; ++j ) {
diff --git a/src/shader/code/vertex.glsl b/src/shader/code/vertex.glsl
index 516f2c6..e132e5f 100644
--- a/src/shader/code/vertex.glsl
+++ b/src/shader/code/vertex.glsl
@@ -7,6 +7,11 @@ out VS_OUT {
 	vec3 color;
 } vs_out;
 
+uniform uint nX;
+uniform uint nY;
+
+const float displayAmplifier = 10.0;
+
 float unit(float x) {
 	return 1.0/(1.0+exp(-x));
 }
@@ -15,8 +20,25 @@ vec3 getColor(float x) {
 	return x*vec3(1.0,0.0,0.0) + (1-x)*vec3(-0.5,0.0,1.0);
 }
 
+vec2 fluidVertexAtIndex(uint i) {
+	const float y = floor(float(i) / float(nX));
+	return vec2(
+		i - nX*y,
+		y
+	);
+
+}
+
 void main() {
-	gl_Position  = vec4(VertexPosition.xy, 0., 1.);
-	vs_out.color = getColor(unit(VertexPosition.z));
+	const vec2 idx = fluidVertexAtIndex(gl_VertexID);
+
+	gl_Position  = vec4(
+		idx.x - nX/2,
+		idx.y - nY/2,
+		0.,
+		1.
+	);
+
+	vs_out.color = getColor(unit(displayAmplifier * VertexPosition.z));
 }
 )";