Skip to main content

KERNEL_MATH_FUNCTION_GPU

Constant KERNEL_MATH_FUNCTION_GPU 

Source
pub const KERNEL_MATH_FUNCTION_GPU: &str = r#"
// GPU math function kernel
// Computes one step of various mathematical functions.
// Mode uniform selects the function.

layout(local_size_x = 256) in;

struct Point {
    vec4 pos_age;    // xyz = position, w = iteration count or age
    vec4 vel_param;  // xyz = velocity/derivative, w = parameter
};

layout(std430, binding = 0) buffer Points {
    Point points[];
};

uniform uint u_point_count;
uniform uint u_function_type;  // 0=Lorenz, 1=Mandelbrot, 2=Julia, 3=Rossler, 4=Aizawa
uniform float u_dt;
uniform float u_time;
uniform float u_param_a;
uniform float u_param_b;
uniform float u_param_c;
uniform float u_param_d;
uniform uint u_max_iterations;
uniform vec2 u_julia_c;  // Julia set constant

// Lorenz attractor: dx/dt = sigma*(y-x), dy/dt = x*(rho-z)-y, dz/dt = x*y - beta*z
vec3 lorenz(vec3 p, float sigma, float rho, float beta) {
    return vec3(
        sigma * (p.y - p.x),
        p.x * (rho - p.z) - p.y,
        p.x * p.y - beta * p.z
    );
}

// Rossler attractor: dx/dt = -y-z, dy/dt = x+a*y, dz/dt = b+z*(x-c)
vec3 rossler(vec3 p, float a, float b, float c) {
    return vec3(
        -p.y - p.z,
        p.x + a * p.y,
        b + p.z * (p.x - c)
    );
}

// Aizawa attractor
vec3 aizawa(vec3 p, float a, float b, float c, float d) {
    float x = p.x, y = p.y, z = p.z;
    return vec3(
        (z - b) * x - d * y,
        d * x + (z - b) * y,
        c + a * z - z * z * z / 3.0 - (x * x + y * y) * (1.0 + 0.25 * z) + 0.1 * z * x * x * x
    );
}

// Complex multiply
vec2 cmul(vec2 a, vec2 b) {
    return vec2(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
}

void main() {
    uint idx = gl_GlobalInvocationID.x;
    if (idx >= u_point_count) return;

    Point pt = points[idx];
    vec3 pos = pt.pos_age.xyz;
    float age = pt.pos_age.w;

    if (u_function_type == 0u) {
        // Lorenz attractor (RK4 integration)
        float sigma = u_param_a;  // default 10.0
        float rho = u_param_b;    // default 28.0
        float beta = u_param_c;   // default 8.0/3.0

        vec3 k1 = lorenz(pos, sigma, rho, beta);
        vec3 k2 = lorenz(pos + 0.5 * u_dt * k1, sigma, rho, beta);
        vec3 k3 = lorenz(pos + 0.5 * u_dt * k2, sigma, rho, beta);
        vec3 k4 = lorenz(pos + u_dt * k3, sigma, rho, beta);

        pos += (u_dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4);
        pt.vel_param.xyz = k1;  // Store derivative for visualization
        age += u_dt;

    } else if (u_function_type == 1u) {
        // Mandelbrot iteration
        // pos.xy = current z, vel_param.xy = c (constant)
        vec2 z = pos.xy;
        vec2 c = pt.vel_param.xy;
        uint iter = uint(age);

        if (iter < u_max_iterations && dot(z, z) < 4.0) {
            z = cmul(z, z) + c;
            pos.xy = z;
            pos.z = dot(z, z);  // magnitude squared for coloring
            age = float(iter + 1u);
        }

    } else if (u_function_type == 2u) {
        // Julia set iteration
        vec2 z = pos.xy;
        vec2 c = u_julia_c;
        uint iter = uint(age);

        if (iter < u_max_iterations && dot(z, z) < 4.0) {
            z = cmul(z, z) + c;
            pos.xy = z;
            pos.z = dot(z, z);
            age = float(iter + 1u);
        }

    } else if (u_function_type == 3u) {
        // Rossler attractor (RK4)
        float a = u_param_a;  // default 0.2
        float b = u_param_b;  // default 0.2
        float c = u_param_c;  // default 5.7

        vec3 k1 = rossler(pos, a, b, c);
        vec3 k2 = rossler(pos + 0.5 * u_dt * k1, a, b, c);
        vec3 k3 = rossler(pos + 0.5 * u_dt * k2, a, b, c);
        vec3 k4 = rossler(pos + u_dt * k3, a, b, c);

        pos += (u_dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4);
        pt.vel_param.xyz = k1;
        age += u_dt;

    } else if (u_function_type == 4u) {
        // Aizawa attractor (RK4)
        float a = u_param_a;  // default 0.95
        float b = u_param_b;  // default 0.7
        float c = u_param_c;  // default 0.6
        float d = u_param_d;  // default 3.5

        vec3 k1 = aizawa(pos, a, b, c, d);
        vec3 k2 = aizawa(pos + 0.5 * u_dt * k1, a, b, c, d);
        vec3 k3 = aizawa(pos + 0.5 * u_dt * k2, a, b, c, d);
        vec3 k4 = aizawa(pos + u_dt * k3, a, b, c, d);

        pos += (u_dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4);
        pt.vel_param.xyz = k1;
        age += u_dt;
    }

    points[idx].pos_age = vec4(pos, age);
    points[idx].vel_param = pt.vel_param;
}
"#;
Expand description

Math function GPU kernel: Lorenz, Mandelbrot, Julia, Rossler, Aizawa.