From 80e00cae4e938050c1ad1d69e8a5dd59846ea95b Mon Sep 17 00:00:00 2001
From: NeilBickford-NV <nbickford@nvidia.com>
Date: Tue, 23 May 2023 14:45:28 -0700
Subject: [PATCH] ray_tracing_gltf: Match samplingHemisphere to Ray Tracing
 Gems implementation and avoid generating directions perpendicular to +z;
 adjust comments on Lambertian PDF and match docs.

---
 ray_tracing_gltf/README.md               | 18 ++++++++++++------
 ray_tracing_gltf/shaders/pathtrace.rchit |  7 +++----
 ray_tracing_gltf/shaders/sampling.glsl   |  9 +++++----
 3 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/ray_tracing_gltf/README.md b/ray_tracing_gltf/README.md
index 7c51c68..f5efc6b 100644
--- a/ray_tracing_gltf/README.md
+++ b/ray_tracing_gltf/README.md
@@ -452,8 +452,8 @@ struct hitPayload
 
 This modification will recursively trace until the `depth`hits 10 (hardcoded) or hit an emissive element (light).
 
-The only information that we will keep from the shader, is the calculation of the hit state: position, normal. So
-all code from `// Vector toward the light` to the end can be remove and be replaced by the following.
+The only information that we will keep from the shader, is the calculation of the hit state: the position and normal. So
+all code from `// Vector toward the light` to the end can be removed and be replaced by the following.
 
 ~~~~C
   // https://en.wikipedia.org/wiki/Path_tracing
@@ -467,12 +467,18 @@ all code from `// Vector toward the light` to the end can be remove and be repla
   vec3 rayOrigin    = world_position;
   vec3 rayDirection = samplingHemisphere(prd.seed, tangent, bitangent, world_normal);
 
-  // Probability of the newRay (cosine distributed)
-  const float p = 1 / M_PI;
+  const float cos_theta = dot(rayDirection, world_normal);
+  // Probability density function of samplingHemisphere choosing this rayDirection
+  const float p = cos_theta / M_PI;
 
   // Compute the BRDF for this ray (assuming Lambertian reflection)
-  float cos_theta = dot(rayDirection, world_normal);
-  vec3  BRDF      = mat.pbrBaseColorFactor.xyz / M_PI;
+  vec3 albedo = mat.pbrBaseColorFactor.xyz;
+  if(mat.pbrBaseColorTexture > -1)
+  {
+    uint txtId = mat.pbrBaseColorTexture;
+    albedo *= texture(texturesMap[nonuniformEXT(txtId)], texcoord0).xyz;
+  }
+  vec3 BRDF = albedo / M_PI;
 
   // Recursively trace reflected light sources.
   if(prd.depth < 10)
diff --git a/ray_tracing_gltf/shaders/pathtrace.rchit b/ray_tracing_gltf/shaders/pathtrace.rchit
index 398176c..b9b2ac1 100644
--- a/ray_tracing_gltf/shaders/pathtrace.rchit
+++ b/ray_tracing_gltf/shaders/pathtrace.rchit
@@ -110,12 +110,11 @@ void main()
   vec3 rayOrigin    = world_position;
   vec3 rayDirection = samplingHemisphere(prd.seed, tangent, bitangent, world_normal);
 
-  // Compute the BRDF for this ray (assuming Lambertian reflection)
-  float cos_theta = dot(rayDirection, world_normal);
-
-  // Probability of the newRay (cosine distributed)
+  const float cos_theta = dot(rayDirection, world_normal);
+  // Probability density function of samplingHemisphere choosing this rayDirection
   const float p = cos_theta / M_PI;
 
+  // Compute the BRDF for this ray (assuming Lambertian reflection)
   vec3 albedo = mat.pbrBaseColorFactor.xyz;
   if(mat.pbrBaseColorTexture > -1)
   {
diff --git a/ray_tracing_gltf/shaders/sampling.glsl b/ray_tracing_gltf/shaders/sampling.glsl
index 67285dd..680e886 100644
--- a/ray_tracing_gltf/shaders/sampling.glsl
+++ b/ray_tracing_gltf/shaders/sampling.glsl
@@ -57,16 +57,17 @@ float rnd(inout uint prev)
 // Sampling
 //-------------------------------------------------------------------------------------------------
 
-// Randomly sampling around +Z
+// Randomly samples from a cosine-weighted hemisphere oriented in the `z` direction.
+// From Ray Tracing Gems section 16.6.1, "Cosine-Weighted Hemisphere Oriented to the Z-Axis"
 vec3 samplingHemisphere(inout uint seed, in vec3 x, in vec3 y, in vec3 z)
 {
-#define M_PI 3.141592
+#define M_PI 3.14159265
 
   float r1 = rnd(seed);
   float r2 = rnd(seed);
-  float sq = sqrt(1.0 - r2);
+  float sq = sqrt(r1);
 
-  vec3 direction = vec3(cos(2 * M_PI * r1) * sq, sin(2 * M_PI * r1) * sq, sqrt(r2));
+  vec3 direction = vec3(cos(2 * M_PI * r2) * sq, sin(2 * M_PI * r2) * sq, sqrt(1. - r1));
   direction      = direction.x * x + direction.y * y + direction.z * z;
 
   return direction;