particle interaction simulation with grid optimization in lobster

grid_based_particle_interaction.lobster

import std
import vec
import color
import gl

// --- System Parameters (Global Constants) ---
let SIM_WIDTH = 400.0
let SIM_HEIGHT = 400.0
let PARTICLE_COUNT = 2000
let INTERACTION_RADIUS = 15.0
let ALPHA_PARAM_DEG = 180.0
let BETA_PARAM_DEG = 17.0
let VELOCITY_PARAM = 0.7
let DT_PARAM = 1.01

// --- Colors ---
let COLOR_GREEN = float4 { 0.0, 1.0, 0.0, 1.0 }
let COLOR_MAGENTA = float4 { 1.0, 0.0, 1.0, 1.0 }
let COLOR_BROWN = float4 { 165.0/255.0, 42.0/255.0, 42.0/255.0, 1.0 }
let COLOR_BLUE = float4 { 0.0, 0.0, 1.0, 1.0 }
let COLOR_YELLOW = float4 { 1.0, 1.0, 0.0, 1.0 }
let COLOR_BLACK = color_black
let COLOR_WHITE = color_white

// --- Font ---
let FONT_PATH = pakfile "/home/izot/.local/share/fonts/InputMono-Regular.ttf"
let FONT_SIZE = 15

// --- Particle Data Structure ---
class ParticleData:
    x:float
    y:float
    heading_deg:float
    current_color:float4

// --- Particle Functions ---
def create_particle_data(px:float, py:float) -> ParticleData:
    return ParticleData {
        x: px,
        y: py,
        heading_deg: rnd_float() * 360.0,
        current_color: COLOR_GREEN
    }

def draw_particle_data(p:ParticleData):
    gl.color(p.current_color)
    gl.translate(float2{p.x, p.y}):
        gl.circle(3.0, 6)

def get_particle_neighbors(p_self:ParticleData, potential_neighbors:[ParticleData]) -> int, int, int:
    var left_count = 0
    var right_count = 0
    var total_count = 0

    for(potential_neighbors) p_other:
        if p_other != p_self:
            var dx = p_other.x - p_self.x
            var dy = p_other.y - p_self.y

            // Toroidal wrapping
            if dx > SIM_WIDTH / 2.0: dx -= SIM_WIDTH
            elif dx < -SIM_WIDTH / 2.0: dx += SIM_WIDTH
            if dy > SIM_HEIGHT / 2.0: dy -= SIM_HEIGHT
            elif dy < -SIM_HEIGHT / 2.0: dy += SIM_HEIGHT

            let distance_sq = dx * dx + dy * dy
            if distance_sq <= INTERACTION_RADIUS * INTERACTION_RADIUS and distance_sq > 0.0001:
                total_count++
                let angle_to_other_deg = atan2(float2{dx, dy})
                var relative_angle_deg = angle_to_other_deg - p_self.heading_deg
                while relative_angle_deg <= -180.0: relative_angle_deg += 360.0
                while relative_angle_deg > 180.0: relative_angle_deg -= 360.0

                if relative_angle_deg > 0.001: right_count++
                elif relative_angle_deg < -0.001: left_count++
    return left_count, right_count, total_count

def update_particle_data(p:ParticleData, left_count:int, right_count:int, total_count:int):
    var turn_direction = 0.0
    if left_count > right_count: turn_direction = 1.0
    elif right_count > left_count: turn_direction = -1.0

    let delta_heading_deg = (ALPHA_PARAM_DEG * turn_direction + BETA_PARAM_DEG * float(total_count))
    p.heading_deg = (p.heading_deg + delta_heading_deg * DT_PARAM)
    while p.heading_deg < 0.0: p.heading_deg += 360.0
    p.heading_deg = p.heading_deg % 360.0

    p.x += VELOCITY_PARAM * cos(p.heading_deg) * DT_PARAM
    p.y += VELOCITY_PARAM * sin(p.heading_deg) * DT_PARAM

    p.x = (p.x % SIM_WIDTH + SIM_WIDTH) % SIM_WIDTH
    p.y = (p.y % SIM_HEIGHT + SIM_HEIGHT) % SIM_HEIGHT

    if 15 < total_count and total_count <= 35: p.current_color = COLOR_BLUE
    elif total_count > 35: p.current_color = COLOR_YELLOW
    elif 13 <= total_count and total_count <= 15: p.current_color = COLOR_BROWN
    elif total_count > 15: p.current_color = COLOR_MAGENTA
    else: p.current_color = COLOR_GREEN


// --- Grid Functions ---
var GRID_CELLS: [[[ParticleData]]] = []
var GRID_WIDTH_CELLS = 0
var GRID_HEIGHT_CELLS = 0
var GRID_CELL_SIZE = 0.0

def create_grid_structure(cell_s:float):
    GRID_CELL_SIZE = cell_s
    GRID_WIDTH_CELLS = ceiling(SIM_WIDTH / GRID_CELL_SIZE)
    GRID_HEIGHT_CELLS = ceiling(SIM_HEIGHT / GRID_CELL_SIZE)

    GRID_CELLS = []
    for(GRID_HEIGHT_CELLS) y_idx:
        var row:[[ParticleData]] = []
        for(GRID_WIDTH_CELLS) x_idx:
            row.push([])
        GRID_CELLS.push(row)

def insert_particles_into_grid_structure(all_particles:[ParticleData]):
    for(GRID_HEIGHT_CELLS) y_idx:
        for(GRID_WIDTH_CELLS) x_idx:
            GRID_CELLS[y_idx][x_idx] = []

    for(all_particles) p:
        let cell_x = min(int(p.x / GRID_CELL_SIZE), GRID_WIDTH_CELLS - 1)
        let cell_y = min(int(p.y / GRID_CELL_SIZE), GRID_HEIGHT_CELLS - 1)
        let final_cell_x = max(0, cell_x)
        let final_cell_y = max(0, cell_y)

        GRID_CELLS[final_cell_y][final_cell_x].push(p)


def get_particles_in_grid_range(p_self:ParticleData) -> [ParticleData]:
    var nearby_particles:[ParticleData] = []
    let self_cell_x = int(p_self.x / GRID_CELL_SIZE)
    let self_cell_y = int(p_self.y / GRID_CELL_SIZE)

    for(3) dy_offset_idx: // -1, 0, 1
        for(3) dx_offset_idx: // -1, 0, 1
            let dx = dx_offset_idx - 1
            let dy = dy_offset_idx - 1
            let neighbor_cell_x = (self_cell_x + dx + GRID_WIDTH_CELLS) % GRID_WIDTH_CELLS
            let neighbor_cell_y = (self_cell_y + dy + GRID_HEIGHT_CELLS) % GRID_HEIGHT_CELLS
            append_into(nearby_particles, GRID_CELLS[neighbor_cell_y][neighbor_cell_x])
            
    return nearby_particles

def main():
    fatal(gl.window("Lobster PPS with Grid", int(SIM_WIDTH), int(SIM_HEIGHT)))
    gl.set_target_delta_time(1.0 / 60.0)
    check(gl.set_font_name(FONT_PATH) and gl.set_font_size(FONT_SIZE), "can\'t load font!")

    var particles:[ParticleData] = []
    for(PARTICLE_COUNT):
        particles.push(create_particle_data(rnd_float() * SIM_WIDTH, rnd_float() * SIM_HEIGHT))

    let cell_s = INTERACTION_RADIUS * 2.0
    create_grid_structure(cell_s)

    while gl.frame() and gl.button("escape") != 1:
        insert_particles_into_grid_structure(particles)

        for(particles) p:
            let potential_neighbors = get_particles_in_grid_range(p)
            let left, right, total = get_particle_neighbors(p, potential_neighbors)
            update_particle_data(p, left, right, total)

        gl.clear(COLOR_BLACK)

        for(particles) p:
            draw_particle_data(p)

        gl.color(COLOR_WHITE)
        gl.translate(float2{5, 5}):
            let fps = if gl.delta_time() > 0.0: 1.0 / gl.delta_time() else: 0.0
            gl.text("FPS: {int(fps)}")
        gl.translate(float2{5, 15}):
             gl.text("Particles: {particles.length}")

main()