BabylonJS MultiRenderTarget particle system example

ParticleSystem.ts

import {
    Constants,
    Color4,
    Mesh,
    Scene,
    RenderTargetTexture,
    Engine,
    Vector3,
    UniversalCamera,
    RawTexture,
    VertexData,
    MeshBuilder,
    Buffer, ShaderMaterial,
    MultiRenderTarget
} from "babylonjs";

import RenderPlane from "@/webgl/renderplane";
import psystemVert from "./shaders/psystem.vert?raw"
import psystemFrag from "./shaders/psystem.frag?raw"


export default class ParticleSystem {

    //@ts-ignore
    mesh: Mesh

    // number of particles
    numParticles: number = 100

    // Flags to differentiate between which buffer is being written to and which is being read from
    ping: number = 0
    pong: number = 1

    // scene for rendering the fbos
    scene: Scene

    // a reference ot the main engine.
    engine: Engine

    camera: UniversalCamera

    // holds fbos / render targets.
    fbos: Array<RenderTargetTexture> = []

    renderPlane: RenderPlane

    texSize: number = 0

    constructor(engine: Engine) {

        this.scene = new Scene(engine);
        this.engine = engine;
        this.renderPlane = new RenderPlane(this.scene)
        this.camera = new UniversalCamera("", new Vector3(0, 0, -40), this.scene)
        this.camera.fov = Math.PI / 4
        this.camera.setTarget(Vector3.Zero());

    }

    /**
     * Sets the number of particles to manage. Will reset FBOs
     * @param num
     */
    setNumParticles(num: number) {
        this.numParticles = num;
        this.fbos = []
        this.setup();
    }

    /**
     * Sets the system up
     * @param clearColor {Color4} a clear color to use.
     */
    setup(clearColor: Color4 = new Color4(0, 0, 0, 0)) {

        let texSize = this._deriveSize()

        for (let i = 0; i < 2; ++i) {
            let rt = new MultiRenderTarget(
                "",
                texSize,
                2,
                this.scene,
                {
                    formats: [
                        Constants.TEXTUREFORMAT_RGBA,
                        Constants.TEXTUREFORMAT_RGBA
                    ],
                    types: [
                        Constants.TEXTURETYPE_FLOAT,
                        Constants.TEXTURETYPE_FLOAT,
                    ],

                    doNotChangeAspectRatio: false
                },
                [
                    "posTex",
                    "metaTex"
                ]
            )
            rt.clearColor = clearColor
            //rt.activeCamera = this.camera

            rt.renderList!.push(this.renderPlane.mesh)

            this.scene.customRenderTargets.push(rt);
            this.fbos.push(rt)
        }

        this.texSize = texSize

        // build position data
        let posData = new Float32Array(texSize * texSize * 4)
        let metaData = new Float32Array(texSize * texSize * 4)

        let size = texSize * texSize * 4;
        for (let i = 0; i < size; i++) {
            posData[i] = Math.random()
        }
        this.setData(posData, {
            texSize
        });
        this.setData(posData, {
            texSize,
            index: 1
        });
        // build metadata
        for (let a = 0; a < size; a += 4) {
            metaData[a] = 1
            metaData[a + 1] = 1
            metaData[a + 2] = 0
            metaData[a + 3] = 1
        }
        this.setData(metaData, {
            texSize,
            layer: 1
        });
        this.setData(metaData, {
            texSize,
            layer: 1,
            index: 1
        });


    }

    /**
     * Sets data onto the FBO
     * @param data {Float32Array} the data to write to the FBO
     * @param texSize {number} the texture size to write to the attachment of the FBO
     * @param index {number} the FBO to write to
     * @param layer {number} for MultiRenderTargets - which attachment to write to
     */
    setData(data: Float32Array, {
        texSize = 512,
        index = 0,
        layer = 0
    } = {}) {
        let size = texSize

        //@ts-ignore
        let tex = new RawTexture(data, size, size, Constants.TEXTUREFORMAT_RGBA, this.scene, false, false, Constants.TEXTURE_TRILINEAR_SAMPLINGMODE, Constants.TEXTURETYPE_FLOAT)

        //@ts-ignore
        this.renderPlane.setData(tex, layer)

        this.fbos[index].render()

    }

    /**
     * Renders the system
     */
    render(engine: Engine) {

        this.renderPlane.setData(this.fbos[this.pong].textures[0])
        this.renderPlane.setData(this.fbos[this.pong].textures[1],1)
        this.fbos[this.ping].render()

        this._swap();

    }

    /**
     * Tries to derive the size of the FBO to use based on the number of particles. Tries to keep power of 2 sizes
     * while being as small as possible
     */
    _deriveSize() {
        let sizes = [512, 256, 128]
        let fSize = 0

        sizes.forEach(size => {
            if ((size * size) > this.numParticles) {
                fSize = size
            }
        })

        return fSize
    }

    /**
     * Builds mesh to render
     * @param scene
     * @param mesh
     */
    buildRenderMesh(scene: Scene, mesh: Mesh | null = null) {

        let vertexData = new VertexData()
        let texSize = this.texSize

        // build uvs when reading textures
        let positions = []

        for (let i = 0; i < texSize; ++i) {
            for (let j = 0; j < texSize; ++j) {
                let ux = i / texSize * 2.0 - 1.0 + .5 / texSize;
                let uy = j / texSize * 2.0 - 1.0 + .5 / texSize;
                //let ux = i / texSize
                //let uy = j / texSize;

                positions.push(ux, uy, 0, 1)
            }
        }

        this.mesh = mesh !== null ? mesh : MeshBuilder.CreateSphere("", {
            segments: 32,
            diameter: 1
        }, scene)

        let meshShader = new ShaderMaterial("", scene, {
                vertexSource: psystemVert,
                fragmentSource: psystemFrag
            },
            {
                attributes: ["position", "iPosition", "uv"],
                uniforms: ["projection", "view", "alpha"],
                samplers: ["positionTexture", "metaTexture"],
                needAlphaBlending: true
            })

        //@ts-ignore
        meshShader.setTexture("positionTexture",this.fbos[1].textures[0])
        //@ts-ignore
        meshShader.setTexture("metaTexture",this.fbos[1].textures[1])

        this.mesh.material = meshShader

        // make instances
        this.mesh.instances = []

        for (let i = 0; i < texSize; ++i) {
            this.mesh.instances.push(
                this.mesh.createInstance(`instance${i}`)
            )
        }

        // Build instanced buffer for lookup ids.
        let posBuffer = new Buffer(this.engine, new Float32Array(positions), true, 0);
        let pBuff = posBuffer.createVertexBuffer("iPosition", 0, 4, 0, true)
        this.mesh.geometry!.setVerticesBuffer(pBuff);

    }

    _swap() {
        let tmp = this.pong;
        this.pong = this.ping;
        this.ping = tmp;
    }
}

psystem.frag

#version 300 es

// basic shader, adjust as needed

in vec4 vColor;
out vec4 glFragColor;

void main() {

    glFragColor = vColor;
}

psystem.vert

#version 300 es

uniform mat4 view;
uniform mat4 projection;
uniform float time;

in vec3 position;
in vec4 iPosition;
in vec2 uv;

uniform sampler2D positionTexture;
uniform sampler2D metaTexture;

out vec2 vUv;
out vec4 vColor;

// basic shader, adjust as needed

void main(){

    vUv = uv;

    vec4 pos = texture(positionTexture,iPosition.xy);
    vec4 meta = texture(metaTexture, iPosition.xy);

    vColor = meta;
    gl_Position =  projection * view * vec4(pos + vec4(position,1.));
}

Renderplane.ts

import {Scene, ShaderMaterial, Mesh, VertexData, Texture} from "babylonjs";

import renderPlaneVert from "./shaders/renderplane.vert?raw"
import renderPlaneFrag from "./shaders/renderplane.frag?raw"

/**
 * Renders a Full screen triangle, use this to write data onto the MRT
 */
export default class RenderPlane {

    scene: Scene
    shader: ShaderMaterial
    mesh: Mesh
    numAttachments:number = 0

    constructor(scene: Scene, {
        numAttachments = 1
    }={}) {

        this.scene = scene;

        this.numAttachments = numAttachments


        this.shader = new ShaderMaterial("", this.scene, {
                vertexSource: renderPlaneVert,
                fragmentSource: renderPlaneFrag
            },
            {
                attributes: ["position"],
                uniforms: [],
                samplers: ["uTex0","uTex1"]
            })

        // build triangle data
        let vertexData = new VertexData();
        vertexData.positions = new Float32Array([
            -1, -1, -1, 4, 4, -1
        ]);

        // need at least one face.
        vertexData.indices = [0, 1, 2];

        this.mesh = new Mesh(`RenderPlaneMesh${Math.random()}`, scene);
        vertexData.applyToMesh(this.mesh);

        this.mesh.material = this.shader;
    }

    setData(tex: Texture, attachment: number = 0) {
        this.shader.setTexture(`uTex${attachment}`, tex);
    }
}