Improve energy conservation in the solari BRDF

crates/bevy_solari/src/pathtracer/pathtracer.wgsl

@@ -1,11 +1,10 @@
 enable wgpu_ray_query;
 
 #import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance
-#import bevy_pbr::pbr_functions::calculate_F0
 #import bevy_pbr::utils::{rand_f, rand_vec2f}
 #import bevy_render::maths::{PI, orthonormalize}
 #import bevy_render::view::View
-#import bevy_solari::brdf::{evaluate_brdf, evaluate_and_sample_brdf, brdf_pdf}
+#import bevy_solari::brdf::{evaluate_brdf, evaluate_and_sample_brdf, brdf_pdf, F_AB}
 #import bevy_solari::sampling::{sample_random_light, random_emissive_light_pdf, ggx_vndf_pdf, power_heuristic}
 #import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, ResolvedRayHitFull, RAY_T_MIN, RAY_T_MAX, MIRROR_ROUGHNESS_THRESHOLD}
 
@@ -45,6 +44,8 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
         if ray.kind != RAY_QUERY_INTERSECTION_NONE {
             let ray_hit = resolve_ray_hit_full(ray);
             let wo = -ray_direction;
+            let NdotV = max(dot(ray_hit.world_normal, wo), 0.0001);
+            let F_ab = F_AB(ray_hit.material.perceptual_roughness, NdotV);
 
             // Emissive contribution
             var mis_weight = 1.0;
@@ -62,16 +63,16 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
 
                 mis_weight = 1.0;
                 if direct_lighting.brdf_rays_can_hit {
-                    let pdf_of_bounce = brdf_pdf(wo, direct_lighting.wi, ray_hit.world_normal, ray_hit.material);
+                    let pdf_of_bounce = brdf_pdf(wo, direct_lighting.wi, ray_hit.world_normal, ray_hit.material, F_ab);
                     mis_weight = power_heuristic(1.0 / direct_lighting.inverse_pdf, pdf_of_bounce);
                 }
 
-                let direct_lighting_brdf = evaluate_brdf(wo, direct_lighting.wi, ray_hit.world_normal, ray_hit.material);
+                let direct_lighting_brdf = evaluate_brdf(wo, direct_lighting.wi, ray_hit.world_normal, ray_hit.material, F_ab);
                 radiance += mis_weight * throughput * direct_lighting.radiance * direct_lighting.inverse_pdf * direct_lighting_brdf;
             }
 
             // Sample new ray direction from the material BRDF for next bounce and apply BRDF
-            let next_bounce = evaluate_and_sample_brdf(wo, ray_hit.world_normal, ray_hit.material, &rng);
+            let next_bounce = evaluate_and_sample_brdf(wo, ray_hit.world_normal, ray_hit.material, F_ab, &rng);
             if next_bounce.pdf == 0.0 { break; }
             ray_direction = next_bounce.wi;
             ray_origin = ray_hit.world_position + (ray_hit.geometric_world_normal * RAY_T_MIN);

crates/bevy_solari/src/realtime/restir_di.wgsl

@@ -8,7 +8,7 @@ enable wgpu_ray_query;
 #import bevy_pbr::utils::{rand_f, rand_range_u, sample_disk}
 #import bevy_render::maths::PI
 #import bevy_render::view::View
-#import bevy_solari::brdf::{evaluate_diffuse_brdf, evaluate_specular_brdf}
+#import bevy_solari::brdf::{evaluate_diffuse_brdf, evaluate_specular_brdf, F_AB}
 #import bevy_solari::gbuffer_utils::{gpixel_resolve, pixel_dissimilar, permute_pixel}
 #import bevy_solari::presample_light_tiles::unpack_resolved_light_sample
 #import bevy_solari::sampling::{LightSample, ResolvedLightSample, NULL_LIGHT_ID, calculate_resolved_light_contribution, resolve_and_calculate_light_contribution, resolve_light_sample, trace_light_visibility, balance_heuristic}
@@ -79,10 +79,12 @@ fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3<u32>) {
 #endif
 
     let wo = normalize(view.world_position - surface.world_position);
-    var brdf = evaluate_diffuse_brdf(wo, merge_result.wi, surface.world_normal, surface.material);
+    let NdotV = max(dot(surface.world_normal, wo), 0.0001);
+    let F_ab = F_AB(surface.material.perceptual_roughness, NdotV);
+    var brdf = evaluate_diffuse_brdf(wo, merge_result.wi, surface.world_normal, surface.material, F_ab);
     // Only consider the specular lobe if the surface is not smooth, else leave it for the specular GI pass to handle
     if surface.material.roughness > SPECULAR_GI_FOR_DI_ROUGHNESS_THRESHOLD {
-        brdf += evaluate_specular_brdf(wo, merge_result.wi, surface.world_normal, surface.material);
+        brdf += evaluate_specular_brdf(wo, merge_result.wi, surface.world_normal, surface.material, F_ab);
     }
 
     var pixel_color = merge_result.selected_sample_radiance * combined_reservoir.unbiased_contribution_weight;

crates/bevy_solari/src/realtime/restir_gi.wgsl

@@ -6,7 +6,7 @@ enable wgpu_ray_query;
 #import bevy_pbr::utils::{rand_f, sample_uniform_hemisphere, uniform_hemisphere_inverse_pdf, sample_disk}
 #import bevy_render::maths::PI
 #import bevy_render::view::View
-#import bevy_solari::brdf::evaluate_diffuse_brdf
+#import bevy_solari::brdf::{evaluate_diffuse_brdf, F_AB}
 #import bevy_solari::gbuffer_utils::{gpixel_resolve, pixel_dissimilar, permute_pixel}
 #import bevy_solari::sampling::{sample_random_light, trace_point_visibility, balance_heuristic, isnan}
 #import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN, RAY_T_MAX}
@@ -80,7 +80,9 @@ fn spatial_and_shade(@builtin(global_invocation_id) global_id: vec3<u32>) {
 #endif
 
     let wo = normalize(view.world_position - surface.world_position);
-    let brdf = evaluate_diffuse_brdf(wo, merge_result.wi, surface.world_normal, surface.material);
+    let NdotV = max(dot(surface.world_normal, wo), 0.0001);
+    let F_ab = F_AB(surface.material.perceptual_roughness, NdotV);
+    let brdf = evaluate_diffuse_brdf(wo, merge_result.wi, surface.world_normal, surface.material, F_ab);
 
     var pixel_color = textureLoad(view_output, global_id.xy);
     pixel_color += vec4(merge_result.selected_sample_radiance * combined_reservoir.unbiased_contribution_weight * view.exposure * brdf, 0.0);

crates/bevy_solari/src/realtime/specular_gi.wgsl

@@ -7,7 +7,7 @@ enable wgpu_ray_query;
 #import bevy_pbr::utils::rand_f
 #import bevy_render::maths::{orthonormalize, PI}
 #import bevy_render::view::View
-#import bevy_solari::brdf::{evaluate_brdf, evaluate_specular_brdf}
+#import bevy_solari::brdf::{evaluate_brdf, evaluate_specular_brdf, F_AB}
 #import bevy_solari::gbuffer_utils::{gpixel_resolve, ResolvedGPixel}
 #import bevy_solari::sampling::{sample_random_light, random_emissive_light_pdf, sample_ggx_vndf, ggx_vndf_pdf, ggx_vndf_sample_invalid, power_heuristic}
 #import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, ResolvedRayHitFull, RAY_T_MIN, RAY_T_MAX, MIRROR_ROUGHNESS_THRESHOLD}
@@ -67,7 +67,9 @@ fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
         }
     }
 
-    let brdf = evaluate_specular_brdf(wo, wi, surface.world_normal, surface.material);
+    let NdotV = max(dot(surface.world_normal, wo), 0.0001);
+    let F_ab = F_AB(surface.material.perceptual_roughness, NdotV);
+    let brdf = evaluate_specular_brdf(wo, wi, surface.world_normal, surface.material, F_ab);
     radiance *= brdf * view.exposure;
 
     var pixel_color = textureLoad(view_output, global_id.xy);
@@ -107,6 +109,8 @@ fn trace_glossy_path(pixel_id: vec2<u32>, primary_surface: ResolvedGPixel, initi
 
         let wo = -wi;
         let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N));
+        let NdotV = max(dot(ray_hit.world_normal, wo), 0.0001);
+        let F_ab = F_AB(ray_hit.material.perceptual_roughness, NdotV);
 
         // Add emissive contribution
         let mis_weight = emissive_mis_weight(i, primary_surface.material.roughness, p_bounce, ray_hit);
@@ -140,7 +144,7 @@ fn trace_glossy_path(pixel_id: vec2<u32>, primary_surface: ResolvedGPixel, initi
         } else if !surface_perfect_mirror {
             // Sample direct lighting (NEE)
             let direct_lighting = sample_random_light(ray_hit.world_position, ray_hit.world_normal, rng);
-            let direct_lighting_brdf = evaluate_brdf(wo, direct_lighting.wi, ray_hit.world_normal, ray_hit.material);
+            let direct_lighting_brdf = evaluate_brdf(wo, direct_lighting.wi, ray_hit.world_normal, ray_hit.material, F_ab);
             let mis_weight = nee_mis_weight(direct_lighting.inverse_pdf, direct_lighting.brdf_rays_can_hit, wo_tangent, direct_lighting.wi, ray_hit, TBN);
             radiance += throughput * mis_weight * direct_lighting.radiance * direct_lighting.inverse_pdf * direct_lighting_brdf;
         }
@@ -153,7 +157,7 @@ fn trace_glossy_path(pixel_id: vec2<u32>, primary_surface: ResolvedGPixel, initi
 
         // Update throughput for next bounce
         p_bounce = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness);
-        throughput *= evaluate_brdf(wo, wi, N, ray_hit.material);
+        throughput *= evaluate_brdf(wo, wi, N, ray_hit.material, F_ab);
         if ray_hit.material.roughness > MIRROR_ROUGHNESS_THRESHOLD {
             throughput /= p_bounce;
         }

crates/bevy_solari/src/scene/brdf.wgsl

@@ -4,7 +4,7 @@ enable wgpu_ray_query;
 
 #import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance
 #import bevy_pbr::lighting::{D_GGX, V_SmithGGXCorrelated, specular_multiscatter}
-#import bevy_pbr::pbr_functions::calculate_F0
+#import bevy_pbr::pbr_functions::calculate_F0_dielectric
 #import bevy_pbr::utils::{rand_f, sample_cosine_hemisphere}
 #import bevy_render::maths::{PI, orthonormalize}
 #import bevy_solari::sampling::{sample_ggx_vndf, ggx_vndf_pdf, ggx_vndf_sample_invalid}
@@ -22,26 +22,28 @@ struct LobeReflectances {
 }
 
 // Hemispherical reflectance of each lobe
-fn lobe_reflectances(F0: vec3<f32>, material: ResolvedMaterial, NdotV: f32) -> LobeReflectances {
-    let F_ab = F_AB(material.perceptual_roughness, NdotV);
+fn lobe_reflectances(F0_metal: vec3<f32>, F0_dielectric: vec3<f32>, material: ResolvedMaterial, F_ab: vec2<f32>) -> LobeReflectances {
     let multiscattering_factor = 1.0 / (F_ab.x + F_ab.y) - 1.0;
-    let rho_specular = (F0 * F_ab.x + F_ab.y) * (1.0 + F0 * multiscattering_factor);
+    let rho_specular_metallic = (F0_metal * F_ab.x + F_ab.y) * (1.0 + F0_metal * multiscattering_factor);
+    let rho_specular_dielectric = (F0_dielectric * F_ab.x + F_ab.y) * (1.0 + F0_dielectric * multiscattering_factor);
     return LobeReflectances(
-        rho_specular,
-        (1.0 - material.metallic) * (1.0 - rho_specular) * material.base_color,
+        material.metallic * rho_specular_metallic + (1.0 - material.metallic) * rho_specular_dielectric,
+        (1.0 - material.metallic) * (1.0 - rho_specular_dielectric) * material.base_color,
     );
 }
 
 fn evaluate_and_sample_brdf(
     wo: vec3<f32>,
     world_normal: vec3<f32>,
     material: ResolvedMaterial,
+    F_ab: vec2<f32>,
     rng: ptr<function, u32>,
 ) -> EvaluateAndSampleBrdfResult {
     let NdotV = dot(world_normal, wo);
     if NdotV < 0.0001 { return EvaluateAndSampleBrdfResult(vec3(0.0), vec3(0.0), 0.0); }
-    let F0 = calculate_F0(material.base_color, material.metallic, vec3(material.reflectance));
-    let rho = lobe_reflectances(F0, material, NdotV);
+    let F0_metal = material.base_color;
+    let F0_dielectric = calculate_F0_dielectric(vec3(material.reflectance));
+    let rho = lobe_reflectances(F0_metal, F0_dielectric, material, F_ab);
     let specular_weight = luminance(rho.specular) / luminance(rho.specular + rho.diffuse);
     let diffuse_weight = 1.0 - specular_weight;
 
@@ -69,7 +71,7 @@ fn evaluate_and_sample_brdf(
         if material.roughness <= MIRROR_ROUGHNESS_THRESHOLD {
             return EvaluateAndSampleBrdfResult(
                 wi,
-                evaluate_specular_brdf(wo, wi, world_normal, material) / specular_weight,
+                evaluate_specular_brdf(wo, wi, world_normal, material, F_ab) / specular_weight,
                 bitcast<f32>(0x7F800000u) // INF
             );
         }
@@ -78,7 +80,7 @@ fn evaluate_and_sample_brdf(
     let diffuse_pdf = wi_tangent.z / PI;
     let specular_pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, material.roughness);
     let pdf = (diffuse_weight * diffuse_pdf) + (specular_weight * specular_pdf);
-    let throughput = evaluate_brdf(wo, wi, world_normal, material) / pdf;
+    let throughput = evaluate_brdf(wo, wi, world_normal, material, F_ab) / pdf;
     return EvaluateAndSampleBrdfResult(wi, throughput, pdf);
 }
 
@@ -87,48 +89,55 @@ fn evaluate_brdf(
     wi: vec3<f32>,
     world_normal: vec3<f32>,
     material: ResolvedMaterial,
+    F_ab: vec2<f32>,
 ) -> vec3<f32> {
-    return evaluate_diffuse_brdf(wo, wi, world_normal, material) + evaluate_specular_brdf(wo, wi, world_normal, material);
+    return evaluate_diffuse_brdf(wo, wi, world_normal, material, F_ab) + evaluate_specular_brdf(wo, wi, world_normal, material, F_ab);
 }
 
-fn evaluate_diffuse_brdf(wo: vec3<f32>, wi: vec3<f32>, world_normal: vec3<f32>, material: ResolvedMaterial) -> vec3<f32> {
+fn evaluate_diffuse_brdf(wo: vec3<f32>, wi: vec3<f32>, world_normal: vec3<f32>, material: ResolvedMaterial, F_ab: vec2<f32>) -> vec3<f32> {
     let NdotL = dot(world_normal, wi);
     let NdotV = dot(world_normal, wo);
     if NdotL < 0.0001 || NdotV < 0.0001 { return vec3(0.0); }
-    let F0 = calculate_F0(material.base_color, material.metallic, vec3(material.reflectance));
-    let rho = lobe_reflectances(F0, material, NdotV);
+    let F0_metal = material.base_color;
+    let F0_dielectric = calculate_F0_dielectric(vec3(material.reflectance));
+    let rho = lobe_reflectances(F0_metal, F0_dielectric, material, F_ab);
     return rho.diffuse / PI * NdotL;
 }
 
-fn evaluate_specular_brdf(wo: vec3<f32>, wi: vec3<f32>, world_normal: vec3<f32>, material: ResolvedMaterial) -> vec3<f32> {
+fn evaluate_specular_brdf(wo: vec3<f32>, wi: vec3<f32>, world_normal: vec3<f32>, material: ResolvedMaterial, F_ab: vec2<f32>) -> vec3<f32> {
     let H = normalize(wi + wo);
     let NdotL = dot(world_normal, wi);
     let NdotH = dot(world_normal, H);
     let LdotH = dot(wi, H);
     let NdotV = dot(world_normal, wo);
     if NdotL < 0.0001 || NdotH < 0.0001 || LdotH < 0.0001 || NdotV < 0.0001 { return vec3(0.0); }
 
-    let F0 = calculate_F0(material.base_color, material.metallic, vec3(material.reflectance));
-    let F = fresnel(F0, LdotH);
+    let F0_metal = material.base_color;
+    let F0_dielectric = calculate_F0_dielectric(vec3(material.reflectance));
 
     if material.roughness <= MIRROR_ROUGHNESS_THRESHOLD {
         if abs(NdotH - 1.0) < 0.0001 {
-            return F;
+            let F_metal = fresnel(F0_metal, LdotH);
+            let F_dielectric = fresnel(F0_dielectric, LdotH);
+            return material.metallic * F_metal + (1.0 - material.metallic) * F_dielectric;
         } else {
             return vec3(0.0);
         }
     }
 
     let D = D_GGX(material.roughness, NdotH);
     let Vs = V_SmithGGXCorrelated(material.roughness, NdotV, NdotL);
-    let F_ab = F_AB(material.perceptual_roughness, NdotV);
-    return specular_multiscatter(D, Vs, F, F0, F_ab, 1.0) * NdotL;
+    let F_metal = fresnel(F0_metal, LdotH);
+    let F_dielectric = fresnel(F0_dielectric, LdotH);
+    return (material.metallic * specular_multiscatter(D, Vs, F_metal, F0_metal, F_ab, 1.0)
+          + (1.0 - material.metallic) * specular_multiscatter(D, Vs, F_dielectric, F0_dielectric, F_ab, 1.0)) * NdotL;
 }
 
-fn brdf_pdf(wo: vec3<f32>, wi: vec3<f32>, world_normal: vec3<f32>, material: ResolvedMaterial) -> f32 {
+fn brdf_pdf(wo: vec3<f32>, wi: vec3<f32>, world_normal: vec3<f32>, material: ResolvedMaterial, F_ab: vec2<f32>) -> f32 {
     let NdotV = max(dot(world_normal, wo), 0.0001);
-    let F0 = calculate_F0(material.base_color, material.metallic, vec3(material.reflectance));
-    let rho = lobe_reflectances(F0, material, NdotV);
+    let F0_metal = material.base_color;
+    let F0_dielectric = calculate_F0_dielectric(vec3(material.reflectance));
+    let rho = lobe_reflectances(F0_metal, F0_dielectric, material, F_ab);
     let specular_weight = luminance(rho.specular) / luminance(rho.specular + rho.diffuse);
     let diffuse_weight = 1.0 - specular_weight;