From 43147f3c5c58d47dbd30c64003133ba78bdf40ce Mon Sep 17 00:00:00 2001
From: Dzmitry Malyshau <kvark@fastmail.com>
Date: Thu, 17 Oct 2024 23:46:39 -0700
Subject: [PATCH] Take texel shape into account for sample environment map

---
 blade-helpers/src/lib.rs                  | 2 +-
 blade-render/code/env-importance.inc.wgsl | 9 ++++++++-
 blade-render/code/ray-trace.wgsl          | 7 +++++--
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/blade-helpers/src/lib.rs b/blade-helpers/src/lib.rs
index e8872ec2..c78124a5 100644
--- a/blade-helpers/src/lib.rs
+++ b/blade-helpers/src/lib.rs
@@ -10,7 +10,7 @@ pub use hud::{populate_debug_selection, ExposeHud};
 pub fn default_ray_config() -> blade_render::RayConfig {
     blade_render::RayConfig {
         num_environment_samples: 1,
-        environment_importance_sampling: false,
+        environment_importance_sampling: true,
         tap_count: 2,
         tap_radius: 20,
         tap_confidence_near: 15,
diff --git a/blade-render/code/env-importance.inc.wgsl b/blade-render/code/env-importance.inc.wgsl
index 43421fce..7c98413f 100644
--- a/blade-render/code/env-importance.inc.wgsl
+++ b/blade-render/code/env-importance.inc.wgsl
@@ -5,9 +5,16 @@ struct EnvImportantSample {
     pdf: f32,
 }
 
+// Returns the range of values proportional to the area, given the texel Y
+fn compute_latitude_area_bounds(texel_y: i32, dim: u32) -> vec2<f32> {
+    return cos(vec2<f32>(vec2<i32>(texel_y, texel_y + 1)) / f32(dim) * PI);
+}
+
 fn compute_texel_solid_angle(itc: vec2<i32>, dim: vec2<u32>) -> f32 {
+    //Note: this has to agree with `map_equirect_uv_to_dir`
     let meridian_solid_angle = 4.0 * PI / f32(dim.x);
-    let meridian_part = 0.5 * (cos(PI * f32(itc.y) / f32(dim.y)) - cos(PI * f32(itc.y + 1) / f32(dim.y)));
+    let bounds = compute_latitude_area_bounds(itc.y, dim.y);
+    let meridian_part = 0.5 * (bounds.x - bounds.y);
     return meridian_solid_angle * meridian_part;
 }
 
diff --git a/blade-render/code/ray-trace.wgsl b/blade-render/code/ray-trace.wgsl
index 24132ea4..df4c3351 100644
--- a/blade-render/code/ray-trace.wgsl
+++ b/blade-render/code/ray-trace.wgsl
@@ -201,8 +201,11 @@ fn sample_light_from_environment(rng: ptr<function, RandomState>) -> LightSample
     // sample the incoming radiance
     ls.radiance = textureLoad(env_map, es.pixel, 0).xyz;
     // for determining direction - offset randomly within the texel
-    // Note: this only works if the texels are sufficiently small
-    ls.uv = (vec2<f32>(es.pixel) + vec2<f32>(random_gen(rng), random_gen(rng))) / vec2<f32>(dim);
+    // this offset has to be uniformly distributed across the surface of the texel
+    let u = (f32(es.pixel.x) + random_gen(rng)) / f32(dim.x);
+    let bounds = compute_latitude_area_bounds(es.pixel.y, dim.y);
+    let v = acos(mix(bounds.x, bounds.y, random_gen(rng))) / PI;
+    ls.uv = vec2<f32>(u, v);
     return ls;
 }