Stratify Sampling

src/camera/mod.rs

@@ -4,11 +4,12 @@ use crate::geometry::HitRecord;
 use crate::geometry::Hittable;
 use crate::interval::Interval;
 use crate::material::Surface;
-use crate::math;
 use crate::ray::Ray;
 use image::RgbImage;
+use itertools::iproduct;
 use nalgebra::Vector3;
 use rand::prelude::*;
+use rand::rngs::ThreadRng;
 use std::f64::consts::PI;
 use std::io::{self, Write};
 use std::sync::Arc;
@@ -22,11 +23,13 @@ pub struct Camera {
     pub image_width: u32,
     pub image_height: u32,
 
+    pub sqrt_spp: u32,
+    pub recip_sqrt_spp: f64,
+
     pub center: Vector3<f64>,
     pub pixel00_loc: Vector3<f64>,
     pub pixel_delta_u: Vector3<f64>,
     pub pixel_delta_v: Vector3<f64>,
-    pub samples: u32,
     pub samples_scale: f64,
     pub max_bounces: u32,
     pub threads: usize,
@@ -38,9 +41,7 @@ pub struct Camera {
     pub background: Vector3<f64>,
 }
 
-fn random_in_unit_disk() -> Vector3<f64> {
-    let mut rng = rand::rng();
-
+fn random_in_unit_disk(rng: &mut ThreadRng) -> Vector3<f64> {
     loop {
         let p = Vector3::new(
             rng.random_range(-1.0f64..1.0f64),
@@ -65,6 +66,21 @@ fn degrees_to_radians(degress: f64) -> f64 {
     degress * PI / 180.0
 }
 
+fn sample_square_stratified(
+    s_i: u32,
+    s_j: u32,
+    recip_sqrt_spp: f64,
+    rng: &mut ThreadRng,
+) -> (f64, f64) {
+    let s_i = s_i as f64;
+    let s_j = s_j as f64;
+
+    (
+        ((s_i + rng.random::<f64>()) * recip_sqrt_spp) - 0.5,
+        ((s_j + rng.random::<f64>()) * recip_sqrt_spp) - 0.5,
+    )
+}
+
 impl Camera {
     pub fn new(options: CameraOptions) -> Self {
         let (image_width, image_height) = options.get_dimensions();
@@ -105,7 +121,9 @@ impl Camera {
         let defocus_disk_v = v * defocus_radius;
 
         let samples = options.samples;
-        let samples_scale = 1.0 / (samples as f64);
+        let sqrt_spp = (samples as f64).sqrt() as u32;
+        let recip_sqrt_spp = 1.0 / (sqrt_spp as f64);
+        let samples_scale = 1.0 / ((sqrt_spp * sqrt_spp) as f64);
 
         let max_bounces = options.max_bounces;
         let threads = options.threads;
@@ -117,14 +135,15 @@ impl Camera {
             pixel00_loc,
             pixel_delta_u,
             pixel_delta_v,
-            samples,
             samples_scale,
             max_bounces,
             threads,
             defocus_disk_u,
             defocus_disk_v,
             defocus_angle,
             background,
+            sqrt_spp,
+            recip_sqrt_spp,
         }
     }
 
@@ -187,8 +206,16 @@ impl Camera {
     }
 
     pub fn get_pixel(&self, world: &Geometry, x: u32, y: u32) -> image::Rgb<u8> {
-        let color = (0..self.samples)
-            .map(|_| self.ray_color(&self.get_ray(x, y), self.max_bounces, world))
+        let mut rng = rand::rng();
+        let color: Vector3<f64> = iproduct!(0..self.sqrt_spp, 0..self.sqrt_spp)
+            .map(|(s_x, s_y)| {
+                self.ray_color(
+                    &self.get_ray(x, y, s_x, s_y, &mut rng),
+                    self.max_bounces,
+                    world,
+                    &mut rng,
+                )
+            })
             .sum::<Vector3<f64>>()
             * self.samples_scale;
 
@@ -200,38 +227,41 @@ impl Camera {
         image::Rgb([r as u8, g as u8, b as u8])
     }
 
-    pub fn get_ray(&self, x: u32, y: u32) -> Ray {
-        let offset = math::random_box(-0.5f64..0.5f64);
+    pub fn get_ray(&self, x: u32, y: u32, s_x: u32, s_y: u32, rng: &mut ThreadRng) -> Ray {
+        let (offset_x, offset_y) = sample_square_stratified(s_x, s_y, self.recip_sqrt_spp, rng);
         let pixel_sample = self.pixel00_loc
-            + (self.pixel_delta_u * (offset.x + x as f64))
-            + (self.pixel_delta_v * (offset.y + y as f64));
+            + (self.pixel_delta_u * (offset_x + x as f64))
+            + (self.pixel_delta_v * (offset_y + y as f64));
 
         let ray_origin = match self.defocus_angle <= 0.0 {
             true => self.center,
-            false => self.defocus_disk_sample(),
+            false => self.defocus_disk_sample(rng),
         };
         let ray_direction = pixel_sample - ray_origin;
 
-        let mut rng = rand::rng();
-        let t = rng.random_range(0.0f64..1.0f64);
-
-        Ray::new(ray_origin, ray_direction, t)
+        Ray::new(ray_origin, ray_direction, rng.random::<f64>())
     }
 
-    pub fn defocus_disk_sample(&self) -> Vector3<f64> {
-        let p = random_in_unit_disk();
+    pub fn defocus_disk_sample(&self, rng: &mut ThreadRng) -> Vector3<f64> {
+        let p = random_in_unit_disk(rng);
         self.center + (p.x * self.defocus_disk_u) + (p.y * self.defocus_disk_v)
     }
 
-    pub fn ray_color(&self, ray: &Ray, depth: u32, world: &Geometry) -> Vector3<f64> {
+    pub fn ray_color(
+        &self,
+        ray: &Ray,
+        depth: u32,
+        world: &Geometry,
+        rng: &mut ThreadRng,
+    ) -> Vector3<f64> {
         if depth == 0 {
             return Vector3::default();
         }
 
         let mut hit_record = HitRecord::default();
         let interval = Interval::new(0.001, f64::INFINITY);
 
-        if !world.hit(ray, &interval, &mut hit_record) {
+        if !world.hit(ray, &interval, &mut hit_record, rng) {
             return self.background;
         }
 
@@ -244,13 +274,13 @@ impl Camera {
 
         if !hit_record
             .material
-            .scatter(ray, &hit_record, &mut attenuation, &mut scattered)
+            .scatter(ray, &hit_record, &mut attenuation, &mut scattered, rng)
         {
             return color_from_emission;
         }
 
         let color_from_scatter =
-            attenuation.component_mul(&self.ray_color(&scattered, depth - 1, world));
+            attenuation.component_mul(&self.ray_color(&scattered, depth - 1, world, rng));
 
         color_from_emission + color_from_scatter
     }

src/geometry/bvh.rs

@@ -5,6 +5,7 @@ use crate::geometry::aabb::Aabb;
 use crate::geometry::empty::Empty;
 use crate::interval::Interval;
 use crate::ray::Ray;
+use rand::rngs::ThreadRng;
 
 #[derive(Debug, Clone)]
 pub struct BvhNode {
@@ -56,12 +57,18 @@ impl BvhNode {
 }
 
 impl Hittable for BvhNode {
-    fn hit(&self, r: &Ray, interval: &Interval, record: &mut HitRecord) -> bool {
+    fn hit(
+        &self,
+        r: &Ray,
+        interval: &Interval,
+        record: &mut HitRecord,
+        rng: &mut ThreadRng,
+    ) -> bool {
         if !self.bbox.hit(r, interval) {
             return false;
         }
 
-        let left_hit = self.left.hit(r, interval, record);
+        let left_hit = self.left.hit(r, interval, record, rng);
         let new_interval = match left_hit {
             true => &Interval {
                 min: interval.min,
@@ -70,7 +77,7 @@ impl Hittable for BvhNode {
             false => interval,
         };
 
-        let right_hit = self.right.hit(r, new_interval, record);
+        let right_hit = self.right.hit(r, new_interval, record, rng);
 
         left_hit || right_hit
     }

src/geometry/cube.rs

@@ -8,6 +8,7 @@ use crate::interval::Interval;
 use crate::material::Material;
 use crate::ray::Ray;
 use nalgebra::Vector3;
+use rand::rngs::ThreadRng;
 
 #[derive(Debug, Clone)]
 pub struct Cube {
@@ -42,8 +43,14 @@ impl Cube {
 }
 
 impl Hittable for Cube {
-    fn hit(&self, r: &Ray, interval: &Interval, record: &mut HitRecord) -> bool {
-        self.children.hit(r, interval, record)
+    fn hit(
+        &self,
+        r: &Ray,
+        interval: &Interval,
+        record: &mut HitRecord,
+        rng: &mut ThreadRng,
+    ) -> bool {
+        self.children.hit(r, interval, record, rng)
     }
 
     fn bounding_box(&self) -> Aabb {

src/geometry/empty.rs

@@ -4,6 +4,7 @@ use crate::geometry::Hittable;
 use crate::geometry::aabb::Aabb;
 use crate::interval::Interval;
 use crate::ray::Ray;
+use rand::rngs::ThreadRng;
 
 #[derive(Debug, Clone)]
 pub struct Empty {}
@@ -15,7 +16,7 @@ impl Empty {
 }
 
 impl Hittable for Empty {
-    fn hit(&self, _: &Ray, _: &Interval, _: &mut HitRecord) -> bool {
+    fn hit(&self, _: &Ray, _: &Interval, _: &mut HitRecord, _: &mut ThreadRng) -> bool {
         false
     }
 

src/geometry/mod.rs

@@ -24,9 +24,16 @@ use crate::material::lambertian::Lambertian;
 use crate::material::texture::SolidColor;
 use crate::ray::Ray;
 use nalgebra::Vector3;
+use rand::rngs::ThreadRng;
 
 pub trait Hittable {
-    fn hit(&self, r: &Ray, interval: &Interval, record: &mut HitRecord) -> bool;
+    fn hit(
+        &self,
+        r: &Ray,
+        interval: &Interval,
+        record: &mut HitRecord,
+        rng: &mut ThreadRng,
+    ) -> bool;
     fn bounding_box(&self) -> Aabb;
 }
 
@@ -43,16 +50,22 @@ pub enum Geometry {
 }
 
 impl Hittable for Geometry {
-    fn hit(&self, ray: &Ray, interval: &Interval, record: &mut HitRecord) -> bool {
+    fn hit(
+        &self,
+        ray: &Ray,
+        interval: &Interval,
+        record: &mut HitRecord,
+        rng: &mut ThreadRng,
+    ) -> bool {
         match self {
-            Geometry::Empty(geomtry) => geomtry.hit(ray, interval, record),
-            Geometry::Quad(geomtry) => geomtry.hit(ray, interval, record),
-            Geometry::Sphere(geomtry) => geomtry.hit(ray, interval, record),
-            Geometry::BvhNode(geomtry) => geomtry.hit(ray, interval, record),
-            Geometry::Cube(geomtry) => geomtry.hit(ray, interval, record),
-            Geometry::Translate(geomtry) => geomtry.hit(ray, interval, record),
-            Geometry::Rotate(geomtry) => geomtry.hit(ray, interval, record),
-            Geometry::Volume(geomtry) => geomtry.hit(ray, interval, record),
+            Geometry::Empty(geomtry) => geomtry.hit(ray, interval, record, rng),
+            Geometry::Quad(geomtry) => geomtry.hit(ray, interval, record, rng),
+            Geometry::Sphere(geomtry) => geomtry.hit(ray, interval, record, rng),
+            Geometry::BvhNode(geomtry) => geomtry.hit(ray, interval, record, rng),
+            Geometry::Cube(geomtry) => geomtry.hit(ray, interval, record, rng),
+            Geometry::Translate(geomtry) => geomtry.hit(ray, interval, record, rng),
+            Geometry::Rotate(geomtry) => geomtry.hit(ray, interval, record, rng),
+            Geometry::Volume(geomtry) => geomtry.hit(ray, interval, record, rng),
         }
     }
 

src/geometry/quad.rs

@@ -6,6 +6,7 @@ use crate::interval::Interval;
 use crate::material::Material;
 use crate::ray::Ray;
 use nalgebra::Vector3;
+use rand::rngs::ThreadRng;
 
 #[derive(Debug, Clone)]
 pub struct Quad {
@@ -63,7 +64,7 @@ fn is_interior(a: f64, b: f64, record: &mut HitRecord) -> bool {
 }
 
 impl Hittable for Quad {
-    fn hit(&self, r: &Ray, interval: &Interval, record: &mut HitRecord) -> bool {
+    fn hit(&self, r: &Ray, interval: &Interval, record: &mut HitRecord, _: &mut ThreadRng) -> bool {
         let denom = self.normal.dot(&r.direction);
         if denom.abs() < 1e-8 {
             return false;

src/geometry/rotate.rs

@@ -7,6 +7,7 @@ use crate::interval::Interval;
 use crate::ray::Ray;
 use nalgebra::Rotation3;
 use nalgebra::Vector3;
+use rand::rngs::ThreadRng;
 
 #[derive(Debug, Clone)]
 pub struct Rotate {
@@ -56,14 +57,20 @@ impl Rotate {
 }
 
 impl Hittable for Rotate {
-    fn hit(&self, r: &Ray, interval: &Interval, record: &mut HitRecord) -> bool {
+    fn hit(
+        &self,
+        r: &Ray,
+        interval: &Interval,
+        record: &mut HitRecord,
+        rng: &mut ThreadRng,
+    ) -> bool {
         // Transform the ray from world space to object space.
         let rotated_origin = self.rotation.inverse() * r.origin;
         let rotated_direction = self.rotation.inverse() * r.direction;
         let rotated_ray = Ray::new(rotated_origin, rotated_direction, r.time);
 
         // Determine whether an intersection exists in object space (and if so, where).
-        if !self.geometry.hit(&rotated_ray, interval, record) {
+        if !self.geometry.hit(&rotated_ray, interval, record, rng) {
             return false;
         }
 

src/geometry/sphere.rs

@@ -6,6 +6,7 @@ use crate::interval::Interval;
 use crate::material::Material;
 use crate::ray::Ray;
 use nalgebra::Vector3;
+use rand::rngs::ThreadRng;
 use std::f64::consts::PI;
 
 #[derive(Debug, Clone)]
@@ -55,7 +56,7 @@ impl Sphere {
 }
 
 impl Hittable for Sphere {
-    fn hit(&self, r: &Ray, interval: &Interval, record: &mut HitRecord) -> bool {
+    fn hit(&self, r: &Ray, interval: &Interval, record: &mut HitRecord, _: &mut ThreadRng) -> bool {
         let current_center = self.center.at(r.time);
         let oc = r.origin - current_center;