shader.comp

// Source code by Shawn Halayka
// Source code is in the public domain
//
// See:
// https://www.cs.cmu.edu/~kmcrane/Projects/QuaternionJulia/QJuliaFragment.html
// http://antongerdelan.net/opengl/compute.html
// https://github.com/daw42/glslcookbook/tree/master/chapter10


// OpenGL 4.3 introduces compute shaders
#version 430 core

// Use a local workgroup size of 1, which is the default size
// Altering the local workgroup size requires a proportional
// reduction in the values passed to glDispatchCompute in main.h
layout(local_size_x = 1, local_size_y = 1) in;

// Single-channel output
layout(binding = 0, r32f) writeonly uniform image2D output_image;

// Four-channel input (because each quaternion has 4 member variables; xyzw)
layout(binding = 1, rgba32f) readonly uniform image2D input_image;

// Quaternion Julia set constants
uniform vec4 C; // Treat quaternions as 4D vectors
uniform int max_iterations;
uniform float threshold;

// Perform the quaternion square operation
vec4 qsquare(vec4 q)
{
	vec4 qout;

	qout.x = q.x * q.x - dot(q.yzw, q.yzw);

	// The cross product of a 3D vector and itself is (0, 0, 0), and so
	// the yzw values are greatly simplified
	qout.yzw = 2 * q.x * q.yzw;

	return qout;
}

// Calculate magnitude of quaternion while underdoing iteration
float iterate(vec4 z)
{
	const float threshold_sq = threshold * threshold;
	float magnitude_sq = dot(z, z);

	for (int i = 0; i < max_iterations; i++)
	{
		// Calculate the traditional quaternion Julia set
		z = qsquare(z) + C;

		// Assign value to magnitude_sq, and then check its value
		// compared to threshold_sq
		// 
		// If the magnitude is greater than or equal to the threshold
		// value, then this point on the grid is considered divergent, 
		// or not in the quaternion Julia set
		if ((magnitude_sq = dot(z, z)) >= threshold_sq)
			break;
	}

	// Return magnitude
	return sqrt(magnitude_sq);
}

void main()
{
	// Get global coordinates
	const ivec2 pixel_coords = ivec2(gl_GlobalInvocationID.xy);

	// Use input image as input quaternion
	vec4 z = imageLoad(input_image, pixel_coords);

	// Test to see whether or not the input quaternion is in the Julia set
	const float magnitude = iterate(z);

	// Assign only the first element, since we're writing to a 
	// single-channel image
	const vec4 output_pixel = vec4(magnitude, 0, 0, 0);

	// Store result in output image
	imageStore(output_image, pixel_coords, output_pixel);
}