The LÖVR implementation of physarum model in a compute shader

physarum.lua

-- https://cargocollective.com/sagejenson/physarum
-- also [Coding Adventure: Ant and Slime Simulations](https://www.youtube.com/watch?v=X-iSQQgOd1A)

-- "agents" move around and leave behind a trail
-- they also detect others' trail and steer towards it
local m = {}

local agent_count = 1e4
local iterations_per_frame = 3

m.gpu = {
  resolution = 1024,     -- size of rectangular texture

  speed_move = 25,      -- forward speed
  speed_turn = 15,      -- maximum steering speed, actual speed is randomized
  time_delta = 1/20,    -- time advancment per iteration

  diffuse_speed = 5.0,    -- amount of trails spreading out
  fadeout_speed = 0.1,    -- dissipation of trail with time

  sensor_distance = 25.,   -- how far ahead agents detect trail
  sensor_span_angle = 0.4, -- left/right sensing angle in rad

  physarum_color =  Vec3(lovr.math.gammaToLinear(0.86, 0.79, 0.33)),
}

local trails_texture = lovr.graphics.newTexture(m.gpu.resolution, m.gpu.resolution,
  {format='rgba32f', mipmaps = false, usage = {'render', 'sample', 'storage', 'transfer'}, linear = false})


local agent_shader = lovr.graphics.newShader([[
layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;

struct Agent {
  vec2 position;
  float angle;
};

layout(std140) buffer agents_buffer {
    Agent agents[];
};

uniform float resolution;
uniform float speed_move;
uniform float speed_turn;
uniform float time_delta;
uniform float sensor_distance;
uniform float sensor_span_angle;

uniform float fadeout_speed;
uniform float diffuse_speed;
uniform vec3 physarum_color;

layout(rgba32f) uniform image2D trails_texture;


uint rnd(uint seed) {
    uint state = seed;
    state = state ^ 2747636419u;
    state = state * 2654435769u;
    state = state ^ state >> 16u;
    state = state * 2654435769u;
    state = state ^ state >> 16u;
    state = state * 2654435769u;
    return state;
}

float sense(vec2 pos, float angle) {
  ivec2 uv = ivec2(pos + sensor_distance * vec2(cos(angle), sin(angle)));
  uv = uv % ivec2(resolution);
  float s = imageLoad(trails_texture, uv).r;
  return s;
}

void lovrmain() {
  ivec3 coord = ivec3(gl_GlobalInvocationID.xyz);

  Agent agent = agents[coord.x];
  uint random = rnd(uint(agent.position.y * resolution + agent.position.x + rnd(coord.x)));
  float randomF = random / 4294967295.;

  float weight_forward = sense(agent.position, agent.angle +  0.);
  float weight_left    = sense(agent.position, agent.angle +  sensor_span_angle);
  float weight_right   = sense(agent.position, agent.angle + -sensor_span_angle);
  if (weight_forward > weight_left && weight_forward > weight_right) {
  } else if (weight_forward < weight_left && weight_forward < weight_right) {
    agent.angle += (randomF - 0.5) * 2 * speed_turn * time_delta;
  } else if (weight_right > weight_left) {
    agent.angle -= randomF * speed_turn * time_delta;
  } else if (weight_right < weight_left) {
    agent.angle += randomF * speed_turn * time_delta;
  }

  vec2 direction = vec2(cos(agent.angle), sin(agent.angle));
  agent.position += direction * speed_move * time_delta;

  agents[coord.x].position = mod(agent.position, vec2(resolution, resolution));
  agents[coord.x].angle = mod(agent.angle, 2 * PI);
  imageStore(trails_texture, ivec2(agent.position), vec4(physarum_color, 1.));
}
]])

local diffuse_shader = lovr.graphics.newShader([[
layout(local_size_x = 32, local_size_y = 32, local_size_z = 1) in;
layout(rgba32f) uniform image2D trails_texture;

uniform float resolution;
uniform float time_delta;
uniform float fadeout_speed;
uniform float diffuse_speed;

void lovrmain() {
  ivec3 coord = ivec3(gl_GlobalInvocationID.xyz);
  if (any(greaterThanEqual(coord, ivec3(resolution))))
    return;

  vec4 sum = vec4(0.0);
  float total_weight = 0.0;

  for (int x = -1; x <= 1; x++) {
    for (int y = -1; y <= 1; y++) {
      ivec2 sampleCoord = (coord.xy + ivec2(x, y)) % ivec2(resolution);
      float weight = 1.0 - sqrt(float(x * x + y * y));
      sum += imageLoad(trails_texture, sampleCoord) * weight;
      total_weight += weight;
    }
  }

  vec4 avg = sum / total_weight;
  vec4 prev = imageLoad(trails_texture, coord.xy);
  vec4 res = mix(prev, avg, diffuse_speed * fadeout_speed * time_delta);
  res -= fadeout_speed * time_delta;
  res = max(vec4(0.), res);
  imageStore(trails_texture, coord.xy, res);
}
]])

m.physarum_pass = lovr.graphics.newPass()
local agent_group_size = agent_shader:getWorkgroupSize()
local diffuse_group_size =  diffuse_shader:getWorkgroupSize()
lovr.graphics.setBackgroundColor(0.002, 0.05, 0.08)

local agents = lovr.graphics.newBuffer(agent_shader:getBufferFormat('agents_buffer'), agent_count)
for i = 1, agents:getLength() do
  agents:setData({{
      position = vec2(
        lovr.math.random() * m.gpu.resolution,
        lovr.math.random() * m.gpu.resolution),
      angle = (-0.5 + math.random()) * 2 * math.pi,
    }}, i)
end


function lovr.update(dt)
  m.physarum_pass:reset()
  m.physarum_pass:setShader(agent_shader)
  m.physarum_pass:send('agents_buffer', agents)
  m.physarum_pass:send('trails_texture', trails_texture)
  for name, value in pairs(m.gpu) do
    m.physarum_pass:send(name, value)
  end
  for i=1, iterations_per_frame do
    m.physarum_pass:compute(agents:getLength() / agent_group_size)
  end
  m.physarum_pass:setShader(diffuse_shader)
  -- we don't need to resend, the uniforms are already bound to the pass
  m.physarum_pass:compute(m.gpu.resolution / diffuse_group_size, m.gpu.resolution / diffuse_group_size)
  lovr.graphics.submit(m.physarum_pass)
end


function lovr.draw(pass)
  pass:setColor(1,1,1)
  pass:setMaterial(trails_texture)
  pass:plane(0, 2, -4, 10)
end

--[[ automatic UI from m.gpu table; code available [here](https://github.com/jmiskovic/grassland)

local panel = require'ui/fromTable'.integrate(m.gpu,  mat4():target(vec3(-8, 0, -4), vec3(0, 1, 8)):rotate(math.pi, 0,1,0):scale(0.5))
--]]

physarum.mp4