Hacky 2d vector field visualization using compute shaders

3 files changed, 286 insertions, 0 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
new file mode 100644
index 0000000..a16dced
--- /dev/null
+++ b/CMakeLists.txt
@@ -0,0 +1,19 @@
+cmake_minimum_required(VERSION 3.10)
+project(shader)
+
+set(
+	CMAKE_CXX_FLAGS
+	"${CMAKE_CXX_FLAGS} -std=c++17 -W -Wall -Wextra -Winline -pedantic"
+)
+
+add_executable(
+	shader
+		main.cc
+)
+
+target_link_libraries(
+	shader
+		GL
+		GLEW
+		glfw
+)
diff --git a/main.cc b/main.cc
new file mode 100644
index 0000000..870006e
--- /dev/null
+++ b/main.cc
@@ -0,0 +1,255 @@
+#include <GL/glew.h>
+#include <GLFW/glfw3.h>
+
+#include <glm/glm.hpp>
+#include <glm/gtc/matrix_transform.hpp>
+
+#include <vector>
+#include <iostream>
+#include <random>
+
+GLFWwindow* window;
+
+GLint getUniform(GLuint program, const std::string& name) {
+	const GLint uniform = glGetUniformLocation(program, name.c_str());
+	if ( uniform == -1 ) {
+		std::cerr << "Could not bind uniform " << name << std::endl;
+	}
+	return uniform;
+}
+
+GLint compileShader(const std::string& source, GLenum type) {
+	GLint local_shader = glCreateShader(type);
+
+	if ( !local_shader ) {
+		std::cerr << "Cannot create a shader of type " << type << std::endl;
+		exit(-1);
+	}
+
+	const char* source_data = source.c_str();
+	const int source_length = source.size();
+
+	glShaderSource(local_shader, 1, &source_data, &source_length);
+	glCompileShader(local_shader);
+
+	GLint compiled;
+	glGetShaderiv(local_shader, GL_COMPILE_STATUS, &compiled);
+	if ( !compiled ) {
+		std::cerr << "Cannot compile shader" << std::endl;
+		GLint maxLength = 0;
+		glGetShaderiv(local_shader, GL_INFO_LOG_LENGTH, &maxLength);
+		std::vector<GLchar> errorLog(maxLength);
+		glGetShaderInfoLog(local_shader, maxLength, &maxLength, &errorLog[0]);
+		for( auto c : errorLog ) {
+			std::cerr << c;
+		}
+		std::cerr << std::endl;
+	}
+
+	return local_shader;
+}
+
+GLint setupShader() {
+	#define GLSL(shader) #shader
+	GLint shader = glCreateProgram();
+
+	glAttachShader(shader, compileShader(GLSL(
+		uniform mat4 MVP;
+
+		float distance(vec2 v, vec2 w) {
+			vec2 u = v - w;
+			return sqrt(u.x*u.x + u.y*u.y);
+		}
+
+		void main() {
+			gl_Position = MVP * vec4(gl_Vertex.xy, 0.0, 1.0);
+			gl_FrontColor = vec4(max(1. - gl_Vertex.z/5., 0.1), 0., 0., 0.);
+		}),
+		GL_VERTEX_SHADER));
+
+	glAttachShader(shader, compileShader(GLSL(
+		void main() {
+			gl_FragColor = gl_Color;
+		}),
+		GL_FRAGMENT_SHADER));
+
+	glLinkProgram(shader);
+	return shader;
+}
+
+GLint setupComputeShader() {
+	GLint shader = glCreateProgram();
+	glAttachShader(shader, compileShader("#version 430\n" GLSL(
+		layout (local_size_x = 1) in;
+		layout (std430, binding=1) buffer bufferA{ float data[]; };
+		uniform vec2 world;
+
+		float rand(vec2 co){
+			return fract(sin(dot(co.xy, vec2(12.9898,78.233))) * 43758.5453);
+		}
+
+		bool insideWorld(vec2 v) {
+			return v.x > -world.x/2.
+			    && v.x <  world.x/2.
+				&& v.y > -world.y/2.
+				&& v.y <  world.y/2.;
+		}
+
+		void main() {
+			uint idx = 3*gl_GlobalInvocationID.x;
+			vec2 v = vec2(data[idx+0], data[idx+1]);
+
+			if ( data[idx+2] < 5. && insideWorld(v) ) {
+				data[idx+0] += 0.01 * cos(v.x*cos(v.y));
+				data[idx+1] += 0.01 * sin(v.x-v.y);
+				data[idx+2] += 0.01;
+			} else {
+				data[idx+0] = -(world.x/2.) + rand(vec2(data[idx+1], data[idx+0])) * world.x;
+				data[idx+1] = -(world.y/2.) + rand(vec2(data[idx+0], data[idx+1])) * world.y;
+				data[idx+2] = uint(rand(v) * 5.);
+			}
+		}),
+		GL_COMPUTE_SHADER));
+	glLinkProgram(shader);
+	return shader;
+}
+
+int window_width  = 800;
+int window_height = 600;
+double mouse_x, mouse_y;
+float world_width, world_height;
+glm::mat4 MVP;
+
+void updateMVP() {
+	world_width  = 20.f;
+	world_height = world_width / window_width * window_height;
+	glm::mat4 projection = glm::ortho(
+		-(world_width /2), world_width/2,
+		-(world_height/2), world_height/2,
+		0.1f, 100.0f
+	);
+	glm::mat4 view = glm::lookAt(
+		glm::vec3(0,0,20),
+		glm::vec3(0,0,0),
+		glm::vec3(0,1,0)
+	);
+	glm::mat4 model = glm::mat4(1.0f);
+	MVP = projection * view * model;
+}
+
+void window_size_callback(GLFWwindow*, int width, int height) {
+	window_width  = width;
+	window_height = height;
+	glViewport(0, 0, width, height);
+	updateMVP();
+}
+
+void cursor_pos_callback(GLFWwindow*, double x, double y) {
+	mouse_x = -(world_width/2) + world_width * (x / window_width);
+	mouse_y = world_height/2 - world_height * (y / window_height);
+}
+
+int main() {
+	if( !glfwInit() ) {
+		std::cerr <<  "Failed to initialize GLFW" << std::endl;
+		return -1;
+	}
+
+	glfwWindowHint(GLFW_SAMPLES, 4);
+	glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
+	glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
+	glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
+
+	window = glfwCreateWindow(window_width, window_height, "shader", NULL, NULL);
+	if( window == NULL ){
+		std::cerr << "Failed to open GLFW window." << std::endl;
+		glfwTerminate();
+		return -1;
+	}
+	glfwSetWindowSizeCallback(window, window_size_callback);
+	glfwSetCursorPosCallback(window, cursor_pos_callback);
+	glfwMakeContextCurrent(window);
+
+	glewExperimental = true;
+	if ( glewInit() != GLEW_OK ) {
+		std::cerr << "Failed to initialize GLEW" << std::endl;
+		glfwTerminate();
+		return -1;
+	}
+
+	glfwSetInputMode(window, GLFW_STICKY_KEYS, GL_TRUE);
+
+	glClearColor(0.0f, 0.0f, 0.4f, 0.0f);
+
+	GLuint VertexArrayID;
+	glGenVertexArrays(1, &VertexArrayID);
+	glBindVertexArray(VertexArrayID);
+
+	GLint ProgramID = setupShader();
+	GLuint MatrixID = glGetUniformLocation(ProgramID, "MVP");
+	GLuint MouseID  = glGetUniformLocation(ProgramID, "mouse");
+
+	GLint ComputeProgramID = setupComputeShader();
+	GLuint WorldID = glGetUniformLocation(ComputeProgramID, "world");
+
+	updateMVP();
+
+	const std::size_t particle_count = 100000; 
+
+	std::vector<GLfloat> vertex_buffer_data;
+	vertex_buffer_data.reserve(3*particle_count);
+
+	std::random_device rd;
+	std::mt19937 gen(rd());
+	std::uniform_real_distribution<GLfloat> distX(-world_width/2.f, world_width/2.f);
+	std::uniform_real_distribution<GLfloat> distY(-world_height/2.f, world_height/2.f);
+	std::uniform_real_distribution<GLfloat> distAge(0.f, 5.f);
+
+	for ( int i = 0; i < particle_count; ++i ) {
+		vertex_buffer_data.emplace_back(distX(gen));
+		vertex_buffer_data.emplace_back(distY(gen));
+		vertex_buffer_data.emplace_back(distAge(gen));
+	}
+
+	GLuint vertexbuffer;
+	glGenBuffers(1, &vertexbuffer);
+	glBindBuffer(GL_ARRAY_BUFFER, vertexbuffer);
+	glBufferData(GL_ARRAY_BUFFER,
+		vertex_buffer_data.size() * sizeof(GLfloat), &vertex_buffer_data[0], GL_STATIC_DRAW);
+
+	do {
+		glUseProgram(ComputeProgramID);
+		glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, vertexbuffer);
+		glUniform2f(WorldID, world_width, world_height);
+		glDispatchCompute(particle_count, 1, 1);
+
+		glClear(GL_COLOR_BUFFER_BIT);
+
+		glUseProgram(ProgramID);
+
+		glUniformMatrix4fv(MatrixID, 1, GL_FALSE, &MVP[0][0]);
+		glUniform2f(MouseID, mouse_x, mouse_y);
+
+		glEnableVertexAttribArray(0);
+		glBindBuffer(GL_ARRAY_BUFFER, vertexbuffer);
+		glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 0, (void*)0);
+
+		glDrawArrays(GL_POINTS, 0, 3*particle_count);
+
+		glDisableVertexAttribArray(0);
+
+		glfwSwapBuffers(window);
+		glfwPollEvents();
+	}
+	while( glfwGetKey(window, GLFW_KEY_ESCAPE ) != GLFW_PRESS &&
+		   glfwWindowShouldClose(window) == 0 );
+
+	glDeleteBuffers(1, &vertexbuffer);
+	glDeleteVertexArrays(1, &VertexArrayID);
+	glDeleteProgram(ProgramID);
+	glDeleteProgram(ComputeProgramID);
+
+	glfwTerminate();
+
+	return 0;
+}
diff --git a/shell.nix b/shell.nix
new file mode 100644
index 0000000..2a16377
--- /dev/null
+++ b/shell.nix
@@ -0,0 +1,12 @@
+with import <nixpkgs> {};
+
+stdenv.mkDerivation rec {
+  name = "env";
+  env = buildEnv { name = name; paths = buildInputs; };
+  buildInputs = [
+    git cmake gcc gdb cgdb
+    glew
+    glfw3
+    glm
+  ];
+}