js/fluid.js

/**

 * 🌿 Ivy's GPU Art Studio

 * Tab 2: Fluid Simulation

 *

 * GPU-accelerated fluid dynamics using compute shaders

 * Based on Jos Stam's "Stable Fluids" algorithm

 * Enhanced with styles, palettes, and effects!

 */

class FluidRenderer {
    constructor() {
        this.device = null;
        this.context = null;
        this.format = null;

        // Simulation parameters
        this.params = {
            style: 0, // 0=classic, 1=ivy, 2=ink, 3=smoke, 4=plasma, 5=watercolor
            palette: 0, // 0=ivy, 1=rainbow, 2=fire, 3=ocean, 4=neon, 5=sunset, 6=cosmic, 7=mono
            viscosity: 0.1,
            diffusion: 0.0001,
            force: 100,
            curl: 30,
            pressure: 0.8,
            bloom: true,
            vortex: false
        };

        // Simulation state
        this.gridSize = 256;
        this.velocityBuffers = [];
        this.densityBuffers = [];
        this.currentBuffer = 0;

        this.input = null;
        this.animationLoop = null;
        this.isActive = false;
        this.time = 0;

        // Previous mouse position for velocity
        this.prevMouseX = 0.5;
        this.prevMouseY = 0.5;
    }

    async init(device, context, format, canvas) {
        this.device = device;
        this.context = context;
        this.format = format;
        this.canvas = canvas;

        // Create simulation buffers
        this.createBuffers();

        // Create pipelines
        await this.createPipelines();

        // Setup input
        this.input = new WebGPUUtils.InputHandler(canvas);

        // Animation loop
        this.animationLoop = new WebGPUUtils.AnimationLoop((dt, totalTime) => {
            this.time = totalTime;
            this.simulate(dt);
            this.render();
        });
    }

    createBuffers() {
        const size = this.gridSize * this.gridSize;

        // Double buffering for velocity (vec2) and density (f32)
        for (let i = 0; i < 2; i++) {
            this.velocityBuffers.push(
                this.device.createBuffer({
                    label: `Velocity Buffer ${i}`,
                    size: size * 8, // vec2f = 8 bytes
                    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
                })
            );

            this.densityBuffers.push(
                this.device.createBuffer({
                    label: `Density Buffer ${i}`,
                    size: size * 16, // vec4f for RGBA = 16 bytes
                    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
                })
            );
        }

        // Uniform buffer
        this.uniformBuffer = this.device.createBuffer({
            label: "Fluid Uniforms",
            size: 64,
            usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
        });

        // Initialize with zeros
        const zeroVelocity = new Float32Array(size * 2);
        const zeroDensity = new Float32Array(size * 4);

        for (let i = 0; i < 2; i++) {
            this.device.queue.writeBuffer(this.velocityBuffers[i], 0, zeroVelocity);
            this.device.queue.writeBuffer(this.densityBuffers[i], 0, zeroDensity);
        }
    }

    async createPipelines() {
        // Compute shader for simulation
        const computeShader = this.device.createShaderModule({
            label: "Fluid Compute Shader",
            code: this.getComputeShaderCode()
        });

        // Render shader for display
        const renderShader = this.device.createShaderModule({
            label: "Fluid Render Shader",
            code: this.getRenderShaderCode()
        });

        // Bind group layouts
        this.computeBindGroupLayout = this.device.createBindGroupLayout({
            entries: [
                { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "uniform" } },
                { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
                { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
                { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
                { binding: 4, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } }
            ]
        });

        this.renderBindGroupLayout = this.device.createBindGroupLayout({
            entries: [
                { binding: 0, visibility: GPUShaderStage.FRAGMENT, buffer: { type: "uniform" } },
                { binding: 1, visibility: GPUShaderStage.FRAGMENT, buffer: { type: "read-only-storage" } },
                { binding: 2, visibility: GPUShaderStage.FRAGMENT, buffer: { type: "read-only-storage" } }
            ]
        });

        // Compute pipeline
        this.computePipeline = this.device.createComputePipeline({
            label: "Fluid Compute Pipeline",
            layout: this.device.createPipelineLayout({
                bindGroupLayouts: [this.computeBindGroupLayout]
            }),
            compute: {
                module: computeShader,
                entryPoint: "main"
            }
        });

        // Render pipeline
        this.renderPipeline = this.device.createRenderPipeline({
            label: "Fluid Render Pipeline",
            layout: this.device.createPipelineLayout({
                bindGroupLayouts: [this.renderBindGroupLayout]
            }),
            vertex: {
                module: renderShader,
                entryPoint: "vertexMain"
            },
            fragment: {
                module: renderShader,
                entryPoint: "fragmentMain",
                targets: [{ format: this.format }]
            },
            primitive: {
                topology: "triangle-list"
            }
        });

        // Create bind groups
        this.updateBindGroups();
    }

    updateBindGroups() {
        const curr = this.currentBuffer;
        const next = 1 - curr;

        this.computeBindGroup = this.device.createBindGroup({
            layout: this.computeBindGroupLayout,
            entries: [
                { binding: 0, resource: { buffer: this.uniformBuffer } },
                { binding: 1, resource: { buffer: this.velocityBuffers[curr] } },
                { binding: 2, resource: { buffer: this.velocityBuffers[next] } },
                { binding: 3, resource: { buffer: this.densityBuffers[curr] } },
                { binding: 4, resource: { buffer: this.densityBuffers[next] } }
            ]
        });

        this.renderBindGroup = this.device.createBindGroup({
            layout: this.renderBindGroupLayout,
            entries: [
                { binding: 0, resource: { buffer: this.uniformBuffer } },
                { binding: 1, resource: { buffer: this.velocityBuffers[next] } },
                { binding: 2, resource: { buffer: this.densityBuffers[next] } }
            ]
        });
    }

    start() {
        this.isActive = true;
        console.log("🌊 FluidRenderer started!");
        this.animationLoop.start();
    }

    stop() {
        this.isActive = false;
        console.log("🌊 FluidRenderer stopped");
        this.animationLoop.stop();
    }

    reset() {
        const size = this.gridSize * this.gridSize;
        const zeroVelocity = new Float32Array(size * 2);
        const zeroDensity = new Float32Array(size * 4);

        for (let i = 0; i < 2; i++) {
            this.device.queue.writeBuffer(this.velocityBuffers[i], 0, zeroVelocity);
            this.device.queue.writeBuffer(this.densityBuffers[i], 0, zeroDensity);
        }
    }

    setViscosity(value) {
        this.params.viscosity = value;
    }

    setDiffusion(value) {
        this.params.diffusion = value;
    }

    setForce(value) {
        this.params.force = value;
    }

    setColorMode(mode) {
        const modes = { ink: 0, fire: 1, rainbow: 2, smoke: 3, ivy: 4 };
        this.params.colorMode = modes[mode] || 0;
    }

    setStyle(style) {
        const styles = { classic: 0, ivy: 1, ink: 2, smoke: 3, plasma: 4, watercolor: 5 };
        this.params.style = styles[style] ?? 0;
    }

    setPalette(palette) {
        const palettes = { ivy: 0, rainbow: 1, fire: 2, ocean: 3, neon: 4, sunset: 5, cosmic: 6, monochrome: 7 };
        this.params.palette = palettes[palette] ?? 0;
    }

    setCurl(value) {
        this.params.curl = value;
    }

    setPressure(value) {
        this.params.pressure = value;
    }

    setBloom(enabled) {
        this.params.bloom = enabled;
    }

    setVortex(enabled) {
        this.params.vortex = enabled;
    }

    simulate(dt) {
        if (!this.isActive) return;

        // Auto-spawn some fluid for visual feedback even without mouse
        const autoSpawn = !this.input.isPressed;
        let mouseX = this.input.mouseX;
        let mouseY = this.input.mouseY;
        let isPressed = this.input.isPressed;

        // Auto animation when not interacting
        if (autoSpawn && this.time > 0) {
            // Create swirling patterns automatically
            const t = this.time * 0.5;
            mouseX = 0.5 + 0.3 * Math.sin(t);
            mouseY = 0.5 + 0.3 * Math.cos(t * 0.7);
            isPressed = true; // Simulate mouse press for auto-spawn
        }

        // Calculate mouse velocity
        const dx = (mouseX - this.prevMouseX) * this.params.force;
        const dy = (mouseY - this.prevMouseY) * this.params.force;
        this.prevMouseX = mouseX;
        this.prevMouseY = mouseY;

        // Update uniforms - expanded for new params
        const uniforms = new Float32Array([
            this.gridSize, // 0: grid size
            dt, // 1: delta time
            this.params.viscosity, // 2: viscosity
            this.params.diffusion, // 3: diffusion
            mouseX, // 4: mouse X
            mouseY, // 5: mouse Y
            dx, // 6: velocity X
            dy, // 7: velocity Y
            isPressed ? 1.0 : 0.0, // 8: is mouse pressed
            this.params.style, // 9: style
            this.params.palette, // 10: palette
            this.params.curl, // 11: curl/vorticity
            this.params.pressure, // 12: pressure
            this.params.bloom ? 1.0 : 0.0, // 13: bloom
            this.params.vortex ? 1.0 : 0.0, // 14: vortex
            this.time // 15: time
        ]);

        this.device.queue.writeBuffer(this.uniformBuffer, 0, uniforms);

        // Update bind groups with current buffer state
        this.updateBindGroups();

        // Run compute shader
        const commandEncoder = this.device.createCommandEncoder();
        const computePass = commandEncoder.beginComputePass();

        computePass.setPipeline(this.computePipeline);
        computePass.setBindGroup(0, this.computeBindGroup);
        computePass.dispatchWorkgroups(Math.ceil(this.gridSize / 8), Math.ceil(this.gridSize / 8));

        computePass.end();
        this.device.queue.submit([commandEncoder.finish()]);

        // Swap buffers
        this.currentBuffer = 1 - this.currentBuffer;
    }

    render() {
        if (!this.isActive) return;

        WebGPUUtils.resizeCanvasToDisplaySize(this.canvas, window.devicePixelRatio);

        // IMPORTANT: Create render bind group to read the LATEST buffer (after compute)
        const curr = this.currentBuffer;
        const renderBindGroup = this.device.createBindGroup({
            layout: this.renderBindGroupLayout,
            entries: [
                { binding: 0, resource: { buffer: this.uniformBuffer } },
                { binding: 1, resource: { buffer: this.velocityBuffers[curr] } },
                { binding: 2, resource: { buffer: this.densityBuffers[curr] } }
            ]
        });

        const commandEncoder = this.device.createCommandEncoder();
        const renderPass = commandEncoder.beginRenderPass({
            colorAttachments: [
                {
                    view: this.context.getCurrentTexture().createView(),
                    clearValue: { r: 0, g: 0, b: 0, a: 1 },
                    loadOp: "clear",
                    storeOp: "store"
                }
            ]
        });

        renderPass.setPipeline(this.renderPipeline);
        renderPass.setBindGroup(0, renderBindGroup);
        renderPass.draw(3);
        renderPass.end();

        this.device.queue.submit([commandEncoder.finish()]);
    }

    getComputeShaderCode() {
        return /* wgsl */ `

            struct Uniforms {

                gridSize: f32,

                dt: f32,

                viscosity: f32,

                diffusion: f32,

                mouseX: f32,

                mouseY: f32,

                velX: f32,

                velY: f32,

                mousePressed: f32,

                style: f32,

                palette: f32,

                curl: f32,

                pressure: f32,

                doBloom: f32,

                doVortex: f32,

                time: f32,

            }



            @group(0) @binding(0) var<uniform> u: Uniforms;

            @group(0) @binding(1) var<storage, read> velIn: array<vec2f>;

            @group(0) @binding(2) var<storage, read_write> velOut: array<vec2f>;

            @group(0) @binding(3) var<storage, read> densIn: array<vec4f>;

            @group(0) @binding(4) var<storage, read_write> densOut: array<vec4f>;



            fn idx(x: i32, y: i32) -> u32 {

                let size = i32(u.gridSize);

                let cx = clamp(x, 0, size - 1);

                let cy = clamp(y, 0, size - 1);

                return u32(cy * size + cx);

            }



            fn getPaletteColor(t: f32, paletteId: i32) -> vec3f {

                let tt = fract(t);



                if (paletteId == 0) { // Ivy Green

                    return vec3f(0.13 * tt + 0.05, 0.77 * tt + 0.2, 0.37 * tt + 0.1);

                } else if (paletteId == 1) { // Rainbow

                    return vec3f(

                        0.5 + 0.5 * sin(tt * 6.28 + 0.0),

                        0.5 + 0.5 * sin(tt * 6.28 + 2.094),

                        0.5 + 0.5 * sin(tt * 6.28 + 4.188)

                    );

                } else if (paletteId == 2) { // Fire

                    return vec3f(tt, tt * 0.4, tt * 0.1);

                } else if (paletteId == 3) { // Ocean

                    return vec3f(0.1 * tt, 0.4 * tt + 0.1, 0.9 * tt + 0.1);

                } else if (paletteId == 4) { // Neon

                    return vec3f(

                        0.5 + 0.5 * sin(tt * 12.0),

                        0.5 + 0.5 * sin(tt * 12.0 + 2.0),

                        0.5 + 0.5 * sin(tt * 12.0 + 4.0)

                    );

                } else if (paletteId == 5) { // Sunset

                    return vec3f(0.9 * tt + 0.1, 0.4 * tt, 0.3 * tt + 0.1);

                } else if (paletteId == 6) { // Cosmic

                    return vec3f(0.3 * tt + 0.1, 0.1 * tt + 0.05, 0.8 * tt + 0.2);

                } else { // Monochrome

                    return vec3f(tt * 0.9 + 0.1);

                }

            }



            @compute @workgroup_size(8, 8)

            fn main(@builtin(global_invocation_id) gid: vec3u) {

                let size = i32(u.gridSize);

                let x = i32(gid.x);

                let y = i32(gid.y);



                if (x >= size || y >= size) {

                    return;

                }



                let i = idx(x, y);

                let paletteId = i32(u.palette);



                // Read previous state

                var newVel = velIn[i];

                var newDens = densIn[i];



                // Get neighbors for diffusion

                let vL = velIn[idx(x - 1, y)];

                let vR = velIn[idx(x + 1, y)];

                let vU = velIn[idx(x, y + 1)];

                let vD = velIn[idx(x, y - 1)];



                let dL = densIn[idx(x - 1, y)];

                let dR = densIn[idx(x + 1, y)];

                let dU = densIn[idx(x, y + 1)];

                let dD = densIn[idx(x, y - 1)];



                // Apply diffusion (controlled by diffusion parameter)

                let diffAmount = u.diffusion * 1000.0;

                newVel = mix(newVel, (vL + vR + vU + vD) * 0.25, diffAmount);

                newDens = mix(newDens, (dL + dR + dU + dD) * 0.25, diffAmount);



                // Apply viscosity (dampens velocity)

                newVel *= (1.0 - u.viscosity * 0.1);



                // Vorticity / curl effect

                if (u.doVortex > 0.5) {

                    let curlAmount = u.curl * 0.0005;

                    let vortex = (vR.y - vL.y) - (vU.x - vD.x);

                    newVel += vec2f(-vortex, vortex) * curlAmount;

                }



                // Add forces from mouse

                let fx = f32(x) / f32(size);

                let fy = f32(y) / f32(size);

                let dist = distance(vec2f(fx, fy), vec2f(u.mouseX, u.mouseY));

                let radius = 0.02 + (u.pressure * 0.1); // Pressure affects brush size



                if (dist < radius && u.mousePressed > 0.5) {

                    let strength = 1.0 - dist / radius;

                    // Force affects velocity strength

                    let forceMultiplier = u.velX * u.velX + u.velY * u.velY;

                    newVel += vec2f(u.velX, u.velY) * strength * u.dt * 2.0;



                    // Add density/color using palette

                    let colorHue = strength + u.time * 0.1;

                    let color = getPaletteColor(colorHue, paletteId);

                    newDens += vec4f(color * strength * 3.0, strength * 3.0);

                }



                // Apply pressure (affects how much velocity is preserved)

                newVel *= u.pressure;



                // Decay

                newVel *= 0.995;

                newDens *= 0.992;



                // Boundary conditions

                if (x <= 1 || x >= size - 2 || y <= 1 || y >= size - 2) {

                    newVel *= 0.5;

                }



                velOut[i] = newVel;

                densOut[i] = newDens;

            }

        `;
    }

    getRenderShaderCode() {
        return /* wgsl */ `

            struct Uniforms {

                gridSize: f32,

                dt: f32,

                viscosity: f32,

                diffusion: f32,

                mouseX: f32,

                mouseY: f32,

                velX: f32,

                velY: f32,

                mousePressed: f32,

                style: f32,

                palette: f32,

                curl: f32,

                pressure: f32,

                doBloom: f32,

                doVortex: f32,

                time: f32,

            }



            @group(0) @binding(0) var<uniform> u: Uniforms;

            @group(0) @binding(1) var<storage, read> velocity: array<vec2f>;

            @group(0) @binding(2) var<storage, read> density: array<vec4f>;



            struct VertexOutput {

                @builtin(position) position: vec4f,

                @location(0) uv: vec2f,

            }



            @vertex

            fn vertexMain(@builtin(vertex_index) vertexIndex: u32) -> VertexOutput {

                var pos = array<vec2f, 3>(

                    vec2f(-1.0, -1.0),

                    vec2f(3.0, -1.0),

                    vec2f(-1.0, 3.0)

                );



                var output: VertexOutput;

                output.position = vec4f(pos[vertexIndex], 0.0, 1.0);

                output.uv = pos[vertexIndex] * 0.5 + 0.5;

                return output;

            }



            fn getPaletteColor(t: f32, paletteId: i32) -> vec3f {

                let tt = fract(t);



                if (paletteId == 0) { // Ivy Green

                    return vec3f(0.1 + 0.2 * tt, 0.5 + 0.5 * tt, 0.2 + 0.3 * tt);

                } else if (paletteId == 1) { // Rainbow

                    return vec3f(

                        0.5 + 0.5 * cos(6.28318 * (tt + 0.0)),

                        0.5 + 0.5 * cos(6.28318 * (tt + 0.33)),

                        0.5 + 0.5 * cos(6.28318 * (tt + 0.67))

                    );

                } else if (paletteId == 2) { // Fire

                    return vec3f(min(1.0, tt * 2.5), tt * tt, tt * tt * tt * 0.3);

                } else if (paletteId == 3) { // Ocean

                    return vec3f(0.0 + 0.2 * tt, 0.3 + 0.4 * tt, 0.6 + 0.4 * tt);

                } else if (paletteId == 4) { // Neon

                    return vec3f(

                        0.5 + 0.5 * sin(tt * 12.56),

                        0.5 + 0.5 * sin(tt * 12.56 + 2.094),

                        0.5 + 0.5 * sin(tt * 12.56 + 4.188)

                    );

                } else if (paletteId == 5) { // Sunset

                    return vec3f(0.9 - 0.2 * tt, 0.3 + 0.4 * tt, 0.3 + 0.5 * tt);

                } else if (paletteId == 6) { // Cosmic

                    return vec3f(

                        0.2 + 0.5 * sin(tt * 6.28),

                        0.1 + 0.3 * sin(tt * 6.28 + 2.0),

                        0.5 + 0.5 * sin(tt * 6.28 + 4.0)

                    );

                } else { // Monochrome

                    return vec3f(tt, tt, tt);

                }

            }



            @fragment

            fn fragmentMain(input: VertexOutput) -> @location(0) vec4f {

                let size = i32(u.gridSize);

                let x = i32(input.uv.x * f32(size));

                let y = i32(input.uv.y * f32(size));

                let i = u32(clamp(y, 0, size - 1) * size + clamp(x, 0, size - 1));



                let d = density[i];

                let v = velocity[i];



                let style = i32(u.style);

                let paletteId = i32(u.palette);

                let speed = length(v);

                let dens = length(d.rgb);



                var color = vec3f(0.0);



                // Show mouse position as a dot for visual feedback

                let mouseDist = distance(input.uv, vec2f(u.mouseX, u.mouseY));

                let mouseGlow = smoothstep(0.08, 0.0, mouseDist) * 0.5;



                // Style-based rendering

                if (style == 0) { // Classic - use density color directly

                    color = d.rgb;

                } else if (style == 1) { // Ivy Flow - organic green tones

                    let hue = dens * 0.3 + speed * 0.1;

                    color = getPaletteColor(hue, paletteId);

                    color *= dens * 1.5;

                } else if (style == 2) { // Ink Drop - high contrast

                    color = getPaletteColor(dens + speed * 0.2, paletteId);

                    color = pow(color * dens, vec3f(0.8));

                } else if (style == 3) { // Smoke - soft gradient

                    let smoke = smoothstep(0.0, 1.0, dens);

                    color = mix(vec3f(0.02), getPaletteColor(speed * 0.5, paletteId), smoke);

                } else if (style == 4) { // Plasma - vibrant swirls

                    let plasma = sin(dens * 10.0 + u.time) * 0.5 + 0.5;

                    color = getPaletteColor(plasma + speed * 0.3, paletteId);

                    color *= dens * 2.0;

                } else { // Watercolor - soft bleeding edges

                    let wc = smoothstep(0.0, 0.5, dens);

                    color = getPaletteColor(dens * 0.5 + u.time * 0.05, paletteId) * wc;

                    color = mix(color, vec3f(1.0), (1.0 - wc) * 0.1);

                }



                // Velocity-based highlights

                color += getPaletteColor(0.8, paletteId) * speed * 0.15;



                // Add mouse indicator

                color += getPaletteColor(u.time * 0.2, paletteId) * mouseGlow;



                // Vortex visualization

                if (u.doVortex > 0.5) {

                    // Approximate curl from velocity

                    let curlVis = abs(v.x - v.y) * 0.5;

                    color += vec3f(curlVis * 0.3, curlVis * 0.1, curlVis * 0.4);

                }



                // Bloom effect

                if (u.doBloom > 0.5) {

                    let bloom = max(0.0, dens - 0.5) * 2.0;

                    color += color * bloom * 0.5;

                    color = color / (1.0 + color * 0.3); // Tone mapping

                }



                return vec4f(color, 1.0);

            }

        `;
    }
}

// Export
window.FluidRenderer = FluidRenderer;