three-gpu-pathtracer 0.0.2 → 0.0.3

package/src/materials/PhysicalPathTracingMaterial.js

@@ -1,22 +1,23 @@
-import { Matrix4, Matrix3, Color } from 'three';
+import { Matrix4, Matrix3, Color, Vector2 } from 'three';
 import { MaterialBase } from './MaterialBase.js';
 import {
 	MeshBVHUniformStruct, FloatVertexAttributeTexture, UIntVertexAttributeTexture,
 	shaderStructs, shaderIntersectFunction,
 } from 'three-mesh-bvh';
 import { shaderMaterialStructs } from '../shader/shaderStructs.js';
-import { MaterialStructArrayUniform } from '../uniforms/MaterialStructArrayUniform.js';
+import { MaterialsTexture } from '../uniforms/MaterialsTexture.js';
 import { RenderTarget2DArray } from '../uniforms/RenderTarget2DArray.js';
 import { shaderMaterialSampling } from '../shader/shaderMaterialSampling.js';
+import { shaderEnvMapSampling } from '../shader/shaderEnvMapSampling.js';
 import { shaderUtils } from '../shader/shaderUtils.js';
 import { PhysicalCameraUniform } from '../uniforms/PhysicalCameraUniform.js';
+import { EquirectHdrInfoUniform } from '../uniforms/EquirectHdrInfoUniform.js';
 
 export class PhysicalPathTracingMaterial extends MaterialBase {
 
-	// three.js relies on this field to add env map functions and defines
-	get envMap() {
+	onBeforeRender() {
 
-		return this.environmentMap;
+		this.setDefine( 'FEATURE_DOF', this.physicalCamera.bokehSize === 0 ? 0 : 1 );
 
 	}
 
@@ -28,13 +29,15 @@ export class PhysicalPathTracingMaterial extends MaterialBase {
 			depthWrite: false,
 
 			defines: {
-				DOF_SUPPORT: 1,
+				FEATURE_MIS: 1,
+				FEATURE_DOF: 1,
+				FEATURE_GRADIENT_BG: 0,
 				TRANSPARENT_TRAVERSALS: 5,
-				MATERIAL_LENGTH: 0,
-				GRADIENT_BG: 0,
 			},
 
 			uniforms: {
+				resolution: { value: new Vector2() },
+
 				bounces: { value: 3 },
 				physicalCamera: { value: new PhysicalCameraUniform() },
 
@@ -43,23 +46,22 @@ export class PhysicalPathTracingMaterial extends MaterialBase {
 				tangentAttribute: { value: new FloatVertexAttributeTexture() },
 				uvAttribute: { value: new FloatVertexAttributeTexture() },
 				materialIndexAttribute: { value: new UIntVertexAttributeTexture() },
-				materials: { value: new MaterialStructArrayUniform() },
+				materials: { value: new MaterialsTexture() },
 				textures: { value: new RenderTarget2DArray().texture },
 				cameraWorldMatrix: { value: new Matrix4() },
 				invProjectionMatrix: { value: new Matrix4() },
-				environmentBlur: { value: 0.2 },
+				backgroundBlur: { value: 0.0 },
 				environmentIntensity: { value: 2.0 },
-				environmentMap: { value: null },
 				environmentRotation: { value: new Matrix3() },
+				envMapInfo: { value: new EquirectHdrInfoUniform() },
+
 				seed: { value: 0 },
 				opacity: { value: 1 },
 				filterGlossyFactor: { value: 0.0 },
 
-				gradientTop: { value: new Color( 0xbfd8ff ) },
-				gradientBottom: { value: new Color( 0xffffff ) },
-
 				bgGradientTop: { value: new Color( 0x111111 ) },
 				bgGradientBottom: { value: new Color( 0x000000 ) },
+				backgroundAlpha: { value: 1.0 },
 			},
 
 			vertexShader: /* glsl */`
@@ -78,14 +80,13 @@ export class PhysicalPathTracingMaterial extends MaterialBase {
 			`,
 
 			fragmentShader: /* glsl */`
-				#define RAY_OFFSET 1e-5
+				#define RAY_OFFSET 1e-4
 
 				precision highp isampler2D;
 				precision highp usampler2D;
 				precision highp sampler2DArray;
 				vec4 envMapTexelToLinear( vec4 a ) { return a; }
 				#include <common>
-				#include <cube_uv_reflection_fragment>
 
 				${ shaderStructs }
 				${ shaderIntersectFunction }
@@ -93,35 +94,27 @@ export class PhysicalPathTracingMaterial extends MaterialBase {
 
 				${ shaderUtils }
 				${ shaderMaterialSampling }
+				${ shaderEnvMapSampling }
 
-				#ifdef USE_ENVMAP
-
-				uniform float environmentBlur;
-				uniform sampler2D environmentMap;
 				uniform mat3 environmentRotation;
+				uniform float backgroundBlur;
+				uniform float backgroundAlpha;
 
-				#else
-
-				uniform vec3 gradientTop;
-				uniform vec3 gradientBottom;
-
-				#endif
-
-				#if GRADIENT_BG
+				#if FEATURE_GRADIENT_BG
 
 				uniform vec3 bgGradientTop;
 				uniform vec3 bgGradientBottom;
 
 				#endif
 
-				#if DOF_SUPPORT
+				#if FEATURE_DOF
 
 				uniform PhysicalCamera physicalCamera;
 
 				#endif
 
+				uniform vec2 resolution;
 				uniform int bounces;
-
 				uniform mat4 cameraWorldMatrix;
 				uniform mat4 invProjectionMatrix;
 				uniform sampler2D normalAttribute;
@@ -133,10 +126,149 @@ export class PhysicalPathTracingMaterial extends MaterialBase {
 				uniform float filterGlossyFactor;
 				uniform int seed;
 				uniform float opacity;
-				uniform Material materials[ MATERIAL_LENGTH ];
+				uniform sampler2D materials;
+
+				uniform EquirectHdrInfo envMapInfo;
+
 				uniform sampler2DArray textures;
 				varying vec2 vUv;
 
+				vec3 sampleBackground( vec3 direction ) {
+
+					#if FEATURE_GRADIENT_BG
+
+					direction = normalize( direction + randDirection() * 0.05 );
+
+					float value = ( direction.y + 1.0 ) / 2.0;
+					value = pow( value, 2.0 );
+
+					return mix( bgGradientBottom, bgGradientTop, value );
+
+					#else
+
+					vec3 sampleDir = normalize( direction + getHemisphereSample( direction, rand2() ) * 0.5 * backgroundBlur );
+					return environmentIntensity * sampleEquirectEnvMapColor( sampleDir, envMapInfo.map );
+
+					#endif
+
+				}
+
+				// step through multiple surface hits and accumulate color attenuation based on transmissive surfaces
+				bool attenuateHit( BVH bvh, vec3 rayOrigin, vec3 rayDirection, int traversals, out vec3 color ) {
+
+					// hit results
+					uvec4 faceIndices = uvec4( 0u );
+					vec3 faceNormal = vec3( 0.0, 0.0, 1.0 );
+					vec3 barycoord = vec3( 0.0 );
+					float side = 1.0;
+					float dist = 0.0;
+
+					color = vec3( 1.0 );
+
+					for ( int i = 0; i < traversals; i ++ ) {
+
+						if ( bvhIntersectFirstHit( bvh, rayOrigin, rayDirection, faceIndices, faceNormal, barycoord, side, dist ) ) {
+
+							// TODO: attenuate the contribution based on the PDF of the resulting ray including refraction values
+							// Should be able to work using the material BSDF functions which will take into account specularity, etc.
+							// TODO: should we account for emissive surfaces here?
+
+							vec2 uv = textureSampleBarycoord( uvAttribute, barycoord, faceIndices.xyz ).xy;
+							uint materialIndex = uTexelFetch1D( materialIndexAttribute, faceIndices.x ).r;
+							Material material = readMaterialInfo( materials, materialIndex );
+
+							// Opacity Test
+
+							// albedo
+							vec4 albedo = vec4( material.color, material.opacity );
+							if ( material.map != - 1 ) {
+
+								albedo *= texture2D( textures, vec3( uv, material.map ) );
+
+							}
+
+							// transmission
+							float transmission = material.transmission;
+							if ( material.transmissionMap != - 1 ) {
+
+								transmission *= texture2D( textures, vec3( uv, material.transmissionMap ) ).r;
+
+							}
+
+							// metalness
+							float metalness = material.metalness;
+							if ( material.metalnessMap != - 1 ) {
+
+								metalness *= texture2D( textures, vec3( uv, material.metalnessMap ) ).b;
+
+							}
+
+							float alphaTest = material.alphaTest;
+							bool useAlphaTest = alphaTest != 0.0;
+							float transmissionFactor = ( 1.0 - metalness ) * transmission;
+							if (
+								transmissionFactor < rand() && ! (
+									// material sidedness
+									material.side != 0.0 && side == material.side
+
+									// alpha test
+									|| useAlphaTest && albedo.a < alphaTest
+
+									// opacity
+									|| ! useAlphaTest && albedo.a < rand()
+								)
+							) {
+
+								return true;
+
+							}
+
+							// only attenuate on the way in
+							bool isBelowSurface = dot( rayDirection, faceNormal ) < 0.0;
+							if ( isBelowSurface ) {
+
+								color *= albedo.rgb;
+
+							}
+
+							// adjust the ray to the new surface
+							vec3 point = rayOrigin + rayDirection * dist;
+							vec3 absPoint = abs( point );
+							float maxPoint = max( absPoint.x, max( absPoint.y, absPoint.z ) );
+							rayOrigin = point + faceNormal * ( maxPoint + 1.0 ) * ( isBelowSurface ? - RAY_OFFSET : RAY_OFFSET );
+
+						} else {
+
+							return false;
+
+						}
+
+					}
+
+					return true;
+
+				}
+
+				// returns whether the ray hit anything, not just the first surface. Could be optimized to not check the full hierarchy.
+				bool anyHit( BVH bvh, vec3 rayOrigin, vec3 rayDirection ) {
+
+					uvec4 faceIndices = uvec4( 0u );
+					vec3 faceNormal = vec3( 0.0, 0.0, 1.0 );
+					vec3 barycoord = vec3( 0.0 );
+					float side = 1.0;
+					float dist = 0.0;
+					return bvhIntersectFirstHit( bvh, rayOrigin, rayDirection, faceIndices, faceNormal, barycoord, side, dist );
+
+				}
+
+				// tentFilter from Peter Shirley's 'Realistic Ray Tracing (2nd Edition)' book, pg. 60
+				// erichlof/THREE.js-PathTracing-Renderer/
+				float tentFilter( float x ) {
+
+					return x < 0.5 ? sqrt( 2.0 * x ) - 1.0 : 1.0 - sqrt( 2.0 - ( 2.0 * x ) );
+
+				}
+
 				void main() {
 
 					rng_initialize( gl_FragCoord.xy, seed );
@@ -146,34 +278,58 @@ export class PhysicalPathTracingMaterial extends MaterialBase {
 					vec3 rayOrigin, rayDirection;
 					ndcToCameraRay( ndc, cameraWorldMatrix, invProjectionMatrix, rayOrigin, rayDirection );
 
-					#if DOF_SUPPORT
+					// Jitter the camera ray by finding a new subpixel point to point to from the camera origin
+					// This is better than just jittering the camera position since it actually results in divergent
+					// rays providing better coverage for the pixel
+					{
 
-					// depth of field
-					vec3 focalPoint = rayOrigin + normalize( rayDirection ) * physicalCamera.focusDistance;
+						// TODO: the complexity here could be improved
+						vec3 cameraOrigin = ( cameraWorldMatrix * vec4( 0.0, 0.0, 0.0, 1.0 ) ).xyz;
+						vec3 ss00, ss01, ss10, temp;
+						ndcToCameraRay( vec2( - 1.0, - 1.0 ), cameraWorldMatrix, invProjectionMatrix, ss00, temp );
+						ndcToCameraRay( vec2( - 1.0, 1.0 ), cameraWorldMatrix, invProjectionMatrix, ss01, temp );
+						ndcToCameraRay( vec2( 1.0, - 1.0 ), cameraWorldMatrix, invProjectionMatrix, ss10, temp );
 
-					// get the aperture sample
-					vec2 apertureSample = sampleAperture( physicalCamera.apertureBlades ) * physicalCamera.bokehSize * 0.5 * 1e-3;
+						vec3 ssdX = ( ss10 - ss00 ) / resolution.x;
+						vec3 ssdY = ( ss01 - ss00 ) / resolution.y;
+						rayOrigin += tentFilter( rand() ) * ssdX + tentFilter( rand() ) * ssdY;
+						rayDirection = normalize( rayOrigin - cameraOrigin );
+
+					}
+
+					#if FEATURE_DOF
+					{
+
+						// depth of field
+						vec3 focalPoint = rayOrigin + normalize( rayDirection ) * physicalCamera.focusDistance;
+
+						// get the aperture sample
+						vec2 apertureSample = sampleAperture( physicalCamera.apertureBlades ) * physicalCamera.bokehSize * 0.5 * 1e-3;
+
+						// rotate the aperture shape
+						float ac = cos( physicalCamera.apertureRotation );
+						float as = sin( physicalCamera.apertureRotation );
+						apertureSample = vec2(
+							apertureSample.x * ac - apertureSample.y * as,
+							apertureSample.x * as + apertureSample.y * ac
+						);
+						apertureSample.x *= saturate( physicalCamera.anamorphicRatio );
+						apertureSample.y *= saturate( 1.0 / physicalCamera.anamorphicRatio );
 
-					// rotate the aperture shape
-					float ac = cos( physicalCamera.apertureRotation );
-					float as = sin( physicalCamera.apertureRotation );
-					apertureSample = vec2(
-						apertureSample.x * ac - apertureSample.y * as,
-						apertureSample.x * as + apertureSample.y * ac
-					);
-					apertureSample.x *= saturate( physicalCamera.anamorphicRatio );
-					apertureSample.y *= saturate( 1.0 / physicalCamera.anamorphicRatio );
+						// create the new ray
+						rayOrigin += ( cameraWorldMatrix * vec4( apertureSample, 0.0, 0.0 ) ).xyz;
+						rayDirection = focalPoint - rayOrigin;
 
+					}
 					#endif
+					rayDirection = normalize( rayDirection );
 
-					// create the new ray
-					rayOrigin += ( cameraWorldMatrix * vec4( apertureSample, 0.0, 0.0 ) ).xyz;
-					rayDirection = focalPoint - rayOrigin;
+					// inverse environment rotation
+					mat3 invEnvironmentRotation = inverse( environmentRotation );
 
-					// Lambertian render
+					// final color
 					gl_FragColor = vec4( 0.0 );
-
-					vec3 throughputColor = vec3( 1.0 );
+					gl_FragColor.a = 1.0;
 
 					// hit results
 					uvec4 faceIndices = uvec4( 0u );
@@ -181,50 +337,60 @@ export class PhysicalPathTracingMaterial extends MaterialBase {
 					vec3 barycoord = vec3( 0.0 );
 					float side = 1.0;
 					float dist = 0.0;
+
+					// path tracing state
 					float accumulatedRoughness = 0.0;
-					int i;
+					bool transmissiveRay = true;
 					int transparentTraversals = TRANSPARENT_TRAVERSALS;
+					vec3 throughputColor = vec3( 1.0 );
+					SampleRec sampleRec;
+					int i;
+
 					for ( i = 0; i < bounces; i ++ ) {
 
 						if ( ! bvhIntersectFirstHit( bvh, rayOrigin, rayDirection, faceIndices, faceNormal, barycoord, side, dist ) ) {
 
-							#if GRADIENT_BG
+							if ( i == 0 || transmissiveRay ) {
 
-							if ( i == 0 ) {
+								gl_FragColor.rgb += sampleBackground( environmentRotation * rayDirection ) * throughputColor;
+								gl_FragColor.a = backgroundAlpha;
 
-								rayDirection = normalize( rayDirection + randDirection() * 0.05 );
-								float value = ( rayDirection.y + 1.0 ) / 2.0;
+							} else {
 
-								value = pow( value, 2.0 );
+								#if FEATURE_MIS
 
-								gl_FragColor = vec4( mix( bgGradientBottom, bgGradientTop, value ), 1.0 );
-								break;
+								// get the PDF of the hit envmap point
+								vec3 envColor;
+								float envPdf = envMapSample( environmentRotation * rayDirection, envMapInfo, envColor );
 
-							}
+								// and weight the contribution
+								float misWeight = misHeuristic( sampleRec.pdf, envPdf );
+								gl_FragColor.rgb += environmentIntensity * envColor * throughputColor * misWeight;
 
-							#endif
+								#else
 
-							#ifdef USE_ENVMAP
+								gl_FragColor.rgb +=
+									environmentIntensity *
+									sampleEquirectEnvMapColor( environmentRotation * rayDirection, envMapInfo.map ) *
+									throughputColor;
 
-                            vec3 skyColor = textureCubeUV( environmentMap, environmentRotation * rayDirection, environmentBlur ).rgb;
+								#endif
 
-							#else
+							}
+							break;
 
-							rayDirection = normalize( rayDirection );
-							float value = ( rayDirection.y + 1.0 ) / 2.0;
-							vec3 skyColor = mix( gradientBottom, gradientTop, value );
+						}
 
-							#endif
+						uint materialIndex = uTexelFetch1D( materialIndexAttribute, faceIndices.x ).r;
+						Material material = readMaterialInfo( materials, materialIndex );
 
-							gl_FragColor += vec4( skyColor * throughputColor * environmentIntensity, 1.0 );
+						if ( material.matte && i == 0 ) {
 
+							gl_FragColor = vec4( 0.0 );
 							break;
 
 						}
 
-						uint materialIndex = uTexelFetch1D( materialIndexAttribute, faceIndices.x ).r;
-						Material material = materials[ materialIndex ];
-
 						vec2 uv = textureSampleBarycoord( uvAttribute, barycoord, faceIndices.xyz ).xy;
 
 						// albedo
@@ -256,7 +422,9 @@ export class PhysicalPathTracingMaterial extends MaterialBase {
 						) {
 
 							vec3 point = rayOrigin + rayDirection * dist;
-							rayOrigin += rayDirection * dist - faceNormal * RAY_OFFSET;
+							vec3 absPoint = abs( point );
+							float maxPoint = max( absPoint.x, max( absPoint.y, absPoint.z ) );
+							rayOrigin = point - ( maxPoint + 1.0 ) * faceNormal * RAY_OFFSET;
 
 							// only allow a limited number of transparency discards otherwise we could
 							// crash the context with too long a loop.
@@ -361,7 +529,7 @@ export class PhysicalPathTracingMaterial extends MaterialBase {
 						mat3 invBasis = inverse( normalBasis );
 
 						vec3 outgoing = - normalize( invBasis * rayDirection );
-						SampleRec sampleRec = bsdfSample( outgoing, surfaceRec );
+						sampleRec = bsdfSample( outgoing, surfaceRec );
 
 						// adjust the hit point by the surface normal by a factor of some offset and the
 						// maximum component-wise value of the current point to accommodate floating point
@@ -374,6 +542,49 @@ export class PhysicalPathTracingMaterial extends MaterialBase {
 						bool isBelowSurface = dot( rayDirection, faceNormal ) < 0.0;
 						rayOrigin = point + faceNormal * ( maxPoint + 1.0 ) * ( isBelowSurface ? - RAY_OFFSET : RAY_OFFSET );
 
+						// direct env map sampling
+						#if FEATURE_MIS
+						{
+
+							// find a sample in the environment map to include in the contribution
+							vec3 envColor, envDirection;
+							float envPdf = randomEnvMapSample( envMapInfo, envColor, envDirection );
+							envDirection = invEnvironmentRotation * envDirection;
+
+							// this env sampling is not set up for transmissive sampling and yields overly bright
+							// results so we ignore the sample in this case.
+							// TODO: this should be improved but how? The env samples could traverse a few layers?
+							bool isSampleBelowSurface = dot( faceNormal, envDirection ) < 0.0;
+							if ( isSampleBelowSurface ) {
+
+								envPdf = 0.0;
+
+							}
+
+							// check if a ray could even reach the surface
+							vec3 attenuatedColor;
+							if (
+								envPdf > 0.0 &&
+								isDirectionValid( envDirection, normal, faceNormal ) &&
+								! attenuateHit( bvh, rayOrigin, envDirection, bounces - i, attenuatedColor )
+							) {
+
+								// get the material pdf
+								vec3 sampleColor;
+								float envMaterialPdf = bsdfResult( outgoing, normalize( invBasis * envDirection ), surfaceRec, sampleColor );
+								if ( envMaterialPdf > 0.0 ) {
+
+									// weight the direct light contribution
+									float misWeight = misHeuristic( envPdf, envMaterialPdf );
+									gl_FragColor.rgb += attenuatedColor * environmentIntensity * envColor * throughputColor * sampleColor * misWeight / envPdf;
+
+								}
+
+							}
+
+						}
+						#endif
+
 						// accumulate a roughness value to offset diffuse, specular, diffuse rays that have high contribution
 						// to a single pixel resulting in fireflies
 						if ( ! isBelowSurface ) {
@@ -381,6 +592,7 @@ export class PhysicalPathTracingMaterial extends MaterialBase {
 							// determine if this is a rough normal or not by checking how far off straight up it is
 							vec3 halfVector = normalize( outgoing + sampleRec.direction );
 							accumulatedRoughness += sin( acos( halfVector.z ) );
+							transmissiveRay = false;
 
 						}
 
@@ -405,7 +617,7 @@ export class PhysicalPathTracingMaterial extends MaterialBase {
 
 					}
 
-					gl_FragColor.a = opacity;
+					gl_FragColor.a *= opacity;
 
 				}
 

package/src/shader/shaderEnvMapSampling.js

@@ -0,0 +1,67 @@
+export const shaderEnvMapSampling = /* glsl */`
+
+vec3 sampleEquirectEnvMapColor( vec3 direction, sampler2D map ) {
+
+	return texture2D( map, equirectDirectionToUv( direction ) ).rgb;
+
+}
+
+float envMapDirectionPdf( vec3 direction ) {
+
+	vec2 uv = equirectDirectionToUv( direction );
+	float theta = uv.y * PI;
+	float sinTheta = sin( theta );
+	if ( sinTheta == 0.0 ) {
+
+		return 0.0;
+
+	}
+
+	return 1.0 / ( 2.0 * PI * PI * sinTheta );
+
+}
+
+float envMapSample( vec3 direction, EquirectHdrInfo info, out vec3 color ) {
+
+	vec2 uv = equirectDirectionToUv( direction );
+	color = texture2D( info.map, uv ).rgb;
+
+	float totalSum = texture2D( info.totalSum, vec2( 0.0 ) ).r;
+	float lum = colorToLuminance( color );
+	ivec2 resolution = textureSize( info.map, 0 );
+	float pdf = lum / totalSum;
+
+	return float( resolution.x * resolution.y ) * pdf * envMapDirectionPdf( direction );
+
+}
+
+float randomEnvMapSample( EquirectHdrInfo info, out vec3 color, out vec3 direction ) {
+
+	// sample env map cdf
+	vec2 r = rand2();
+	float v = texture2D( info.marginalWeights, vec2( r.x, 0.0 ) ).x;
+	float u = texture2D( info.conditionalWeights, vec2( r.y, v ) ).x;
+	vec2 uv = vec2( u, v );
+
+	vec3 derivedDirection = equirectUvToDirection( uv );
+	direction = derivedDirection;
+	color = texture2D( info.map, uv ).rgb;
+
+	float totalSum = texture2D( info.totalSum, vec2( 0.0 ) ).r;
+	float lum = colorToLuminance( color );
+	ivec2 resolution = textureSize( info.map, 0 );
+	float pdf = lum / totalSum;
+
+	return float( resolution.x * resolution.y ) * pdf * envMapDirectionPdf( direction );
+
+}
+
+float misHeuristic( float a, float b ) {
+
+	float aa = a * a;
+	float bb = a * b;
+	return aa / ( bb + aa );
+
+}
+
+`;

package/src/shader/shaderMaterialSampling.js

@@ -278,6 +278,13 @@ vec3 bsdfColor( vec3 wo, vec3 wi, SurfaceRec surf ) {
 
 }
 
+float bsdfResult( vec3 wo, vec3 wi, SurfaceRec surf, out vec3 color ) {
+
+	color = bsdfColor( wo, wi, surf );
+	return bsdfPdf( wo, wi, surf );
+
+}
+
 SampleRec bsdfSample( vec3 wo, SurfaceRec surf ) {
 
 	float ior = surf.ior;