halton sampling

Cargo.lock

@@ -41,6 +41,12 @@ dependencies = [
  "bytemuck",
 ]
 
+[[package]]
+name = "hamming"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "65043da274378d68241eb9a8f8f8aa54e349136f7b8e12f63e3ef44043cc30e1"
+
 [[package]]
 name = "heck"
 version = "0.5.0"
@@ -66,6 +72,7 @@ dependencies = [
  "bytemuck",
  "enum_dispatch",
  "glam",
+ "primal",
 ]
 
 [[package]]
@@ -86,6 +93,24 @@ dependencies = [
  "autocfg",
 ]
 
+[[package]]
+name = "num-integer"
+version = "0.1.46"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-traits"
+version = "0.2.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
+dependencies = [
+ "autocfg",
+]
+
 [[package]]
 name = "once_cell"
 version = "1.21.3"
@@ -98,6 +123,52 @@ version = "1.11.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "350e9b48cbc6b0e028b0473b114454c6316e57336ee184ceab6e53f72c178b3e"
 
+[[package]]
+name = "primal"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b1e5f354948532e6017fc91f9a5ff5ba1be0dabd3a0c9e9c417969cd4c1ad6e8"
+dependencies = [
+ "primal-check",
+ "primal-estimate",
+ "primal-sieve",
+]
+
+[[package]]
+name = "primal-bit"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "252429dbb8aeacc3233df500dc3a6a367bf28eb3a711272884d7540a7b636055"
+dependencies = [
+ "hamming",
+]
+
+[[package]]
+name = "primal-check"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc0d895b311e3af9902528fbb8f928688abbd95872819320517cc24ca6b2bd08"
+dependencies = [
+ "num-integer",
+]
+
+[[package]]
+name = "primal-estimate"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a432100a0b3a61085e75b5f89e9f42de73c0acb7dea5038b893697918105d822"
+
+[[package]]
+name = "primal-sieve"
+version = "0.3.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8e982796d82203351983d3602a8d6372d1d7894e86960047ba0d4b7426a5edd3"
+dependencies = [
+ "primal-bit",
+ "primal-estimate",
+ "smallvec",
+]
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.95"
@@ -179,6 +250,12 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "smallvec"
+version = "1.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8917285742e9f3e1683f0a9c4e6b57960b7314d0b08d30d1ecd426713ee2eee9"
+
 [[package]]
 name = "syn"
 version = "2.0.100"

core/Cargo.toml

@@ -14,3 +14,6 @@ glam = "0.30.2"
 [features]
 fast-math = ["glam/fast-math"]
 bytemuck = ["dep:bytemuck", "glam/bytemuck"]
+
+[build-dependencies]
+primal = "0.3.3"

core/build.rs

@@ -0,0 +1,17 @@
+use std::{env, fs::File, io::Write, path::PathBuf};
+
+fn main() {
+    let dest: PathBuf = env::var("OUT_DIR").unwrap().into();
+    let mut file = File::create(dest.join("primes.rs")).unwrap();
+
+    write_primes(&mut file).unwrap();
+}
+
+fn write_primes(file: &mut File) -> std::io::Result<()> {
+    write!(file, "[")?;
+    let sieve = primal::Sieve::new(7919);
+    for p in sieve.primes_from(2).take_while(|x| *x <= 7919) {
+        write!(file, "{p},")?;
+    }
+    write!(file, "]")
+}

core/src/geometry.rs

@@ -1,6 +1,6 @@
 use std::{
     fmt,
-    ops::{Add, Deref, Div, Mul, Sub},
+    ops::{Add, AddAssign, Deref, Div, DivAssign, Mul, Sub},
 };
 
 pub mod shapes;
@@ -94,6 +94,12 @@ impl Div<f32> for Vector {
     }
 }
 
+impl DivAssign<f32> for Vector {
+    fn div_assign(&mut self, rhs: f32) {
+        *self = *self / rhs;
+    }
+}
+
 impl Add for Vector {
     type Output = Self;
 
@@ -102,6 +108,12 @@ impl Add for Vector {
     }
 }
 
+impl AddAssign for Vector {
+    fn add_assign(&mut self, rhs: Self) {
+        *self = *self + rhs;
+    }
+}
+
 #[repr(transparent)]
 #[derive(Clone, Copy)]
 pub struct Point(glam::Vec3A);

core/src/lib.rs

@@ -5,6 +5,7 @@ pub mod color;
 
 mod camera;
 mod geometry;
+mod sampling;
 
 use core::f32;
 
@@ -16,6 +17,7 @@ use geometry::{
     Ray,
     shapes::{Hit, Hittable},
 };
+use sampling::Sampler;
 
 pub struct Scene {
     camera: Camera,
@@ -27,23 +29,36 @@ impl Scene {
         Self { camera, contents }
     }
 
-    pub fn render(&self, width: usize, height: usize) -> Box<[Rgba]> {
-        let ray = self.camera.ray(0.0, 0.0);
-        dbg!(&ray);
-        dbg!(self.intersect(ray));
-        let mut data = Box::new_uninit_slice(width * height);
+    pub fn render(&self, width: u32, height: u32, samples: u32) -> Box<[Rgba]> {
+        let mut sampler = Sampler::new(width, height);
+        let mut data = Box::new_uninit_slice(width as usize * height as usize);
         let w = width as f32;
         let h = height as f32;
         let aspect_ratio = h / w;
 
         for (i, pixel) in data.iter_mut().enumerate() {
-            let x = ((i % width) as f32 * 2.0 / w) - 1.0;
-            let y = (1.0 - (i / width) as f32 * 2.0 / h) * aspect_ratio;
+            let bx = i as u32 % width;
+            let by = i as u32 / width;
 
-            let rgb = self
-                .intersect(self.camera.ray(x, y))
-                .map(|h| h.normal.normalize() * 0.5 + Vector::new(0.5, 0.5, 0.5))
-                .unwrap_or(Vector::ZERO);
+            let mut rgb = Vector::ZERO;
+            // TODO: there's a way to make this more idiomatic with iterators
+            sampler.start_pixel(bx, by);
+            for _ in 0..samples {
+                let (offset_x, offset_y) = sampler.get_pixel_2d();
+
+                let x = bx as f32 + offset_x;
+                let y = by as f32 + offset_y;
+
+                let cx = 2.0 * x / w - 1.0;
+                let cy = (1.0 - 2.0 * y / h) * aspect_ratio;
+
+                if let Some(h) = self.intersect(self.camera.ray(cx, cy)) {
+                    rgb += h.normal.normalize() * 0.5 + Vector::new(0.5, 0.5, 0.5);
+                }
+                sampler.next_sample();
+            }
+
+            rgb /= samples as f32;
 
             pixel.write(Rgba::new(rgb.x, rgb.y, rgb.z, 1.0));
         }

core/src/sampling.rs

@@ -0,0 +1,207 @@
+//! Halton sampling
+
+const PRIMES: [u32; 1000] = include!(concat!(env!("OUT_DIR"), "/primes.rs"));
+const MAX_HALTON_RESOLUTION: u32 = 128;
+
+pub struct Sampler {
+    x_scale: u8,
+    x_exponent: u8,
+    x_inverse: u8,
+    y_scale: u8,
+    y_exponent: u8,
+    y_inverse: u8,
+    index: u32,
+    dimension: u16,
+}
+
+impl Sampler {
+    pub fn new(width: u32, height: u32) -> Self {
+        let mut x_scale = 1u8;
+        let mut x_exponent = 0u8;
+        let x_res = u32::min(width, MAX_HALTON_RESOLUTION) as u8;
+        while x_scale < x_res {
+            x_scale *= 2;
+            x_exponent += 1;
+        }
+
+        let mut y_scale = 1u8;
+        let mut y_exponent = 0u8;
+        let y_res = u32::min(height, MAX_HALTON_RESOLUTION) as u8;
+        while y_scale < y_res {
+            y_scale *= 3;
+            y_exponent += 1;
+        }
+
+        let x_inverse = mod_inverse(y_scale, x_scale);
+        let y_inverse = mod_inverse(x_scale, y_scale);
+
+        Self {
+            dimension: 2,
+            index: 0,
+            x_scale,
+            x_exponent,
+            x_inverse,
+            y_scale,
+            y_exponent,
+            y_inverse,
+        }
+    }
+
+    pub fn start_pixel(&mut self, x: u32, y: u32) {
+        let mut index = 0;
+        let sample_stride = self.y_scale as u32 * self.x_scale as u32;
+
+        if sample_stride > 1 {
+            let mx = x % MAX_HALTON_RESOLUTION;
+            let my = y % MAX_HALTON_RESOLUTION;
+
+            let x_offset = inverse_radical_inverse(mx, 2, self.x_exponent as u32);
+            let y_offset = inverse_radical_inverse(my, 3, self.y_exponent as u32);
+
+            index += x_offset * (sample_stride / self.x_scale as u32) * self.x_inverse as u32;
+            index += y_offset * (sample_stride / self.y_scale as u32) * self.y_inverse as u32;
+
+            index %= sample_stride;
+        }
+
+        self.index = index;
+    }
+
+    pub fn next_sample(&mut self) {
+        self.index += self.x_scale as u32 * self.y_scale as u32;
+    }
+
+    pub fn get_pixel_2d(&self) -> (f32, f32) {
+        (
+            radical_inverse(0, self.index >> self.x_exponent),
+            radical_inverse(1, self.index / self.y_scale as u32),
+        )
+    }
+
+    pub fn get_1d(&mut self) -> f32 {
+        if self.dimension as usize >= PRIMES.len() {
+            self.dimension = 2;
+        }
+        let dim = self.dimension as u64;
+        self.dimension += 1;
+        scrambled_radical_inverse(dim as usize, self.index, mix_bits(1 + (dim << 4)))
+    }
+
+    pub fn get_2d(&mut self) -> (f32, f32) {
+        if self.dimension as usize + 1 >= PRIMES.len() {
+            self.dimension = 2;
+        }
+        let mut dim = self.dimension as u64;
+        self.dimension += 2;
+        let x = scrambled_radical_inverse(dim as usize, self.index, mix_bits(1 + (dim << 4)));
+        dim += 1;
+        let y = scrambled_radical_inverse(dim as usize, self.index, mix_bits(1 + (dim << 4)));
+        (x, y)
+    }
+}
+
+fn egcd(a: i16, b: i16) -> (i16, i16) {
+    if a == 0 {
+        (0, 1)
+    } else {
+        let (x, y) = egcd(b % a, a);
+        (y - (b / a) * x, x)
+    }
+}
+
+fn mod_inverse(a: u8, m: u8) -> u8 {
+    let m = m as i16;
+    let (x, _) = egcd(a as i16, m);
+    ((x % m + m) % m) as u8
+}
+
+fn radical_inverse(base_index: usize, mut i: u32) -> f32 {
+    let base = PRIMES[base_index];
+    let inv_base = 1.0 / base as f32;
+    let mut inv_base_m = 1.0;
+    let mut reversed = 0;
+
+    while i > 0 {
+        let next = i / base;
+        let digit = i - next * base;
+        reversed = reversed * base + digit;
+        inv_base_m *= inv_base;
+        i = next
+    }
+
+    f32::min(reversed as f32 * inv_base_m, 1.0 - f32::EPSILON)
+}
+
+fn inverse_radical_inverse(mut inverse: u32, base: u32, digits: u32) -> u32 {
+    let mut i = 0;
+    for _ in 0..digits {
+        let digit = inverse % base;
+        inverse /= base;
+        i = i * base + digit;
+    }
+    i
+}
+
+fn scrambled_radical_inverse(base_index: usize, mut i: u32, hash: u64) -> f32 {
+    let base = PRIMES[base_index];
+    let inv_base = 1.0 / base as f32;
+    let mut inv_base_m = 1.0;
+    let mut reversed = 0;
+
+    while 1.0 - inv_base_m < 1.0 {
+        let next = i / base;
+        let digit_value = i - next * base;
+        let digit_hash = mix_bits(hash ^ digit_value as u64);
+        let digit = permutation_element(digit_value, base, digit_hash as u32);
+        reversed = reversed * base + digit;
+        inv_base_m *= inv_base;
+        i = next;
+    }
+
+    f32::min(reversed as f32 * inv_base_m, 1.0 - f32::EPSILON)
+}
+
+// magic. Taken from pbrt/util/hash.h in pbrt
+#[inline(always)]
+fn mix_bits(mut val: u64) -> u64 {
+    val ^= val >> 31;
+    val *= 0x7fb5d329728ea185;
+    val ^= val >> 27;
+    val *= 0x81dadef4bc2dd44d;
+    val ^ (val >> 33)
+}
+
+// more magic. Taken from pbrt/util/math.h in pbrt
+#[inline(always)]
+fn permutation_element(mut i: u32, n: u32, seed: u32) -> u32 {
+    let mut w = n - 1;
+    w |= w >> 1;
+    w |= w >> 2;
+    w |= w >> 4;
+    w |= w >> 8;
+    w |= w >> 16;
+    loop {
+        i ^= seed;
+        i *= 0xe170893d;
+        i ^= seed >> 16;
+        i ^= (i & w) >> 4;
+        i ^= seed >> 8;
+        i *= 0x0929eb3f;
+        i ^= seed >> 23;
+        i ^= (i & w) >> 1;
+        i *= 1 | seed >> 27;
+        i *= 0x6935fa69;
+        i ^= (i & w) >> 11;
+        i *= 0x74dcb303;
+        i ^= (i & w) >> 2;
+        i *= 0x9e501cc3;
+        i ^= (i & w) >> 2;
+        i *= 0xc860a3df;
+        i &= w;
+        i ^= i >> 5;
+        if i < n {
+            break;
+        }
+    }
+    (i + seed) % n
+}

examples/basic.py

@@ -4,6 +4,7 @@ from mechthild import Scene, Shape, Camera
 
 WIDTH = 800
 HEIGHT = 600
+SAMPLES = 20
 
 s1 = Shape.sphere((0.75, 1.0, 1.0), 1.0)
 s2 = Shape.sphere((-1.5, 2.0, -2.0), 2.0)
@@ -13,7 +14,7 @@ camera = Camera((0.0, 1.5, 5.0), (0.0, 1.0, 0.0), (0.0, 1.0, 0.0), 80.0)
 
 scene = Scene(camera, [s1, s2, ground])
 
-render = scene.render(WIDTH, HEIGHT)
+render = scene.render(WIDTH, HEIGHT, SAMPLES)
 rgba = bytes([int(v * 255) for v in render])
 
 image = Image.frombuffer('RGBA', (WIDTH, HEIGHT), rgba)

py/src/lib.rs

@@ -55,14 +55,13 @@ impl PyScene {
         Self(Scene::new(camera.0.clone(), contents))
     }
 
-    pub fn render(&self, width: usize, height: usize) -> Vec<f32> {
-        let result = self.0.render(width, height);
+    pub fn render(&self, width: u32, height: u32, samples: u32) -> Vec<f32> {
+        let result = self.0.render(width, height, samples);
 
         cast_slice_box::<Rgba, f32>(result).into_vec()
     }
 }
 
-/// A Python module implemented in Rust.
 #[pymodule]
 fn mechthild(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<PyCamera>()?;