From 15d5513974fb2f0c6f82e933e97d8a3d05a30899 Mon Sep 17 00:00:00 2001
From: Adrien Vannson <adrien.vannson@protonmail.com>
Date: Tue, 19 Mar 2024 23:41:37 +0100
Subject: [PATCH] Stratified sampler

---
 src/renderers/monte_carlo_renderer.rs |  5 +-
 src/sampler.rs                        |  3 -
 src/samplers/independent_sampler.rs   |  4 +-
 src/samplers/mod.rs                   |  1 +
 src/samplers/stratified_sampler.rs    | 83 +++++++++++++++++++++++++++
 5 files changed, 90 insertions(+), 6 deletions(-)
 create mode 100644 src/samplers/stratified_sampler.rs

diff --git a/src/renderers/monte_carlo_renderer.rs b/src/renderers/monte_carlo_renderer.rs
index c10ea08..2a647c7 100644
--- a/src/renderers/monte_carlo_renderer.rs
+++ b/src/renderers/monte_carlo_renderer.rs
@@ -144,6 +144,7 @@ fn color<S: Sampler>(
     let mut sum = (0., 0., 0.);
 
     for _ in 0..nb_samples {
+        sampler.new_sample();
         let color = one_color(ray, scene, sampler.next2d(), ambient_occlusion);
 
         sum.0 += color.red;
@@ -183,7 +184,7 @@ impl<S: Sampler + 'static> Renderer for MonteCarloRenderer<S> {
         let mut tx_workers = Vec::new();
         let mut handles = Vec::new();
 
-        let workers_count = 32;
+        let workers_count = 4; // TODO choose wisely
 
         for worker_id in 0..workers_count {
             let scene = Arc::clone(&scene);
@@ -196,7 +197,7 @@ impl<S: Sampler + 'static> Renderer for MonteCarloRenderer<S> {
 
             handles.push(thread::spawn(move || {
                 while let Some(request) = rx_worker.recv().unwrap() {
-                    sampler.prepare(1, 1, iterations_per_pixel as usize);
+                    sampler.prepare(0, 1, iterations_per_pixel as usize);
 
                     let (x, y) = (request.x, request.y);
 
diff --git a/src/sampler.rs b/src/sampler.rs
index 462e869..f4093f2 100644
--- a/src/sampler.rs
+++ b/src/sampler.rs
@@ -1,7 +1,4 @@
 pub trait Sampler: Send {
-    /// Create a new instance of the sampler.
-    fn new(seed: u64) -> Self;
-
     /// Initializes the sampler. It will need to generate a fixed number of samples.
     /// Each sample will be composed of at most a fixed number of unique random values and a fixed
     /// number of couples of random values.
diff --git a/src/samplers/independent_sampler.rs b/src/samplers/independent_sampler.rs
index b3359ba..c0f2c57 100644
--- a/src/samplers/independent_sampler.rs
+++ b/src/samplers/independent_sampler.rs
@@ -5,13 +5,15 @@ pub struct IndependentSampler {
     rng: Rng,
 }
 
-impl Sampler for IndependentSampler {
+impl IndependentSampler {
     fn new(seed: u64) -> Self {
         IndependentSampler {
             rng: fastrand::Rng::with_seed(seed),
         }
     }
+}
 
+impl Sampler for IndependentSampler {
     fn prepare(&mut self, nb_1d: usize, nb_2d: usize, nb_samples: usize) {}
 
     fn new_sample(&mut self) {}
diff --git a/src/samplers/mod.rs b/src/samplers/mod.rs
index beaf27d..174f57f 100644
--- a/src/samplers/mod.rs
+++ b/src/samplers/mod.rs
@@ -1 +1,2 @@
 pub mod independent_sampler;
+pub mod stratified_sampler;
diff --git a/src/samplers/stratified_sampler.rs b/src/samplers/stratified_sampler.rs
new file mode 100644
index 0000000..e9b42c7
--- /dev/null
+++ b/src/samplers/stratified_sampler.rs
@@ -0,0 +1,83 @@
+use crate::sampler::Sampler;
+use fastrand::Rng;
+
+pub struct StratifiedSampler {
+    rng: Rng,
+    with_jittering: bool,
+    samples_2d: Vec<Vec<[f64; 2]>>, // (nb_samples, nb_2d)
+    current_sample: usize,
+    current_2d_dim: usize,
+}
+
+impl StratifiedSampler {
+    fn new(seed: u64) -> Self {
+        StratifiedSampler {
+            rng: fastrand::Rng::with_seed(seed),
+            with_jittering: true,
+            samples_2d: Vec::new(),
+            current_sample: 0,
+            current_2d_dim: 0,
+        }
+    }
+
+    fn new_without_jittering(seed: u64) -> Self {
+        StratifiedSampler {
+            rng: fastrand::Rng::with_seed(seed),
+            with_jittering: false,
+            samples_2d: Vec::new(),
+            current_sample: 0,
+            current_2d_dim: 0,
+        }
+    }
+}
+
+impl Sampler for StratifiedSampler {
+    fn prepare(&mut self, nb_1d: usize, nb_2d: usize, nb_samples: usize) {
+        // TODO 1D
+        // TODO several 2d dimensions
+        assert_eq!(nb_2d, 1);
+
+        let root = (nb_samples as f64).sqrt() as usize;
+        assert_eq!(root * root, nb_samples);
+
+        let mut samples_2d = Vec::new();
+
+        for i in 0..root {
+            for j in 0..root {
+                // Not 0.5 to prevent rays from being parallel to the walls
+                // TODO 0.5
+                let (dx, dy) = if self.with_jittering {
+                    (self.rng.f64(), self.rng.f64())
+                } else {
+                    (0.501, 0.501)
+                };
+
+                samples_2d.push(vec![[
+                    (i as f64 + dx) / root as f64,
+                    (j as f64 + dy) / root as f64,
+                ]]);
+            }
+        }
+
+        self.samples_2d = samples_2d;
+        self.current_sample = 0;
+        self.current_2d_dim = 0;
+    }
+
+    fn new_sample(&mut self) {
+        if self.current_2d_dim != 0 {
+            // TODO and 1D dimension
+            self.current_sample += 1;
+            self.current_2d_dim = 0;
+        }
+    }
+
+    fn next1d(&mut self) -> f64 {
+        unimplemented!();
+    }
+
+    fn next2d(&mut self) -> [f64; 2] {
+        self.current_2d_dim += 1;
+        self.samples_2d[self.current_sample][self.current_2d_dim - 1]
+    }
+}