Applying fixes to Sobol samplers

shared/src/core/sampler.rs

@@ -2,15 +2,15 @@ use crate::core::filter::FilterTrait;
 use crate::core::geometry::{Bounds2f, Point2f, Point2i, Vector2f};
 use crate::core::pbrt::{Float, ONE_MINUS_EPSILON, PI, PI_OVER_2, PI_OVER_4};
 use crate::utils::math::{
-    clamp, encode_morton_2, inverse_radical_inverse, lerp, log2_int,
+    BinaryPermuteScrambler, DigitPermutation, FastOwenScrambler, NoRandomizer, OwenScrambler,
+    PRIME_TABLE_SIZE, Scrambler, clamp, encode_morton_2, inverse_radical_inverse, lerp, log2_int,
     owen_scrambled_radical_inverse, permutation_element, radical_inverse, round_up_pow2,
-    scrambled_radical_inverse, sobol_interval_to_index, sobol_sample, BinaryPermuteScrambler,
+    scrambled_radical_inverse, sobol_interval_to_index, sobol_sample,
-    DigitPermutation, FastOwenScrambler, NoRandomizer, OwenScrambler, Scrambler, PRIME_TABLE_SIZE,
 };
 use crate::utils::rng::Rng;
 use crate::utils::sobol::N_SOBOL_DIMENSIONS;
 use crate::utils::{hash::*, sobol};
-use crate::{gvec, GVec, Ptr};
+use crate::{GVec, Ptr, gvec};
 use enum_dispatch::enum_dispatch;
 
 #[repr(C)]
@@ -530,6 +530,18 @@ pub struct ZSobolSampler {
     dim: u32,
 }
 
+/// pbrt writes `0x55555555u * dimension` -- a 32-bit unsigned product that wraps.
+#[inline]
+fn scramble_seed(dim: u32) -> u64 {
+    0x5555_5555u32.wrapping_mul(dim) as u64
+}
+
+/// pbrt's `Hash(dimension, seed)`: both are `int`, so exactly 8 packed bytes.
+#[inline]
+fn dim_seed_hash(dim: u32, seed: u64) -> u64 {
+    hash_buffer(&[dim, seed as u32], 0)
+}
+
 impl ZSobolSampler {
     pub fn new(
         samples_per_pixel: i32,
@@ -537,10 +549,12 @@ impl ZSobolSampler {
         randomize: RandomizeStrategy,
         seed: Option<u64>,
     ) -> Self {
-        let log2_samples_per_pixel = log2_int(samples_per_pixel as Float) as u32;
+        // pbrt calls the integer Log2Int overload; the float one disagrees for
+        // non-power-of-two sample counts.
+        let log2_samples_per_pixel = (samples_per_pixel.max(1) as u32).ilog2();
         let res = round_up_pow2(full_resolution.x().max(full_resolution.y()));
         let log4_samples_per_pixel = log2_samples_per_pixel.div_ceil(2);
-        let n_base4_digits = log2_int(res as Float) as u32 + log4_samples_per_pixel as u32;
+        let n_base4_digits = (res.max(1) as u32).ilog2() + log4_samples_per_pixel;
         Self {
             randomize,
             seed: seed.unwrap_or(0),
@@ -590,7 +604,7 @@ impl ZSobolSampler {
 
             let higher_digits = self.morton_index >> (digit_shift + 2);
 
-            let mix_input = higher_digits ^ (0x55555555 * self.dim as u64);
+            let mix_input = higher_digits ^ scramble_seed(self.dim);
             let p = (mix_bits(mix_input) >> 24) % 24;
 
             digit = PERMUTATIONS[p as usize][digit as usize] as u64;
@@ -599,8 +613,9 @@ impl ZSobolSampler {
         }
 
         if pow2_samples {
-            let lsb = self.morton_index & 1;
+            let digit = self.morton_index & 1;
-            sample_index |= lsb;
+            sample_index |=
+                digit ^ (mix_bits((self.morton_index >> 1) ^ scramble_seed(self.dim)) & 1);
         }
 
         sample_index
@@ -609,8 +624,9 @@ impl ZSobolSampler {
 
 impl SamplerTrait for ZSobolSampler {
     fn samples_per_pixel(&self) -> i32 {
-        todo!()
+        1 << self.log2_samples_per_pixel
     }
+
     fn start_pixel_sample(&mut self, p: Point2i, sample_index: i32, dim: Option<u32>) {
         self.dim = dim.unwrap_or(0);
         self.morton_index = (encode_morton_2(p.x() as u32, p.y() as u32)
@@ -620,31 +636,25 @@ impl SamplerTrait for ZSobolSampler {
 
     fn get1d(&mut self) -> Float {
         let sample_index = self.get_sample_index();
-        let hash_input = [self.dim as u64, self.seed];
-        let hash = hash_buffer(&hash_input, 0) as u32;
         self.dim += 1;
-        if self.randomize == RandomizeStrategy::None {
+        let hash = dim_seed_hash(self.dim, self.seed) as u32;
-            return sobol_sample(sample_index, self.dim, NoRandomizer);
+        // Always Sobol dimension 0 -- decorrelation comes from the hash.
-        }
         match self.randomize {
+            RandomizeStrategy::None => sobol_sample(sample_index, 0, NoRandomizer),
             RandomizeStrategy::PermuteDigits => {
-                sobol_sample(sample_index, self.dim, BinaryPermuteScrambler::new(hash))
+                sobol_sample(sample_index, 0, BinaryPermuteScrambler::new(hash))
             }
             RandomizeStrategy::FastOwen => {
-                sobol_sample(sample_index, self.dim, FastOwenScrambler::new(hash))
+                sobol_sample(sample_index, 0, FastOwenScrambler::new(hash))
             }
-            RandomizeStrategy::Owen => {
+            RandomizeStrategy::Owen => sobol_sample(sample_index, 0, OwenScrambler::new(hash)),
-                sobol_sample(sample_index, self.dim, OwenScrambler::new(hash))
-            }
-            RandomizeStrategy::None => unreachable!(),
         }
     }
 
     fn get2d(&mut self) -> Point2f {
         let sample_index = self.get_sample_index();
         self.dim += 2;
-        let hash_input = [self.dim as u64, self.seed];
+        let hash = dim_seed_hash(self.dim, self.seed);
-        let hash = hash_buffer(&hash_input, 0);
         let sample_hash = [hash as u32, (hash >> 32) as u32];
         if self.randomize == RandomizeStrategy::None {
             return Point2f::new(
@@ -675,16 +685,92 @@ impl SamplerTrait for ZSobolSampler {
 }
 
 #[derive(Default, Debug, Clone)]
-pub struct MLTSampler;
+struct PrimarySample {
+    value: Float,
+    last_mod_iteration: i64,
+    value_backup: Float,
+    mod_backup: i64,
+}
+
+impl PrimarySample {
+    fn backup(&mut self) {
+        self.value_backup = self.value;
+        self.mod_backup = self.last_mod_iteration;
+    }
+
+    fn restore(&mut self) {
+        self.value = self.value_backup;
+        self.last_mod_iteration = self.mod_backup;
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct MLTSampler {
+    mutations_per_pixel: i32,
+    rng: Rng,
+    sigma: Float,
+    large_step_prob: Float,
+    stream_count: i32,
+    x: GVec<PrimarySample>,
+    current_iter: i64,
+    large_step: bool,
+    last_large_step_iter: i64,
+    stream_ind: i32,
+    sample_ind: i32,
+    seed: u64,
+}
+
+impl MLTSampler {
+    pub fn new(
+        mutations_per_pixel: i32,
+        rng_seq_ind: i32,
+        sigma: Float,
+        large_step_prob: Float,
+        stream_count: i32,
+        seed: u64,
+    ) -> Self {
+        Self {
+            mutations_per_pixel,
+            rng: Rng::new(mix_bits(rng_seq_ind.try_into().unwrap()) ^ mix_bits(seed)),
+            seed,
+            sigma,
+            large_step_prob,
+            stream_count,
+            x: gvec(),
+            current_iter: 0,
+            large_step: true,
+            last_large_step_iter: 0,
+            stream_ind: 0,
+            sample_ind: 0,
+        }
+    }
+
+    pub fn get_next_index(&mut self) -> i32 {
+        self.sample_ind += 1;
+        self.stream_ind + self.stream_count * self.sample_ind
+    }
+}
+
 impl SamplerTrait for MLTSampler {
     fn samples_per_pixel(&self) -> i32 {
-        todo!()
+        self.mutations_per_pixel
     }
-    fn start_pixel_sample(&mut self, _p: Point2i, _sample_index: i32, _dim: Option<u32>) {
+
-        todo!()
+    fn start_pixel_sample(&mut self, p: Point2i, sample_index: i32, dim: Option<u32>) {
+        let hash_input = [p.x() as u64, p.y() as u64, self.seed];
+        let sequence_index = hash_buffer(&hash_input, 0);
+        self.rng.set_sequence(sequence_index);
+        self.rng
+            .advance((sample_index as u64) * 65536 + (dim.unwrap_or(0) as u64));
     }
+
     fn get1d(&mut self) -> Float {
-        todo!()
+        #[cfg(not(any(feature = "cuda", feature = "vulkan")))]
+        {
+            return 0.;
+        }
+
+        let ind = self.get_next_index();
     }
     fn get2d(&mut self) -> Point2f {
         todo!()