Add a color coordinate to the chaos game

Also refactors the gasket example to avoid using shader entry points; the entry point signatures are not stable, and I don't want to keep refactoring the gasket while adding more advanced color/image features.

enkou-shaders/examples/gasket.rs

@@ -1,36 +1,39 @@
 use anyhow::{Context, Result};
 use enkou_shaders::Coefficients2;
 use enkou_shaders::camera::Camera;
-use enkou_shaders::camera::entry::main_camera;
+use enkou_shaders::chaos_game::ChaosGame;
-use enkou_shaders::chaos_game::entry::main_chaos_game;
 use enkou_shaders::transform::Transform;
 use enkou_shaders::variation::Variation;
-use glam::{Affine2, IVec2, UVec2, Vec2, uvec2};
+use glam::{Affine2, UVec2, Vec2, uvec2, vec2};
 use image::{GrayImage, Luma};
+use rand::SeedableRng;
+use rand_xoshiro::Xoshiro256StarStar;
 use std::mem;
 use std::process::Command;
 use tempfile::NamedTempFile;
 
-const ITERATIONS_DISCARD: u32 = 20;
+const ITERATIONS_DISCARD: usize = 20;
-const ITERATIONS: u32 = 50_000;
+const ITERATIONS: usize = 50_000;
 const IMAGE_DIMENSION: UVec2 = uvec2(600, 600);
 
 pub fn main() -> Result<()> {
+    let mut rng = Xoshiro256StarStar::from_seed([4u8; 32]);
+
     let transforms = [
         {
             // F_0: (x / 2, y / 2)
             let coefficients = Affine2::from_coefficients(0.5, 0.0, 0.0, 0.0, 0.5, 0.0);
-            Transform::new(coefficients, Affine2::IDENTITY, uvec2(0, 1))
+            Transform::new(coefficients, Affine2::IDENTITY, uvec2(0, 1), vec2(0.0, 1.0))
         },
         {
             // F_1: ((x + 1) / 2, y / 2)
             let coefficients = Affine2::from_coefficients(0.5, 0.0, 0.5, 0.0, 0.5, 0.0);
-            Transform::new(coefficients, Affine2::IDENTITY, uvec2(0, 1))
+            Transform::new(coefficients, Affine2::IDENTITY, uvec2(0, 1), vec2(0.0, 1.0))
         },
         {
             // F_2: (x / 2, (y + 1) / 2)
             let coefficients = Affine2::from_coefficients(0.5, 0.0, 0.0, 0.0, 0.5, 0.5);
-            Transform::new(coefficients, Affine2::IDENTITY, uvec2(0, 1))
+            Transform::new(coefficients, Affine2::IDENTITY, uvec2(0, 1), vec2(0.0, 1.0))
         },
     ];
 
@@ -38,18 +41,6 @@ pub fn main() -> Result<()> {
 
     let variations = [Variation::IDENTITY];
 
-    let mut output_points_ifs = Vec::new();
-    output_points_ifs.resize(ITERATIONS as usize, Vec2::ZERO);
-
-    main_chaos_game(
-        ITERATIONS_DISCARD,
-        &[4u8; 32],
-        &transforms,
-        &weights,
-        &variations,
-        &mut output_points_ifs,
-    );
-
     // The gasket is defined on the range [0, 1] for both X and Y
     let camera = Camera::new(
         IMAGE_DIMENSION,
@@ -59,20 +50,15 @@ pub fn main() -> Result<()> {
         IMAGE_DIMENSION.as_vec2(),
     );
 
-    let mut output_points_pixel = Vec::new();
-    output_points_pixel.resize(ITERATIONS as usize, IVec2::ZERO);
-
-    main_camera(&camera, &output_points_ifs, &mut output_points_pixel);
-
     let mut image = GrayImage::new(IMAGE_DIMENSION.x, IMAGE_DIMENSION.y);
-    let dimensions = image.dimensions();
+
-    output_points_pixel
+    let chaos_game = ChaosGame::new(&mut rng, &transforms, &weights, &variations);
-        .iter()
+
-        .skip_while(|p| {
+    chaos_game
-            p.x < 0 || (p.x as u32) > dimensions.0 || p.y < 0 || (p.y as u32) > dimensions.1
+        .skip(ITERATIONS_DISCARD)
-        })
+        .take(ITERATIONS)
-        .map(|p| (p.x as u32, p.y as u32))
+        .filter_map(|(point_ifs, _)| camera.transform_point_to_image(point_ifs))
-        .for_each(|(x, y)| image.put_pixel(x, y, Luma([255u8])));
+        .for_each(|point_pixel| image.put_pixel(point_pixel.x, point_pixel.y, Luma([255])));
 
     let temp = NamedTempFile::with_suffix(".png").context("Unable to create file for image")?;
     image.save(temp.path()).context("Unable to save image")?;

enkou-shaders/src/chaos_game.rs

@@ -12,7 +12,6 @@
 //!
 //! This algorithm is also known as the ["chaos game"](https://en.wikipedia.org/wiki/Chaos_game),
 //! and it forms the basic system for producing images.
-
 use crate::transform::Transform;
 use crate::variation::Variation;
 use rand::distr::{Distribution, StandardUniform};
@@ -36,11 +35,12 @@ impl Distribution<f32> for BiUnit {
 /// * `weights` - Weights are assumed to be normalized; adding all elements together should return the value 1
 pub fn step_chaos_game<R: Rng>(
     point: Vec2,
+    color: f32,
     rng: &mut R,
     transforms: &[Transform],
     weights: &[f32],
     variations: &[Variation],
-) -> (Vec2, u32) {
+) -> (Vec2, f32, u32) {
     let mut choice_weight = rng.sample::<f32, _>(StandardUniform);
     let mut transform_index: u32 = 0;
 
@@ -53,8 +53,11 @@ pub fn step_chaos_game<R: Rng>(
         transform_index += 1;
     }
 
+    let transform = transforms[transform_index as usize];
+
     (
-        transforms[transform_index as usize].transform_point(rng, variations, point),
+        transform.transform_point(rng, variations, point),
+        transform.transform_color(color),
         transform_index,
     )
 }
@@ -65,6 +68,7 @@ pub fn step_chaos_game<R: Rng>(
 /// New points in the chaos game are produced by iterating on the chaos game.
 pub struct ChaosGame<'a, R: Rng> {
     current_point: Vec2,
+    current_color: f32,
     rng: &'a mut R,
     transforms: &'a [Transform],
     weights: &'a [f32],
@@ -79,9 +83,14 @@ impl<'a, R: Rng> ChaosGame<'a, R> {
         weights: &'a [f32],
         variations: &'a [Variation],
     ) -> Self {
-        let current_point = vec2(rng.sample(BiUnit), rng.sample(BiUnit));
+        let current_point = vec2(
+            rng.sample(BiUnit),
+            rng.sample(BiUnit),
+        );
+        let current_color = rng.sample(StandardUniform);
         ChaosGame {
             current_point,
+            current_color,
             rng,
             transforms,
             weights,
@@ -91,54 +100,20 @@ impl<'a, R: Rng> ChaosGame<'a, R> {
 }
 
 impl<'a, R: Rng> Iterator for ChaosGame<'a, R> {
-    type Item = Vec2;
+    type Item = (Vec2, f32);
 
     fn next(&mut self) -> Option<Self::Item> {
-        let (next_point, _) = step_chaos_game(
+        let (next_point, next_color, _) = step_chaos_game(
             self.current_point,
+            self.current_color,
             self.rng,
             self.transforms,
             self.weights,
             self.variations,
         );
         self.current_point = next_point;
+        self.current_color = next_color;
 
-        Some(next_point)
+        Some((next_point, next_color))
-    }
-}
-
-/// Shader entry point for running the chaos game to produce new IFS coordinates
-pub mod entry {
-    use crate::chaos_game::ChaosGame;
-    use crate::rng::xoshiro256starstar_from_seed;
-    use crate::transform::Transform;
-    use crate::variation::Variation;
-    use glam::Vec2;
-    use spirv_std::spirv;
-
-    /// Given a set of fractal flame parameters, generate new IFS coordinates
-    /// and store them in the output array.
-    #[spirv(compute(entry_point_name = "main_chaos_game", threads(1)))]
-    pub fn main_chaos_game(
-        #[spirv(spec_constant(id = 1, default = 20))] iteration_discard: u32,
-        #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] rng_seed: &[u8],
-        #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] transforms: &[Transform],
-        #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] weights: &[f32],
-        #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] variations: &[Variation],
-        #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] output: &mut [Vec2],
-    ) {
-        let mut rng_seed_actual = [0u8; 32];
-        (0..32).for_each(|i| rng_seed_actual[i] = rng_seed[i]);
-
-        let mut rng = xoshiro256starstar_from_seed(rng_seed_actual);
-        let mut chaos_game = ChaosGame::new(&mut rng, transforms, weights, variations);
-
-        for _ in 0..iteration_discard {
-            chaos_game.next().unwrap();
-        }
-
-        for i in 0..output.len() {
-            output[i] = chaos_game.next().unwrap();
-        }
     }
 }

enkou-shaders/src/lib.rs

@@ -1,6 +1,6 @@
 //! # Enkou
 #![no_std]
-#![deny(missing_docs)]
+#![warn(missing_docs)]
 #![allow(clippy::needless_range_loop)] // SPIR-V backend has issues with iteration over items
 
 pub mod camera;

enkou-shaders/src/transform.rs

@@ -5,7 +5,7 @@
 //! but produce more interesting images once we add variations.
 use crate::variation::Variation;
 use bytemuck::{Pod, Zeroable};
-use glam::{Affine2, UVec2, Vec2};
+use glam::{Affine2, FloatExt, UVec2, Vec2};
 use rand::Rng;
 
 /// Affine transform for use in the [`chaos_game`](crate::chaos_game).
@@ -15,15 +15,22 @@ pub struct Transform {
     coefficients: Affine2,
     coefficients_post: Affine2,
     variation_range: UVec2,
+    color: Vec2,
 }
 
 impl Transform {
     /// Create a new transform from an affine transformation matrix
-    pub fn new(coefficients: Affine2, coefficients_post: Affine2, variation_range: UVec2) -> Self {
+    pub fn new(
+        coefficients: Affine2,
+        coefficients_post: Affine2,
+        variation_range: UVec2,
+        color: Vec2,
+    ) -> Self {
         Transform {
             coefficients,
             coefficients_post,
             variation_range,
+            color,
         }
     }
 
@@ -47,6 +54,11 @@ impl Transform {
 
         self.coefficients_post.transform_point2(point)
     }
+
+    /// Apply this transform to a color
+    pub fn transform_color(&self, color: f32) -> f32 {
+        color.lerp(self.color.x, self.color.y)
+    }
 }
 
 #[cfg(test)]
@@ -54,7 +66,7 @@ mod test {
     use crate::rng::xoshiro256starstar_from_seed;
     use crate::transform::Transform;
     use crate::variation::{Variation, VariationKind};
-    use glam::{Affine2, uvec2, vec2};
+    use glam::{Affine2, uvec2, vec2, Vec2};
 
     #[test]
     fn transform_scaling() {
@@ -63,6 +75,7 @@ mod test {
             Affine2::from_scale(scale_coefficients),
             Affine2::IDENTITY,
             uvec2(0, 1),
+            Vec2::ZERO,
         );
 
         let mut rng = xoshiro256starstar_from_seed([0; 32]);
@@ -78,11 +91,12 @@ mod test {
     #[test]
     fn transform_scaling_post() {
         let scale_coefficients = vec2(2.0, 0.5);
-        let transform_pdj = Transform::new(Affine2::IDENTITY, Affine2::IDENTITY, uvec2(0, 1));
+        let transform_pdj = Transform::new(Affine2::IDENTITY, Affine2::IDENTITY, uvec2(0, 1), Vec2::ZERO);
         let transform_pdj_post = Transform::new(
             Affine2::IDENTITY,
             Affine2::from_scale(scale_coefficients),
             uvec2(0, 1),
+            Vec2::ZERO,
         );
 
         let mut rng = xoshiro256starstar_from_seed([0; 32]);
@@ -94,4 +108,23 @@ mod test {
 
         assert_eq!(point_pdj * scale_coefficients, point_pdj_post);
     }
+    
+    #[test]
+    fn transform_color() {
+        // Color 0.5, color speed 1.0, so color value will always be 0.5 after transform
+        let color = vec2(0.5, 1.0);
+        let transform = Transform::new(Affine2::IDENTITY, Affine2::IDENTITY, uvec2(0, 1), color);
+
+        assert_eq!(transform.transform_color(0.0), 0.5);
+        assert_eq!(transform.transform_color(1.0), 0.5);
+        assert_eq!(transform.transform_color(2.0), 0.5);
+        
+        // Color 1.0, color speed 0.5, so color value moves to halfway between current and 1.0
+        let color = vec2(1.0, 0.5);
+        let transform = Transform::new(Affine2::IDENTITY, Affine2::IDENTITY, uvec2(0, 1), color);
+        
+        assert_eq!(transform.transform_color(0.0), 0.5);
+        assert_eq!(transform.transform_color(1.0), 1.0);
+        assert_eq!(transform.transform_color(2.0), 1.0);
+    }
 }