Custom PBR Shader Equations in Three.js

This guide explains how to implement custom Physically Based Rendering (PBR) equations in Three.js by extending its built-in shader chunks. By leveraging the onBeforeCompile callback on standard materials, you can inject custom GLSL code to modify light-material interactions, such as customizing the Bidirectional Reflectance Distribution Function (BRDF) for specular or diffuse reflections, without rebuilding the entire rendering engine from scratch.

Understanding Three.js Shader Chunks

Three.js constructs its complex PBR materials, such as MeshStandardMaterial and MeshPhysicalMaterial, using modular GLSL templates called Shader Chunks. These chunks are assembled at runtime into complete vertex and fragment shaders.

To customize PBR calculations, you must modify the chunks responsible for lighting equations. The primary chunk for physical rendering calculations is lights_physical_pars_fragment. This chunk defines key functions like: * D_GGX: The microfacet distribution function. * V_SmithGGXCorrelated: The geometric shadowing/masking function. * F_Schlick: The Fresnel reflection coefficient. * RE_Direct_Physical: The main function that evaluates direct light using these BRDF terms.

Step 1: Hooking into onBeforeCompile

To modify shader chunks before Three.js compiles them to the GPU, use the onBeforeCompile callback of your material. This function exposes the raw shader object, containing the vertex shader, fragment shader, and uniforms.

const material = new THREE.MeshStandardMaterial({
    color: 0x3498db,
    roughness: 0.5,
    metalness: 0.5
});

material.onBeforeCompile = (shader) => {
    // We will inject custom code and uniforms here
    console.log(shader.fragmentShader);
};

Step 2: Defining Custom Uniforms

If your custom PBR equation requires user-defined parameters, you must inject custom uniforms into the shader.

// Define custom uniform values
const customUniforms = {
    uCustomRoughnessExponent: { value: 2.0 }
};

material.onBeforeCompile = (shader) => {
    // Merge custom uniforms into the material's shader program
    shader.uniforms.uCustomRoughnessExponent = customUniforms.uCustomRoughnessExponent;

    // Declare the uniform in the fragment shader header
    shader.fragmentShader = `
        uniform float uCustomRoughnessExponent;
    ` + shader.fragmentShader;
};

Step 3: Modifying the PBR Shader Chunks

To change how light reflects, locate the target GLSL function in lights_physical_pars_fragment and replace it using JavaScript’s string replacement API.

In this example, we will replace the default Trowbridge-Reitz (GGX) specular distribution function (D_GGX) with a customized variation that uses our injected exponent uniform to control the shape of the specular highlight.

material.onBeforeCompile = (shader) => {
    // 1. Inject custom uniforms
    shader.uniforms.uCustomRoughnessExponent = customUniforms.uCustomRoughnessExponent;
    
    shader.fragmentShader = `
        uniform float uCustomRoughnessExponent;
    \n` + shader.fragmentShader;

    // 2. Define the new custom D_GGX function replacement
    const customD_GGX = `
    float D_GGX( const in float alpha, const in float dotNH ) {
        float a2 = pow( alpha, uCustomRoughnessExponent ); // Custom exponent applied here
        float d = ( dotNH * a2 - dotNH ) * dotNH + 1.0;
        return PI2 * a2 / ( d * d );
    }
    `;

    // 3. Find the original D_GGX definition and replace it
    // The default definition in Three.js matches: float D_GGX( const in float alpha, const in float dotNH ) { ... }
    const targetSignature = /float D_GGX[\s\S]*?return[\s\S]*?;[\s\S]*?}/;

    if (targetSignature.test(shader.fragmentShader)) {
        shader.fragmentShader = shader.fragmentShader.replace(targetSignature, customD_GGX);
        console.log("D_GGX successfully modified!");
    } else {
        console.warn("Could not find the default D_GGX function in the shader chunk.");
    }
};

Step 4: Forcing Shader Recompilation (If Needed)

If you modify the properties of customUniforms at runtime, Three.js will update the GPU values automatically. However, if you dynamically change the custom GLSL code structure, you must force a recompile by setting the material’s needsUpdate property to true:

// Run this only if you modify the GLSL logic dynamically after the initial render
material.needsUpdate = true;

By targeting specific mathematical functions inside lights_physical_pars_fragment, you can build non-photorealistic shading models, stylized anisotropic reflections, or specialized material models (like velvet or silk) while retaining the built-in Three.js features like environment maps, shadows, and post-processing compatibility.