Merge pull request #228 from Walther/hero-wavelength-sampling

feat: add Hero Wavelength Spectral Sampling

Justfile

@@ -31,6 +31,14 @@ test:
 bench *ARGS:
   cargo bench --quiet {{ARGS}}
 
+compare *ARGS:
+  git stash; \
+  cargo b -r -q; \
+  just cli render {{ARGS}} --output renders/before; \
+  git stash pop; \
+  cargo b -r -q; \
+  just cli render {{ARGS}} --output renders/after;
+
 # Verify no_std compatibility
 nostd:
   cargo clippy -q --release --package clovers --lib --no-default-features

clovers-cli/src/draw_cpu.rs

@@ -1,6 +1,8 @@
 //! An opinionated method for drawing a scene using the CPU for rendering.
 
-use clovers::wavelength::{random_wavelength, wavelength_into_xyz};
+use clovers::wavelength::{
+    random_wavelength, rotate_wavelength, wavelength_into_xyz, WAVE_SAMPLE_COUNT,
+};
 use clovers::Vec2;
 use clovers::{ray::Ray, scenes::Scene, Float};
 use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle};
@@ -123,13 +125,23 @@ fn render_pixel(
         let ray: Ray = scene
             .camera
             .get_ray(pixel_uv, lens_offset, time, wavelength);
-        let spectral_power: Float = trace(&ray, scene, 0, max_depth, rng, sampler);
-        // TODO: find and debug more sources of NaNs!
-        if spectral_power.is_normal() && spectral_power.is_sign_positive() {
-            let sample_color = wavelength_into_xyz(ray.wavelength);
-            pixel_color += sample_color * spectral_power;
+        let waves = rotate_wavelength(wavelength);
+        let spectral_powers = trace(&ray, scene, 0, max_depth, rng, sampler);
+        // Does our path have terminated wavelengths, i.e. does the path include a dispersive material?
+        if spectral_powers[1..].iter().all(|&p| p == 0.0) {
+            // Yes; colorize based on hero wavelength only
+            pixel_color += wavelength_into_xyz(waves[0]) * spectral_powers[0];
+        } else {
+            // No; colorize by all wavelengths, dividing by count
+            for i in 0..WAVE_SAMPLE_COUNT {
+                if spectral_powers[i].is_normal() && spectral_powers[i].is_sign_positive() {
+                    pixel_color += wavelength_into_xyz(waves[i]) * spectral_powers[i]
+                        / WAVE_SAMPLE_COUNT as Float;
+                }
+            }
         }
     }
+
     pixel_color / opts.samples as Float
 }
 

clovers-cli/src/trace.rs

@@ -6,7 +6,8 @@ use clovers::{
     pdf::{HitablePDF, MixturePDF, PDFTrait, PDF},
     ray::Ray,
     scenes::Scene,
-    spectrum::spectrum_xyz_to_p,
+    spectrum::{spectral_power, spectral_powers},
+    wavelength::{rotate_wavelength, WAVE_SAMPLE_COUNT},
     Float, EPSILON_SHADOW_ACNE,
 };
 use nalgebra::Unit;
@@ -25,14 +26,16 @@ pub fn trace(
     max_depth: u32,
     rng: &mut SmallRng,
     sampler: &dyn SamplerTrait,
-) -> Float {
-    let wavelength = ray.wavelength;
-    let bg: Float = spectrum_xyz_to_p(wavelength, scene.background);
+) -> [Float; WAVE_SAMPLE_COUNT] {
+    let hero = ray.wavelength;
+    let wavelengths = rotate_wavelength(hero);
+
+    let bg = spectral_powers(scene.background, wavelengths);
 
     // Have we reached the maximum recursion i.e. ray bounce depth?
     if depth > max_depth {
         // Ray bounce limit reached, early return zero emissivity
-        return 0.0;
+        return [0.0; WAVE_SAMPLE_COUNT];
     }
 
     // Send the ray to the scene, and see if it hits anything.
@@ -47,30 +50,31 @@ pub fn trace(
 
     // Get the emitted color from the surface that we just hit
     let emitted: Xyz<E> = hit_record.material.emit(ray, &hit_record);
-    let emitted: Float = spectrum_xyz_to_p(wavelength, emitted);
+    let emitted = spectral_powers(emitted, wavelengths);
 
     // Do we scatter?
     let Some(scatter_record) = hit_record.material.scatter(ray, &hit_record, rng) else {
         // No scatter, early return the emitted color only
         return emitted;
     };
     // We have scattered, and received an attenuation from the material
-    let attenuation = spectrum_xyz_to_p(wavelength, scatter_record.attenuation);
+    // Are we on a dispersive material? If so, terminate other wavelengths
+    let attenuations = match hit_record.material.is_wavelength_dependent() {
+        true => {
+            let mut ret = [0.0; WAVE_SAMPLE_COUNT];
+            ret[0] = spectral_power(scatter_record.attenuation, hero);
+            ret
+        }
+        false => spectral_powers(scatter_record.attenuation, wavelengths),
+    };
 
     // Check the material type and recurse accordingly:
     match scatter_record.material_type {
         MaterialType::Specular => {
-            // If we hit a specular material, generate a specular ray, and multiply it with the attenuation
-            let specular = trace(
-                // a scatter_record from a specular material should always have this ray
-                &scatter_record.specular_ray.unwrap(),
-                scene,
-                depth + 1,
-                max_depth,
-                rng,
-                sampler,
-            );
-            specular * attenuation
+            // If we hit a specular material, recurse with a specular ray, and multiply it with the attenuation
+            let scatter_ray = scatter_record.specular_ray.unwrap();
+            let specular = trace(&scatter_ray, scene, depth + 1, max_depth, rng, sampler);
+            std::array::from_fn(|i| specular[i] * attenuations[i])
         }
         MaterialType::Diffuse => {
             // Multiple Importance Sampling:
@@ -95,12 +99,7 @@ pub fn trace(
 
             // Get the distribution value for the PDF
             // TODO: improve correctness & optimization!
-            let pdf_val = mixture_pdf.value(scatter_ray.direction, ray.wavelength, ray.time, rng);
-            if pdf_val <= 0.0 {
-                // scattering impossible, prevent division by zero below
-                // for more ctx, see https://github.com/RayTracing/raytracing.github.io/issues/979#issuecomment-1034517236
-                return emitted;
-            }
+            let mis_pdf_value = mixture_pdf.value(direction, hero, ray.time, rng);
 
             // Calculate the PDF weighting for the scatter
             // TODO: improve correctness & optimization!
@@ -114,10 +113,9 @@ pub fn trace(
 
             // Recurse for the scattering ray
             let recurse = trace(&scatter_ray, scene, depth + 1, max_depth, rng, sampler);
-            // Tint and weight it according to the PDF
-            let scattered = attenuation * scattering_pdf * recurse / pdf_val;
-            // Blend it all together
-            emitted + scattered
+            std::array::from_fn(|i| {
+                emitted[i] + recurse[i] * attenuations[i] * scattering_pdf / mis_pdf_value
+            })
         }
     }
 }

clovers/benches/spectrum.rs

@@ -20,9 +20,9 @@ fn xyz_to_p(bencher: divan::Bencher) {
         .with_inputs(|| {
             let mut rng = SmallRng::from_entropy();
             let wave = random_wavelength(&mut rng);
-            let xyz: Xyz<E> = Xyz::new(1.0, 1.0, 1.0);
-            (wave, xyz)
+            let color: Xyz<E> = Xyz::new(1.0, 1.0, 1.0);
+            (wave, color)
         })
         .counter(1u32)
-        .bench_values(|(wave, xyz)| black_box(spectrum_xyz_to_p(wave, xyz)))
+        .bench_values(|(wave, color)| black_box(spectral_power(color, wave)))
 }

clovers/src/materials.rs

@@ -91,6 +91,10 @@ impl MaterialTrait for Material {
     fn emit(&self, ray: &Ray, hit_record: &HitRecord) -> Xyz<E> {
         self.kind.emit(ray, hit_record)
     }
+
+    fn is_wavelength_dependent(&self) -> bool {
+        self.thin_film.is_some() || self.kind.is_wavelength_dependent()
+    }
 }
 
 #[enum_dispatch]
@@ -113,6 +117,11 @@ pub trait MaterialTrait: Debug {
     fn emit(&self, _ray: &Ray, _hit_record: &HitRecord) -> Xyz<E> {
         Xyz::new(0.0, 0.0, 0.0)
     }
+
+    /// Returns true if the material has wavelength-dependent scattering, like dispersion or iridescence.
+    fn is_wavelength_dependent(&self) -> bool {
+        false
+    }
 }
 
 #[enum_dispatch(MaterialTrait)]

clovers/src/materials/dispersive.rs

@@ -117,4 +117,8 @@ impl MaterialTrait for Dispersive {
     }
 
     // TODO: should this material provide a `scattering_pdf` function?
+
+    fn is_wavelength_dependent(&self) -> bool {
+        true
+    }
 }

clovers/src/spectrum.rs

@@ -4,7 +4,10 @@
 
 use palette::{white_point::E, Xyz};
 
-use crate::{wavelength::Wavelength, Float};
+use crate::{
+    wavelength::{Wavelength, WAVE_SAMPLE_COUNT},
+    Float,
+};
 
 use self::spectra_xyz_5nm_380_780_097::equal_energy_reflectance;
 
@@ -13,30 +16,39 @@ pub mod spectrum_grid;
 
 /// Evaluate the spectrum at the given wavelength for the given XYZ color
 #[must_use]
-pub fn spectrum_xyz_to_p(lambda: Wavelength, xyz: Xyz<E>) -> Float {
+pub fn spectral_power(color: Xyz<E>, lambda: Wavelength) -> Float {
     // Currently, the data is only built for 5nm intervals
     // TODO: generate a file with 1nm intervals?
     let lambda: f64 = lambda as f64;
-    let xyz: [f64; 3] = [f64::from(xyz.x), f64::from(xyz.y), f64::from(xyz.z)];
+    let xyz: [f64; 3] = [f64::from(color.x), f64::from(color.y), f64::from(color.z)];
     let p = spectrum_grid::spectrum_xyz_to_p(lambda, xyz) / equal_energy_reflectance;
 
     #[allow(clippy::cast_possible_truncation)]
     let p = p as Float;
     p
 }
 
+/// Evaluate the spectral powers at the given multiple wavelengths for the given single XYZ color
+#[must_use]
+pub fn spectral_powers(
+    color: Xyz<E>,
+    wavelengths: [usize; WAVE_SAMPLE_COUNT],
+) -> [Float; WAVE_SAMPLE_COUNT] {
+    std::array::from_fn(|i| spectral_power(color, wavelengths[i]))
+}
+
 #[cfg(test)]
 #[allow(clippy::float_cmp)]
 mod unit {
     use palette::{white_point::E, Xyz};
 
-    use super::spectrum_xyz_to_p;
+    use super::spectral_power;
 
     #[test]
     fn equal_energy() {
         let lambda = 600;
         let xyz: Xyz<E> = Xyz::new(1.0, 1.0, 1.0);
-        let p = spectrum_xyz_to_p(lambda, xyz);
+        let p = spectral_power(xyz, lambda);
         // FIXME: floating point accuracy
         assert_eq!(1.000_000_7, p);
     }
@@ -46,7 +58,7 @@ mod unit {
     fn zero() {
         let lambda = 600;
         let xyz: Xyz<E> = Xyz::new(0.0, 0.0, 0.0);
-        let p = spectrum_xyz_to_p(lambda, xyz);
+        let p = spectral_power(xyz, lambda);
         // FIXME: floating point accuracy
         assert_eq!(0.0, p);
     }
@@ -56,7 +68,7 @@ mod unit {
     fn two() {
         let lambda = 600;
         let xyz: Xyz<E> = Xyz::new(2.0, 2.0, 2.0);
-        let p = spectrum_xyz_to_p(lambda, xyz);
+        let p = spectral_power(xyz, lambda);
         // FIXME: floating point accuracy
         assert_eq!(2.000_001_4, p);
     }

clovers/src/wavelength.rs

@@ -49,8 +49,8 @@ pub fn sample_wavelength(sample: Float) -> Wavelength {
 #[must_use]
 pub fn rotate_wavelength(hero: Wavelength) -> [Wavelength; WAVE_SAMPLE_COUNT] {
     from_fn(|j| {
-        (hero - MIN_WAVELENGTH + (j * SPECTRUM_SIZE / WAVE_SAMPLE_COUNT)) % SPECTRUM_SIZE
-            + MIN_WAVELENGTH
+        let step = j * SPECTRUM_SIZE / WAVE_SAMPLE_COUNT;
+        MIN_WAVELENGTH + (hero - MIN_WAVELENGTH + step) % SPECTRUM_SIZE
     })
 }
 

clovers/tests/spectrum.rs

@@ -6,23 +6,23 @@ proptest! {
   #[test]
   fn converts_all_wavelengths_black(lambda in SPECTRUM) {
     let xyz: Xyz<E> = Xyz::new(0.0, 0.0, 0.0);
-      let _ = spectrum_xyz_to_p(lambda, xyz);
+      let _ = spectral_power(xyz, lambda);
   }
 }
 
 proptest! {
   #[test]
   fn converts_all_wavelengths_grey(lambda in SPECTRUM) {
     let xyz: Xyz<E> = Xyz::new(0.5, 0.5, 0.5);
-      let _ = spectrum_xyz_to_p(lambda, xyz);
+      let _ = spectral_power(xyz, lambda);
   }
 }
 
 proptest! {
   #[test]
   fn converts_all_wavelengths_white(lambda in SPECTRUM) {
     let xyz: Xyz<E> = Xyz::new(1.0, 1.0, 1.0);
-      let _ = spectrum_xyz_to_p(lambda, xyz);
+      let _ = spectral_power(xyz, lambda);
   }
 }
 

clovers/tests/wavelength.rs

@@ -6,9 +6,7 @@ proptest! {
   fn converts_all_wavelengths(lambda in SPECTRUM) {
       let _ = wavelength_into_xyz(lambda);
   }
-}
 
-proptest! {
   #[test]
   fn rotates_all_wavelengths(lambda in SPECTRUM) {
       let mut waves = rotate_wavelength(lambda);
@@ -19,4 +17,11 @@ proptest! {
       prop_assert_eq!(diff, b.abs_diff(c));
       prop_assert_eq!(diff, c.abs_diff(d));
   }
+
+
+  #[test]
+  fn hero_always_first(lambda in SPECTRUM) {
+      let waves = rotate_wavelength(lambda);
+      prop_assert_eq!(lambda, waves[0]);
+  }
 }