three.js/examples/jsm/pmrem/PMREMGenerator.js

/**
 * @author Emmett Lalish / elalish
 *
 * This class generates a Prefiltered, Mipmapped Radiance Environment Map
 * (PMREM) from a cubeMap environment texture. This allows different levels of
 * blur to be quickly accessed based on material roughness. It is packed into a
 * special CubeUV format that allows us to perform custom interpolation so that
 * we can support nonlinear formats such as RGBE. Unlike a traditional mipmap
 * chain, it only goes down to the LOD_MIN level (above), and then creates extra
 * even more filtered 'mips' at the same LOD_MIN resolution, associated with
 * higher roughness levels. In this way we maintain resolution to smoothly
 * interpolate diffuse lighting while limiting sampling computation.
 */

import {
	BufferAttribute,
	BufferGeometry,
	CubeUVReflectionMapping,
	GammaEncoding,
	LinearEncoding,
	LinearToneMapping,
	Mesh,
	NearestFilter,
	NoBlending,
	OrthographicCamera,
	PerspectiveCamera,
	RGBDEncoding,
	RGBEEncoding,
	RGBEFormat,
	RGBM16Encoding,
	RGBM7Encoding,
	RawShaderMaterial,
	Scene,
	UnsignedByteType,
	Vector2,
	Vector3,
	WebGLRenderTarget,
	sRGBEncoding
} from "../../../build/three.module.js";

var PMREMGenerator = ( function () {

	var LOD_MIN = 4;
	var LOD_MAX = 8;
	var SIZE_MAX = Math.pow( 2, LOD_MAX );
	// The standard deviations (radians) associated with the extra mips. These are
	// chosen to approximate a Trowbridge-Reitz distribution function times the
	// geometric shadowing function. These sigma values squared must match the
	// variance #defines in cube_uv_reflection_fragment.glsl.js.
	var EXTRA_LOD_SIGMA = [ 0.125, 0.215, 0.35, 0.446, 0.526, 0.582 ];
	var TOTAL_LODS = LOD_MAX - LOD_MIN + 1 + EXTRA_LOD_SIGMA.length;
	// The maximum length of the blur for loop. Smaller sigmas will use fewer
	// samples and exit early, but not recompile the shader.
	var MAX_SAMPLES = 20;
	var ENCODINGS = {
		[ LinearEncoding ]: 0,
		[ sRGBEncoding ]: 1,
		[ RGBEEncoding ]: 2,
		[ RGBM7Encoding ]: 3,
		[ RGBM16Encoding ]: 4,
		[ RGBDEncoding ]: 5,
		[ GammaEncoding ]: 6
	  };

	var _flatCamera = new OrthographicCamera();
	var _blurMaterial = _getBlurShader( MAX_SAMPLES );

	var { _lodPlanes, _sizeLods, _sigmas } = _createPlanes();
	var _pingPongRenderTarget = null;
	var _renderer = null;

	// Golden Ratio
	var PHI = ( 1 + Math.sqrt( 5 ) ) / 2;
	var INV_PHI = 1 / PHI;
	// Vertices of a dodecahedron (except the opposites, which represent the
	// same axis), used as axis directions evenly spread on a sphere.
	var _axisDirections = [
		new Vector3( 1, 1, 1 ),
		new Vector3( - 1, 1, 1 ),
		new Vector3( 1, 1, - 1 ),
		new Vector3( - 1, 1, - 1 ),
		new Vector3( 0, PHI, INV_PHI ),
		new Vector3( 0, PHI, - INV_PHI ),
		new Vector3( INV_PHI, 0, PHI ),
		new Vector3( - INV_PHI, 0, PHI ),
		new Vector3( PHI, INV_PHI, 0 ),
		new Vector3( - PHI, INV_PHI, 0 ) ];

	var PMREMGenerator = function ( renderer ) {

		_renderer = renderer;

	};

	PMREMGenerator.prototype = {

		constructor: PMREMGenerator,

		/**
		 * Generates a PMREM from a supplied Scene, which can be faster than using an
		 * image if networking bandwidth is low. Optional sigma specifies a blur radius
		 * in radians to be applied to the scene before PMREM generation. Optional near
		 * and far planes ensure the scene is rendered in its entirety (the cubeCamera
		 * is placed at the origin).
		 */
		fromScene: function ( scene, sigma = 0, near = 0.1, far = 100 ) {

			var cubeUVRenderTarget = _allocateTargets();
			_sceneToCubeUV( scene, near, far, cubeUVRenderTarget );
			if ( sigma > 0 ) {

				_blur( cubeUVRenderTarget, 0, 0, sigma );

			}
			_applyPMREM( cubeUVRenderTarget );
			_cleanup();

			return cubeUVRenderTarget;

		},

		/**
		 * Generates a PMREM from an equirectangular texture, which can be either LDR
		 * (RGBFormat) or HDR (RGBEFormat). The ideal input image size is 1k (1024 x 512),
		 * as this matches best with the 256 x 256 cubemap output.
		 */
		fromEquirectangular: function ( equirectangular ) {

			equirectangular.magFilter = NearestFilter;
			equirectangular.minFilter = NearestFilter;
			equirectangular.generateMipmaps = false;

			var cubeUVRenderTarget = _allocateTargets( equirectangular );
			_equirectangularToCubeUV( equirectangular, cubeUVRenderTarget );
			_applyPMREM( cubeUVRenderTarget );
			_cleanup();

			return cubeUVRenderTarget;

		},

	};

	function _createPlanes() {

		var _lodPlanes = [];
		var _sizeLods = [];
		var _sigmas = [];

		var lod = LOD_MAX;
		for ( var i = 0; i < TOTAL_LODS; i ++ ) {

			var sizeLod = Math.pow( 2, lod );
			_sizeLods.push( sizeLod );
			var sigma = 1.0 / sizeLod;
			if ( i > LOD_MAX - LOD_MIN ) {

				sigma = EXTRA_LOD_SIGMA[ i - LOD_MAX + LOD_MIN - 1 ];

			} else if ( i == 0 ) {

				sigma = 0;

			}
			_sigmas.push( sigma );

			var texelSize = 1.0 / ( sizeLod - 1 );
			var min = - texelSize / 2;
			var max = 1 + texelSize / 2;
			var uv1 = [ min, min, max, min, max, max, min, min, max, max, min, max ];

			var cubeFaces = 6;
			var vertices = 6;
			var positionSize = 3;
			var uvSize = 2;
			var faceIndexSize = 1;

			var position = new Float32Array( positionSize * vertices * cubeFaces );
			var uv = new Float32Array( uvSize * vertices * cubeFaces );
			var faceIndex = new Float32Array( faceIndexSize * vertices * cubeFaces );

			for ( var face = 0; face < cubeFaces; face ++ ) {

				var x = ( face % 3 ) * 2 / 3 - 1;
				var y = face > 2 ? 0 : - 1;
				var coordinates = [
					[ x, y, 0 ],
					[ x + 2 / 3, y, 0 ],
					[ x + 2 / 3, y + 1, 0 ],
					[ x, y, 0 ],
					[ x + 2 / 3, y + 1, 0 ],
					[ x, y + 1, 0 ]
				];
				position.set( [].concat( ...coordinates ),
					positionSize * vertices * face );
				uv.set( uv1, uvSize * vertices * face );
				var fill = [ face, face, face, face, face, face ];
				faceIndex.set( fill, faceIndexSize * vertices * face );

			}
			var planes = new BufferGeometry();
			planes.setAttribute(
				'position', new BufferAttribute( position, positionSize ) );
			planes.setAttribute( 'uv', new BufferAttribute( uv, uvSize ) );
			planes.setAttribute(
				'faceIndex', new BufferAttribute( faceIndex, faceIndexSize ) );
			_lodPlanes.push( planes );

			if ( lod > LOD_MIN ) {

				lod --;

			}

		}
		return { _lodPlanes, _sizeLods, _sigmas };

	}

	function _allocateTargets( equirectangular ) {

		var params = {
		  magFilter: NearestFilter,
		  minFilter: NearestFilter,
		  generateMipmaps: false,
		  type: equirectangular ? equirectangular.type : UnsignedByteType,
		  format: equirectangular ? equirectangular.format : RGBEFormat,
		  encoding: equirectangular ? equirectangular.encoding : RGBEEncoding,
		  depthBuffer: false,
		  stencilBuffer: false
		};
		var cubeUVRenderTarget = _createRenderTarget(
			{ ...params, depthBuffer: ( equirectangular ? false : true ) } );
		_pingPongRenderTarget = _createRenderTarget( params );
		return cubeUVRenderTarget;

	}

	function _cleanup() {

		_pingPongRenderTarget.dispose();
		_renderer.setRenderTarget( null );
		var size = _renderer.getSize( new Vector2() );
		_renderer.setViewport( 0, 0, size.x, size.y );

	}

	function _sceneToCubeUV(
		scene, near, far,
		cubeUVRenderTarget ) {

	  var fov = 90;
	  var aspect = 1;
	  var cubeCamera = new PerspectiveCamera( fov, aspect, near, far );
	  var upSign = [ 1, 1, 1, 1, - 1, 1 ];
	  var forwardSign = [ 1, 1, - 1, - 1, - 1, 1 ];

	  var gammaOutput = _renderer.gammaOutput;
	  var toneMapping = _renderer.toneMapping;
	  var toneMappingExposure = _renderer.toneMappingExposure;

	  _renderer.toneMapping = LinearToneMapping;
	  _renderer.toneMappingExposure = 1.0;
	  _renderer.gammaOutput = false;
	  scene.scale.z *= - 1;

	  _renderer.setRenderTarget( cubeUVRenderTarget );
	  for ( var i = 0; i < 6; i ++ ) {

			var col = i % 3;
			if ( col == 0 ) {

		  cubeCamera.up.set( 0, upSign[ i ], 0 );
		  cubeCamera.lookAt( forwardSign[ i ], 0, 0 );

			} else if ( col == 1 ) {

		  cubeCamera.up.set( 0, 0, upSign[ i ] );
		  cubeCamera.lookAt( 0, forwardSign[ i ], 0 );

			} else {

		  cubeCamera.up.set( 0, upSign[ i ], 0 );
		  cubeCamera.lookAt( 0, 0, forwardSign[ i ] );

			}
			_setViewport(
				col * SIZE_MAX, i > 2 ? SIZE_MAX : 0, SIZE_MAX, SIZE_MAX );
			_renderer.render( scene, cubeCamera );

		}

	  _renderer.toneMapping = toneMapping;
	  _renderer.toneMappingExposure = toneMappingExposure;
	  _renderer.gammaOutput = gammaOutput;
	  scene.scale.z *= - 1;

	}

	function _equirectangularToCubeUV(
		equirectangular, cubeUVRenderTarget ) {

	  var scene = new Scene();
	  var equirectMaterial = _getEquirectShader();
	  scene.add( new Mesh( _lodPlanes[ 0 ], equirectMaterial ) );
	  var uniforms = equirectMaterial.uniforms;

	  uniforms[ 'envMap' ].value = equirectangular;
	  uniforms[ 'texelSize' ].value.set(
		  1.0 / equirectangular.image.width, 1.0 / equirectangular.image.height );
	  uniforms[ 'inputEncoding' ].value = ENCODINGS[ equirectangular.encoding ];
	  uniforms[ 'outputEncoding' ].value = ENCODINGS[ equirectangular.encoding ];

	  _renderer.setRenderTarget( cubeUVRenderTarget );
	  _setViewport( 0, 0, 3 * SIZE_MAX, 2 * SIZE_MAX );
	  _renderer.render( scene, _flatCamera );

	}

	function _createRenderTarget( params ) {

	  var cubeUVRenderTarget =
		  new WebGLRenderTarget( 3 * SIZE_MAX, 3 * SIZE_MAX, params );
	  cubeUVRenderTarget.texture.mapping = CubeUVReflectionMapping;
	  cubeUVRenderTarget.texture.name = 'PMREM.cubeUv';
	  return cubeUVRenderTarget;

	}

	function _setViewport( x, y, width, height ) {

		var dpr = _renderer.getPixelRatio();
		_renderer.setViewport( x / dpr, y / dpr, width / dpr, height / dpr );

	}

	function _applyPMREM( cubeUVRenderTarget ) {

		var autoClear = _renderer.autoClear;
		_renderer.autoClear = false;

	  	for ( var i = 1; i < TOTAL_LODS; i ++ ) {

			var sigma = Math.sqrt(
				_sigmas[ i ] * _sigmas[ i ] -
			_sigmas[ i - 1 ] * _sigmas[ i - 1 ] );
			var poleAxis =
			_axisDirections[ ( i - 1 ) % _axisDirections.length ];
			_blur( cubeUVRenderTarget, i - 1, i, sigma, poleAxis );

		}

		_renderer.autoClear = autoClear;

	}

	/**
   * This is a two-pass Gaussian blur for a cubemap. Normally this is done
   * vertically and horizontally, but this breaks down on a cube. Here we apply
   * the blur latitudinally (around the poles), and then longitudinally (towards
   * the poles) to approximate the orthogonally-separable blur. It is least
   * accurate at the poles, but still does a decent job.
   */
	function _blur(
		cubeUVRenderTarget, lodIn, lodOut,
		sigma, poleAxis ) {

		_halfBlur(
	  cubeUVRenderTarget,
	  _pingPongRenderTarget,
	  lodIn,
	  lodOut,
	  sigma,
	  'latitudinal',
	  poleAxis );

		_halfBlur(
	  _pingPongRenderTarget,
	  cubeUVRenderTarget,
	  lodOut,
	  lodOut,
	  sigma,
	  'longitudinal',
	  poleAxis );

	}

	function _halfBlur(
		targetIn, targetOut, lodIn,
		lodOut, sigmaRadians, direction,
		poleAxis ) {

		if ( direction !== 'latitudinal' && direction !== 'longitudinal' ) {

			console.error(
				'blur direction must be either latitudinal or longitudinal!' );

		}

		// Number of standard deviations at which to cut off the discrete approximation.
		var STANDARD_DEVIATIONS = 3;

		var blurScene = new Scene();
		blurScene.add( new Mesh( _lodPlanes[ lodOut ], _blurMaterial ) );
		var blurUniforms = _blurMaterial.uniforms;

		var pixels = _sizeLods[ lodIn ] - 1;
		var radiansPerPixel = isFinite( sigmaRadians ) ? Math.PI / ( 2 * pixels ) : 2 * Math.PI / ( 2 * MAX_SAMPLES - 1 );
		var sigmaPixels = sigmaRadians / radiansPerPixel;
		var samples = isFinite( sigmaRadians ) ? 1 + Math.floor( STANDARD_DEVIATIONS * sigmaPixels ) : MAX_SAMPLES;

		if ( samples > MAX_SAMPLES ) {

			console.warn( `sigmaRadians, ${
				sigmaRadians}, is too large and will clip, as it requested ${
				samples} samples when the maximum is set to ${MAX_SAMPLES}` );

		}

		var weights = [];
		var sum = 0;
		for ( var i = 0; i < MAX_SAMPLES; ++ i ) {

			var x = i / sigmaPixels;
			var weight = Math.exp( - x * x / 2 );
			weights.push( weight );
			if ( i == 0 ) {

	  			 sum += weight;

			} else if ( i < samples ) {

	  			sum += 2 * weight;

			}

		}
		weights = weights.map( w => w / sum );

		blurUniforms[ 'envMap' ].value = targetIn.texture;
		blurUniforms[ 'samples' ].value = samples;
		blurUniforms[ 'weights' ].value = weights;
		blurUniforms[ 'latitudinal' ].value = direction === 'latitudinal';
		if ( poleAxis ) {

			blurUniforms[ 'poleAxis' ].value = poleAxis;

		}
		blurUniforms[ 'dTheta' ].value = radiansPerPixel;
		blurUniforms[ 'mipInt' ].value = LOD_MAX - lodIn;
		blurUniforms[ 'inputEncoding' ].value = ENCODINGS[ targetIn.texture.encoding ];
		blurUniforms[ 'outputEncoding' ].value = ENCODINGS[ targetIn.texture.encoding ];

		var outputSize = _sizeLods[ lodOut ];
		var x = 3 * Math.max( 0, SIZE_MAX - 2 * outputSize );
		var y = ( lodOut === 0 ? 0 : 2 * SIZE_MAX ) +
	  2 * outputSize *
		  ( lodOut > LOD_MAX - LOD_MIN ? lodOut - LOD_MAX + LOD_MIN : 0 );

		_renderer.setRenderTarget( targetOut );
		_setViewport( x, y, 3 * outputSize, 2 * outputSize );
		_renderer.render( blurScene, _flatCamera );

	}

	function _getBlurShader( maxSamples ) {

		var weights = new Float32Array( maxSamples );
		var poleAxis = new Vector3( 0, 1, 0 );
		var shaderMaterial = new RawShaderMaterial( {

			defines: { 'n': maxSamples },

			uniforms: {
				'envMap': { value: null },
				'samples': { value: 1 },
				'weights': { value: weights },
				'latitudinal': { value: false },
				'dTheta': { value: 0 },
				'mipInt': { value: 0 },
				'poleAxis': { value: poleAxis },
				'inputEncoding': { value: ENCODINGS[ LinearEncoding ] },
				'outputEncoding': { value: ENCODINGS[ LinearEncoding ] }
			},

			vertexShader: _getCommonVertexShader(),

			fragmentShader: `
precision mediump float;
precision mediump int;
varying vec3 vOutputDirection;
uniform sampler2D envMap;
uniform int samples;
uniform float weights[n];
uniform bool latitudinal;
uniform float dTheta;
uniform float mipInt;
uniform vec3 poleAxis;

${_getEncodings()}

#define ENVMAP_TYPE_CUBE_UV
#include <cube_uv_reflection_fragment>

void main() {
	gl_FragColor = vec4(0.0);
    for (int i = 0; i < n; i++) {
      if (i >= samples)
        break;
      for (int dir = -1; dir < 2; dir += 2) {
        if (i == 0 && dir == 1)
          continue;
        vec3 axis = latitudinal ? poleAxis : cross(poleAxis, vOutputDirection);
        if (all(equal(axis, vec3(0.0))))
          axis = cross(vec3(0.0, 1.0, 0.0), vOutputDirection);
        axis = normalize(axis);
        float theta = dTheta * float(dir * i);
        float cosTheta = cos(theta);
        // Rodrigues' axis-angle rotation
        vec3 sampleDirection = vOutputDirection * cosTheta 
            + cross(axis, vOutputDirection) * sin(theta) 
            + axis * dot(axis, vOutputDirection) * (1.0 - cosTheta);
        gl_FragColor.rgb +=
            weights[i] * bilinearCubeUV(envMap, sampleDirection, mipInt);
      }
    }
  	gl_FragColor = linearToOutputTexel(gl_FragColor);
}
     		`,

			blending: NoBlending,
			depthTest: false,
	   		depthWrite: false

		} );

		shaderMaterial.type = 'SphericalGaussianBlur';

		return shaderMaterial;

	}

	function _getEquirectShader() {

		var texelSize = new Vector2( 1, 1 );
		var shaderMaterial = new RawShaderMaterial( {

			uniforms: {
				'envMap': { value: null },
				'texelSize': { value: texelSize },
				'inputEncoding': { value: ENCODINGS[ LinearEncoding ] },
				'outputEncoding': { value: ENCODINGS[ LinearEncoding ] }
			},

			vertexShader: _getCommonVertexShader(),

			fragmentShader: `
precision mediump float;
precision mediump int;
varying vec3 vOutputDirection;
uniform sampler2D envMap;
uniform vec2 texelSize;

${_getEncodings()}

#define RECIPROCAL_PI 0.31830988618
#define RECIPROCAL_PI2 0.15915494

void main() {
	gl_FragColor = vec4(0.0);
	vec3 outputDirection = normalize(vOutputDirection);
	vec2 uv;
	uv.y = asin(clamp(outputDirection.y, -1.0, 1.0)) * RECIPROCAL_PI + 0.5;
	uv.x = atan(outputDirection.z, outputDirection.x) * RECIPROCAL_PI2 + 0.5;
	vec2 f = fract(uv / texelSize - 0.5);
	uv -= f * texelSize;
	vec3 tl = envMapTexelToLinear(texture2D(envMap, uv)).rgb;
	uv.x += texelSize.x;
	vec3 tr = envMapTexelToLinear(texture2D(envMap, uv)).rgb;
	uv.y += texelSize.y;
	vec3 br = envMapTexelToLinear(texture2D(envMap, uv)).rgb;
	uv.x -= texelSize.x;
	vec3 bl = envMapTexelToLinear(texture2D(envMap, uv)).rgb;
	vec3 tm = mix(tl, tr, f.x);
	vec3 bm = mix(bl, br, f.x);
	gl_FragColor.rgb = mix(tm, bm, f.y);
  	gl_FragColor = linearToOutputTexel(gl_FragColor);
}
     		`,

			blending: NoBlending,
			depthTest: false,
	   		depthWrite: false

		} );

		shaderMaterial.type = 'EquirectangularToCubeUV';

		return shaderMaterial;

	}

	function _getCommonVertexShader() {

		return `
precision mediump float;
precision mediump int;
attribute vec3 position;
attribute vec2 uv;
attribute float faceIndex;
varying vec3 vOutputDirection;
vec3 getDirection(vec2 uv, float face) {
	uv = 2.0 * uv - 1.0;
	vec3 direction = vec3(uv, 1.0);
	if (face == 0.0) {
		direction = direction.zyx;
		direction.z *= -1.0;
	} else if (face == 1.0) {
		direction = direction.xzy;
		direction.z *= -1.0;
	} else if (face == 3.0) {
		direction = direction.zyx;
		direction.x *= -1.0;
	} else if (face == 4.0) {
		direction = direction.xzy;
		direction.y *= -1.0;
	} else if (face == 5.0) {
		direction.xz *= -1.0;
	}
	return direction;
}
void main() {
	vOutputDirection = getDirection(uv, faceIndex);
	gl_Position = vec4( position, 1.0 );
}
		`;

	}

	function _getEncodings() {

		return `
uniform int inputEncoding;
uniform int outputEncoding;

#include <encodings_pars_fragment>

vec4 inputTexelToLinear(vec4 value){
    if(inputEncoding == 0){
        return value;
    }else if(inputEncoding == 1){
        return sRGBToLinear(value);
    }else if(inputEncoding == 2){
        return RGBEToLinear(value);
    }else if(inputEncoding == 3){
        return RGBMToLinear(value, 7.0);
    }else if(inputEncoding == 4){
        return RGBMToLinear(value, 16.0);
    }else if(inputEncoding == 5){
        return RGBDToLinear(value, 256.0);
    }else{
        return GammaToLinear(value, 2.2);
    }
}

vec4 linearToOutputTexel(vec4 value){
    if(outputEncoding == 0){
        return value;
    }else if(outputEncoding == 1){
        return LinearTosRGB(value);
    }else if(outputEncoding == 2){
        return LinearToRGBE(value);
    }else if(outputEncoding == 3){
        return LinearToRGBM(value, 7.0);
    }else if(outputEncoding == 4){
        return LinearToRGBM(value, 16.0);
    }else if(outputEncoding == 5){
        return LinearToRGBD(value, 256.0);
    }else{
        return LinearToGamma(value, 2.2);
    }
}

vec4 envMapTexelToLinear(vec4 color) {
  return inputTexelToLinear(color);
}
		`;

	}

	return PMREMGenerator;

} )();

export { PMREMGenerator };