Lattice Boltzman WebGL translation results in NaNs everywhere - glsl

Question

I came across this implementation of Lattice Boltzmann fluid solver from this blog which goes over its implementation. I decided I wanted to translate this to ShaderToy with webgl. Instead of computing each step necessary for equilibrium, the limitations of shader toy force me to use multiple texture frames.
My algorithm is basically like this:
Calculate bulk velocity and density for each cell in texture A from textures B, C, and D which hold the 9 direction components (NW,N,NE,W in each vec4 in B, E,SW,S,SE in each vec4 in C, center in each vec4 in D).
After these components are calculated, recalculate stream/advection for each necessary cell for each texture frame, B,C,D, and take the density and velocity components from A to compute the equilibrium.
set the final values for each direction to new_direction - (new_direction - new_direction_equilibrium), ie(north_west - (north_west - north_west_eq))`. This isn't any different from the reference code.
Additionally I wrap coordinates around, so there are no boundary conditions in order to avoid dealing with boundary logic, and the user clicks to cause disturbances in the lattice, which sets the "not moving" direction value to a number.
In my code however, I end up getting lots of NaNs (colored here in white, red represents the density) and I'm not sure why. I put safe guards in locations in my code to avoid zero densities causing issues, but that didn't seem to do anything.
You can test this out on shadertoy, but my code is as follows:
//COMMON functions
const int DIRECTION_COUNT = 9;
const int DIMENSION_COUNT = 2;
const float LATTICE_SPEED = 0.1;
const float TAU = 0.9;
const vec2 north_offset = vec2(0.0,1.0);
const vec2 north_west_offset = vec2(-1.0,1.0);
const vec2 north_east_offset = vec2(1.0,1.0);
const vec2 west_offset = vec2(-1.0,0.0);
const vec2 east_offset = vec2(1.0,0.0);
const vec2 south_offset = vec2(0.0,-1.0);
const vec2 south_west_offset = vec2(-1.0,-1.0);
const vec2 south_east_offset = vec2(1.0,-1.0);
const vec2 center_offset = vec2(0.0,0.0);
const vec2 offsets[DIRECTION_COUNT] = vec2[DIRECTION_COUNT](
north_west_offset,
north_offset,
north_east_offset,
west_offset,
center_offset,
east_offset,
south_west_offset,
south_offset,
south_east_offset);
const int north_west_tex_idx = 0;
const int north_tex_idx = 1;
const int north_east_tex_idx = 2;
const int west_tex_idx = 3;
const int east_tex_idx = 0;
const int south_west_tex_idx = 1;
const int south_tex_idx = 2;
const int south_east_tex_idx = 3;
const int center_tex_idx = 0;
float textureN(sampler2D NW_N_NE_W_channel, vec2 coord, vec2 resolution){
vec2 offset_coord = coord + south_offset;
return texture(NW_N_NE_W_channel, offset_coord/resolution)[north_tex_idx];
}
float textureNW(sampler2D NW_N_NE_W_channel, vec2 coord, vec2 resolution){
vec2 offset_coord = coord + south_east_offset;
return texture(NW_N_NE_W_channel, offset_coord/resolution)[north_west_tex_idx];
}
float textureNE(sampler2D NW_N_NE_W_channel, vec2 coord, vec2 resolution){
vec2 offset_coord = coord + south_west_offset;
return texture(NW_N_NE_W_channel, offset_coord/resolution)[north_east_tex_idx];
}
float textureW(sampler2D NW_N_NE_W_channel, vec2 coord, vec2 resolution){
vec2 offset_coord = coord + east_offset;
return texture(NW_N_NE_W_channel, offset_coord/resolution)[west_tex_idx];
}
float textureS(sampler2D E_SW_S_SE_channel, vec2 coord, vec2 resolution){
vec2 offset_coord = coord + north_offset;
return texture(E_SW_S_SE_channel, offset_coord/resolution)[south_tex_idx];
}
float textureSW(sampler2D E_SW_S_SE_channel, vec2 coord, vec2 resolution){
vec2 offset_coord = coord + north_east_offset;
return texture(E_SW_S_SE_channel, offset_coord/resolution)[south_west_tex_idx];
}
float textureSE(sampler2D E_SW_S_SE_channel, vec2 coord, vec2 resolution){
vec2 offset_coord = coord + north_west_offset;
return texture(E_SW_S_SE_channel, offset_coord/resolution)[south_east_tex_idx];
}
float textureE(sampler2D E_SW_S_SE_channel, vec2 coord, vec2 resolution){
vec2 offset_coord = coord + west_offset;
return texture(E_SW_S_SE_channel, offset_coord/resolution)[east_tex_idx];
}
float textureC(sampler2D C_channel, vec2 coord, vec2 resolution){
vec2 offset_coord = coord + center_offset;
return texture(C_channel, offset_coord/resolution)[center_tex_idx];
}
float calc_equilibrium(const in float density,
const in vec2 velocity,
const in ivec2 ij) {
int i = ij.x;
int j = ij.y;
// u . u
float velmag = dot(velocity, velocity);
// Compute the weight.
float weight;
if(i == 0 && j == 0) {
weight = 4.0 / 9.0;
} else if(i == 0 || j == 0) {
weight = 1.0 / 9.0;
} else {
weight = 1.0 / 36.0;
}
// e_i . u
float dotprod = float(i) * velocity.x + float(j) * velocity.y;
float sum = 1.0;
sum += (3.0 / LATTICE_SPEED) * dotprod;
sum += (4.5 / (LATTICE_SPEED * LATTICE_SPEED)) * dotprod * dotprod;
sum -= (1.5 / (LATTICE_SPEED * LATTICE_SPEED)) * velmag;
if(density == 0.0){
return 0.0;
}
return weight * density * sum;
}
|
//Buffer A, takes in B, C, and D as in put in that order
float[DIRECTION_COUNT] stream_all(
sampler2D NW_N_NE_W_channel,
sampler2D E_SW_S_SE_channel,
sampler2D C_channel,
in vec2 ifragCoord){
float north_west = textureNW(NW_N_NE_W_channel, ifragCoord, iResolution.xy);
float north = textureN(NW_N_NE_W_channel, ifragCoord, iResolution.xy);
float north_east = textureNE(NW_N_NE_W_channel, ifragCoord, iResolution.xy);
float west = textureW(NW_N_NE_W_channel, ifragCoord, iResolution.xy);
float east = textureE(E_SW_S_SE_channel, ifragCoord, iResolution.xy);
float south_west = textureSW(E_SW_S_SE_channel, ifragCoord, iResolution.xy);
float south = textureS(E_SW_S_SE_channel, ifragCoord, iResolution.xy);
float south_east = textureSE(E_SW_S_SE_channel, ifragCoord, iResolution.xy);
float center = textureC(C_channel, ifragCoord, iResolution.xy);
return float[DIRECTION_COUNT](
north_west, north, north_east, west, center, east, south_west, south, south_east
);
}
float calc_density(const in float new_directions[DIRECTION_COUNT]) {
float density;
for(int i = 0; i < DIRECTION_COUNT; ++i){
density += new_directions[i];
}
return density;
}
vec2 calc_velocity(const in float new_directions[DIRECTION_COUNT], const in float density) {
if(density == 0.0){
return vec2(0.0);
}
if(isinf(density)){
return vec2(0.0);
}
// Compute target indices.
vec2 velocity = vec2(0.0);
for(int idx = 0; idx < DIRECTION_COUNT; ++idx){
vec2 ij = offsets[idx];
float i = ij.x;
float j = ij.y;
velocity.x += new_directions[idx] * (i);
velocity.y += new_directions[idx] * (j);
}
return velocity * (LATTICE_SPEED/density);
}
void mainImage( out vec4 fragColor, in vec2 fragCoord )
{
ivec2 ifragCoord = ivec2(fragCoord);
float new_directions[DIRECTION_COUNT] = stream_all(iChannel0, iChannel1, iChannel2, fragCoord);
float density = calc_density(new_directions);
vec2 velocity = calc_velocity(new_directions, density);
fragColor = vec4(density,velocity.x,velocity.y,0.0);
float center = textureC(iChannel2, fragCoord, iResolution.xy);
float debug = center;
if(isnan(density)){
debug = 1.0;
fragColor.w = debug;
}
//fragColor = vec4(1.0);
}
|
//Buffer B, takes in B, and A in that order
void mainImage( out vec4 fragColor, in vec2 fragCoord )
{
if(iFrame < 10){
fragColor = vec4(0.0);
return;
}
ivec2 ifragCoord = ivec2(fragCoord - 0.5);
float north_west = textureNW(iChannel0, fragCoord, iResolution.xy);
float north = textureN(iChannel0, fragCoord, iResolution.xy);
float north_east = textureNE(iChannel0, fragCoord, iResolution.xy);
float west = textureW(iChannel0, fragCoord, iResolution.xy);
vec4 density_velocity = texelFetch(iChannel1, ifragCoord, 0);
float density = density_velocity.x;
vec2 velocity = density_velocity.yz;
float north_west_eq = calc_equilibrium(density, velocity, ivec2(north_west_offset));
float north_eq = calc_equilibrium(density, velocity, ivec2(north_offset));
float north_east_eq = calc_equilibrium(density, velocity, ivec2(north_east_offset));
float west_eq = calc_equilibrium(density, velocity, ivec2(west_offset));
fragColor = vec4((north_west - (north_west - north_west_eq) / TAU),
(north - (north - north_eq) / TAU),
(north_east - (north_east - north_east_eq) / TAU),
(west - (west - west_eq) / TAU));
}
|
//Buffer C, takes in C and A in that order.
void mainImage( out vec4 fragColor, in vec2 fragCoord )
{
if(iFrame < 10){
fragColor = vec4(0.0);
return;
}
if(iFrame < 30 && fragCoord.y < -1.0){
fragColor = vec4(10.0, 0.0,10.0,0.0);
return;
}
ivec2 ifragCoord = ivec2(fragCoord - 0.5);
float east = textureE(iChannel0, fragCoord, iResolution.xy);
float south_west = textureSW(iChannel0, fragCoord, iResolution.xy);
float south = textureS(iChannel0, fragCoord, iResolution.xy);
float south_east = textureSE(iChannel0, fragCoord, iResolution.xy);
vec4 density_velocity = texelFetch(iChannel1, ifragCoord, 0);
float density = density_velocity.x;
vec2 velocity = density_velocity.yz;
float east_eq = calc_equilibrium(density, velocity, ivec2(east_offset));
float south_west_eq = calc_equilibrium(density, velocity, ivec2(south_west_offset));
float south_eq = calc_equilibrium(density, velocity, ivec2(south_offset));
float south_east_eq = calc_equilibrium(density, velocity, ivec2(south_east_offset));
fragColor = vec4((east - (east - east_eq) / TAU),
(south_west - (south_west - south_west_eq) / TAU),
(south - (south - south_eq) / TAU),
(south_east - (south_east - south_east_eq) / TAU));
}
|
//Buffer D takes in D and A in that order
void mainImage( out vec4 fragColor, in vec2 fragCoord )
{
if(iFrame < 10){
fragColor = vec4(1, 0.0,0.0,0.0);
return;
}
ivec2 ifragCoord = ivec2(fragCoord - 0.5);
float center = textureC(iChannel0, fragCoord, iResolution.xy);
vec4 density_velocity = texelFetch(iChannel1, ifragCoord, 0);
float density = density_velocity.x;
vec2 velocity = density_velocity.yz;
float center_eq = calc_equilibrium(density, velocity, ivec2(center_offset));
fragColor = vec4((center - (center - center_eq) / TAU),
0.0,
0.0,
0.0);
vec2 mouse = vec2(iMouse.zw);
if(mouse.x > 0.0 && mouse.y > 0.0){
vec2 current_mouse = vec2(iMouse.xy);
if(distance(fragCoord, current_mouse) < 3.0){
fragColor.r = vec4(10.0).r;
}
}
}
|
//main image output, only takes in A as an iChannel
void mainImage( out vec4 fragColor, in vec2 fragCoord )
{
ivec2 ifragCoord = ivec2(fragCoord-0.5);
vec4 density_velocity = texelFetch(iChannel0, ifragCoord, 0);
float density = density_velocity.r;
vec2 velocity = density_velocity.gb;
float vel_length = length(velocity);
velocity = normalize(velocity);
//Output to screen
//fragColor = vec4(abs(velocity),density/100.0,vel_length/100.0);
//fragColor = vec4(abs(velocity),0.0,1.0);
fragColor = vec4(density/10.0,0.0,0.0,1.0);
//
if(density_velocity.w == 1.0){
fragColor = vec4(1.0);
}
}
What have I done incorrectly to result in all of these Nans? Is there a way to stop them?

Answer 1

Clamping the return value from calc_equilibrium should avoid the white NaN blooms.
return clamp(weight * density * sum, -1000.0, 1000.0);
Preventing the red/black noise blooms does not appear to be so simple.
For every frame that occurs while the mouse button is held down, a lot of energy is being added to the system, and at some point it is bound to boil over.