diff --git a/shared/src/core/sampler.rs b/shared/src/core/sampler.rs index 293aeff..2efe4fd 100644 --- a/shared/src/core/sampler.rs +++ b/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, - DigitPermutation, FastOwenScrambler, NoRandomizer, OwenScrambler, Scrambler, PRIME_TABLE_SIZE, + scrambled_radical_inverse, sobol_interval_to_index, sobol_sample, }; 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, ) -> 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; - sample_index |= lsb; + let digit = self.morton_index & 1; + 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) { 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 { - return sobol_sample(sample_index, self.dim, NoRandomizer); - } + let hash = dim_seed_hash(self.dim, self.seed) as u32; + // 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 => { - sobol_sample(sample_index, self.dim, OwenScrambler::new(hash)) - } - RandomizeStrategy::None => unreachable!(), + RandomizeStrategy::Owen => sobol_sample(sample_index, 0, OwenScrambler::new(hash)), } } 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 = hash_buffer(&hash_input, 0); + let hash = dim_seed_hash(self.dim, self.seed); 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, + 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) { - todo!() + + fn start_pixel_sample(&mut self, p: Point2i, sample_index: i32, dim: Option) { + 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!()