enkou/enkou-shaders/src/transform.rs

//! # Transform
//!
//! Transforms are the "functions" in an iterated function system. They take in a point,
//! and generate a new point. For fractal flames, transforms are always affine,
//! but produce more interesting images once we add variations.
use crate::variation::Variation;
use bytemuck::{Pod, Zeroable};
use glam::{Affine2, FloatExt, UVec2, Vec2};
use rand::Rng;

/// Affine transform for use in the [`chaos_game`](crate::chaos_game).
#[derive(Copy, Clone, Pod, Zeroable)]
#[repr(C)]
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,
        color: Vec2,
    ) -> Self {
        Transform {
            coefficients,
            coefficients_post,
            variation_range,
            color,
        }
    }

    /// Apply this transform to a point in IFS coordinates, producing a new point
    pub fn transform_point<R: Rng>(
        &self,
        rng: &mut R,
        variations: &[Variation],
        point: Vec2,
    ) -> Vec2 {
        let point = self.coefficients.transform_point2(point);

        let mut point_output = Vec2::ZERO;

        let variation_start = self.variation_range.x;
        let variation_end = self.variation_range.y;
        for variation_index in variation_start..variation_end {
            let variation = &variations[variation_index as usize];
            point_output += variation.transform_point(point, rng, &self.coefficients)
        }

        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)]
mod test {
    use crate::rng::xoshiro256starstar_from_seed;
    use crate::transform::Transform;
    use crate::variation::{Variation, VariationKind};
    use glam::{Affine2, uvec2, vec2, Vec2};

    #[test]
    fn transform_scaling() {
        let scale_coefficients = vec2(2.0, 0.5);
        let transform = Transform::new(
            Affine2::from_scale(scale_coefficients),
            Affine2::IDENTITY,
            uvec2(0, 1),
            Vec2::ZERO,
        );

        let mut rng = xoshiro256starstar_from_seed([0; 32]);
        let variations = [Variation::IDENTITY];
        let point = vec2(1.0, 1.0);

        assert_eq!(
            transform.transform_point(&mut rng, &variations, point),
            scale_coefficients
        );
    }

    #[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), 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]);
        let variations = [Variation::new(VariationKind::Pdj, 1.0, [0.0f32; 4].into())];
        let point = vec2(1.0, 1.0);

        let point_pdj = transform_pdj.transform_point(&mut rng, &variations, point);
        let point_pdj_post = transform_pdj_post.transform_point(&mut rng, &variations, point);

        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);
    }
}