Slime simulation

This is a compute shader that simulates thousands or millions of agents and observe emergent behaviours as you change their parameters. Largely inspired by Sebastian Lague's video on the subject you can watch here.

The folder containing the project can be downloaded and run from here or the demo code at the bottom of the page (limited to only one colour palette and no shortcuts for starting patterns) can be run provided you have the necessary python extensions.

import moderngl
import numpy as np
from math import *
np.set_printoptions(threshold=np.inf)
import struct
import glfw
glfw.init()
window = glfw.create_window(1000,1000, "Slime Simulation", None, None)
glfw.make_context_current(window)
ctx = moderngl.create_context()
# draw a fullscreen quad to the framebuffer to write the texture to the buffer
program = ctx.program(vertex_shader='''
#version 430

in vec2 in_vert;
out vec2 pos;

void main() {
    pos = in_vert * 0.5 + 0.5;
    gl_Position = vec4(in_vert, 0.0, 1.0);
}
''', fragment_shader='''
#version 430

in vec2 pos;
out vec4 out_color;

uniform bool to_screen;
uniform bool left_click;
uniform bool right_click;
uniform float radius;
uniform vec2 Mpos;
uniform vec2 res;
uniform float fade;
uniform float blur;
uniform sampler2D u_texture;

void main() {
    vec2 lum = texture(u_texture, pos).xy;
    if (to_screen){
        out_color = vec4(0,0.1,0.2,1) + vec4(sin(lum.y*3),0,cos(lum.y*9),1)*lum.x + vec4(0.7,0.8,0.8,1)*lum.x;
    }
    else{
        lum *= (1/blur);
        lum += texture(u_texture, pos + vec2(1,0)/res.x).xy;
        lum += texture(u_texture, pos + vec2(0,1)/res.y).xy;
        lum += texture(u_texture, pos + vec2(-1,0)/res.x).xy;
        lum += texture(u_texture, pos + vec2(0,-1)/res.y).xy;
        lum /= 4+1/blur;
        if (left_click && distance(Mpos,pos)<radius){
            out_color.x = 40;
        }
        else if (right_click && distance(Mpos,pos)<radius){
            out_color.x = -10;
        }
        else if (lum.x==1.1){
            out_color.x = 1;
            }
        else{
            out_color.xy = lum*(1.0-fade);
            }
    }
}
''')
vbo = ctx.buffer(struct.pack('8f', -1.0, -1.0, -1.0, 1.0, 1.0, -1.0, 1.0, 1.0))
vao = ctx.vertex_array(program, vbo, 'in_vert')

compute_shader = ctx.compute_shader("""
#version 430

layout (local_size_x = 64, local_size_y = 1) in;
layout(std430,binding = 0) buffer agent_buffer{ vec3 agents[];};
layout(std430,binding = 1) buffer pixel_buffer{ vec2 pixels[1000][1000];};
layout(binding=2) uniform sampler2D u_texture;

#define TAU 3.1415926538*2
uniform float turn_rate;
uniform float sensor_angle;
uniform float sensor_dist;
uniform float randomness;
uniform float stepsize;
uniform uvec2 resolution;
uniform uint frame;

float rand(vec2 co){
    return fract(sin(dot(co, vec2(12.9898, 78.233))) * 43758.5453)-0.5;
}

void main() {
    uint ID = gl_GlobalInvocationID.x;
    vec3 agent = agents[ID];
    // Update agent position and velocity
    if (randomness != 0){
        agent.z += rand(vec2(ID,frame))/(1/randomness*100)*TAU;
        }
    vec2 sensor_a_pos;
    sensor_a_pos.x = (sin((agent.z+sensor_angle)*TAU)*sensor_dist+agent.x*resolution.x);
    sensor_a_pos.y = (cos((agent.z+sensor_angle)*TAU)*sensor_dist+agent.y*resolution.y);
    float sensor_a = pixels[int(sensor_a_pos.x)][int(sensor_a_pos.y)].x;
    vec2 sensor_b_pos;
    sensor_b_pos.x = (sin((agent.z-sensor_angle)*TAU)*sensor_dist+agent.x*resolution.x);
    sensor_b_pos.y = (cos((agent.z-sensor_angle)*TAU)*sensor_dist+agent.y*resolution.y);
    float sensor_b = pixels[int(sensor_b_pos.x)][int(sensor_b_pos.y)].x;
    agent.z -= (sensor_b-sensor_a)*turn_rate;
    
    vec2 step = vec2(sin(agent.z*TAU),cos(agent.z*TAU))*stepsize/1000;
    agent.xy += step;
    if (agent.x >= 1||agent.x <= 0){
        //agent.z = 1 - agent.z;
        agent.x -= 1*sign(agent.x);
    }
    if (agent.y >= 1||agent.y <= 0){
        agent.y -= 1*sign(agent.y);
    }
    // Store the updated particle data
    agents[ID] = agent;
    pixels[int(agent.x*resolution.x)][int(agent.y*resolution.y)].x = 1.1;
    pixels[int(agent.x*resolution.x)][int(agent.y*resolution.y)].y = agent.z;
}
"""
)
#Initialises window
glfw.set_time(0)
program["u_texture"] = 2
resolution = glfw.get_window_size(window)
compute_shader["resolution"] = resolution
program["res"] = resolution
#------------------------------
compute_shader["randomness"] = 0.1 #Controls randomness of agents
sensor_dist = 50 #Generally controls scale of paths formed
turn_rate = 0.015 #Controls smoothness of paths formed
sensor_angle = 0.1#Sensor angle in number of turns away from direction
num_agents = 65535 #max is 65535 on laptops it is the max 16bit number
#Sets fade and blur amount
program["blur"] = 1
fade = 0.015
stepsize = 2
program["radius"] = 0.013 #Radius of mouse pixel influence
#------------------------------

#initialises buffer of agents positions and velocities to be read by the compute shader
agent_data = np.random.rand(num_agents, 3).astype(dtype=np.float32)
#for agent in agent_data:
 #   agent[0] = cos(agent[2]*2*pi)/2+0.5
  #  agent[1] = sin(agent[2]*2*pi)/2+0.5

agent_buffer = ctx.buffer(agent_data.tobytes())
agent_buffer.bind_to_storage_buffer(0)
  
#Sets a buffer of all pixels to be modified by the compute shader  
pixel_data = np.zeros((resolution[0],resolution[1],2),np.float32)
pixel_buffer = ctx.buffer(pixel_data.tobytes())
pixel_buffer.bind_to_storage_buffer(1)
#Creates a texture for the pixels
pixel_texture = ctx.texture(resolution, 2, dtype = "f4")
pixel_texture.use(2)

#Sets the framebuffer as the render target and renders to the last_frame texture
last_frame = ctx.texture(resolution, 2, dtype = "f4")
fbo = ctx.framebuffer(last_frame)
fbo.use()

frame = 0
time = 0
while not glfw.window_should_close(window):
    delttime = glfw.get_time() - time
    print(round(delttime*1000,3))
    time = glfw.get_time()
    glfw.poll_events()
    resolution = glfw.get_window_size(window)
    #Get cursor
    x,y = glfw.get_cursor_pos(window)
    program["Mpos"] = [max(0, min(x/resolution[0],1)),max(0, min(1-y/resolution[1],1))]
    program["left_click"] = glfw.get_mouse_button(window, 0)
    program["right_click"] = glfw.get_mouse_button(window, 1)
    
    if glfw.get_key(window,glfw.KEY_MINUS):
        turn_rate *= 0.99
        print("Turn_rate:",round(turn_rate,5))
    if glfw.get_key(window,glfw.KEY_EQUAL):
        turn_rate *= 1.01
        print("Turn_rate:",round(turn_rate,5))
    compute_shader["turn_rate"] = turn_rate
    
    if glfw.get_key(window,glfw.KEY_UP):
        fade *= 0.99
        print("Fade:",round(fade,5))
    if glfw.get_key(window,glfw.KEY_DOWN):
        fade *= 1.01
        print("Fade:",round(fade,5))
    program["fade"] = fade
    if glfw.get_key(window,glfw.KEY_LEFT):
        sensor_dist *= 0.99
        print("Sensor_dist:",round(sensor_dist,2))
    if glfw.get_key(window,glfw.KEY_RIGHT):
        sensor_dist *= 1.01
        print("Sensor_dist:",round(sensor_dist,2))
    compute_shader["sensor_dist"] = sensor_dist

    if glfw.get_key(window,glfw.KEY_COMMA):
        sensor_angle *= 0.99
        print("Sensor_angle:",round(sensor_angle,5))
    if glfw.get_key(window,glfw.KEY_PERIOD):
        sensor_angle *= 1.01
        print("Sensor_angle:",round(sensor_angle,5))
    compute_shader["sensor_angle"] = sensor_angle
    

    compute_shader["stepsize"] = stepsize#*(sin(frame/30)+1.5)#*delttime*10
    compute_shader["frame"] = frame
    ctx.viewport = (0,0) + resolution
    #Dispatch compute shader
    compute_shader.run(int(num_agents/64)+1,1,1)
    pixel_texture.write(pixel_buffer)#Write the pixel positions to a texture


    #Renders to framebuffer
    fbo.use()
    program["to_screen"] = False
    vao.render(mode=moderngl.TRIANGLE_STRIP)
    #Sets the pixel buffer to the dissapated and blurred texture
    pixel_buffer.write(last_frame.read())
    #Renders to screen
    ctx.screen.use()
    program["to_screen"] = True
    vao.render(mode=moderngl.TRIANGLE_STRIP)
    glfw.swap_buffers(window)
    frame += 1
glfw.terminate()