Implement importance sampling for materials

src/integrators.rs

@@ -1,6 +1,7 @@
 use crate::math::Vec3;
 
 use super::colour::{ColourRgbF, Photon, Spectrum};
+use super::materials::MaterialSampleResult;
 use super::raycasting::{IntersectionInfo, Ray};
 use super::sampler::Sampler;
 use super::util::algebra_utils::try_change_of_basis_matrix;
@@ -59,7 +60,7 @@ impl Integrator for WhittedIntegrator {
                     .material
                     .sample(&(world_to_bsdf_space * info.retro), &photon)]
                 .iter()
-                .map(|direction| {
+                .map(|MaterialSampleResult { direction, pdf: _ }| {
                     let world_space_direction = bsdf_to_world_space * direction;
                     match sampler
                         .sample(&Ray::new(info.location, world_space_direction).bias(0.000_000_1))
@@ -118,11 +119,14 @@ impl Integrator for SimpleRandomIntegrator {
             .expect("Expected matrix to be invertable.");
         let world_space_w_i = info.retro;
         let w_i = world_to_bsdf_space * world_space_w_i;
-        let w_o = info.material.sample(&w_i, &photon);
+        let MaterialSampleResult {
+            direction: w_o,
+            pdf: w_o_pdf,
+        } = info.material.sample(&w_i, &photon);
         let world_space_w_o = bsdf_to_world_space * w_o;
         info.material.bsdf()(
-            &w_o.normalize(),
-            &w_i.normalize(),
+            &w_o,
+            &w_i,
             &match sampler.sample(&Ray::new(info.location, world_space_w_o).bias(0.000_000_1)) {
                 None => photon.set_intensity(test_lighting_environment(
                     &world_space_w_o,
@@ -132,6 +136,7 @@ impl Integrator for SimpleRandomIntegrator {
                     self.integrate(&sampler, &recursive_hit, &photon, recursion_limit - 1)
                 }
             }
+            .scale_intensity(w_o_pdf)
             .scale_intensity(world_space_w_o.dot(&info.normal).abs()),
         )
     }
@@ -142,7 +147,7 @@ pub fn test_lighting_environment(w_o: &Vec3, wavelength: f64) -> f64 {
     if w_o.dot(&sun_direction) >= 0.99 {
         300.0
     } else {
-        let sky_colour = ColourRgbF::new(w_o.y(), w_o.y(), 1.0) * 0.1;
+        let sky_colour = ColourRgbF::new(w_o.y(), w_o.y(), 1.0);
         Spectrum::reflection_from_linear_rgb(&sky_colour).intensity_at_wavelength(wavelength)
     }
 }

src/materials/lambertian_material.rs

@@ -1,8 +1,12 @@
 use crate::colour::{Photon, Spectrum};
 use crate::math::Vec3;
 
-use super::Material;
+use super::{Material, MaterialSampleResult};
 
+use rand::distributions::Open01;
+use rand::{thread_rng, Rng};
+
+use std::f64::consts::PI;
 use std::fmt::Debug;
 
 #[derive(Debug)]
@@ -28,4 +32,29 @@ impl Material for LambertianMaterial {
             result
         })
     }
+
+    fn sample(&self, _w_i: &Vec3, _photon: &Photon) -> MaterialSampleResult {
+        let mut rng = thread_rng();
+        let mut w_o = Vec3::new(
+            2.0 * rng.sample::<f64, _>(Open01) - 1.0,
+            2.0 * rng.sample::<f64, _>(Open01) - 1.0,
+            0.0,
+        );
+        while w_o.norm_squared() > 1.0 {
+            w_o = Vec3::new(
+                2.0 * rng.sample::<f64, _>(Open01) - 1.0,
+                2.0 * rng.sample::<f64, _>(Open01) - 1.0,
+                0.0,
+            );
+        }
+        w_o.coords[2] = (1.0 - w_o.x() * w_o.x() - w_o.y() * w_o.y())
+            .sqrt()
+            .max(0.0);
+        let cos_theta = w_o.dot(&Vec3::unit_z());
+        let sin_theta = (1.0 - cos_theta * cos_theta).sqrt();
+        MaterialSampleResult {
+            direction: w_o.normalize(),
+            pdf: (cos_theta * sin_theta) / PI,
+        }
+    }
 }

src/materials/mod.rs

@@ -4,6 +4,7 @@ use super::colour::Photon;
 
 use rand::distributions::{Open01, OpenClosed01};
 use rand::{thread_rng, Rng};
+use std::f64::consts::PI;
 use std::fmt::Debug;
 
 pub mod lambertian_material;
@@ -18,10 +19,15 @@ pub use reflective_material::ReflectiveMaterial;
 pub mod smooth_transparent_dialectric;
 pub use smooth_transparent_dialectric::SmoothTransparentDialectric;
 
+pub struct MaterialSampleResult {
+    pub direction: Vec3,
+    pub pdf: f64,
+}
+
 pub trait Material: Debug + Sync + Send {
     fn bsdf<'a>(&'a self) -> Box<dyn Fn(&Vec3, &Vec3, &Photon) -> Photon + 'a>;
 
-    fn sample(&self, _w_i: &Vec3, _photon: &Photon) -> Vec3 {
+    fn sample(&self, _w_i: &Vec3, _photon: &Photon) -> MaterialSampleResult {
         let mut rng = thread_rng();
         let mut w_o = Vec3::new(
             2.0 * rng.sample::<f64, _>(Open01) - 1.0,
@@ -35,10 +41,9 @@ pub trait Material: Debug + Sync + Send {
                 rng.sample::<f64, _>(OpenClosed01),
             );
         }
-        w_o
-    }
-
-    fn pdf(&self, _w_i: &Vec3, _w_o: &Vec3) -> f64 {
-        1.0
+        MaterialSampleResult {
+            direction: w_o.normalize(),
+            pdf: 1.0 / (2.0 * PI),
+        }
     }
 }

src/materials/reflective_material.rs

@@ -3,7 +3,7 @@ use crate::math::Vec3;
 
 use std::fmt::Debug;
 
-use super::Material;
+use super::{Material, MaterialSampleResult};
 
 #[derive(Debug)]
 pub struct ReflectiveMaterial {
@@ -39,7 +39,10 @@ impl Material for ReflectiveMaterial {
         })
     }
 
-    fn sample(&self, w_o: &Vec3, _photon: &Photon) -> Vec3 {
-        Vec3::new(-w_o.x(), -w_o.y(), w_o.z())
+    fn sample(&self, w_o: &Vec3, _photon: &Photon) -> MaterialSampleResult {
+        MaterialSampleResult {
+            direction: Vec3::new(-w_o.x(), -w_o.y(), w_o.z()),
+            pdf: 1.0,
+        }
     }
 }

src/materials/smooth_transparent_dialectric.rs

@@ -1,5 +1,5 @@
 use crate::colour::{Photon, Spectrum};
-use crate::materials::Material;
+use crate::materials::{Material, MaterialSampleResult};
 use crate::math::Vec3;
 
 use rand::random;
@@ -83,17 +83,12 @@ impl Material for SmoothTransparentDialectric {
             } else if (*w_o - fresnel.transmission_direction).norm_squared() < 0.0000000001 {
                 photon_in.scale_intensity(fresnel.transmission_strength)
             } else {
-                /*dbg!(
-                    w_o,
-                    fresnel.reflection_direction,
-                    fresnel.transmission_direction
-                );*/
                 photon_in.set_intensity(0.0)
             }
         })
     }
 
-    fn sample(&self, w_i: &Vec3, photon: &Photon) -> Vec3 {
+    fn sample(&self, w_i: &Vec3, photon: &Photon) -> MaterialSampleResult {
         let (eta1, eta2) = if w_i.z() >= 0.0 {
             (1.0, self.eta.intensity_at_wavelength(photon.wavelength))
         } else {
@@ -101,15 +96,20 @@ impl Material for SmoothTransparentDialectric {
         };
         let fresnel = fresnel(w_i, eta1, eta2);
         if fresnel.transmission_strength <= 0.0000000001 {
-            fresnel.reflection_direction
+            MaterialSampleResult {
+                direction: fresnel.reflection_direction,
+                pdf: 0.5,
+            }
         } else if fresnel.reflection_strength <= 0.0000000001 || random() {
-            fresnel.transmission_direction
+            MaterialSampleResult {
+                direction: fresnel.transmission_direction,
+                pdf: 0.5,
+            }
         } else {
-            fresnel.reflection_direction
+            MaterialSampleResult {
+                direction: fresnel.reflection_direction,
+                pdf: 0.5,
             }
         }
-
-    fn pdf(&self, _w_i: &Vec3, _w_o: &Vec3) -> f64 {
-        0.0
     }
 }