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.
This commit is contained in:
2026-07-12 21:12:01 -04:00
parent 7ff19631ba
commit fd49ae7256
4 changed files with 75 additions and 81 deletions
+19 -33
View File
@@ -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::entry::main_chaos_game;
use enkou_shaders::chaos_game::ChaosGame;
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: u32 = 50_000;
const ITERATIONS_DISCARD: usize = 20;
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
.iter()
.skip_while(|p| {
p.x < 0 || (p.x as u32) > dimensions.0 || p.y < 0 || (p.y as u32) > dimensions.1
})
.map(|p| (p.x as u32, p.y as u32))
.for_each(|(x, y)| image.put_pixel(x, y, Luma([255u8])));
let chaos_game = ChaosGame::new(&mut rng, &transforms, &weights, &variations);
chaos_game
.skip(ITERATIONS_DISCARD)
.take(ITERATIONS)
.filter_map(|(point_ifs, _)| camera.transform_point_to_image(point_ifs))
.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")?;
+18 -43
View File
@@ -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)
}
}
/// 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();
}
Some((next_point, next_color))
}
}
+1 -1
View File
@@ -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;
+37 -4
View File
@@ -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);
}
}